jaxdem.colliders#

Collision-detection interfaces and implementations.

Classes

Collider()

The base interface for defining how contact detection and force computations are performed in a simulation.
class jaxdem.colliders.Collider[source]#

Bases: Factory, ABC

The base interface for defining how contact detection and force computations are performed in a simulation.

Concrete subclasses of Collider implement the specific algorithms for calculating the interactions.

Notes

Self-interaction (i.e., calling the force/energy computation for i=j) is allowed, and the underlying force_model is responsible for correctly handling or ignoring this case.

Example

To define a custom collider, inherit from Collider, register it and implement its abstract methods:

>>> @Collider.register("CustomCollider")
>>> @jax.tree_util.register_dataclass
>>> @dataclass(slots=True, frozen=True)
>>> class CustomCollider(Collider):
        ...

Then, instantiate it:

>>> jaxdem.Collider.create("CustomCollider", **custom_collider_kw)
abstractmethod static compute_force(state: State, system: System) → Tuple['State', 'System'][source][source]#

Abstract method to compute the total force acting on each particle in the simulation.

Implementations should calculate inter-particle forces based on the current state and system configuration, then update the accel attribute of the state object with the resulting total acceleration for each particle.

TO DO: DEFINE HOW TO RESET THE FORCE AND HOW TO ADD FORCE EXTERNALLY

Parameters:

  • state (State) – The current state of the simulation.

  • system (System) – The configuration of the simulation.

Returns:

A tuple containing the updated State object (with computed accelerations) and the System object.

Return type:

Tuple[State, System]

abstractmethod static compute_potential_energy(state: State, system: System) → jax.Array[source][source]#

Abstract method to compute the total potential energy of the system.

Implementations should calculate the sum per particle of all potential energies present in the system based on the current state and system configuration.

Parameters:

  • state (State) – The current state of the simulation.

  • system (System) – The configuration of the simulation.

Returns:

A scalar JAX array representing the total potential energy of each particle.

Return type:

jax.Array

Example

>>> potential_energy = system.collider.compute_potential_energy(state, system)
>>> print(f"Potential energy per particle: {potential_energy:.4f}")
>>> print(potential_energy.shape") # (N, 1)
class jaxdem.colliders.NaiveSimulator[source]#

Bases: Collider

Implementation that computes forces and potential energies using a naive \(O(N^2)\) all-pairs interaction loop.

Notes

Due to its \(O(N^2)\) complexity, NaiveSimulator is suitable for simulations with a relatively small number of particles. For larger systems, a more efficient spatial partitioning collider should be used. However, this collider should be the fastest option for small systems (\(<1k-5k\) spheres depending on the GPU).

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

Computes the total force acting on each particle using a naive \(O(N^2)\) all-pairs loop.

This method sums the force contributions from all particle pairs (i, j) as computed by the system.force_model and updates the particle accelerations.

Parameters:

  • state (State) – The current state of the simulation.

  • system (System) – The configuration of the simulation.

Returns:

A tuple containing the updated State object with computed accelerations and the unmodified System object.

Return type:

Tuple[State, System]

static compute_potential_energy(state: State, system: System) → jax.Array[source][source]#

Computes the potential energy associated with each particle using a naive \(O(N^2)\) all-pairs loop.

This method iterates over all particle pairs (i, j) and sums the potential energy contributions computed by the system.force_model.

Parameters:

  • state (State) – The current state of the simulation.

  • system (System) – The configuration of the simulation.

Returns:

One-dimensional array containing the total potential energy contribution for each particle.

Return type:

jax.Array

classmethod registry_name() → str[source]#
property type_name: str[source]#

Modules

naive

Naive \(O(N^2)\) collider implementation.