Reference

waxMorph’s top-level modules default to the PyTorch backend and expose the graph-network-based simulator, contact-graph construction, shape losses, and the training entry point. Explicit JAX counterparts live under waxmorph.jax.

Module map

waxmorph.data

Loads, normalizes, and volumetrically samples meshes. Provides helpers for source-target shape pairs and for the intermediate target sequences used as supervision frames.

waxmorph.graph

Builds the contact graph induced by spatial proximity. Node features carry per-cell signaling-molecule concentrations; edge features encode pairwise distance and relative polarity. Re-exports the PyTorch implementation.

waxmorph.gnn

The graph-network-based simulator (GNS) that learns local, neighbor-dependent update rules over the contact graph.

waxmorph.train

Trains the emulator with a fixed agent count by coupling learned updates with differentiable physical constraints. Re-exports the PyTorch training API.

waxmorph.losses

Point-cloud shape losses for assigned (squared) and unassigned (Chamfer, optimal-transport) target morphologies.

waxmorph.simulator and waxmorph.emulator

Warp kernels for explicit mechanochemical forward simulation and for the differentiable physics corrections used in rollouts with a fixed agent count. The simulator integrates soft-sphere mechanics, reaction-diffusion, growth, and division; the emulator freezes neighbor topology within a step so gradients propagate through the physics.

waxmorph.render

Static, interactive, and movie renderers for trajectories of spheroidal agents.

The backend-specific APIs live under waxmorph.torch and waxmorph.jax.

State and notation

waxMorph represents a tissue as spheroidal agents carrying positions X, polarities P, radii R, and signaling-molecule concentrations c; the graph-network processor predicts increments dX, dP, and dc over these fields. See Cells as spheroidal agents for the field shapes and index conventions.

For a worked build_graph → GNS → train flow, see the mesh-to-emulator guide. JAX provides backend-specific graph returns and padding, loss families, PRNG, optimizer and checkpoint behavior, and CUDA execution for Warp physics; see Choose a learning backend.

Training and external APIs

• Training APIs
  • Post-update selection
  • Checkpoint and backend contracts
• External API references

Module index

waxmorph

Differentiable morphogenesis and shape assembly with Warp physics.