waxmorph.jax.warp_autograd.warp_mech_step

Contents

waxmorph.jax.warp_autograd.warp_mech_step#

waxmorph.jax.warp_autograd.warp_mech_step(X, R, pair_i, pair_j, dt, *, num_pairs=None, device='cuda')[source]#

Apply one differentiable sticky-sphere mechanics step.

For a frozen pair \((i,j)\), define raw displacement \(\mathbf d=\mathbf x_i-\mathbf x_j\), stabilized distance \(r=\lVert\mathbf d\rVert+\epsilon_n\), regularized direction \(\mathbf u=\mathbf d/r\), and \(s=R_i+R_j\):

\[\begin{split}f_{rep} &= k_{rep}\,\max(s - \epsilon_d - r,\; 0) \\ f_{att} &= k_{att}\,\max(s + \epsilon_d - r,\; 0) \cdot \mathbb{1}[r > s - \epsilon_d] \\ \mathbf{F}_{i} &= (f_{rep} - f_{att})\,\mathbf{u}, \quad \mathbf{F}_{j} = -\mathbf{F}_{i}\end{split}\]

Here \(k_{rep}=2\), \(k_{att}=0.5\), \(\epsilon_d\) is EPS_DIST, and \(\epsilon_n\) is EPS_NORM. JAX scatter-adds endpoint forces and applies \(\mathbf x\leftarrow\mathbf x+\mathrm{dt}\,\mathbf F\).

The custom VJP differentiates X and R through per-pair forces while pair selection remains a fixed forward input. max has zero derivative on inactive branches; the strict adhesion indicator is constant in backward. Values at their nonsmooth boundaries follow Warp’s branch convention.

pair_i and pair_j are unordered frozen endpoints. For a padded pair buffer, num_pairs masks indices beyond the real count; omitted means all pairs are real. An empty buffer returns X unchanged.

Raises:

RuntimeError – If CUDA-backed Warp/JAX device validation fails.

Parameters:
Return type:

Array

Contents