Skip to content

jax-morph core walkthrough#

A minimal tour of the hardened core built so far (Phase 1): a typed state generated from a model, the uniform step contract, the hybrid macro-step, score-function logp, and end-to-end differentiability. No physics or control modules yet - those come in later phases - so the "cell behaviours" here are deliberately toy steps whose only job is to exercise the machinery.

We cover, in order: steps and the model, generating a state, accessing fields, functional updates, geometry, autodiff utilities, dataflow validation, simulate, transition scoring, and gradients.

import jax
import jax.numpy as jnp

# The library is written for float64; enable it before anything touches JAX.
jax.config.update('jax_enable_x64', True)

import equinox as eqx

import jax_morph as jxm

jxm.__version__

'0.4.0'

1. Steps and the model#

A step is an eqx.Module that declares its step_type (a StepType) and its data dependencies, and implements one uniform method, __call__(state, *, dt, key) -> state (dt and key are keyword-only). The three step types differ only in what the returned state means, which the model interprets:

  • QUASISTATIC - the returned state is the new state (an instantaneous constraint/solve).
  • DISCRETE - the returned state is the new state (a discrete jump).
  • DYNAMIC - the step returns a sparse delta state via state.deltas(field=increment): the written fields hold that field's dt-scaled increment and every untouched field is None. The model sums these across dynamic writers, so a dynamic step returns only its own contribution.

A step documents its inputs and outputs as tuples of StateFieldSpec (each spec carries its own name): state_reads() for the fields it consumes and state_writes() for the fields it produces. For base fields (position, radius, celltype, alive, t) reuse the prefilled specs jxm.POSITION, jxm.RADIUS, ...; for a new field pass a jxm.StateFieldSpec('name', shape=...). state_requires() is the derived union. The step does not care who else touches a field; consistency is the Model's job. dt and key are keyword-only, so a step that ignores one or both still accepts them and the model calls every step the same way.

Below are two toy steps: a quasistatic Sense that writes a per-cell signal (distance from the cluster centroid) and a dynamic Grow that grows every radius at a constant rate.

class Sense(jxm.SimulationStep):
    step_type = jxm.StepType.QUASISTATIC

    def state_reads(self):
        # inputs as a tuple of specs; base fields reuse their prefilled specs
        return (jxm.POSITION, jxm.RADIUS, jxm.ALIVE)

    def state_writes(self):
        # a StateFieldSpec introduces a new field; heritable=False -> a daughter does not
        # inherit the mother's value on division (it starts at the spec default instead)
        return (jxm.StateFieldSpec('signal', shape=(), heritable=False),)

    def __call__(self, state, *, dt, key):
        # quasistatic: return the new state (this step ignores dt and key)
        alive = state.alive.astype(state.radius.dtype)
        n_alive = jnp.clip(alive.sum(), 1.0, None)
        com = (state.position * alive[:, None]).sum(0) / n_alive
        dist = jxm.ad_utils.safe_norm(state.position - com, axis=-1)
        return state.set('signal', dist * alive)


class Grow(jxm.SimulationStep):
    step_type = jxm.StepType.DYNAMIC
    rate: float  # plain field: a Python float is static, a jax.Array is a trainable parameter

    def state_writes(self):
        return (jxm.RADIUS,)  # a base field: reuse its prefilled spec

    def __call__(self, state, *, dt, key):
        # dynamic: return a sparse delta state holding this step's own dt-scaled increment
        return state.deltas(radius=self.rate * dt * jnp.ones_like(state.radius))


model = jxm.Model([Sense(), Grow(rate=0.5)])
model

Model(steps=(Sense(), Grow(rate=0.5)))

2. Generate a state, then set an initial condition#

build_state_from_model(model, name) generates (and caches) a jxm.BaseState subclass whose fields are the base fields plus everything the model requires - all as real attributes. name is capitalized into the class name ('cluster' -> ClusterState). It has two constructors: the plain field-wise ClusterState(position=..., radius=..., ...) (pass every field array directly), and ClusterState.init_empty(capacity=..., n_space_dim=..., n_types=...), which allocates every field to its default (per-cell fields get a leading capacity axis, globals like time t do not), giving an empty (no live cells) state. Base and custom fields are treated identically; there is no privileged base-field seeding.

Setting an initial condition is a separate, simulation-specific step - just state.update(...). The common "seed a live cluster" pattern is a tiny helper you keep in your own script (seed below); we deliberately don't bake it into the library.

Note the distinction: model.state_requires() is what the model touches, while the state also always carries the base fields (they're the state's, not the model's).

def seed(state, n, *, radius=0.5, key=None):
    # A simulation-specific helper (kept in user code, not the library): seed the first n slots as a
    # live cluster. Every field - base or custom - is just set with update().
    updates = dict(
        alive=state.alive.at[:n].set(True),
        radius=state.radius.at[:n].set(radius),
        celltype=state.celltype.at[:n, 0].set(1.0),
    )
    if key is not None:
        updates['position'] = state.position.at[:n].set(
            radius * jax.random.normal(key, (n, state.n_space_dim))
        )
    return state.update(**updates)


key = jax.random.PRNGKey(0)
ClusterState = jxm.build_state_from_model(model, name='cluster')  # capitalized -> ClusterState
empty = ClusterState.init_empty(capacity=8, n_space_dim=2, n_types=1)  # allocate an empty state
print('right after init_empty, live cells:', int(empty.alive.sum()))  # 0 - an empty state

state = seed(empty, 4, radius=0.5, key=key)  # the initial condition: a 4-cell live cluster
print('generated class:', ClusterState.__name__, '| jxm.BaseState subclass:',
      issubclass(ClusterState, jxm.BaseState))
print('model requires:', sorted(sp.name for sp in model.state_requires()))  # the model's own fields
print('full schema    :', sorted(state.specs))              # + always-present base fields
print('n_cells:', state.n_cells, '| n_space_dim:', state.n_space_dim, '| t:', float(state.t))
print('live cells:', int(state.alive.sum()))

right after init_empty, live cells: 0

generated class: ClusterState | jxm.BaseState subclass: True
model requires: ['alive', 'position', 'radius', 'signal']
full schema    : ['alive', 'celltype', 'position', 'radius', 'signal', 't']
n_cells: 8 | n_space_dim: 2 | t: 0.0
live cells: 4

Two ways to build a state#

init_empty allocates the arrays for you. The field-wise constructor takes them directly - every schema field (base and custom, like signal) passed by keyword, space optional and the static schema filled in automatically. It's the low-level primitive init_empty is built on, and what you'd reach for when you already hold the arrays (a checkpoint, a hand-built initial condition, a test). Below we round-trip the state we just seeded straight through its own arrays.

# every field must be supplied - omitting one (or passing an unknown name) is a clear TypeError
direct = ClusterState(
    position=state.position,
    radius=state.radius,
    celltype=state.celltype,
    alive=state.alive,
    t=state.t,
    signal=state.signal,   # the custom field is passed exactly like a base one
)
print('field-wise build -> live cells:', int(direct.alive.sum()))
print('round-trips the seeded state:',
      bool(jnp.array_equal(direct.alive, state.alive)),
      bool(jnp.allclose(direct.position, state.position)))

field-wise build -> live cells: 4
round-trips the seeded state: True True

Accessing fields#

Every field - base or declared - is a real attribute: state.position, state.signal. It is also reachable dynamically by name with state['name'] (handy for generic code). The static schema lives at state.specs.

print('state.position shape:', state.position.shape)
print('state.signal (declared, starts zero):', state.signal)   # <- attribute access
print("dict-style also works, state['signal']:", state['signal'].shape)
print('schema (state.specs):', sorted(state.specs))

state.position shape: (8, 2)
state.signal (declared, starts zero): [0. 0. 0. 0. 0. 0. 0. 0.]
dict-style also works, state['signal']: (8,)
schema (state.specs): ['alive', 'celltype', 'position', 'radius', 'signal', 't']

3. Functional updates#

The state is an immutable jxm.BaseState pytree. set(name, value) and update(**fields) return a new state; the original is untouched. This is what lets a whole rollout be a jax.lax.scan and stay differentiable.

A dynamic step instead returns a sparse delta state with deltas(**fields): the named fields are set and every other field is None (the Equinox "updates" idiom). The model sums these and applies them with eqx.apply_updates, so untouched fields are left exactly as they were - and bool/discrete base fields like alive are never accidentally integrated.

s2 = state.set('radius', state.radius + 1.0)
print('original radius[0]:', float(state.radius[0]))   # unchanged
print('updated  radius[0]:', float(s2.radius[0]))      # +1

s3 = state.update(radius=state.radius * 2, signal=state.signal + 1.0)
print('multi-field update radius[0]:', float(s3.radius[0]), '| signal[0]:', float(s3.signal[0]))

# a sparse delta state: only 'radius' is populated, everything else is None
d = state.deltas(radius=jnp.ones_like(state.radius))
print('delta radius set?', d.radius is not None, '| delta position None?', d.position is None)
print('apply_updates adds only the touched field:', float(eqx.apply_updates(state, d).radius[0]))

original radius[0]: 0.5
updated  radius[0]: 1.5
multi-field update radius[0]: 1.0 | signal[0]: 1.0
delta radius set? True | delta position None? True
apply_updates adds only the touched field: 1.5

4. Geometry#

The space (boundary conditions) is injected at state construction: free_space() by default, or periodic(box). pairwise_displacements(positions, displacement) is the fundamental primitive - the dense (N, N, dim) displacements; distances are safe_norm(disp, axis=-1) at the call site. neighbor_sum reduces a per-pair tensor over live neighbours (excluding self and dead cells).

disp = jxm.geometry.pairwise_displacements(state.position, state.displacement)
dist = jxm.ad_utils.safe_norm(disp, axis=-1)
print('pairwise displacements:', disp.shape, '| distances:', dist.shape)

# sum a constant vector over each cell's live neighbours -> (n_live - 1) for live cells
contrib = jnp.ones((state.n_cells, state.n_cells, 2))
print('neighbor_sum for cell 0:', jxm.geometry.neighbor_sum(contrib, state.alive)[0])

# a periodic box uses the minimum-image convention (space passed to the constructor)
pstate = ClusterState.init_empty(capacity=4, n_space_dim=2, n_types=1, space=jxm.geometry.periodic(10.0))
print('periodic displacement of (9,0)->(1,0):', pstate.displacement(jnp.array([9.0, 0.0]),
                                                                    jnp.array([1.0, 0.0])))

pairwise displacements: (8, 8, 2) | distances: (8, 8)

neighbor_sum for cell 0: [3. 3.]
periodic displacement of (9,0)->(1,0): [-2.  0.]

5. Autodiff utilities (jxm.ad_utils)#

Safe numerics, forward-exact / backward-smooth straight-through estimators, and the log-density helpers that pair with the samplers (sample_bernoulli_st <-> bernoulli_logp, sample_categorical_st <-> categorical_logp). The straight-through primitives are forward-exact / backward-smooth: the forward output is bit-identical to the hard operation and independent of any surrogate temperature; only the gradient uses the smooth surrogate.

ad = jxm.ad_utils

# safe_norm has value AND gradient 0 at the zero vector (a bare jnp.linalg.norm gives NaN grad)
print('grad of safe_norm at 0:', jax.grad(lambda v: ad.safe_norm(v))(jnp.zeros(3)))

# heaviside_st: forward is exactly (x > 0), for any temperature
print('heaviside_st([-1, 0.5, 2]):', ad.heaviside_st(jnp.array([-1.0, 0.5, 2.0])))

# an exact Bernoulli sample whose gradient w.r.t. p is the identity surrogate
print('sample_bernoulli_st(p=[0, 1, 0.5]):', ad.sample_bernoulli_st(key,
                                                                     jnp.array([0.0, 1.0, 0.5])))

# its paired log-density scores an observed outcome under p (safe at p = 0 or 1)
print('bernoulli_logp(outcome=[1,0], p=[0.3,0.3]):', ad.bernoulli_logp(jnp.array([1.0, 0.0]),
                                                                       jnp.array([0.3, 0.3])))

grad of safe_norm at 0: [0. 0. 0.]
heaviside_st([-1, 0.5, 2]): [0. 1. 1.]
sample_bernoulli_st(p=[0, 1, 0.5]): [0. 1. 0.]

bernoulli_logp(outcome=[1,0], p=[0.3,0.3]): [-1.2039728  -0.35667494]

6. Dataflow validation#

Model validates the field graph when it is constructed: a quasistatic field admits exactly one writer (dynamic writers instead accumulate), and no two steps may declare the same field name with different specs. Reads are unrestricted - a field may be read-only (fixed by the initial condition). Here two Sense steps both write signal, which is rejected up front.

try:
    jxm.Model([Sense(), Sense()])
except ValueError as err:
    print('ValueError:', err)

ValueError: Field 'signal' is written by Sense and again by Sense; a quasistatic field admits exactly one writer

7. Simulate#

simulate rolls the model out for n_steps macro-steps of size dt (physical time T = n_steps * dt), advancing the simulation time t. Each macro-step is one model(state, dt=dt, key=key) call - three phases: quasistatic constraints, then the summed dynamic deltas, then discrete events. dt is a rollout argument, not a property of the model. key is optional (last, defaulting to None): this model is deterministic so we omit it; a model with stochastic steps requires one. simulate returns the final BaseState, or with history=True the complete state history s0, s1, ..., sn stacked along a leading step axis. The forward pass is a pure sampler; scoring is a separate top-level driver call (next section).

final = jxm.simulate(model, state, n_steps=20, dt=0.1)  # deterministic model: no key needed
print('radius[0]: 0.5 ->', float(final.radius[0]), '| t ->', float(final.t))

traj = jxm.simulate(model, state, n_steps=20, dt=0.1, history=True)
print('trajectory radius shape (steps, cells):', traj.radius.shape)
print('radius[0] over initial state and first 5 steps:', traj.radius[:6, 0])

radius[0]: 0.5 -> 1.5000000000000009 | t -> 2.0000000000000004
trajectory radius shape (steps, cells): (21, 8)
radius[0] over initial state and first 5 steps: [0.5  0.55 0.6  0.65 0.7  0.75]

8. Scoring stochastic traces#

The forward pass samples; a separate evaluator scores. A stochastic step subclasses jxm.StochasticStep: instead of a forward __call__, it samples a trace - the exogenous noise plus the realized action - and writes it into declared trace fields (trace_writes(), analogous to state_writes() but ephemeral, reset each macro-step), so the trace rides in the trajectory automatically. _dist(state, dt) is the single source of the distribution's params, shared by sampling and scoring. transition_logp(model, state, next_state, dt) re-runs one macro-step on the detached observed state with the model parameters live, then replays each stochastic step from the trace recorded in next_state and sums the scored steps' logp. Each logp recomputes its params from the re-run state and scores the recorded action, which is stop_gradient-ed - so only the recomputed distribution carries gradient. Actions are co-emitted with the post-step state, so a transition is the pair (s_t, s_{t+1}).

For a sampled or observed complete history, use trajectory_logp(model, history, dt) instead. It reconstructs a live state across macro-step boundaries and returns one term per transition; call .sum() for the joint trajectory score. The initial frame is not an event record and may retain trace values from the input state, so actions belong to destination frames history[1:]. This one-step example uses the explicitly conditional driver because it is demonstrating a single transition.

Here MaybeDivide gives each live cell an independent chance p to divide. sample_trace draws the exact 0/1 outcome and records it (with a divide_eligible mask - which cells were alive to decide) into the trace fields; replay writes that recorded action back; and logp scores it over the eligible cells with the paired bernoulli_logp against p recomputed live. The score-function gradient d logp / d p reaches the policy parameter p.

class MaybeDivide(jxm.StochasticStep):
    step_type = jxm.StepType.DISCRETE
    p: jax.Array  # trainable division probability

    def trace_writes(self):
        # the recorded action and the alive-at-decision mask: ephemeral trace fields, reset each
        # macro-step, co-emitted into the post-step state so they ride the trajectory
        return (
            jxm.StateFieldSpec('divided', shape=(), heritable=False),
            jxm.StateFieldSpec('divide_eligible', shape=(), heritable=False),
        )

    def _dist(self, state, dt):
        # the single source of the distribution's params, used by BOTH sampling and scoring
        return self.p * jnp.ones_like(state.radius)

    def sample_trace(self, state, *, dt, key):
        divided = jxm.ad_utils.sample_bernoulli_st(key, self._dist(state, dt))  # exact 0/1 sample
        alive_f = state.alive.astype(divided.dtype)
        return {'divided': divided * alive_f, 'divide_eligible': alive_f}

    def replay(self, state, trace, *, dt, pathwise):
        # discrete: nothing to reparameterize, so pathwise is ignored; write the recorded action
        return state.update(divided=trace['divided'], divide_eligible=trace['divide_eligible'])

    def logp(self, state, trace, dt):
        prob = self._dist(state, dt)                 # recomputed live through the param
        # mask with the RECORDED eligibility, not state.alive: who was eligible to sample is a fact
        # of the forward pass (a real division changes alive mid-step), so it rides in the trace
        return jnp.sum(jxm.ad_utils.bernoulli_logp(trace['divided'], prob) * trace['divide_eligible'])


stoch = jxm.Model([MaybeDivide(p=jnp.array(0.7))])

DivideState = jxm.build_state_from_model(stoch, name='divide')
sstate = DivideState.init_empty(capacity=6, n_space_dim=2, n_types=1)
sstate = seed(sstate, 4, key=key)

s_next = stoch(sstate, dt=1.0, key=key)             # forward: sample + record the action
print('sampled divisions:', s_next.divided)
# score the transition (s_t, s_t+1): condition on sstate, read the action from s_next
print('logp of that transition:', float(jxm.transition_logp(stoch, sstate, s_next, 1.0)))

# score-function (REINFORCE-style) gradient: d logp / d p reaches the policy parameter
g = eqx.filter_grad(lambda m: jxm.transition_logp(m, sstate, s_next, 1.0))(stoch)
print('d logp / d p:', float(g.steps[0].p))

sampled divisions: [0. 1. 1. 0. 0. 0.]
logp of that transition: -3.1212954965293367
d logp / d p: -3.809523809523809

9. Differentiability#

The forward rollout is differentiable end to end. Two flavours:

  • w.r.t. the initial state - here d(final radius) / d(initial radius) = 1.
  • w.r.t. model parameters - make a parameter a jax.Array and it becomes trainable. A Python scalar stays static; a jax.Array is traced. eqx.filter_grad returns a Model-shaped pytree whose parameter leaves hold the gradients.
def final_radius(r0):
    s = state.set('radius', state.radius.at[0].set(r0))
    f = jxm.simulate(model, s, n_steps=10, dt=0.1)
    return f.radius[0]


print('d(final radius) / d(initial radius):', float(jax.grad(final_radius)(0.5)))

# a trainable growth rate: pass a jax.Array instead of a Python float
trainable = jxm.Model([Sense(), Grow(rate=jnp.array(0.5))])


def total_mass(m):
    f = jxm.simulate(m, state, n_steps=10, dt=0.1)
    return (f.radius * f.alive).sum()


grads = eqx.filter_grad(total_mass)(trainable)
print('d(total mass) / d(growth rate):', float(grads.steps[1].rate))

d(final radius) / d(initial radius): 1.0
d(total mass) / d(growth rate): 3.9999999999999996