A differentiable deep-brain-stimulation (DBS) field solver built on
NVIDIA Warp. It solves the quasi-static
volume-conductor equation ∇·(σ∇φ) = −q around a 4-contact DBS lead in
inhomogeneous brain tissue with a matrix-free preconditioned conjugate-gradient
solver written entirely in Warp kernels, then backpropagates through the
physics (wp.Tape) to optimize the four contact currents so the stimulated
region covers a target nucleus while sparing an avoidance region — under a
hard Shannon charge-density safety envelope.
Every kernel on the hot path (7-point stencil, CG updates, E-field gradient,
soft-VTA objective) is the same Warp source on CPU and CUDA. The entire
validation suite and demo run and are verified on CPU (this repo was built on
an Apple-silicon Mac); on an NVIDIA GPU, Warp compiles the identical kernels
to CUDA — device="cuda" is the only change.
Synthetic inhomogeneous anatomy (gray/white/CSF/encapsulation at literature conductivities), Medtronic-3389-like lead, STN-like ellipsoidal target, 96³ grid at 0.25 mm:
| currents (4 contacts) | target coverage | avoid-region overlap |
|---|---|---|
| initial guess (0.9/0.9/0.6/0.6 mA) | 3.2% | 10.9 mm³ |
| equal split of same total (2×3.99 mA equiv.) | 45.7% | 66.6 mm³ |
| optimized (projected Adam, 250 steps) | 74.1% | 27.9 mm³ |
The optimizer lands exactly on the safety boundary: the two active contacts
sit at 30 µC/cm²/phase — the classic Shannon limit (k = 1.75) — verified as a
constraint, not a coincidence (tests/test_safety.py).
The PDE is linear in the injected currents, so
E(c) = Σₖ cₖ·Eₖ where the unit-current basis fields Eₖ come from one CG
solve per contact (4 solves, ~25 s each on CPU, cached). The optimization
loop never re-solves the PDE: wp.Tape records only the cheap combination
step (E = Σ cₖEₖ → |E| → soft-VTA objective) and backpropagates to the four
scalar currents — 137 ms per gradient step instead of ~100 s. Superposition
is verified to machine precision: recombined basis fields match a direct solve
to 6.6 × 10⁻¹² relative (tests/test_superposition.py).
- Analytic — point source in a homogeneous medium vs the exact
φ = I/(4πσr), within a few percent in the 2–8 mm shell. - Interface physics — harmonic-mean face conductivities preserve normal
current continuity
σ₁E₁ₙ = σ₂E₂ₙacross tissue boundaries. - Grid convergence — shell error decreases at order ≥ 1.5 under refinement (second-order stencil, point source on one voxel).
- Superposition — 6.6e-12 relative, above.
- Gradient check —
wp.Tapeadjoint vs central finite differences, agreement ~1e-8 (all-float64 pipeline). - Safety envelope — charge-density projection is enforced and binding.
pip install warp-lang numpy matplotlib pytest
./run.sh test # 10 validation tests, ~75 s CPU
./run.sh demo # basis solves + optimization, figures + metrics to results/
notebooks/cuda_verify.ipynb runs the full validation suite and the steering
demo on a free Colab T4. warpfield.solver picks the default Warp device, so
on a CUDA runtime every kernel executes on the GPU unchanged; the notebook
asserts the device, reruns all 10 tests, and prints a verification verdict.
- The VTA is a field-threshold proxy (|E| ≥ 0.2 V/mm, standard in the DBS modeling literature, cf. Åström et al. 2015) — not multi-compartment axon-cable modeling. No fiber orientation, diameter distribution, or pulse-shape dependence; "activation" numbers are a field-strength heuristic.
- The anatomy is synthetic (procedural gray/white/CSF/encapsulation with literature conductivities, isotropic, quasi-static). Nothing here is patient-specific or clinical; this is a methods/systems artifact about differentiable PDE optimization on Warp, not a medical device.
- CUDA execution is untested in this repo (no NVIDIA GPU on the build machine); what is verified is that every hot-path kernel is pure Warp with no CPU-only constructs, and Warp's programming model makes the CPU/CUDA switch a device string.
Warp gives this project the three things it actually needs: kernels in Python
with C-like performance, one source for CPU and CUDA, and built-in adjoints
(wp.Tape) so the safety-constrained optimization is gradient-based rather
than a grid search over currents. The same pattern — differentiable simulation
for treatment planning — is where GPU-accelerated computational medicine is
heading, and Warp is the most direct open-source path to it.
Åström et al., IEEE TBME 2015 (E-field threshold VTA) · Shannon 1992 (charge-density safety) · Butson et al. 2006 (encapsulation impedance) · Baumann et al. 1997 (CSF conductivity) · NVIDIA Warp (Macklin, 2022).
MIT © 2026 Azra Bano
- disagg-sim — discrete-event simulator for disaggregated LLM serving and KV-aware routing
- llm-roofline — first-principles roofline model of LLM inference on NVIDIA GPUs