The Simulation System

Now that we know how to use and manipulate the simulation state (State), it’s time to delve into the simulation configuration in System.

A System holds the “static” configuration of a simulation, such as the domain, integrator settings, and force model. Although we call it “static”, many fields (e.g., the time step \(\Delta t\), domain dimensions, boundary conditions) can be changed at runtime—even inside a JIT-compiled function—because both State and System are JAX pytrees.

System Creation

By default, create() initializes unspecified attributes (e.g., domain, force_model, \(\Delta t\)) with sensible defaults.

import jax
import jax.numpy as jnp
import jaxdem as jdem

The system’s dimension must match the state’s dimension. Some components (e.g., domains) transform arrays of shape \((N, d)\) and require \(d\) to agree with the system.

state = jdem.State.create(pos=jnp.zeros((1, 2)))
system = jdem.System.create(state.shape)
state, system = system.step(state, system)  # one step

Instead of state.shape, you can pass the state itself with state=state: the shape is inferred, and the state is forwarded automatically to colliders whose Create method needs one (cell lists, neighbor lists).

system = jdem.System.create(state=state)

A note on static methods

Every operation on State and System (step, trajectory_rollout, merge, stack, etc.) is a static method. That means system.step(state, system) and jdem.System.step(state, system) are equivalent. Static methods make it straightforward to use these operations inside jax.jit(), jax.vmap(), and other JAX transforms.

Configuring the System

You can configure submodules when creating the system via keyword arguments.

system = jdem.System.create(state.shape, domain_type="periodic")
print("periodic domain:", system.domain)

periodic domain: PeriodicDomain(box_size=Array([1., 1.], dtype=float64), inv_box_size=Array([1., 1.], dtype=float64), anchor=Array([0., 0.], dtype=float64))

You can also pass constructor arguments to submodules via *_kw dictionaries.

system = jdem.System.create(
    state.shape,
    domain_type="periodic",
    domain_kw={"box_size": 10.0 * jnp.ones(2), "anchor": jnp.zeros(2)},
)
print("periodic domain (10x10):", system.domain)

periodic domain (10x10): PeriodicDomain(box_size=Array([10., 10.], dtype=float64), inv_box_size=Array([0.1, 0.1], dtype=float64), anchor=Array([0., 0.], dtype=float64))

Passing Module Objects Directly

Internally, create() builds each submodule and performs sanity checks. You can also build a component yourself and pass the instance directly. Every component slot accepts a pre-built instance (domain, collider, linear_integrator, rotation_integrator, force_model, force_manager, bonded_force_model, mat_table), which overrides the corresponding *_type / *_kw arguments.

domain = jdem.Domain.create("free", dim=2)
collider = jdem.Collider.create("naive")
system = jdem.System.create(state.shape, domain=domain, collider=collider)
print("free default domain:", system.domain)
print("directly assigned collider:", type(system.collider).__name__)

free default domain: FreeDomain(box_size=Array([1., 1.], dtype=float64), inv_box_size=Array([1., 1.], dtype=float64), anchor=Array([0., 0.], dtype=float64))
directly assigned collider: NaiveSimulator

This works for instances of your own custom components too (see the custom modules guide).

Post-hoc replacement (system.domain = domain) also works, since System is a mutable dataclass, but passing the instances to create() is preferred because the factory can validate and wire the components together. For swapping dynamics setup), you can use dataclasses.replace to swap the integrators while preserving everything else — including the domain’s current box.

import dataclasses
import jax.numpy as jnp
from jaxdem.integrators import LinearIntegrator

system_dyn = dataclasses.replace(
    system,
    linear_integrator=LinearIntegrator.create("verlet"),
    dt=jnp.asarray(1e-3, dtype=float),
)
print("swapped integrator:", type(system_dyn.linear_integrator).__name__)

swapped integrator: VelocityVerlet

Summary: three ways to build a system component

The previous sections showed three equivalent ways to get a component into a System. From most to least common:

1. Let System.create build it — pass the registered name and its constructor arguments:

   system = jdem.System.create(state.shape, domain_type="periodic",
                               domain_kw={"box_size": box})
   

2. Build it yourself, then pass the instance — use the component’s own factory (jdem.Domain.create("periodic", box_size=box)) or its constructor, and hand the object to System.create via the slot name (domain=domain). Equivalent to 1; useful when you want to inspect or reuse the component.

3. Assign it to an existing system — system.domain = domain. Use this for swapping components after creation; prefer 1 or 2 when first building the system so the factory can validate the combination.

Two things live outside this scheme: writers are plain objects you construct directly (jdem.VTKWriter(...)), and minimizers are optax constructor functions (e.g. jdem.fire) passed via minimizer= / minimizer_kw= — see the integrator guide.

Time stepping

The system controls how the simulation advances in time. You can take a single step or multiple steps at once. Multi-step calls use jax.lax.fori_loop() internally for speed.

state = jdem.State.create(jnp.zeros((1, 2)))
state, system = system.step(state, system)  # 1 step

# Multiple steps in a single call:
state, system = system.step(state, system, n=10)  # 10 steps

Trajectory rollout

If you want to store snapshots along the way, use trajectory_rollout(). It records n snapshots separated by stride integration steps each, for a total of \(n \times \text{stride}\) steps. Each snapshot is taken after its integration steps, so the initial (step-0) state is not stored; to record it, save it yourself before the rollout, or pass per-frame strides with a leading 0 entry.

state = jdem.State.create(jnp.zeros((1, 2)))

state, system, trajectory = system.trajectory_rollout(
    state, system, n=10, stride=2  # total steps = 20
)

The trajectory is a Tuple[State, System] with an extra leading axis of length n.

traj_state, traj_system = trajectory
print("trajectory pos shape:", traj_state.pos.shape)  # (n, N, d)

trajectory pos shape: (10, 1, 2)

Batched simulations with vmap

You can run many independent simulations in parallel with jax.vmap(). Make sure the initialization returns per-simulation State/System pairs.

def initialize(i):
    st = jdem.State.create(jnp.zeros((1, 2)))
    sys = jdem.System.create(
        st.shape,
        domain_type="reflect",
        domain_kw={"box_size": (2 + i) * jnp.ones(2), "anchor": jnp.zeros(2)},
    )
    return st, sys


# Create a batch of 5 simulations
state_b, system_b = jax.vmap(initialize)(jnp.arange(5))
print(system_b.domain)  # batched variable domain

ReflectDomain(box_size=Array([[2., 2.],
       [3., 3.],
       [4., 4.],
       [5., 5.],
       [6., 6.]], dtype=float64), inv_box_size=Array([[0.5       , 0.5       ],
       [0.33333333, 0.33333333],
       [0.25      , 0.25      ],
       [0.2       , 0.2       ],
       [0.16666667, 0.16666667]], dtype=float64), anchor=Array([[0., 0.],
       [0., 0.],
       [0., 0.],
       [0., 0.],
       [0., 0.]], dtype=float64), restitution_coefficient=Array([1., 1., 1., 1., 1.], dtype=float64))

Advance each simulation by 10 steps. Use the class method (or a small wrapper) to avoid variable shadowing.

state_b, system_b = jax.vmap(lambda st, sys: jdem.System.step(st, sys, n=10))(
    state_b, system_b
)
print("batched pos shape:", state_b.pos.shape)  # (batch, N, d)

batched pos shape: (5, 1, 2)

Another way to create batch systems is the stack method:

state = jdem.State.create(jnp.zeros((1, 2)))
system = jdem.System.create(
    state.shape,
)

system = system.stack([system, system, system])
print("stacked system:", system)

stacked system: System(linear_integrator=VelocityVerlet(), rotation_integrator=VelocityVerletSpiral(), collider=NaiveSimulator(overflow=Array([False, False, False], dtype=bool)), domain=FreeDomain(box_size=Array([[1., 1.],
       [1., 1.],
       [1., 1.]], dtype=float64), inv_box_size=Array([[1., 1.],
       [1., 1.],
       [1., 1.]], dtype=float64), anchor=Array([[0., 0.],
       [0., 0.],
       [0., 0.]], dtype=float64)), force_manager=ForceManager(gravity=Array([[0., 0.],
       [0., 0.],
       [0., 0.]], dtype=float64), external_force=Array([[[0., 0.]],

       [[0., 0.]],

       [[0., 0.]]], dtype=float64), external_force_com=Array([[[0., 0.]],

       [[0., 0.]],

       [[0., 0.]]], dtype=float64), external_torque=Array([[[0.]],

       [[0.]],

       [[0.]]], dtype=float64), is_com_force=(), force_functions=(), energy_functions=()), bonded_force_model=None, force_model=SpringForce(laws=()), mat_table=MaterialTable(props={'density': Array([[0.27],
       [0.27],
       [0.27]], dtype=float64), 'poisson': Array([[0.3],
       [0.3],
       [0.3]], dtype=float64), 'young': Array([[10000.],
       [10000.],
       [10000.]], dtype=float64)}, pair={'density_eff': Array([[[0.27]],

       [[0.27]],

       [[0.27]]], dtype=float64), 'poisson_eff': Array([[[0.3]],

       [[0.3]],

       [[0.3]]], dtype=float64), 'young_eff': Array([[[10000.]],

       [[10000.]],

       [[10000.]]], dtype=float64)}, matcher=HarmonicMaterialMatchmaker()), dt=Array([0.005, 0.005, 0.005], dtype=float64), time=Array([0., 0., 0.], dtype=float64), dim=Array([2, 2, 2], dtype=int64), step_count=Array([0, 0, 0], dtype=int64), key=Array([[0, 0],
       [0, 0],
       [0, 0]], dtype=uint32), interact_same_bond_id=Array([False, False, False], dtype=bool), user_pre_step_actions=<PjitFunction of <function _save_state_system at 0x7fb88cc2d4e0>>, user_post_step_actions=<PjitFunction of <function _save_state_system at 0x7fb88cc2d4e0>>, minimizer=CustomGradientTransformation(init=<function fire.<locals>.init at 0x7fb773f9aa20>, update=<function fire.<locals>.update at 0x7fb7703e2160>), target_fn=None)

Deactivating Components

Some modules can be deactivated when creating the system. For the integrators, pass None (preferred; the empty string "" is equivalent) to select the base no-op integrator.

Component	Deactivation value	Effect
linear_integrator_type	None (or "")	No position/velocity updates.
rotation_integrator_type	None (or "")	No orientation/angular-velocity updates.
bonded_force_model_type	None (default)	No bonded forces.
force_manager_kw -> gravity	None (default)	No gravitational acceleration.
force_manager_kw -> force_functions	() (default)	No custom external force and torque functions.

Note: the domain (domain_type), collider (collider_type), and force model (force_model_type) cannot be deactivated — a valid type must always be provided.

# No integration — a "frozen" system:
system_frozen = jdem.System.create(
    state.shape,
    linear_integrator_type=None,
    rotation_integrator_type=None,
)
print("Integrator:", type(system_frozen.linear_integrator).__name__)

Integrator: LinearIntegrator

Random Number Generation

create() accepts a seed (integer) or a key (jax.random.PRNGKey()) that initializes the system’s JAX PRNG state. An explicit key takes precedence over seed. The key is stored in system.key and is available for stochastic integrators or custom force functions.

system_rng = jdem.System.create(state.shape, seed=42)
print("PRNG key:", system_rng.key)

PRNG key: [ 0 42]