add MeshCNN conv/pool/unpool layers

point2mesh/layers.py

@@ -0,0 +1,269 @@
+"""
+MeshCNN layers — convolution, pooling and unpooling on triangle meshes.
+
+Convolution works on edges: each edge has 4 topological neighbors from its
+two incident faces.  Symmetric aggregation removes the face-ordering
+ambiguity.  Pooling collapses edges by L2-norm priority; unpooling restores
+them from stored history.
+"""
+
+from __future__ import annotations
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import numpy as np
+from typing import List, Optional
+from .mesh import Mesh
+
+
+# ──────────────────────────────────────────────────────────────────────
+#  MeshConv
+# ──────────────────────────────────────────────────────────────────────
+class MeshConv(nn.Module):
+    """
+    Edge-based convolution (MeshCNN §4.1).
+
+    For each edge *e* with 4 neighbors (a, b, c, d) we form 5 inputs:
+        [e, |a−c|, a+c, |b−d|, b+d]
+    and apply a learned linear combination (via Conv2d with kernel (1,5)).
+    """
+
+    def __init__(self, in_ch: int, out_ch: int, bias: bool = True):
+        super().__init__()
+        self.conv = nn.Conv2d(in_ch, out_ch, kernel_size=(1, 5), bias=bias)
+
+    def forward(self, x: torch.Tensor, mesh: Mesh) -> torch.Tensor:
+        """
+        x    : (1, C_in, N_e)
+        mesh : Mesh with .gemm_edges [N_e, 4]
+        Returns: (1, C_out, N_e)
+        """
+        # Gather the 4 neighbor features
+        G = self._gather_neighbors(x, mesh)  # (1, C, N_e, 4)
+
+        # Symmetric aggregation
+        a, b, c, d = G[..., 0], G[..., 1], G[..., 2], G[..., 3]
+        sym = torch.stack([
+            torch.abs(a - c),
+            a + c,
+            torch.abs(b - d),
+            b + d,
+        ], dim=-1)  # (1, C, N_e, 4)
+
+        # Concatenate center edge + 4 symmetric descriptors → width-5
+        x_5 = torch.cat([x.unsqueeze(-1), sym], dim=-1)  # (1, C, N_e, 5)
+        out = self.conv(x_5)  # (1, C_out, N_e, 1)
+        return out.squeeze(-1)
+
+    @staticmethod
+    def _gather_neighbors(x: torch.Tensor, mesh: Mesh) -> torch.Tensor:
+        """Gather features of the 4 neighbor edges for every edge."""
+        # x: (1, C, N_e)
+        B, C, N_e = x.shape
+        gemm = mesh.gemm_edges  # (N_e, 4) on mesh.device
+        # Clamp to handle any −1 (boundary mirror already filled, but be safe)
+        gemm = gemm.clamp(min=0)
+        flat = gemm.reshape(-1)  # (N_e*4,)
+        gathered = x[:, :, flat]  # (1, C, N_e*4)
+        return gathered.view(B, C, N_e, 4)
+
+
+# ──────────────────────────────────────────────────────────────────────
+#  MeshPool  (edge collapse)
+# ──────────────────────────────────────────────────────────────────────
+class MeshPool(nn.Module):
+    """
+    Mesh pooling via edge collapse (MeshCNN §4.2).
+
+    Edges are prioritised by L2-norm of their feature vector; the
+    smallest-norm edges are collapsed first.  After each collapse the
+    features of the two resulting edges are set to the average of the
+    three merged edges.
+
+    The collapse history is stored so that `MeshUnpool` can invert the
+    operation.
+    """
+
+    def __init__(self, target: int):
+        """target : number of edges to keep after pooling."""
+        super().__init__()
+        self.target = target
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        mesh: Mesh,
+    ) -> tuple[torch.Tensor, Mesh, dict]:
+        """
+        x     : (1, C, N_e)
+        mesh  : current Mesh
+        Returns
+        -------
+        x_pooled : (1, C, target)
+        mesh_new : Mesh with updated topology (MUTATED)
+        history  : dict consumed by MeshUnpool
+        """
+        B, C, N_e = x.shape
+        device = x.device
+
+        # Work on CPU numpy for topology manipulation (small meshes are fast)
+        gemm = mesh.gemm_edges.cpu().numpy().copy()  # (N_e, 4)
+        edges_np = mesh.edges.cpu().numpy().copy()
+        feat = x.squeeze(0).detach().cpu().numpy().copy()  # (C, N_e)
+
+        active = np.ones(N_e, dtype=bool)
+        n_active = int(active.sum())
+
+        # Priority: L2 norm of each edge's feature vector
+        norms = np.linalg.norm(feat, axis=0)  # (N_e,)
+
+        # Sorted order of edges by ascending norm
+        order = np.argsort(norms)
+
+        # History bookkeeping for unpooling
+        # Maps: new_edge_idx → set of old edge indices that contributed
+        merge_log: List[tuple] = []  # (surviving_edge, [merged edges], [merge weights])
+
+        collapse_map = np.arange(N_e)  # edge redirect after collapses
+
+        idx = 0
+        while n_active > self.target and idx < len(order):
+            e = order[idx]
+            idx += 1
+            if not active[e]:
+                continue
+
+            a, b, c, d = gemm[e]
+            # Validity checks
+            if a < 0 or b < 0 or c < 0 or d < 0:
+                continue
+            if not (active[a] and active[b] and active[c] and active[d]):
+                continue
+            # Non-manifold guard: skip if collapsing would merge two boundary verts
+            # (simplified: skip if any neighbor is already dead or re-targeted)
+            if a == b or c == d:
+                continue
+
+            # Collapse edge e:
+            #   Face-0 edges (a, b) → surviving edge p
+            #   Face-1 edges (c, d) → surviving edge q
+            #   Merged features: p = avg(e, a, b), q = avg(e, c, d)
+            p, q = b, d  # surviving edge labels (keep the "second" edge of each face)
+
+            # Update features (on numpy)
+            feat[:, p] = (feat[:, e] + feat[:, a] + feat[:, b]) / 3.0
+            feat[:, q] = (feat[:, e] + feat[:, c] + feat[:, d]) / 3.0
+
+            merge_log.append((p, [e, a, b]))
+            merge_log.append((q, [e, c, d]))
+
+            # Deactivate collapsed edges
+            active[e] = False
+            active[a] = False
+            active[c] = False
+            n_active -= 3
+
+            # Redirect any neighbor pointers that point to a or c
+            gemm[gemm == a] = p
+            gemm[gemm == c] = q
+            gemm[gemm == e] = p  # default redirect to p
+
+        # Build new compact edge set
+        kept = np.where(active)[0]
+        old2new = np.full(N_e, -1, dtype=np.int64)
+        for new_i, old_i in enumerate(kept):
+            old2new[old_i] = new_i
+
+        # Re-index gemm for surviving edges
+        new_gemm = gemm[kept].copy()
+        for i in range(new_gemm.shape[0]):
+            for j in range(4):
+                mapped = old2new[new_gemm[i, j]]
+                new_gemm[i, j] = mapped if mapped >= 0 else i  # self-loop fallback
+
+        # Build new feature tensor (differentiable path)
+        kept_t = torch.tensor(kept, dtype=torch.long, device=device)
+        x_pooled = x[:, :, kept_t]  # (1, C, n_kept)
+
+        # Overwrite collapsed features differentiably
+        # We re-run the averaging in torch for grad flow
+        x_work = x.squeeze(0)  # (C, N_e)
+        new_feats = []
+        for old_i in kept:
+            new_feats.append(x_work[:, old_i])
+        # Override with merged averages
+        merge_map = {}
+        for surv, sources in merge_log:
+            if surv in merge_map:
+                continue  # keep first
+            merge_map[surv] = sources
+        new_feat_list = []
+        for ni, old_i in enumerate(kept):
+            if old_i in merge_map:
+                srcs = merge_map[old_i]
+                avg = sum(x_work[:, s] for s in srcs) / len(srcs)
+                new_feat_list.append(avg)
+            else:
+                new_feat_list.append(x_work[:, old_i])
+        x_pooled = torch.stack(new_feat_list, dim=1).unsqueeze(0)  # (1, C, n_kept)
+
+        # Construct new Mesh from surviving edges/faces
+        # (for simplicity we update the mesh in-place rather than rebuild faces)
+        mesh_new = mesh.clone()
+        mesh_new.gemm_edges = torch.tensor(new_gemm, dtype=torch.long, device=mesh.device)
+        mesh_new.n_edges = len(kept)
+        # Keep edges array updated (vertex indices of surviving edges)
+        mesh_new.edges = mesh.edges[kept_t.to(mesh.device)]
+        mesh_new.edge_v0 = mesh_new.edges[:, 0]
+        mesh_new.edge_v1 = mesh_new.edges[:, 1]
+
+        history = {
+            "kept": kept,           # indices of surviving edges in old ordering
+            "old2new": old2new,     # old_edge → new_edge mapping
+            "merge_log": merge_log, # how features were merged
+            "n_old": N_e,
+        }
+        return x_pooled, mesh_new, history
+
+
+# ──────────────────────────────────────────────────────────────────────
+#  MeshUnpool  (restore topology from history)
+# ──────────────────────────────────────────────────────────────────────
+class MeshUnpool(nn.Module):
+    """
+    Restore the pre-pooling edge topology using stored history.
+
+    Unpooled edge features are set to the feature of the surviving edge
+    they were merged into (broadcast).
+    """
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        history: dict,
+    ) -> torch.Tensor:
+        """
+        x       : (1, C, N_pooled)
+        history : dict from MeshPool.forward
+        Returns : (1, C, N_old)
+        """
+        B, C, N_pooled = x.shape
+        N_old = history["n_old"]
+        device = x.device
+        kept = history["kept"]
+
+        out = torch.zeros(B, C, N_old, device=device, dtype=x.dtype)
+
+        # Place surviving edge features at their original indices
+        kept_t = torch.tensor(kept, dtype=torch.long, device=device)
+        out[:, :, kept_t] = x
+
+        # For collapsed edges, copy from the surviving edge they merged into
+        for surv, sources in history["merge_log"]:
+            surv_new = int(history["old2new"][surv])
+            for s in sources:
+                if s not in kept:
+                    out[:, :, s] = x[:, :, surv_new]
+
+        return out