jaxdem.domains#

Simulation domains and boundary-condition implementations.

Classes

Domain(box_size, inv_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, inv_box_size: Array, anchor: Array)#

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

  • **kw (Any) – Extra keyword arguments forwarded verbatim 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]#

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]#

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 domains do not need to wrap coordinates during time stepping, so their apply is a no-op; wrapping is done by shift() (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]#

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, inv_box_size: Array, anchor: Array)#

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 apply method to encompass all particles. Some hashing tools require the domain size.

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

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, periodic images are offset along the shear-flow axis alpha by the current shear strain gamma.

gamma is a plain state field that both displacement() and shift() read directly. The domain does not advance it; the shear protocol is imposed 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. Updated externally to impose the desired 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]#

Computes 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]#

Wraps 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 displacement vector between particles is computed using 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]#

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

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

Wraps 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.

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

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]#

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

Verlet Time-of-Collision Correction

The collision time fraction \(\alpha \in [0, 1]\) is obtained per clump from the shared Verlet-consistent solver jaxdem.domains._toc.verlet_collision_fraction() (also used by ReflectSphereDomain), evaluated at the contact point. The contact-point velocity and angular velocity at the moment of collision are reconstructed as \(v_{col} = v + (\alpha - 1) \Delta t\, a\) before the impulse is applied, and the post-impulse velocity change is then integrated 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. 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, modulated 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]#

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 (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 collision time fraction \(\alpha \in [0, 1]\) and the velocity at the moment of collision are obtained from the shared Verlet-consistent solver jaxdem.domains._toc.verlet_collision_fraction() (also used by ReflectDomain), and the pre-collision velocity is reconstructed as \(v_{col} = v + (\alpha - 1) \Delta t\, a\).

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

  • 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.