remove torch-scatter dependency

Author: timkucera

cosmic/cosmo.py

@@ -1,6 +1,7 @@
 import torch
 from torch import nn
-from torch_scatter import scatter_add, scatter_mean, scatter_softmax
+
+from .utilities import scatter_add, scatter_mean, scatter_softmax
 
 """
 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.

cosmic/lift.py

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

cosmic/utilities.py

@@ -131,3 +131,64 @@ 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)

pyproject.toml

@@ -30,8 +30,3 @@ include = ["cosmic", "cosmic.*"]
 
 [tool.setuptools_scm]
 
-[project.optional-dependencies]
-scatter = [
-  "torch_scatter>=2.1",
-]
-