def __init__(self, bipartite_graphs, beta, k=4, lam_max=2.0, device=None):
h0_c, g0_c, self.orthogonality = design_biorth_kernel(k)
h0 = partial(polyval, torch.from_numpy(h0_c))
h0.__name__ = "h0"
g0 = partial(polyval, torch.from_numpy(g0_c))
g0.__name__ = "g0"
def h1(x):
return g0(2 - x)
def g1(x):
return h0(2 - x)
self.analysis_krn = np.array([[[h0], [h1]]])
self.synthesis_krn = np.array([[[g0], [g1]]])
ana_coeff = cheby_coeff(
self.analysis_krn, K=2 * k, lam_max=lam_max
).squeeze_() # (1,2,1,K) --> (2,K)
syn_coeff = cheby_coeff(
self.synthesis_krn, K=2 * k, lam_max=lam_max
).squeeze_() # (1,2,1,K) --> (2,K)
operator = QmfOperator.compute_basis(bipartite_graphs, ana_coeff, beta, lam_max)
self.operator = SparseTensor.from_scipy(operator).to(device)
inv_operator = QmfOperator.compute_basis(
bipartite_graphs, syn_coeff, beta, lam_max
)
self.inv_operator = SparseTensor.from_scipy(inv_operator).to(device)
def init_adj(self, edge_index):
""" cache normalized adjacency and normalized strict two-hop adjacency,
neither has self loops
"""
n = self.num_nodes
if isinstance(edge_index, SparseTensor):
dev = adj_t.device
adj_t = edge_index
adj_t = scipy.sparse.csr_matrix(adj_t.to_scipy())
adj_t[adj_t > 0] = 1
adj_t[adj_t < 0] = 0
adj_t = SparseTensor.from_scipy(adj_t).to(dev)
elif isinstance(edge_index, torch.Tensor):
row, col = edge_index
adj_t = SparseTensor(row=col, col=row, value=None, sparse_sizes=(n, n))
adj_t.remove_diag(0)
adj_t2 = matmul(adj_t, adj_t)
adj_t2.remove_diag(0)
adj_t = scipy.sparse.csr_matrix(adj_t.to_scipy())
adj_t2 = scipy.sparse.csr_matrix(adj_t2.to_scipy())
adj_t2 = adj_t2 - adj_t
adj_t2[adj_t2 > 0] = 1
adj_t2[adj_t2 < 0] = 0
adj_t = SparseTensor.from_scipy(adj_t)
adj_t2 = SparseTensor.from_scipy(adj_t2)
adj_t = gcn_norm(adj_t, None, n, add_self_loops=False)
adj_t2 = gcn_norm(adj_t2, None, n, add_self_loops=False)
self.adj_t = adj_t.to(edge_index.device)
self.adj_t2 = adj_t2.to(edge_index.device)
def __init__(
self,
G: Graph,
k: int = 8,
in_channels: int = 1,
order: int = 16,
strategy: str = "admm",
level: int = 1,
lam_max: float = 2.0,
zero_dc: bool = False,
**kwargs,
):
self.adj = G
N = self.adj.size(-1)
self.lam_max = lam_max
self.strategy = strategy
self.num_bgraph = level
device = self.adj.device()
dtype = self.adj.dtype()
if strategy == "admm":
if N < 80:
bipartite_graphs_dense = admm_bga(
self.adj.to_dense().to(torch.double), M=level, **kwargs
)
beta = bipartite_graphs_dense.new_zeros(N, level).to(torch.bool)
bipartite_graphs = []
bipartite_graphs_dense = bipartite_graphs_dense.to(dtype).to(device)
for i, B in enumerate(bipartite_graphs_dense):
_, vtx_color, _ = is_bipartite_fix(B, fix_flag=True)
beta[:, i] = torch.as_tensor(vtx_color)
bipartite_graphs.append(SparseTensor.from_dense(B))
else:
bipartite_graphs, beta, self.partptr, self.perm = admm_lbga_ray(
self.adj, level, **kwargs
)
bipartite_graphs = [
SparseTensor.from_scipy(B).to(dtype).to(device)
for B in bipartite_graphs
]
elif strategy == "amfs":
bipartite_graphs, beta = amfs(self.adj, level=self.num_bgraph, **kwargs)
bipartite_graphs = [
SparseTensor.from_scipy(B).to(dtype).to(device)
for B in bipartite_graphs
]
else:
raise RuntimeError(
f"{strategy} is not a valid numerical decomposition "
f"algorithm supported at present."
)
super().__init__(
bipartite_graphs, beta, k, in_channels, order, lam_max, zero_dc
)
def prune_adj(oriadj, non_zero_idx: int, percent: int):
original_prune_num = int(
((non_zero_idx - oriadj.size()[0]) / 2) * (percent / 100))
adj = SparseTensor.from_torch_sparse_coo_tensor(oriadj).to_scipy()
# find the lower half of the matrix
low_adj = tril(adj, -1)
non_zero_low_adj = low_adj.data[low_adj.data != 0]
low_pcen = np.percentile(abs(non_zero_low_adj), percent)
under_threshold = abs(low_adj.data) < low_pcen
before = len(non_zero_low_adj)
low_adj.data[under_threshold] = 0
non_zero_low_adj = low_adj.data[low_adj.data != 0]
after = len(non_zero_low_adj)
rest_pruned = original_prune_num - (before - after)
if rest_pruned > 0:
mask_low_adj = (low_adj.data != 0)
low_adj.data[low_adj.data == 0] = 2000000
flat_indices = np.argpartition(low_adj.data,
rest_pruned - 1)[:rest_pruned]
low_adj.data = np.multiply(low_adj.data, mask_low_adj)
low_adj.data[flat_indices] = 0
low_adj.eliminate_zeros()
new_adj = low_adj + low_adj.transpose()
new_adj = new_adj + sparse.eye(new_adj.shape[0])
return SparseTensor.from_scipy(new_adj).to_torch_sparse_coo_tensor().to(
device)
def process(self):
import pandas as pd
x = sp.load_npz(self.raw_paths[0])
if x.shape[-1] > 10000 or self.name == 'mag':
x = SparseTensor.from_scipy(x).to(torch.float)
else:
x = torch.from_numpy(x.todense()).to(torch.float)
df = pd.read_csv(self.raw_paths[1],
header=None,
sep=None,
engine='python')
edge_index = torch.from_numpy(df.values).t().contiguous()
with open(self.raw_paths[2], 'r') as f:
ys = f.read().split('\n')[:-1]
ys = [[int(y) - 1 for y in row.split()[1:]] for row in ys]
multilabel = max([len(y) for y in ys]) > 1
if not multilabel:
y = torch.tensor(ys).view(-1)
else:
num_classes = max([y for row in ys for y in row]) + 1
y = torch.zeros((len(ys), num_classes), dtype=torch.float)
for i, row in enumerate(ys):
for j in row:
y[i, j] = 1.
data = Data(x=x, edge_index=edge_index, y=y)
data = data if self.pre_transform is None else self.pre_transform(data)
torch.save(self.collate([data]), self.processed_paths[0])
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 get_bga(adj, strategy, num_bgraph, **kwargs):
num_node = adj.size(-1)
dv = adj.device()
if strategy == "admm":
if num_node < 80:
bipartite_graphs_dense = admm_bga(
adj.to_dense().to(torch.double), M=num_bgraph, **kwargs
)
beta = torch.zeros(num_node, num_bgraph, dtype=torch.bool, device=dv)
bipartite_graphs = []
for i, B in enumerate(bipartite_graphs_dense):
_, vtx_color, _ = is_bipartite_fix(B, fix_flag=True)
beta[:, i] = torch.as_tensor(vtx_color)
bipartite_graphs.append(SparseTensor.from_dense(B))
return bipartite_graphs, beta
else:
bipartite_graphs, beta, partptr, perm = admm_lbga_ray(
adj, num_bgraph, **kwargs
)
elif strategy == "amfs":
bipartite_graphs, beta = amfs(adj, level=num_bgraph, **kwargs)
else:
raise RuntimeError(
f"{str(strategy)} is not a valid numerical decomposition algorithm "
f"supported at present."
)
bipartite_graphs = [SparseTensor.from_scipy(B).to(dv) for B in bipartite_graphs]
return bipartite_graphs, beta
def rand_bipartite(N1, N2, p=0.2, dtype=None, device=None, return_partition=False):
G = nx.bipartite.random_graph(N1, N2, p)
csr_sci_adj = nx.adjacency_matrix(G)
A = SparseTensor.from_scipy(csr_sci_adj).to(device, dtype)
bpg = Graph(A)
if return_partition:
beta = [1 if d["bipartite"] == 1 else 0 for n, d in G.nodes(data=True)]
beta = torch.as_tensor(beta, dtype=torch.bool)
return bpg, beta
return bpg
def __collate__(self, batch):
# Call parent's collate function
data = super().__collate__(batch)
# Get the node indexes from batch
if not isinstance(batch, torch.Tensor):
batch = torch.tensor(batch)
start = self.cluster_data.partptr[batch].tolist()
end = self.cluster_data.partptr[batch + 1].tolist()
node_idx = torch.cat([torch.arange(s, e) for s, e in zip(start, end)])
# Convert edge_index to sparse adjacency matrix
row, col = data.edge_index
edge_attr = np.ones(row.size(0), dtype=np.float32)
adj_matrix = coo_matrix((edge_attr, (row, col)),
(data.num_nodes, data.num_nodes))
adj_matrix_tensor = SparseTensor.from_scipy(adj_matrix)
# adj_matrix_tensor = SparseTensor(row=row, col = col, value = torch.ones(row.size(0), dtype=torch.float32), sparse_sizes = (data.num_nodes, data.num_nodes))
# Normalize the adjacency matrix
if self.normalize_adj_matrix == True:
adj_matrix = aug_normalized_adjacency(adj_matrix)
# Attach the adjacency matrix to the data
data['adj_matrix'] = adj_matrix
data['adj_t'] = adj_matrix_tensor
# If the node features is a zero tensor with dimension one
# re-create it here as a (num_nodes, num_nodes) sparse identity matrix
if data.num_node_features == 1 and torch.equal(
data['x'], torch.zeros(data.num_nodes,
data.num_node_features)):
node_features = sp.identity(data.num_nodes)
node_features = sparse_mx_to_torch_sparse_tensor(
node_features).float()
data['x'] = node_features
# If the node features is a zero tensor with dimension one
# re-create it here as a (num_nodes, num_nodes) sparse identity matrix
if data.num_node_features == 1 and torch.equal(
data['x'], torch.zeros(data.num_nodes,
data.num_node_features)):
node_features = sp.identity(data.num_nodes)
node_features = sparse_mx_to_torch_sparse_tensor(
node_features).float()
data['x'] = node_features
# Return the enhanced data
return data
def __init__(self, bipartite_graphs, beta, order=24, lam_max=2.0, device=None):
N, M = beta.shape
assert len(bipartite_graphs) == M
self.num_node, self.num_bgraph = N, M
self.order = order
self.device = device
self.lam_max = lam_max
# 1(n_graph) x 2(Cout) x 1(Cin) kernel
krn = np.array([[[meyer_kernel], [meyer_mirror_kernel]]])
# (1,2,1,K) --> (2,K)
coeff = cheby_coeff(krn, K=order, lam_max=lam_max).squeeze_()
operator = self.compute_basis(bipartite_graphs, coeff, beta, lam_max)
self.operator = SparseTensor.from_scipy(operator).to(device)
self.dtype = self.operator.dtype()
self.device = self.operator.device()
def to_torch_sparse(mat):
if isinstance(mat, torch.Tensor):
if not mat.is_sparse:
stm = SparseTensor.from_dense(mat)
else:
stm = SparseTensor.from_torch_sparse_coo_tensor(mat)
elif isinstance(mat, np.ndarray):
stm = SparseTensor.from_dense(torch.as_tensor(mat))
elif isinstance(mat, spmatrix):
stm = SparseTensor.from_scipy(mat)
elif isinstance(mat, SparseTensor):
stm = mat
else:
raise TypeError(f"{type(mat)} not supported now")
return stm
def get_bga(self, adj, strategy, vtx_color, **kwargs):
dv = adj.device()
if strategy == "harary":
bipartite_graphs, beta, beta_dist, new_vtx_color, mapper = harary(
adj, vtx_color=vtx_color, **kwargs
)
elif strategy == "osglm":
bipartite_graphs, beta, append_nodes, new_vtx_color = osglm(
adj, vtx_color=vtx_color, **kwargs
)
self.append_nodes = append_nodes # noqa
else:
raise RuntimeError(
f"{strategy} is not a valid color-based decomposition algorithm."
)
self.vtx_color = new_vtx_color # noqa
self.num_node = adj.size(-1) # noqa
self.strategy = strategy # noqa
self.dtype = adj.dtype() # noqa
bipartite_graphs = [SparseTensor.from_scipy(B).to(dv) for B in bipartite_graphs]
return bipartite_graphs, beta
import torch
from torch_sparse import SparseTensor
import scipy.sparse as sp
import torch_sparse_utils as tsu
n = 10
b = 5
a = sp.rand(n, n, density=0.5, format='csr')
a[a > 0] = 1
adj = SparseTensor.from_scipy(a).type_as(torch.FloatTensor())
batch_nodes = torch.randperm(n)[:b]
num_neighbors = 10
num_proc = 2
sa = tsu.sample_neighbors(adj,
batch_nodes,
num_neighbors,
num_proc,
replace=False)
def tree_decomposition(mol, return_vocab=False):
r"""The tree decomposition algorithm of molecules from the
`"Junction Tree Variational Autoencoder for Molecular Graph Generation"
<https://arxiv.org/abs/1802.04364>`_ paper.
Returns the graph connectivity of the junction tree, the assignment
mapping of each atom to the clique in the junction tree, and the number
of cliques.
Args:
mol (rdkit.Chem.Mol): A :obj:`rdkit` molecule.
return_vocab (bool, optional): If set to :obj:`True`, will return an
identifier for each clique (ring, bond, bridged compounds, single).
(default: :obj:`False`)
:rtype: (LongTensor, LongTensor, int)
"""
import rdkit.Chem as Chem
# Cliques = rings and bonds.
cliques = [list(x) for x in Chem.GetSymmSSSR(mol)]
xs = [0] * len(cliques)
for bond in mol.GetBonds():
if not bond.IsInRing():
cliques.append([bond.GetBeginAtomIdx(), bond.GetEndAtomIdx()])
xs.append(1)
# Generate `atom2clique` mappings.
atom2clique = [[] for i in range(mol.GetNumAtoms())]
for c in range(len(cliques)):
for atom in cliques[c]:
atom2clique[atom].append(c)
# Merge rings that share more than 2 atoms as they form bridged compounds.
for c1 in range(len(cliques)):
for atom in cliques[c1]:
for c2 in atom2clique[atom]:
if c1 >= c2 or len(cliques[c1]) <= 2 or len(cliques[c2]) <= 2:
continue
if len(set(cliques[c1]) & set(cliques[c2])) > 2:
cliques[c1] = set(cliques[c1]) | set(cliques[c2])
xs[c1] = 2
cliques[c2] = []
xs[c2] = -1
cliques = [c for c in cliques if len(c) > 0]
xs = [x for x in xs if x >= 0]
# Update `atom2clique` mappings.
atom2clique = [[] for i in range(mol.GetNumAtoms())]
for c in range(len(cliques)):
for atom in cliques[c]:
atom2clique[atom].append(c)
# Add singleton cliques in case there are more than 2 intersecting
# cliques. We further compute the "initial" clique graph.
edges = {}
for atom in range(mol.GetNumAtoms()):
cs = atom2clique[atom]
if len(cs) <= 1:
continue
# Number of bond clusters that the atom lies in.
bonds = [c for c in cs if len(cliques[c]) == 2]
# Number of ring clusters that the atom lies in.
rings = [c for c in cs if len(cliques[c]) > 4]
if len(bonds) > 2 or (len(bonds) == 2 and len(cs) > 2):
cliques.append([atom])
xs.append(3)
c2 = len(cliques) - 1
for c1 in cs:
edges[(c1, c2)] = 1
elif len(rings) > 2:
cliques.append([atom])
xs.append(3)
c2 = len(cliques) - 1
for c1 in cs:
edges[(c1, c2)] = 99
else:
for i in range(len(cs)):
for j in range(i + 1, len(cs)):
c1, c2 = cs[i], cs[j]
count = len(set(cliques[c1]) & set(cliques[c2]))
edges[(c1, c2)] = min(count, edges.get((c1, c2), 99))
# Update `atom2clique` mappings.
atom2clique = [[] for i in range(mol.GetNumAtoms())]
for c in range(len(cliques)):
for atom in cliques[c]:
atom2clique[atom].append(c)
if len(edges) > 0:
edge_index_T, weight = zip(*edges.items())
row, col = torch.tensor(edge_index_T).t()
inv_weight = 100 - torch.tensor(weight)
clique_graph = SparseTensor(row=row,
col=col,
value=inv_weight,
sparse_sizes=(len(cliques), len(cliques)))
junc_tree = minimum_spanning_tree(clique_graph.to_scipy('csr'))
row, col, _ = SparseTensor.from_scipy(junc_tree).coo()
edge_index = torch.stack([row, col], dim=0)
edge_index = to_undirected(edge_index, num_nodes=len(cliques))
else:
edge_index = torch.empty((2, 0), dtype=torch.long)
rows = [[i] * len(atom2clique[i]) for i in range(mol.GetNumAtoms())]
row = torch.tensor(list(chain.from_iterable(rows)))
col = torch.tensor(list(chain.from_iterable(atom2clique)))
atom2clique = torch.stack([row, col], dim=0).to(torch.long)
if return_vocab:
vocab = torch.tensor(xs, dtype=torch.long)
return edge_index, atom2clique, len(cliques), vocab
else:
return edge_index, atom2clique, len(cliques)
def img2graph(img, threshold: int = None, grid=False):
img = np.asarray(img)
shape = img.shape
if len(shape) == 2:
H, W = shape
pixels = img.reshape(-1)
elif len(shape) == 3:
weights = np.array([0.3, 0.59, 0.11]).reshape(3, 1)
C, H, W = shape
pixels = img.reshape(C, -1)
pixels = (weights * pixels).sum(0)
else:
raise RuntimeError(
"RGB(3-dim) or Gray(2-dim) expected, but got {} array(or tensor)".
format(img.shape))
def filter_edges(r, c):
if threshold:
diff = pixels[r] - pixels[c]
idx = abs(diff) < threshold
r = r[idx]
c = c[idx]
i = np.concatenate([r, c])
j = np.concatenate([c, r])
return coo_matrix((np.ones(i.shape), (i, j)), shape=(N, N))
N = H * W
pixel_order = np.arange(N).reshape(H, W)
row_h = pixel_order[:, :-1].reshape(-1)
col_h = row_h + 1
row_v = pixel_order[:-1, :].reshape(-1)
col_v = row_v + W
row_r = np.concatenate([row_h, row_v])
col_r = np.concatenate([col_h, col_v])
Ar = filter_edges(row_r, col_r)
row_br = pixel_order[:-1, :-1].reshape(-1)
col_br = row_br + W + 1
row_bl = pixel_order[:-1, 1:].reshape(-1)
col_bl = row_bl + W - 1
row_d = np.concatenate([row_br, row_bl])
col_d = np.concatenate([col_br, col_bl])
Ad = filter_edges(row_d, col_d)
Ar = SparseTensor.from_scipy(Ar)
Ad = SparseTensor.from_scipy(Ad)
beta_r = np.zeros((H, W), dtype=bool)
beta_r[::2, ::2] = 1
beta_r[1::2, 1::2] = 1
beta_r = beta_r.reshape(-1)
beta_d = np.zeros((H, W), dtype=bool)
beta_d[::2] = 1
beta_d = beta_d.reshape(-1)
xy = None
if grid:
x, y = np.meshgrid(np.arange(W), np.arange(H - 1, -1, -1))
xy = np.hstack([x.reshape(-1, 1), y.reshape(-1, 1)])
return Ar, Ad, beta_r, beta_d, pixels, xy
from thgsp import loadmat
from thgsp.sampling.rsbs import recon_rsbs
from thgsp.utils import mse, snr
def snr_and_mse(x, target):
s, m = snr(x, target), mse(x, target)
print(f"SNR: {s:.4f} | MSE: {m:.4e}")
return s, m
num_sig = 10 # repeat the signal for num_sig times
# Load Data captured from one trial of official code
# http://grsamplingbox.gforge.inria.fr.
# with the official code, the reconstruction SNR is about 30 dB, which is consistent
# with `rsbs` reconstruction in thgsp.
data = loadmat("rsbs.mat")
coh1 = data["weight"].ravel() # distribution of nodes being sampled
S1 = data["ind_obs"].ravel().astype(int) - 1 # sampling node set
mu1 = data["mu"].item() # the factor of regularization term
f = torch.as_tensor(data["x"]) # the original bandlimited signal
y = torch.as_tensor(data["ynoise_init"]).repeat(
1, num_sig) # the contaminated signal
L1 = SparseTensor.from_scipy(data["L"]) # the combinatorial laplacian used
f_hat = recon_rsbs(y, S=S1, L=L1, cum_coh=coh1, mu=mu1, reg_order=1)
s, m = snr_and_mse(f_hat.view(-1, num_sig), f)
