jaxdem.colliders.naive#

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

Classes

NaiveSimulator()

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

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_potential_energy(state: State, system: System) → jax.Array[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:

Scalar containing the total potential energy of the system.

Return type:

jax.Array

static create_neighbor_list(state: State, system: System, cutoff: float, max_neighbors: int) → tuple[State, System, jax.Array, jax.Array][source]#

Computes a neighbor list using a naive \(O(N^2)\) all-pairs search.

Parameters:

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

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

  • cutoff (float) – The interaction radius (force cutoff).

  • max_neighbors (int) – Maximum number of neighbors to store per particle.

Returns:

A tuple containing: - state: The simulation state. - system: The simulation system. - neighbor_list: Array of shape (N, max_neighbors) containing neighbor indices. - overflow: Boolean flag indicating if any particle exceeded max_neighbors.

Return type:

Tuple[State, System, jax.Array, jax.Array]

static compute_force(state: State, system: System) → tuple[State, System][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 forces.

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 forces and the unmodified System object.

Return type:

Tuple[State, System]