Metadata-Version: 2.4
Name: jaxnu
Version: 0.2.3
Summary: Differentiable neutrino oscillation calculator in JAX (vacuum, constant density, PREM Earth, atmospheric, solar, NSI, sterile)
Author: pgranger23
License: MIT
Project-URL: Homepage, https://github.com/pgranger23/jaxnu-osc
Project-URL: Repository, https://github.com/pgranger23/jaxnu-osc
Project-URL: Issues, https://github.com/pgranger23/jaxnu-osc/issues
Keywords: neutrino,oscillation,jax,differentiable,PREM,atmospheric,matter effects,MSW,NSI,sterile neutrino
Classifier: Development Status :: 4 - Beta
Classifier: Intended Audience :: Science/Research
Classifier: License :: OSI Approved :: MIT License
Classifier: Programming Language :: Python :: 3
Classifier: Topic :: Scientific/Engineering :: Physics
Requires-Python: >=3.10
Description-Content-Type: text/markdown
License-File: LICENSE
Requires-Dist: jax>=0.4.20
Requires-Dist: numpy>=1.24
Provides-Extra: test
Requires-Dist: pytest>=7; extra == "test"
Provides-Extra: examples
Requires-Dist: matplotlib>=3.6; extra == "examples"
Requires-Dist: blackjax>=1.0; extra == "examples"
Provides-Extra: diffrax
Requires-Dist: diffrax>=0.5; extra == "diffrax"
Dynamic: license-file

# jaxnu

[![CI](https://github.com/pgranger23/jaxnu-osc/actions/workflows/ci.yml/badge.svg)](https://github.com/pgranger23/jaxnu-osc/actions/workflows/ci.yml)
[![License: MIT](https://img.shields.io/badge/License-MIT-blue.svg)](LICENSE)

**A differentiable neutrino oscillation calculator in [JAX](https://docs.jax.dev).**

`jaxnu` 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](https://github.com/joaoabcoelho/OscProb) and
[NuFast](https://github.com/PeterDenton), and reproduces the reference plots of the
[nu-waves](https://github.com/nadrino/nu-waves) library.

![atmospheric oscillogram](examples/nuwaves/figures/matter_prem_test.jpg)

*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. `jaxnu` 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](https://arxiv.org/abs/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](examples/oscillograds_atm.jpg)

*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`](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=)`:

![core electron-density oscillograds](examples/oscillograds_density.jpg)

*`∂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:

```python
model = jaxnu.prem_layered(n_sub=3)   # constant-density-shell Earth sampled from PREM
P = jaxnu.probability_earth(params, E, cz, earth_model=model)
jac = jax.jacrev(lambda m: jaxnu.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 -> ...
```

![Earth-model oscillograds](examples/oscillograds_earth_model.jpg)

*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/`](analyses/dune) reproduces the **DUNE TDR oscillation analysis**
using jaxnu 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 jaxnu**
(`jax.value_and_grad` → L-BFGS). See its [README](analyses/dune/README.md).

## Installation

```bash
git clone https://github.com/pgranger23/jaxnu-osc.git
cd jaxnu-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 jaxnu`.

---

## Quick start

```python
import jax, jax.numpy as jnp
import jaxnu
from jaxnu 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 = jaxnu.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 jaxnu.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

```python
# vacuum
jaxnu.probability_vacuum(params, E, baseline_km=295.0)                      # T2K-like

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

# PREM Earth (+ optional atmosphere); single entrypoint
jaxnu.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)
jaxnu.probability_profile(params, E, density_gcc=[...], ye=[...], length_km=[...])

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

### NSI and sterile neutrinos

```python
# matter NSI: an epsilon_{alpha beta} object on any probability function
P = jaxnu.probability_constant(params, E, 1300.0, density=2.8,
                               nsi=jaxnu.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 = jaxnu.Sterile3plus1(theta12, theta13, theta23, theta14, theta24, theta34,
                         delta13, delta24, dm21, dm31, dm41)
P = jaxnu.probability_earth(st, E, cz)            # (n_cz, n_E, 4, 4)

# arbitrary 3+N via the generic builder
U = jaxnu.pmns_nflavor(5, [(3, 4, 0.1), (2, 4, 0.05), ..., (0, 1, theta12)])
params5 = jaxnu.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

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

# non-unitary mixing N = (1 - alpha) U: zero-distance effect + non-cancelling
# NC matter term included
P = jaxnu.nonunitarity.probability(params, jaxnu.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.

![decoherence](examples/decoherence.jpg)

*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](examples/decoherence_plot.py))*

![non-unitarity](examples/nonunitarity.jpg)

*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](examples/nonunitarity_plot.py))*

### Loading GLoBES experiments (a "differentiable GLoBES")

`jaxnu.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:

```python
exp = jaxnu.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.

![DUNE spectra from the GLoBES loader](examples/dune_globes.jpg)

*The classic DUNE TDR figure straight from `jaxnu.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](examples/dune_globes_plot.py))*

### Fisher forecasts and Bayesian fits (`jaxnu.stat`)

```python
F = jaxnu.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](examples/hmc_posterior.py)
runs a blackjax **NUTS** posterior of (θ₂₃, θ₁₃, Δm²₃₁, δ_CP) for a DUNE-like
experiment using exact jaxnu gradients and overlays the analytic Fisher ellipses:

![HMC posterior vs Fisher](examples/hmc_posterior.jpg)

---

## 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), `jaxnu.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](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;
jaxnu 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/`](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`](examples/nsi_sterile.py) demonstrates
the NSI and sterile front-ends.

---

## Performance (Apple M3 Pro, CPU, float64)

| workload | jaxnu | 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 |

`jaxnu` 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. jaxnu's Earth value is differentiability + GPU/TPU batching
(where the gap closes). See [DESIGN.md](DESIGN.md) §11 for the GPU-oriented
NuFast-Earth roadmap item.

### Reproducing the paper's benchmark tables

[`benchmarks/`](benchmarks) has standalone, self-contained scripts (only
`jaxnu`/`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

```
jaxnu/
  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](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.
  `jaxnu/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. `jaxnu/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](https://github.com/PeterDenton/NuFast-LBL) — the intermediate
  variable naming in `jaxnu/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 `jaxnu/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`](THIRD_PARTY_LICENSES.md), as its terms require.
- Validated against [OscProb](https://github.com/joaoabcoelho/OscProb) and
  [NuFast](https://github.com/PeterDenton); reproduces
  [nu-waves](https://github.com/nadrino/nu-waves) reference plots (see
  [`examples/nuwaves/README.md`](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`](examples/data/bs05_agsop.dat) is that
  published model's own data file, used as input to `jaxnu.solar.load_bs05`
  (`jaxnu/solar.py`), which is otherwise an independent implementation of the
  textbook averaged-adiabatic MSW formula.

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

## Releasing to PyPI

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

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

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

**To cut a release:**

1. Bump the version in `pyproject.toml` **and** `jaxnu/__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 jaxnu`.

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

## License

MIT — see [LICENSE](LICENSE).
