jaxdem.domains#

Simulation domains and boundary-condition implementations.

Classes

Domain(box_size, anchor)

The base interface for defining the simulation domain and the effect of its boundary conditions.
class jaxdem.domains.Domain(box_size: Array, anchor: Array)[source]#

Bases: Factory, ABC

The base interface for defining the simulation domain and the effect of its boundary conditions.

The Domain class defines how:

  • Relative displacement vectors between particles are calculated.

  • Particles’ positions are “shifted” or constrained to remain within the defined simulation boundaries based on the boundary condition type.

Example

To define a custom domain, inherit from Domain and implement its abstract methods:

>>> @Domain.register("my_custom_domain")
>>> @jax.tree_util.register_dataclass
>>> @dataclass(slots=True)
>>> class MyCustomDomain(Domain):
        ...
box_size: jax.Array#

Length of the simulation domain along each dimension.

anchor: jax.Array#

Anchor position (minimum coordinate) of the simulation domain.

property periodic: bool[source]#

Whether the domain enforces periodic boundary conditions.

classmethod Create(dim: int, box_size: Array | None = None, anchor: Array | None = None) → Self[source][source]#

Default factory method for the Domain class.

This method constructs a new Domain instance with a box-shaped domain of the given dimensionality. If box_size or anchor are not provided, they are initialized to default values.

Parameters:

  • dim (int) – The dimensionality of the domain (e.g., 2, 3).

  • box_size (jax.Array, optional) – The size of the domain along each dimension. If not provided, defaults to an array of ones with shape (dim,).

  • anchor (jax.Array, optional) – The anchor (origin) of the domain. If not provided, defaults to an array of zeros with shape (dim,).

Returns:

A new instance of the Domain subclass with the specified or default configuration.

Return type:

Domain

Raises:

AssertionError – If box_size and anchor do not have the same shape.

static displacement(ri: jax.Array, rj: jax.Array, system: System) → jax.Array[source][source]#

Computes the displacement vector between two particles \(r_i\) and \(r_j\), considering the domain’s boundary conditions.

Parameters:

  • ri (jax.Array) – Position vector of the first particle \(r_i\). Shape (dim,).

  • rj (jax.Array) – Position vector of the second particle \(r_j\). Shape (dim,).

  • system (System) – The configuration of the simulation, containing the domain instance.

Returns:

The displacement vector \(r_{ij} = r_i - r_j\), adjusted for boundary conditions. Shape (dim,).

Return type:

jax.Array

Example

>>> rij = system.domain.displacement(ri, rj, system)
static apply(state: State, system: System) → Tuple['State', 'System'][source][source]#

Applies boundary conditions during the simulation step.

This method updates the state based on the domain’s rules, ensuring particles handle interactions at boundaries appropriately (e.g., reflection).

Parameters:

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

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

Returns:

A tuple containing the updated State object adjusted by the boundary conditions and the System object.

Return type:

Tuple[State, System]

Note

  • Periodic boundary conditions dont require wrapping of the coordinates during time stepping,

but reflective boundaries require changing positions and velocities. To wrap positions for periodic boundaries so they are displayed correctly when saving, and other algorithms use the shift method. - This method donates state and system

Example

>>> state, system = system.domain.apply(state, system)
static shift(state: State, system: System) → Tuple['State', 'System'][source][source]#

This method updates the state based on the domain’s rules, ensuring particles remain within the simulation box or handle interactions at boundaries appropriately (e.g., reflection, wrapping).

Parameters:

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

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

Returns:

A tuple containing the updated State object adjusted by the boundary conditions and the System object.

Return type:

Tuple[State, System]

Example

>>> state, system = system.domain.shift(state, system)
class jaxdem.domains.FreeDomain(box_size: Array, anchor: Array)[source]#

Bases: Domain

A Domain implementation representing an unbounded, “free” space.

In a FreeDomain, there are no explicit boundary conditions applied to particles. Particles can move indefinitely in any direction, and the concept of a “simulation box” is only used to define the bounding box of the system.

Notes

  • The box_size and anchor attributes are dynamically updated in the shift method to encompass all particles. Some hashing tools require the domain size.

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

Updates the System’s domain anchor and box_size to encompass all particles. Does not apply any transformations to the state.

Parameters:

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

  • system (System) – The current system configuration.

Returns:

The original State object (unchanged) and the System object with updated domain.anchor and domain.box_size.

Return type:

Tuple[State, System]

Note

  • This method donates state and system

class jaxdem.domains.PeriodicDomain(box_size: Array, anchor: Array)[source]#

Bases: Domain

A Domain implementation that enforces periodic boundary conditions.

Particles that move out of one side of the simulation box re-enter from the opposite side. The displacement vector between particles is computed using the minimum image convention.

static displacement(ri: jax.Array, rj: jax.Array, system: System) → jax.Array[source][source]#

Computes the minimum image displacement vector between two particles \(r_i\) and \(r_j\).

For periodic boundary conditions, the displacement is calculated as the shortest vector that connects \(r_j\) to \(r_i\), potentially by crossing periodic boundaries.

Parameters:

  • ri (jax.Array) – Position vector of the first particle \(r_i\).

  • rj (jax.Array) – Position vector of the second particle \(r_j\).

  • system (System) – The configuration of the simulation, containing the domain instance with anchor and box_size for periodicity.

Returns:

The minimum image displacement vector:

\[\begin{split}& r_{ij} = (r_i - a) - (r_j - a) \\ & r_{ij} = r_{ij} - B \cdot \text{round}(r_{ij}/B)\end{split}\]
where:

Return type:

jax.Array

property periodic: bool[source]#

Whether the domain enforces periodic boundary conditions.

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

Wraps particles back into the primary simulation box.

\[\begin{split}r = r - B \cdot \text{floor}((r - a)/B) \\\end{split}\]
where:
Parameters:

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

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

Returns:

The updated State object with wrapped particle positions, and the System object.

Return type:

Tuple[State, System]

class jaxdem.domains.ReflectDomain(box_size: Array, anchor: Array, restitution_coefficient: Array)[source]#

Bases: Domain

A Domain implementation that enforces reflective boundary conditions.

Particles that attempt to move beyond the defined box_size will have their positions reflected back into the box and their velocities reversed in the direction normal to the boundary.

restitution_coefficient: jax.Array#
classmethod Create(dim: int, box_size: Array | None = None, anchor: Array | None = None, restitution_coefficient: float = 1.0) → Self[source][source]#

Default factory method for the Domain class.

This method constructs a new Domain instance with a box-shaped domain of the given dimensionality. If box_size or anchor are not provided, they are initialized to default values.

Parameters:

  • dim (int) – The dimensionality of the domain (e.g., 2, 3).

  • box_size (jax.Array, optional) – The size of the domain along each dimension. If not provided, defaults to an array of ones with shape (dim,).

  • anchor (jax.Array, optional) – The anchor (origin) of the domain. If not provided, defaults to an array of zeros with shape (dim,).

  • restitution_coefficient (float) – Restitution coefficient between 0 and 1 to modulate energy conservation with wall.

Returns:

A new instance of the Domain subclass with the specified or default configuration.

Return type:

ReflectDomain

Raises:

AssertionError – If box_size and anchor do not have the same shape.

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

Applies reflective boundary conditions to particles.

Particles are checked against the domain boundaries. If a particle attempts to move beyond a boundary, it is reflected. The reflection is governed by the impulse-momentum equations for rigid bodies.

Velocity Update (Impulse)

\[\begin{split}\vec{v}' &= \vec{v} + \frac{1}{m}\vec{J} \\ \vec{\omega}' &= \vec{\omega} + \mathbf{I}^{-1} (\vec{r}_{p} \times \vec{J})\end{split}\]

where the impulse vector \(J\) is:

\[\vec{J} = \frac{-(1+e)(\vec{v}_{contact} \cdot \hat{n})}{\frac{1}{m} + [\mathbf{I}^{-1} (\vec{r}_{p} \times \hat{n})] \cdot (\vec{r}_{p} \times \hat{n})} \hat{n}\]

and the velocity of the contact point \(\vec{v}_{contact}\) is:

\[\vec{v}_{contact} = \vec{v} + \vec{\omega} \times \vec{r}_{p}\]

Position Update

Finally, the particle is moved out of the boundary by reflecting its position based on the penetration depth \(\delta\):

\[\vec{r}_c' = \vec{r}_c + 2 \delta \hat{n}\]

Definitions

  • \(\vec{r}_c\): Particle center of mass position (jaxdem.State.pos_c).

  • \(\vec{r}_{p}\): Vector from COM to contact sphere in the lab frame (jaxdem.State.pos_p).

  • \(\vec{v}\): Particle linear velocity (jaxdem.State.vel).

  • \(\vec{\omega}\): Particle angular velocity (jaxdem.State.angVel).

  • \(\hat{n}\): Boundary normal vector (pointing into the domain).

  • \(\delta\): Penetration depth (positive value).

  • \(e\): Coefficient of restitution.

Parameters:

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

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

Returns:

The updated State object with reflected positions and velocities, and the System object.

Return type:

Tuple[State, System]

Note

  • This method donates state and system

https://www.myphysicslab.com/engine2D/collision-en.html

class jaxdem.domains.ReflectSphereDomain(box_size: Array, anchor: Array)[source]#

Bases: Domain

A Domain implementation that enforces reflective boundary conditions only for spheres. We have this dedicated version for performance reasons.

Particles that attempt to move beyond the defined box_size will have their positions reflected back into the box and their velocities reversed in the direction normal to the boundary.

Notes

  • The reflection occurs at the boundaries defined by anchor and anchor + box_size.

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

Applies reflective boundary conditions to particles.

Particles are checked against the domain boundaries. If a particle attempts to move beyond a boundary, its position is reflected back into the box, and its velocity component normal to that boundary is reversed.

\[\begin{split}l &= a + R \\\\ u &= a + B - R \\\\ v' &= \\begin{cases} -v & \\text{if } r < l \\text{ or } r > u \\\\ v & \\text{otherwise} \\end{cases} \\\\ r' &= \\begin{cases} 2l - r & \\text{if } r < l \\\\ r & \\text{otherwise} \\end{cases} \\\\ r'' &= \\begin{cases} 2u - r' & \\text{if } r' > u \\\\ r' & \\text{otherwise} \\end{cases} r = r''\end{split}\]
where:

TO DO: Ensure correctness when adding different types of shapes and angular vel

Parameters:

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

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

Returns:

The updated State object with reflected positions and velocities, and the System object.

Return type:

Tuple[State, System]

Note

  • This method donates state and system

  • Only works for states with ONLY spheres.

Modules

free

Unbounded (free) simulation domain.

periodic

Periodic boundary-condition domain.

reflect

Reflective boundary-condition domain.

reflect_sphere

Reflective boundary-condition domain.