Potentials

jax_morph.physics.Potential #

Potential()

Interaction-potential protocol.

Implement total_energy(positions, state) -> scalar; forces follows via autodiff. Any custom energy (not just pairwise) can subclass this - a relaxation or a Brownian drift consumes only total_energy / forces.

Methods:

total_energy –

Evaluate total interaction energy at supplied positions.
forces –

Differentiate total energy into per-cell forces.
state_reads –

Declare non-base state fields read by the energy.

total_energy #

total_energy(positions, state)

Evaluate total interaction energy at supplied positions.

Parameters:

positions –

Cell positions of shape (capacity, n_space_dim).
state –

State supplying all non-position energy inputs.

Returns:

–

Scalar total interaction energy.

Raises:

NotImplementedError –

Always, until a subclass supplies an energy.

forces #

forces(state)

Differentiate total energy into per-cell forces at state.position.

Parameters:

state –

State supplying positions and any energy parameters.

Returns:

–

Force array of shape (capacity, n_space_dim).

state_reads #

state_reads()

State fields this energy consumes beyond the always-present base fields (default: none).

The wrapping steps (relaxation, Brownian, stress) merge this into their own state_reads so that a custom field a potential reads is allocated by build_state_from_model. Override to declare any such field; PairwisePotential already returns its from-state couplings.

Returns:

–

Tuple of additional StateFieldSpec objects. Defaults to an empty tuple.

jax_morph.physics.NoForce #

NoForce()

The trivial potential: no interactions, so every cell feels zero energy and zero force.

The neutral default for the position-moving steps (BrownianDynamics, ActiveBrownianDynamics2D, MechanicalRelaxation) when they are built without a potential - a free Brownian / active-Brownian gas, or a relaxation that leaves positions untouched. It overrides forces to return zeros directly, skipping the autodiff of total_energy and so also the dense O(N^2) pairwise sum: the right base for large-N active-matter / flocking models whose only coupling lives in a separate step (e.g. a neighbour-alignment torque).

total_energy #

total_energy(positions, state)

Always zero (no interactions).

forces #

forces(state)

Zero force on every cell (overrides the autodiff path - no energy to differentiate).

jax_morph.physics.PairwisePotential #

PairwisePotential()

A pair potential defined by pair_energy(r, *params) and per-pair pair_params(state).

Implement the elementwise pair_energy and the per-pair pair_params (a tuple of (N, N) arrays); total_energy, forces, and per-cell virial virial_pressure come for free - the standard way to add a soft-matter interaction (harmonic, Lennard-Jones, Hertzian, ...). Reductions go through neighbor_sum (the sparse seam), which masks self-pairs and dead cells.

Methods:

pair_energy –

Evaluate an elementwise energy for each pair separation.
pair_params –

Construct per-pair parameter matrices from a state.
mix –

Convert one per-cell coupling into a symmetric pair coupling.
virial_pressure –

Evaluate per-cell virial pressure.

pair_energy #

pair_energy(r, *params)

Evaluate elementwise pair energy at a separation.

Parameters:

r –

Pair separation array of shape (capacity, capacity).
*params –

Per-pair parameter arrays broadcastable with r.

Returns:

–

Pair-energy array of shape (capacity, capacity).

Raises:

NotImplementedError –

Always, until a subclass supplies the energy.

pair_params #

pair_params(state)

Construct per-pair parameter arrays from state.

Parameters:

state –

State supplying any field-based couplings.

Returns:

–

Tuple of arrays, each with shape (capacity, capacity).

Raises:

NotImplementedError –

Always, until a subclass supplies the parameterization.

mix #

mix(v)

Combine a per-cell coupling v (N,) into a per-pair coupling (N, N): the arithmetic mean.

Override for a different symmetric rule (e.g. geometric mean / Lorentz-Berthelot). The arithmetic mean has a finite gradient everywhere (no sqrt), so it stays safe at zeros and dead cells under debug_nans. The contact distance keeps its own additive rule sigma = r_i + r_j.

Parameters:

v –

Per-cell scalar coupling array of shape (capacity,).

Returns:

–

Symmetric pair coupling array of shape (capacity, capacity).

state_reads #

state_reads()

The per-cell state fields this potential sources its couplings from (the spec params).

total_energy #

total_energy(positions, state)

Sum pair energy over live, non-self pairs, counting each pair once.

Parameters:

positions –

Cell positions of shape (capacity, n_space_dim).
state –

State supplying liveness, space, radii, and any couplings.

Returns:

–

Scalar total pair energy.

virial_pressure #

virial_pressure(state)

Per-cell virial pressure p_i = -(1 / (2 d V_i)) sum_j r_ij (dU/dr)(r_ij).

The Irving-Kirkwood one-half bond split (each pair's virial is shared between its two cells); the minus sign makes repulsion (dU/dr < 0, compression) give p > 0. V_i is the cell's d-ball volume (2 r, pi r**2, 4/3 pi r**3 for d = 1, 2, 3). Dead cells score 0.

Parameters:

state –

State supplying cell geometry, liveness, space, and couplings.

Returns:

–

Per-cell pressure array of shape (capacity,).

jax_morph.physics.Morse #

Morse(
    epsilon: object = 3.0,
    alpha: object = 2.8,
    r_onset_frac: float = 1.5,
    r_cutoff_frac: float = 2.5,
)

Morse pair potential with a sigma-relative smooth cutoff.

With well depth \(\epsilon\), steepness \(\alpha\), and contact distance \(\sigma = r_i + r_j\) (the well minimum, where \(U = -\epsilon\)), the energy is

\[U(r) = \epsilon\left[\left(1 - e^{-\alpha (r - \sigma)}\right)^2 - 1\right],\]

multiplied by a smooth cutoff that turns the energy off between r_onset_frac * sigma and r_cutoff_frac * sigma. Larger alpha narrows the well. epsilon and alpha are each a shared scalar or a per-cell StateFieldSpec.

Attributes:

epsilon –

Well depth, scalar or per-cell StateFieldSpec. Defaults to 3.0.
alpha –

Well steepness (larger values narrow the well), scalar or per-cell StateFieldSpec. Defaults to 2.8.
r_onset_frac –

Smooth-cutoff onset as a multiple of contact distance. Defaults to 1.5.
r_cutoff_frac –

Smooth-cutoff end as a multiple of contact distance. Defaults to 2.5.

pair_params #

pair_params(state)

Return (sigma, epsilon, alpha) as (N, N) arrays; sigma = r_i + r_j is per pair.

pair_energy #

pair_energy(r, sigma, eps, alpha)

Morse energy at separation r times the sigma-relative smooth cutoff.

jax_morph.physics.SoftSphere #

SoftSphere(epsilon: object = 1.0)

Harmonic soft-sphere repulsion: a purely repulsive, compact excluded-volume interaction.

Writing \(\sigma = r_i + r_j\) for the contact distance, the pair energy is

\[U(r) = \frac{\epsilon}{2}\left(1 - \frac{r}{\sigma}\right)^2\]

for \(r < \sigma\) and zero at and beyond contact, so both energy and force vanish there (C1) - the canonical soft-disk / active-matter excluded-volume model. epsilon is a shared scalar or a per-cell StateFieldSpec; there is no cutoff to set.

Attributes:

epsilon –

Repulsion strength, scalar or per-cell StateFieldSpec. Defaults to 1.0.

pair_params #

pair_params(state)

Return (sigma, epsilon) as (N, N) arrays; sigma = r_i + r_j is per pair.

pair_energy #

pair_energy(r, sigma, eps)

Harmonic overlap energy, compact at the contact distance.

jax_morph.physics.Hertzian #

Hertzian(epsilon: object = 1.0)

Hertzian elastic-contact repulsion: purely repulsive and compact, with a soft onset.

With stiffness \(\epsilon\) and contact distance \(\sigma = r_i + r_j\), the pair energy is

\[U(r) = \frac{2}{5}\,\epsilon \left(1 - \frac{r}{\sigma}\right)^{5/2}\]

for \(r < \sigma\) and zero beyond. Like SoftSphere but softer at contact - both the force and its slope vanish there - modelling deformable elastic spheres / cells. epsilon is a shared scalar or a per-cell StateFieldSpec.

Attributes:

epsilon –

Contact stiffness, scalar or per-cell StateFieldSpec. Defaults to 1.0.

pair_params #

pair_params(state)

Return (sigma, epsilon) as (N, N) arrays; sigma = r_i + r_j is per pair.

pair_energy #

pair_energy(r, sigma, eps)

Hertzian contact energy (exponent 5/2), compact at the contact distance.

jax_morph.physics.Harmonic #

Harmonic(k: object = 1.0, r_cutoff_frac: float = 2.5)

Finite-range harmonic spring: a shifted parabola with its minimum at contact.

Writing \(k\) for the stiffness, \(\sigma = r_i + r_j\) for the contact distance, and \(r_c\) for the cutoff r_cutoff_frac * sigma, the pair energy is

\[U(r) = \frac{k}{2}\left[(r - \sigma)^2 - (r_c - \sigma)^2\right]\]

for \(r < r_c\) and zero beyond - the parabola shifted down so it vanishes at the cutoff. It is a well of depth \(\frac{k}{2}(r_c - \sigma)^2\) minimised at contact: repulsive when compressed, adhesive when stretched. The energy is only C0 at the cutoff (the force jumps there), harmless because the cutoff sits well beyond the resting contact distance. k is a shared scalar or a per-cell StateFieldSpec.

Attributes:

k –

Spring stiffness, scalar or per-cell StateFieldSpec. Defaults to 1.0.
r_cutoff_frac –

Cutoff distance as a multiple of contact distance. Defaults to 2.5.

pair_params #

pair_params(state)

Return (sigma, k) as (N, N) arrays; sigma = r_i + r_j is per pair.

pair_energy #

pair_energy(r, sigma, k)

Shifted harmonic well: negative inside the cutoff, 0 at and beyond it (truncated).

jax_morph.physics.LennardJones #

LennardJones(
    epsilon: object = 1.0,
    r_onset_frac: float = 1.5,
    r_cutoff_frac: float = 2.5,
)

Lennard-Jones potential (r_min form): minimum -epsilon at contact, with a sigma-relative cutoff.

With well depth \(\epsilon\) and contact distance \(\sigma = r_i + r_j\), the pair energy is

\[U(r) = \epsilon\left[\left(\frac{\sigma}{r}\right)^{12} - 2\left(\frac{\sigma}{r}\right)^{6}\right],\]

a hard \(r^{-12}\) core plus an \(r^{-6}\) adhesive tail with its minimum \(-\epsilon\) exactly at contact. A smooth cutoff truncates the tail between r_onset_frac * sigma and r_cutoff_frac * sigma. epsilon is a shared scalar or a per-cell StateFieldSpec.

Attributes:

epsilon –

Well depth, scalar or per-cell StateFieldSpec. Defaults to 1.0.
r_onset_frac –

Smooth-cutoff onset as a multiple of contact distance. Defaults to 1.5.
r_cutoff_frac –

Smooth-cutoff end as a multiple of contact distance. Defaults to 2.5.

pair_params #

pair_params(state)

Return (sigma, epsilon) as (N, N) arrays; sigma = r_i + r_j is per pair.

pair_energy #

pair_energy(r, sigma, eps)

LJ energy (r_min form) times the sigma-relative smooth cutoff; safe at r=0.