waxmorph.jax.warp_autograd.warp_diffusion_step#
- waxmorph.jax.warp_autograd.warp_diffusion_step(c, pair_i, pair_j, D_emu, dt, *, num_pairs=None, device='cuda')[source]#
Apply one differentiable graph-Laplacian signaling-molecule diffusion step.
Each frozen pair contributes antisymmetric flux \(\phi_{ij,g}=c_{j,g}-c_{i,g}\). JAX scatter-adds these values into \((Lc)_{i,g}=\sum_{j\in\mathcal N(i)}(c_{j,g}-c_{i,g})\), then applies
\[c_{i,g} \leftarrow \max\!\big(c_{i,g} + \mathrm{dt}\,D\,(L c)_{i,g},\; 0\big)\]where \(D\) is
D_emu. The custom VJP differentiates concentrations through the per-pair flux, while pair selection remains fixed. At the JAX output clamp, the derivative with respect to the pre-clamp value is zero below zero, one above zero, and one half at exactly zero.num_pairsmasks padded indices beyond the real frozen-pair count; omitted means every pair is real. An empty pair buffer returnscunchanged.