Step primitives

jax_morph.SimulationStep #

SimulationStep()

Base class for all simulation steps. Subclasses set the class var step_type.

A step documents its inputs and outputs as tuples of StateFieldSpec (each spec carries its own name):

state_reads() -> the fields it consumes.
state_writes() -> the fields it produces.

Use the prefilled base specs (POSITION, RADIUS, ...) for base fields; the step does not care whether a field is base or who else touches it. state_requires() is the derived union state_reads U state_writes. Consistency and merging across steps is the Model's job.

Every step implements the same method, __call__(state, *, dt, key) -> state (dt and key are keyword-only); its return is interpreted by step_type:

QUASISTATIC / DISCRETE: return the new state (the constraint solve, or the jump, applied).
DYNAMIC: return state.deltas(**{field: increment}) - a sparse delta state whose written fields hold that field's increment over dt (the step bakes in its own dt-scaling) and whose every other field is None. The model sums these across dynamic writers, so a dynamic step returns only its own contribution, never a full state.

A deterministic step is exactly this. A stochastic step additionally samples and scores an action; it subclasses StochasticStep (below), which supplies is_stochastic = True and the trace / replay / logp contract.

Methods:

state_reads –

Declare state fields consumed by this step.
state_writes –

Declare physical state fields produced by this step.
state_requires –

Return the merged read/write schema.
__call__ –

Apply the step according to its declared step type.

is_stochastic `property` #

is_stochastic

Return whether this step samples an action and can contribute a score term.

Returns:

–

False for a deterministic step.

state_reads #

state_reads() -> tuple

Return field specifications this step reads.

Returns:

tuple –

A tuple of StateFieldSpec objects. Defaults to an empty tuple.

state_writes #

state_writes() -> tuple

Return physical field specifications this step writes.

Returns:

tuple –

A tuple of StateFieldSpec objects. Defaults to an empty tuple.

state_requires #

state_requires() -> tuple

Return the deduplicated union of read and written field specifications.

Returns:

tuple –

A tuple of StateFieldSpec objects, preserving first-seen order.

jax_morph.StochasticStep #

StochasticStep(*, score_by_default: bool = True)

Mixin for a stochastic step: it records a trace, replays it, and scores it.

A stochastic step samples a trace - the exogenous noise (a parameter-free draw, e.g. a standard-normal xi) plus the realized action (e.g. the 0/1 divided outcome, or a displacement dx) - and a single replay turns a trace into the step's effect. The forward __call__ samples a trace and replays it pathwise; the top-level scoring drivers read a recorded trace back out of the post-step state and replay it again with parameters live, scoring the selected steps.

The one per-instance knob is score_by_default: whether this step contributes its logp when a scoring driver is called with no score= override. Whether the step is reparameterizable - and so whether an unscored replay can carry a pathwise gradient - is not a knob but intrinsic to the distribution and lives inside replay: replay(..., pathwise=True) recomputes a reparameterizable value from the trace's exogenous noise with live parameters, while a discrete step has nothing to reparameterize and uses the realized action regardless. The model always passes pathwise = not scored.

Subclasses implement trace_writes (the recorded fields), sample_trace, replay(state, trace, *, dt, pathwise) and logp(state, trace, dt); trace_from_state has a default (read the declared trace fields back out of a state) and is overridden only for a bespoke layout. _dist is the recommended single source of the distribution's parameters, shared by sample_trace and logp so a step cannot sample under one distribution and score under another.

Attributes:

score_by_default –

Whether a scorer with score=None includes this stochastic step. Defaults to True.

Methods:

trace_writes –

Declare ephemeral trace fields for one macro-step.
sample_trace –

Draw exogenous noise and realized discrete actions.
trace_from_state –

Recover the recorded trace from a post-step state.
replay –

Apply the effect of a recorded trace.
logp –

Score a recorded trace under the live state.

state_requires #

state_requires() -> tuple

Return read, physical-write, and ephemeral-trace specifications.

Returns:

tuple –

A deduplicated tuple of StateFieldSpec objects.

trace_writes #

trace_writes() -> tuple

Field specs for the recorded trace: ephemeral, reset to their defaults each macro-step.

Trace fields are ordinary state fields (declared, allocated, co-emitted into the post-step state) but exist only to reconstruct and score this macro-step; the model resets them at each macro-step start. They are disjoint from state_writes - a stochastic output that must persist as physical state (read by a later macro-step) is a separate state_writes field written by replay, never a trace field, which the reset would wipe. A dynamic (additive) trace field must default to the additive identity 0 so that reset-then- accumulate records exactly that step's own value.

sample_trace #

sample_trace(
    state: BaseState, *, dt: float | Array, key: Array
) -> dict[str, jax.jaxlib._jax.Array]

Draw the trace entries the pathwise replay consumes; return them (a dict of arrays).

Must contain every entry replay(..., pathwise=True) reads: the exogenous noise, plus any directly-sampled (discrete) action. A derived realized value (a Brownian dx = mean + std * xi) may be omitted - the pathwise replay recomputes it from the noise and records it, so computing it here too would only duplicate the reparameterization. Used only by the forward __call__.

Parameters:

state (BaseState) –

Live pre-step state conditioning the draw.
dt (float | Array) –

Macro-step duration.
key (Array) –

JAX PRNG key for the draw.

Returns:

dict[str, Array] –

A dictionary from trace-field names to their sampled arrays.

Raises:

NotImplementedError –

Always, until a subclass supplies the sampler.

trace_from_state #

trace_from_state(
    state: BaseState,
) -> dict[str, jax.jaxlib._jax.Array]

Read this step's recorded trace back out of a (post-step) state.

The default reads every declared trace_writes field back by name; override only for a bespoke trace layout.

Parameters:

state (BaseState) –

Post-step state containing the recorded trace.

Returns:

dict[str, Array] –

A dictionary from declared trace-field names to their arrays.

replay #

replay(
    state: BaseState,
    trace: dict[str, Array],
    *,
    dt: float | Array,
    pathwise: bool,
) -> BaseState

Apply the step's effect from trace and write the trace fields - the single effect source.

Returns the same shape as __call__ (a new state for quasistatic/discrete, a sparse delta for dynamic). pathwise=True recomputes reparameterizable quantities from the trace's exogenous noise with live parameters (dx = mean(state) + std(theta) * xi); pathwise =False uses the realized value frozen. A step with nothing to reparameterize (a discrete node) uses the realized action regardless of pathwise.

Parameters:

state (BaseState) –

Live pre-step state.
trace (dict[str, Array]) –

Recorded exogenous noise and realized actions.
dt (float | Array) –

Macro-step duration.
pathwise (bool) –

Whether to recompute reparameterizable values with live parameters.

Returns:

BaseState –

A new state for a quasistatic or discrete step, or a sparse delta state for a dynamic
BaseState –

step.

Raises:

NotImplementedError –

Always, until a subclass supplies replay semantics.

logp #

logp(
    state: BaseState,
    trace: dict[str, Array],
    dt: float | Array,
) -> jaxlib._jax.Array

Log-density this step assigns to trace's action under _dist(state, dt).

The conditioning state is live during replay; the caller has already stop_gradient-ed the trace, so only the recomputed distribution carries gradient.

Parameters:

state (BaseState) –

Live pre-step state conditioning the distribution.
trace (dict[str, Array]) –

Detached recorded trace to score.
dt (float | Array) –

Macro-step duration.

Returns:

Array –

Scalar log-probability contribution for the trace.

Raises:

NotImplementedError –

Always, until a subclass supplies the density.

jax_morph.StepType #

The type of dynamics a step contributes.

DISCRETE `class-attribute` #

DISCRETE = <StepType.DISCRETE: 'discrete'>

The type of dynamics a step contributes.

DYNAMIC `class-attribute` #

DYNAMIC = <StepType.DYNAMIC: 'dynamic'>

The type of dynamics a step contributes.

QUASISTATIC `class-attribute` #

QUASISTATIC = <StepType.QUASISTATIC: 'quasistatic'>

The type of dynamics a step contributes.

jax_morph.check_stochastic_step #

check_stochastic_step(
    step: StochasticStep,
    state: BaseState,
    *,
    dt: float | Array = 1.0,
    key: Array,
) -> dict[str, jax.jaxlib._jax.Array]

Assert a stochastic step's replay co-emits its trace, so scoring reads it back.

The one contract a StochasticStep can silently violate is a replay that forgets to write a trace field into the state (or delta) it returns: the forward then records the reset default, trace_from_state reads that default back, and logp scores garbage - with no error. This helper drives the step and checks the round-trip in two ways:

Co-emission (discrete and dynamic alike). The step is driven twice from trace fields pre-filled with two different sentinels. A replay that co-emits a field derives it from the sampled trace and the physical state, so it records the same value both times; a field it forgets keeps the differing pre-fill - absent (None) in a dynamic delta, or the leaked sentinel in a discrete full state - so the two recorded values disagree and the check fires. (This is why resetting to the default and inspecting one run is not enough: a discrete replay that forgets a field leaves the default there, indistinguishable from writing it.)
Fidelity. Every entry sample_trace drew must survive into the recorded trace unchanged, so the exogenous noise the pathwise replay consumes is preserved. This assumes the documented convention that sample_trace returns each entry as it will be recorded (any eligibility masking happens inside sample_trace, as in the MaybeDivide example) - a step that instead masks a draw inside replay should not list that entry from sample_trace.

Returns the recorded trace on success; raises AssertionError on a round-trip failure and TypeError / ValueError for a misused (non-stochastic or trace-less) step.

Parameters:

step (StochasticStep) –

Stochastic step whose trace contract to validate.
state (BaseState) –

Pre-step state with allocated trace fields.
dt (float | Array, default: 1.0 ) –

Macro-step duration. Defaults to 1.0.
key (Array) –

JAX PRNG key for the repeatable trace draw.

Returns:

dict[str, Array] –

Recorded trace dictionary on success.

Raises:

TypeError –

If step is not stochastic.
ValueError –

If step declares no trace fields.
AssertionError –

If replay fails to co-emit or preserve a trace field.

Step primitives

jax_morph.SimulationStep #

is_stochastic property #

state_reads #

state_writes #

state_requires #

jax_morph.StochasticStep #

state_requires #

trace_writes #

sample_trace #

trace_from_state #

replay #

logp #

jax_morph.StepType #

DISCRETE class-attribute #

DYNAMIC class-attribute #

QUASISTATIC class-attribute #

jax_morph.check_stochastic_step #

is_stochastic `property` #

DISCRETE `class-attribute` #

DYNAMIC `class-attribute` #

QUASISTATIC `class-attribute` #