Skip to content

Birth and death

jax_morph.physics.Death #

Death(*, score_by_default=True)

Discrete stochastic removal: each alive cell dies as an independent Bernoulli event.

An alive cell is removed over a macro-step with probability

\[p = 1 - e^{-\lambda\,\Delta t},\]

where \(\lambda\) is the per-cell death_rate. death records each slot removed during the macro-step for lineage postprocessing, while died and die_eligible are the ephemeral scored trace. When composed with Division, this step must follow it: the discrete phase performs divide-then-die, and deferring reuse of slots freed by death until the next macro-step preserves the death lineage record.

Attributes:

  • score_by_default –

    Whether default trajectory scoring includes death. Defaults to True.

Methods:

  • sample_trace –

    Draw death actions for eligible cells.

  • replay –

    Apply the recorded removal events.

  • logp –

    Score recorded death actions.

Parameters:

  • score_by_default –

    Whether default trajectory scoring includes this step. Defaults to True.

state_reads #

state_reads()

Read the per-cell death hazard rate.

state_writes #

state_writes()

Write the alive mask and persistent per-step death record.

trace_writes #

trace_writes()

Record the scored action and alive-at-decision eligibility mask.

sample_trace #

sample_trace(state, *, dt, key)

Draw forward-exact, backward-straight-through death actions.

Parameters:

  • state –

    Pre-step state with alive and death_rate arrays of shape (capacity,).

  • dt –

    Macro-step duration.

  • key –

    JAX PRNG key.

Returns:

  • –

    Trace dictionary with died and die_eligible arrays of shape (capacity,).

replay #

replay(state, trace, *, dt, pathwise)

Apply recorded hard death actions and co-emit their trace.

Parameters:

  • state –

    Pre-step state.

  • trace –

    Recorded death action and eligibility arrays.

  • dt –

    Unused macro-step duration.

  • pathwise –

    Unused because the discrete event replays its recorded action.

Returns:

  • –

    Post-death state with alive, death, and trace fields updated.

logp #

logp(state, trace, dt)

Score recorded actions over cells alive at the death decision.

Parameters:

  • state –

    Live pre-step state conditioning death probabilities.

  • trace –

    Recorded action and eligibility arrays.

  • dt –

    Macro-step duration.

Returns:

  • –

    Scalar sum of eligible-cell Bernoulli log-probabilities.

jax_morph.physics.Division #

Division(
    *,
    n_space_dim,
    orientation_snr=0.0,
    score_by_default=True,
)

Discrete stochastic step: each alive cell divides as an independent Bernoulli event.

A cell divides over a macro-step with probability

\[p = 1 - e^{-\lambda\,\Delta t},\]

where \(\lambda\) is the per-cell division_rate (a constant initial condition, or an upstream controller's output - size / crowding / morphogen gating all live there). Daughters conserve volume, each taking radius \(r\,2^{-1/d}\), and are placed just touching along an orientable direction, inheriting the mother's heritable fields. orientation_snr (the amplitude signal-to-noise ratio; 0 -> isotropic) sharpens placement toward the per-cell division_axis. Lineage is recorded as born + mother; overflow past capacity is capped into the global division_overflow counter. When composed with Death, this step must precede it: the discrete phase performs divide-then-die, and deferring reuse of slots freed by death until the next macro-step preserves the death lineage record.

Attributes:

  • n_space_dim –

    Static spatial dimension used for division_axis and division_dir, both shaped (capacity, n_space_dim).

  • orientation_snr –

    Placement-orientation signal-to-noise amplitude. Defaults to 0.0.

  • score_by_default –

    Whether default trajectory scoring includes division. Defaults to True.

Methods:

  • sample_trace –

    Draw division actions and placement directions.

  • replay –

    Apply the recorded division events.

  • logp –

    Score recorded division actions.

Parameters:

  • n_space_dim –

    Spatial dimension; must be 1, 2, or 3.

  • orientation_snr –

    Placement-orientation signal-to-noise amplitude. Defaults to 0.0.

  • score_by_default –

    Whether default trajectory scoring includes this step. Defaults to True.

Raises:

  • ValueError –

    If n_space_dim is not 1, 2, or 3.

state_reads #

state_reads()

Reads the per-cell hazard rate and the optional per-cell orientation axis.

state_writes #

state_writes()

Writes the base fields, the overflow diagnostic, and the lineage fields.

trace_writes #

trace_writes()

Ephemeral trace: the 0/1 action, the eligibility mask, and the realized direction.

sample_trace #

sample_trace(state, *, dt, key)

Draw division actions and oriented placement directions.

Parameters:

  • state –

    Pre-step state with cell arrays of leading shape (capacity,).

  • dt –

    Macro-step duration.

  • key –

    JAX PRNG key.

Returns:

  • –

    Trace dictionary containing divided and divide_eligible arrays of shape

  • –

    (capacity,) and division_dir of shape (capacity, n_space_dim).

replay #

replay(state, trace, *, dt, pathwise)

Apply recorded division events and co-emit their trace.

Parameters:

  • state –

    Pre-step state.

  • trace –

    Recorded action, eligibility, and direction arrays.

  • dt –

    Unused macro-step duration.

  • pathwise –

    Unused because the discrete event always replays its recorded action.

Returns:

  • –

    Post-division state with physical outputs and trace fields updated.

Raises:

  • ValueError –

    If the state's spatial dimension differs from n_space_dim.

logp #

logp(state, trace, dt)

Score recorded division actions over cells alive at the decision.

Parameters:

  • state –

    Live pre-step state conditioning division probabilities.

  • trace –

    Recorded action and eligibility arrays.

  • dt –

    Macro-step duration.

Returns:

  • –

    Scalar sum of eligible-cell Bernoulli log-probabilities.

jax_morph.physics.reconstruct_lineage #

reconstruct_lineage(born, mother, alive, death=None)

Rebuild a cell-lineage graph from a recorded simulation history (postprocessing, not jitted).

A cell's identity in a running simulation is its state-array slot, and slots are reused after a cell dies, so a stable lineage needs the per-step birth record. A complete simulation history begins with its initial state: every slot alive in frame zero is a founder, and every later born slot mints a fresh node whose parent is the node currently occupying the recorded mother slot. When a matching death history is provided, each node is annotated with the macro-step when it was removed. Births are processed before deaths within a frame, matching the required Division-then-Death order.

Parameters:

  • born –

    (n_steps + 1, capacity) new-daughter flags over a complete history.

  • mother –

    (n_steps + 1, capacity) parent slot indices (sentinel -1).

  • alive –

    (n_steps + 1, capacity) alive masks (frame zero seeds founders).

  • death –

    Optional (n_steps + 1, capacity) per-step death flags. Defaults to None.

Returns:

  • –

    A list of nodes ordered by id, each a dict with id (unique int), parent (parent id, or

  • –

    None for a founder), slot (the array slot it was born into), birth_step, and

  • –

    death_step (or None when no death was recorded).