Relaxation

jax_morph.physics.MechanicalRelaxation

MechanicalRelaxation(
    potential, *, max_steps=500, f_tol=0.001, ridge=1e-06
)

Quasistatic step: relax positions to mechanical equilibrium under potential each step.

Each step drives positions to a force balance \(\nabla_x U = 0\) with FIRE and differentiates that equilibrium by the implicit function theorem, not the solver path.

Gradient behaviour (easy to misread): the equilibrium sensitivity is taken on the physical (deformation) subspace only. The rigid-body null modes of a free-space equilibrium - the dead cells' degrees of freedom, a global translation, and a global rotation of the alive cells - are projected out of the adjoint and carry no gradient at all. This is deliberate: the energy is flat along those gauge modes (the relaxed cluster's absolute position and orientation are fixed by the initial condition, not the parameters). A shape objective is invariant to them and unaffected; a translation- or rotation-sensitive objective simply gets no gradient for that sensitivity (rather than a spurious ridge-scaled one). To make the cluster's position or orientation optimizable, break the symmetry physically (pin a cell, add an external field / substrate).

Attributes:

• potential –

  Interaction potential defining the equilibrium.

• max_steps –

  Static maximum number of FIRE iterations. Defaults to 500.

• f_tol –

  Maximum absolute force tolerated at convergence. Defaults to 1e-3.

• ridge –

  Regularization for residual soft modes in the implicit adjoint. Defaults to 1e-6.

Pass potential=None for NoForce, whose every configuration is already an equilibrium, so relaxation is a no-op (positions pass through unchanged).

Parameters:

• potential –

  Interaction potential, or None for NoForce.

• max_steps –

  Static maximum number of FIRE iterations. Defaults to 500.

• f_tol –

  Maximum absolute force tolerated at convergence. Defaults to 1e-3.

• ridge –

  Regularization for residual soft modes in the implicit adjoint. Defaults to 1e-6.

state_reads

state_reads()

Reads positions (the FIRE start), plus any state field the potential sources params from.

state_writes

state_writes()

Writes the relaxed positions.

---

jax_morph.physics.relax_equilibrium

relax_equilibrium(
    potential,
    state,
    *,
    max_steps=500,
    f_tol=0.001,
    ridge=1e-06,
)

Relax positions to mechanical equilibrium, with implicit-diff (not solver-path) gradients.

The forward runs FIRE to a genuine force tolerance |grad U| <= f_tol (a real equilibrium, not a fixed step count). The backward is the implicit-function-theorem sensitivity of that equilibrium, restricted to the physical subspace: the rigid-body null modes (dead-cell DOFs, a global translation, and a global rotation of the alive cells) are projected out, and a tiny ridge regularizes only residual near-zero soft modes. Differentiates w.r.t. both the potential's parameters and the incoming state (radii -> sigma, ...).

If FIRE hits max_steps before reaching f_tol it emits a runtime RuntimeWarning and returns the last iterate anyway (behaviour is otherwise unchanged); raise max_steps or f_tol to silence it.

Pass potential=None for NoForce, whose every configuration is already an equilibrium, so the positions are returned unchanged (matching MechanicalRelaxation).

Parameters:

• potential –

  Interaction potential, or None for NoForce.

• state –

  Input state whose positions provide the FIRE initial condition.

• max_steps –

  Static maximum number of FIRE iterations. Defaults to 500.

• f_tol –

  Maximum absolute force tolerated at convergence. Defaults to 1e-3.

• ridge –

  Regularization for residual soft modes in the implicit adjoint. Defaults to 1e-6.

Returns:

• –

  Relaxed position array with shape (capacity, n_space_dim).