jaxdem#

JaxDEM module.

class jaxdem.BondedForceModel#

Bases: Factory, ABC

Abstract interface for bonded interaction containers.

Design intent#

System stores one concrete bonded container instance.
The container exposes bonded force/energy callables through force_and_energy_fns.
The ForceManager gets and runs these callables to compute bonded contributions at each time step.
The bonded data stays accessible through System, so the force/energy callables can read what they need.

property force_and_energy_fns: tuple[ForceFunction, EnergyFunction, bool][source]#

Build the bonded force/energy callables for the force manager.

Returns:

(force_fn, energy_fn, is_com_force) where:

force_fn computes bonded force and torque contributions.
energy_fn computes bonded potential-energy contributions.
is_com_force tells where the force acts: True for the center of mass, False for the contact point. It has no effect on spheres.

Return type:

Tuple[ForceFunction, EnergyFunction, bool]

abstractmethod compute_potential_energy(pos: jax.Array, state: State, system: System) → jax.Array[source]#: Compute the total bonded potential energy of the system.

abstractmethod static merge(model1: BondedForceModel, model2: BondedForceModel | Sequence[BondedForceModel]) → BondedForceModel[source]#

Merge two or more bonded-force models into one.

The merge concatenates topology, reference, and coefficient arrays. It shifts vertex indices and body IDs automatically so references stay consistent. When one side has a term that the other does not, the merge pads missing coefficients with 0 and missing reference values with 1.

Parameters:

model1 (BondedForceModel) – Base model.
model2 (BondedForceModel or Sequence[BondedForceModel]) – Model(s) to merge into model1.

Returns:

A new model containing all bodies from both sides.

Return type:

BondedForceModel

abstractmethod static add(model: BondedForceModel, **kwargs: Any) → BondedForceModel[source]#

Create a new body from raw arrays and merge it into an existing model.

This is a convenience wrapper equivalent to calling the concrete Create constructor followed by merge().

Parameters:

model (BondedForceModel) – Existing model to extend.
**kwargs – Constructor arguments forwarded to the concrete Create method (e.g. vertices, elements, coefficients, …).

Returns:

The extended model.

Return type:

BondedForceModel

classmethod stack(models: Sequence[BondedT]) → BondedT[source]#

classmethod unstack(model: BondedT) → list[BondedT][source]#

class jaxdem.CheckpointLoader(directory: Path | str = PosixPath('checkpoints'))#

Bases: BaseCheckpointManager

Thin wrapper around Orbax checkpoint restoring for jaxdem.state and jaxdem.system.

load(step: int | None = None, *, strict: bool = True) → tuple[State, System][source]#

Restore a checkpoint.

Parameters:

step (Optional[int]) –
- If None, load the latest checkpoint.
- Otherwise, load the specified step.
strict (bool, optional) – If True (default), raise a RuntimeError when the loader cannot re-import a custom force function recorded in the checkpoint. If False, skip force functions that do not load, with a warning.

Returns:

The restored State and System.

Return type:

Tuple[State, System]

latest_step() → int | None[source]#

class jaxdem.CheckpointModelLoader(directory: Path | str = PosixPath('checkpoints'))#

Bases: BaseCheckpointManager

Thin wrapper around Orbax checkpoint restoring for jaxdem.rl.models.Model.

load(step: int | None = None) → Model[source]#: Load the model at the given step. If step is None, load the latest checkpoint.

latest_step() → int | None[source]#

class jaxdem.CheckpointModelWriter(directory: Path | str = PosixPath('checkpoints'), max_to_keep: int | None = None, save_every: int = 1, clean: bool = False)#

Bases: BaseCheckpointManager

Thin wrapper around Orbax checkpoint saving for jaxdem.rl.models.Model.

max_to_keep: int | None = None#: Keep the last max_to_keep checkpoints. If None, keep all checkpoints.

save_every: int = 1#: How often to write. The writer saves on every save_every-th call to save().

clean: bool = False#: If True, erase and recreate the target directory on construction. If False (the default), keep existing checkpoints in the directory, so a resumed run does not destroy earlier checkpoints.

save(model: Model, step: int) → None[source]#

Save the model at the given step.

Stores model_state and JSON metadata. model.metadata must contain JSON-serializable fields. The writer adds model_type to the metadata.

class jaxdem.CheckpointWriter(directory: Path | str = PosixPath('checkpoints'), max_to_keep: int | None = None, save_every: int = 1, clean: bool = False)#

Bases: BaseCheckpointManager

Thin wrapper around Orbax checkpoint saving.

Notes

The writer serializes custom force functions passed via force_manager_kw by their fully-qualified module path (e.g. mypackage.forces.trap). A different script cannot restore functions defined in the top-level script (__main__). The writer emits a warning at save time if any force function lives in __main__. To keep checkpoints portable, define force functions in an importable module.

max_to_keep: int | None = None#: Keep the last max_to_keep checkpoints. If None, keep all checkpoints.

save_every: int = 1#: How often to write. The writer saves on every save_every-th call to save().

clean: bool = False#: If True, erase and recreate the target directory on construction. If False (the default), keep existing checkpoints in the directory, so a resumed run does not destroy earlier checkpoints.

save(state: State, system: System) → None[source]#

Save a checkpoint of the given state and system at the current step.

Parameters:

state (State) – The current state of the simulation.
system (System) – The current system configuration.

class jaxdem.Collider(*, overflow: Array = <factory>)#

Bases: Factory, ABC

The base interface for contact detection and force computation in a simulation.

Concrete subclasses of Collider implement the interaction algorithms.

Notes:#

Self-interaction (calling the force/energy computation for i=j) is allowed. The force_model must handle or ignore this case correctly.

Example:#

To define a custom collider, inherit from Collider, register it, and implement its abstract methods:

>>> @Collider.register("CustomCollider")
>>> @jax.tree_util.register_dataclass
>>> @dataclass(slots=True)
>>> class CustomCollider(Collider):
        ...

Then, instantiate it:

>>> jaxdem.Collider.create("CustomCollider", **custom_collider_kw)

overflow: Array#: True when a collider overflow occurred.

static compute_force(state: State, system: System) → tuple[State, System][source]#

Compute the total force acting on each particle in the simulation.

This base implementation is a concrete no-op: it zeroes the force and torque attributes of the state and returns. It backs the "" (empty-string) no-op collider registration for systems whose dynamics come only from bonded forces or user force functions.

Subclasses override it to compute inter-particle forces and torques from the current state and system configuration. They write the total force and torque of each particle to the force and torque attributes of the state object.

Parameters:

state (State) – The current state of the simulation.
system (System) – The configuration of the simulation.

Returns:

A tuple containing the updated State object (with computed forces) and the System object.

Return type:

Tuple[State, System]

static compute_potential_energy(state: State, system: System) → tuple[State, System, jax.Array][source]#

Compute the total (scalar) non-bonded potential energy of the system.

Implementations sum every pair-interaction contribution defined by system.force_model and return a single scalar. They weight pair energies with the standard 0.5 factor, so each pair counts once even when the neighbor list visits (i, j) and (j, i) separately.

Parameters:

state (State) – The current state of the simulation.
system (System) – The configuration of the simulation.

Returns:

A tuple of (state, system, potential_energy). The potential_energy is a scalar JAX array (shape ()) with the total non-bonded potential energy of the system.

Return type:

Tuple[State, System, jax.Array]

Example

>>> state, system, potential_energy = system.collider.compute_potential_energy(state, system)
>>> print(f"Total potential energy: {float(potential_energy):.4f}")
>>> print(potential_energy.shape)  # ()

static create_neighbor_list(state: State, system: System, cutoff: float, max_neighbors: int) → tuple[State, System, jax.Array, jax.Array][source]#

Build a neighbor list for the current collider.

Neighbor-list-based algorithms and diagnostics use this list. Implementations match the cell-list semantics:

Return a neighbor list of shape (N, max_neighbors) padded with -1.
Neighbor indices refer to the returned state.
Also return an overflow boolean flag. The flag is True when any particle has more than max_neighbors neighbors within the cutoff.

static create_cross_neighbor_list(pos_a: jax.Array, pos_b: jax.Array, system: System, cutoff: float, max_neighbors: int) → tuple[jax.Array, jax.Array][source]#

Build a cross-neighbor list between two sets of positions.

For each point in pos_a, find all neighbors from pos_b within the cutoff distance. Use this to couple different particle systems or to compute interactions between distinct sets of objects.

The default implementation runs a naive \(O(N_A \times N_B)\) all-pairs search. Subclasses can override it with faster algorithms.

Parameters:

pos_a (jax.Array) – Query positions, shape (N_A, dim).
pos_b (jax.Array) – Database positions, shape (N_B, dim).
system (System) – The configuration of the simulation (used for domain displacement).
cutoff (float) – Search radius.
max_neighbors (int) – Maximum number of neighbors to store per query point.

Returns:

A tuple containing:

neighbor_list: Array of shape (N_A, max_neighbors) containing indices into pos_b, padded with -1.
overflow: Boolean flag. True when any query point has more than max_neighbors neighbors within the cutoff.

Return type:

Tuple[jax.Array, jax.Array]

class jaxdem.Domain(box_size: Array, inv_box_size: Array, anchor: Array)#

Bases: Factory, ABC

The base interface for the simulation domain and its boundary conditions.

The Domain class defines:

How to compute relative displacement vectors between particles.
How to “shift” or constrain particle positions so they stay within the simulation boundaries.

Example:#

To define a custom domain, inherit from Domain and implement its abstract methods:

>>> @Domain.register("my_custom_domain")
>>> @jax.tree_util.register_dataclass
>>> @dataclass(slots=True)
>>> class MyCustomDomain(Domain):
        ...

box_size: Array#: Length of the simulation domain along each dimension.

inv_box_size: Array#: Inverse length of the simulation domain along each dimension.

anchor: Array#: Anchor position (minimum coordinate) of the simulation domain.

property periodic: bool[source]#: Whether the domain enforces periodic boundary conditions.

classmethod Create(dim: int, box_size: Array | None = None, anchor: Array | None = None, **kw: Any) → Self[source]#

Default factory method for the Domain class.

This method constructs a new Domain instance with a box-shaped domain of the given dimensionality. If you do not provide box_size or anchor, they default to the values below.

Parameters:

dim (int) – The dimensionality of the domain (e.g., 2, 3).
box_size (jax.Array, optional) – The size of the domain along each dimension. If not provided, defaults to an array of ones with shape (dim,).
anchor (jax.Array, optional) – The anchor (origin) of the domain. If not provided, defaults to an array of zeros with shape (dim,).
**kw (Any) – Extra keyword arguments passed to the subclass constructor (e.g. restitution_coefficient for reflective domains).

Returns:

A new instance of the Domain subclass with the specified or default configuration.

Return type:

Domain

Raises:

ValueError – If box_size or anchor do not have shape (dim,).

static displacement(ri: jax.Array, rj: jax.Array, system: System) → jax.Array[source]#

Compute the displacement vector between two particles \(r_i\) and \(r_j\), respecting the domain’s boundary conditions.

Parameters:

ri (jax.Array) – Position vector of the first particle \(r_i\). Shape (dim,).
rj (jax.Array) – Position vector of the second particle \(r_j\). Shape (dim,).
system (System) – The configuration of the simulation, containing the domain instance.

Returns:

The displacement vector \(r_{ij} = r_i - r_j\), adjusted for boundary conditions. Shape (dim,).

Return type:

jax.Array

Example

>>> rij = system.domain.displacement(ri, rj, system)

static apply(state: State, system: System) → tuple[State, System][source]#

Apply boundary conditions during the simulation step.

This method updates the state with the domain’s rules so particles handle boundary interactions (e.g., reflection).

Parameters:

state (State) – The current state of the simulation.
system (System) – The configuration of the simulation.

Returns:

A tuple containing the updated State object adjusted by the boundary conditions and the System object.

Return type:

Tuple[State, System]

Note

Periodic domains do not need to wrap coordinates during time stepping, so their apply is a no-op. shift() wraps the coordinates instead (e.g. when saving, so positions are displayed inside the box). Reflective domains, in contrast, must update positions and velocities here.

Example

>>> state, system = system.domain.apply(state, system)

static shift(state: State, system: System) → tuple[State, System][source]#

Shift particles according to the domain’s boundary-condition rules.

This method updates the state with the domain’s rules so particles stay within the simulation box or handle boundary interactions (e.g., reflection, wrapping).

Parameters:

state (State) – The current state of the simulation.
system (System) – The configuration of the simulation.

Returns:

A tuple containing the updated State object adjusted by the boundary conditions and the System object.

Return type:

Tuple[State, System]

Example

>>> state, system = system.domain.shift(state, system)

class jaxdem.Factory#

Bases: ABC

Base class for components that register and create subclasses by a string key.

Notes:#

Each concrete subclass gets its own private registry. The factory normalizes keys before use: lookup is case-insensitive and ignores spaces, underscores, and hyphens ("CellList", "cell_list", and "celllist" are the same key).

Example:#

Use Factory as a base class for a specific component type (e.g., Foo):

>>> class Foo(Factory["Foo"], ABC):
>>>   ...

>>> @Foo.register("bar")
>>> class bar:
>>>     ...

To create an instance of the subclass:

>>> Foo.create("bar", **bar_kw)

classmethod registry_name() → str[source]#: Return the registry key of this class.

property type_name: str[source]#: Return the registry key of the class of this instance.

property metadata: dict[str, Any][source]#: Serialize the component’s dataclass fields for checkpointing and restoration.

classmethod register(key: str | None = None) → Callable[[type[SubT]], type[SubT]][source]#

This method returns a decorator that registers a class under a specific key.

Parameters:: key (str or None, optional) – The string key under which to register the subclass. If None, the method uses the lowercase subclass name as the key. The method normalizes keys (lowercase, without spaces, underscores, and hyphens), so "CellList", "cell_list", and "celllist" all denote the same key.
Returns:: A decorator that registers the class and returns it unchanged.
Return type:: Callable[[Type[T]], Type[T]]
Raises:: ValueError – If the provided key (or the default class name) is already registered in the factory’s registry for a different class. Registering the same class under the same key again (for example when you re-run a notebook cell) works and is idempotent.

Example

>>> @MyFactory.register("mycomp")
>>> class MyComponent:
>>>     ...

>>> @MyFactory.register()
>>> class DefaultComponent:
>>>     ...

classmethod create(key: str, /, **kw: Any) → RootT[source]#

Create and return an instance of a registered subclass.

This method looks up the subclass registered under the given key and calls its constructor with the provided arguments. If the subclass defines a Create method (capitalized), the factory calls that method instead of the constructor. This lets subclasses validate or preprocess arguments before the factory creates the instance.

Parameters:

key (str) – The registration key of the subclass to create.
**kw (Any) – Keyword arguments passed to the constructor of the registered subclass.

Returns:

An instance of the registered subclass.

Return type:

Raises:

KeyError – If the factory’s registry does not contain the provided key.
TypeError – If the provided **kw arguments do not match the signature of the registered subclass’s constructor.

Example

Given Foo factory and Bar registered:

>>> bar_instance = Foo.create("bar", value=42)
>>> print(bar_instance)
Bar(value=42)

class jaxdem.ForceManager(gravity: jax.Array, external_force: jax.Array, external_force_com: jax.Array, external_torque: jax.Array, is_com_force: tuple[bool, ...] = (), force_functions: tuple[ForceFunction, ...] = (), energy_functions: tuple[EnergyFunction | None, ...] = ())#

Bases: object

Manage custom force contributions outside the collider.

After the collider runs, apply() adds these contributions to the state forces and aggregates them over rigid bodies.

gravity: jax.Array#: Constant acceleration applied to all particles. Shape (dim,).

external_force: jax.Array#: Accumulated external force applied to all particles (at particle position). apply() clears this buffer.

external_force_com: jax.Array#: Accumulated external force applied to the center of mass (induces no torque). apply() clears this buffer.

external_torque: jax.Array#: Accumulated external torque applied to all particles. apply() clears this buffer.

is_com_force: tuple[bool, ...]#: Boolean array corresponding to force_functions with shape (n_forces,). If True, apply the force to the center of mass (no induced torque). If False, apply the force at the particle position (induces torque through the lever arm).

force_functions: tuple[ForceFunction, ...]#: Tuple of callables with signature (pos, state, system) returning per-particle force and torque arrays.

energy_functions: tuple[EnergyFunction | None, ...]#: Tuple of callables (or None) with signature (pos, state, system) returning per-particle potential energy arrays. Corresponds to force_functions.

Create a ForceManager for a state with the given shape.

Parameters:

state_shape – Shape of the state position array, typically (..., dim).
gravity – Optional initial gravitational acceleration. Defaults to zeros of shape (dim,).
force_functions –
Sequence of callables or tuples. Signature of ForceFunc: (pos, state, system) -> (Force, Torque). Signature of EnergyFunc: (pos, state, system) -> Energy. Supported formats:
- func -> (func, None, False)
- (func,) -> (func, None, False)
- (func, bool) -> (func, None, bool)
- (func, energy) -> (func, energy, False)
- (func, energy, bool) -> (func, energy, bool)
- (func, None, bool) -> (func, None, bool)

static add_force(state: State, system: System, force: jax.Array, *, is_com: bool = False) → System[source]#

Buffer an external force on all particles for the next apply call.

The method returns only system. The state does not change because the force waits in the ForceManager buffer until apply() runs.

Parameters:

state (State) – Current state of the simulation. Used to normalize COM forces by the clump member count.
system (System) – Simulation system configuration.
force (jax.Array) – External force to add to every particle (in the state’s current particle order).
is_com (bool, optional) – If True, apply the force to the center of mass (no induced torque). The force goes to every clump member, so each clump receives force in total, not force per member. If False (default), apply the force at the particle position (induces torque).

static add_force_at(state: State, system: System, force: jax.Array, idx: jax.Array, *, is_com: bool = False) → System[source]#

Buffer an external force on the particles with array index idx for the next apply call.

The method returns only system. The state does not change because the force waits in the ForceManager buffer until apply() runs.

Parameters:

state (State) – Current state of the simulation.
system (System) – Simulation system configuration.
force (jax.Array) – External force to add to the particles with array index idx.
idx (jax.Array) – Array indices of the particles the external force acts on.
is_com (bool, optional) – If True, apply the force to the center of mass (no induced torque). If False (default), apply the force at the particle position (induces torque).

static add_torque(state: State, system: System, torque: jax.Array) → System[source]#

Buffer an external torque on all particles for the next apply call.

The method returns only system. The state does not change because the torque waits in the ForceManager buffer until apply() runs.

Parameters:

state (State) – Current state of the simulation. Used to normalize the torque by the clump member count.
system (System) – Simulation system configuration.
torque (jax.Array) – External torque to add to every particle (in the state’s current particle order). The torque goes to every clump member, so each clump receives torque in total, not torque per member.

static add_torque_at(state: State, system: System, torque: jax.Array, idx: jax.Array) → System[source]#

Buffer an external torque on the particles with array index idx for the next apply call.

The method returns only system. The state does not change because the torque waits in the ForceManager buffer until apply() runs.

Parameters:

state (State) – Current state of the simulation.
system (System) – Simulation system configuration.
torque (jax.Array) – External torque to add to the particles with array index idx.
idx (jax.Array) – Array indices of the particles the external torque acts on.

static apply(state: State, system: System) → tuple[State, System][source]#

Add the managed per-particle contributions to the collider forces, then aggregate over clumps and broadcast back to the members.

Parameters:

state (State) – Current state of the simulation.
system (System) – Simulation system configuration.

Returns:

The updated state and system.

Return type:

Tuple[State, System]

static compute_potential_energy(state: State, system: System) → jax.Array[source]#

Compute the total potential energy of the system.

Notes

The energy of clump members is divided by the number of spheres in the clump.

Parameters:

state (State) – Current state of the simulation.
system (System) – Simulation system configuration.

Returns:

Scalar total potential energy.

Return type:

jax.Array

class jaxdem.ForceModel(laws: tuple[ForceModel, ...] = ())#

Bases: Factory, ABC

Abstract base class for inter-particle force laws and their potential energies.

Concrete subclasses implement specific force and energy models, such as linear springs and Hertzian contacts.

Notes:#

The force() and energy() methods must handle the case where i and j refer to the same particle (i == j). Self-interaction calls can occur.

Example:#

To define a custom force model, inherit from ForceModel and implement its abstract methods:

>>> @ForceModel.register("myCustomForce")
>>> @jax.tree_util.register_dataclass
>>> @dataclass(slots=True)
>>> class MyCustomForce(ForceModel):
        ...

laws: tuple[ForceModel, ...]#

Static tuple of other ForceModel instances that compose this force model.

Use it to build composite force models, for example a spring force plus a damping force.

abstractmethod static force(i: int, j: int, pos: jax.Array, state: State, system: System) → tuple[jax.Array, jax.Array][source]#

Compute the force and torque on particle \(i\) from particle \(j\).

Parameters:

i (int) – Index of the first particle (on which the interaction acts).
j (int) – Index of the second particle (which exerts the interaction).
pos (jax.Array) – Particle positions.
state (State) – Current state of the simulation.
system (System) – Simulation system configuration.

Returns:

A tuple (force, torque) where force has shape (dim,) and torque has shape (1,) in 2D or (3,) in 3D.

Return type:

Tuple[jax.Array, jax.Array]

abstractmethod static energy(i: int, j: int, pos: jax.Array, state: State, system: System) → jax.Array[source]#

Compute the potential energy of the interaction between particle \(i\) and particle \(j\).

Parameters:

i (int) – Index of the first particle.
j (int) – Index of the second particle.
pos (jax.Array) – Particle positions.
state (State) – Current state of the simulation.
system (System) – Simulation system configuration.

Returns:

Scalar potential energy of the interaction between particles \(i\) and \(j\).

Return type:

jax.Array

property requires_history: bool[source]#: Whether this force model needs persistent pair history.

init_history(shape: tuple[int, ...]) → Any[source]#

Initialize the history variables for this force model.

Parameters:: shape (tuple[int, ...]) – The expected shape for pair-wise quantities, typically (…, N, max_neighbors).
Returns:: A PyTree of initialized JAX arrays, or None by default.
Return type:: Any

static force_and_history(i: int, j: int, pos: jax.Array, state: State, system: System, history: Any) → tuple[jax.Array, jax.Array, Any][source]#

Compute the force and torque, and update history.

By default, this calls force and returns history unchanged.

property required_material_properties: tuple[str, ...][source]#

Names of the material properties this force model needs.

Each name (for example ‘young_eff’ or ‘restitution’) must be present in System.mat_table. Used for validation.

class jaxdem.ForceRouter(laws: tuple[ForceModel, ...] = (), table: tuple[tuple[ForceModel, ...], ...] = ())#

Bases: ForceModel

A ForceModel that selects the force law from the species of the interacting particles.

The router holds a symmetric \(S \times S\) lookup table of force laws, where \(S\) is the number of species. For a particle pair \((i, j)\), the router evaluates the law at table[species_id[i]][species_id[j]].

Notes

Use from_dict() to build the table from a mapping of species pairs. Pairs not present in the mapping default to an empty LawCombiner, which produces zero force, torque, and energy.
Dispatch evaluates every law in the table and selects the result with jax.lax.select_n(). The cost grows quadratically with the number of species, for scalar and batched calls alike.
required_material_properties is the union of the requirements of all laws in the table.

table: tuple[tuple[ForceModel, ...], ...]#: A symmetric \(S \times S\) table where entry table[a][b] is the ForceModel that governs interactions between species a and b.

property requires_history: bool[source]#: Whether this force model needs persistent pair history.

init_history(shape: tuple[int, ...]) → Any[source]#

static force_and_history(i: int, j: int, pos: jax.Array, state: State, system: System, history: Any) → tuple[jax.Array, jax.Array, Any][source]#

property required_material_properties: tuple[str, ...][source]#

Names of the material properties this force model needs.

The sorted union of the material properties required by all laws in the table. Each name must be present in System.mat_table. Used for validation.

static from_dict(S: int, mapping: dict[tuple[int, int], ForceModel]) → ForceRouter[source]#

Build a ForceRouter from a mapping of species pairs to force laws.

The router symmetrizes the mapping: entry (a, b) also fills (b, a). Pairs not present in the mapping default to an empty LawCombiner (zero force, torque, and energy).

Parameters:

S (int) – Number of species. The resulting table has shape S x S.
mapping (dict[tuple[int, int], ForceModel]) – Mapping from species-index pairs to the force law that governs interactions between those species.

Returns:

A router with the fully populated, symmetric lookup table.

Return type:

ForceRouter

static force(i: int, j: int, pos: jax.Array, state: State, system: System) → tuple[jax.Array, jax.Array][source]#

Compute the force and torque on particle \(i\) from particle \(j\) with the law their species select.

Parameters:

i (int) – Index of the first particle.
j (int) – Index of the second particle.
pos (jax.Array) – Particle positions used to evaluate the interaction.
state (State) – Current state of the simulation.
system (System) – Simulation system configuration.

Returns:

A tuple (force, torque) computed by the law at table[species_id[i]][species_id[j]].

Return type:

Tuple[jax.Array, jax.Array]

static energy(i: int, j: int, pos: jax.Array, state: State, system: System) → jax.Array[source]#

Compute the potential energy of the interaction between particle \(i\) and particle \(j\) with the law their species select.

Parameters:

i (int) – Index of the first particle.
j (int) – Index of the second particle.
pos (jax.Array) – Particle positions used to evaluate the interaction.
state (State) – Current state of the simulation.
system (System) – Simulation system configuration.

Returns:

Scalar potential energy computed by the law at table[species_id[i]][species_id[j]].

Return type:

jax.Array

class jaxdem.Integrator#

Bases: Factory, ABC

Abstract base class that defines the interface for time-stepping.

Example:#

To define a custom integrator, inherit from Integrator and implement its abstract methods:

>>> @Integrator.register("myCustomIntegrator")
>>> @jax.tree_util.register_dataclass
>>> @dataclass(slots=True)
>>> class MyCustomIntegrator(Integrator):
        ...

static step_before_force(state: State, system: System) → tuple[State, System][source]#

Advance the simulation state before the force evaluation.

Parameters:

state (State) – Current state of the simulation.
system (System) – Simulation system configuration.

Returns:

The updated state and system.

Return type:

Tuple[State, System]

static step_after_force(state: State, system: System) → tuple[State, System][source]#

Advance the simulation state after the force evaluation.

Parameters:

state (State) – Current state of the simulation.
system (System) – Simulation system configuration.

Returns:

The updated state and system.

Return type:

Tuple[State, System]

static initialize(state: State, system: System) → tuple[State, System][source]#

Initialize the integrator.

Some integration methods need an initialization step, for example LeapFrog. The default implementation returns the state and system unchanged.

Parameters:

state (State) – Current state of the simulation.
system (System) – Simulation system configuration.

Returns:

The state and system after initialization.

Return type:

Tuple[State, System]

Example

>>> state, system = system.integrator.initialize(state, system)

class jaxdem.LawCombiner(laws: tuple[ForceModel, ...] = ())#

Bases: ForceModel

A ForceModel that sums a tuple of elementary force laws.

The total force, torque, and potential energy of the interaction between particles \(i\) and \(j\) are the sums over the contained laws:

\[F_{ij} = \sum_k F^{(k)}_{ij}, \qquad \tau_{ij} = \sum_k \tau^{(k)}_{ij}, \qquad E_{ij} = \sum_k E^{(k)}_{ij}\]

Notes

The combiner evaluates each sub-law with a system whose force_model is the sub-law itself. Laws that read their own configuration from jaxdem.System.force_model (including nested combiners) work correctly.
An empty combiner (laws=()) returns zero force, torque, and energy. ForceRouter.from_dict() uses it as the default no-interaction law.
required_material_properties is the union of the requirements of all contained laws.

property requires_history: bool[source]#: Whether this force model needs persistent pair history.

init_history(shape: tuple[int, ...]) → Any[source]#

static force_and_history(i: int, j: int, pos: jax.Array, state: State, system: System, history: Any) → tuple[jax.Array, jax.Array, Any][source]#

property required_material_properties: tuple[str, ...][source]#

Names of the material properties this force model needs.

The sorted union of the material properties required by all contained laws. Each name must be present in System.mat_table. Used for validation.

static force(i: int, j: int, pos: jax.Array, state: State, system: System) → tuple[jax.Array, jax.Array][source]#

Compute the total force and torque on particle \(i\) from particle \(j\) by summing all contained laws.

Parameters:

i (int) – Index of the first particle.
j (int) – Index of the second particle.
pos (jax.Array) – Particle positions used to evaluate the interaction.
state (State) – Current state of the simulation.
system (System) – Simulation system configuration.

Returns:

A tuple (force, torque) with the sums of the forces and torques of all contained laws acting on particle \(i\) from particle \(j\).

Return type:

Tuple[jax.Array, jax.Array]

static energy(i: int, j: int, pos: jax.Array, state: State, system: System) → jax.Array[source]#

Compute the total potential energy of the interaction between particle \(i\) and particle \(j\) by summing all contained laws.

Parameters:

i (int) – Index of the first particle.
j (int) – Index of the second particle.
pos (jax.Array) – Particle positions used to evaluate the interaction.
state (State) – Current state of the simulation.
system (System) – Simulation system configuration.

Returns:

Scalar total potential energy of the interaction between particles \(i\) and \(j\).

Return type:

jax.Array

class jaxdem.LinearIntegrator#

Bases: Integrator

Namespace for translation/linear-time integrators.

Purpose#

Groups integrators that update linear state (e.g., position and velocity). Concrete methods (e.g., DirectEuler) subclass this to register with the Factory and to show that they operate on linear kinematics.

class jaxdem.Material(density: float)#

Bases: Factory

Abstract base class for materials.

Concrete subclasses of Material define scalar or vector fields (e.g., young, poisson, mu) for the physical properties of a material. The MaterialTable collects and manages these fields.

Notes:#

Each field of a concrete Material subclass becomes a named property in the MaterialTable.props dictionary.

Example:#

To define a custom material, inherit from Material

>>> @Material.register("my_custom_material")
>>> @jax.tree_util.register_dataclass
>>> @dataclass(slots=True)
>>> class MyCustomMaterial(Material):
        ...

density: float#

class jaxdem.MaterialMatchmaker#

Bases: Factory, ABC

Abstract base class for material property mixing rules.

Notes:#

The jaxdem.MaterialTable uses these matchmakers to pre-compute interaction matrices.

Example:#

To define a custom matchmaker, inherit from MaterialMatchmaker and implement its abstract methods:

>>> @MaterialMatchmaker.register("myCustomForce")
>>> @jax.tree_util.register_dataclass
>>> @dataclass(slots=True)
>>> class MyCustomMatchmaker(MaterialMatchmaker):
        ...

abstractmethod static get_effective_property(prop1: Array, prop2: Array) → Array[source]#

Compute the effective property from two material properties.

Concrete implementations define the specific mixing rule.

Parameters:

prop1 (jax.Array) – The property value from the first material. Can be a scalar or an array.
prop2 (jax.Array) – The property value from the second material. Can be a scalar or an array.

Returns:

The effective property, computed from prop1 and prop2 with the matchmaker’s mixing rule.

Return type:

jax.Array

class jaxdem.MaterialTable(props: dict[str, Array], pair: dict[str, Array], matcher: MaterialMatchmaker)#

Bases: object

A container for material properties, organized as Structures of Arrays (SoA) and pre-computed effective pair properties.

The table gives direct access to the scalar properties of each material and to the pre-computed effective properties for material pairs.

Notes:#

Access scalar properties directly with dot notation (e.g., material_table.young).
Access effective pair properties directly with dot notation (e.g., material_table.young_eff).

Example:#

Creating a MaterialTable from multiple material types:

>>> import jax.numpy as jnp
>>> import jaxdem as jdem
>>>
>>> # Define different material instances
>>> mat1 = jdem.Material.create("elastic", density=2500.0, young=1.0e4, poisson=0.3)
>>> mat2 = jdem.Material.create("elasticfrict", density=7800.0, young=2.0e4, poisson=0.4, mu=0.5, e=1.0)
>>>
>>> # Create a MaterialTable using a linear matcher
>>> matcher_instance = jdem.MaterialMatchmaker.create("linear")
>>> mat_table = jdem.MaterialTable.from_materials(
>>>     [mat1, mat2],
>>>     matcher=matcher_instance
>>> )

props: dict[str, Array]#: A dictionary mapping scalar material property names (e.g., “young”, “poisson”, “mu”) to JAX arrays. Each array has shape (M,), where M is the total number of distinct material types present in the table.

pair: dict[str, Array]#: A dictionary mapping effective pair property names (e.g., “young_eff”, “mu_eff”) to JAX arrays. Each array has shape (M, M) and holds the effective property for interactions between any two material types (M_i, M_j).

matcher: MaterialMatchmaker#: The jaxdem.MaterialMatchmaker instance that computed the effective pair properties stored in the pair dictionary.

static from_materials(mats: Sequence[Material], *, matcher: MaterialMatchmaker | None = None, fill: float = 0.0) → MaterialTable[source]#

Construct a MaterialTable from a sequence of Material instances.

Parameters:

mats (Sequence[Material]) – A sequence of concrete Material instances. Each instance represents a distinct material type in the simulation. The order in this sequence defines their material IDs (0 to len(mats)-1).
matcher (MaterialMatchmaker) – The jaxdem.MaterialMatchmaker instance used to compute effective pair properties (e.g., harmonic mean, arithmetic mean). If None, defaults to the harmonic matchmaker.
fill (float, optional) – Fill value for material properties that a Material subclass does not define. For example, if an Elastic material appears with an ElasticFriction material, mu takes this value. Defaults to 0.0.

Returns:

A new MaterialTable instance containing the scalar properties and pre-computed effective pair properties for all provided materials.

Return type:

MaterialTable

Raises:

TypeError – If mats is not a sequence of Material instances.

property metadata: dict[str, Any][source]#: MaterialTable configuration parameters for serialization and restoration.

jaxdem.minimize(state: State, system: System, max_steps: int = 10000, pe_tol: float = 1e-16, pe_diff_tol: float = 1e-16, force_tol: float = 0.0) → tuple[State, System, int, float | jax.Array][source]#

Minimize the energy of the system using the configured optax optimizer.

This function runs a JAX-compatible optimization loop using the minimizer in system.minimizer. The function packs the positions and orientations into a parameter dictionary, optimizes them, and unpacks them into the returned State. The function re-anchors the rotation parameters at the current orientation each iteration (delta rotation vectors), so the torque-as-gradient identity stays exact regardless of the accumulated rotation.

The loop performs exactly one force and energy evaluation per iteration, plus one initial evaluation. It carries the value and the gradient through the loop state.

The optimization loop terminates when any of the following conditions are met:

The number of steps reaches max_steps.
The magnitude of the potential energy per particle drops below pe_tol (or of the overall objective if system.target_fn is defined): \(|E_k| \le \text{pe\_tol}\).
The relative change in potential energy between successive steps drops below pe_diff_tol (with a safe denominator, so a zero-energy state does not produce NaN):

\[\frac{|E_k - E_{k-1}|}{\max(|E_k|, |E_{k-1}|, \epsilon)} < \text{pe\_diff\_tol}\]
The maximum absolute gradient component (force/torque) drops to force_tol or below: \(\max_i |g_i| \le \text{force\_tol}\).

Parameters:

state (State) – The state of the system.
system (System) – The system to minimize.
max_steps (int, default 10000) – The maximum number of optimization steps to take.
pe_tol (float, default 1e-16) – The absolute potential energy tolerance (applied to the magnitude, so negative-energy objectives such as Lennard-Jones do not exit prematurely).
pe_diff_tol (float, default 1e-16) – The relative potential energy difference tolerance for convergence.
force_tol (float, default 0.0) – Force-norm (max absolute gradient component) tolerance. The default of 0.0 only triggers for an exactly force-free configuration.

Returns:

A tuple containing: - The energy-minimized State. - The updated System. - The number of steps actually taken. - The final potential energy.

Return type:

Tuple[State, System, int, float | jax.Array]

jaxdem.fire(dt: float, alpha_init: float = 0.1, f_inc: float = 1.1, f_dec: float = 0.5, f_alpha: float = 0.99, N_min: int = 5, N_bad_max: int = 10, dt_max_scale: float = 10.0, dt_min_scale: float = 0.001) → Any[source]#

Fast Inertial Relaxation Engine (FIRE) custom optax optimizer.

The FIRE algorithm accelerates or decelerates the dynamics based on the power of the force and the velocity. It minimizes the energy of granular particles.

Mathematical Formulation#

At each step:

Update the velocities and positions:

\[\begin{split}v_{old} &= v(t) + F(t) \cdot \frac{dt}{2} \\ P &= F(t) \cdot v_{old}\end{split}\]
Update the algorithm parameters depending on the power \(P\):
- Downhill Step (:math:`P > 0`):
  
  \[\begin{split}N_{good} &\to N_{good} + 1 \\ N_{bad} &\to 0 \\ dt &\to \begin{cases} \min(dt \cdot f_{inc}, dt_{max}) & \text{if } N_{good} > N_{min} \\ dt & \text{otherwise} \end{cases} \\ \alpha &\to \begin{cases} \alpha \cdot f_{\alpha} & \text{if } N_{good} > N_{min} \\ \alpha & \text{otherwise} \end{cases}\end{split}\]
- Uphill Step (:math:`P le 0`):
  
  \[\begin{split}N_{good} &\to 0 \\ N_{bad} &\to N_{bad} + 1 \\ dt &\to \max(dt \cdot f_{dec}, dt_{min}) \\ \alpha &\to \alpha_{init} \\ v_{old} &\to 0\end{split}\]
Perform velocity mixing:

\[\begin{split}v_{half} &= v_{old} \cdot (1 - \alpha) + \hat{F}(t) \cdot |v_{old}| \cdot \alpha \\ v(t + dt) &= v_{half} + F(t) \cdot \frac{dt}{2}\end{split}\]

param dt:

The base time step.

type dt:

float

param alpha_init:

The initial mixing coefficient.

type alpha_init:

float, default 0.1

param f_inc:

The factor by which the time step increases on downhill steps.

type f_inc:

float, default 1.1

param f_dec:

The factor by which the time step decreases on uphill steps.

type f_dec:

float, default 0.5

param f_alpha:

The decay factor for the mixing coefficient.

type f_alpha:

float, default 0.99

param N_min:

The number of consecutive downhill steps required to increase the time step.

type N_min:

int, default 5

param N_bad_max:

The maximum number of uphill steps before a reset.

type N_bad_max:

int, default 10

param dt_max_scale:

The maximum time step scale limit: \(dt_{max} = dt \cdot dt_{max\_scale}\).

type dt_max_scale:

float, default 10.0

param dt_min_scale:

The minimum time step scale limit: \(dt_{min} = dt \cdot dt_{min\_scale}\).

type dt_min_scale:

float, default 1e-3

returns:

CustomGradientTransformation – An optax gradient transformation for the FIRE algorithm.
Reference
———
Bitzek et al., Structural Relaxation Made Simple, Phys. Rev. Lett. 97, 170201 (2006)

jaxdem.damped_newtonian(dt: float, gamma: float = 0.5) → Any[source]#

Damped Newtonian dynamics custom optax optimizer.

This optimizer advances the parameters with a velocity-verlet-like scheme and a linear velocity damping term to minimize the energy of the system.

Mathematical Formulation#

At each step \(k\), the optimizer advances the parameters with:

\[\begin{split}v_{k} &= \frac{v_{half} + F(t) \cdot \frac{dt}{2}}{1 + \gamma \cdot \frac{dt}{2}} \\ v(t+dt) &= v_{k} \cdot \left(1 - \gamma \cdot \frac{dt}{2}\right) + F(t) \cdot \frac{dt}{2} \\ x(t+dt) &= x(t) + v(t+dt) \cdot dt\end{split}\]

param dt:: The time step.
type dt:: float
param gamma:: The damping coefficient.
type gamma:: float, default 0.5
returns:: An optax gradient transformation for the damped Newtonian algorithm.
rtype:: CustomGradientTransformation

class jaxdem.RotationIntegrator#

Bases: Integrator

Namespace for rotation/angular-time integrators.

Purpose#

Groups integrators that update angular state (e.g., orientation, angular velocity). Concrete methods (e.g., DirectEulerRotation) subclass this to register with the Factory and to show that they operate on rotational kinematics.

final class jaxdem.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

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

Notes:#

State supports these 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 the batch dimension (`B`). This is what batch_size returns.
- All leading dimensions before it (T_1, T_2, … T_k) are trajectory dimensions.
  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. Particles with the same clump_id belong to the same rigid body. Shape is (…, N). IDs must be between 0 and N - 1.

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. Connected particles do not interact. Shape is (…, N, max_num_neighbors). Empty slots contain -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]#: Return 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 fills in default values and makes the array shapes of all state attributes consistent.

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, the method creates an empty state. 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, the method creates an empty state 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. The method ignores mass when you provide mat_table. 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. Pass a nested list (possibly with uneven lengths), or a 2D array. The method symmetrizes the connections and pads them 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).
facet_vertices (jax.typing.ArrayLike or None, optional) – 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 you provide mat_table, the method ignores the mass argument and computes particle masses 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]#

Merge one or more State instances into a single new State.

This method concatenates the particles from the provided state(s) onto state1. The method shifts clump_ids, bond_ids, and facet IDs to keep them unique 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 an input state is invalid, or if the spatial dimension (dim) or batch size (batch_size) does not match between states.
ValueError – If the merged state is not valid.

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]

Add new particles to an existing State instance and return a new State.

Parameters:

state (State) – The existing State to which the method adds particles.
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. The method ignores mass when you provide mat_table.
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, the method generates new IDs.
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. Pass a nested list (possibly with uneven lengths), or a 2D array. The method symmetrizes the connections and pads them 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 you provide mat_table, the method computes masses 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]#

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

Use this method to collect simulation snapshots over time into a single State object where the leading dimension represents time, or to prepare 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

The method does not shift clump_ids, 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 method is the inverse of State.stack():

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

Notes

The method splits along axis 0 (the leading axis).
This method cannot split a single snapshot State (e.g. pos.shape == (N, dim)), because axis 0 would refer to particles, not snapshots.

Add a new clump of multiple spheres to an existing State. All spheres in the new clump share the rigid body properties (center of mass position pos_c, velocity, mass, orientation q, angular velocity, force, torque, inertia, fixed, and clump_id). Only pos_p (offsets in the body reference frame), rad, and the ID fields (mat_id, species_id, and bond_id) can vary within a rigid clump.

Parameters:

state (State) – The existing State to which the method adds particles.
pos (jax.typing.ArrayLike) – If pos_p is None, pos gives the absolute coordinates of the spheres in the clump. If pos_p is not None, pos gives 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, the method computes it from the absolute coordinates in pos and computes the center of mass with 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. Pass a nested list (possibly with uneven lengths), or a 2D array. The method symmetrizes the connections and pads them 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

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

Note: Facets that this method adds do not share vertices. Each facet has its own copy of the 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: all vertices share one clump ID, and the method computes the clump inertia and orientation. If False, the facet is flexible: its vertices behave like individual spheres (sphere moment of inertia, identity orientation, and unique clump IDs).
safety_factor (float, default 1.0) – Factor that scales _rad to enlarge the broad-phase detection box.

Add a new mesh (a collection of facets) 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: all vertices share one clump ID, and the method computes the clump inertia and orientation. If False, the mesh is flexible.
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 that scales _rad to enlarge the broad-phase detection box.

Add a new facet that connects existing vertices, new vertices, or both, to the State.

Parameters:

state (State) – The existing state.
vertex_specs (list of int or ArrayLike) – Each spec is one vertex of the new facet. A scalar integer (or scalar array) spec is the index of an existing vertex. Any other spec is a position array of shape (dim,) for a new vertex.

final class jaxdem.System(linear_integrator: LinearIntegrator, rotation_integrator: RotationIntegrator, collider: Collider, domain: Domain, force_manager: ForceManager, bonded_force_model: BondedForceModel | None, force_model: ForceModel, mat_table: MaterialTable, dt: jax.Array, time: jax.Array, dim: jax.Array, step_count: jax.Array, key: jax.Array, interact_same_bond_id: jax.Array, user_pre_step_actions: Callable[[State, System], tuple[State, System]] = <PjitFunction of <function _save_state_system>>, user_post_step_actions: Callable[[State, System], tuple[State, System]] = <PjitFunction of <function _save_state_system>>, minimizer: Any = None, target_fn: Callable[[State, System], jax.Array] | None = None)#

Bases: object

The full simulation configuration.

Notes:#

The System object supports JIT compilation for efficient execution.
The System dataclass is compatible with jax.jit(), so every field should remain JAX arrays for best performance.

Example:#

Creating a basic 2D simulation system:

>>> import jaxdem as jdem
>>> import jax.numpy as jnp
>>>
>>> # Create a System instance
>>> sim_system = jdem.System.create(
>>>     state_shape=state.shape,
>>>     dt=0.001,
>>>     linear_integrator_type="euler",
>>>     rotation_integrator_type="spiral",
>>>     collider_type="naive",
>>>     domain_type="free",
>>>     force_model_type="spring",
>>>     # You can pass keyword arguments to component constructors via '_kw' dicts
>>>     domain_kw=dict(box_size=jnp.array([5.0, 5.0]), anchor=jnp.array([0.0, 0.0]))
>>> )
>>>
>>> print(f"System integrator: {sim_system.linear_integrator.__class__.__name__}")
>>> print(f"System force model: {sim_system.force_model.__class__.__name__}")
>>> print(f"Domain box size: {sim_system.domain.box_size}")

linear_integrator: LinearIntegrator#: Instance of jaxdem.LinearIntegrator that advances the simulation linear state in time.

rotation_integrator: RotationIntegrator#: Instance of jaxdem.RotationIntegrator that advances the simulation angular state in time.

collider: Collider#: Instance of jaxdem.Collider that performs contact detection and computes inter-particle forces and potential energies.

domain: Domain#: Instance of jaxdem.Domain that defines the simulation boundaries, displacement rules, and boundary conditions.

force_manager: ForceManager#: Instance of jaxdem.ForceManager that handles per particle forces like external forces and resets forces.

bonded_force_model: BondedForceModel | None#: Optional instance of jaxdem.BondedForceModel that defines bonded interactions by passing a force and energy function to the ForceManager.

force_model: ForceModel#: Instance of jaxdem.ForceModel that defines the physical laws for inter-particle interactions.

mat_table: MaterialTable#: Instance of jaxdem.MaterialTable holding material properties and pairwise interaction parameters.

dt: jax.Array#: The global simulation time step \(\Delta t\).

time: jax.Array#: Elapsed simulation time.

dim: jax.Array#: Spatial dimension of the system.

step_count: jax.Array#: Number of integration steps that have been performed.

key: jax.Array#: PRNG key for stochastic operations. Always update it with split so each use gets new random numbers.

interact_same_bond_id: jax.Array#

Boolean scalar controlling interactions between particles with the same bond_id.

If False (default), colliders mask out these pairs. If True, these pairs interact.

user_pre_step_actions(system: System) → tuple[State, System][source]#: Function called before every step to perform user-defined actions.

user_post_step_actions(system: System) → tuple[State, System][source]#: Function called after every step to perform user-defined actions.

minimizer: Any = None#: An optax GradientTransformation wrapped in CustomGradientTransformation, used for target_fn minimization.

target_fn: Callable[[State, System], jax.Array] | None = None#: Optional custom target evaluation function for minimization.

static create(state_shape: tuple[int, ...] | None = None, *, state: State | None = None, dt: float = 0.005, time: float = 0.0, linear_integrator_type: str | None = 'verlet', rotation_integrator_type: str | None = 'verletspiral', collider_type: str = 'naive', domain_type: str = 'free', bonded_force_model_type: str | None = None, bonded_force_model_kw: dict[str, Any] | None = None, bonded_force_manager_kw: dict[str, Any] | None = None, bonded_force_model: BondedForceModel | None = None, force_model_type: str = 'spring', force_manager_kw: dict[str, Any] | None = None, mat_table: MaterialTable | None = None, linear_integrator: LinearIntegrator | None = None, rotation_integrator: RotationIntegrator | None = None, collider: Collider | None = None, domain: Domain | None = None, force_model: ForceModel | None = None, force_manager: ForceManager | None = None, linear_integrator_kw: dict[str, Any] | None = None, rotation_integrator_kw: dict[str, Any] | None = None, collider_kw: dict[str, Any] | None = None, domain_kw: dict[str, Any] | None = None, force_model_kw: dict[str, Any] | None = None, seed: int = 0, key: jax.Array | None = None, interact_same_bond_id: bool = False, user_pre_step_actions: Callable[[State, System], tuple[State, System]] | None = None, user_post_step_actions: Callable[[State, System], tuple[State, System]] | None = None, minimizer: Any = None, minimizer_kw: dict[str, Any] | None = None, target_fn: Callable[[State, System], jax.Array] | None = None) → System[source]#

Factory method to create a System instance with specified components.

Every component slot accepts either a pre-built instance (linear_integrator, rotation_integrator, collider, domain, force_model, force_manager, bonded_force_model, mat_table) or a registered type string plus keyword dict (<component>_type / <component>_kw). When you provide an instance, the method uses it as-is and ignores the corresponding *_type / *_kw arguments.

Parameters:

state_shape (Tuple, optional) – Shape of the state tensors handled by the simulation. The penultimate dimension corresponds to the number of particles N and the last dimension corresponds to the spatial dimension dim. You can omit it when you provide state.
state (State, optional) – The initial simulation state. When provided, the method infers state_shape from it and forwards the state to colliders whose Create method requires one (e.g. "CellList", "NeighborList"), so collider_kw={"state": state} is not needed.
dt (float, optional) – The global simulation time step.
linear_integrator_type (str or None, optional) – The registered type string for the jaxdem.integrators.LinearIntegrator used to evolve translational degrees of freedom. None (or the empty string) disables linear integration (no-op integrator).
rotation_integrator_type (str or None, optional) – The registered type string for the jaxdem.integrators.RotationIntegrator used to evolve angular degrees of freedom. None (or the empty string) disables rotational integration (no-op integrator).
collider_type (str, optional) – The registered type string for the jaxdem.Collider to use.
domain_type (str, optional) – The registered type string for the jaxdem.Domain to use.
bonded_force_model_type (str or None, optional) – The registered type string for the jaxdem.BondedForceModel to use.
bonded_force_model_kw (Dict[str, Any] or None, optional) – Keyword arguments forwarded to BondedForceModel.create.
bonded_force_manager_kw (Dict[str, Any] or None, optional) – Deprecated alias of bonded_force_model_kw (the dict has always been forwarded to the bonded force model, not the manager).
force_model_type (str, optional) – The registered type string for the jaxdem.ForceModel to use.
force_manager_kw (Dict[str, Any] or None, optional) – Keyword arguments to pass to the constructor of ForceManager.
mat_table (MaterialTable or None, optional) – An optional pre-configured jaxdem.MaterialTable. If None, the method creates a default jaxdem.MaterialTable with one generic elastic material and the “harmonic” jaxdem.MaterialMatchmaker.
linear_integrator (LinearIntegrator, optional) – Pre-built linear integrator instance. Overrides linear_integrator_type / linear_integrator_kw.
rotation_integrator (RotationIntegrator, optional) – Pre-built rotation integrator instance. Overrides rotation_integrator_type / rotation_integrator_kw.
collider (Collider, optional) – Pre-built collider instance. Overrides collider_type / collider_kw.
domain (Domain, optional) – Pre-built domain instance. Overrides domain_type / domain_kw.
force_model (ForceModel, optional) – Pre-built force model instance. Overrides force_model_type / force_model_kw.
force_manager (ForceManager, optional) – Pre-built force manager instance. Overrides force_manager_kw. Cannot be combined with a bonded force model (the bonded force functions must already be part of the provided manager).
linear_integrator_kw (Dict[str, Any] or None, optional) – Keyword arguments forwarded to the constructor of the selected LinearIntegrator type.
rotation_integrator_kw (Dict[str, Any] or None, optional) – Keyword arguments forwarded to the constructor of the selected RotationIntegrator type.
collider_kw (Dict[str, Any] or None, optional) – Keyword arguments to pass to the constructor of the selected Collider type.
domain_kw (Dict[str, Any] or None, optional) – Keyword arguments to pass to the constructor of the selected Domain type.
force_model_kw (Dict[str, Any] or None, optional) – Keyword arguments to pass to the constructor of the selected ForceModel type.
seed (int, optional) – Integer seed used for random number generation. Defaults to 0. Used only when key is not provided.
key (jax.Array, optional) – Key for JAX random number generation. When you provide key, the method ignores seed.
interact_same_bond_id (bool, optional) – Whether particles with the same bond_id interact. Defaults to False.
user_pre_step_actions (Callable, optional) – A function called before every time step to perform user-defined actions.
user_post_step_actions (Callable, optional) – A function called after every time step to perform user-defined actions.
minimizer (Callable, optional) – Optimizer factory used by System.minimize(). Called as minimizer(**minimizer_kw) and must return an optax-style GradientTransformation. Defaults to FIRE (jaxdem.minimizers.fire()).
minimizer_kw (Dict[str, Any] or None, optional) – Keyword arguments passed to minimizer. If the minimizer’s signature accepts a dt parameter and none is given here, the method passes the system dt automatically.
target_fn (Callable, optional) – Custom objective (state, system) -> scalar minimized by System.minimize(). When None, System.minimize() uses the total potential energy.

Returns:

A fully configured System instance ready for simulation.

Return type:

System

Raises:

KeyError – If a specified *_type is not registered in its respective factory, or if the mat_table is missing properties required by the force_model.
TypeError – If constructor keyword arguments are invalid for any component.
ValueError – If the domain_kw ‘box_size’ or ‘anchor’ shapes do not match the dim.

Example

Creating a 3D system with reflective boundaries and a custom dt:

>>> import jaxdem as jdem
>>> import jax.numpy as jnp
>>>
>>> system_reflect = jdem.System.create(
>>>     state_shape=(N, 3),
>>>     dt=0.0005,
>>>     domain_type="reflect",
>>>     domain_kw=dict(box_size=jnp.array([20.0, 20.0, 20.0]), anchor=jnp.array([-10.0, -10.0, -10.0])),
>>>     force_model_type="spring",
>>> )
>>> print(f"System dt: {system_reflect.dt}")
>>> print(f"Domain type: {system_reflect.domain.__class__.__name__}")

Creating a system with a pre-defined MaterialTable:

>>> custom_mat_kw = dict(young=2.0e5, poisson=0.25)
>>> custom_material = jdem.Material.create("custom_mat", **custom_mat_kw)
>>> custom_mat_table = jdem.MaterialTable.from_materials(
...     [custom_material], matcher=jdem.MaterialMatchmaker.create("linear")
... )
>>>
>>> system_custom_mat = jdem.System.create(
...     state_shape=(N, 2),
...     mat_table=custom_mat_table,
...     force_model_type="spring"
... )

static trajectory_rollout(state: State, system: System, *, n: int | None = None, stride: int = 1, strides: jax.Array | None = None, save_fn: Callable[[State, System], Any] = <PjitFunction of <function _save_state_system>>, unroll: int = 2) → tuple[State, System, Any][source]#

Roll the system forward while collecting saved outputs at each frame.

The rollout always stores one output per frame via save_fn(state, system). The output of save_fn must be a pytree. Frame spacing can be either: - constant (stride), or - variable (strides jax.Array).

The rollout saves each frame after its integration steps, so it does not store the initial (step-0) state. To record it, save it before the rollout, or pass a leading 0 entry in strides.

Parameters:

state (State) – Initial state.
system (System) – Initial system configuration.
n (int, optional) – Number of saved frames. Required when strides is None. The method ignores n when you provide strides.
stride (int, optional) – Constant number of integration steps between consecutive saves. Used only when strides is None. Defaults to 1.
strides (jax.Array, optional) – Integer 1D array of per-frame integration strides. When provided, this overrides stride, and the method infers n from len(strides).
save_fn (Callable[[State, System], Any], optional) – Function called after each saved frame. The rollout stacks its return pytree along axis 0 across frames. Defaults to returning (state, system).
unroll (int, optional) – Unroll factor passed to the outer jax.lax.scan. Defaults to 2.

Returns:

(final_state, final_system, trajectory_like) where trajectory_like is the stacked output of save_fn.

Return type:

Tuple[State, System, Any]

Raises:

ValueError – If n is missing while strides is None, or if strides is not 1D.

Example

>>> import jaxdem as jdem
>>> import jax.numpy as jnp
>>>
>>> state = jdem.utils.grid_state(n_per_axis=(1, 1), spacing=1.0, radius=0.1)
>>> system = jdem.System.create(state_shape=state.shape, dt=0.01)
>>>
>>> # Constant stride: n is required
>>> final_state, final_system, traj = jdem.System.trajectory_rollout(
...     state, system, n=10, stride=5
... )
>>>
>>> # Variable strides: n inferred from len(strides)
>>> deltas = jnp.array([1, 2, 4, 8])
>>> final_state, final_system, traj = jdem.System.trajectory_rollout(
...     state, system, strides=deltas
... )

static step(state: State, system: System, *, n: int | jax.Array = 1) → tuple[State, System][source]#

Advance the simulation by n integration steps.

Parameters:

state (State) – Current state.
system (System) – Current system configuration.
n (int or jax.Array, optional) – Number of integration steps. May be a Python int or a scalar JAX array. Defaults to 1.

Returns:

(final_state, final_system) after n steps.

Return type:

Tuple[State, System]

Example

>>> # Advance by 10 steps
>>> state_after_10_steps, system_after_10_steps = jdem.System.step(state, system, n=10)

Notes

This method does not check collider overflow, to avoid a host synchronization per step.

static stack(systems: Sequence[System]) → System[source]#

Concatenate a sequence of System snapshots into a trajectory or batch along axis 0.

Use this method to collect simulation snapshots over time into a single System object where the leading dimension represents time, or to prepare a batched system.

Parameters:: systems (Sequence[System]) – A sequence (e.g., list, tuple) of System instances to be stacked.
Returns:: A new System 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:: System

static unstack(system: System) → list[System][source]#

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

This method is the inverse of System.stack():

If stacked = System.stack([sys0, sys1, …]), then System.unstack(stacked) returns [sys0, sys1, …].

Notes

The method splits along axis 0 (the leading axis).
This method cannot split a single snapshot System.

static minimize(state: State, system: System, *, max_steps: int = 10000, pe_tol: float = 1e-16, pe_diff_tol: float = 1e-16) → tuple[State, System, int, float][source]#

Minimize the energy of the system using the configured minimizer.

Parameters:

state (State) – The state of the simulation.
system (System) – The system configuration.
max_steps (int, optional) – The maximum number of steps to take. Defaults to 10000.
pe_tol (float, optional) – The tolerance for the potential energy. Defaults to 1e-16.
pe_diff_tol (float, optional) – The tolerance for the difference in potential energy. Defaults to 1e-16.

Returns:

The final state, system, number of steps, and potential energy (per particle when no custom target_fn is set).

Return type:

Tuple[State, System, int, float]

Notes

The loop stops as soon as any convergence criterion is met (energy tolerance, relative energy change, or force tolerance) — see jaxdem.minimizers.minimize() for the full list.

property metadata: dict[str, Any][source]#: System configuration parameters for serialization and restoration.

class jaxdem.VTKBaseWriter#

Bases: Factory, ABC

Abstract base class for writers that output simulation data.

Concrete subclasses implement the write method to convert a snapshot (jaxdem.State, jaxdem.System pair) into a specific file format.

Example:#

To define a custom VTK writer, inherit from VTKBaseWriter and implement its abstract methods:

>>> @VTKBaseWriter.register("my_custom_vtk_writer")
>>> @dataclass(slots=True)
>>> class MyCustomVTKWriter(VTKBaseWriter):
        ...

classmethod is_active(state: State, system: System) → bool[source]#: Check whether this writer has data to write for the given state and system.

abstractmethod classmethod write(state: State, system: System, filename: Path, binary: bool) → None[source]#

Write information from a simulation snapshot to a VTK PolyData file.

Concrete writers implement this method. The caller converts all JAX arrays to NumPy arrays before it calls write.

Parameters:

state (State) – The simulation jaxdem.State snapshot to write.
system (System) – The simulation jaxdem.System configuration.
filename (Path) – Target path of the VTK file. The caller guarantees that the parent directory exists.
binary (bool) – If True, write the VTK file in binary mode. If False, write it in ASCII (human-readable) mode.

class jaxdem.VTKWriter(directory: Path = PosixPath('frames'), save_every: int = 1, clean: bool = True, max_workers: int = 8, max_queue_size: int = 512, writers: list[str] = <factory>, binary: bool = True)#

Bases: BaseAsyncWriter

High-level front end for writing simulation data to VTK files.

This class converts JAX-based jaxdem.State and jaxdem.System pytrees into VTK files. It handles batches, trajectories, and dispatch to registered jaxdem.VTKBaseWriter subclasses.

How leading axes are interpreted#

Let particle positions have shape (..., N, dim), where N is the number of particles and dim is 2 or 3. Define L = state.pos_c.ndim - 2, i.e., the number of leading axes before (N, dim).

L == 0 — single snapshot
The input is one frame. It is written directly into frames/batch_00000000/ (no batching, no trajectory).
trajectory=False (default)
The writer treats all leading axes as batch axes (not time). If multiple batch axes exist, the writer flattens them into a single batch axis: (B, N, dim) with B = prod(shape[:L]). The writer writes each batch b as a single snapshot under its own subdirectory frames/batch_XXXXXXXX/. No trajectory is implied.
- Example: (B, N, dim) → B separate directories with one frame each.
- Example: (B1, B2, N, dim) → flatten to (B1*B2, N, dim) and treat as above.
trajectory=True
The writer swaps the axis given by trajectory_axis to the front (axis 0) and treats it as time T. Any remaining leading axes are batch axes. If more than one non-time leading axis exists, the writer flattens them into a single batch axis. The data becomes (T, B, N, dim) with B = prod(other leading axes).
- If there is only time (L == 1): (T, N, dim) — a single batch
  directory frames/batch_00000000/ contains a time series with T frames.
- If there is time plus batching (L >= 2): (T, B, N, dim) — each
  batch b gets its own directory frames/batch_XXXXXXXX/ containing a time series (T frames) for that batch.

After these swaps/reshapes, dispatch is: - (N, dim) → single snapshot - (B, N, dim) → batches (no time) - (T, N, dim) → single batch with a trajectory - (T, B, N, dim) → per-batch trajectories

Concrete writers receive per-frame NumPy arrays. The writer slices and broadcasts System leaves to match the current frame and batch.

writers: list[str]#: Names of the registered VTKBaseWriter subclasses to use for writing. If empty, use all registered subclasses. Name matching follows registry keys: case-insensitive, and spaces, underscores, and hyphens are ignored. The spelling given here sets the output file and .pvd names.

binary: bool = True#: If True, write VTK files in binary format. If False, write files in ASCII format.

save(state: State, system: System, *, trajectory: bool = False, trajectory_axis: int = 0, batch0: int = 0) → None[source]#

Schedule writing of a jaxdem.State / jaxdem.System pair to VTK files.

This public entry point interprets the leading axes as batch or trajectory axes, swaps and flattens axes as needed, and pushes the data to the background writer queue.

Parameters:

state (State) – The simulation jaxdem.State object to save.
system (System) – The jaxdem.System object corresponding to state.
trajectory (bool, optional) – If True, interpret trajectory_axis as time.
trajectory_axis (int, optional) – The axis in state/system to treat as the trajectory axis.
batch0 (int, optional) – The starting batch index for the input data.

directory: Path#: The root directory where the writer saves simulation frames.

save_every: int#: Save frequency. The writer pushes a frame to the queue on the first call and on every save_every-th call to the save() method.

clean: bool#: If True, the writer deletes and recreates directory on initialization. Safety checks prevent deleting the current working directory or the system root.

max_workers: int#: The number of background worker threads to use for parallel I/O.

max_queue_size: int#: Maximum number of pending tasks in the background queue. When the queue is full, submit() blocks until a worker frees a slot. This backpressure keeps memory bounded when the simulation outruns disk I/O. Set to 0 for an unbounded queue.

Modules

`analysis`	Post-processing and analysis utilities.
`bonded_forces`	Bonded-force interfaces independent of the collider.
`colliders`	Collision-detection interfaces and implementations.
`domains`	Simulation domains and boundary-condition implementations.
`factory`	The factory registers and creates simulation components.
`forces`	Force-law interfaces.
`integrators`	Time-integration interfaces and implementations.
`material_matchmakers`	Material mix rules and implementations.
`materials`	Interface for defining materials and the MaterialTable.
`minimizers`	Energy-minimizer interfaces and implementations.
`rl`	JaxDEM reinforcement learning (RL) module.
`state`	Defines the simulation State.
`system`	The simulation configuration and the tools that drive the simulation.
`utils`	Utility functions used to set up simulations and analyze the output.
`writers`	Interface for defining data writers.