jaxdem.colliders.sweep_and_prune#

Sweep and prune \(O(N log N)\) collider implementation.

Functions

compute_hash(state, proj_perp, aabb, shift)

compute_virtual_shift(m, M, HASH)

force(i, j, state, system)

Compute linear spring-like interaction force acting on particle \(i\) due to particle \(j\).

pad_state(state)

pad_to_power2(x)

Pad odd-dimensional vectors to an even size (Pallas kernel limitation).

sap_kernel_full(state_ref, system_ref, ...)

sort(state, iota, m, M)

Classes

SweeAPrune

SweepAndPrune()

jaxdem.colliders.sweep_and_prune.pad_to_power2(x: Array) → Array[source][source]#

Pad odd-dimensional vectors to an even size (Pallas kernel limitation).

jaxdem.colliders.sweep_and_prune.sap_kernel_full(state_ref: Any, system_ref: Any, aabb_ref: Array, m_ref: Array, M_ref: Array, HASH_ref: Array, forces_ref: Array) → None[source][source]#
jaxdem.colliders.sweep_and_prune.compute_hash(state: Any, proj_perp: Array, aabb: Array, shift: Array) → Array[source][source]#
jaxdem.colliders.sweep_and_prune.compute_virtual_shift(m: Array, M: Array, HASH: Array) → Tuple[Array, Array][source][source]#
jaxdem.colliders.sweep_and_prune.sort(state: State, iota: jax.Array, m: jax.Array, M: jax.Array) → Tuple[State, jax.Array, jax.Array, jax.Array][source][source]#
jaxdem.colliders.sweep_and_prune.pad_state(state: State) → State[source][source]#
jaxdem.colliders.sweep_and_prune.force(i: int, j: int, state: State, system: System) → Tuple[jax.Array, jax.Array][source][source]#

Compute linear spring-like interaction force acting on particle \(i\) due to particle \(j\).

Returns zero when \(i = j\).

Parameters:

  • i (int) – Index of the first particle.

  • j (int) – Index of the second particle.

  • state (State) – Current state of the simulation.

  • system (System) – Simulation system configuration.

Returns:

Force vector acting on particle \(i\) due to particle \(j\).

Return type:

jax.Array

class jaxdem.colliders.sweep_and_prune.SweepAndPrune[source]#

Bases: Collider

static create_neighbor_list(state: State, system: System, cutoff: float, max_neighbors: int) → Tuple[State, System, jax.Array, jax.Array][source][source]#
static compute_force(state: State, system: System) → Tuple[State, System][source][source]#
jaxdem.colliders.sweep_and_prune.SweeAPrune#

alias of SweepAndPrune