nereids_fitting/
joint_resolution.rs

1//! Fit a sample and the calibrant that measured its resolution together.
2//!
3//! Pinning a calibrated resolution into a sample fit reports the temperature
4//! as more certain than it is: resolution width and temperature broaden the
5//! line the same way, so the uncertainty that belongs to their degeneracy is
6//! dropped. Carrying the calibration forward as a Gaussian prior does not
7//! recover it — the calibration's own uncertainty is neither Gaussian nor
8//! separable. Its objective is flat where the kernel is narrower than the
9//! line it broadens and a wall above, and for the Gaussian family the two
10//! width coordinates trade off almost exactly, so a per-parameter sigma
11//! describes a direction the calibration never moves in.
12//!
13//! What has none of those problems is the calibrant's residuals themselves.
14//! This model evaluates the sample and the calibrant against ONE resolution
15//! drawn from the shared parameter vector and returns both predictions, so
16//! the optimizer sees a single objective
17//!
18//! ```text
19//! chi^2(T, n, w) = chi^2_sample(T, n, w) + chi^2_calibrant(w)
20//! ```
21//!
22//! whose temperature uncertainty already contains what the calibrant failed
23//! to pin down.
24//!
25//! The objective is exact: no part of the calibration is summarized, so the
26//! shape the summary would have lost is still there. The uncertainty READ OFF
27//! it is not. `temperature_k_unc` comes from the optimizer's local curvature
28//! at the solution and is a Gaussian approximation like any other, so on a
29//! surface with a flat side and a wall it describes the solution's
30//! neighbourhood rather than the whole interval. What the joint objective
31//! fixes is that the neighbourhood is now the right one — it includes the
32//! resolution's freedom instead of holding it fixed.
33//!
34//! The calibrant's own density and temperature are what make it a calibrant
35//! and stay fixed; only the resolution is shared.
36//!
37//! The two shared slots hold the SQUARED widths, in µs² and m². The kernel
38//! combines the timing and flight-path terms in quadrature,
39//! `W² = timing(Δt)² + path(ΔL)²`, so a width itself enters `W` quadratically
40//! and `∂W/∂ΔL` is exactly zero at `ΔL = 0`. `W²` is linear in the squares,
41//! so `∂W/∂(ΔL²)` is finite there and a width seeded at zero is still fitted.
42
43use std::sync::Arc;
44
45use nereids_endf::resonance::ResonanceData;
46use nereids_physics::resolution::ResolutionParams;
47use nereids_physics::transmission::{self, InstrumentParams, SampleParams};
48
49use crate::error::FittingError;
50use crate::lm::FitModel;
51
52/// Where a spectrum's areal densities come from.
53pub enum Densities {
54    /// `params[i]` for each index: the fit determines them.
55    Fitted(Vec<usize>),
56    /// One value per isotope, held at what the caller knows it to be.
57    Known(Vec<f64>),
58}
59
60impl Densities {
61    fn len(&self) -> usize {
62        match self {
63            Self::Fitted(indices) => indices.len(),
64            Self::Known(values) => values.len(),
65        }
66    }
67
68    fn at(&self, i: usize, params: &[f64]) -> f64 {
69        match self {
70            Self::Fitted(indices) => params[indices[i]],
71            Self::Known(values) => values[i],
72        }
73    }
74}
75
76/// One spectrum in a joint fit: its grid, what is in the beam, and how its
77/// free parameters are read out of the shared vector.
78pub struct SpectrumSpec {
79    /// Energy grid (eV), ascending.
80    pub energies: Vec<f64>,
81    /// One entry per isotope in the beam.
82    pub resonance_data: Vec<ResonanceData>,
83    /// Areal densities, fitted or known.
84    pub densities: Densities,
85    /// `params[temperature_index]` is the temperature, else `temperature_k`.
86    pub temperature_index: Option<usize>,
87    /// Temperature (K) when it is not fitted.
88    pub temperature_k: f64,
89}
90
91impl SpectrumSpec {
92    fn sample_params(&self, params: &[f64]) -> Result<SampleParams, FittingError> {
93        let temperature_k = match self.temperature_index {
94            Some(i) => params[i],
95            None => self.temperature_k,
96        };
97        let isotopes = self
98            .resonance_data
99            .iter()
100            .enumerate()
101            .map(|(i, rd)| (rd.clone(), self.densities.at(i, params)))
102            .collect();
103        SampleParams::new(temperature_k, isotopes)
104            .map_err(|e| FittingError::EvaluationFailed(format!("sample params: {e:?}")))
105    }
106
107    fn predict(
108        &self,
109        params: &[f64],
110        instrument: &InstrumentParams,
111    ) -> Result<Vec<f64>, FittingError> {
112        let sample = self.sample_params(params)?;
113        transmission::forward_model(&self.energies, &sample, Some(instrument))
114            .map_err(|e| FittingError::EvaluationFailed(format!("forward: {e:?}")))
115    }
116}
117
118/// A sample and its calibrant, sharing one Gaussian resolution.
119///
120/// [`FitModel::evaluate`] returns the sample's predictions followed by the
121/// calibrant's, so the caller fits against the two spectra concatenated in
122/// that order.
123pub struct JointResolutionModel {
124    sample: SpectrumSpec,
125    calibrant: SpectrumSpec,
126    flight_path_m: f64,
127    /// `params[delta_t_sq_index]` / `params[delta_l_sq_index]` are the SQUARED
128    /// Gaussian widths, in µs² and m², shared by both spectra.
129    delta_t_sq_index: usize,
130    delta_l_sq_index: usize,
131}
132
133impl JointResolutionModel {
134    /// Build the joint model.
135    ///
136    /// `delta_t_sq_index` / `delta_l_sq_index` are the shared slots holding
137    /// the squared widths. Everything else about each arm, including whether
138    /// its densities and temperature are fitted, is in its
139    /// [`SpectrumSpec`] — the calibrant's being known is what makes it a
140    /// calibrant.
141    ///
142    /// # Errors
143    /// [`FittingError::InvalidConfig`] when a spectrum's grid is empty, when
144    /// its density count does not match its isotope count, or when two
145    /// parameters share a slot.
146    pub fn new(
147        sample: SpectrumSpec,
148        calibrant: SpectrumSpec,
149        flight_path_m: f64,
150        delta_t_sq_index: usize,
151        delta_l_sq_index: usize,
152    ) -> Result<Self, FittingError> {
153        for (label, spec) in [("sample", &sample), ("calibrant", &calibrant)] {
154            if spec.energies.is_empty() {
155                return Err(FittingError::InvalidConfig(format!(
156                    "the {label} needs a non-empty energy grid"
157                )));
158            }
159            if spec.densities.len() != spec.resonance_data.len() {
160                return Err(FittingError::InvalidConfig(format!(
161                    "the {label} has {} densities for {} isotopes",
162                    spec.densities.len(),
163                    spec.resonance_data.len(),
164                )));
165            }
166        }
167        // Every slot the model reads must name one quantity. Two of them
168        // sharing an index makes a single optimizer coordinate move two
169        // different physical things at once.
170        let mut slots = vec![delta_t_sq_index, delta_l_sq_index];
171        for spec in [&sample, &calibrant] {
172            if let Densities::Fitted(indices) = &spec.densities {
173                slots.extend(indices);
174            }
175            slots.extend(spec.temperature_index);
176        }
177        let mut seen = slots.clone();
178        seen.sort_unstable();
179        seen.dedup();
180        if seen.len() != slots.len() {
181            return Err(FittingError::InvalidConfig(format!(
182                "two parameters share a slot: {slots:?}"
183            )));
184        }
185        Ok(Self {
186            sample,
187            calibrant,
188            flight_path_m,
189            delta_t_sq_index,
190            delta_l_sq_index,
191        })
192    }
193
194    /// Number of data points the sample contributes, i.e. where the
195    /// calibrant's predictions start in [`FitModel::evaluate`]'s output.
196    #[must_use]
197    pub fn sample_len(&self) -> usize {
198        self.sample.energies.len()
199    }
200
201    /// Total length of the concatenated prediction.
202    #[must_use]
203    pub fn len(&self) -> usize {
204        self.sample.energies.len() + self.calibrant.energies.len()
205    }
206
207    /// Whether the joint prediction is empty. Never true: both grids are
208    /// checked non-empty at construction.
209    #[must_use]
210    pub fn is_empty(&self) -> bool {
211        self.len() == 0
212    }
213
214    /// The resolution both spectra are evaluated with at this probe.
215    fn instrument(&self, params: &[f64]) -> Result<Arc<InstrumentParams>, FittingError> {
216        let resolution = ResolutionParams::new(
217            self.flight_path_m,
218            params[self.delta_t_sq_index].max(0.0).sqrt(),
219            params[self.delta_l_sq_index].max(0.0).sqrt(),
220            0.0,
221        )
222        .map_err(|e| FittingError::EvaluationFailed(format!("shared resolution: {e:?}")))?;
223        Ok(Arc::new(InstrumentParams {
224            resolution: nereids_physics::resolution::ResolutionFunction::Gaussian(resolution),
225        }))
226    }
227}
228
229impl FitModel for JointResolutionModel {
230    fn evaluate(&self, params: &[f64]) -> Result<Vec<f64>, FittingError> {
231        let instrument = self.instrument(params)?;
232        let mut out = self.sample.predict(params, &instrument)?;
233        out.extend(self.calibrant.predict(params, &instrument)?);
234        Ok(out)
235    }
236}
237
238#[cfg(test)]
239mod tests {
240    use super::*;
241    use crate::lm::{LmConfig, levenberg_marquardt};
242    use crate::parameters::{FitParameter, ParameterSet};
243    use crate::resolution_calib::{CalibrationConfig, ResolutionFamily, calibrate_resolution};
244    use nereids_endf::resonance::test_support::synthetic_isotope;
245    use nereids_physics::resolution::ResolutionFunction;
246    use rand::SeedableRng;
247    use rand_chacha::ChaCha12Rng;
248    use rand_distr::{Distribution, Normal};
249
250    const L: f64 = 25.0;
251    const T_TRUE: f64 = 300.0;
252    const DENSITY: f64 = 2.0e-3;
253    const W_TRUE: f64 = 0.30;
254    const DL_TRUE: f64 = 0.05;
255    const NOISE: f64 = 0.002;
256
257    /// The two arms for a test: the sample's density and temperature fitted
258    /// at the given slots, the calibrant's known.
259    fn arms(
260        iso: &ResonanceData,
261        energies: &[f64],
262        temperature_index: Option<usize>,
263    ) -> (SpectrumSpec, SpectrumSpec) {
264        (
265            SpectrumSpec {
266                energies: energies.to_vec(),
267                resonance_data: vec![iso.clone()],
268                densities: Densities::Fitted(vec![0]),
269                temperature_index,
270                temperature_k: T_TRUE,
271            },
272            SpectrumSpec {
273                energies: energies.to_vec(),
274                resonance_data: vec![iso.clone()],
275                densities: Densities::Known(vec![DENSITY]),
276                temperature_index: None,
277                temperature_k: T_TRUE,
278            },
279        )
280    }
281
282    fn fixture() -> (ResonanceData, Vec<f64>, Vec<f64>) {
283        let iso = synthetic_isotope(72, 178, 20.0, 0.05, 0.06);
284        let energies: Vec<f64> = (0..120).map(|i| 18.0 + i as f64 * 0.04).collect();
285        let sample = SampleParams::new(T_TRUE, vec![(iso.clone(), DENSITY)]).unwrap();
286        let inst = InstrumentParams {
287            resolution: ResolutionFunction::Gaussian(
288                ResolutionParams::new(L, W_TRUE, DL_TRUE, 0.0).unwrap(),
289            ),
290        };
291        let truth = transmission::forward_model(&energies, &sample, Some(&inst)).unwrap();
292        (iso, energies, truth)
293    }
294
295    /// The joint fit reports the temperature uncertainty the two-stage
296    /// procedure actually has; pinning the calibrated kernel reports less.
297    ///
298    /// The oracle is repetition. Calibrate on a fresh noisy calibrant, fit a
299    /// fresh noisy sample with that resolution pinned, and the fitted
300    /// temperature scatters by the full two-stage uncertainty. The pinned fit
301    /// cannot see the calibration noise and reports only its own; the joint
302    /// fit has the calibrant residuals in its objective and should report
303    /// both.
304    #[test]
305    #[ignore = "slow; runs nightly"]
306    fn the_joint_fit_reports_the_temperature_uncertainty_pinning_drops() {
307        const REALIZATIONS: usize = 16;
308        let (iso, energies, truth) = fixture();
309        let unc = vec![NOISE; energies.len()];
310        let mut rng = ChaCha12Rng::seed_from_u64(20260917);
311        let normal = Normal::new(0.0, NOISE).unwrap();
312        let mut noisy =
313            || -> Vec<f64> { truth.iter().map(|t| t + normal.sample(&mut rng)).collect() };
314
315        let cfg = CalibrationConfig {
316            ic_n_energies: 8,
317            ic_n_tau: 32,
318            max_iter: 400,
319            ..Default::default()
320        };
321        let calibrant_sample = SampleParams::new(T_TRUE, vec![(iso.clone(), DENSITY)]).unwrap();
322
323        let pinned_fit = |data: &[f64], w: f64, dl: f64| -> Option<(f64, f64)> {
324            let model = {
325                let (sample, calibrant) = arms(&iso, &energies, Some(1));
326                JointResolutionModel::new(sample, calibrant, L, 2, 3)
327            }
328            .unwrap();
329            // Sample arm only: the calibrant half is masked out by fitting
330            // against the sample data alone.
331            let sample_only = SampleOnly { inner: model };
332            let mut params = ParameterSet::new(vec![
333                FitParameter::non_negative("density", DENSITY),
334                FitParameter::non_negative("temperature_k", 285.0),
335                FitParameter::fixed("delta_t_us_sq", w * w),
336                FitParameter::fixed("delta_l_m_sq", dl * dl),
337            ]);
338            let r = levenberg_marquardt(
339                &sample_only,
340                data,
341                &unc,
342                &mut params,
343                &LmConfig {
344                    compute_covariance: true,
345                    max_iter: 100,
346                    ..Default::default()
347                },
348            )
349            .ok()?;
350            let sigma = r.uncertainties.as_ref()?.get(1).copied()?;
351            Some((r.params[1], sigma))
352        };
353
354        let mut fitted = Vec::new();
355        let mut pinned_sigmas = Vec::new();
356        for _ in 0..REALIZATIONS {
357            let cal_data = noisy();
358            let Ok(cal) = calibrate_resolution(
359                ResolutionFamily::Gaussian,
360                &energies,
361                &cal_data,
362                &unc,
363                &calibrant_sample,
364                &cfg,
365            ) else {
366                continue;
367            };
368            let ResolutionFunction::Gaussian(p) = &cal.resolution else {
369                unreachable!()
370            };
371            let sample_data = noisy();
372            if let Some((t, s)) = pinned_fit(&sample_data, p.delta_t_us(), p.delta_l_m()) {
373                fitted.push(t);
374                pinned_sigmas.push(s);
375            }
376        }
377        assert!(
378            fitted.len() >= REALIZATIONS / 2,
379            "only {} of {REALIZATIONS} two-stage realizations produced a fit",
380            fitted.len()
381        );
382        let n = fitted.len() as f64;
383        let mean = fitted.iter().sum::<f64>() / n;
384        let observed = (fitted.iter().map(|t| (t - mean).powi(2)).sum::<f64>() / (n - 1.0)).sqrt();
385        let pinned = pinned_sigmas.iter().sum::<f64>() / n;
386
387        // Joint fit on one realization: resolution free, calibrant in the
388        // objective.
389        let cal_data = noisy();
390        let sample_data = noisy();
391        let model = {
392            let (sample, calibrant) = arms(&iso, &energies, Some(1));
393            JointResolutionModel::new(sample, calibrant, L, 2, 3)
394        }
395        .unwrap();
396        let mut joint_data = sample_data.clone();
397        joint_data.extend_from_slice(&cal_data);
398        let joint_unc = vec![NOISE; joint_data.len()];
399        let mut params = ParameterSet::new(vec![
400            FitParameter::non_negative("density", DENSITY),
401            FitParameter::non_negative("temperature_k", 285.0),
402            FitParameter::non_negative("delta_t_us_sq", W_TRUE * W_TRUE),
403            FitParameter::non_negative("delta_l_m_sq", DL_TRUE * DL_TRUE),
404        ]);
405        let r = levenberg_marquardt(
406            &model,
407            &joint_data,
408            &joint_unc,
409            &mut params,
410            &LmConfig {
411                compute_covariance: true,
412                max_iter: 200,
413                ..Default::default()
414            },
415        )
416        .expect("joint fit runs");
417        let joint = r.uncertainties.as_ref().expect("joint covariance")[1];
418
419        eprintln!(
420            "observed {observed:.4}  pinned {pinned:.4}  joint {joint:.4}  T {:.3}",
421            r.params[1]
422        );
423        assert!(
424            pinned < 0.9 * observed,
425            "pinned sigma_T {pinned:.4} does not understate the two-stage scatter \
426             {observed:.4}; without that gap this test cannot show the joint fit \
427             recovering anything"
428        );
429        // The claim, and the only assertion a joint fit that ignored its
430        // calibrant would fail: dropping the calibrant half leaves the
431        // sample-only fit, whose sigma_T is the pinned one exactly.
432        assert!(
433            joint > pinned,
434            "the joint fit reports sigma_T {joint:.4}, no more than the pinned \
435             {pinned:.4}; the calibrant residuals are not reaching the objective"
436        );
437        let ratio = joint / observed;
438        assert!(
439            (0.5..=2.0).contains(&ratio),
440            "the joint fit reports sigma_T {joint:.4} against an observed \
441             two-stage scatter of {observed:.4} (ratio {ratio:.2})"
442        );
443    }
444
445    /// A flight-path width seeded at zero is still fitted.
446    ///
447    /// The kernel combines the two terms in quadrature, so `W` depends on a
448    /// width through its square and `dW/d(dL)` is exactly zero at `dL = 0`. A
449    /// finite-difference optimizer probing that coordinate sees only the
450    /// second-order term and the width never moves. The measure is the
451    /// Jacobian column the optimizer actually gets at that point, against the
452    /// timing column as the scale of a column it can follow.
453    #[test]
454    fn a_zero_flight_path_width_still_has_a_usable_jacobian_column() {
455        const FD_STEP: f64 = 1.0e-6;
456
457        let (iso, energies, _) = fixture();
458        let model = {
459            let (sample, calibrant) = arms(&iso, &energies, None);
460            JointResolutionModel::new(sample, calibrant, L, 1, 2)
461        }
462        .unwrap();
463
464        // params = [density, delta_t^2, delta_l^2], the flight-path width at
465        // zero and the timing width at its usual scale.
466        let base = [DENSITY, W_TRUE * W_TRUE, 0.0];
467        let column = |slot: usize| -> f64 {
468            let mut probed = base;
469            probed[slot] += FD_STEP * (1.0 + base[slot].abs());
470            let (a, b) = (
471                model.evaluate(&base).expect("base evaluates"),
472                model.evaluate(&probed).expect("probe evaluates"),
473            );
474            a.iter()
475                .zip(&b)
476                .map(|(x, y)| (x - y).abs())
477                .fold(0.0_f64, f64::max)
478                / FD_STEP
479        };
480
481        let timing = column(1);
482        let path = column(2);
483        assert!(
484            timing > 0.0,
485            "the timing column is {timing}, so there is no scale to compare against"
486        );
487        assert!(
488            path > 0.01 * timing,
489            "at a zero flight-path width the column is {path:.4e} against a \
490             timing column of {timing:.4e}; a column that small is below the \
491             noise of any real measurement and the width would never move"
492        );
493    }
494
495    /// No two parameters may name the same slot.
496    ///
497    /// A shared index does not fail loudly — it makes one optimizer
498    /// coordinate move two physical quantities at once, and the fit returns a
499    /// number.
500    #[test]
501    fn parameters_sharing_an_index_are_rejected() {
502        let (iso, energies, _) = fixture();
503        let build = |density: usize, temperature: Option<usize>, dt: usize, dl: usize| {
504            let (mut sample, calibrant) = arms(&iso, &energies, temperature);
505            sample.densities = Densities::Fitted(vec![density]);
506            JointResolutionModel::new(sample, calibrant, L, dt, dl)
507        };
508        assert!(
509            build(0, Some(1), 2, 3).is_ok(),
510            "the distinct layout is legal"
511        );
512        for (label, density, temperature, dt, dl) in [
513            ("temperature on the density slot", 0, Some(0), 2, 3),
514            ("width on the density slot", 0, Some(1), 0, 3),
515            ("width on the temperature slot", 0, Some(1), 1, 3),
516            ("the two widths on one slot", 0, Some(1), 2, 2),
517        ] {
518            assert!(
519                build(density, temperature, dt, dl).is_err(),
520                "{label} must be rejected"
521            );
522        }
523    }
524
525    /// The sample arm alone, for the pinned comparison.
526    struct SampleOnly {
527        inner: JointResolutionModel,
528    }
529
530    impl FitModel for SampleOnly {
531        fn evaluate(&self, params: &[f64]) -> Result<Vec<f64>, FittingError> {
532            let mut all = self.inner.evaluate(params)?;
533            all.truncate(self.inner.sample_len());
534            Ok(all)
535        }
536    }
537}
nereids_fitting/joint_resolution.rs

nereids_fitting/
joint_resolution.rs
