jaxdem.domains.lees_edwards#

Lees-Edwards shear-periodic boundary-condition domain.

Classes

LeesEdwardsDomain(box_size, inv_box_size, ...)

A Domain implementation that enforces Lees-Edwards boundary conditions.
class jaxdem.domains.lees_edwards.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:

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

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:

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

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

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]