MANGO

MANGO: A Neutrino Gradient Oscillator

MANGO (A Neutrino Gradient Oscillator) is an autodiff-first neutrino oscillation calculator in JAX.

mango computes neutrino oscillation probabilities in vacuum, constant-density matter, non-constant density (the Earth via the PREM model, with an atmospheric production height), arbitrary user profiles, and the Sun (adiabatic MSW). Everything is jit-compilable, vmap-batchable, and differentiable end-to-end — with respect to the oscillation parameters and the geometry (zenith angle, baseline, energy). Non-standard interactions (NSI) and sterile neutrinos (3+N) are first-class.

It is validated to machine precision against the established codes OscProb and NuFast, and reproduces the reference plots of the nu-waves library.

Atmospheric oscillograms through the PREM Earth: ν/ν̄ × (μ→e, μ→μ) over (energy, cos θ_z), showing the MSW resonance, the matter ν/ν̄ asymmetry, the core–mantle boundary at cos θ_z ≈ −0.84, and near-horizon atmospheric oscillations.

---

Why another oscillation code?

The established fast codes (NuFast, OscProb, Prob3++, …) are excellent at forward evaluation but are not differentiable. mango is built autodiff-first: you get exact gradients of any probability with respect to any input for free, which is what you want inside a JAX fitting / HMC / neural-network pipeline. The cost of all six oscillation-parameter gradients is ~3–5× a single forward evaluation (vs ~12× noisy evaluations for finite differences), and the same code runs batched on GPU/TPU.

---

Showcase: oscillograds

An oscillograd (arXiv:2512.16427) is a map of the gradient of an oscillation probability with respect to a parameter, ∂P/∂ζ — it shows where in parameter space a probability is most sensitive to ζ. This is exactly what a differentiable simulator is for: jax.jacrev returns the oscillograd for every parameter in a single reverse-mode pass.

Atmospheric oscillograds of P(νμ→νe) through the PREM Earth: the probability and its exact derivatives w.r.t. δ_CP, θ₂₃, θ₁₃ and Δm²₃₁ over (E, cos θ_z). θ₁₃ drives the MSW resonance (~6 GeV), θ₂₃ sets the appearance amplitude, δ_CP sensitivity sits at low energy — and the core–mantle boundary at cos θ_z ≈ −0.84 is visible in the gradients. Computed through fully differentiable matter propagation.

examples/oscillograds.py generates these (and the DUNE-baseline constant-density version, paper-style) and verifies every gradient against central finite differences — agreement is ~1e−9 (constant density) and ~1e−8 (PREM Earth), for all six parameters.

You can also differentiate w.r.t. the Earth model itself — the two-zone electron fraction Y_e (∝ electron density N_e = ρ Y_e N_A) is a traced input of probability_earth(..., ye_core=, ye_mantle=):

∂P/∂Y_e^core (center) is exactly zero for mantle-only trajectories and turns on only below the core-crossing threshold cos θ_z ≈ −0.84 — the autodiff gradient correctly "knows" the chord geometry — while ∂P/∂Y_e^mantle (right) is non-zero for all up-going paths.

For full control, a parametric LayeredEarth makes the shell boundary radii, densities and Y_e all differentiable inputs:

model = mango.prem_layered(n_sub=3)   # constant-density-shell Earth sampled from PREM
P = mango.probability_earth(params, E, cz, earth_model=model)
jac = jax.jacrev(lambda m: mango.probability_earth(params, E, cz, earth_model=m,
                          flavor_in=Flavor.MU, flavor_out=Flavor.E))(model)
# jac.outer -> dP/d(boundary radii), jac.density -> dP/d(shell densities), jac.ye -> ...

Differentiating w.r.t. the Earth structure itself: ∂P/∂R_cmb (the core–mantle boundary position, center) and ∂P/∂ρ for a deep core shell at 1735 km (right). Each is non-zero only for trajectories that reach that radius — the boundary sensitivity below cos θ_z ≈ −0.84, the 1735 km shell only below ≈ −0.96 — so the gradients exactly trace the radial geometry. (Matches finite differences to ~1e−8.)

Application: DUNE long-baseline sensitivity

analyses/dune/ reproduces the DUNE TDR oscillation analysis using mango for the oscillation probabilities and DUNE's official GLoBES configuration (arXiv:2103.04797) for the flux, cross-sections, per-channel migration matrices, efficiencies, and systematics. The far-detector spectra match the reference essentially exactly, and the CP-violation and mass-ordering sensitivities reproduce the DUNE curves (correct shapes, peak positions, and zeros) — with the χ² profiled over ~15 parameters using the exact gradient from mango (jax.value_and_grad → L-BFGS). See its README.

Installation

git clone https://github.com/pgranger23/mango-osc.git
cd mango-osc
pip install -e .                 # core: jax + numpy
pip install -e ".[examples]"     # + matplotlib for the example scripts
pip install -e ".[diffrax]"      # + diffrax for the optional stiff ODE backend
pip install -e ".[test]"         # + pytest

Float64 is required for oscillation phases (Δm² ~ 1e−3 eV²) and is enabled automatically on import mango.

---

Quick start

import jax, jax.numpy as jnp
import mango
from mango import nufit_no, Flavor

params = nufit_no()                       # a NuFIT-like normal-ordering point
E  = jnp.linspace(1.0, 25.0, 200)         # GeV
cz = jnp.linspace(-1.0, -0.05, 150)       # cos(zenith), up-going

# Atmospheric oscillogram P(νμ → νe) through the Earth, shape (n_cz, n_E)
P = mango.probability_earth(params, E, cz,
                            flavor_in=Flavor.MU, flavor_out=Flavor.E)

# Exact gradient w.r.t. every oscillation parameter (returns an OscParams of grads)
def loss(p):
    return mango.probability_earth(p, jnp.asarray(4.0), jnp.asarray(-1.0),
                                   flavor_in=Flavor.MU, flavor_out=Flavor.E)
grads = jax.grad(loss)(params)            # grads.theta23, grads.dm31, grads.deltacp, ...

Probability matrices have shape (..., N, N) indexed P[β, α] = P(ν_α → ν_β); pass flavor_in / flavor_out to select one channel. Antineutrinos: anti=True. Flavor.E, Flavor.MU, Flavor.TAU = 0, 1, 2 (sterile flavors are indices ≥ 3).

The five regimes

# vacuum
mango.probability_vacuum(params, E, baseline_km=295.0)                      # T2K-like

# constant-density matter
mango.probability_constant(params, E, baseline_km=1300.0, density=2.8, ye=0.5)  # DUNE-like

# PREM Earth (+ optional atmosphere); single entrypoint
mango.probability_earth(params, E, cz,
                        ye_core=0.466, ye_mantle=0.494,   # configurable two-zone Y_e
                        h_atm_km=15.0,                    # production height -> full cosz
                        n_sub=4, det_depth_km=0.0)

# arbitrary piecewise-constant profile (source -> detector ordering)
mango.probability_profile(params, E, density_gcc=[...], ye=[...], length_km=[...])

# solar adiabatic MSW: vacuum mass-state fractions vs radius
prof = mango.solar.load_bs05("examples/data/bs05_agsop.dat")
F = mango.solar.adiabatic_mass_fractions(params, E_GeV, prof, r_km, r_emit_km)

NSI and sterile neutrinos

# matter NSI: an epsilon_{alpha beta} object on any probability function
P = mango.probability_constant(params, E, 1300.0, density=2.8,
                               nsi=mango.NSI(eps_emu=0.1 + 0.05j, eps_ee=0.05))

# 3+1 sterile (flavors e, mu, tau, s); the matter potential automatically gives
# sterile flavors only the relative neutral-current term; backend auto-switches to eigh
st = mango.Sterile3plus1(theta12, theta13, theta23, theta14, theta24, theta34,
                         delta13, delta24, dm21, dm31, dm41)
P = mango.probability_earth(st, E, cz)            # (n_cz, n_E, 4, 4)

# arbitrary 3+N via the generic builder
U = mango.pmns_nflavor(5, [(3, 4, 0.1), (2, 4, 0.05), ..., (0, 1, theta12)])
params5 = mango.NFlavorParams(U=U, msq=jnp.array([0, dm21, dm31, dm41, dm51]), n_active=3)

Both NSI and sterile parameters are differentiable leaves — jax.grad works through NSI(...) and Sterile3plus1(...).

Decoherence and non-unitary mixing

# Lindblad decoherence Gamma_ij = gamma_ij (E/E0)^n, or wave-packet separation
P = mango.decoherence.probability(params, E, L, mango.Decoherence(gamma21=1e-14, n=0),
                                  density=2.8)
P = mango.decoherence.probability(params, E, L, mango.WavePacket(sigma_x_m=2e-13))

# non-unitary mixing N = (1 - alpha) U: zero-distance effect + non-cancelling
# NC matter term included
P = mango.nonunitarity.probability(params, mango.NonUnitarity(alpha21=0.02 + 0.01j),
                                   E, L, density=2.8)

Both are validated against their exact limits (γ→0 / α→0 reproduce the standard probabilities to ~1e-15; full damping gives the interference-averaged rates; the L→0 flavor violation matches the analytic zero-distance formula) and are differentiable through gamma / sigma_x / alpha — decoherence and unitarity-violation sensitivity forecasts are one jax.grad away.

Left: Lindblad damping washes the DUNE-baseline νμ survival oscillations out toward the interference-averaged rate as γL grows. Right: the classic reactor wave-packet signature at a JUNO-like baseline — the fast Δm²₃₁ wiggles disappear as σ_x shrinks while the slow solar oscillation survives. (examples/decoherence_plot.py)

Left: |α₂₁| = 0.02 shifts the DUNE appearance probability with a phase φ₂₁ dependence that mimics δ_CP (the NU–CP degeneracy). Right: the zero-distance effect — P(νμ→νe) plateaus at the analytic |(NN†)_eμ|² value (dotted) instead of vanishing as L→0. (examples/nonunitarity_plot.py)

Loading GLoBES experiments (a "differentiable GLoBES")

mango.globes parses a GLoBES/AEDL experiment definition — includes, channels, rules, migration matrices or analytic Gaussian smearing, efficiencies, normalization systematics, matter profile — and turns it into a differentiable forward model:

exp = mango.globes.load("DUNE_GLoBES.glb", scale=1.21)
spectra = exp.spectra(params)                 # dict rule -> reco spectrum
chi2 = exp.chi2(params, xi, data)             # Poisson + norm systematics
grad = jax.grad(lambda p: exp.spectra(p)["nue_app"].sum())(params)

Validated end-to-end on the official DUNE TDR configuration (it reproduces the bespoke DUNE analysis in analyses/dune to ~1e-5 — and parsing the file's @energy_window even caught a hand-coded window bug in the bespoke version). See the module docstring for the supported AEDL subset.

The classic DUNE TDR figure straight from mango.globes.load: νe / ν̄e appearance for δ_CP = −π/2, 0, +π/2 (note the CP effect flipping sign between neutrino and antineutrino modes), backgrounds shaded. (examples/dune_globes_plot.py)

Fisher forecasts and Bayesian fits (mango.stat)

F = mango.stat.fisher_matrix(model, params)     # exact Poisson Fisher via autodiff
F = F.with_prior("theta13", 0.002)              # reactor constraint
F.sigma("deltacp")                              # marginalized 1-sigma forecast
F.fixed("dm31").sigma("deltacp")                # conditional bound

The Fisher matrix equals the Asimov-χ² curvature (tested), and pairs naturally with gradient-based samplers: examples/hmc_posterior.py runs a blackjax NUTS posterior of (θ₂₃, θ₁₃, Δm²₃₁, δ_CP) for a DUNE-like experiment using exact mango gradients and overlays the analytic Fisher ellipses:

---

Algorithm choices

This section documents why each numerical method was chosen.

Units and precision

The numerical core runs entirely in natural units (ħ = c = 1): energies in eV, lengths in eV⁻¹, masses² in eV². Unit conversions (GeV, km, g/cm³) and the matter potential happen only at the API boundary (constants.py), derived from first-principles CODATA/PDG values — there are no opaque literals like 7.6e-14. The charged-current potential is V_CC = √2 G_F N_e; sterile neutrinos additionally feel the relative neutral-current term ½√2 G_F N_n. float64/complex128 is mandatory and set on import: at float32 the large oscillation phases lose precision (float32 buys only ~1.06× on CPU and ~1e−6 accuracy loss — not worth it except possibly on GPU).

Propagation backends

The evolution of a constant-density layer is S = exp(−iHL); P_{αβ} = |S_{βα}|². Four backends, agreeing to ~1e−14:

backend	method	used for
nufast	analytic NuFast-LBL formula (default for constant density / vacuum)	single layer, standard 3-flavor
cayley	Cayley–Hamilton via Newton divided differences (default for the Earth path)	layered Earth, all cases
eigh	generic Hermitian eigendecomposition	any N (steriles), validation
expm	jax.scipy.linalg.expm	reference cross-check

• nufast ports Parke & Denton's NuFast-LBL ("DMP Rosetta") algorithm: the matter mixing-matrix magnitudes |Uᵐ_{αi}|² and the matter Jarlskog are computed analytically from the eigenvalues and plugged straight into the standard sin² probability formula — no eigenvectors, no matrix exponential, no matmuls. It is exact (matches the matrix-exponential backends to ~5e−15), differentiable, and several times faster, putting constant-density throughput on par with NuFast's C++. It applies only to a single constant-density layer and standard 3-flavor, so it falls back automatically to cayley/eigh when NSI or steriles are present.
• cayley is the default for the multi-layer Earth path because there you need the evolution operator S (to chain across layers), not just probabilities. It uses the analytic 3×3 Hermitian eigenvalues (stable trigonometric solution of the characteristic cubic) and the Newton divided-difference (Hermite) form of exp(−iM) — which is eigenvector-free, handles repeated / zero eigenvalues exactly (so zero-length Earth segments collapse to identity cleanly), and avoids the ill-conditioning of a naive Vandermonde solve (the raw Hamiltonian eigenvalues are ~1e−13 eV, so it works with the dimensionless M = HL instead).
• eigh handles N ≠ 3 (sterile models) and is the generic validation oracle. The backend is auto-switched to eigh for N ≠ 3 since the analytic kernels are 3×3-specific.

Earth model and geometry (earth.py)

• PREM (Dziewonski & Anderson 1981) density as piecewise polynomials in radius; a two-zone Y_e (core/mantle, boundary at 3480 km) that is user-configurable (ye_core / ye_mantle) to match whatever convention you are comparing against.
• One unified entrypoint (probability_earth → earth.chord_segments) covers up-going through the Earth, down-going, and the atmospheric production height h_atm_km (so down-going / near-horizon directions get the correct vacuum baseline). Geometry uses the distance-from-closest-approach coordinate d(r) = √(r²−r_min²); PREM shell boundaries are respected exactly.
• Layer placement, NuFast-style. Shells are allocated proportional to radial thickness (not n_sub per PREM region), so resolution goes to the thick core/mantle where the oscillation phase accumulates rather than to thin crust layers; density is sampled at each segment's path midpoint (not the shell's geometric midpoint). Together these converge ~3–5× faster in n_sub: e.g. n_sub=2 (25 shells) reaches ~3e−3 accuracy, better than naive uniform subdivision at 121 shells.

Layer product (layers.py), device-aware

The per-layer operators are chained S = S_N ⋯ S_1. Two strategies, auto-selected:

• CPU: a sequential lax.scan — least total work and memory; fastest where there is no parallel hardware to exploit.
• GPU/TPU: build all per-layer propagators with vmap (vectorized eigensolves, hoisting the layer-independent vacuum Hamiltonian) and combine with lax.associative_scan (parallel-prefix, O(log N) depth).

(The parallel-prefix product is slower on CPU because it computes all prefixes; benchmarking confirmed the sequential scan wins on CPU, so it is the CPU default.)

Differentiability

• The cayley divided differences and the analytic eigensolver are guarded with double-where so degenerate / zero-length segments carry no NaN gradients (sqrt has an infinite derivative at 0).
• The Earth chord geometry uses a gradient-safe sqrt at the turning point (closest approach), where the geometric vertical tangent would otherwise NaN dP/d cos θ_z. Autodiff matches finite differences for parameters and geometry.
• Antineutrinos are a conjugation/sign flag (U → U*, V → −V, ε → ε*), not a separate code path, so gradients flow identically.

Continuous (ODE) backend (ode.py)

For arbitrary smooth profiles (and as an independent cross-check of the layer method), mango.ode integrates dS/ds = −iH(s)S directly. Two solvers: odeint (default, no extra dependency) and an optional diffrax backend (backend="diffrax", solver="dopri8"/"kvaerno5"/…) with stiff/implicit solvers and checkpointed adjoints — most useful for hard profiles (e.g. supernovae). The diffrax path evolves a real [Re S, Im S] state (diffrax complex support is experimental). These are only approximately unitary (RK drifts off the unitary group), so they are not on the default path; the layer method is exact per layer.

Solar adiabatic MSW (solar.py)

A BS05-style standard-solar-model loader (or an analytic exponential profile) plus the averaged adiabatic mass-state fractions F_i(r) = Σ_k |⟨ν_i^vac|ν_k^m(r)⟩|² · |⟨ν_k^m(r_emit)|ν_e⟩|² — the textbook LMA-MSW result (a ν_e produced in the dense core emerges predominantly as ν₂).

---

Validation

python run_tests.py (or pytest) — 132 checks, including:

• vacuum two-flavor analytic limit; unitarity in vacuum & matter; CP asymmetry; MSW;
• exact Earth chord geometry; oscillogram unitarity, no NaNs; down-going → vacuum;
• cross-backend agreement (nufast / cayley / eigh / expm) to ~1e−14;
• autodiff vs finite differences for oscillation parameters and cos θ_z; jit;
• NSI(0) → standard, sterile(θ_i4=0) → 3-flavor, reactor-sterile RAA depth, 3+2 unitarity, gradients through NSI / sterile params;
• ODE (odeint & diffrax) vs the layer method.

External cross-code validation (see validation/README.md):

comparison	constant density	PREM Earth
vs OscProb (ROOT/Eigen)	~1e−9	~1e−7 (identical path)
vs NuFast (C++)	5e−13 (constants aligned)	~1e−5 (converged)

The residual vs NuFast is fully explained by NuFast's 6-significant-figure constants; mango and OscProb (which share first-principles constants) agree to ~1e−9 despite completely different propagators. Reference values are embedded as regression tests (tests/test_oscprob_reference.py, tests/test_nufast_reference.py).

Reproductions of the nu-waves reference plots

examples/nuwaves/ reproduces all six figures from the nu-waves README (vacuum PMNS, sterile + energy smearing, constant-density NO/IO, (E,L) maps, the atmospheric oscillogram, solar adiabatic MSW) — five match essentially exactly; the solar one reproduces the correct MSW physics (see that folder's README for the convention note). examples/nsi_sterile.py demonstrates the NSI and sterile front-ends.

---

Performance (Apple M3 Pro, CPU, float64)

workload	mango	NuFast (C++)	OscProb (C++)
constant density, P(νμ→νe)	37–94 Mevals/s (nufast backend)	28	14
PREM Earth oscillogram	~0.04–0.07 Mevals/s	51	0.5
all 6 param gradients	2.9–5.4× a forward eval (exact)	n/a (finite diff ≈ 12×)	n/a

mango matches or beats the hand-optimized C++ for the constant-density/LBL case (the nufast backend is a single analytic formula, no scan). For PREM Earth on CPU it is much slower: the cost is XLA per-op overhead on the sequential scan of tiny 3×3 ops, not the algorithm — NuFast-Earth's analytic per-layer machinery is in a different class on CPU. mango's Earth value is differentiability + GPU/TPU batching (where the gap closes). See DESIGN.md §11 for the GPU-oriented NuFast-Earth roadmap item.

Reproducing the paper's benchmark tables

benchmarks/ has standalone, self-contained scripts (only mango/numpy/jax/matplotlib) that reproduce every performance and validation table in the accompanying paper, plus a README mapping each script to the table/figure it reproduces. The paper's own absolute numbers were measured on 8 cores of an AMD EPYC 7542 (CPU) and one NVIDIA A100 (GPU); expect different absolute timings on other hardware, but the same relative ordering of the backends.

---

Module map

mango/
  constants.py   physical constants + unit conversions + matter potentials
  params.py      OscParams PyTree (differentiable leaves) + a benchmark point
  pmns.py        generic N-flavor PMNS construction
  hamiltonian.py vacuum + matter Hamiltonian (generic N, NSI + sterile)
  eigensolve.py  analytic 3x3 Hermitian eigenvalues
  propagator.py  exp(-iHL): cayley / eigh / expm
  nufast.py      fast analytic constant-density probability (NuFast-LBL port)
  layers.py      device-aware layered propagation
  earth.py       PREM + chord geometry + thickness-proportional shells
  oscillator.py  high-level API (vacuum / constant / earth / profile)
  nsi.py         non-standard interactions
  sterile.py     3+N sterile front-end
  decoherence.py Lindblad / wave-packet decoherence
  nonunitarity.py non-unitary mixing (alpha parametrization)
  globes.py      generic GLoBES/AEDL experiment loader
  stat.py        Fisher information & chi2 utilities
  solar.py       solar profile + adiabatic MSW
  ode.py         continuous-density backend (odeint / diffrax)

---

Limitations & roadmap

• 3-flavor is the core; sterile (3+N) and matter NSI are implemented.
• CPU Earth throughput is far below the specialized C++ codes (see Performance). A GPU-oriented NuFast-Earth propagator (eigensystem caching across cos θ_z + reduced-basis real per-shell amplitudes + symmetric-trajectory halving) is the roadmap item to close that gap on accelerators — DESIGN.md §11.
• A supernova module would be the natural consumer of the diffrax stiff backend.

---

References & acknowledgements

• PREM. Dziewonski & Anderson, Phys. Earth Planet. Inter. 25 (1981) 297. mango/earth.py implements the PREM density as piecewise polynomials in radius, from the coefficients tabulated in that paper.
• NuFast / DMP. Denton & Parke, Phys. Rev. D 110 (2024) 073005 (arXiv:2405.02400); see also Denton, Minakata & Parke (arXiv:1604.08167) for the underlying "Rosetta" relations. mango/nufast.py's nufast backend is a close, line-by-line JAX transcription of the reference C++ implementation that the authors publish alongside the paper as NuFast-LBL — the intermediate variable naming in mango/nufast.py tracks that source closely enough that this is a port, not an independent re-derivation from the formulas with new structure. NuFast-LBL is itself MIT-licensed (Copyright (c) 2024 Peter B. Denton); this port is used under the terms of that license, with the original authors' copyright credited here and in mango/nufast.py's module docstring. If you use the nufast backend, please also cite Denton & Parke (2024) as the original authors request. The NuFast-LBL MIT licence is reproduced verbatim in THIRD_PARTY_LICENSES.md, as its terms require.
• Validated against OscProb and NuFast; reproduces nu-waves reference plots (see examples/nuwaves/README.md for per-figure credit and the one physics-convention caveat on the solar plot).
• Solar profile. Bahcall, Serenelli & Basu, BS05(AGS,OP) standard solar model (astro-ph/0412440); the table shipped at examples/data/bs05_agsop.dat is that published model's own data file, used as input to mango.solar.load_bs05 (mango/solar.py), which is otherwise an independent implementation of the textbook averaged-adiabatic MSW formula.

See CITATION.cff for how to cite this software (and the accompanying paper, once available), and CHANGELOG.md for release history.

Releasing to PyPI

Publishing is automated via GitHub Actions (.github/workflows/publish.yml) using PyPI trusted publishing (OIDC) — no API tokens or secrets.

One-time setup (on pypi.org, needs your PyPI login). Account → Publishing → Add a pending publisher with:

field	value
PyPI Project Name	mango
Owner	pgranger23
Repository name	mango-osc
Workflow name	publish.yml
Environment name	pypi

To cut a release:

1. Bump the version in pyproject.toml and mango/__init__.py, commit, push.
2. Tag and publish a GitHub Release: e.g. git tag v0.1.0 && git push origin v0.1.0, then create a Release for that tag (or gh release create v0.1.0 --generate-notes).
3. The workflow builds the sdist + wheel and publishes to PyPI → pip install mango.

(You can dry-run the build locally with python -m build && twine check dist/*.)

License

MIT — see LICENSE.