jaxdem.domains#

Simulation domains and boundary-condition implementations.

Classes

Domain(box_size, inv_box_size, anchor)

The base interface for the simulation domain and its boundary conditions.
class jaxdem.domains.Domain(box_size: Array, inv_box_size: Array, anchor: Array)#

Bases: Factory, ABC

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

The Domain class defines:

  • How to compute relative displacement vectors between particles.

  • How to “shift” or constrain particle positions so they stay within the simulation boundaries.

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: Array#

Length of the simulation domain along each dimension.

inv_box_size: Array#

Inverse length of the simulation domain along each dimension.

anchor: 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, **kw: Any) → Self[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 you do not provide box_size or anchor, they default to the values below.

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,).

  • **kw (Any) – Extra keyword arguments passed to the subclass constructor (e.g. restitution_coefficient for reflective domains).

Returns:

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

Return type:

Domain

Raises:

ValueError – If box_size or anchor do not have shape (dim,).

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

Compute the displacement vector between two particles \(r_i\) and \(r_j\), respecting 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]#

Apply boundary conditions during the simulation step.

This method updates the state with the domain’s rules so particles handle boundary interactions (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 domains do not need to wrap coordinates during time stepping, so their apply is a no-op. shift() wraps the coordinates instead (e.g. when saving, so positions are displayed inside the box). Reflective domains, in contrast, must update positions and velocities here.

Example

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

Shift particles according to the domain’s boundary-condition rules.

This method updates the state with the domain’s rules so particles stay within the simulation box or handle boundary interactions (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, inv_box_size: Array, anchor: Array)#

Bases: Domain

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

A FreeDomain applies no explicit boundary conditions to particles. Particles can move indefinitely in any direction. The “simulation box” only defines the bounding box of the system.

Notes

  • The apply method updates the box_size and anchor attributes to encompass all particles. Some hashing tools require the domain size.

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

Update the domain anchor and box_size of the System to encompass all particles.

This method does not transform 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]

class jaxdem.domains.LeesEdwardsDomain(box_size: Array, inv_box_size: Array, anchor: Array, gamma: Array, alpha_axis: Array, beta_axis: Array, alpha: int = 0, beta: int = 1)#

Bases: Domain

A Domain implementation that enforces Lees-Edwards boundary conditions.

The domain is periodic in all directions. Across the shear-gradient axis beta, the current shear strain gamma offsets the periodic images along the shear-flow axis alpha.

gamma is a plain state field that both displacement() and shift() read directly. The domain does not advance it. You impose the shear protocol externally by updating gamma between steps (e.g. in user_post_step_actions). For example, constant-rate shear is:

from dataclasses import replace

def shear(state, system):
    gamma = system.domain.gamma + gamma_dot * system.dt
    return state, replace(system, domain=replace(system.domain, gamma=gamma))

while oscillatory shear sets gamma = gamma_amp * jnp.sin(omega * system.time).

gamma: Array#

Current shear strain. The Lees-Edwards image offset along alpha is gamma * L_beta. Update it externally to impose the shear protocol.

alpha_axis: Array#

One-hot vector for the shear-flow coordinate.

beta_axis: Array#

One-hot vector for the shear-gradient coordinate.

alpha: int#

Index of the shear-flow coordinate.

beta: int#

Index of the shear-gradient coordinate.

classmethod Create(dim: int, box_size: Array | None = None, anchor: Array | None = None, gamma: float | Array = 0.0, alpha: int = 0, beta: int = 1, **kwargs: Any) → LeesEdwardsDomain[source]#

Construct a Lees-Edwards domain with validated shear axes.

property periodic: bool[source]#

Whether the domain enforces periodic boundary conditions.

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

Compute the shear-periodic minimum image displacement vector.

When the minimum image crosses the shear-gradient axis beta, the displacement is shifted along the shear-flow axis alpha by \(\gamma L_\beta\) per crossed image.

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 box_size and Lees-Edwards shear parameters.

Returns:

The shear-periodic minimum image displacement vector:

\[\begin{split}& r_{ij} = r_i - r_j \\\\ & r_{ij,\alpha} = r_{ij,\alpha} - \operatorname{round}(r_{ij,\beta}/L_\beta)\gamma L_\beta \\\\ & r_{ij} = r_{ij} - L \left\lfloor 0.5 + r_{ij}/L \right\rfloor\end{split}\]
where:

Return type:

jax.Array

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

Wrap particles back into the primary shear-periodic simulation box.

\[\begin{split}& n_\beta = \left\lfloor (r_\beta - a_\beta)/L_\beta \right\rfloor \\\\ & r_\alpha = r_\alpha - n_\beta \gamma L_\beta \\\\ & r = r - L \left\lfloor (r-a)/L \right\rfloor\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.PeriodicDomain(box_size: Array, inv_box_size: Array, anchor: Array)#

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 domain computes the displacement vector between particles with the minimum image convention.

property periodic: bool[source]#

Whether the domain enforces periodic boundary conditions.

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

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

For periodic boundary conditions, the displacement is the shortest vector that connects \(r_j\) to \(r_i\), possibly 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

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

Wrap particles back into the primary simulation box.

\[r = r - B \cdot \text{floor}((r - a)/B)\]
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, inv_box_size: Array, anchor: Array, restitution_coefficient: Array)#

Bases: Domain

A Domain implementation that enforces reflective boundary conditions.

When a particle moves beyond the defined box_size, the domain reflects its position back into the box. It also reverses the velocity component normal to the boundary.

restitution_coefficient: Array#
classmethod Create(dim: int, box_size: Array | None = None, anchor: Array | None = None, restitution_coefficient: float = 1.0, **kw: Any) → Self[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 you do not provide box_size or anchor, they default to the values below.

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:

ValueError – If box_size or anchor have the wrong shape, or if restitution_coefficient is outside (0, 1].

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

Apply reflective boundary conditions to particles.

The method checks particles against the domain boundaries. When a particle moves beyond a boundary, the method reflects it. The impulse-momentum equations for rigid bodies govern the reflection.

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}\]

Verlet Time-of-Collision Correction

The shared Verlet-consistent solver jaxdem.domains._toc.verlet_collision_fraction() (also used by ReflectSphereDomain) computes the collision time fraction \(\alpha \in [0, 1]\) per clump at the contact point. Before the impulse, the method reconstructs the contact-point velocity and angular velocity at the moment of collision as \(v_{col} = v + (\alpha - 1) \Delta t\, a\). It then integrates the post-impulse velocity change over the remaining \((1 - \alpha) \Delta t\) to correct positions and orientations.

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.ang_vel).

  • \(\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:

  • Tuple[State, System] – The updated State object with reflected positions and velocities, and the System object.

  • Reference

  • ———-

  • https (//www.myphysicslab.com/engine2D/collision-en.html)

class jaxdem.domains.ReflectSphereDomain(box_size: Array, inv_box_size: Array, anchor: Array, restitution_coefficient: Array)#

Bases: Domain

A Domain implementation that enforces reflective boundary conditions only for spheres. This dedicated version exists for performance.

When a particle moves beyond the defined box_size, the domain reflects its position back into the box. It also reverses the velocity component normal to the boundary, scaled by restitution_coefficient.

Notes

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

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

Default factory method for the ReflectSphereDomain class.

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 with the specified or default configuration.

Return type:

ReflectSphereDomain

Raises:

ValueError – If box_size or anchor have the wrong shape, or if restitution_coefficient is outside (0, 1].

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

Apply reflective boundary conditions to particles.

The method checks particles against the domain boundaries. When a particle moves beyond a boundary, the method reflects its position back into the box. It also reverses the velocity component normal to that boundary, scaled by the restitution coefficient \(e\).

\[\begin{split}l &= a + R \\ u &= a + B - R \\ v' &= \begin{cases} -e\,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:

  • \(r\) is the current particle position (jaxdem.State.pos)

  • \(v\) is the current particle velocity (jaxdem.State.vel)

  • \(a\) is the domain anchor (Domain.anchor)

  • \(B\) is the domain box size (Domain.box_size)

  • \(R\) is the particle radius (jaxdem.State.rad)

  • \(l\) is the lower boundary for the particle center

  • \(u\) is the upper boundary for the particle center

  • \(e\) is the restitution coefficient.

Verlet Time-of-Collision Correction

The shared Verlet-consistent solver jaxdem.domains._toc.verlet_collision_fraction() (also used by ReflectDomain) computes the collision time fraction \(\alpha \in [0, 1]\) and the velocity at the moment of collision. The method reconstructs the pre-collision velocity as \(v_{col} = v + (\alpha - 1) \Delta t\, a\).

TO DO: Check correctness when adding different shape types and angular velocity

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

  • Only works for states with ONLY spheres.

Modules

free

Unbounded (free) simulation domain.

lees_edwards

Lees-Edwards shear-periodic boundary-condition domain.

periodic

Periodic boundary-condition domain.

reflect

Reflective boundary-condition domain.

reflect_sphere

Reflective boundary-condition domain.