jaxdem.state#

Defines the simulation State.

Module Attributes

INERTIA_REGULARIZATION

Diagonal regularization added to inertia computed from facet/mesh geometry, as mass * INERTIA_REGULARIZATION (units of mass, so the bias is exactly 1e-4 relative only for bodies of unit length scale).

Classes

State(pos_c, pos_p, vel, force, q, ang_vel, ...)

Represents the complete simulation state for a system of N particles in 2D or 3D.

jaxdem.state.INERTIA_REGULARIZATION = 0.0001#: Diagonal regularization added to inertia computed from facet/mesh geometry, as mass * INERTIA_REGULARIZATION (units of mass, so the bias is exactly 1e-4 relative only for bodies of unit length scale). Degenerate geometry (collinear or coplanar vertices) yields zero inertia about some axis, which would make 1/I in the rotation integrators infinite; this floor keeps angular accelerations finite at the cost of a small bias in the affected principal moment.

final class jaxdem.state.State(pos_c: Array, pos_p: Array, vel: Array, force: Array, q: Quaternion, ang_vel: Array, torque: Array, rad: Array, _rad: Array, volume: Array, mass: Array, inertia: Array, clump_id: Array, bond_id: Array, mat_id: Array, species_id: Array, fixed: Array, facet_id: Array = <factory>, facet_vertices: Array = <factory>, _pos_p_rot: Array = <factory>)#

Bases: object

Represents the complete simulation state for a system of N particles in 2D or 3D.

Notes:#

State is designed to support various data layouts:

Single snapshot:
pos.shape = (N, dim) for particle properties (e.g., pos, vel, force), and (N,) for scalar properties (e.g., rad, mass). In this case, batch_size is 1.
Batched states:
pos.shape = (B, N, dim) for particle properties, and (B, N) for scalar properties. Here, B is the batch dimension (batch_size = pos.shape[0]).
Trajectories of a single simulation:
pos.shape = (T, N, dim) for particle properties, and (T, N) for scalar properties. Here, T is the trajectory dimension.
Trajectories of batched states:
pos.shape = (T_1, T_2, …, T_k, B, N, dim) for particle properties, and (T_1, T_2, …, T_k, B, N) for scalar properties.
- The dimension immediately preceding N (i.e., pos.shape[-3]) is always interpreted as the batch dimension (`B`) — this is what batch_size returns.
- All leading dimensions before it (T_1, T_2, … T_k) are interpreted as trajectory dimensions
  and they are flattened at save time if there is more than 1 trajectory dimension.

The class is final and cannot be subclassed.

Example:#

Creating a simple 2D state for 4 particles:

>>> import jaxdem as jdem
>>> import jax.numpy as jnp
>>> import jax
>>>
>>> positions = jnp.array([[0.0, 0.0], [1.0, 0.0], [0.0, 1.0], [1.0, 1.0]])
>>> state = jdem.State.create(pos=positions)
>>>
>>> print(f"Number of particles (N): {state.N}")
>>> print(f"Spatial dimension (dim): {state.dim}")
>>> print(f"Positions: {state.pos}")

Creating a batched state:

>>> batched_state = jax.vmap(lambda _: jdem.State.create(pos=positions))(jnp.arange(10))
>>>
>>> print(f"Batch size: {batched_state.batch_size}")  # 10
>>> print(f"Positions shape: {batched_state.pos.shape}")

pos_c: Array#

Array of particle center of mass positions. Shape is (…, N, dim).

For rigid clump members this is the clump’s center of mass, identical on every member; the sphere’s own center is the computed property pos.

pos_p: Array#: Vector relative to the center of mass (pos_p = pos - pos_c) in the principal reference frame. This field should be constant. Shape is (…, N, dim).

vel: Array#

Array of particle center of mass velocities. Shape is (…, N, dim).

For rigid clump members this is the clump COM velocity, identical on every member.

force: Array#

Array of particle forces. Shape is (…, N, dim).

For rigid clump members, every member stores the total aggregated clump force.

q: Quaternion#

Quaternion representing the orientation of the particle.

For rigid clump members this is the clump orientation, identical on every member.

ang_vel: Array#

Array of particle center of mass angular velocities. Shape is (…, N, 1 | 3) depending on 2D or 3D simulations.

For rigid clump members this is the clump angular velocity, identical on every member.

torque: Array#

Array of particle torques. Shape is (…, N, 1 | 3) depending on 2D or 3D simulations.

For rigid clump members, every member stores the total aggregated clump torque about the clump COM.

rad: Array#: Array of particle radii. Shape is (…, N).

volume: Array#

Array of particle volumes (or areas if 2D). Shape is (…, N).

Per-sphere by default; when an explicit clump volume is passed to State.add_clump(), it is replicated on every member.

mass: Array#

Array of particle masses. Shape is (…, N).

For rigid clump members, every member stores the total clump mass, not a per-member share.

inertia: Array#

Inertia tensor in the principal axis frame (…, N, 1 | 3) depending on 2D or 3D simulations.

For rigid clump members, every member stores the total clump inertia tensor.

clump_id: Array#: Array of clump identifiers. Bodies with the same clump_id are treated as part of the same rigid body. Shape is (…, N). IDs need to be between 0 and N.

bond_id: Array#: Array of connected neighbors for contact filtering. For each particle, it stores the array indices of the neighbor particles it is connected to. Interactions between connected particles are disabled. Shape is (…, N, max_num_neighbors). Empty slots/non-connections are padded with -1.

mat_id: Array#: Array of material IDs for each particle. Shape is (…, N).

species_id: Array#: Array of species IDs for each particle. Shape is (…, N).

fixed: Array#

Boolean array indicating if a particle is fixed (immobile). Shape is (…, N).

Identical on every member of a rigid clump.

facet_id: Array#: Array of facet identifiers. Shape is (…, N).

facet_vertices: Array#: Array of vertex indices for each facet. Shape is (…, N, dim).

property N: int[source]#: Number of particles in the state.

property dim: int[source]#: Spatial dimension of the simulation.

property shape: tuple[int, ...][source]#: Shape of the position array pos_c, e.g. (N, dim) or (B, N, dim).

property batch_size: int[source]#: Return the batch size of the state.

property pos: Array[source]#: Returns the position of each sphere in the state. pos_c is the center of mass and pos_p is the vector relative to the center of mass in the principal reference frame, such that pos = pos_c + R(q) @ pos_p where R(q) rotates pos_p to the lab frame.

property is_valid: bool[source]#

Check if the internal representation of the State is consistent.

Verifies that:

The spatial dimension (dim) is either 2 or 3.

All position-like arrays (pos_c, pos_p, vel, force) have the same shape.

All angular-like arrays (ang_vel, torque, inertia) have the same shape.

All scalar-per-particle arrays (rad, mass, clump_id, bond_id, mat_id, species_id, fixed) have a shape consistent with pos.shape[:-1].

Returns:: True if the state is internally consistent, False otherwise.
Return type:: bool

Factory method to create a new State instance.

This method handles default values and ensures consistent array shapes for all state attributes.

Parameters:

pos (jax.typing.ArrayLike or None, optional) – Array of particle center of mass positions, equivalent to state.pos_c. Expected shape: (…, N, dim). If None, an empty state is created. With dim=None, shape is (0, 0) (wildcard empty); with dim=2|3, shape is (0, dim).
dim (int or None, optional) – Spatial dimension used only when pos is None to create an empty state. Must be 2 or 3. If None, an empty state is created with wildcard dimension semantics (it can merge with 2D or 3D states).
pos_p (jax.typing.ArrayLike) – Vector relative to the center of mass (pos_p = pos - pos_c) in the principal reference frame. This field should be constant. Shape is (…, N, dim).
vel (jax.typing.ArrayLike or None, optional) – Initial velocities of particles. If None, defaults to zeros. Expected shape: (…, N, dim).
force (jax.typing.ArrayLike or None, optional) – Initial forces on particles. If None, defaults to zeros. Expected shape: (…, N, dim).
q (Quaternion or array-like, optional) – Initial particle orientations. If None, defaults to identity quaternions. Accepted shapes: quaternion objects or arrays of shape (…, N, 4) with components ordered as (w, x, y, z).
ang_vel (jax.typing.ArrayLike or None, optional) – Initial angular velocities of particles. If None, defaults to zeros. Expected shape: (…, N, 1) in 2D or (…, N, 3) in 3D.
torque (jax.typing.ArrayLike or None, optional) – Initial torques on particles. If None, defaults to zeros. Expected shape: (…, N, 1) in 2D or (…, N, 3) in 3D.
rad (jax.typing.ArrayLike or None, optional) – Radii of particles. If None, defaults to ones. Expected shape: (…, N).
_rad (jax.typing.ArrayLike or None, optional) – Broad-phase search radii of particles, used by cell-list/neighbor colliders to size the contact-detection box. If None, defaults to rad. Expected shape: (…, N).
volume (jax.typing.ArrayLike or None, optional) – Volume of particles (or area in 2D). If None, defaults to hypersphere volumes of the radii. Expected shape: (…, N).
mass (jax.typing.ArrayLike or None, optional) – Masses of particles. If None, defaults to ones. Ignored when mat_table is provided. Expected shape: (…, N).
inertia (jax.typing.ArrayLike or None, optional) – Moments of inertia in the principal axes frame. If None, defaults to solid disks (2D) or spheres (3D). Expected shape: (…, N, 1) in 2D or (…, N, 3) in 3D.
clump_id (jax.typing.ArrayLike or None, optional) – Unique identifiers for clumps. If None, defaults to jnp.arange(). Expected shape: (…, N).
bond_id (jax.typing.ArrayLike or None, optional) – List of connected index values for each particle, storing the indices of the particles it is connected to. Can be passed as a nested list (potentially with uneven lengths), or a 2D array. Connections are automatically symmetrized and padded with -1. If None, defaults to no connections (shape (…, N, 1) filled with -1).
mat_id (jax.typing.ArrayLike or None, optional) – Material IDs for particles. If None, defaults to zeros. Expected shape: (…, N).
species_id (jax.typing.ArrayLike or None, optional) – Species IDs for particles. If None, defaults to zeros. Expected shape: (…, N).
fixed (jax.typing.ArrayLike or None, optional) – Boolean array indicating fixed particles. If None, defaults to all False. Expected shape: (…, N).
facet_id (jax.typing.ArrayLike or None, optional) – Facet identifiers; -1 marks particles that are not facet vertices. If None, defaults to all -1. Expected shape: (…, N). Vertex indices for the facet each particle belongs to; -1 marks non-facet particles. If None, defaults to all -1. Expected shape: (…, N, dim).
mat_table (MaterialTable or None, optional) – Optional material table providing per-material densities. When provided, the mass argument is ignored and particle masses are computed from density and particle volume.

Returns:

A new State instance with all attributes correctly initialized and shaped.

Return type:

State

Raises:

ValueError – If the created State is not valid.

Example

Creating a 3D state for 5 particles:

>>> import jaxdem as jdem
>>> import jax.numpy as jnp
>>>
>>> my_pos = jnp.array([[0.,0.,0.], [1.,0.,0.], [0.,1.,0.], [0.,0.,1.], [1.,1.,1.]])
>>> my_rad = jnp.array([0.5, 0.5, 0.5, 0.5, 0.5])
>>> my_mass = jnp.array([1.0, 1.0, 1.0, 1.0, 1.0])
>>>
>>> state_5_particles = jdem.State.create(pos=my_pos, rad=my_rad, mass=my_mass)
>>> print(f"Shape of positions: {state_5_particles.pos.shape}")
>>> print(f"Radii: {state_5_particles.rad}")

static merge(state1: State, state2: State | Sequence[State]) → State[source]#

Merges multiple State instances into a single new State.

This method concatenates the particles from the provided state(s) onto state1. Particle clump_ids and bond_ids are shifted to ensure uniqueness across the merged system.

Parameters:

state1 (State) – The first State instance. Its particles will appear first in the merged state.
state2 (State or Sequence[State]) – The second State or a list/tuple of State instances to append.

Returns:

A new State instance containing all particles from both input states.

Return type:

State

Raises:

AssertionError – If either input state is invalid, or if there is a mismatch in spatial dimension (dim) or batch size (batch_size).
ValueError – If the resulting merged state is somehow invalid.

Example

>>> import jaxdem as jdem
>>> import jax.numpy as jnp
>>>
>>> state_a = jdem.State.create(pos=jnp.array([[0.0, 0.0], [1.0, 1.0]]), clump_id=jnp.array([0, 1]))
>>> state_b = jdem.State.create(pos=jnp.array([[2.0, 2.0], [3.0, 3.0]]), clump_id=jnp.array([0, 1]))
>>> merged_state = jdem.State.merge(state_a, [state_b, state_b, state_b])
>>> merged_state = jdem.State.merge(state_a, state_b)
>>>
>>> print(f"Merged state N: {merged_state.N}")  # Expected: 4
>>> print(f"Merged state positions:\\n{merged_state.pos}")
>>> print(f"Merged state clump_ids: {merged_state.clump_id}")  # Expected: [0, 1, 2, 3]

Adds new particles to an existing State instance, returning a new State.

Parameters:

state (State) – The existing State to which particles will be added.
pos (jax.typing.ArrayLike) – Array of particle center of mass positions, equivalent to state.pos_c. Expected shape: (…, N, dim).
pos_p (jax.typing.ArrayLike) – Vector relative to the center of mass (pos_p = pos - pos_c) in the principal reference frame. This field should be constant. Shape is (…, N, dim).
vel (jax.typing.ArrayLike or None, optional) – Velocities of the new particle(s). Defaults to zeros.
force (jax.typing.ArrayLike or None, optional) – Forces of the new particle(s). Defaults to zeros.
q (Quaternion or array-like, optional) – Initial orientations of the new particle(s). Defaults to identity quaternions.
ang_vel (jax.typing.ArrayLike or None, optional) – Angular velocities of the new particle(s). Defaults to zeros.
torque (jax.typing.ArrayLike or None, optional) – Torques of the new particle(s). Defaults to zeros.
rad (jax.typing.ArrayLike or None, optional) – Radii of the new particle(s). Defaults to ones.
_rad (jax.typing.ArrayLike or None, optional) – Broad-phase search radii of the new particle(s), used by cell-list/neighbor colliders to size the contact-detection box. Defaults to rad.
volume (jax.typing.ArrayLike or None, optional) – Volume of the new particle(s) (or area in 2D). Defaults to hypersphere volumes of the radii.
mass (jax.typing.ArrayLike or None, optional) – Masses of the new particle(s). Defaults to ones. Ignored when a mat_table is provided.
inertia (jax.typing.ArrayLike or None, optional) – Moments of inertia of the new particle(s). Defaults to solid disks (2D) or spheres (3D).
clump_id (jax.typing.ArrayLike or None, optional) – clump_ids of the new clump(s). If None, new IDs are generated.
bond_id (jax.typing.ArrayLike or None, optional) – List of connected index values for each particle, storing the indices of the particles it is connected to. Can be passed as a nested list (potentially with uneven lengths), or a 2D array. Connections are automatically symmetrized and padded with -1. If None, defaults to no connections.
mat_id (jax.typing.ArrayLike or None, optional) – Material IDs of the new particle(s). Defaults to zeros.
species_id (jax.typing.ArrayLike or None, optional) – Species IDs of the new particle(s). Defaults to zeros.
fixed (jax.typing.ArrayLike or None, optional) – Fixed status of the new particle(s). Defaults to all False.
mat_table (MaterialTable or None, optional) – Optional material table providing per-material densities. When provided, masses are computed from density and particle volume.

Returns:

A new State instance containing all particles from the original state plus the newly added particles.

Return type:

State

Raises:

ValueError – If the created new particle state or the merged state is invalid.
AssertionError – If batch size or dimension mismatch between existing state and new particles.

Example

>>> import jaxdem as jdem
>>> import jax.numpy as jnp
>>>
>>> # Initial state with 4 particles
>>> state = jdem.State.create(pos=jnp.zeros((4, 2)))
>>> print(f"Original state N: {state.N}, clump_ids: {state.clump_id}")
>>>
>>> # Add a single new particle
>>> state_with_added_particle = jdem.State.add(
...     state,
...     pos=jnp.array([[10.0, 10.0]]),
...     rad=jnp.array([0.5]),
...     mass=jnp.array([2.0]),
... )
>>> print(f"New state N: {state_with_added_particle.N}, clump_ids: {state_with_added_particle.clump_id}")
>>> print(f"New particle position: {state_with_added_particle.pos[-1]}")
>>>
>>> # Add multiple new particles
>>> state_multiple_added = jdem.State.add(
...     state,
...     pos=jnp.array([[10.0, 10.0], [11.0, 11.0], [12.0, 12.0]]),
... )
>>> print(f"State with multiple added N: {state_multiple_added.N}, clump_ids: {state_multiple_added.clump_id}")

static stack(states: Sequence[State]) → State[source]#

Concatenates a sequence of State snapshots into a trajectory or batch along axis 0.

This method is useful for collecting simulation snapshots over time into a single State object where the leading dimension represents time or when preparing a batched state.

Parameters:

states (Sequence[State]) – A sequence (e.g., list, tuple) of State instances to be stacked.

Returns:

A new State instance where each attribute is a JAX array with an additional leading dimension representing the stacked trajectory. For example, if input pos was (N, dim), output pos will be (T, N, dim).

Return type:

State

Raises:

ValueError – If the input states sequence is empty. If the stacked State is invalid.
AssertionError – If any input state is invalid, or if there is a mismatch in spatial dimension (dim), batch size (batch_size), or number of particles (N) between the states in the sequence.

Notes

No clump_id shifting is performed because the leading axis represents time (or another batch dimension), not new particles.

Example

>>> import jaxdem as jdem
>>> import jax.numpy as jnp
>>>
>>> # Create a sequence of 3 simple 2D snapshots
>>> snapshot1 = jdem.State.create(pos=jnp.array([[0.,0.], [1.,1.]]), vel=jnp.array([[0.1,0.], [0.0,0.1]]))
>>> snapshot2 = jdem.State.create(pos=jnp.array([[0.1,0.], [1.,1.1]]), vel=jnp.array([[0.1,0.], [0.0,0.1]]))
>>> snapshot3 = jdem.State.create(pos=jnp.array([[0.2,0.], [1.,1.2]]), vel=jnp.array([[0.1,0.], [0.0,0.1]]))
>>>
>>> trajectory_state = State.stack([snapshot1, snapshot2, snapshot3])
>>>
>>> print(f"Trajectory positions shape: {trajectory_state.pos.shape}") # Expected: (3, 2, 2)
>>> print(f"Positions at time step 0:\\n{trajectory_state.pos[0]}")
>>> print(f"Positions at time step 1:\\n{trajectory_state.pos[1]}")

static unstack(state: State) → list[State][source]#

Split a stacked/batched State along the leading axis into a Python list.

This is the convenient inverse of State.stack():

If stacked = State.stack([s0, s1, …]), then State.unstack(stacked) returns [s0, s1, …].

Notes

The split is performed along axis 0 (the leading axis).
A single snapshot State (e.g. pos.shape == (N, dim)) cannot be unstacked with this method, because axis 0 would refer to particles, not snapshots.

Adds a new clump consisting of multiple spheres to an existing State. Rigid body properties (center of mass position pos_c, velocity, mass, orientation q, angular velocity, force, torque, inertia, fixed, and clump_id) are broadcasted/shared by all spheres in the new clump. The per-sphere properties that can vary within a rigid clump are pos_p (offsets in the body reference frame), rad, and individual ID fields (mat_id, species_id, and bond_id).

Parameters:

state (State) – The existing State to which particles will be added.
pos (jax.typing.ArrayLike) – If pos_p is None, this represents the absolute coordinates of the spheres in the clump. If pos_p is not None, this represents the center of mass (COM) position of the clump. Expected shape: (…, N, dim).
pos_p (jax.typing.ArrayLike or None, optional) – Vector relative to the center of mass (pos_p = pos - pos_c) in the principal reference frame. This field should be constant. Shape is (…, N, dim). If None, it is computed from the absolute coordinates provided in pos and the center of mass is calculated using the sphere volume weights.
vel (jax.typing.ArrayLike or None, optional) – Velocities of the new particle(s). Defaults to zeros.
force (jax.typing.ArrayLike or None, optional) – Forces of the new particle(s). Defaults to zeros.
q (Quaternion or array-like, optional) – Initial orientations of the new particle(s). Defaults to identity quaternions.
ang_vel (jax.typing.ArrayLike or None, optional) – Angular velocities of the new particle(s). Defaults to zeros.
torque (jax.typing.ArrayLike or None, optional) – Torques of the new particle(s). Defaults to zeros.
rad (jax.typing.ArrayLike or None, optional) – Radii of the new particle(s). Defaults to ones.
volume (jax.typing.ArrayLike or None, optional) – Volume of the new particle(s) (or area in 2D). Defaults to hypersphere volumes of the radii.
mass (jax.typing.ArrayLike or None, optional) – Masses of the new particle(s). Defaults to ones.
inertia (jax.typing.ArrayLike or None, optional) – Moments of inertia of the new particle(s). Defaults to solid disks (2D) or spheres (3D).
bond_id (jax.typing.ArrayLike or None, optional) – List of connected index values for each particle, storing the indices of the particles it is connected to. Can be passed as a nested list (potentially with uneven lengths), or a 2D array. Connections are automatically symmetrized and padded with -1. If None, defaults to no connections.
mat_id (jax.typing.ArrayLike or None, optional) – Material IDs of the new particle(s). Defaults to zeros.
species_id (jax.typing.ArrayLike or None, optional) – Species IDs of the new particle(s). Defaults to zeros.
fixed (jax.typing.ArrayLike or None, optional) – Fixed status of the new particle(s). Defaults to all False.

Returns:

A new State instance containing all particles from the original state plus the newly added particles.

Return type:

State

Adds a new facet clump (2D line segment or 3D triangle) consisting of vertex spheres to an existing State.

Note: Facets added via this method do not share vertices (i.e. each facet has its own independent copy of vertex particles).

Parameters:

state (State) – The existing state.
vertices (ArrayLike) – Vertices of the facets, shape (…, V, dim).
vel (ArrayLike or None, optional) – Initial linear velocity.
force (ArrayLike or None, optional) – Initial force.
q (Quaternion or ArrayLike or None, optional) – Initial orientations.
ang_vel (ArrayLike or None, optional) – Initial angular velocities.
torque (ArrayLike or None, optional) – Initial torques.
thickness (float, optional) – Physical thickness/radius of the facet vertex spheres.
mass (ArrayLike or None, optional) – Mass of the facets.
mat_id (ArrayLike or None, optional) – Material IDs.
species_id (ArrayLike or None, optional) – Species IDs.
fixed (ArrayLike or None, optional) – Whether the facet vertices are fixed in space.
rigid (bool, default True) – If True, the facet is rigid, grouping all vertices under the same clump ID, with proper inertia and orientation calculations. If False, the facet is flexible/deformable, meaning its vertices behave like individual spheres (having standard sphere moment of inertia, identity orientation, and unique clump IDs).
safety_factor (float, default 1.0) – Factor to scale/multiply _rad to enlarge the broad-phase detection box.

Adds a new mesh (collection of facets) consisting of vertex spheres to an existing State.

Parameters:

state (State) – The existing state.
vertices (ArrayLike) – Vertices of the mesh, shape (…, V_mesh, dim).
faces (ArrayLike) – Faces of the mesh (indices into vertices), shape (…, F, dim).
vel (ArrayLike or None, optional) – Initial linear velocity.
force (ArrayLike or None, optional) – Initial force.
q (Quaternion or ArrayLike or None, optional) – Initial orientations.
ang_vel (ArrayLike or None, optional) – Initial angular velocities.
torque (ArrayLike or None, optional) – Initial torques.
thickness (float, optional) – Physical thickness/radius of the facet vertex spheres.
mass (ArrayLike or None, optional) – Mass of the mesh.
mat_id (ArrayLike or None, optional) – Material IDs.
species_id (ArrayLike or None, optional) – Species IDs.
fixed (ArrayLike or None, optional) – Whether the facet vertices are fixed in space.
rigid (bool, default True) – If True, the mesh is rigid, grouping all vertices under the same clump ID, with proper inertia and orientation calculations. If False, the mesh is flexible/deformable.
filled (bool, default True) – If True, the mesh represents a filled solid polyhedron/polygon. If False, it represents a hollow boundary shell.
safety_factor (float, default 1.0) – Factor to scale/multiply _rad to enlarge the broad-phase detection box.

Adds a new facet connecting existing vertices and/or newly added vertices to the State.

Parameters:

state (State) – The existing state.
vertex_specs (list of int or ArrayLike) – Each spec represents a vertex of the new facet. If a spec is a scalar integer (or scalar array), it is treated as the index of an existing vertex. Otherwise, it is treated as a position array of shape (dim,) for a new vertex.