API Reference

This reference maps the curated public surface of jax_morph. For end-to-end examples, start with Basic Usage or the example notebooks.

Core#

StateFieldSpec and BaseState declare and hold model state; build_state_from_model synthesizes a typed state class.
Geometry utilities: displacement, distance, and neighbourhood helpers.
SimulationStep, StochasticStep, and StepType define validated steps; check_stochastic_step validates a stochastic step's trace contract.
Model: compose steps into a validated simulation pipeline.
simulate: sample a pathwise-differentiable rollout.
trajectory_logp and transition_logp: score sampled rollouts.
Autodiff utilities: straight-through and numerical helpers.
Field constants: prefilled state and physics field specs for custom steps.

Interaction potentials: force-free, pairwise, Morse, compact-repulsion, harmonic, and Lennard-Jones energy laws.
Relaxation: MechanicalRelaxation and relax_equilibrium, FIRE relaxation with implicit equilibrium gradients.
Dynamics: BrownianDynamics and ActiveBrownianDynamics2D, scorable stochastic motion.
VirialStress: per-cell mechanical stress sensing.
FreeScreenedDiffusion: screened-diffusion signalling.
SaturatingCellGrowth: saturating radial growth.
Birth and death: Death, Division, and reconstruct_lineage, discrete population events and lineage recovery.

ODEController: base class for continuous-time controllers.
GeneNetworkConnectionist, NeuralODE, and GeneNetworkMWC: concrete controller families.

TrajectoryRecord: a loaded complete trajectory and its step size.
save_model and load_model: persist model numeric leaves against a caller-provided template.
save_state and load_state: persist generated state snapshots.
save_trajectory and load_trajectory: persist complete histories with metadata.

Visualization: static cell-cluster rendering, trajectory animation, and per-cell field time series. Install the optional viz extra to render with matplotlib.