jaxdem.system#

The simulation configuration and the tools that drive the simulation.

Classes

System(linear_integrator, ...)

The full simulation configuration.

final class jaxdem.system.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, search_overflow: jax.Array = <factory>, 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.
Dynamic numeric leaves are JAX arrays. Callbacks, optimizers, targets, and component configuration can carry explicit static metadata.

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 connected by bond_id adjacency.

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

search_overflow: jax.Array#

a force evaluation used an incomplete spatial search.

Check with check_overflow at host chunk boundaries. Retry from the last valid state after resizing; clearing this flag cannot repair a trajectory.

Type:: Sticky flag

user_pre_step_actions: Callable[[State, System], tuple[State, System]]#

Function called before every step to perform user-defined actions.

The callback sees the tentative end time using the entering dt. The recorded end time uses the callback’s dt; any callback change to time is retained as an offset from the tentative end time.

user_post_step_actions: Callable[[State, System], tuple[State, System]]#: Function called after every step to perform user-defined actions.

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

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

static create(state_shape: tuple[int, ...] | None = None, *, state: State | None = None, dt: ArrayLike = 0.005, time: float = 0.0, linear_integrator_type: str | None = 'verlet', rotation_integrator_type: str | None = 'verletspiral', collider_type: str | None = '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 (ArrayLike, optional) – The scalar global simulation time step, including a traced JAX scalar.
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 or None, optional) – The registered type string for the jaxdem.Collider to use. None (or the empty string) disables non-bonded particle interactions by selecting the no-op collider. Bonded forces, external forces, and domain boundaries remain active.
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 initialize(state: State, system: System) → tuple[State, System][source]#

Compute initial forces and initialize both configured integrators.

Call explicitly after constructing the state and system, before the first integration step. This method supports one snapshot or matching stacked states and systems, like step.

Bounds and collider caches are updated, pair and managed forces are evaluated, then the linear and rotational initialize hooks run. An integrator can use its hook to stagger velocities, for example by a backward half-kick. The default integrator hook leaves state unchanged.

Initialization does not advance time, step count, contact/plastic history, or timestep callbacks. Managed force buffers follow their ordinary apply-and-clear semantics. No initialized flag is stored: callers must initialize new trajectories themselves and must not repeat integrator setup when continuing an initialized checkpoint.

Example

>>> state, system = System.initialize(state, system)
>>> state, system = System.step(state, system, n=100)

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

Inspect instantaneous forces without advancing physical state.

Return updated forces and search caches while preserving contact and plastic history, queued loads, time, and integrator state. Queued loads contribute to the returned forces but remain available for the next physical step. This is an observation, not integrator initialization. Custom force callbacks must themselves be pure evaluations.

check_overflow() → None[source]#

Raise at a host boundary if this trajectory used truncated searches.

This synchronizes the status flag, so call once per completed chunk, outside compiled loops. A successful later query does not clear it.

validate(state: State, *, strict_clumps: bool = True) → None[source]#: Run opt-in host validation for a state against this system.

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 = 1) → tuple[State, System][source]#

Advance the simulation by n integration steps.

Parameters:

state (State) – Current state.
system (System) – Current system configuration.
n (int, optional) – Static number of integration steps. Defaults to 1. Use step_dynamic for a traced scalar count.

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

Call initialize explicitly before stepping a new trajectory. This method performs no automatic force or integrator initialization. Checkpoint continuation uses its saved forces and integrator state without repeating initialization.

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

static step_checked(state: State, system: System, *, n: int = 1) → StepResult[source]#

Advance at most n steps, stopping and rolling back on failure.

Unlike step, this opt-in device loop checks finite particle state and spatial-search overflow after every step. result.steps counts accepted steps; result.status preserves failure even though the returned state/system precede the failed step. An invalid input is returned unchanged with zero accepted steps. Check result.check() on the host before continuing or writing successful output.

Initialization and physical setup validation remain explicit. A batch returns one status/count per snapshot. This dynamic stopping loop does not support reverse-mode differentiation. No automatic retry changes the timestep, neighbor capacity, or physical history.

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

Advance by a scalar integer count that may be traced at run time.

This path uses lax.fori_loop with dynamic bounds. It is intended for compiled control flow such as variable trajectory strides. Reverse mode differentiation through dynamic loop bounds is not supported by JAX; use step with a Python integer for differentiable rollouts.

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, force_tol: float = 1e-12, torque_tol: float | None = None, advance_history: bool = False) → MinimizationResult[source]#

Relax until both free-body residual norms meet tolerance.

Returns an iterable four-value result with reason and info diagnostics. Energy is reported per constituent sphere for the physical objective. Pair history is frozen unless advance_history=True requests displacement-driven updates after accepted optimizer steps. Residuals are mechanical forces and torques for force-based relaxation, or objective gradients for energies and targets differentiated by AD. See jaxdem.minimizers.minimize for convergence and optimizer behavior.

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