jaxdem.integrators

Time-integration interfaces and implementations.

Classes

Integrator()	Abstract base class for defining the interface for time-stepping.
LinearIntegrator()	Namespace for translation/linear-time integrators.
RotationIntegrator()	Namespace for rotation/angular-time integrators.

class jaxdem.integrators.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) should subclass this to register via the Factory and to signal they operate on linear kinematics.

class jaxdem.integrators.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) should subclass this to register via the Factory and to signal they operate on rotational kinematics.

class jaxdem.integrators.DirectEuler

Bases: LinearIntegrator

Implements the explicit (forward) Euler integration method.

static step_after_force(state: State, system: System) → Tuple[State, System]

Advances the simulation state by one time step after the force calculation using the Direct Euler method.

The update equations are:

\[\begin{split}& v(t + \Delta t) &= v(t) + \Delta t a(t) \\ & r(t + \Delta t) &= r(t) + \Delta t v(t + \Delta t)\end{split}\]

where:

• \(r\) is the particle position (jaxdem.State.pos)

• \(v\) is the particle velocity (jaxdem.State.vel)

• \(a\) is the particle acceleration computed from forces (jaxdem.State.force)

• \(\Delta t\) is the time step (jaxdem.System.dt)

Parameters:

• state (State) – Current state of the simulation.

• system (System) – Simulation system configuration.

Returns:

The updated state and system after one time step.

Return type:

Tuple[State, System]

Note

• This method donates state and system

class jaxdem.integrators.Spiral

Bases: RotationIntegrator

Non-leapfrog spiral integrator for angular velocities.

The implementation follows the velocity update described in del Valle et al. (2023).

static step_after_force(state: State, system: System) → Tuple[State, System]

Advance angular velocities by a single time step.

A third-order Runge–Kutta scheme (SSPRK3) integrates the rigid-body angular momentum equations in the principal axis frame. The quaternion is updated based on the spiral non-leapfrog algorithm.

• SPIRAL algorithm:

\[q(t + \Delta t) = q(t) \cdot e^\left(\frac{\Delta t}{2}\omega\right) \cdot e^\left(\frac{\Delta t^2}{4}\dot{\omega}\right)\]

Where the angular velocity and its derivative are purely imaginary quaternions (scalar part is zero and the vector part is equal to the vector). The exponential map of a purely imaginary quaternion is

\[e^u = \cos(|u|) + \frac{\vec{u}}{|u|}\sin(|u|)\]

Angular velocity is then updated using SSPRK3:

\[\begin{split}& \vec{\omega}(t + \Delta t) = \vec{\omega}(t) + \frac{1}{6}(k_1 + k_2 + 4k_3) \\ & k_1 = \Delta t\; \dot{\vec{\omega}}(\vec{\omega}(t + \Delta t / 2), \vec{\tau}(t + \Delta t)) \\ & k_2 = \Delta t\; \dot{\vec{\omega}}(\vec{\omega}(t + \Delta t / 2) + k1, \vec{\tau}(t + \Delta t)) \\ & k_3 = \Delta t\; \dot{\vec{\omega}}(\vec{\omega}(t + \Delta t / 2) + (k1 + k2)/4, \vec{\tau}(t + \Delta t)) \\\end{split}\]

Where the angular velocity derivative is a function of the torque and angular velocity:

\[\dot{\vec{\omega}} = (\tau - \vec{\omega} \times (I \vec{\omega}))I^{-1}\]

Parameters:

• state (State) – Current state of the simulation.

• system (System) – Simulation system configuration.

Returns:

The updated state and system after one time step.

Return type:

Tuple[State, System]

Reference

del Valle et. al, SPIRAL: An efficient algorithm for the integration of the equation of rotational motion, https://doi.org/10.1016/j.cpc.2023.109077.

Note

• This method donates state and system

class jaxdem.integrators.VelocityVerlet

Bases: LinearIntegrator

Implements the Velocity Verlet integration method.

static step_before_force(state: State, system: System) → Tuple[State, System]

Advances the simulation state by one half-step before the force calculation using the Velocity Verlet scheme.

The update equations are:

\[\begin{split}& v(t + \Delta t / 2) &= v(t) + \Delta t a(t) / 2 \\ & r(t + \Delta t) &= r(t) + \Delta t v(t + \Delta t / 2) \Delta t\end{split}\]

where:

• \(r\) is the particle position (jaxdem.State.pos)

• \(v\) is the particle velocity (jaxdem.State.vel)

• \(a\) is the particle acceleration computed from forces (jaxdem.State.force)

• \(\Delta t\) is the time step (jaxdem.System.dt)

Parameters:

• state (State) – Current state of the simulation.

• system (System) – Simulation system configuration.

Returns:

The updated state and system after one time step.

Return type:

Tuple[State, System]

Note

• This method donates state and system

static step_after_force(state: State, system: System) → Tuple[State, System]

Advances the simulation state by one half-step after the force calculation using the Velocity Verlet scheme.

The update equations are:

\[\begin{split}& v(t + \Delta t) &= v(t + \Delta t / 2) + \Delta t a(t) / 2 \\\end{split}\]

where:

• \(v\) is the particle velocity (jaxdem.State.vel)

• \(a\) is the particle acceleration computed from forces (jaxdem.State.force)

• \(\Delta t\) is the time step (jaxdem.System.dt)

Parameters:

• state (State) – Current state of the simulation.

• system (System) – Simulation system configuration.

Returns:

The updated state and system after one time step.

Return type:

Tuple[State, System]

Note

• This method donates state and system

class jaxdem.integrators.VelocityVerletSpiral

Bases: RotationIntegrator

Leapfrog spiral integrator for angular velocities adapted to Velocity Verlet.

The implementation follows the velocity update described in del Valle et al. (2023).

static step_before_force(state: State, system: System) → Tuple[State, System]

Advances the simulation state by one half-step before the force calculation using the Velocity Verlet scheme.

A third-order Runge–Kutta scheme (SSPRK3) integrates the rigid-body angular momentum equations in the principal axis frame. The quaternion is updated with the spiral leapfrog algorithm to implement a Velocity Verlet-like method.

• SPIRAL algorithm:

\[q(t + \Delta t) = q(t) \cdot e^\left(\frac{\Delta t}{2}\omega(t + \Delta t/2)\right)\]

Where the angular velocity and its derivative are purely imaginary quaternions (scalar part is zero and the vector part is equal to the vector). The exponential map of a purely imaginary quaternion is

\[e^u = \cos(|u|) + \frac{\vec{u}}{|u|}\sin(|u|)\]

Angular velocity is then updated using SSPRK3 which we:

\[\begin{split}& \vec{\omega}(t + \Delta t/2) = \vec{\omega}(t) + \frac{1}{6}(k_1 + k_2 + 4k_3) \\ & k_1 = \Delta t/2\; \dot{\vec{\omega}}(\vec{\omega}(t), \vec{\tau}(t)) \\ & k_2 = \Delta t/2\; \dot{\vec{\omega}}(\vec{\omega}(t) + k1, \vec{\tau}(t)) \\ & k_3 = \Delta t/2\; \dot{\vec{\omega}}(\vec{\omega}(t) + (k1 + k2)/4, \vec{\tau}(t)) \\\end{split}\]

Where the angular velocity derivative is a function of the torque and angular velocity:

\[\dot{\vec{\omega}} = (\tau - \vec{\omega} \times (I \vec{\omega}))I^{-1}\]

Parameters:

• state (State) – Current state of the simulation.

• system (System) – Simulation system configuration.

Returns:

• Tuple[State, System] – The updated state and system after one time step.

• Reference

• ———–

• del Valle et. al, SPIRAL (An efficient algorithm for the integration of the equation of rotational motion, https://doi.org/10.1016/j.cpc.2023.109077.)

Note

• This method donates state and system.

static step_after_force(state: State, system: System) → Tuple[State, System]

Advances the simulation state by one half-step after the force calculation using the Velocity Verlet scheme.

A third-order Runge–Kutta scheme (SSPRK3) integrates the rigid-body angular momentum equations in the principal axis frame. The quaternion is updated with the spiral leapfrog algorithm to implement a Velocity Verlet-like method.

\[\begin{split}& \vec{\omega}(t + \Delta t) = \vec{\omega}(t + \Delta t/2) + \frac{1}{6}(k_1 + k_2 + 4k_3) \\ & k_1 = \Delta t/2\; \dot{\vec{\omega}}(\vec{\omega}(t), \vec{\tau}(t)) \\ & k_2 = \Delta t/2\; \dot{\vec{\omega}}(\vec{\omega}(t) + k1, \vec{\tau}(t)) \\ & k_3 = \Delta t/2\; \dot{\vec{\omega}}(\vec{\omega}(t) + (k1 + k2)/4, \vec{\tau}(t)) \\\end{split}\]

Where the angular velocity derivative is a function of the torque and angular velocity:

\[\dot{\vec{\omega}} = (\tau - \vec{\omega} \times (I \vec{\omega}))I^{-1}\]

Parameters:

• state (State) – Current state of the simulation.

• system (System) – Simulation system configuration.

Returns:

• Tuple[State, System] – The updated state and system after one time step.

• Reference

• ———–

• del Valle et. al, SPIRAL (An efficient algorithm for the integration of the equation of rotational motion, https://doi.org/10.1016/j.cpc.2023.109077.)

Note

• This method donates state and system.

class jaxdem.integrators.VicsekExtrinsic(neighbor_radius: Array, eta: Array, v0: Array, max_neighbors: int)

Bases: LinearIntegrator

Vicsek-model integrator with extrinsic (vectorial) noise.

This integrator implements a Vicsek-like update rule by directly setting the translational velocity magnitude to v0 each step, based on the direction of a vector that combines:

• the current accumulated force vector (from colliders + force functions),

• the average neighbor velocity direction (including self),

• an additive random unit vector scaled by eta (extrinsic noise).

Notes

• Noise is generated per clump (one sample per rigid body) and then broadcast to all clump members so clumps move coherently.

• Neighbor lists may be cached (e.g., NeighborList collider) or may sort the state (e.g., some cell-list builders). This integrator uses the returned state from create_neighbor_list for consistency.

neighbor_radius: Array

eta: Array

v0: Array

max_neighbors: int

static step_after_force(state: State, system: System) → Tuple[State, System]

class jaxdem.integrators.VicsekIntrinsic(neighbor_radius: Array, eta: Array, v0: Array, max_neighbors: int)

Bases: LinearIntegrator

Vicsek-model integrator with intrinsic noise.

This variant perturbs the direction of the desired motion by applying a random rotation to the normalized base direction (rather than adding a random vector in force space as in the extrinsic / vectorial-noise variant).

The base direction is computed from: - the current accumulated force vector (from colliders + force functions), - the average neighbor velocity direction (including self), then noise is applied per clump and broadcast to all clump members so clumps move coherently.

neighbor_radius: Array

eta: Array

v0: Array

max_neighbors: int

static step_after_force(state: State, system: System) → Tuple[State, System]

Modules

direct_euler	Direct (forward) Integrator.
spiral	Angular-velocity integrator based on the spiral scheme.
velocity_verlet	Velocity Verlet Integrator.
velocity_verlet_spiral	Angular-velocity integrator based on the spiral scheme.
vicsek	Vicsek-style integrators (extrinsic and intrinsic noise).