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

• TO DO: make it work without padding the vectors

class jaxdem.integrators.VelocityVerlet

Bases: LinearIntegrator

Implements the Velocity Verlet integration method.

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

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

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_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 based on the spiral leapfrog algorithm to implement a velocity verlet like version.

\[\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

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 based on the spiral leapfrog algorithm to implement a velocity verlet like version.

• 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 wich 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

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.