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, gamma_dot: Array = <factory>, *, box_rate: Array = <factory>)#

Bases: PeriodicDomain

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 the current shear strain and gamma_dot its time derivative. Time stepping does not advance these fields automatically. Set the initial values before initializing forces and update the protocol in user_pre_step_actions, before each force evaluation. For constant-rate shear, set gamma_dot when creating the domain and use:

from dataclasses import replace

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

For oscillatory shear, set gamma = gamma_amp * jnp.sin(omega * system.time) and gamma_dot = gamma_amp * omega * jnp.cos(omega * system.time) in the same callback. Contact laws use relative_velocity to account for the moving image even when particle coordinates remain unwrapped.

For discrete affine shear, call shear to apply one strain increment to particle centers, periodic images, and any pair history. Quasistatic protocols can then minimize the returned state at fixed strain. Box rescaling goes through Domain.resize; box_rate is the box length rate for continuous protocols and is zero otherwise.

box_rate is inherited from PeriodicDomain. Periodic images move at image * box_rate when the box changes, and the shear offset gamma * L_beta of an image across the gradient boundary moves with L_beta as well. relative_velocity subtracts both so a pair interacting across a boundary sees only its true relative motion. See PeriodicDomain.box_rate.

For fixed integer image counts \(\mathbf n\), the image position and velocity are

\[\begin{split}\mathbf r_{j,\mathrm{image}} &= \mathbf r_j + \mathbf n \odot \mathbf L + n_\beta \gamma L_\beta \mathbf e_\alpha, \\ \mathbf v_{\mathrm{image}} &= \mathbf n \odot \dot{\mathbf L} + n_\beta (\dot\gamma L_\beta + \gamma\dot L_\beta) \mathbf e_\alpha.\end{split}\]

Here \(\mathbf v_{\mathrm{image}}\) is the velocity offset relative to particle j, so the contact relative velocity is \(\mathbf v_i - \mathbf v_j - \mathbf v_{\mathrm{image}}\).

search_geometry: ClassVar[SearchGeometry | None] = 'shear_periodic'#

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#: Legacy serialized one-hot cache; runtime geometry uses canonical alpha.

beta_axis: Array#: Legacy serialized one-hot cache; runtime geometry uses canonical beta.

alpha: int#: Index of the shear-flow coordinate.

beta: int#: Index of the shear-gradient coordinate.

gamma_dot: Array#: Current shear rate, used for image velocities. Keep it consistent with the externally imposed derivative of gamma; it does not advance strain.

rates_between(old: Domain, dt: Array) → Domain[source]#: Set gamma_dot and box_rate from the geometry change since old.

search_geometry_snapshot() → Array[source]#: Return box and canonical shear geometry used by search caches.

search_geometry_drift(previous: Array, old_pos: Array) → Array[source]#: Bound box and shear image motion, including unwrapped coordinates.

neighbor_cells(coords: Array, image: Array, grid_dims: Array, cell_size: Array) → Array[source]#: Cover both flow-axis cells touched by each fractional shear image.

classmethod Create(dim: int, box_size: Array | None = None, anchor: Array | None = None, gamma: float | Array = 0.0, alpha: int = 0, beta: int = 1, gamma_dot: float | Array = 0.0, box_rate: Array | None = None, **kwargs: Any) → LeesEdwardsDomain[source]#: Construct a Lees-Edwards domain with validated shear axes.

static shear(state: State, system: System, dgamma: float | jax.Array) → tuple[State, System][source]#

Apply an affine strain increment about the domain anchor.

Move body centers along alpha by dgamma * (r_beta - anchor_beta) and increment gamma by dgamma. Rigid-body offsets and orientations are preserved. Pair history is advanced once for the resulting relative motion, including that of periodic shear images, without time passing: velocities, dt, gamma_dot, and rate-dependent forces play no part. Caches rebuild when the strain change exhausts their skin. For quasistatic shear, minimize the returned state at its new strain before applying the next increment.

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\operatorname{round}(r_{ij}/L)\end{split}\]

where:

\(L\) is the domain box size (Domain.box_size)
Rounding uses ties to even at half-box displacements.

Return type:

jax.Array

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

Subtract the velocity of the same image used by displacement.

The image of j sits at rj + n * L + n_beta * gamma * L_beta e_alpha with integer image counts n. Its velocity is n * box_rate plus n_beta * (gamma_dot * L_beta + gamma * box_rate_beta) along alpha. Image counts use the supplied unwrapped particle centers, including when they span multiple boxes.