fix scatter

Author: timkucera

README.md

@@ -4,7 +4,7 @@
 
 This repository contains the Cosmo neural network lift and convolution layers. For a usage example and reproduction of the results of the RECOMB 2026 submission "Gaining mechanistic insight from geometric deep learning on molecule structures through equivariant convolution", see https://github.com/BorgwardtLab/RECOMB2026Cosmo.
 
-Installation: `pip install cosmic-torch` or `pip install git+https://github.com/BorgwardtLab/Cosmo`
+Installation: `pip install cosmic-torch` or `pip install git+https://github.com/BorgwardtLab/Cosmo`. Make sure to before install [torch](https://pytorch.org/get-started/locally/) and [torch-scatter](https://pypi.org/project/torch-scatter/) according to their instructions.
 
 ### Cosmo
 

cosmic/cosmo.py

@@ -1,7 +1,6 @@
 import torch
 from torch import nn
-
-from .utilities import scatter_mean, scatter_softmax, scatter_sum
+from torch_scatter import scatter_mean, scatter_softmax, scatter_sum
 
 """
 Cosmo can be implemented with various filter functions. The underlying principle is always to compute the filter under transformation of a local reference frame (hood_coords) which is derived from neighboring input points. The forward signature of the layer is always the same and inputs can be obtained from a Lift2D or Lift3D module.
@@ -41,7 +40,7 @@ def forward(
         w = self.w[:, nn_idx]  # use closest kernel point
         f = features[source]
         out_channels = torch.einsum("ni,oni->no", f, w)  # m x out
-        features = scatter_sum(out_channels, target, m)
+        features = scatter_sum(out_channels, target, dim_size=m, dim=0)
         return features  # Updated features of shape (m, out_channels)
 
 
@@ -94,7 +93,7 @@ def forward(
         )
         f = features[source]
         out_channels = torch.einsum("ni,noi->no", f, w)  # m x out
-        features = scatter_mean(out_channels, target, m)
+        features = scatter_mean(out_channels, target, dim_size=m, dim=0)
         return features  # Updated features of shape (m, out_channels)
 
 
@@ -136,6 +135,6 @@ def forward(
         w1 = self.w1(features)
         w2 = self.w2(features)
         w3 = self.w3(features)
-        a = scatter_softmax(w1[target] - w2[source] + d, target, m)
-        features = scatter_sum(a * (w3[source] + d), target, m)
+        a = scatter_softmax(w1[target] - w2[source] + d, target, dim_size=m, dim=0)
+        features = scatter_sum(a * (w3[source] + d), target, dim_size=m, dim=0)
         return features  # Updated features of shape (m, out_channels)

cosmic/lift.py

@@ -1,6 +1,7 @@
 from types import SimpleNamespace
 
 import torch
+from torch_scatter import scatter_max, scatter_mean, scatter_softmax, scatter_sum
 
 from .utilities import *
 
@@ -123,12 +124,12 @@ def __init__(self, agg="mean"):
 
     def __call__(self, features, index, size, return_index=False):
         if self.agg == "sum":
-            return scatter_sum(features, index, size)
+            return scatter_sum(features, index, dim_size=size, dim=0)
         elif self.agg == "mean":
-            return scatter_mean(features, index, size)
+            return scatter_mean(features, index, dim_size=size, dim=0)
         elif self.agg == "max":
-            result = scatter_max(features, index, size)
+            result = scatter_max(features, index, dim_size=size, dim=0)
             return result[0] if not return_index else result
         elif self.agg == "softmax":
-            a = scatter_softmax(features, index, size)
-            return scatter_sum(a * features, index, size)
+            a = scatter_softmax(features, index, dim_size=size, dim=0)
+            return scatter_sum(a * features, index, dim_size=size, dim=0)

cosmic/utilities.py

@@ -131,64 +131,3 @@ def filter_bases(bases, minimum_angle, coords, triangles):
         valid = angles_deg >= minimum_angle
         bases[~valid] = torch.nan
     return bases
-
-
-def scatter_sum(src, index, size):
-    assert index.dtype == torch.long
-    assert index.shape[0] == src.shape[0]
-    out_shape = (size,) + src.shape[1:]
-    out = torch.zeros(out_shape, dtype=src.dtype, device=src.device)
-    out.index_add_(0, index, src)
-    return out
-
-
-def scatter_mean(src, index, size):
-    sum_out = scatter_sum(src, index, size)
-    count = scatter_sum(torch.ones_like(src), index, size)
-    return sum_out / count.clamp_min(1)
-
-
-def scatter_max(src, index, size):
-    assert src.is_floating_point() and index.dtype == torch.long
-    N = src.size(0)
-    idx_exp = index.view(N, *([1] * (src.ndim - 1))).expand_as(src)
-    vals = torch.full(
-        (size,) + src.shape[1:],
-        torch.finfo(src.dtype).min,
-        dtype=src.dtype,
-        device=src.device,
-    )
-    vals.scatter_reduce_(0, idx_exp, src, reduce="amax")
-    pos = (
-        torch.arange(N, device=src.device)
-        .view(N, *([1] * (src.ndim - 1)))
-        .expand_as(src)
-    )
-    pos_mask = torch.where(src == vals[index], pos, torch.full_like(pos, N))
-    argmax = torch.full_like(vals, N, dtype=torch.long)
-    argmax.scatter_reduce_(0, idx_exp, pos_mask, reduce="amin")
-    argmax[argmax == N] = -1
-    count = torch.zeros_like(vals, dtype=torch.long)
-    count.scatter_reduce_(
-        0, idx_exp, torch.ones_like(src, dtype=torch.long), reduce="sum"
-    )
-    vals[count == 0] = 0
-    return vals, argmax
-
-
-def scatter_softmax(src, index, size, eps=1e-12):
-    assert src.is_floating_point()
-    assert index.dtype == torch.long
-    N = src.shape[0]
-    idx_exp = index.view(N, *([1] * (src.ndim - 1))).expand_as(src)
-    max_vals = torch.full(
-        (size,) + src.shape[1:],
-        torch.finfo(src.dtype).min,
-        dtype=src.dtype,
-        device=src.device,
-    )
-    max_vals.scatter_reduce_(0, idx_exp, src, reduce="amax")
-    ex = (src - max_vals[index]).exp()
-    sum_per_group = torch.zeros_like(max_vals)
-    sum_per_group.scatter_reduce_(0, idx_exp, ex, reduce="sum")
-    return ex / sum_per_group[index].clamp_min(eps)