The Simulation System

Now that we know how to use and modify the simulation state (State), we move to the simulation configuration in System.

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

System Creation

By default, create() fills unspecified attributes (e.g., domain, force_model, \(\Delta t\)) with default values.

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. JaxDEM then infers the shape and forwards the state 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 easy to use these operations inside jax.jit(), jax.vmap(), and other JAX transforms.

Configuring the System

You can configure submodules when you create the system with 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 with *_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 runs 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). The instance 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, because System is a mutable dataclass. Passing the instances to create() is better, because the factory validates the components and wires them together. To swap the dynamics setup, use dataclasses.replace to change the integrators and keep 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 with 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 to swap 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(...)). 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. It takes each snapshot after its integration steps, so it does not store the initial (step-0) state. 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 0x7fbfd10fd1c0>>, user_post_step_actions=<PjitFunction of <function _save_state_system at 0x7fbfd10fd1c0>>, minimizer=CustomGradientTransformation(init=<function fire.<locals>.init at 0x7fbf0ea71120>, update=<function fire.<locals>.update at 0x7fbf037cc360>), target_fn=None)

Deactivating Components

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

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: you cannot deactivate the domain (domain_type), collider (collider_type), or force model (force_model_type). You must always provide a valid type.

# 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()) to initialize the system’s JAX PRNG state. An explicit key takes precedence over seed. JaxDEM stores the key in system.key 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]