add basic tests + standalone scatter/radius_graph

Author: kleinhenz

.gitignore

@@ -0,0 +1 @@
+__pycache__

pyproject.toml

@@ -9,9 +9,18 @@ authors = [
 requires-python = ">=3.11"
 dependencies = [
     "e3nn>=0.5.5",
+    "jaxtyping>=0.2.38",
     "torch>=2.4.1",
 ]
 
 [build-system]
 requires = ["hatchling"]
 build-backend = "hatchling.build"
+
+[dependency-groups]
+dev = [
+    "pytest>=8.3.4",
+]
+
+[tool.ruff.lint]
+ignore = ["F722"]

src/e3tools/__init__.py

@@ -0,0 +1,4 @@
+from ._scatter import scatter
+from ._radius import radius, radius_graph
+
+__all__ = ["scatter", "radius", "radius_graph"]

src/e3tools/_radius.py

@@ -0,0 +1,79 @@
+import math
+
+import torch
+from torch import Tensor
+
+from jaxtyping import Float, Int64
+
+
+# ref https://github.com/rusty1s/pytorch_cluster/blob/master/torch_cluster/radius.py
+def radius(
+    x: Float[Tensor, "N D"],
+    y: Float[Tensor, "M D"],
+    r: float,
+    batch_x: Int64[Tensor, " N"] | None = None,
+    batch_y: Int64[Tensor, " M"] | None = None,
+    ignore_same_index: bool = True,
+    chunk_size: int | None = None,
+) -> Int64[Tensor, "2 E"]:
+    """For each element in `y` find all points in `x` within distance `r`"""
+    N, _ = x.shape
+    M, _ = y.shape
+
+    if chunk_size is None:
+        chunk_size = N + 1
+
+    if batch_x is None:
+        batch_x = torch.zeros(N, dtype=torch.int64, device=x.device)
+
+    if batch_y is None:
+        batch_y = torch.zeros(N, dtype=torch.int64, device=x.device)
+
+    if x.numel() == 0 or y.numel() == 0:
+        return torch.empty(2, 0, dtype=torch.long, device=x.device)
+
+    x = x.view(-1, 1) if x.dim() == 1 else x
+    y = y.view(-1, 1) if y.dim() == 1 else y
+    x, y = x.contiguous(), y.contiguous()
+
+    batch_size = int(batch_x.max()) + 1
+    batch_size = max(batch_size, int(batch_y.max()) + 1)
+    assert batch_size > 0
+
+    r2 = torch.as_tensor(r * r, dtype=x.dtype, device=x.device)
+
+    n_chunks = math.ceil(N / chunk_size)
+
+    rows = []
+    cols = []
+
+    for y_chunk, batch_y_chunk, index_y_chunk in zip(
+        torch.chunk(y, n_chunks),
+        torch.chunk(batch_y, n_chunks),
+        torch.chunk(torch.arange(M, device=x.device), n_chunks),
+    ):
+        pdist = (x[:, None] - y_chunk).pow(2).sum(dim=-1)
+        same_batch = batch_x[:, None] == batch_y_chunk
+        same_index = torch.arange(N, device=x.device)[:, None] == index_y_chunk
+
+        connected = (pdist <= r2) & same_batch
+        if ignore_same_index:
+            connected = connected & ~same_index
+
+        row, col = torch.nonzero(connected, as_tuple=True)
+        cols.append(col + index_y_chunk[0])
+        rows.append(row)
+
+    row = torch.cat(rows, dim=0)
+    col = torch.cat(cols, dim=0)
+
+    return torch.stack((col, row), dim=0)
+
+
+def radius_graph(
+    x: Float[Tensor, "N D"],
+    r: float,
+    batch: Int64[Tensor, " N"] | None = None,
+    chunk_size: int | None = None,
+) -> Int64[Tensor, "2 E"]:
+    return radius(x, x, r, batch, batch, ignore_same_index=True, chunk_size=chunk_size)

src/e3tools/_scatter.py

@@ -0,0 +1,39 @@
+import torch
+from torch import Tensor
+
+
+def broadcast(src: Tensor, other: Tensor, dim: int):
+    if dim < 0:
+        dim = other.dim() + dim
+    if src.dim() == 1:
+        for _ in range(0, dim):
+            src = src.unsqueeze(0)
+    for _ in range(src.dim(), other.dim()):
+        src = src.unsqueeze(-1)
+    src = src.expand(other.size())
+    return src
+
+
+def scatter(src, index, dim, dim_size: int | None = None, reduce="sum"):
+    in_shape = src.shape
+
+    if dim < 0:
+        dim = src.dim() + dim
+
+    if dim_size is None:
+        if index.numel() == 0:
+            dim_size = 0
+        else:
+            dim_size = int(index.max()) + 1
+
+    index = broadcast(index, src, dim)
+
+    assert src.ndim == index.ndim, f"{src.ndim=}, {index.ndim=}"
+
+    out_shape = (*in_shape[:dim], dim_size, *in_shape[dim + 1 :])
+    out = torch.zeros(*out_shape, dtype=src.dtype, device=src.device)
+
+    assert out.ndim == index.ndim, (
+        f"{out.ndim=}, {index.ndim=} {out_shape=}, {in_shape=}, {dim=}"
+    )
+    return torch.scatter_reduce(out, dim, index, src, reduce, include_self=False)

src/e3tools/nn/_conv.py

@@ -4,7 +4,8 @@
 import e3nn
 import torch
 from e3nn import o3
-from torch_scatter import scatter
+
+from e3tools import scatter
 
 from ._gate import Gated
 from ._interaction import LinearSelfInteraction

src/e3tools/nn/_transformer.py

@@ -4,7 +4,8 @@
 import e3nn
 import torch
 from e3nn import o3
-from torch_scatter import scatter
+
+from e3tools import scatter
 
 from ._conv import Conv
 from ._interaction import LinearSelfInteraction

tests/test_basic.py

@@ -0,0 +1,207 @@
+import functools
+
+import e3nn
+import pytest
+import torch
+from e3nn import o3
+from e3tools.nn import (
+    Attention,
+    Conv,
+    ConvBlock,
+    EquivariantMLP,
+    ExperimentalConv,
+    Gated,
+    LayerNorm,
+    MultiheadAttention,
+    SeparableConv,
+    TransformerBlock,
+)
+from e3tools import radius_graph
+
+torch.set_default_dtype(torch.float64)
+
+CONV_LAYERS = [Conv, SeparableConv, ExperimentalConv]
+
+
+def apply_layer_rot(layer):
+    N = 20
+    edge_attr_dim = 10
+    max_radius = 1.3
+
+    pos = torch.randn(N, 3)
+    node_attr = layer.irreps_in.randn(N, -1)
+
+    edge_index = radius_graph(pos, max_radius)
+    edge_vec = pos[edge_index[0]] - pos[edge_index[1]]
+    edge_length = (edge_vec).norm(dim=1)
+    edge_attr = e3nn.math.soft_one_hot_linspace(
+        edge_length,
+        start=0.0,
+        end=max_radius,
+        number=edge_attr_dim,
+        basis="smooth_finite",
+        cutoff=True,
+    )
+
+    edge_sh = o3.spherical_harmonics(
+        layer.irreps_sh, edge_vec, True, normalization="component"
+    )
+
+    rot = o3.rand_matrix()
+
+    D_node_attr = layer.irreps_in.D_from_matrix(rot)
+    D_edge_sh = layer.irreps_sh.D_from_matrix(rot)
+
+    D_out = layer.irreps_out.D_from_matrix(rot)
+
+    out_1 = layer(
+        node_attr @ D_node_attr.T, edge_index, edge_attr, edge_sh @ D_edge_sh.T
+    )
+    out_2 = layer(node_attr, edge_index, edge_attr, edge_sh) @ D_out.T
+
+    return out_1, out_2
+
+
+@pytest.mark.parametrize("conv", CONV_LAYERS)
+def test_conv(conv):
+    irreps_in = o3.Irreps("10x0e + 10x1o + 10x2e")
+    irreps_sh = irreps_in.spherical_harmonics(2)
+    edge_attr_dim = 10
+
+    layer = conv(irreps_in, irreps_in, irreps_sh, edge_attr_dim=edge_attr_dim)
+
+    out_1, out_2 = apply_layer_rot(layer)
+    assert torch.allclose(out_1, out_2, atol=1e-10)
+
+
+@pytest.mark.parametrize("conv", CONV_LAYERS)
+def test_gated_conv(conv):
+    irreps_in = o3.Irreps("10x0e + 10x1o + 10x2e")
+    irreps_sh = irreps_in.spherical_harmonics(2)
+    edge_attr_dim = 10
+
+    wrapped = functools.partial(conv, irreps_sh=irreps_sh, edge_attr_dim=edge_attr_dim)
+
+    layer = Gated(wrapped, irreps_in=irreps_in, irreps_out=irreps_in)
+
+    out_1, out_2 = apply_layer_rot(layer)
+    assert torch.allclose(out_1, out_2, atol=1e-10)
+
+
+@pytest.mark.parametrize("conv", CONV_LAYERS)
+def test_conv_block(conv):
+    irreps_in = o3.Irreps("10x0e + 10x1o + 10x2e")
+    irreps_sh = irreps_in.spherical_harmonics(2)
+    edge_attr_dim = 10
+
+    layer = ConvBlock(
+        irreps_in=irreps_in,
+        irreps_out=irreps_in,
+        irreps_sh=irreps_sh,
+        edge_attr_dim=edge_attr_dim,
+        conv=conv,
+    )
+
+    out_1, out_2 = apply_layer_rot(layer)
+    assert torch.allclose(out_1, out_2, atol=1e-10)
+
+
+@pytest.mark.parametrize("conv", CONV_LAYERS)
+def test_attention(conv):
+    irreps_in = o3.Irreps("10x0e + 10x1o + 10x2e")
+    irreps_out = irreps_in
+    irreps_sh = irreps_in.spherical_harmonics(2)
+    irreps_key = irreps_in
+    irreps_query = irreps_in
+    edge_attr_dim = 10
+
+    layer = Attention(
+        irreps_in,
+        irreps_out,
+        irreps_sh,
+        irreps_query,
+        irreps_key,
+        edge_attr_dim,
+        conv=conv,
+    )
+
+    out_1, out_2 = apply_layer_rot(layer)
+    assert torch.allclose(out_1, out_2, atol=1e-10)
+
+
+@pytest.mark.parametrize("conv", [Conv, SeparableConv])
+def test_multihead_attention(conv):
+    irreps_in = o3.Irreps("10x0e + 10x1o + 10x2e")
+    irreps_out = irreps_in
+    irreps_sh = irreps_in.spherical_harmonics(2)
+    irreps_key = irreps_in
+    irreps_query = irreps_in
+    edge_attr_dim = 10
+    n_head = 2
+
+    layer = MultiheadAttention(
+        irreps_in,
+        irreps_out,
+        irreps_sh,
+        irreps_query,
+        irreps_key,
+        edge_attr_dim,
+        n_head,
+        conv=conv,
+    )
+
+    out_1, out_2 = apply_layer_rot(layer)
+    assert torch.allclose(out_1, out_2, atol=1e-10)
+
+
+def test_layer_norm():
+    irreps = o3.Irreps("10x0e + 10x1o + 10x2e")
+
+    layer = LayerNorm(irreps)
+    rot = o3.rand_matrix()
+    D = irreps.D_from_matrix(rot)
+
+    x = irreps.randn(10, -1)
+
+    out_1 = layer(x @ D.T)
+    out_2 = layer(x) @ D.T
+
+    assert torch.allclose(out_1, out_2, atol=1e-10)
+
+
+def test_equivariant_mlp():
+    irreps = o3.Irreps("10x0e + 10x1o + 10x2e")
+    irreps_hidden = o3.Irreps([(4 * mul, ir) for mul, ir in irreps])
+
+    layer = EquivariantMLP(
+        irreps, irreps, [irreps_hidden, irreps_hidden], norm_layer=LayerNorm
+    )
+
+    rot = o3.rand_matrix()
+    D = irreps.D_from_matrix(rot)
+
+    x = irreps.randn(10, -1)
+
+    out_1 = layer(x @ D.T)
+    out_2 = layer(x) @ D.T
+
+    assert torch.allclose(out_1, out_2, atol=1e-10)
+
+
+def test_transformer():
+    irreps_in = o3.Irreps("10x0e + 10x1o + 10x2e")
+    irreps_out = irreps_in
+    irreps_sh = irreps_in.spherical_harmonics(2)
+    edge_attr_dim = 10
+    n_head = 2
+
+    layer = TransformerBlock(
+        irreps_in=irreps_in,
+        irreps_out=irreps_out,
+        irreps_sh=irreps_sh,
+        edge_attr_dim=edge_attr_dim,
+        n_head=n_head,
+    )
+
+    out_1, out_2 = apply_layer_rot(layer)
+    assert torch.allclose(out_1, out_2, atol=1e-10)