fdmsim · roadmap · research proposal

Differentiable Adaptive-Mesh Schrödinger–Poisson for Fuzzy Dark Matter

A solver that is both differentiable and adaptive — enabling gradient-based inference of the fuzzy-dark-matter boson mass m directly from resolved solitonic cores. It unifies our two codes: JAXiON's differentiability and GAMER's AMR resolution, in one engine.

proposal · 2026-06-28 phases · 0 → 3 total · ≈ weeks (AI-time) compute · idle CPU → 1–8 GPU
Proposed

PROPOSED · not yet started. This is a forward-looking research plan, not active work. It is held until the current GAMER / JAXiON / GADGET-4 pipelines finish. Nothing on this page has been executed. Each phase below carries an explicit go/no-go gate; the program advances only if Phase 0 passes.

A solver that is both differentiable and adaptive — so the boson mass can be inferred, with gradients, directly from cores we can finally resolve.

the right hard problem

Why this is worth doing

positioning

Novelty — what exists vs. the gap

what exists in the literature vs. the gap this work fills
ExistsGap (this work)
Differentiable AMR in JAX — JANC / JAX-AMR (2025), reacting CFDnot for Schrödinger–Poisson
Differentiable cosmology — DISCO-DJ, pmwd, FlowPM, jFoFnone do FDM / wave dark matter
Adaptive-wavelet Schrödinger solversnot differentiable, not FDM / cosmology
GAMER ELBDM AMRnot differentiable
→ The first differentiable, adaptive Schrödinger–Poisson solver for FDM. A methods paper (filling the DISCO-DJ / pmwd gap) plus new science.
technical approach

How it works

The forward model is the Schrödinger–Poisson (Gross–Pitaevskii) system for the FDM wavefunction $\psi$:

$i\hbar\,\partial_t \psi = -\dfrac{\hbar^2}{2m}\nabla^2\psi + m\,\Phi\,\psi, \qquad \nabla^2 \Phi = 4\pi G\,\big(|\psi|^2 - \bar\rho\big).$

Core building blocks — all individually de-risked in the literature:

  1. Differentiable split-step. JAXiON already has this on a uniform grid: alternate kinetic half-steps in Fourier space with potential kicks in real space, each step a differentiable operation.
  2. Fixed-shape multi-level blocks + masking. Solves XLA's static-shape constraint (the JAX-AMR pattern): no dynamic allocation — nested levels are dense arrays with active-masks rather than ragged refinement trees.
  3. Soft (differentiable) refinement. A sigmoid threshold replaces the hard refine / don't-refine decision; or a frozen-mesh adjoint where the mesh is a non-differentiated structural parameter.
  4. Per-block solver. Local spectral (GRAMFE-style) or finite-difference, with level-interface buffers to suppress reflections at refinement boundaries.
  5. Adjoint-through-time + gradient checkpointing. Memory-efficient reverse-mode autodiff over many steps, proven in pmwd.

Architecture options — evaluated, not committed blindly

ABlock-structured differentiable AMR

JANC-style. The most direct route: port the proven CFD pattern to Schrödinger–Poisson.

primary path
BDifferentiable adaptive-wavelet

$\psi$ in a wavelet basis; smooth coefficient thresholding gives differentiable refinement; multiresolution is intrinsic.

evaluated · Phase 1.5
CLearned (GNN / RL) refinement

Refinement policy learned through the differentiable solver. Deferred until A or B works.

deferred
DStatic nested grids in JAXiON

Fixed high-resolution sub-boxes. The de-risking stepping-stone.

stepping-stone · Phase 0/1
phase-wise plan

The game plan, gated phase by phase

On "AI-time." AI-time is agent execution time if the science cooperates; research risk is separate and is exactly what each gate tests. All phases run isolated on idle compute. This is not the 3–6 months a human team would need — it is also not a promise the physics works. Each phase has a concrete go/no-go; we do not advance on hope.

Phase 0KILL-TEST — the gate~hours – 1 day
Objective
Answer the single question that decides the whole program — do usable gradients flow through adaptive refinement?
Method
A 1D Schrödinger–Poisson split-step in JAX; a 2-level static nested grid with a soft refinement mask; compute $\partial(\text{observable})/\partial m$ through the adaptive solver; compare to a high-uniform-resolution reference; attempt to invert for $m$ from a mock.
Success
(a) the adaptive solver matches the uniform reference on $\rho_c$ / $r_c$ to within a few %; (b) gradients $\partial\text{obs}/\partial m$ agree with finite-difference to $\sim 10^{-3}$ and are smooth (not dominated by refinement-boundary noise); (c) a 1-parameter inversion for $m$ recovers truth.
Kill if
gradients through refinement are too noisy or biased to invert. Then: AMR-in-JAXiON is not viable; the two-code bridge stays the answer — and we have learned that cheaply.
Phase 13D differentiable block-AMR soliton~days
Objective
Show in 3D what uniform JAXiON cannot — a converged core ($r_c/\mathrm{d}x \geq 8$) with working gradients.
Method
JANC-style fixed multi-level blocks around a single soliton or binary merger; per-block local-spectral or FD split-step; soft refinement; differentiate with respect to $m$ and the soliton parameters.
Success
core converges ($\rho_c$ plateaus, $r_c/\mathrm{d}x \geq 8$); $\partial\text{obs}/\partial m$ valid; runtime and memory acceptable on one GPU.
Phase 1.5
(optional, ~days) prototype the wavelet variant (B) on the same test; compare accuracy, cost, and gradient quality vs. block-AMR.
Phase 2Cosmological AMR + adjoint-through-time long pole · real risk~weeks
Objective
A cosmological box with AMR-resolved cores, fully differentiable end-to-end.
Method
Add comoving evolution + Poisson (FFT base + multilevel correction); adjoint + reverse-time + checkpointing across the run; refine on collapsing halos; infer $m$ from a resolved core–halo system.
Success
matches GAMER on the core–halo relation (validation); beats uniform JAXiON on core resolution; the end-to-end gradient $\partial(\text{core–halo obs})/\partial m$ is usable for inference.
Phase 3Science — differentiable inference on the tension, at core level~days – weeks
Objective
The payoff — gradient-based / SBI inference of $m$ (and self-interaction, core–halo scatter) from resolved cores, applied to the dwarf–Lyman-$\alpha$ tension.
Method
Plug the differentiable adaptive cores into the unified inference framework (the tension analysis); marginalize the core–halo scatter the solver now resolves; run a controlled CDM–FDM twin.
Output
a differentiable, core-resolved constraint on $m$ + a principled statement on the tension. The original scientific result.
risk register

Risks & kill criteria

#1
Gradients through refinement too noisy. Phase 0 tests this first, cheaply. hard kill gate.
#2
Level-interface artifacts — reflections at refinement boundaries corrupting the wavefunction. interface buffers; validated against the uniform reference each phase.
#3
Adjoint memory blow-up over many steps. gradient checkpointing (pmwd-proven); cap steps if needed.
#4
Wavelet–Poisson coupling complexity (path B). pursued only if block-AMR (A) underperforms.
#5
Scope creep. Each phase has a concrete go/no-go; we do not advance on hope.
compute & timeline

Footprint and schedule

Compute. Phases 0–1: idle CPU or a single GPU (cheap). Phase 2: 1–8 GPUs for the cosmological adjoint runs. Far below the GAMER campaign's footprint.

Phase 0
hours – 1 day
the kill-test gate
Phase 1
days
3D converged core + gradients
Phase 2
weeks
cosmological adjoint · long pole
Phase 3
days – weeks
science: inference on the tension

Total ≈ weeks of agent execution (AI-time, execution-if-it-works), gated phase-by-phase by go/no-go — not the "3–6 months" a human team would need.

references

References

JANC / JAX-AMR (CPC 2025) · DISCO-DJ (arXiv 2311.03291) · pmwd + adjoint (2211.09958, 2211.09815) · jFoF (2510.26851) · adaptive-multiresolution Schrödinger (2007.01471) · G-Adaptive GNN + difFEM (2407.04516) · Schive 2014 (FDM solitons) · our convergence-impossibility & dwarf–Lyman-$\alpha$ tension results.