Three-brick state pattern

The three-brick state pattern¶

A simulated stellar system has two distinct kinds of “state”:

1. Evolving state — positions, velocities, internal stellar properties. These change every timestep.

2. Static configuration — gravity model, softening parameters, physical units, integration scheme. These do not change during the simulation.

Mixing the two into a single mutable object is the most common N-body-code structure (NBODY6, COSMIC, McLuster all use it). progenax deliberately does not: it separates them into three Equinox modules called the three bricks, each immutable, each a PyTree. This separation is the keystone of differentiable simulation. This chapter documents the three bricks, why they exist, and the rules for using them correctly.

The three bricks¶

@jax.tree_util.register_pytree_node_class
class State(eqx.Module):
    positions: Float[Array, "N 3"]
    velocities: Float[Array, "N 3"]
    masses: Float[Array, "N"]
    time: Float[Array, ""]

class SystemParams(eqx.Module):
    gravity_policy: GravityPolicy           # e.g. NewtonianPlummer
    softening_policy: SofteningPolicy       # e.g. ConstantSoftening
    units: UnitSystem                        # e.g. STELLAR
    G: Float[Array, ""]

class ParticleSystem(eqx.Module):
    state: State
    params: SystemParams

Both State and SystemParams are Equinox modules, so they are PyTrees and JAX-traceable. ParticleSystem is a composite that exposes the two bricks together — the user-facing shape that builders construct and integrators consume.

Why separate evolving state from static config¶

The separation matters for gradient flow. Consider an HMC chain that infers cluster parameters $\boldsymbol{\theta}$ — say, the half-mass radius $r_h$ or the IMF slope $\alpha$ — from a final snapshot. The forward model is

The gradient $\partial \mathcal{O} / \partial \boldsymbol{\theta}$ flows through State but not through SystemParams. Things in SystemParams are not the inference target; freezing them in a separate brick clarifies the gradient graph.

How the bricks interact during evolution¶

A single integration step is

@jax.jit
def integrate_step(particle_system, dt):
    # Read both bricks
    state = particle_system.state
    params = particle_system.params

    # Compute forces from both
    forces = params.gravity_policy.compute_forces(
        state.positions, state.masses, params.G,
        softening=params.softening_policy.compute_softening(state.positions),
    )

    # Update state (immutable — produces new state)
    new_state = State(
        positions=state.positions + dt * state.velocities,
        velocities=state.velocities + dt * forces / state.masses[:, None],
        masses=state.masses,
        time=state.time + dt,
    )

    # Wrap back into ParticleSystem with same params
    return ParticleSystem(state=new_state, params=params)

The pattern is: read state and params, compute new state from both, return a new ParticleSystem with the new state and the same params. SystemParams never gets mutated; the integrator threads it through the call chain unchanged.

Builder pattern¶

ICs are constructed in two steps that map cleanly onto the three bricks. progenax builds the physical state — an ICResult (positions, velocities, masses, stellar radii) — and gravax assembles the ParticleSystem (State + SystemParams) from it:

from progenax import build_plummer_cluster   # -> ICResult (pure physical state)
from gravax import ParticleSystem
from jaxstro.units import STELLAR

# 1. progenax builder: parameters + PRNG key -> ICResult (the State data).
ic = build_plummer_cluster(n=1000, r_h=1.0, key=key)   # differentiable in r_h

# 2. gravax assembles the bricks: State from the ICResult, SystemParams from the
#    units + default policies (e.g. ConstantSoftening eps ~ 0.05 * d_mean for the
#    collisionless integrators, gravity_policy=NewtonianPlummer).
system = ParticleSystem.from_ic(ic, units=STELLAR)

The builder is JIT-compatible and differentiable in the profile parameters (r_h, plus the IMF / anisotropy / tidal / rotation knobs) — they flow through the ICResult’s State data. It is not differentiable in units, gravity_policy, or other SystemParams fields, which is the right behaviour: those are configuration choices, not inference targets.

Gotchas and rules¶

Rule 1: State is owned by progenax (and downstream)¶

When code outside progenax (a custom integrator, a renderer) consumes a ParticleSystem, it should treat state as the input data and params as the configuration. Reading either is fine; writing either requires producing a new ParticleSystem rather than mutating in place.

Rule 2: Don’t put traced data in SystemParams¶

If you want to infer the softening-length, do not put it in SofteningPolicy.eps. Put it in State (or a new dedicated brick). SystemParams is for compile-time-static values; placing traced parameters there will retrigger JIT every time the parameter changes.

Rule 3: Don’t put static config in State¶

Conversely, do not put gravity policy, units, or G in State. State should contain only the data that changes over the integration. A mistake here doesn’t break correctness but bloats the PyTree and slows JIT compilation.

Rule 4: ParticleSystem is the only public type¶

User-facing code (builders, integrators, renderers) returns and accepts ParticleSystem. The two bricks are accessed via attribute lookup (ps.state.positions) but should not appear in public function signatures. This decouples the user-facing API from internal restructuring of the bricks.

When to add a fourth brick¶

The three-brick pattern handles ~95% of progenax’s needs. Two situations call for a fourth:

1. Per-mass-bin state for multi-mass clusters where each bin needs its own integration substep. A hypothetical fourth brick BinState would carry per-bin auxiliary state. (On the IC side, per-component structure is already handled without a new brick: MultiComponentCluster.sample_cluster labels each star’s generating component via ICResult.component_id.)

2. External-field state for clusters in time-varying galactic potentials. The galactic potential is technically SystemParams but its time evolution is State; the fourth brick ExternalState carries the time-varying piece.

Adding a fourth brick is a coordinated change across builders, integrators, and renderers — it is not done lightly. The three-brick pattern is the default.

References¶

The state-vs-config separation is a standard pattern in differentiable programming; see Kidger & Garcia (2021) for the Equinox-Module idiom that makes it ergonomic. progenax’s specific naming follows Bradbury et al. (2018)’s PyTree convention. The pattern is consistent with Diffrax’s “diffeq state” / “diffeq solver state” split.