Merge pull request #62 from SherylHYX/no_torch_sparse

remove dependence on torch_sparse

.github/workflows/main.yml

@@ -38,8 +38,8 @@ jobs:
     - name: Run installation.
       run: |
         python -m pip install torch==2.3.0 torchvision torchaudio -f https://download.pytorch.org/whl/cpu/torch_stable.html
-        python -m pip install torch-scatter -f https://data.pyg.org/whl/torch-2.3.0+cpu.html
         python -m pip install torch-sparse -f https://data.pyg.org/whl/torch-2.3.0+cpu.html
+        python -m pip install torch-scatter -f https://data.pyg.org/whl/torch-2.3.0+cpu.html
         python -m pip install torch-geometric
         python -m pip install sphinx sphinx_rtd_theme
         python -m pip install -e .[test]

setup.py

@@ -8,7 +8,6 @@
 
 install_requires = [
     "torch",
-    "torch_sparse",
     "scikit-learn",
     "torch_geometric",
     "numpy",

torch_geometric_signed_directed/data/directed/WikipediaNetwork.py

@@ -1,6 +1,6 @@
 from typing import Optional, Callable
 import os.path as osp
-from torch_sparse import coalesce
+from torch_geometric.utils import coalesce
 
 import torch
 import numpy as np
@@ -69,7 +69,7 @@ def process(self):
             data = f.read().split('\n')[1:-1]
             data = [[int(v) for v in r.split('\t')] for r in data]
             edge_index = torch.tensor(data, dtype=torch.long).t().contiguous()
-            edge_index, _ = coalesce(edge_index, None, x.size(0), x.size(0))
+            edge_index, _ = coalesce(edge_index, None, x.size(0))
 
         train_masks, val_masks, test_masks = [], [], []
         for f in self.raw_paths[2:]:

torch_geometric_signed_directed/nn/directed/DGCNConv.py

@@ -2,7 +2,8 @@
 from torch_geometric.typing import Adj, OptTensor
 
 from torch import Tensor
-from torch_sparse import SparseTensor, matmul
+from torch_geometric.typing import SparseTensor
+from torch_geometric.utils import spmm
 from torch_geometric.nn.conv import MessagePassing
 from torch_geometric.nn.conv.gcn_conv import gcn_norm
 
@@ -99,4 +100,4 @@ def message(self, x_j: Tensor, edge_weight: OptTensor) -> Tensor:
         return x_j if edge_weight is None else edge_weight.view(-1, 1) * x_j
 
     def message_and_aggregate(self, adj_t: SparseTensor, x: Tensor) -> Tensor:
-        return matmul(adj_t, x, reduce=self.aggr)
+        return spmm(adj_t, x, reduce=self.aggr)

torch_geometric_signed_directed/nn/general/conv_base.py

@@ -3,10 +3,9 @@
 import torch
 from torch import Tensor
 from torch_geometric.typing import Adj, OptTensor
-from torch_scatter import scatter_add
-from torch_sparse import SparseTensor
+from torch_geometric.typing import SparseTensor
 from torch_geometric.nn.conv import MessagePassing
-from torch_geometric.utils import add_remaining_self_loops
+from torch_geometric.utils import add_remaining_self_loops, scatter
 from torch_geometric.utils.num_nodes import maybe_num_nodes
 
 
@@ -26,7 +25,7 @@ def conv_norm_rw(edge_index, fill_value=0.5, edge_weight=None, num_nodes=None,
         edge_weight = tmp_edge_weight
 
     row = edge_index[0]
-    row_deg = scatter_add(edge_weight, row, dim=0, dim_size=num_nodes)
+    row_deg = scatter(edge_weight, row, dim=0, dim_size=num_nodes, reduce='sum')
     deg_inv = row_deg.pow_(-1)
     deg_inv.masked_fill_(deg_inv == float('inf'), 0)
     return edge_index, deg_inv[row] * edge_weight

torch_geometric_signed_directed/nn/signed/SGCNConv.py

@@ -5,7 +5,8 @@
 from torch import Tensor
 import torch.nn.functional as F
 from torch_geometric.nn.dense.linear import Linear
-from torch_sparse import SparseTensor, matmul
+from torch_geometric.typing import SparseTensor
+from torch_geometric.utils import spmm
 from torch_geometric.nn.conv import MessagePassing
 
 
@@ -130,7 +131,7 @@ def message(self, x_j: Tensor) -> Tensor:
     def message_and_aggregate(self, adj_t: SparseTensor,
                               x: PairTensor) -> Tensor:
         adj_t = adj_t.set_value(None, layout=None)
-        return matmul(adj_t, x[0], reduce=self.aggr)
+        return spmm(adj_t, x[0], reduce=self.aggr)
 
     def __repr__(self) -> str:
         return (f'{self.__class__.__name__}({self.in_dim}, '

torch_geometric_signed_directed/utils/directed/get_adjs_DiGCN.py

@@ -1,8 +1,7 @@
 from typing import Union, Optional, Tuple
 
 import torch
-from torch_geometric.utils import add_self_loops
-from torch_scatter import scatter_add
+from torch_geometric.utils import add_self_loops, scatter
 import scipy
 import numpy as np
 import scipy.sparse as sp
@@ -104,7 +103,7 @@ def cal_fast_appr(alpha: float, edge_index: torch.LongTensor,
 
     # sys normalization
     row, col = edge_index
-    deg = scatter_add(edge_weight, row, dim=0, dim_size=num_nodes)
+    deg = scatter(edge_weight, row, dim=0, dim_size=num_nodes, reduce='sum')
     deg_inv_sqrt = deg.pow(-0.5)
     deg_inv_sqrt[deg_inv_sqrt == float('inf')] = 0
 
@@ -137,7 +136,7 @@ def get_appr_directed_adj(alpha: float, edge_index: torch.LongTensor,
     edge_index, edge_weight = add_self_loops(
         edge_index, edge_weight, fill_value, num_nodes)
     row, col = edge_index
-    deg = scatter_add(edge_weight, row, dim=0, dim_size=num_nodes)
+    deg = scatter(edge_weight, row, dim=0, dim_size=num_nodes, reduce='sum')
     deg_inv = deg.pow(-1)
     deg_inv[deg_inv == float('inf')] = 0
     p = deg_inv[row] * edge_weight
@@ -189,7 +188,7 @@ def get_appr_directed_adj(alpha: float, edge_index: torch.LongTensor,
 
     # row normalization
     row, col = edge_index
-    deg = scatter_add(edge_weight, row, dim=0, dim_size=num_nodes)
+    deg = scatter(edge_weight, row, dim=0, dim_size=num_nodes, reduce='sum')
     deg_inv_sqrt = deg.pow(-0.5)
     deg_inv_sqrt[deg_inv_sqrt == float('inf')] = 0
 
@@ -221,7 +220,7 @@ def get_second_directed_adj(edge_index: torch.LongTensor,
     edge_index, edge_weight = add_self_loops(
         edge_index, edge_weight, fill_value, num_nodes)
     row, col = edge_index
-    deg = scatter_add(edge_weight, row, dim=0, dim_size=num_nodes)
+    deg = scatter(edge_weight, row, dim=0, dim_size=num_nodes, reduce='sum')
     deg_inv = deg.pow(-1)
     deg_inv[deg_inv == float('inf')] = 0
     p = deg_inv[row] * edge_weight
@@ -248,7 +247,7 @@ def get_second_directed_adj(edge_index: torch.LongTensor,
 
     # row normalization
     row, col = edge_index
-    deg = scatter_add(edge_weight, row, dim=0, dim_size=num_nodes)
+    deg = scatter(edge_weight, row, dim=0, dim_size=num_nodes, reduce='sum')
     deg_inv_sqrt = deg.pow(-0.5)
     deg_inv_sqrt[deg_inv_sqrt == float('inf')] = 0
 

torch_geometric_signed_directed/utils/directed/get_magnetic_Laplacian.py

@@ -1,9 +1,7 @@
 from typing import Optional
 
 import torch
-from torch_scatter import scatter_add
-from torch_sparse import coalesce
-from torch_geometric.utils import add_self_loops, remove_self_loops, to_scipy_sparse_matrix
+from torch_geometric.utils import add_self_loops, remove_self_loops, to_scipy_sparse_matrix, coalesce, scatter
 from torch_geometric.utils.num_nodes import maybe_num_nodes
 import numpy as np
 from scipy.sparse.linalg import eigsh
@@ -59,14 +57,13 @@ def get_magnetic_Laplacian(edge_index: torch.LongTensor, edge_weight: Optional[t
     sym_attr = torch.cat([edge_weight, edge_weight], dim=0)
     edge_attr = torch.stack([sym_attr, theta_attr], dim=1)
 
-    edge_index_sym, edge_attr = coalesce(edge_index, edge_attr, num_nodes,
-                                         num_nodes, "add")
+    edge_index_sym, edge_attr = coalesce(edge_index, edge_attr, num_nodes, "add")
 
     edge_weight_sym = edge_attr[:, 0]
     edge_weight_sym = edge_weight_sym/2
 
     row, col = edge_index_sym[0], edge_index_sym[1]
-    deg = scatter_add(edge_weight_sym, row, dim=0, dim_size=num_nodes)
+    deg = scatter(edge_weight_sym, row, dim=0, dim_size=num_nodes, reduce='sum')
 
     edge_weight_q = torch.exp(1j * 2 * np.pi * q * edge_attr[:, 1])
 

torch_geometric_signed_directed/utils/general/get_magnetic_signed_Laplacian.py

@@ -1,9 +1,7 @@
 from typing import Optional
 
 import torch
-from torch_scatter import scatter_add
-from torch_sparse import coalesce
-from torch_geometric.utils import add_self_loops, remove_self_loops, to_scipy_sparse_matrix
+from torch_geometric.utils import add_self_loops, remove_self_loops, to_scipy_sparse_matrix, coalesce, scatter
 from torch_geometric.utils.num_nodes import maybe_num_nodes
 import numpy as np
 from scipy.sparse.linalg import eigsh
@@ -61,8 +59,7 @@ def get_magnetic_signed_Laplacian(edge_index: torch.LongTensor, edge_weight: Opt
     sym_abs_attr = torch.cat([torch.abs(edge_weight), torch.abs(edge_weight)], dim=0)
     edge_attr = torch.stack([sym_attr, theta_attr, sym_abs_attr], dim=1)
 
-    edge_index_sym, edge_attr = coalesce(edge_index, edge_attr, num_nodes,
-                                     num_nodes, "add")
+    edge_index_sym, edge_attr = coalesce(edge_index, edge_attr, num_nodes, "add")
 
     row, col = edge_index_sym[0], edge_index_sym[1]
     edge_weight_sym = edge_attr[:, 0]
@@ -71,9 +68,9 @@ def get_magnetic_signed_Laplacian(edge_index: torch.LongTensor, edge_weight: Opt
     if absolute_degree:
         edge_weight_abs_sym = edge_attr[:, 2]
         edge_weight_abs_sym = edge_weight_abs_sym/2
-        deg = scatter_add(edge_weight_abs_sym, row, dim=0, dim_size=num_nodes)
+        deg = scatter(edge_weight_abs_sym, row, dim=0, dim_size=num_nodes, reduce='sum')
     else:
-        deg = scatter_add(torch.abs(edge_weight_sym), row, dim=0, dim_size=num_nodes) # absolute values for edge weights
+        deg = scatter(torch.abs(edge_weight_sym), row, dim=0, dim_size=num_nodes, reduce='sum') # absolute values for edge weights
 
     edge_weight_q = torch.exp(1j * 2* np.pi * q * edge_attr[:, 1])
 

torch_geometric_signed_directed/utils/signed/create_spectral_features.py

@@ -2,7 +2,7 @@
 import torch
 import scipy.sparse as sp
 from sklearn.decomposition import TruncatedSVD
-from torch_sparse import coalesce
+from torch_geometric.utils import coalesce
 
 
 def create_spectral_features(
@@ -27,7 +27,7 @@ def create_spectral_features(
     edge_index = torch.cat([edge_index, torch.stack([col, row])], dim=1)
     val = torch.cat([val, val], dim=0)
 
-    edge_index, val = coalesce(edge_index, val, N, N)
+    edge_index, val = coalesce(edge_index, val, N)
     val = val - 1
 
     # Borrowed from: