jaxdem.integrators#

Time-integration interfaces and implementations.

Classes

Integrator()

Abstract base class for defining the interface for time-stepping.

class jaxdem.integrators.Integrator[source]#

Bases: Factory, ABC

Abstract base class for defining 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, frozen=True)
>>> class MyCustomIntegrator(Integrator):
        ...

abstractmethod static step(state: State, system: System) → Tuple['State', 'System'][source][source]#

Advance the simulation state by one time step using a specific numerical integration method.

Parameters:

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

Returns:

A tuple containing the updated State and System after one time step of integration.

Return type:

Tuple[State, System]

Notes

This method performs the following updates:
1. Applies boundary conditions using jaxdem.Domain.shift().
2. Computes forces and accelerations using jaxdem.Collider.compute_force().
3. Updates velocities based on current acceleration.
4. Updates positions based on the newly updated velocities.
Particles with state.fixed set to True will have their velocities and positions unaffected by the integration step.

Example

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

abstractmethod static initialize(state: State, system: System) → Tuple['State', 'System'][source][source]#

Some integration methods require an initialization step, for example LeapFrog. This function implements the interface for the initialization.

Parameters:

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

Returns:

A tuple containing the updated State and System after the initialization.

Return type:

Tuple[State, System]

Example

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

class jaxdem.integrators.DirectEuler[source]#

Bases: Integrator

Implements the explicit (forward) Euler integration method.

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

The Direct Euler integrator does not require a specific initialization step.

Parameters:

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

Returns:

The original State and System objects.

Return type:

Tuple[State, System]

classmethod registry_name() → str[source]#

static step(state: State, system: System) → Tuple['State', 'System'][source][source]#

Advances the simulation state by one time step 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 (jaxdem.State.accel)
\(\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]

property type_name: str[source]#

Modules

direct_euler

Direct (forward) Euler integrator.