def sparse_xcipy_logdet(A, beta=0):
xp, xcipy, xsplin = get_array_module(1)
if not isinstance(
A,
(xp.ndarray,
xcipy.sparse.spmatrix)): # cupy is installed but A is a cpu array
xp, xcipy, xsplin = get_array_module(0)
betaI = beta * xcipy.sparse.eye(A.shape[-1], dtype=A.dtype)
slu = xsplin.splu(A + betaI)
ld = xp.log(slu.U.diagonal()).sum()
return ld.item()
def test_cheby_op_basis(dtype, device):
N = 8
A = torch.rand(N, N, dtype=dtype, device=device)
A = A + A.t()
A.fill_diagonal_(0)
Ats = SparseTensor.from_dense(A)
xp, xcipy, xsplin = get_array_module(Ats.is_cuda())
x = torch.ones(N, 1, dtype=dtype, device=device)
xxp = xp.asarray(x)
from thgsp.graphs.laplace import laplace as ts_laplace
Lts = ts_laplace(Ats, "sym")
krn = np.array([[[meyer_kernel], [meyer_mirror_kernel]]
]) # 1(n_graph) x 2(Cout) x 1(Cin) kernel
coeff = cheby_coeff(krn, K=6, device=device,
dtype=dtype) # M x Co x Ci x K+1
yts = cheby_op(x, Lts, coeff[0])
yts0, yts1 = yts[..., 0]
c0, c1 = coeff[0].squeeze_()
H0 = cheby_op_basis(Lts, c0, return_ts=False)
H1 = cheby_op_basis(Lts, c1, return_ts=True)
assert H1.dtype() == Lts.dtype()
assert H1.device() == Lts.device()
y0 = H0 @ xxp
y1 = H1 @ x
err = 1e-5 if dtype is float_dtypes[0] else 1e-12
assert np.abs(yts0 - y0.ravel()).max() < err
assert torch.abs(yts1 - y1.ravel()).max() < err
def test_construct_dia(dt, dv, layout, inverse, narrow):
on_gpu = "cuda" == dv.type
xp, xcipy, _ = get_array_module(on_gpu)
N = 7
S = [5, 1, 3]
M = len(S)
diag_data = xp.arange(N)
P = construct_dia(S,
diag_data=diag_data,
ps=narrow,
inverse=inverse,
dtype=dt,
layout=layout)
if narrow:
assert P.shape == (M, M)
if inverse:
assert xp.allclose(
P.diagonal(),
1 / diag_data[S],
)
else:
assert xp.allclose(P.diagonal(), diag_data[S])
else:
assert P.shape == (N, N)
if inverse:
assert xp.allclose(P.diagonal(), 1 / diag_data)
else:
assert xp.allclose(P.diagonal(), diag_data)
def test_construct_sampling_matrix(dtype, dv, layout):
on_gpu = "cuda" == dv.type
xp, xcipy, _ = get_array_module(on_gpu)
S = [5, 4, 1]
N = 6
H = construct_sampling_matrix(N, S, dtype, dv, layout)
assert H.format == layout
assert H.shape == (len(S), N)
D = xcipy.sparse.diags(xp.arange(1, N + 1)[S])
HtHD0 = H.T @ D @ H
HtH0 = H.T @ H
print("\n", H.A)
print(HtH0.A)
print(HtHD0.A)
HtH1 = construct_hth(N, S, dtype=dtype, device=dv, layout=layout)
HtHD1 = construct_hth(N,
S,
xp.arange(1, N + 1)[S],
dtype=dtype,
device=dv,
layout=layout)
# print(HtH1.A)
# print(HtHD1.A)
print(type(HtH1), HtH1.shape, "on_gpu: ", on_gpu)
print(type(HtHD1), HtHD1.shape, "on_gpu: ", on_gpu)
assert np.allclose(HtH0.A, HtH1.A)
assert np.allclose(HtHD0.A, HtHD1.A)
def recon_ess(y, S, U, bd, **kwargs):
"""Naive implementation of ESS sampling reconstruction.
Parameters
----------
y: Tensor
Dense Shape: :obj:`(N)`
S: List
The sampling set
U: Tensor
Dense :obj:`(N, bd)`
bd: int
The bandwidth of target signal
kwargs: dict
The optional arguments of `xp.linalg.lstsq`
Returns
-------
f_hat: Tensor
The reconstructed signal
"""
assert bd > 1
assert len(S) > 1
dv = U.device
xp, _, _ = get_array_module(U.is_cuda)
tmp = xp.linalg.lstsq(xp.asarray(U[S, :bd]),
xp.asarray(y),
**kwargs,
rcond=None)[0]
f_hat = U[:, :bd] @ torch.as_tensor(tmp, device=dv)
return f_hat
def recon_bsgda(y,
S,
L: SparseTensor,
mu: float = 0.01,
reg_order: int = 1,
**kwargs):
N = L.size(-1)
M = len(S)
assert N >= M > 0
if y.ndim == 1:
y = y.view(-1, 1)
if y.shape[0] != M:
raise RuntimeError(
f"y is expected to have a shape ({M},num_signal) or ({M},), not {y.shape}"
)
dt, dv, density, on_gpu = get_ddd(L)
xp, xcipy, xsplin = get_array_module(on_gpu)
L = to_xcipy(L)
Ht = construct_sampling_matrix(N, S, dtype=dt, device=dv, layout="csr").T
HtH = construct_hth(N, S, dtype=dt, device=dv, layout="csr")
B = HtH + mu * L**reg_order
Hty = Ht @ xp.asarray(y)
if xp == np:
x_hat = xsplin.spsolve(B, Hty, **kwargs)
else:
num_sig = Hty.shape[-1]
x_hat = xp.empty((N, num_sig), dtype=Hty.dtype)
for j in range(num_sig):
x_hat[:, j] = xsplin.spsolve(B, Hty[:, j], **kwargs)
return torch.as_tensor(x_hat, device=dv)
def cheby_op_basis(L, coeff, lam_max=2.0, return_ts=False):
assert coeff.ndim == 1
K = len(coeff)
if isinstance(L, SparseTensor):
dt, dv, density, on_gpu = get_ddd(L)
xp, xcipy, _ = get_array_module(on_gpu)
elif isinstance(L, spmatrix):
import scipy
xcipy = scipy
xp = np
else:
raise TypeError
L = to_xcipy(L)
N = L.shape[-1]
coeff = xp.asarray(coeff)
Im = xcipy.sparse.eye(N, dtype=L.dtype, format="csr")
Ln = L * (2 / lam_max) - Im
Tl_old = Im
Tl_cur = Ln
Hl = 0.5 * coeff[0] * Tl_old + coeff[1] * Tl_cur
for k in range(2, K):
Tl_new = 2 * Ln * Tl_cur - Tl_old
Hl = Hl + coeff[k] * Tl_new
Tl_old = Tl_cur
Tl_cur = Tl_new
if return_ts:
result = SparseTensor.from_scipy(Hl) if xp == np else from_cpx(Hl)
else:
result = Hl
return result
def test_get_array_module(on_gpu):
xp, xcipy, _ = get_array_module(on_gpu=on_gpu)
if on_gpu:
try:
import cupy as cp
import cupyx.scipy as xscipy
assert xp == cp
assert xscipy == xscipy
except ImportError:
pytest.skip("CuPy is not installed, use numpy and scipy instead")
else:
assert xp == np
assert xcipy == scipy
def recon_rsbs(y,
S,
L: SparseTensor,
cum_coh,
mu: float = 0.01,
reg_order: int = 1,
**kwargs):
N = L.size(-1)
M = len(S)
assert M > 0
if y.shape[0] != M:
raise RuntimeError(
f"y is expected to have a shape ({M},num_signal) or ({M},), not {y.shape}"
)
dt, dv, density, on_gpu = get_ddd(L)
xp, xcipy, xsplin = get_array_module(on_gpu)
yp = to_xp(y).astype("d")
L = to_xcipy(L, layout="csr").astype("d")
H = construct_sampling_matrix(N, S, torch.double, dv)
Psinv = construct_dia(S,
cum_coh,
ps=True,
inverse=True,
dtype=torch.double,
device=dv)
Bl = H.T * Psinv
B = Bl * H + mu * L**reg_order
HPy = Bl @ yp
if xp == np:
x_hat = xsplin.spsolve(B, HPy, **kwargs)
else:
num_sig = yp.shape[-1]
x_hat = xp.empty((N, num_sig), dtype=yp.dtype)
for j in range(num_sig):
x_hat[:, j] = xsplin.spsolve(B, HPy[:, j], **kwargs)
return torch.as_tensor(x_hat, device=dv)
def recon_fastssss(y, S, T, order, sd=0.5):
"""A primary implementation of reconstruction method associated with
"FastSSS" sampling algorithm.
Parameters
----------
y: Tensor
The measurements on sampling set :obj:S:. If the localization operator :obj:`T`
has a density greater than the threshold :obj:`sd`, :obj:`y` has a shape of
either :obj:`(M,)`，:obj:`(M,1)`，or :obj:`(M,C)`; otherwise :obj:`y` could only
be either :obj:`(M,)` or :obj:`(M,1)`.
S: List
A list consisting of all indices of sampled nodes.
T: Tensor, SparseTensor
The localization operator
order: int
sd: float
The threshold of :obj:`T`'s density that controls when we use a dense or sparse
linear solver.
Returns
-------
Tensor
The signal recovered from the measurements :obj:`y` .
"""
T_k = matrix_power(T, order)
dt, dv, density, on_gpu = get_ddd(T_k)
if density > sd:
T_k = T_k.to_dense()
tmp = torch.linalg.solve(T_k[np.ix_(S, S)], torch.as_tensor(y))
x_hat = T_k[:, S] @ tmp
else:
T_k = to_xcipy(T_k)
xp, xcipy, xsplin = get_array_module(on_gpu)
tmp = xsplin.spsolve(T_k[xp.ix_(S, S)], xp.asarray(y))
x_hat = T_k[:, S] * tmp
return torch.as_tensor(x_hat, device=dv)