Dynamics

jax_morph.physics.BrownianDynamics

BrownianDynamics(
    potential,
    *,
    n_space_dim,
    gamma=1.0,
    kT=0.1,
    tag='brownian',
    score_by_default=False,
)

Overdamped Langevin (Brownian) dynamics as a reparameterized dynamic stochastic step.

Each macro-step applies one Euler-Maruyama displacement

\[\Delta x = -\frac{\nabla U}{\gamma}\,\Delta t + \sqrt{\frac{2\,kT\,\Delta t}{\gamma}}\;\xi,\]

with \(\xi\) standard normal. The noise is recorded, so the step is both pathwise-differentiable (recompute \(\Delta x\) from the frozen noise with live parameters) and scorable: logp is the Gaussian log-density of the recorded \(\Delta x\) under \(\mathcal{N}\!\left(-\frac{\nabla U}{\gamma}\Delta t,\ \frac{2\,kT\,\Delta t}{\gamma}\right)\), with parameters (the potential's, kT, gamma) recomputed live.

The trace fields {tag}_xi and {tag}_dx are additive dynamic fields (default 0), namespaced by tag so several instances coexist; tag='brownian' gives brownian_xi / brownian_dx.

Attributes:

• potential –

  Interaction potential supplying mechanical forces.

• n_space_dim –

  Static spatial dimension; position and trace arrays have trailing shape (n_space_dim,).

• gamma –

  Translational drag coefficient. Defaults to 1.0.

• kT –

  Thermal energy controlling translational noise. Defaults to 0.1.

• tag –

  Static namespace for trace fields. Defaults to 'brownian'.

• score_by_default –

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

Pass potential=None for NoForce (a free Brownian gas: pure noise, no drift).

Parameters:

• potential –

  Interaction potential, or None for NoForce.

• n_space_dim –

  Spatial dimension used for trace-field shapes.

• gamma –

  Translational drag coefficient. Defaults to 1.0.

• kT –

  Thermal energy controlling translational noise. Defaults to 0.1.

• tag –

  Trace-field namespace. Defaults to 'brownian'.

• score_by_default –

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

state_reads

state_reads()

Reads positions (the drift), plus any state field the potential sources params from.

state_writes

state_writes()

Writes the position displacement (a dynamic delta).

trace_writes

trace_writes()

The exogenous noise {tag}_xi and the realized displacement {tag}_dx (both default 0).

sample_trace

sample_trace(state, *, dt, key)

Draw standard-normal displacement noise for pathwise replay.

Parameters:

• state –

  Pre-step state with positions of shape (capacity, n_space_dim).

• dt –

  Unused macro-step duration.

• key –

  JAX PRNG key.

Returns:

• –

  Trace dictionary containing {tag}_xi with shape (capacity, n_space_dim).

Raises:

• ValueError –

  If the state's spatial dimension differs from n_space_dim.

replay

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

Apply and record a displacement, recomputing it from noise when pathwise.

Parameters:

• state –

  Live pre-step state.

• trace –

  Recorded noise and displacement arrays of shape (capacity, n_space_dim).

• dt –

  Macro-step duration.

• pathwise –

  Whether to recompute displacement from the live transition kernel.

Returns:

• –

  Sparse delta containing position and the two trace arrays.

Raises:

• ValueError –

  If the state's spatial dimension differs from n_space_dim.

logp

logp(state, trace, dt)

Score recorded displacement under the live transition kernel over alive cells.

Parameters:

• state –

  Live pre-step state conditioning the drift and noise scale.

• trace –

  Recorded displacement trace.

• dt –

  Macro-step duration.

Returns:

• –

  Scalar sum of alive-cell Gaussian log-densities, or zero for deterministic noise.

---

jax_morph.physics.ActiveBrownianDynamics2D

ActiveBrownianDynamics2D(
    potential,
    *,
    n_space_dim,
    gamma=1.0,
    kT=0.1,
    rot_diffusion=1.0,
    tag='active',
    score_by_default=False,
)

Overdamped active Brownian (self-propelled) dynamics as a reparameterized dynamic step.

On top of the passive Langevin drift and noise, each cell self-propels at speed \(v_0\) = active_speed along a persistent heading \(\theta\), adding \(v_0\,(\cos\theta, \sin\theta)\) to the drift; the heading rotationally diffuses, \(\Delta\theta = \sqrt{2 D_r\,\Delta t}\;\xi_\theta\) at rate \(D_r\) = rot_diffusion. Both exogenous noises are recorded, so the step is pathwise-differentiable and scorable (a Gaussian kernel on the displacement and on the heading increment). 2-D only (the heading is an angle).

Attributes:

• potential –

  Interaction potential supplying mechanical forces.

• n_space_dim –

  Static spatial dimension, fixed at 2; position traces have shape (capacity, 2).

• gamma –

  Translational drag coefficient. Defaults to 1.0.

• kT –

  Thermal energy controlling translational noise. Defaults to 0.1.

• rot_diffusion –

  Heading rotational-diffusion coefficient. Defaults to 1.0.

• tag –

  Static namespace for trace fields. Defaults to 'active'.

• score_by_default –

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

Pass potential=None for NoForce (self-propulsion + heading noise only, no force); this is the natural base for a Vicsek-type model, whose alignment lives in a separate step.

Parameters:

• potential –

  Interaction potential, or None for NoForce.

• n_space_dim –

  Spatial dimension; must be 2.

• gamma –

  Translational drag coefficient. Defaults to 1.0.

• kT –

  Thermal energy controlling translational noise. Defaults to 0.1.

• rot_diffusion –

  Heading rotational-diffusion coefficient. Defaults to 1.0.

• tag –

  Trace-field namespace. Defaults to 'active'.

• score_by_default –

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

Raises:

• ValueError –

  If n_space_dim is not 2.

state_reads

state_reads()

Reads positions (drift), the per-cell speed and heading, and the potential's fields.

state_writes

state_writes()

Writes the position delta and the shared, persistent (heritable) heading angle.

trace_writes

trace_writes()

Exogenous noises and realized increments for the translation and the heading (default 0).

sample_trace

sample_trace(state, *, dt, key)

Draw standard-normal translational and rotational noise for pathwise replay.

Parameters:

• state –

  Pre-step state with positions of shape (capacity, 2).

• dt –

  Unused macro-step duration.

• key –

  JAX PRNG key.

Returns:

• –

  Trace dictionary with translational noise shape (capacity, 2) and rotational noise

• –

  shape (capacity,).

Raises:

• ValueError –

  If the state is not two-dimensional.

replay

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

Apply and record displacement and heading increments from the recorded trace.

Parameters:

• state –

  Live pre-step state.

• trace –

  Recorded translational and rotational noise/increment arrays.

• dt –

  Macro-step duration.

• pathwise –

  Whether to recompute increments from live transition parameters.

Returns:

• –

  Sparse delta containing position, heading, and trace increments.

Raises:

• ValueError –

  If the state is not two-dimensional.

logp

logp(state, trace, dt)

Score recorded displacement and heading increments over alive cells.

Parameters:

• state –

  Live pre-step state conditioning the transition kernels.

• trace –

  Recorded translational and rotational increments.

• dt –

  Macro-step duration.

Returns:

• –

  Scalar sum of alive-cell Gaussian log-densities, with zero contribution for any

• –

  deterministic component.