impact-adjoint¶
Exact gradients through impact events, across a Julia and JAX Tesseract boundary.
Simulations of contact are easy to write and hard to differentiate. impact-adjoint computes parameter sensitivities across impact events that are exact for the hybrid system at fixed event topology, using the classical saltation matrix, and serves them from a Julia solver through a Tesseract component boundary, so a JAX program obtains the true derivative without doing any event handling of its own. The resulting gradients design passive structures whose function exists only because of impacts.
jax.grad returns
d x(T)/d v0y = 0.0 exactly on flat terrain.
The true value is +0.09.
The gradient is not noisy or approximate. It is confidently, silently wrong, and refining the step size does not help, because the step at which the event fires is an integer and autodiff cannot differentiate it.
Demo¶
results.json, and the
terminal segment is a real run captured when the video is built.
See it fail, then see it fixed¶
Both snippets below run as written. What differs is where the impact is applied; the second also enables x64, because the component returns Float64 and JAX refuses to hand those back without it.
The impact is applied at the integrator step where it is detected.
def bounce(v0y, n=2000):
dt = 2.0 / n
def step(q, _):
qn = rk4(q, dt)
hit = (qn[1] < 0) & (q[1] > 0)
qn = jnp.where(hit, reset(qn), qn)
return qn, None
q0 = jnp.array([0., 1., 2., v0y])
qf, _ = jax.lax.scan(
step, q0, None, length=n)
return qf[0]
jax.grad(bounce)(0.5)
# 0.0
The step at which the event fires is an integer. Autodiff differentiates a
staircase, and returns zero at every dt. Hand-writing an interpolated event
inside the program does recover a converging gradient, to within 0.05% at
this step size and 0.3% across the whole sweep;
what the component removes is the derivation, not the last three digits.
The impact is applied at the crossing, and its time carries a derivative.
jax.config.update("jax_enable_x64", True)
API = ("tesseracts/contact_sim"
"/tesseract_api.py")
sim = Tesseract.from_tesseract_api(API)
def bounce(v0y):
cfg = {**CFG,
"v0": jnp.array([2.0, v0y])}
out = apply_tesseract(sim, cfg)
return out["qf"][0]
jax.grad(bounce)(0.5)
# 0.0903777362581133
The saltation matrix supplies the missing term at each impact. Exact at any
dt, with no event handling in the JAX program.
Move the slider¶
Left is coarse, right is fine: the slider runs from
4.00e-3 down to
6.25e-5. All three columns are measured at
the step size shown, including saltation, which is solved afresh at each
one rather than carried over. Grid reset never leaves zero; the
interpolated repair wanders; saltation moves by
6.02e-14 across the whole sweep, against a
reference value of 0.09037774. Read from
experiments/e3_rows.npy. The sweep is flat-terrain ballistic
flight, where RK4 is exact by construction; that is what isolates the
event term from integrator error.
What the gradients build¶
- 3.99
- observed order of accuracy, against an analytic solution
- 7.1e-12
- multi-bounce agreement with a symbolic closed form
- 0 / 150
- failures in Tesseract's gradient checker across its three gradient endpoints, at 10-4 relative tolerance
- 4.6e-9
- analytic Jacobian vs an independent scipy reimplementation
- 26 s
- to re-run the four verification checks, warm
Reproduce it¶
CPU only, no Docker, no GPU. The four checks take about 26 seconds on a
warm Julia depot, give or take what else the machine is doing. The first run
adds a few minutes of pip install and a Julia bootstrap that runs in two
stages, one on the first call and one the first time ForwardDiff is needed,
so budget nearer fifteen minutes from a genuinely cold machine.
git clone https://github.com/singhharsh1708/impact-adjoint
cd impact-adjoint
python3 -m venv .venv && source .venv/bin/activate
pip install -r docs/requirements-repro.txt
python scripts/proof_local.py # boundary proof
python scripts/validate_closed_form.py # symbolic oracle
python scripts/validate_reference.py # scipy oracle
python experiments/e3_naive_vs_saltation.py
The last command is the claim at the top of this page: it prints the
grid-reset gradient as exactly 0.0 at every step size, next to the
saltation gradient. Full instructions are in
getting started.
Where to start¶
Why autodiff returns a wrong gradient at an impact, what the pure-JAX repair costs, and what Diffrax can and cannot do about it today.
The saltation jump condition, the component boundary that carries it, and the termination semantics that keep optimizers honest.
Seven verification studies and eight experiments, every number generated from a committed artifact.
How it is built¶
Item |
Detail |
|---|---|
Track |
1, inverse design and shape optimization |
Components |
|
Gradient method |
forward variational sensitivities with analytic saltation updates at events |
Checked against |
a symbolic closed form and an independent scipy reimplementation, plus golden regression tests on every push |