Deprecatted torch_scatter. Now we can use native pytorch implementation

pyproject.toml

@@ -1,12 +1,34 @@
 [project]
 name = "hypergraph-wavelets"
 version = "0.1.0"
-description = "Add your description here"
+description = ""
 readme = "README.md"
 requires-python = ">=3.10"
 dependencies = [
-    "pip>=25.2",
-    "torch==2.8",
-    "torchvision>=0.23.0",
+    "anndata",
+    "einops",
+    "graphtools",
+    "matplotlib",
+    "networkx",
+    "numpy",
+    "opencv-python",
+    "pandas",
+    "phate",
+    "pyarrow",
+    "pytorch-lightning",
+    "scanpy",
+    "scikit-learn",
+    "scipy",
+    "torch",
+    "torch-geometric",
+    "torchmetrics",
+    "torchvision",
+    "tqdm",
 ]
 
+[project.optional-dependencies]
+dev = [
+    "ipykernel",
+    "jupyter",
+    "jupyterlab",
+]

src/hypergraphs/hyperedge_aggregator.py

@@ -1,35 +1,25 @@
 import torch
-from torch_scatter import scatter_mean
+from torch_geometric.utils import scatter
 
 def compute_mean(x, edge_index):
     row, col = edge_index
-    mean_features = scatter_mean(x[row], col, dim=0)
+    mean_features = scatter(x[row], col, dim=0, reduce='mean')
     return mean_features
 
 def compute_var(x, edge_index):
     row, col = edge_index
-    mean_features = scatter_mean(x[row], col, dim=0)
-    mean_sq = scatter_mean((x*x)[row], col, dim=0)
+    mean_features = scatter(x[row], col, dim=0, reduce='mean')
+    mean_sq = scatter((x*x)[row], col, dim=0, reduce='mean')
     var_features = mean_sq - mean_features * mean_features
     return var_features
 
 def compute_cov(x, edge_index):
     row, col = edge_index
-    mean_features = scatter_mean(x[row], col, dim=0)
+    mean_features = scatter(x[row], col, dim=0, reduce='mean')
     xxT = torch.einsum('bi,bj->bij', x, x)
     xxT = xxT[row,...]
     shapes = xxT.size()
     xxT = xxT.reshape(-1, x.size(1) * x.size(1))
-    mean_xxT = scatter_mean(xxT, col, dim=0).reshape(-1, x.size(1), x.size(1))
+    mean_xxT = scatter(xxT, col, dim=0, reduce='mean').reshape(-1, x.size(1), x.size(1))
     cov = mean_xxT - torch.einsum('bi,bj->bij', mean_features, mean_features)
-    return cov
-
-"""
-(manual, for testing, slow)
-"""
-def mean_manual(x, edge_index):
-    means = torch.zeros(edge_index[1].max()+1, x.size(1))
-    for i in range(edge_index[1].max()+1):
-        ids = edge_index[0][edge_index[1] == i]
-        means[i,:] = x[ids,:].mean(dim=0)
-    return means
+    return cov

src/models/hypergraph_scattering.py

@@ -21,8 +21,8 @@
 from torch_geometric.nn.pool import global_mean_pool
 from torch_geometric.nn.conv import MessagePassing
 from torch_geometric.utils import scatter
+from torch_geometric.utils import softmax as pyg_softmax
 from torch_geometric.nn import global_mean_pool, global_max_pool, global_add_pool
-from torch_scatter import scatter_softmax, scatter_sum
 
 
 class LazyLayer(torch.nn.Module):
@@ -269,9 +269,9 @@ def forward(self, x, batch, return_attn=False):
         The shape of `x` is [num_nodes, num_features].
         NOTE: In our case, `x` is the hyperedge features.
         '''
-        gate_scores = self.gate_nn(x).squeeze(-1)                        # [num_nodes]
-        attn_weights = scatter_softmax(gate_scores, batch)               # [num_nodes]
-        out = scatter_sum(x * attn_weights.unsqueeze(-1), batch, dim=0)  # [B, num_features]
+        gate_scores = self.gate_nn(x).squeeze(-1)                              # [num_nodes]
+        attn_weights = pyg_softmax(gate_scores, batch)                         # [num_nodes]
+        out = scatter(x * attn_weights.unsqueeze(-1), batch, dim=0, reduce='sum')  # [B, num_features]
 
         if return_attn:
             return out, attn_weights