Skip to content

State

jax_morph.StateFieldSpec #

StateFieldSpec(
    name: str,
    shape: tuple = (),
    dtype: object = None,
    heritable: bool = True,
    default: float = 0.0,
    scope: str = 'cell',
)

Declaration of one state field: its name, how to allocate it, and its division behaviour.

A spec is self-describing - it carries its own name - so steps declare inputs/outputs as plain tuples of specs. Calling a spec returns a copy with the given attributes overridden, so a prefilled base spec can be customised inline, e.g. POSITION(heritable=False).

Attributes:

  • name –

    The field's name; becomes a real attribute on the generated state class.

  • shape –

    Trailing array shape. A cell-scope field has full shape (capacity, *shape); a global field has shape shape. Defaults to ().

  • dtype –

    Array dtype. Defaults to None (the JAX default float).

  • heritable –

    Whether a daughter cell inherits the mother's value on division ('cell' scope only). True marks a heritable property of the cell itself (size, type, internal chemical state): the daughter is born with a copy of the mother's value. False marks a transient or contextual quantity (a recorded action, a per-step signal): the daughter starts at default instead. Defaults to True.

  • default –

    Fill value used to allocate empty states, to reset ephemeral trace fields each macro-step, and to initialise non-inherited (heritable=False) fields of newborn daughters. Defaults to 0.0.

  • scope –

    'cell' (per-cell: leading capacity axis, alive-masked, resized on division) or 'global' (not indexed by cell: simulation time, scalars, grids). Defaults to 'cell'.

__call__ #

__call__(**overrides: object) -> StateFieldSpec

Return a copy of this specification with selected attributes overridden.

Parameters:

  • **overrides (object) –

    Field attributes to replace.

Returns:

  • StateFieldSpec –

    A new specification containing the requested replacements.

Raises:

  • TypeError –

    If an override does not name a specification field.

jax_morph.BaseState #

BaseState(
    specs: Mapping, displacement: Callable, shift: Callable
)

Base class for a model's state; generated subclasses add the fields as real attributes.

Carries the static schema (specs) and space (displacement/shift) and the shared field-access helpers. Instances are always of a build_state_from_model subclass whose per-cell / global fields are real jax.Array attributes (state.position, state.custom_field), so this base declares none of them itself.

Attributes:

  • specs –

    Static mapping from field names to StateFieldSpec declarations. Generated subclasses add one array attribute per entry, with shape (capacity, *spec.shape) for cell-scope fields and spec.shape for global fields.

  • displacement –

    Static binary function mapping two positions of shape (dim,) to a displacement of shape (dim,).

  • shift –

    Static binary function applying a displacement of shape (dim,) to a position of shape (dim,).

Methods:

  • set –

    Functionally replace one field.

  • update –

    Functionally replace several fields.

  • deltas –

    Build a sparse dynamic-step update.

n_cells property #

n_cells

Number of cell slots (the allocated capacity).

n_space_dim property #

n_space_dim

Spatial dimension.

set #

set(name: str, value: Array) -> BaseState

Return a copy with one field replaced.

Parameters:

  • name (str) –

    Name of the state field to replace.

  • value (Array) –

    Replacement array.

Returns:

  • BaseState –

    A same-typed state containing the replacement.

update #

update(**kw: Array) -> BaseState

Return a copy with several fields replaced by keyword.

Parameters:

  • **kw (Array) –

    Field names and replacement arrays.

Returns:

  • BaseState –

    A same-typed state containing all replacements.

deltas #

deltas(**fields: Array) -> BaseState

Return a sparse delta state: the named fields set, every other field None.

The result is a same-typed pytree that carries this state's static schema and space but holds an array only for the fields passed in (each a per-dt increment) and None everywhere else - the Equinox "updates" idiom. A dynamic step returns one of these; the model sums them across dynamic writers and applies them with eqx.apply_updates, so untouched fields (None) are left exactly as they were and bool/discrete base fields like alive are never accidentally integrated.

Parameters:

  • **fields (Array) –

    Per-dt increments keyed by field name.

Returns:

  • BaseState –

    A same-typed sparse state with None in every untouched array leaf.

jax_morph.build_state_from_model #

build_state_from_model(
    model: object, name: str | None = None
) -> type[jax_morph.core.state.BaseState]

Generate (and cache) a BaseState subclass for model with fields as real attributes.

The full schema is the always-present base fields (the state's) merged with the model's state_requires(). The returned class has two constructors: the plain field-wise StateCls(**field_arrays, space=...) and StateCls.init_empty(capacity=..., n_space_dim=..., n_types=..., space=...) (allocate an empty state); seed an initial condition with state.update(...). name sets the class name and is capitalized ('elongation' -> ElongationState); default is State.

Parameters:

  • model (object) –

    Model whose declared field requirements define the schema.

  • name (str | None, default: None ) –

    Optional generated class name or prefix. Defaults to None (the class is named State).

Returns:

  • type[BaseState] –

    A cached BaseState subclass with one array attribute per required field.

Raises:

  • ValueError –

    If model steps declare conflicting specifications for one field.