waxmorph.torch.gnn

PyTorch Encode-Process-Decode graph network for cell-state updates.

Contacts are directed COO edges. waxmorph.jax.gnn is a behavioral counterpart with backend-specific architecture and serialization. Architecture: Sanchez-Gonzalez et al., “Learning to Simulate Complex Physics with Graph Networks” (ICML 2020).

Classes

GNS(node_feature_dim, edge_feature_dim[, ...])	Encode-process-decode graph network.
GraphNetworkBlock([node_latent_dim, ...])	Residual edge-then-node message passing.