Differentiation test suite and tracer leak fix

[astanziola]

Differentiation test suite and tracer leak fix

Summary

• Add a differentiation test suite (tests/test_differentiation.py) verifying jax.grad works through all solvers w.r.t. all differentiable inputs
• Fix tracer leak in laplacian_with_pml that caused helmholtz_solver with checkpoint=False to fail under jax.grad

Tracer leak fix

FourierSeries laplacian_with_pml lazily computed gradient.default_params(rho0) inside the function body (line 224 of operators.py). When GMRES re-traced helm_func for its adjoint with checkpoint=False, this created new arrays inside the custom_linear_solve trace scope that escaped — causing an UnexpectedTracerError.

Fix: moved fft_rho0 computation into fourier_laplacian_with_pml_init so it's pre-computed before tracing. Also added init_params to the OnGrid variant for consistency.

Verified bitwise identical results (forward field, loss, gradient) before and after the fix.

Known issues documented as xfail

• born_series uses lax.while_loop which doesn't support reverse-mode AD (CBS adjoint #189)

Related issues

• Incorrect gradients for 3D Helmholtz simulations using checkpointing #150 — Incorrect gradients for 3D Helmholtz with checkpointing
• NaNs when batching simulations in tutorial example. #168 — NaNs when batching simulations
• Allow float64 for helmholtz solver #187 — Allow float64 for Helmholtz solver
• CBS adjoint #189 — CBS adjoint

[codecov]

Codecov Report

❌ Patch coverage is 46.66667% with 8 lines in your changes missing coverage. Please review.
⚠️ Please upload report for BASE (main@7e16754). Learn more about missing BASE report.
⚠️ Report is 1 commits behind head on main.

Files with missing lines	Patch %	Lines
jwave/acoustics/operators.py	46.66%	8 Missing ⚠️

Additional details and impacted files

@@           Coverage Diff           @@
##             main     #243   +/-   ##
=======================================
  Coverage        ?   80.42%           
=======================================
  Files           ?       15           
  Lines           ?     1216           
  Branches        ?        0           
=======================================
  Hits            ?      978           
  Misses          ?      238           
  Partials        ?        0           

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.