odmr-estimator

A 1D residual CNN that reads a magnetic field out of an NV-center ODMR spectrum, replacing nonlinear least squares. 180,270 parameters, ~0.4 ms per spectrum batched, with a trained heteroscedastic uncertainty head.

Inputs and outputs

Input — 512-point normalised fluorescence sweep over 2700–3040 MHz.

Output — 7 means and 7 log-variances:

target	units	notes
p1–p4	mT	|B·nᵢ|, field projection magnitude onto each NV axis
D_shift	MHz	zero-field-splitting shift; maps to temperature at −74 kHz/K
E_strain	MHz	transverse strain term
linewidth	MHz	Lorentzian FWHM

Field magnitude follows exactly from Σᵢ(B·nᵢ)² = (4/3)|B|².

The spectrum does not determine the sign of each projection — the diamond's point-group symmetry makes it unobservable — so the model regresses projection magnitudes, which are uniquely identified. Recovering the Cartesian vector requires a known bias field.

Results

split	MAE B∥ (mT)	median AE	fail >0.5 mT	MAE |B|	MAE Γ (MHz)
validation	0.0287	0.0217	0.0%	0.035	0.274
test_hard	0.0516	0.0431	0.0%	0.061	0.962

Against classical estimators on validation: NLLS warm start 0.505 mT (25.2% failures), NLLS cold start 0.701 mT (42.8%), peak-find 1.163 mT (54.0%). Speedup is ~200× batched against a cold-start fit.

The NLLS baseline is honest — it converges rather than running out of iterations, and its failures are local-minima trapping in a 24-line overlapping fit. Its median error (0.085 mT) is far better than its mean (0.505 mT), and quoting only the mean would overstate the win.

Uncertainty calibration — read this before trusting the error bars

The uncertainty head is trained by Gaussian negative log-likelihood, so predicted σ should track actual error. Measured 1σ coverage:

• validation: 66.2% against an ideal 68.3% — well calibrated
• test_hard: 53.2% against 68.3% — overconfident

On out-of-distribution spectra the model understates its own error. If you deploy this in a regime unlike the training distribution, inflate σ or recalibrate. This is the model's main known weakness and it is not fixed.

Usage

infer_numpy.py is a dependency-free forward pass that mirrors the JAX model exactly (verified parity: max abs difference 1.17e-06). No deep-learning runtime required, which is the point — this is meant to run on a sensor.

import numpy as np
from infer_numpy import load, predict

params = load("weights.npz")
mean, sigma = predict(params, spectrum)      # spectrum: (512,) or (N, 512)
b_magnitude = np.sqrt(0.75 * np.sum(mean[:4] ** 2))

Training

40,000 synthetic spectra, 40 epochs, Adam, Gaussian NLL loss. Roughly 15 minutes on one CPU core. generate.py and train.py reproduce it end to end; train.log is the full run.

Limitations

Trained entirely on synthetic spectra from the NV ground-state Hamiltonian. Standard practice for ODMR — the lineshape is known physics and labels are exact — but real instruments add microwave power broadening, sweep-correlated laser noise, frequency calibration error, and diamond-specific strain distributions that are not simulated here. No validation on measured spectra has been performed. Treat the numbers as an upper bound until that is done.