jaxdem.domains#
Simulation domains and boundary-condition implementations.
Classes
|
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, frozen=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.
- periodic: ClassVar[bool] = False#
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:
- Raises:
AssertionError – If box_size and anchor do not have the same shape.
- abstractmethod 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)
- abstractmethod static shift(state: State, system: System) Tuple['State', 'System'] [source][source]#
Applies boundary conditions to particles state.
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:
- Returns:
A tuple containing the updated State object adjusted by the boundary conditions and the System object.
- Return type:
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 displacement(ri: jax.Array, rj: jax.Array, _: System) jax.Array [source][source]#
Computes the displacement vector between two particles.
In a free domain, the displacement is simply the direct vector difference between the particle positions.
- Parameters:
ri (jax.Array) – Position vector of the first particle \(r_i\).
rj (jax.Array) – Position vector of the second particle \(r_j\).
_ (System) – The system object.
- Returns:
The direct displacement vector \(r_i - r_j\).
- Return type:
jax.Array
- 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.
Notes
This domain type is periodic (periodic = True).
- 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:
\(a\) is the domain anchor (
Domain.anchor
)\(B\) is the domain box size (
Domain.box_size
)
- Return type:
jax.Array
- periodic: ClassVar[bool] = True#
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:
\(a\) is the domain anchor (
Domain.anchor
)\(B\) is the domain box size (
Domain.box_size
)
- class jaxdem.domains.ReflectDomain(box_size: Array, anchor: 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.
Notes
The reflection occurs at the boundaries defined by anchor and anchor + box_size.
- static displacement(ri: jax.Array, rj: jax.Array, _: System) jax.Array [source][source]#
Computes the displacement vector between two particles.
In a reflective domain, the displacement is simply the direct vector difference.
- Parameters:
ri (jax.Array) – Position vector of the first particle \(r_i\).
rj (jax.Array) – Position vector of the second particle \(r_j\).
_ (System) – The system object.
- Returns:
The direct displacement vector \(r_i - r_j\).
- Return type:
jax.Array
- static shift(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:
\(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
TO DO: Ensure correctness when adding different types of shapes and angular vel
Modules