update readme

Author: timkucera

README.md

@@ -1,4 +1,79 @@
+<p align="center">
+    <img src="assets/logo.png" width="100"/>
+</p>
+
 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`
+
+### Cosmo
+
+Cosmo is a neural network architecture based on message passing on geometric graphs of molecule structures. It applies a convolutional filter by translating it to vertices and rotating it towards neighbors. The resulting feature activation (message) is passed to the neighbor that the filter was pointed at. This way, large geometric patterns can be modeled with a template-matching objective by using multiple Cosmo layers. A Cosmo network is equivariant to translation and rotation, and highly interpretable as its weight matrices can be linearly combined and its filter poses can be reconstructed geometrically. For more details, please see the paper.
+
+### Example Usage
+
+Cosmo layers operate on lifted geometric graphs. These are computed from an adjacency matrix of the data, either given by e.g. atomic bond connectivity, or constructed by e.g. k-NN:
+
+`adj = torch_geometric.nn.knn_graph(coords, k, batch_index)`
+
+where `coords` are the input point coordinates of the data, `k` is a hyperparameter, and `batch_index` assigns each node to an instance in the batch (compare the computing principles of [PyG](https://pytorch-geometric.readthedocs.io/en/2.4.0/index.html), which we highly recommend to use).
+
+Given coordinates, node features (e.g. one-hot encoded atom type), and the adjacency we can lift the input graph:
+
+```
+L = Lift2D()(features, coords, adj, batch_index) # or Lift3D()
+```
+
+The `L` namespace contains everything that we need to compute in subsequent Cosmo layers:
+
+```
+features = layer(L.source, L.target, L.features, L.hood_coords)
+```
+
+After the Cosmo layers we need to undo the lift operation (lowering) to obtain features on the input graph. This is done by aggregating the edge/triangle features to the nodes, which yields a standard graph object that can be further computed on with PyG layers, for example.
+
+```
+node_features = Lower(agg="max")(features, L.lifted2node, num_nodes)
+```
+
+Or, if features should be aggregated directly to the instance (graph) level:
+
+```
+graph_features = Lower(agg="max")(features, L.lifted2inst, num_instances)
+```
+
+An entire Cosmo network for a node classification task could look like this:
+
+```
+from cosmic import *
+import torch.nn as nn
+
+class CosmoModel(nn.Module):
+
+    def __init__(self):
+        self.lift = Lift3D()
+        self.lower = Lower()
+        self.cosmo_layers = nn.ModuleList([
+            NeuralFieldCosmo(in_channels=5, out_channels=128, dim=3),
+            NeuralFieldCosmo(in_channels=128, out_channels=128, dim=3),
+            NeuralFieldCosmo(in_channels=128, out_channels=10, dim=3)
+        ])
+
+    def forward(self, node_features, coords, adj, batch_index, num_nodes):
+        L = self.lift(node_features, coords, adj, batch_index)
+        features = L.features
+        for layer in self.cosmo_layers:
+            features = layer(L.source, L.target, features, L.hood_coords)
+        node_features = self.lower(features, L.lifted2node, num_nodes)
+        # there could be some classic GNN-layers here, or an MLP head
+        return node_features
+```
+
+
+### Citation
+
+TBD
+
+### License
 
-Installation: This package depends on torch and torch-scatter. Please install according to your system and their instructions.
+TBD

assets/logo.png

@@ -1,80 +1,114 @@
 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.
+"""
 
 
 class KernelPointCosmo(nn.Module):
+    """
+    Implements Kernel Point Convolution (Thomas et al. 2019) in the Cosmo framework.
+    Note that we implement an optimization trick from KPConvX (Thomas et al. 2024) which uses only the closest kernel point for each input point.
+    """
 
-    def __init__(self, in_channels, out_channels, filter):
+    def __init__(
+        self,
+        in_channels,  # Number of input channels
+        out_channels,  # Number of output channels
+        kernel_points,  # Kernel points of shape (k, dim)
+    ):
         super().__init__()
         self.out_channels = out_channels
-        mu = filter.unsqueeze(0).float()  # out_channels x k x d
+        mu = kernel_points.unsqueeze(0).float()  # out_channels x k x d
         self.register_buffer("mu", mu)
-        self.w = nn.Parameter(
-            torch.rand(out_channels, mu.shape[1], in_channels)
-        )  # out_channels x k x in_channels
+        self.w = nn.Parameter(torch.rand(out_channels, mu.shape[1], in_channels))
         nn.init.xavier_uniform_(self.w)
 
-    def forward(self, ijk, jkl, triangle_features, hood_coords):
+    def forward(
+        self,
+        source,  # Source edges ij or triangles ijk (m,)
+        target,  # Target edges jk or triangles jkl (m,)
+        features,  # Edge or triangle features of shape (m, in_channels)
+        hood_coords,  # Locally transformed coordinates of shape (m, dim)
+    ):
+        m = features.shape[0]
         with torch.no_grad():
-            dist = torch.cdist(hood_coords.unsqueeze(0), self.mu)  # n x k
+            dist = torch.cdist(hood_coords.unsqueeze(0), self.mu)  # m x k
             nn_idx = dist.argmin(dim=2).squeeze(0)
         w = self.w[:, nn_idx]  # use closest kernel point
-        f = triangle_features[ijk]
-        out_channels = torch.einsum("ni,oni->no", f, w)  # n x out
-        triangle_features = scatter_add(
-            out_channels, jkl, dim=0, dim_size=triangle_features.shape[0]
-        )
-        return triangle_features
+        f = features[source]
+        out_channels = torch.einsum("ni,oni->no", f, w)  # m x out
+        features = scatter_add(out_channels, target, dim=0, dim_size=m)
+        return features  # Updated features of shape (m, out_channels)
 
 
 class NeuralFieldCosmo(nn.Module):
+    """
+    Implements Neural Field Convolution (Proposed with Cosmo in Kucera et al. 2026) in the Cosmo framework. Weight matrices are computed from input coordinates in the local reference frame using a neural field (parameterized by a neural network).
+    """
+
     def __init__(
         self,
-        in_channels,
-        out_channels,
-        hidden_channels=32,
-        num_layers=3,
-        dropout=0.0,
-        radius=1.0,
-        dim=3,
-        field_activation=nn.Tanh,
+        in_channels,  # Number of input channels
+        out_channels,  # Number of output channels
+        field_channels=32,  # Number of channels in the neural field
+        field_layers=3,  # Number of layers in the neural field
+        field_dropout=0.0,  # Dropout rate in the neural field
+        field_activation=nn.Tanh,  # Activation function in the neural field
+        radius=1.0,  # Scale parameter for input coordinates
+        dim=3,  # Dimension of the input data (2 or 3)
     ):
         super().__init__()
         self.register_buffer("radius", torch.tensor(radius))
         self.register_buffer("in_channels", torch.tensor(in_channels))
         self.register_buffer("out_channels", torch.tensor(out_channels))
         self.neural_field = nn.Sequential(
-            nn.Linear(dim, hidden_channels),
-            nn.LayerNorm(hidden_channels),
+            nn.Linear(dim, field_channels),
+            nn.LayerNorm(field_channels),
             nn.ReLU(),
-            nn.Dropout(dropout),
+            nn.Dropout(field_dropout),
             *[
-                nn.Linear(hidden_channels, hidden_channels),
-                nn.LayerNorm(hidden_channels),
+                nn.Linear(field_channels, field_channels),
+                nn.LayerNorm(field_channels),
                 nn.ReLU(),
-                nn.Dropout(dropout),
+                nn.Dropout(field_dropout),
             ]
-            * (num_layers - 2),
-            nn.Linear(hidden_channels, in_channels * out_channels),
+            * (field_layers - 2),
+            nn.Linear(field_channels, in_channels * out_channels),
             field_activation(),
         )
 
-    def forward(self, in_edges, out_edges, edge_features, hood_coords):
+    def forward(
+        self,
+        source,  # Source edges ij or triangles ijk (m,)
+        target,  # Target edges jk or triangles jkl (m,)
+        features,  # Edge or triangle features of shape (m, in_channels)
+        hood_coords,  # Locally transformed coordinates of shape (m, dim)
+    ):
+        m = features.shape[0]
         w = self.neural_field(hood_coords / self.radius).view(
             -1, self.out_channels, self.in_channels
         )
-        f = edge_features[in_edges]
-        out_channels = torch.einsum("ni,noi->no", f, w)  # n x out
-        edge_features = scatter_mean(
-            out_channels, out_edges, dim_size=edge_features.shape[0], dim=0
-        )
-        return edge_features
+        f = features[source]
+        out_channels = torch.einsum("ni,noi->no", f, w)  # m x out
+        features = scatter_mean(out_channels, target, dim_size=m, dim=0)
+        return features  # Updated features of shape (m, out_channels)
 
 
 class PointTransformerCosmo(nn.Module):
-    def __init__(self, in_channels, out_channels, radius, dim=3):
+    """
+    Implements Point Transformer Convolution (Zhao et al. 2020) in the Cosmo framework.
+    """
+
+    def __init__(
+        self,
+        in_channels,  # Number of input channels
+        out_channels,  # Number of output channels
+        radius=1.0,  # Scale parameter for input coordinates
+        dim=3,  # Dimension of the input data (2 or 3)
+    ):
         super().__init__()
         self.register_buffer("radius", torch.tensor(radius))
         self.delta = nn.Sequential(
@@ -89,12 +123,18 @@ def __init__(self, in_channels, out_channels, radius, dim=3):
         self.w2 = nn.Linear(in_channels, out_channels, bias=False)
         self.w3 = nn.Linear(in_channels, out_channels, bias=False)
 
-    def forward(self, ijk, jkl, tri_features, hood_coords):
-        n = tri_features.shape[0]
+    def forward(
+        self,
+        source,  # Source edges ij or triangles ijk (m,)
+        target,  # Target edges jk or triangles jkl (m,)
+        features,  # Edge or triangle features of shape (m, in_channels)
+        hood_coords,  # Locally transformed coordinates of shape (m, dim)
+    ):
+        m = features.shape[0]
         d = self.delta(hood_coords / self.radius)
-        w1 = self.w1(tri_features)
-        w2 = self.w2(tri_features)
-        w3 = self.w3(tri_features)
-        a = scatter_softmax(w1[jkl] - w2[ijk] + d, jkl, dim=0, dim_size=n)
-        tri_features = scatter_add(a * (w3[ijk] + d), jkl, dim=0, dim_size=n)
-        return tri_features
+        w1 = self.w1(features)
+        w2 = self.w2(features)
+        w3 = self.w3(features)
+        a = scatter_softmax(w1[target] - w2[source] + d, target, dim=0, dim_size=m)
+        features = scatter_add(a * (w3[source] + d), target, dim=0, dim_size=m)
+        return features  # Updated features of shape (m, out_channels)

cosmic/cosmo.py

@@ -1,14 +1,24 @@
 from types import SimpleNamespace
 
 import torch
+from torch_scatter import scatter_max, scatter_mean, scatter_softmax, scatter_sum
 
 from .utilities import *
 
 
 class Lift2D:
+    """
+    Parameter-free module to lift a 2D geometric graph. Given node features, global coordinates, and a graph adjacency matrix it computes the lifted adjacency and coordinates of neighborhoods in the local reference frame of the edges, together with some helper variables. These build the input to a Cosmo layer.
+    """
 
     @torch.compiler.disable
-    def __call__(self, node_features, coords, edge_index, node2inst):
+    def __call__(
+        self,
+        node_features,  # Node features of shape (n, in_channels)
+        coords,  # Global coordinates of shape (n, 2)
+        edge_index,  # Edge index of shape (2, m)
+        node2inst,  # Node-to-instance mapping of shape (n,)
+    ):
         n = coords.shape[0]
         adj = (
             torch.sparse_coo_tensor(
@@ -32,26 +42,36 @@ def __call__(self, node_features, coords, edge_index, node2inst):
         edge2inst = node2inst[edge2node]
         edge_features = node_features[edge2node]
         return SimpleNamespace(
-            adj=adj,
-            ij=ij,
-            jk=jk,
-            coords=coords,
-            hood_coords=hood_coords,
-            edge_features=edge_features,
-            bases=bases,
-            i=i,
-            j=j,
-            edges=edges,
-            node2inst=node2inst,
-            edge2node=edge2node,
-            edge2inst=edge2inst,
+            adj=adj,  # Sorted adjacency matrix of shape (n, n)
+            source=ij,  # Lifted source edges (m,)
+            target=jk,  # Lifted target edges (m,)
+            coords=coords,  # Global coordinates of shape (n, 2)
+            hood_coords=hood_coords,  # Local coordinates of shape (m, 2)
+            features=edge_features,  # Edge features of shape (m, in_channels)
+            bases=bases,  # Bases of shape (m, 2, 2)
+            i=i,  # Node indices i (m,)
+            j=j,  # Node indices j (m,)
+            edges=edges,  # Tuples of i,j (m, 2)
+            node2inst=node2inst,  # Node-to-instance mapping of shape (n,)
+            lifted2node=edge2node,  # Edge-to-node mapping of shape (m,)
+            lifted2inst=edge2inst,  # Edge-to-instance mapping of shape (m,)
         )
 
 
 class Lift3D:
+    """
+    Parameter-free module to lift a 3D geometric graph. Given node features, global coordinates, and a graph adjacency matrix it computes the lifted adjacency and coordinates of neighborhoods in the local reference frame of the triangles, together with some helper variables. These build the input to a Cosmo layer.
+    """
 
     @torch.compiler.disable
-    def __call__(self, node_features, coords, edge_index, node2inst, minimum_angle=0.0):
+    def __call__(
+        self,
+        node_features,  # Node features of shape (n, in_channels)
+        coords,  # Global coordinates of shape (n, 3)
+        edge_index,  # Edge index of shape (2, m)
+        node2inst,  # Node-to-instance mapping of shape (n,)
+        minimum_angle=0.0,  # Minimum angle to filter nearly colinear triangles (default: 0.0)
+    ):
         n = coords.shape[0]
         adj = (
             torch.sparse_coo_tensor(
@@ -76,18 +96,43 @@ def __call__(self, node_features, coords, edge_index, node2inst, minimum_angle=0
         tri2inst = node2inst[tri2node]
         tri_features = node_features[tri2node]
         return SimpleNamespace(
-            adj=adj,
-            ijk=ijk,
-            jkl=jkl,
-            coords=coords,
-            hood_coords=hood_coords,
-            tri_features=tri_features,
-            bases=bases,
-            i=i,
-            j=j,
-            triangles=triangles,
-            node2inst=node2inst,
-            tri2node=tri2node,
-            tri2edge=tri2edge,
-            tri2inst=tri2inst,
+            adj=adj,  # Sorted adjacency matrix of shape (n, n)
+            source=ijk,  # Lifted source triangles (m,)
+            target=jkl,  # Lifted target triangles (m,)
+            coords=coords,  # Global coordinates of shape (n, 3)
+            hood_coords=hood_coords,  # Local coordinates of shape (m, 3)
+            features=tri_features,  # Triangle features of shape (m, in_channels)
+            bases=bases,  # Bases of shape (m, 3, 3)
+            i=i,  # Node indices i (m,)
+            j=j,  # Node indices j (m,)
+            triangles=triangles,  # Tuples of i,j,k (m, 3)
+            node2inst=node2inst,  # Node-to-instance mapping of shape (n,)
+            lifted2node=tri2node,  # Triangle-to-node mapping of shape (m,)
+            lifted2edge=tri2edge,  # Triangle-to-edge mapping of shape (m,)
+            lifted2inst=tri2inst,  # Triangle-to-instance mapping of shape (m,)
         )
+
+
+class Lower:
+    """
+    Parameter-free module to lower a lifted geometric graph back to the input graph. Given edge/triangle features and the corresponding index it aggregates the features to the input graph.
+    """
+
+    def __init__(self, agg="mean"):
+        assert agg in ["sum", "mean", "max", "softmax"]
+        self.agg = agg
+
+    def __call__(self, features, index, size, return_index=False):
+        if self.agg == "sum":
+            return scatter_sum(features, index, dim_size=size, dim=0)
+        elif self.agg == "mean":
+            return scatter_mean(features, index, dim_size=size, dim=0)
+        elif self.agg == "max":
+            val, idx = scatter_max(features, index, dim_size=size, dim=0)
+            if return_index:
+                return val, idx
+            else:
+                return val
+        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)