sheaf-laplacian-coordination

Sheaf Laplacian Coordination

Trit: 0 (ERGODIC - coordinator) Color: Green (#26D826)

Overview

Implements sheaf neural network coordination using graph Laplacians for:

• Distributed consensus via sheaf diffusion
• Harmonic extension/restriction operators
• Spectral clustering on sheaf sections
• Multi-agent coordination with vector space representations

Key Papers

• Sheaf Neural Networks - Hansen & Gebhart 2020
• Neural Sheaf Diffusion - Bodnar et al. 2022
• Cooperative Sheaf Neural Networks - Ribeiro et al. 2025
• Sheaf Diffusion Goes Nonlinear - Zaghen et al. 2024

Core Concepts

Sheaf Laplacian

The sheaf Laplacian generalizes the graph Laplacian by associating vector spaces to nodes and linear maps to edges:

L_F = D^\top D

where D is the coboundary operator:
(Df)_e = F_{e,t} f_t - F_{e,s} f_s

F_{e,v} : F(v) → F(e)  (restriction maps)

Diffusion Process

Sheaf diffusion for consensus:

\frac{dx}{dt} = -L_F x

At equilibrium: L_F x = 0 (harmonic sections)

In/Out Degree Laplacians (Cooperative SNNs)

For directed graphs with cooperative behavior:

L_{in} = D_{in}^\top D_{in}   (gathering information)
L_{out} = D_{out}^\top D_{out} (conveying information)

API

Python Implementation

import torch
import torch.nn as nn

class SheafLaplacian(nn.Module):
    """Learnable sheaf Laplacian for graph coordination."""
    
    def __init__(self, num_nodes, stalk_dim, edge_index):
        super().__init__()
        self.num_nodes = num_nodes
        self.stalk_dim = stalk_dim
        self.edge_index = edge_index
        
        # Learnable restriction maps F_{e,v}
        num_edges = edge_index.shape[1]
        self.restriction_maps = nn.Parameter(
            torch.randn(num_edges, 2, stalk_dim, stalk_dim)
        )
    
    def build_laplacian(self):
        """Construct sheaf Laplacian from restriction maps."""
        L = torch.zeros(
            self.num_nodes * self.stalk_dim,
            self.num_nodes * self.stalk_dim
        )
        
        for e, (s, t) in enumerate(self.edge_index.T):
            F_es = self.restriction_maps[e, 0]  # Source restriction
            F_et = self.restriction_maps[e, 1]  # Target restriction
            
            # Add edge contribution to Laplacian
            # L[s,s] += F_es^T F_es, L[t,t] += F_et^T F_et
            # L[s,t] -= F_es^T F_et, L[t,s] -= F_et^T F_es
            
        return L
    
    def diffuse(self, x, steps=10, dt=0.1):
        """Run sheaf diffusion for consensus."""
        L = self.build_laplacian()
        for _ in range(steps):
            x = x - dt * (L @ x)
        return x
    
    def harmonic_extension(self, boundary_values, boundary_mask):
        """Extend boundary values harmonically."""
        L = self.build_laplacian()
        # Solve L_interior x_interior = -L_boundary x_boundary
        return solve_harmonic(L, boundary_values, boundary_mask)


class CooperativeSheafNN(nn.Module):
    """Cooperative SNN with in/out degree control."""
    
    def __init__(self, in_dim, hidden_dim, out_dim, edge_index):
        super().__init__()
        self.sheaf = SheafLaplacian(
            num_nodes=edge_index.max() + 1,
            stalk_dim=hidden_dim,
            edge_index=edge_index
        )
        self.encoder = nn.Linear(in_dim, hidden_dim)
        self.decoder = nn.Linear(hidden_dim, out_dim)
        
        # Cooperative gates: control gather vs convey
        self.gather_gate = nn.Parameter(torch.ones(1))
        self.convey_gate = nn.Parameter(torch.ones(1))
    
    def forward(self, x, edge_index):
        h = self.encoder(x)
        
        # Cooperative diffusion
        h_diffused = self.sheaf.diffuse(h)
        
        # Apply cooperative gates
        h_out = self.gather_gate * h + self.convey_gate * h_diffused
        
        return self.decoder(h_out)