jaxdem.integrator#

Interface for defining time integrators. The time integrator performs one simulation step.

Classes

DirectEuler()

Implements the explicit (forward) Euler integration method.

Integrator()

Abstract base class for defining the interface for time-stepping.
class jaxdem.integrator.Integrator[source][source]#

Bases: Factory[Integrator], ABC

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

Notes

  • Must Support 2D and 3D domains.

  • Must be jit compatible

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):
        ...
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]

Raises:

NotImplementedError – This is an abstract method and must be implemented by subclasses.

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]

Raises:

NotImplementedError – This is an abstract method and must be implemented by subclasses.

Example

>>> state, system = system.integrator.initialize(state, system)
classmethod create(key: str, /, **kw: Any) → T[source]#

Creates and returns an instance of a registered subclass.

This method looks up the subclass associated with the given key in the factory’s registry and then calls its constructor with the provided arguments.

Parameters:

  • key (str) – The registration key of the subclass to be created.

  • **kw (Any) – Arbitrary keyword arguments to be passed directly to the constructor of the registered subclass.

Returns:

An instance of the registered subclass.

Return type:

T

Raises:

  • KeyError – If the provided key is not found in the factory’s registry.

  • 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)
classmethod register(key: str | None = None) → Callable[[Type[T]], Type[T]][source]#

Registers a subclass with the factory’s registry.

This method returns a decorator that can be applied to a class to register it under a specific key.

Parameters:

key (str or None, optional) – The string key under which to register the subclass. If None, the lowercase name of the subclass itself will be used as the key.

Returns:

A decorator function that takes a class and registers it, returning the class 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.

Example

Register a class named “MyComponent” under the key “mycomp”:

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

Register a class named “DefaultComponent” using its own name as the key:

>>> @MyFactory.register()
>>> class DefaultComponent:
>>>     ...
class jaxdem.integrator.DirectEuler[source][source]#

Bases: Integrator

Implements the explicit (forward) Euler integration method.

Notes

  • This method performs the following updates:

    1. Applies boundary conditions using system.domain.shift().

    2. Computes forces and accelerations using system.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.

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 (state.pos)

  • \(v\) is the particle velocity (state.vel)

  • \(a\) is the particle acceleration (state.accel)

  • \(\Delta t\) is the time step (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]

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 create(key: str, /, **kw: Any) → T[source]#

Creates and returns an instance of a registered subclass.

This method looks up the subclass associated with the given key in the factory’s registry and then calls its constructor with the provided arguments.

Parameters:

  • key (str) – The registration key of the subclass to be created.

  • **kw (Any) – Arbitrary keyword arguments to be passed directly to the constructor of the registered subclass.

Returns:

An instance of the registered subclass.

Return type:

T

Raises:

  • KeyError – If the provided key is not found in the factory’s registry.

  • 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)
classmethod register(key: str | None = None) → Callable[[Type[T]], Type[T]][source]#

Registers a subclass with the factory’s registry.

This method returns a decorator that can be applied to a class to register it under a specific key.

Parameters:

key (str or None, optional) – The string key under which to register the subclass. If None, the lowercase name of the subclass itself will be used as the key.

Returns:

A decorator function that takes a class and registers it, returning the class 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.

Example

Register a class named “MyComponent” under the key “mycomp”:

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

Register a class named “DefaultComponent” using its own name as the key:

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