The Simulation State

This example covers the jaxdem.state.State object, a core component of JaxDEM that holds all particle data in a simulation.

JaxDEM stores particle data in a Structure-of-Arrays (SoA) layout, which suits JAX’s vectorized and parallel computations. This layout also makes trajectories and batched simulations easy to handle.

This guide shows how to create, modify, and extend the simulation state.

State Creation

We start with a simple 2D state with a single particle at the origin. By default, jaxdem.state.State.create() fills non-specified attributes (like velocity, radius, mass) with default values.

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

state = jdem.State.create(pos=jnp.array([[0.0, 0.0]]))
print(f"Dimension of state: {state.dim}")
print(f"Initial position: {state.pos}")

Dimension of state: 2
Initial position: [[0. 0.]]

To create a 3D state, pass 3D coordinates. JaxDEM infers the dimension from the position data (JaxDEM supports only 2D and 3D).

state = jdem.State.create(pos=jnp.array([[0.0, 0.0, 0.0]]))
print(f"Dimension of state: {state.dim}")
print(f"Initial position: {state.pos}")

Dimension of state: 3
Initial position: [[0. 0. 0.]]

Understanding Positions: pos, pos_c, and pos_p

Note: state.pos is not a stored field. It is a computed property: pos = pos_c + R(q) @ pos_p, where R(q) is the rotation given by the particle’s quaternion orientation.

The stored fields are:

• pos_c — the center-of-mass position of each particle (or clump).

• pos_p — the offset from the center of mass in the principal (body) frame. For simple spheres pos_p is zero, so pos == pos_c.

For clumps (rigid bodies made of multiple spheres), every sphere in the same clump shares the same center of mass position pos_c, orientation q, velocity vel, angular velocity ang_vel, force, torque, mass, inertia, fixed, and clump_id. The per-sphere fields that can vary within a rigid clump are:

• pos_p — the body-frame offset relative to the COM

• rad — individual sphere radius

• volume — stored per sphere: State.create defaults it to each sphere’s own hypersphere volume, while add_clump() broadcasts a provided clump volume to every member sphere. Packing-fraction utilities read one value per clump (via a segment max).

• Identifiers: clump_id, bond_id, mat_id, species_id

This design allows vectorized operations over all spheres without branching on clump membership.

Particle Sizes: rad and _rad

JaxDEM distinguishes between the physical size of a particle and its search/interaction radius:

• rad — the “true” physical radius of the particle. Narrow-phase contact force and energy calculations use this radius, as do plotting and visualization. For facets, rad represents the physical thickness of the segment or triangle.

• _rad — the private “search” radius of the particle. The broad-phase colliders use it to build neighbor/cell lists. For standard spheres, _rad is equal to rad. For facet vertices, _rad is the maximum distance from the vertex to the center of mass (COM) of the facet, so the broad-phase candidate list covers the entire facet.

JaxDEM computes the broad-phase search radius _rad internally. User-facing constructors (like State.create or State.add_facet) do not expose it as a parameter.

Modifying State Attributes

We have two primary ways to set or modify particle attributes:

1. Direct assignment: You can assign new JAX arrays to attributes like state.vel. This is flexible, but you must keep the shapes consistent.

state.vel = jnp.ones_like(state.pos)
print(state.vel)

[[1. 1. 1.]]

Because these are JAX arrays, code like

state.vel[i] = jnp.asarray([1, 2, 3], dtype=float)

raises an error. The correct way is

i = 0
state.vel = state.vel.at[i].set(jnp.asarray([1, 2, 3], dtype=float))
print(state.vel)

[[1. 2. 3.]]

This is inefficient. Prefer vectorized operations.

2. Constructor arguments: This is the safer approach. The jaxdem.state.State.create() constructor validates shapes and types and keeps all attributes consistent.

state = jdem.State.create(pos=jnp.zeros((1, 2)), vel=jnp.ones((1, 2)))
print(state.vel)

[[1. 1.]]

Fixed (Immobile) Particles

The boolean field state.fixed marks particles with a prescribed motion. Forces do not drive them. The integrator multiplies velocity updates by (1 - fixed), which masks the acceleration. A fixed particle keeps whatever velocity it currently has, no matter which forces act on it. With zero initial velocity it stays put. With a nonzero initial velocity it keeps moving at that prescribed velocity. This is useful for walls, obstacles, or driven boundary particles.

state = jdem.State.create(
    pos=jnp.array([[0.0, 0.0], [2.0, 0.0]]),
    rad=jnp.array([1.0, 1.0]),
    fixed=jnp.array([True, False]),
)
print("Fixed mask:", state.fixed)

Fixed mask: [ True False]

Identifier Fields

Each particle carries several integer identifiers:

• clump_id — groups particles into rigid bodies (see Clumps (Rigid Bodies)). Particles with the same clump_id never interact via contact forces and move as one body. By default every particle has a unique clump_id.

• bond_id — connectivity masking array (see Deformable Particles). For each particle, it stores the array indices of the neighbor particles it connects to. JaxDEM disables (masks out) interactions between connected particles. It has shape (N, max_num_neighbors), and JaxDEM pads it with -1.

• mat_id — indexes into the MaterialTable to look up material properties (density, Young’s modulus, …).

• species_id — selects which force law applies to a pair when using a ForceRouter (see Force Models).

print("clump_id :", state.clump_id)
print("bond_id  :", state.bond_id)
print("mat_id   :", state.mat_id)
print("species_id:", state.species_id)

clump_id : [0 1]
bond_id  : [[-1]
 [-1]]
mat_id   : [0 0]
species_id: [0 0]

Setting Up Connections with bond_id

For each particle, the bond_id array stores the array indices of other particles it connects to.

By default, the collider uses these connections to ignore contact/non-bonded interactions (so connected particles do not collide with each other). The interact_same_bond_id parameter in the system creation options controls this behavior (see jaxdem.system.System.create()). Set interact_same_bond_id=True to let particles connected via bond_id still feel contact forces.

This connectivity masking is useful with bonded models (where interactions are permanent), such as deformable particle models and cohesive/bonded networks. For details on bonded interactions, see the Deformable Particles and the Colliders.

When you call jaxdem.state.State.create(), you can define connections by passing a list of lists (which can have uneven lengths) for the bond_id argument. JaxDEM symmetrizes these connections (if particle A connects to B, then B connects to A) and pads the array with -1 up to the maximum number of connections. If a particle has no connections, its row contains only -1. If you pass no connections at all, bond_id defaults to a shape of (N, 1) filled with -1.

# Create a state with 4 particles:
# - Particle 0 connects to 1 and 2
# - Particle 1 connects to 0
# - Particle 2 connects to 0
# - Particle 3 has no connections
positions = jnp.array([[0.0, 0.0], [1.0, 0.0], [0.0, 1.0], [2.0, 2.0]])
state_bonded = jdem.State.create(pos=positions, bond_id=[[1, 2], [0], [0], []])
print("Bond IDs for each particle:\n", state_bonded.bond_id)

# Default behavior when bond_id is not passed
state_no_bonds = jdem.State.create(pos=positions)
print("Default bond IDs (no bonds):\n", state_no_bonds.bond_id)

Bond IDs for each particle:
 [[ 1  2]
 [ 0 -1]
 [ 0 -1]
 [-1 -1]]
Default bond IDs (no bonds):
 [[-1]
 [-1]
 [-1]
 [-1]]

Extending the State

Working directly with SoA structures can feel less intuitive than Array-of-Structures (AoS) when you add or modify individual particles. To help with this, JaxDEM provides utility methods like jaxdem.state.State.add().

jaxdem.state.State.add() appends new particles to an existing state. It assigns unique clump_ids and checks dimension consistency.

state = jdem.State.create(pos=jnp.array([[0.0, 0.0]]), rad=jnp.array([0.5]))
print(f"Initial state (N={state.N}, clump_ids={state.clump_id}):\npos={state.pos}")

state = jdem.State.add(
    state,
    pos=jnp.array([[1.0, 1.0]]),
    vel=2 * jnp.ones((1, 2)),
    rad=10 * jnp.ones(1),
)
print(
    f"\nState after addition (N={state.N}, clump_ids={state.clump_id}):\npos={state.pos}"
)
print(f"New particle velocity: {state.vel[-1]}")
print(f"New particle radius: {state.rad[-1]}")

Initial state (N=1, clump_ids=[0]):
pos=[[0. 0.]]

State after addition (N=2, clump_ids=[0 1]):
pos=[[0. 0.]
 [1. 1.]]
New particle velocity: [2. 2.]
New particle radius: 10.0

You can also add several particles at once with arrays of the appropriate shape. jaxdem.state.State.add() checks that the dimensions of the new particles match the existing state.

state = jdem.State.add(
    state,
    pos=jnp.array([[2.0, 0.0], [0.0, 2.0]]),
    vel=jnp.zeros((2, 2)),
    rad=jnp.array([0.8, 0.3]),
    clump_id=jnp.array([2, 3]),
)
print(
    f"\nState after adding multiple particles (N={state.N}, clump_ids={state.clump_id}):\n{state.pos}"
)

State after adding multiple particles (N=4, clump_ids=[0 1 2 3]):
[[0. 0.]
 [1. 1.]
 [2. 0.]
 [0. 2.]]

Note that we provided explicit clump_id values here. jaxdem.state.State.add() adds jnp.max(state.clump_id) + 1 to the provided IDs to avoid overlaps. The resulting sequence may not be contiguous, and that is valid.

Merging Two States

To combine two State objects, use jaxdem.state.State.merge(). It concatenates the particles from the second state onto the first. This is useful for assembling initial configurations from smaller parts.

state_a = jdem.State.create(
    pos=jnp.array([[0.0, 0.0], [1.0, 1.0]]),
)
state_b = jdem.State.create(
    jnp.array([[2.0, 2.0], [3.0, 3.0], [5.0, 2.0]]),
)
state = jdem.State.merge(state_a, state_b)

print(f"State A (N={state_a.N}, clump_ids={state_a.clump_id}):\npos={state_a.pos}")
print(f"State B (N={state_b.N}, clump_ids={state_b.clump_id}):\npos={state_b.pos}")
print(f"Merged state (N={state.N}, clump_ids={state.clump_id}):\npos={state.pos}")

State A (N=2, clump_ids=[0 1]):
pos=[[0. 0.]
 [1. 1.]]
State B (N=3, clump_ids=[0 1 2]):
pos=[[2. 2.]
 [3. 3.]
 [5. 2.]]
Merged state (N=5, clump_ids=[0 1 2 3 4]):
pos=[[0. 0.]
 [1. 1.]
 [2. 2.]
 [3. 3.]
 [5. 2.]]

Stacking States for Trajectories or Batches

JaxDEM can handle batched states. A batch can represent trajectories (multiple snapshots over time) or independent simulations (multiple distinct initial conditions).

This is useful for performance. JaxDEM targets throughput: if your GPU is not saturated, you are leaving performance on the table. A common DEM task is running parameter sweeps. JaxDEM lets you run many independent simulations in parallel. Until you saturate the GPU, all of them can finish in about the time of one.

Also, with trajectory support you do not have to interrupt the GPU for I/O (e.g., saving state to disk). You can accumulate a full trajectory in memory and save everything at the end. This is often faster at the cost of more memory.

To manage simulation trajectories or run batched simulations, use jaxdem.state.State.stack(). It takes a sequence of jaxdem.state.State snapshots and concatenates them along a new leading axis. This creates a multi-dimensional state where the first axis can represent time steps, batch elements, or other high-level groupings. Note that stacking does not shift particle clump_ids, because it assumes the particles are the same entities across the stacked dimension. jaxdem.state.State.stack() makes sure shapes are consistent.

snapshot1 = jdem.State.create(pos=jnp.array([[0.0, 0.0]]), rad=jnp.array([2.0]))
snapshot2 = jdem.State.create(pos=jnp.array([[0.1, 0.0]]), vel=jnp.array([[0.1, 0.0]]))
snapshot3 = jdem.State.create(pos=jnp.array([[0.2, 0.0]]), mass=jnp.array([3.3]))

batched_state = jdem.State.stack([snapshot1, snapshot2, snapshot3])

print(f"Shape of stacked positions (B, N, dim): {batched_state.pos.shape}")
print(f"Batch size: {batched_state.batch_size}")

Shape of stacked positions (B, N, dim): (3, 1, 2)
Batch size: 3

Another way to create batch states is JAX’s vmap:

batched_state = jax.vmap(
    lambda i: jdem.State.create(
        i
        * jnp.ones(
            (1, 2),
        )
    )
)(jnp.arange(4))
print(f"Shape of stacked positions (B, N, dim): {batched_state.pos.shape}")
print(f"Batch size: {batched_state.batch_size}")
print(f"Position at batch 0: {batched_state.pos[0]}")
print(f"Position at batch 1: {batched_state.pos[1]}")
print(f"Position at batch 2: {batched_state.pos[2]}")

Shape of stacked positions (B, N, dim): (4, 1, 2)
Batch size: 4
Position at batch 0: [[0. 0.]]
Position at batch 1: [[1. 1.]]
Position at batch 2: [[2. 2.]]

A more realistic way to get a batched state:

def initialize(i: jax.Array) -> tuple[jdem.State, jdem.System]:
    state = jdem.State.create(i * jnp.ones((4, 2)))
    system = jdem.System.create(state.shape)
    return state, system


N_batches = 10
state, system = jax.vmap(initialize)(jnp.arange(N_batches))

Then, to run this simulation:

state, system = system.step(state, system, n=10)
print(f"Shape of positions (B, N, dim): {state.pos.shape}")

Shape of positions (B, N, dim): (10, 4, 2)

The system can change over time, so each state needs its own system.

Trajectories of Batches

JaxDEM can also accumulate trajectories of batched states.

This is useful for parameter sweeps: you run several independent simulations (a batch) and capture their full time evolution (a trajectory) without frequent I/O.

jaxdem.writers.VTKWriter.save() understands these multi-dimensional states.

By convention, when dealing with State attributes of shape (…, N, dim):

• For a single snapshot (no batch, no trajectory), the shape is (N, dim).

• For a batched state (a single snapshot across multiple independent simulations), the shape is (B, N, dim).

• For a single trajectory (multiple snapshots over time of a single simulation), the shape is (T, N, dim).

• For a trajectory of batches (multiple snapshots over time of multiple parallel simulations), the shape is (T, B, N, dim).

In JaxDEM, the batch dimension B (if present) is always at shape[-3]. This is the axis before the particle dimension N at shape[-2]. So batch_size returns shape[-3] when ndim >= 3. This gives B for both (B, N, dim) and (T, B, N, dim) shapes.

When collecting trajectories (via trajectory_rollout()), each snapshot is stacked along the first axis (axis 0), producing a state of shape (T, B, N, dim) for batched trajectories.

jaxdem.writers.VTKWriter.save() understands these layouts. By default (trajectory=False) it treats all leading axes as independent batches. Pass trajectory=True to tell the writer which axis is time (trajectory_axis, default 0). The writer swaps that axis to the front, keeps it as T, and flattens any remaining leading axes into a single batch axis B. The result is (T, B, N, dim) internally.

batched_state = jdem.State.stack([batched_state, batched_state, batched_state])
print(f"Shape of stacked positions (T, B, N, dim): {batched_state.pos.shape}")
print(f"Batch size: {batched_state.batch_size}")

Shape of stacked positions (T, B, N, dim): (3, 4, 1, 2)
Batch size: 4

As in the previous section, you can get a trajectory of batches like this:

N_batches = 9
state, system = jax.vmap(initialize)(jnp.arange(N_batches))

state, system, (state_traj, system_traj) = system.trajectory_rollout(
    state, system, n=10
)

print(f"Shape of positions (T, B, N, dim): {state_traj.pos.shape}")

Shape of positions (T, B, N, dim): (10, 9, 4, 2)

Utilities

JaxDEM includes utility functions in jaxdem.utils for setting up simulations. For example, you can create a state with randomized attributes:

from jaxdem import utils as utils

state = utils.random_state(dim=3, N=10)
print(state)

State(pos_c=Array([[9.69947998, 3.1518186 , 6.6109639 ],
       [9.91386812, 2.82521542, 3.55488773],
       [5.25257116, 3.2260835 , 0.58677741],
       [6.15111094, 9.54725682, 8.38766458],
       [7.92837724, 1.78128245, 3.73091685],
       [7.69204925, 8.30547677, 3.51584082],
       [7.63565438, 9.26869417, 0.50223628],
       [4.82317485, 2.7366294 , 7.91500943],
       [3.35842191, 9.16958235, 2.76594463],
       [1.19191249, 7.09289908, 7.46735741]], dtype=float64), pos_p=Array([[0., 0., 0.],
       [0., 0., 0.],
       [0., 0., 0.],
       [0., 0., 0.],
       [0., 0., 0.],
       [0., 0., 0.],
       [0., 0., 0.],
       [0., 0., 0.],
       [0., 0., 0.],
       [0., 0., 0.]], dtype=float64), vel=Array([[ 0.50113777, -0.04070654,  0.95583301],
       [-0.61279657, -0.72622868, -0.4203813 ],
       [-0.63960196,  0.73330472,  0.01790463],
       [-0.25902228, -0.83200616, -0.12118167],
       [-0.59683588, -0.85513092,  0.97808852],
       [ 0.40873141, -0.47646233, -0.68906414],
       [-0.39218955,  0.95561441, -0.01519757],
       [ 0.91524195,  0.99416172,  0.89142842],
       [-0.00991449, -0.36131967,  0.47242913],
       [ 0.71953161,  0.97946472, -0.41548261]], dtype=float64), force=Array([[0., 0., 0.],
       [0., 0., 0.],
       [0., 0., 0.],
       [0., 0., 0.],
       [0., 0., 0.],
       [0., 0., 0.],
       [0., 0., 0.],
       [0., 0., 0.],
       [0., 0., 0.],
       [0., 0., 0.]], dtype=float64), q=Quaternion(w=Array([[1.],
       [1.],
       [1.],
       [1.],
       [1.],
       [1.],
       [1.],
       [1.],
       [1.],
       [1.]], dtype=float64), xyz=Array([[0., 0., 0.],
       [0., 0., 0.],
       [0., 0., 0.],
       [0., 0., 0.],
       [0., 0., 0.],
       [0., 0., 0.],
       [0., 0., 0.],
       [0., 0., 0.],
       [0., 0., 0.],
       [0., 0., 0.]], dtype=float64)), ang_vel=Array([[0., 0., 0.],
       [0., 0., 0.],
       [0., 0., 0.],
       [0., 0., 0.],
       [0., 0., 0.],
       [0., 0., 0.],
       [0., 0., 0.],
       [0., 0., 0.],
       [0., 0., 0.],
       [0., 0., 0.]], dtype=float64), torque=Array([[0., 0., 0.],
       [0., 0., 0.],
       [0., 0., 0.],
       [0., 0., 0.],
       [0., 0., 0.],
       [0., 0., 0.],
       [0., 0., 0.],
       [0., 0., 0.],
       [0., 0., 0.],
       [0., 0., 0.]], dtype=float64), rad=Array([0.54031218, 0.96197442, 0.86697617, 0.78972512, 0.7988015 ,
       0.85001331, 0.90281005, 0.94473502, 0.53813944, 0.93642903],      dtype=float64), _rad=Array([0.54031218, 0.96197442, 0.86697617, 0.78972512, 0.7988015 ,
       0.85001331, 0.90281005, 0.94473502, 0.53813944, 0.93642903],      dtype=float64), volume=Array([0.66072827, 3.72888661, 2.72966967, 2.06308191, 2.13503613,
       2.57256167, 3.08232027, 3.53198336, 0.65278938, 3.43964177],      dtype=float64), mass=Array([1., 1., 1., 1., 1., 1., 1., 1., 1., 1.], dtype=float64), inertia=Array([[0.1167749 , 0.1167749 , 0.1167749 ],
       [0.37015791, 0.37015791, 0.37015791],
       [0.30065907, 0.30065907, 0.30065907],
       [0.24946631, 0.24946631, 0.24946631],
       [0.25523354, 0.25523354, 0.25523354],
       [0.28900905, 0.28900905, 0.28900905],
       [0.32602639, 0.32602639, 0.32602639],
       [0.3570097 , 0.3570097 , 0.3570097 ],
       [0.11583762, 0.11583762, 0.11583762],
       [0.35075973, 0.35075973, 0.35075973]], dtype=float64), clump_id=Array([0, 1, 2, 3, 4, 5, 6, 7, 8, 9], dtype=int64), bond_id=Array([[-1],
       [-1],
       [-1],
       [-1],
       [-1],
       [-1],
       [-1],
       [-1],
       [-1],
       [-1]], dtype=int64), mat_id=Array([0, 0, 0, 0, 0, 0, 0, 0, 0, 0], dtype=int64), species_id=Array([0, 0, 0, 0, 0, 0, 0, 0, 0, 0], dtype=int64), fixed=Array([False, False, False, False, False, False, False, False, False,
       False], dtype=bool), facet_id=Array([-1, -1, -1, -1, -1, -1, -1, -1, -1, -1], dtype=int64), facet_vertices=Array([[-1, -1, -1],
       [-1, -1, -1],
       [-1, -1, -1],
       [-1, -1, -1],
       [-1, -1, -1],
       [-1, -1, -1],
       [-1, -1, -1],
       [-1, -1, -1],
       [-1, -1, -1],
       [-1, -1, -1]], dtype=int64), _pos_p_rot=Array([[0., 0., 0.],
       [0., 0., 0.],
       [0., 0., 0.],
       [0., 0., 0.],
       [0., 0., 0.],
       [0., 0., 0.],
       [0., 0., 0.],
       [0., 0., 0.],
       [0., 0., 0.],
       [0., 0., 0.]], dtype=float64))