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.
`SearchGeometry`(*values)	Spatial-search geometry implemented by a domain.

class jaxdem.domains.SearchGeometry(*values)[source]#

Bases: Enum

Spatial-search geometry implemented by a domain.

ORTHOGONAL = 'orthogonal'#

SHEAR_PERIODIC = 'shear_periodic'#

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.

search_geometry: ClassVar[SearchGeometry | None] = None#

search_geometry_snapshot() → Array[source]#: Return fixed-shape geometry data used to validate search caches.

search_geometry_drift(previous: Array, old_pos: Array) → Array[source]#: Bound image motion since the geometry snapshot of the last cache build.

neighbor_cells(coords: Array, image: Array, grid_dims: Array, cell_size: Array) → Array[source]#

Adjust wrapped stencil cells for the domain’s periodic image rule.

coords and signed image counts have shape (N, stencil, dim). The base rule keeps the orthogonal stencil; shear-periodic domains may expand it to cover fractional image offsets.

rates_between(old: Domain, dt: Array) → Domain[source]#

Return this domain with image rates for the change from old over dt.

Domains with moving periodic images (shear strain, box length) set the rate fields that relative_velocity() reads, so a finite geometry change can be expressed as velocities. The base domain has no moving images and returns itself.

property periodic: bool[source]#: Whether the domain enforces periodic boundary conditions.

static update_bounds(pos: jax.Array, system: System, padding: float | jax.Array = 0.0) → System[source]#: Update query bounds when the domain derives them from particle positions.

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 relative_velocity(ri: jax.Array, rj: jax.Array, vi: jax.Array, vj: jax.Array, system: System) → jax.Array[source]#

Return relative velocity for the image selected by displacement.

ri and rj are the particle centers used to select that image. vi and vj may include rotational contact-point velocities. Domains with moving periodic images override this operation to add the image’s translational velocity offset.

static resize(state: State, system: System, box_size: jax.Array, group_id: jax.Array | None = None) → tuple[State, System][source]#

Rescale the box to box_size and move bodies affinely about the anchor.

Clumps translate with their stored centers of mass, keeping member offsets and orientations fixed. Explicit groups translate together using the mean of their pos_c values. Pair history is advanced once for the resulting relative motion without time passing, so velocities, dt, and rate-dependent forces play no part. Caches rebuild when the geometry change exhausts their skin.

The history update is not a no-op even though the map is affine. Bodies move rigidly, so two member spheres of different bodies are displaced by the difference of their body offsets, which is not parallel to their contact normal; the tangential spring of such a contact changes. For single spheres under isotropic scaling the relative displacement is purely normal and the stored spring is unchanged. Periodic images move with the box, so the domain’s box_rate supplies their velocity and a pair interacting across a boundary sees only its true relative motion.

Parameters:

state (State, System) – Configuration to rescale.
system (State, System) – Configuration to rescale.
box_size (jax.Array) – New box lengths, shape (dim,).
group_id (jax.Array, optional) – Optional per-particle group labels, shape (N,). Each group translates using its mean pos_c. Without labels, stored centers move directly. Passing bond-graph component labels preserves a deformable body’s shape during the resize.

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)

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.