Skip to main content

nereids_fitting/
transmission_model.rs

1//! Transmission forward model adapter for fitting.
2//!
3//! Wraps the physics `forward_model` function into a `FitModel` trait object
4//! that the LM optimizer can call. The fit parameters are the areal densities
5//! (thicknesses) of each isotope in the sample.
6
7use std::cell::{Cell, RefCell};
8use std::rc::Rc;
9use std::sync::Arc;
10
11use nereids_core::constants::{EV_TO_JOULES, NEUTRON_MASS_KG};
12use nereids_endf::resonance::ResonanceData;
13use nereids_physics::resolution::{self, ResolutionFunction, ResolutionPlan};
14use nereids_physics::surrogate::{ScalarSurrogatePlan, SparseEmpiricalCubaturePlan};
15use nereids_physics::transmission::{self, InstrumentParams, SampleParams};
16
17use crate::error::FittingError;
18use crate::lm::{FitModel, FlatMatrix};
19
20/// Absolute-magnitude threshold for `L_scale` division safety in the
21/// partial-GAL rank-1 derivation of the energy-scale Jacobian.  When
22/// `|l_scale| < L_SCALE_EPSILON`, the per-bin
23/// `(tof_i - t0_clamped) / l_scale` factor in
24/// [`EnergyScaleTransmissionModel::analytical_jacobian`] blows up,
25/// and combined with the FD-based reference t0 column (which goes to
26/// ~0 at the same boundary) produces NaN entries in the L_scale
27/// Jacobian column.
28///
29/// Below this threshold the L_scale column falls through to the
30/// per-coordinate central-FD path that already follows the partial-GAL
31/// block in the same function.
32///
33/// **Note:** the literal `1.0e-12` matches the `1e-12` factor in the
34/// t0 clamp at [`EnergyScaleTransmissionModel::corrected_energies`]
35/// and the partial-GAL t0 FD precompute, but the semantic role
36/// differs — the t0 clamp is *relative* (`min_tof_us * (1 - 1e-12)`)
37/// while this constant is an *absolute* magnitude bound.  Both
38/// guards protect against the same `(tof - t0) / l_eff` blow-up at
39/// the energy-scale-degenerate corner; the value choice is
40/// coincident, not tied.  See issue #500 for the L_scale gap
41/// closure.
42const L_SCALE_EPSILON: f64 = 1.0e-12;
43
44/// Transmission model backed by precomputed Doppler-broadened cross-sections.
45///
46/// The expensive physics steps (resonance → σ(E), Doppler broadening) are
47/// computed once and stored.  Each `evaluate()` call performs Beer-Lambert
48/// and, when `instrument` is present, resolution broadening on the total
49/// transmission:
50///
51///   T(E) = R ⊗ exp(−Σᵢ nᵢ · σ_{D,i}(E))
52///
53/// Issue #442: resolution broadening is applied to T(E) after Beer-Lambert,
54/// not to σ(E) before.
55///
56/// Construct via `nereids_physics::transmission::broadened_cross_sections`,
57/// then wrap in `Arc` so the same precomputed data is shared read-only
58/// across all rayon worker threads.
59pub struct PrecomputedTransmissionModel {
60    /// Doppler-broadened cross-sections σ_D(E) per isotope, shape
61    /// \[n_isotopes\]\[n_grid_energies\].
62    ///
63    /// **The grid these σ live on is determined by
64    /// [`work_layout`](Self::work_layout):**
65    ///
66    /// * `work_layout` is `Some` (Gaussian resolution → auxiliary extended
67    ///   grid): σ live on the **working grid**, i.e.
68    ///   `work_layout.energies`, with `n_grid_energies ==
69    ///   work_layout.energies.len()`.  `evaluate()` / `analytical_jacobian()`
70    ///   apply Beer-Lambert + resolution on this working grid and extract the
71    ///   data points LAST via `work_layout.extract(..)` — matching
72    ///   `forward_model` (issue #608).
73    /// * `work_layout` is `None` (tabulated resolution, or no resolution): the
74    ///   working grid IS the data grid, so σ live on the **data grid**
75    ///   (`energies`), with `n_grid_energies == energies.len()`.  No extraction
76    ///   is needed and the surrogate fast paths + data-grid `resolution_plan`
77    ///   behave exactly as before.
78    pub cross_sections: Arc<Vec<Vec<f64>>>,
79    /// Mapping: `params[density_indices[i]]` is the density of isotope `i`.
80    ///
81    /// Wrapped in `Arc` so that parallel pixel loops can share one copy
82    /// via cheap reference-count increments instead of deep-cloning per pixel.
83    ///
84    /// Kept `pub` (not `pub(crate)`) because the Python bindings
85    /// (`nereids-python`) construct and access this field directly.
86    pub density_indices: Arc<Vec<usize>>,
87    /// Energy grid (eV), required for resolution broadening.
88    /// `None` when resolution is disabled — Beer-Lambert only.
89    pub energies: Option<Arc<Vec<f64>>>,
90    /// Instrument resolution parameters.
91    /// When `Some`, resolution broadening is applied to the total
92    /// transmission after Beer-Lambert in `evaluate()`.
93    pub instrument: Option<Arc<InstrumentParams>>,
94    /// Optional pre-built broadening plan for `(energies, resolution)`.
95    ///
96    /// When a caller builds the plan once (e.g. spatial dispatch for
97    /// a grid shared across every pixel) and passes it via
98    /// `with_resolution_plan`, `evaluate()` and `analytical_jacobian()`
99    /// skip the per-call kernel-interp / bracket / trap-weight work
100    /// and reduce each broadening call to a gather + multiply-add.
101    /// `None` ⇒ fall back to the per-call broadening path, byte-
102    /// identical output.
103    pub resolution_plan: Option<Arc<ResolutionPlan>>,
104    /// Optional sparse empirical cubature plan.
105    ///
106    /// When the plan is present AND its `target_energies` match this
107    /// model's energy grid AND `cubature.k() == n_density_params`
108    /// AND no temperature / energy-scale fitting is active, the
109    /// `evaluate()` / `analytical_jacobian()` fast path calls
110    /// `cubature.forward_and_jacobian(n)` directly instead of
111    /// `exp(-Σ n σ) + apply_resolution`.  Any guard failure falls
112    /// back to the exact path, so installing a plan cannot change
113    /// results unless every guard passes.
114    pub sparse_cubature_plan: Option<Arc<SparseEmpiricalCubaturePlan>>,
115    /// Optional scalar (k = 1) surrogate plan.
116    ///
117    /// Mutually exclusive with `sparse_cubature_plan` in practice —
118    /// the cubature dispatch fires only for `k ≥ 2` and the scalar
119    /// plan only for `k == 1`.  The type alias
120    /// `ScalarSurrogatePlan = ScalarChebyshevPlan` is kept as a
121    /// stable public name so a future scalar surrogate can swap in
122    /// without touching this field or any dispatch call site.
123    /// Chebyshev-in-density was picked over Lanczos Gauss
124    /// quadrature after a real-VENUS bench-off (Chebyshev won on
125    /// both the accuracy and wall-time axes; see
126    /// `nereids_physics::surrogate` module docs).
127    pub sparse_scalar_plan: Option<Arc<ScalarSurrogatePlan>>,
128    /// Working-grid layout matching [`cross_sections`](Self::cross_sections).
129    ///
130    /// Issue #608: when `cross_sections` is stored on the auxiliary extended
131    /// grid (Gaussian resolution), this maps the working grid back to the data
132    /// grid so `evaluate()` / `analytical_jacobian()` apply resolution on the
133    /// working grid and extract the data points last.  `None` ⇒ the working
134    /// grid is the data grid (tabulated / no resolution): Beer-Lambert and
135    /// resolution run directly on `energies` and no extraction is needed, which
136    /// keeps the surrogate fast paths and the data-grid `resolution_plan`
137    /// byte-identical to before.
138    pub work_layout: Option<Arc<transmission::WorkingGridLayout>>,
139}
140
141/// Deduplicate `density_indices` and return the distinct density-
142/// parameter indices **sorted ascending by value** — e.g.
143/// `[0,0,0,0,0,0]` (grouped) → `[0]`; `[0,1,2,3,4,5]` (ungrouped) →
144/// `[0,1,2,3,4,5]`; `[1,0,1]` (non-monotonic group layout) →
145/// `[0,1]` (NOT first-appearance order `[1,0]`).
146///
147/// **Why sorted-by-value, not first-appearance?** The cubature
148/// dispatch maps `n[j] = params[result[j]]` onto the cubature's
149/// j-th atom column.  The cubature was built from a σ stack
150/// indexed by density-param index (`sigmas[j * n_rows + ℓ] =
151/// σ_{param_j}(E'_ℓ)`) — so atom column `j` corresponds to
152/// density param `j`.  Using sorted-by-value output keeps the
153/// dispatched `params[result[j]]` aligned with `cubature.atoms()`
154/// at column `j` regardless of the user's `density_indices`
155/// ordering.  First-appearance order would swap columns for
156/// non-monotonic mappings, returning wrong transmissions and
157/// wrong Jacobians.
158fn density_param_indices(density_indices: &[usize]) -> Vec<usize> {
159    // `sort_unstable` + `dedup` is O(n log n) and avoids the O(n²)
160    // cost of repeated `Vec::contains` scans.  This runs on every
161    // `evaluate()` / `analytical_jacobian()` call, so the linear-
162    // scan version showed up in spatial-map profiling once the
163    // per-pixel cubature dispatch started firing.
164    let mut seen: Vec<usize> = density_indices.to_vec();
165    seen.sort_unstable();
166    seen.dedup();
167    seen
168}
169
170/// Check whether a cubature-based forward evaluation is eligible
171/// given the plan, the model's energy grid, the model's active
172/// resolution plan, and density-param structure.  Centralized so
173/// `evaluate`, `analytical_jacobian`, and both model types share a
174/// single predicate.
175///
176/// **Grid identity** (not just length) matters: a cached plan from a
177/// previous spatial call on a different grid with the same bin count
178/// would silently return forward/Jacobian values for the stale grid.
179/// We compare `plan.target_energies()` against the model's `energies`
180/// via `to_bits()` per element (same contract
181/// `apply_resolution_with_plan` already enforces).
182///
183/// **Tabulated-kernel tie**: the cubature fast path folds
184/// `apply_resolution*` into its atom sweep — skipping it when the
185/// model otherwise would have applied a Gaussian kernel is a
186/// silent wrong-answer path.  We require
187/// `matches!(instrument_resolution, ResolutionFunction::Tabulated(_))`
188/// so Gaussian-resolution models never hit the cubature path (a
189/// plan is only ever built against a tabulated kernel).
190///
191/// **Optional `resolution_plan` cross-check**: when a prebuilt
192/// `ResolutionPlan` is attached (e.g., via
193/// `spatial_map_typed`'s plan-hoist pathway), we additionally
194/// verify its grid matches the cubature plan's grid — defence-in-
195/// depth against a
196/// `with_precomputed_resolution_plan(plan_A) +
197/// with_precomputed_sparse_cubature_plan(plan_B_on_different_grid)`
198/// mis-configuration.  When no resolution plan is attached (the
199/// default on the single-spectrum entrypoint, where
200/// `fit_spectrum_typed` / `build_transmission_model` don't
201/// synthesize one), eligibility falls back to the cubature-plan
202/// grid check alone; this keeps the `with_precomputed_sparse_cubature_plan`
203/// API usable on the single-spectrum surface without the caller
204/// having to pre-build a matching `ResolutionPlan` just to unlock
205/// the fast path.
206///
207/// **Known caveat (same-grid kernel swap)**: if a caller rebuilds
208/// the tabulated resolution plan for a *different kernel* on the
209/// same energy grid without rebuilding the cubature, the grid
210/// bit-check here passes but the atom weights still encode the
211/// OLD operator.  Guarding against this requires a kernel
212/// fingerprint on the cubature plan, which is not implemented
213/// here.  Upstream callers are
214/// responsible for clearing the cubature when they swap kernels;
215/// in spatial dispatch this is enforced by
216/// `UnifiedFitConfig::with_precomputed_cross_sections` /
217/// `with_precomputed_base_xs` / `with_groups` all clearing the
218/// cached cubature (see pipeline.rs), so a refit through the
219/// standard surface cannot hit this case.
220/// Check whether a scalar (k = 1) surrogate plan is eligible given
221/// the model's energy grid, active tabulated resolution,
222/// attached `ResolutionPlan`, current σ row, and
223/// `n_density_params == 1`.  Parallels [`cubature_eligible`] for
224/// the multi-isotope path on grid-identity + `Tabulated(_)` guard,
225/// and **additionally** enforces content identity via the
226/// source-`ResolutionPlan` `Arc::ptr_eq` check and a σ
227/// fingerprint — closing a same-grid stale-plan correctness hole:
228/// a plan built from different σ or a different kernel but
229/// attached on the same energy grid must never dispatch the
230/// surrogate.
231fn scalar_eligible(
232    plan: &ScalarSurrogatePlan,
233    energies: &[f64],
234    instrument_resolution: &ResolutionFunction,
235    resolution_plan: Option<&Arc<ResolutionPlan>>,
236    sigma_row: &[f64],
237    n_density_params: usize,
238) -> bool {
239    if n_density_params != 1 {
240        return false;
241    }
242    if plan.len() != energies.len() {
243        return false;
244    }
245    if !matches!(instrument_resolution, ResolutionFunction::Tabulated(_)) {
246        return false;
247    }
248    let plan_grid = plan.target_energies();
249    for (e_cur, e_plan) in energies.iter().zip(plan_grid) {
250        if e_cur.to_bits() != e_plan.to_bits() {
251            return false;
252        }
253    }
254    // Source-`ResolutionPlan` identity via `Arc::ptr_eq` — O(1)
255    // check that the plan was built from the SAME resolution
256    // kernel the model is currently using.  The grid-only
257    // check was insufficient: a plan built for a different
258    // tabulated kernel on an identical grid would silently
259    // dispatch and return transmissions shifted by ~0.13
260    // absolute (measured).  Requiring the model to attach the exact same
261    // `Arc<ResolutionPlan>` the scalar plan was built from
262    // closes that hole.
263    let Some(model_plan) = resolution_plan else {
264        return false;
265    };
266    if !Arc::ptr_eq(model_plan, plan.source_resolution_plan()) {
267        return false;
268    }
269    // Transitive grid-identity on `resolution_plan` (retained from
270    // the previous check — catches an `Arc::ptr_eq`-true pair whose
271    // inner grid has been mutated out from under us, e.g. a
272    // `Mutex<ResolutionPlan>` unsafe pattern; defence-in-depth).
273    if model_plan.target_energies().len() != energies.len() {
274        return false;
275    }
276    for (e_cur, e_res) in energies.iter().zip(model_plan.target_energies()) {
277        if e_cur.to_bits() != e_res.to_bits() {
278            return false;
279        }
280    }
281    // σ fingerprint: same-grid-different-σ would otherwise pass
282    // every grid check.  FNV-1a-64 over `to_bits()` is fast
283    // (~3 µs for 3471-point VENUS grid) and cryptographically
284    // sufficient for catching unintentional mismatch; matched-bit
285    // collisions would require an adversarial σ, which isn't a
286    // threat model here (the wrong-σ bug surfaces from
287    // copy-paste caller errors).
288    if nereids_physics::surrogate::fingerprint_f64_slice(sigma_row) != plan.sigma_fingerprint() {
289        return false;
290    }
291    true
292}
293
294/// Check whether the scalar iterate `n` is inside the surrogate's
295/// recorded training box `[0, train_max]` — **strict** `n ≤ train_max`,
296/// unlike the cubature's 1.5× tolerance.
297///
298/// Chebyshev-in-density is a polynomial interpolant.  Inside
299/// `[0, n_max]` it is exact at the M = 16 nodes and tight (≤ 1e-15
300/// rel err) between them; outside, the interpolant diverges
301/// exponentially in `(n - n_max) / n_max` — measured:
302/// **73 % relative error at `1.5 × n_max`** and catastrophic
303/// divergence beyond — exactly the "silently wrong forward"
304/// failure mode that would corrupt a fit without the solver
305/// ever seeing an error flag.
306///
307/// The cubature's 1.5× tolerance is safe because LP-matched atoms
308/// moment-match the σ-pushforward measure and generalize gracefully
309/// past the box; Chebyshev polynomials do not.  So the scalar
310/// box is a **hard boundary**: the solver must either stay inside
311/// or trigger the exact-path fallback.  Because the spatial build
312/// site sets `n_max = 2 × initial_density`, the initial iterate
313/// sits at 50 % of the box — with plenty of room for solver
314/// exploration up to 2× the initial density before the guard
315/// fires.
316fn scalar_density_within_box(plan: &ScalarSurrogatePlan, n: f64) -> bool {
317    let Some(train_max) = plan.density_box() else {
318        return true;
319    };
320    if !n.is_finite() || n < 0.0 {
321        return false;
322    }
323    n <= train_max
324}
325
326/// Check whether the current density iterate `n` is inside the
327/// training region recorded on the cubature plan, with a 50 %
328/// expansion tolerance to avoid thrashing at the box boundary.
329/// When the plan has no recorded box, accepts unconditionally
330/// (caller is responsible; legacy code path).
331///
332/// Returns `false` when any component escapes the tolerance-
333/// expanded box OR is negative, OR is not finite.  Without this,
334/// a spatial fit whose per-pixel
335/// optimum drifts beyond `2 × initial_densities` silently runs the
336/// surrogate out of domain.
337fn density_within_box(plan: &SparseEmpiricalCubaturePlan, n: &[f64]) -> bool {
338    let Some(train_max) = plan.density_box() else {
339        // No box recorded — caller accepts the risk.
340        return true;
341    };
342    if train_max.len() != n.len() {
343        return false;
344    }
345    const TOLERANCE: f64 = 1.5; // 50 % slack above train_max
346    for (&n_i, &max_i) in n.iter().zip(train_max) {
347        if !n_i.is_finite() || n_i < 0.0 {
348            return false;
349        }
350        if n_i > max_i * TOLERANCE {
351            return false;
352        }
353    }
354    true
355}
356
357fn cubature_eligible(
358    plan: &SparseEmpiricalCubaturePlan,
359    energies: &[f64],
360    instrument_resolution: &ResolutionFunction,
361    resolution_plan: Option<&ResolutionPlan>,
362    n_density_params: usize,
363) -> bool {
364    // k ≥ 2: the scalar k=1 branch handles the grouped case.
365    if n_density_params < 2 {
366        return false;
367    }
368    if plan.k() != n_density_params {
369        return false;
370    }
371    if plan.len() != energies.len() {
372        return false;
373    }
374    // Gaussian-resolution models must NOT hit the cubature path:
375    // the cubature was built against a TabulatedResolution kernel
376    // (it's the only kernel `ResolutionPlan::compile_to_matrix`
377    // accepts), so firing it on a Gaussian-active model would
378    // silently replace Gaussian broadening with a tabulated
379    // surrogate.
380    if !matches!(instrument_resolution, ResolutionFunction::Tabulated(_)) {
381        return false;
382    }
383    // Per-element `to_bits()` grid identity check catches `-0.0` vs
384    // `+0.0` and NaN-bit differences that float `==` silently
385    // accepts or rejects.  The cubature plan's own grid is the
386    // primary reference (atoms are indexed against it).
387    let cub_grid = plan.target_energies();
388    for (e_cur, e_cub) in energies.iter().zip(cub_grid) {
389        if e_cur.to_bits() != e_cub.to_bits() {
390            return false;
391        }
392    }
393    // Defense-in-depth: when a ResolutionPlan is ALSO attached,
394    // verify transitive grid identity.  Catches the
395    // `with_precomputed_resolution_plan(plan_A) +
396    // with_precomputed_sparse_cubature_plan(plan_B_on_different_grid)`
397    // mis-configuration case.  When no resolution plan is attached
398    // (typical single-spectrum entrypoint —
399    // `fit_spectrum_typed` / `build_transmission_model` don't
400    // synthesize one by default), the in-model resolution broaden
401    // path falls back to per-call `apply_resolution` and the
402    // cubature's self-check above is the grid guard.  An earlier
403    // "resolution_plan.is_some() required" rule was over-strict and
404    // silently disabled the fast path on the single-spectrum
405    // surface.
406    if let Some(res_plan) = resolution_plan {
407        if res_plan.target_energies().len() != energies.len() {
408            return false;
409        }
410        let res_grid = res_plan.target_energies();
411        for (e_cur, e_res) in energies.iter().zip(res_grid) {
412            if e_cur.to_bits() != e_res.to_bits() {
413                return false;
414            }
415        }
416    }
417    true
418}
419
420impl PrecomputedTransmissionModel {
421    /// Working-grid energies for resolution broadening (issue #608).
422    ///
423    /// Returns the auxiliary extended grid when `work_layout` is set (Gaussian
424    /// resolution), otherwise the data grid (`energies`).  Returns `None` only
425    /// when no instrument is configured (Beer-Lambert-only path).
426    fn work_energies(&self) -> Option<&[f64]> {
427        match (&self.work_layout, &self.energies) {
428            (Some(layout), _) => Some(layout.energies.as_slice()),
429            (None, Some(energies)) => Some(energies.as_slice()),
430            (None, None) => None,
431        }
432    }
433
434    /// Extract the data-grid points from a working-grid spectrum (issue #608).
435    ///
436    /// When `work_layout` is `None` the working grid IS the data grid, so this
437    /// is the identity (a plain clone).
438    fn extract_data_points(&self, working: &[f64]) -> Vec<f64> {
439        match &self.work_layout {
440            Some(layout) => layout.extract(working),
441            None => working.to_vec(),
442        }
443    }
444}
445
446impl FitModel for PrecomputedTransmissionModel {
447    fn evaluate(&self, params: &[f64]) -> Result<Vec<f64>, FittingError> {
448        if self.cross_sections.is_empty() {
449            return Err(FittingError::InvalidConfig(
450                "PrecomputedTransmissionModel.cross_sections must not be empty".into(),
451            ));
452        }
453        let n_e = self.cross_sections[0].len();
454
455        // Cubature fast path: when the plan is installed, matches
456        // the grid + isotope count, and instrument resolution is
457        // enabled (cubature folds both `exp(-Σ n σ)` and `apply_R`
458        // into a single per-row atom sweep).
459        if let (Some(cubature), Some(inst), Some(energies)) =
460            (&self.sparse_cubature_plan, &self.instrument, &self.energies)
461        {
462            let params_indices = density_param_indices(&self.density_indices);
463            if cubature_eligible(
464                cubature,
465                energies,
466                &inst.resolution,
467                self.resolution_plan.as_deref(),
468                params_indices.len(),
469            ) {
470                let n: Vec<f64> = params_indices.iter().map(|&i| params[i]).collect();
471                if density_within_box(cubature, &n) {
472                    return Ok(cubature.forward(&n));
473                }
474                // Density escaped the training box — fall through
475                // to the exact path (cubature accuracy degrades
476                // quickly outside the trained region).
477            }
478        }
479
480        // Scalar (k = 1) surrogate fast path — same eligibility
481        // stack as the cubature, gated on `n_density_params == 1`.
482        // The content-identity guards
483        // (σ-fingerprint + Arc::ptr_eq on source resolution plan)
484        // close the same-grid stale-plan hole.
485        if let (Some(scalar), Some(inst), Some(energies)) =
486            (&self.sparse_scalar_plan, &self.instrument, &self.energies)
487        {
488            let params_indices = density_param_indices(&self.density_indices);
489            // Only fire when the σ stack is the single collapsed
490            // row the scalar plan was built from (spatial's
491            // post-grouping shape).  Non-collapsed k = 1 flows
492            // cannot safely dispatch here.
493            if self.cross_sections.len() == 1
494                && self.density_indices.len() == 1
495                && self.density_indices[0] == params_indices[0]
496                && scalar_eligible(
497                    scalar,
498                    energies,
499                    &inst.resolution,
500                    self.resolution_plan.as_ref(),
501                    &self.cross_sections[0],
502                    params_indices.len(),
503                )
504            {
505                let n = params[params_indices[0]];
506                if scalar_density_within_box(scalar, n) {
507                    return Ok(scalar.forward_scalar(n));
508                }
509            }
510        }
511
512        // Beer-Lambert on the WORKING grid (issue #608): `cross_sections` are
513        // stored on the working grid (auxiliary extended grid for Gaussian
514        // resolution; the data grid for tabulated / no resolution), so `n_e`
515        // is the working-grid length.
516        let mut neg_opt = vec![0.0f64; n_e];
517        // #109.1: No density > 0 guard — let Beer-Lambert handle all densities
518        // naturally.  exp(−n·σ) is well-defined for negative n (gives T > 1,
519        // which is unphysical but the optimizer will reject it via chi2
520        // increase).  Removing the guard makes evaluate() consistent with
521        // the analytical Jacobian, which always computes ∂T/∂n = −σ·T
522        // regardless of the sign of n.
523        for (i, xs) in self.cross_sections.iter().enumerate() {
524            let density = params[self.density_indices[i]];
525            for (j, &sigma) in xs.iter().enumerate() {
526                neg_opt[j] -= density * sigma;
527            }
528        }
529        let transmission: Vec<f64> = neg_opt.iter().map(|&d| d.exp()).collect();
530
531        // Issue #442 + #608: apply resolution broadening to the total
532        // transmission AFTER Beer-Lambert, on the WORKING grid, then extract
533        // the data points LAST.  When `work_layout` is `None` (tabulated / no
534        // resolution) the working grid IS the data grid (`self.energies`), the
535        // extraction is the identity, and the data-grid `resolution_plan` still
536        // matches — byte-identical to the pre-#608 path.
537        // Resolution applies iff there is an instrument AND a working grid to
538        // apply it on (`work_energies()` = the layout grid when present, else the
539        // data grid).  `evaluate` and `analytical_jacobian` share this exact
540        // guard so the two paths cannot diverge (issue #608).
541        if let (Some(inst), Some(work_energies)) = (&self.instrument, self.work_energies()) {
542            let t_broadened = resolution::apply_resolution_with_plan(
543                self.resolution_plan.as_deref(),
544                work_energies,
545                &transmission,
546                &inst.resolution,
547            )
548            .map_err(|e| FittingError::EvaluationFailed(format!("resolution broadening: {e}")))?;
549            Ok(self.extract_data_points(&t_broadened))
550        } else {
551            Ok(self.extract_data_points(&transmission))
552        }
553    }
554
555    /// Analytical Jacobian for the Beer-Lambert transmission model.
556    ///
557    /// Without resolution:
558    ///   T(E) = exp(-Σᵢ nᵢ · σᵢ(E))
559    ///   ∂T/∂nᵢ = -σᵢ(E) · T(E)
560    ///
561    /// With resolution (R is a linear operator):
562    ///   T_obs(E) = R\[T\](E) = R\[exp(-Σᵢ nᵢ · σᵢ)\](E)
563    ///   ∂T_obs/∂nᵢ = R\[-σᵢ(E) · T(E)\]
564    ///
565    /// For grouped isotopes sharing density parameter N_g:
566    ///   ∂T_obs/∂N_g = R\[-(Σ_{i∈g} σᵢ(E)) · T(E)\]
567    fn analytical_jacobian(
568        &self,
569        params: &[f64],
570        free_param_indices: &[usize],
571        y_current: &[f64],
572    ) -> Option<FlatMatrix> {
573        let n_e = if self.cross_sections.is_empty() {
574            return None;
575        } else {
576            self.cross_sections[0].len()
577        };
578        let n_free = free_param_indices.len();
579
580        // Cubature fast path: same eligibility as `evaluate()` plus
581        // the requirement that every free param is a density param
582        // (cubature can't produce Jacobian columns for non-density
583        // params like background / normalization, which are the
584        // calling layer's responsibility).
585        if let (Some(cubature), Some(inst), Some(energies)) =
586            (&self.sparse_cubature_plan, &self.instrument, &self.energies)
587        {
588            let params_indices = density_param_indices(&self.density_indices);
589            if cubature_eligible(
590                cubature,
591                energies,
592                &inst.resolution,
593                self.resolution_plan.as_deref(),
594                params_indices.len(),
595            ) {
596                // Map each free param to its column in the cubature
597                // Jacobian.  `None` for any free param that isn't a
598                // density param → fall through to the exact path.
599                // Wrappers (`NormalizedTransmissionModel`,
600                // `TransmissionKLBackgroundModel`) ensure only density
601                // slots reach this layer; non-density free params here
602                // would indicate a wrapper bypass.
603                let col_map: Option<Vec<usize>> = free_param_indices
604                    .iter()
605                    .map(|&fp| params_indices.iter().position(|&i| i == fp))
606                    .collect();
607                if let Some(col_map) = col_map {
608                    let n: Vec<f64> = params_indices.iter().map(|&i| params[i]).collect();
609                    if density_within_box(cubature, &n) {
610                        let (_t, jac_flat) = cubature.forward_and_jacobian(&n);
611                        // jac_flat[i * k + ell] = ∂T_i / ∂n_ell
612                        let k = params_indices.len();
613                        let mut jacobian = FlatMatrix::zeros(n_e, n_free);
614                        for (col, &ell) in col_map.iter().enumerate() {
615                            for i in 0..n_e {
616                                *jacobian.get_mut(i, col) = jac_flat[i * k + ell];
617                            }
618                        }
619                        return Some(jacobian);
620                    }
621                    // Density outside box → fall through to exact.
622                }
623            }
624        }
625
626        // Scalar (k = 1) surrogate Jacobian fast path.  For a
627        // scalar fit `free_param_indices = [0]`, so
628        // the Jacobian has one column.
629        if let (Some(scalar), Some(inst), Some(energies)) =
630            (&self.sparse_scalar_plan, &self.instrument, &self.energies)
631        {
632            let params_indices = density_param_indices(&self.density_indices);
633            if self.cross_sections.len() == 1
634                && self.density_indices.len() == 1
635                && self.density_indices[0] == params_indices[0]
636                && scalar_eligible(
637                    scalar,
638                    energies,
639                    &inst.resolution,
640                    self.resolution_plan.as_ref(),
641                    &self.cross_sections[0],
642                    params_indices.len(),
643                )
644                && free_param_indices.len() == 1
645                && free_param_indices[0] == params_indices[0]
646            {
647                let n = params[params_indices[0]];
648                if scalar_density_within_box(scalar, n) {
649                    let (_t, dt) = scalar.forward_and_derivative_scalar(n);
650                    let mut jacobian = FlatMatrix::zeros(n_e, 1);
651                    for (i, &v) in dt.iter().enumerate() {
652                        *jacobian.get_mut(i, 0) = v;
653                    }
654                    return Some(jacobian);
655                }
656            }
657        }
658
659        // For each free parameter, sum the cross-sections of every isotope
660        // tied to that parameter index.  σ (and the sums) are on the WORKING
661        // grid (issue #608), so `n_e` is the working-grid length.
662        //   ∂T/∂N_g = -(Σ_{iso∈g} σ_iso(E)) · T(E)
663        let fp_xs_sums: Vec<Vec<f64>> = free_param_indices
664            .iter()
665            .map(|&fp_idx| {
666                let mut sum = vec![0.0f64; n_e];
667                for (iso, &di) in self.density_indices.iter().enumerate() {
668                    if di == fp_idx {
669                        for (j, &sigma) in self.cross_sections[iso].iter().enumerate() {
670                            sum[j] += sigma;
671                        }
672                    }
673                }
674                sum
675            })
676            .collect();
677
678        // The Jacobian has one row per DATA point; `y_current` is on the data
679        // grid.  When resolution is enabled the inner derivative is formed on
680        // the working grid, resolution-broadened there, and the data points
681        // extracted last (issue #608).
682        let n_data = y_current.len();
683
684        // When resolution is enabled, we need the UNRESOLVED T(E) = exp(-Σnσ)
685        // on the WORKING grid to form the inner derivative -σ·T, then apply
686        // resolution on the working grid and extract the data points.
687        // y_current is T_obs = R[T] on the DATA grid, which is NOT the same.
688        // Same resolution guard as `evaluate` (issue #608) so the two paths
689        // agree by construction; the else branch is the no-resolution Jacobian.
690        if let (Some(inst), Some(work_energies)) = (&self.instrument, self.work_energies()) {
691            // Recompute unresolved T on the working grid from σ and params.
692            let mut neg_opt = vec![0.0f64; n_e];
693            for (i, xs) in self.cross_sections.iter().enumerate() {
694                let density = params[self.density_indices[i]];
695                for (j, &sigma) in xs.iter().enumerate() {
696                    neg_opt[j] -= density * sigma;
697                }
698            }
699            let t_unresolved: Vec<f64> = neg_opt.iter().map(|&d| d.exp()).collect();
700
701            // ∂T_obs/∂N_g = extract(R[-σ_sum(E) · T_unresolved(E)])
702            let mut jacobian = FlatMatrix::zeros(n_data, n_free);
703            for (col, xs_sum) in fp_xs_sums.iter().enumerate() {
704                let inner_deriv: Vec<f64> =
705                    (0..n_e).map(|i| -xs_sum[i] * t_unresolved[i]).collect();
706                let resolved_deriv = resolution::apply_resolution_with_plan(
707                    self.resolution_plan.as_deref(),
708                    work_energies,
709                    &inner_deriv,
710                    &inst.resolution,
711                )
712                .ok()?;
713                let resolved_deriv = self.extract_data_points(&resolved_deriv);
714                for (i, &val) in resolved_deriv.iter().enumerate() {
715                    *jacobian.get_mut(i, col) = val;
716                }
717            }
718            Some(jacobian)
719        } else {
720            // No resolution → no auxiliary grid: the working grid IS the data
721            // grid (`n_e == n_data`), and y_current IS T(E) directly.
722            //   ∂T/∂N_g = -σ_sum(E) · T(E)
723            let mut jacobian = FlatMatrix::zeros(n_data, n_free);
724            for i in 0..n_data {
725                for (j, xs_sum) in fp_xs_sums.iter().enumerate() {
726                    *jacobian.get_mut(i, j) = -xs_sum[i] * y_current[i];
727                }
728            }
729            Some(jacobian)
730        }
731    }
732}
733
734/// Forward model for fitting isotopic areal densities from transmission data.
735///
736/// The model computes T(E) for a set of isotopes with variable areal densities.
737/// Each isotope's resonance data and the energy grid are fixed; only the
738/// areal densities are adjusted during fitting.
739///
740/// Optionally, the sample temperature can also be fitted by setting
741/// `temperature_index` to the parameter slot holding the temperature value.
742/// When `temperature_index` is `Some(idx)`, the Doppler broadening kernel
743/// is recomputed at `params[idx]` when the temperature changes (cached
744/// across calls at the same temperature), and the analytical Jacobian
745/// provides density columns directly plus a single FD column for temperature.
746///
747/// `instrument` uses `Arc` so that parallel pixel loops can share one copy
748/// of a potentially large tabulated resolution kernel via cheap
749/// reference-count increments instead of deep-cloning per pixel.
750pub struct TransmissionFitModel {
751    /// Energy grid (eV), ascending.
752    energies: Vec<f64>,
753    /// Resonance data for each isotope.
754    resonance_data: Vec<ResonanceData>,
755    /// Sample temperature in Kelvin (used when `temperature_index` is `None`).
756    temperature_k: f64,
757    /// Optional instrument resolution parameters (Arc-shared for parallel use).
758    instrument: Option<Arc<InstrumentParams>>,
759    /// Index mapping: which `params` indices correspond to areal densities.
760    /// params[density_indices[i]] = areal density of isotope i.
761    ///
762    /// Uses `Vec<usize>` (not `Arc<Vec<usize>>`) because `TransmissionFitModel`
763    /// is constructed fresh per pixel (via `fit_spectrum`) and never shared
764    /// across threads.  `PrecomputedTransmissionModel` uses `Arc<Vec<usize>>`
765    /// for its density_indices because it _is_ shared across rayon workers.
766    density_indices: Vec<usize>,
767    /// Fractional ratio of each member isotope within its group.
768    /// For ungrouped isotopes, all values are 1.0.
769    /// When groups are active: `effective_density_i = params[density_indices[i]] * density_ratios[i]`
770    density_ratios: Vec<f64>,
771    /// If `Some(idx)`, `params[idx]` is treated as the sample temperature (K)
772    /// and included as a free parameter in the fit. The Doppler broadening
773    /// kernel is recomputed at each `evaluate()` call.
774    temperature_index: Option<usize>,
775    /// Cached unbroadened (Reich-Moore) cross-sections, computed once in
776    /// `new()` when `temperature_index` is `Some`. Eliminates redundant
777    /// O(N_energy × N_resonances) computation on every `evaluate()` call.
778    /// Wrapped in `Arc` so `spatial_map` can share a single allocation across
779    /// all per-pixel `TransmissionFitModel` instances without deep cloning.
780    base_xs: Option<Arc<Vec<Vec<f64>>>>,
781    /// Cached broadened cross-sections from the last `evaluate()` call, on the
782    /// **working grid** (auxiliary extended grid when Gaussian resolution is
783    /// active, else the data grid).  Used by `analytical_jacobian()` to provide
784    /// density columns without rebroadening AND to build the inner derivative
785    /// `−σ·T` on the working grid before resolution + data-point extraction
786    /// (issue #608).  Interior mutability via `RefCell` is needed because
787    /// `FitModel::evaluate` takes `&self`.  Safe because `TransmissionFitModel`
788    /// is constructed per-pixel and never shared across threads.
789    cached_broadened_xs: RefCell<Option<Rc<Vec<Vec<f64>>>>>,
790    /// Cached analytical temperature derivative ∂σ/∂T, on the **working grid**,
791    /// computed on-demand by `analytical_jacobian()` when the temperature
792    /// column is needed.  Invalidated when temperature changes (cleared in
793    /// `evaluate()`).
794    cached_dxs_dt: RefCell<Option<Rc<Vec<Vec<f64>>>>>,
795    /// Working-grid layout (energies + data-index map) matching
796    /// `cached_broadened_xs` / `cached_dxs_dt`.  Resolution broadening is
797    /// applied on `layout.energies` and the data points are extracted last
798    /// (issue #608).  Set in `evaluate()` alongside the broadened σ cache.
799    cached_work_layout: RefCell<Option<Rc<transmission::WorkingGridLayout>>>,
800    /// Temperature at which `cached_broadened_xs` was computed.
801    /// `Cell` is sufficient because `f64` is `Copy`.
802    cached_temperature: Cell<f64>,
803    /// Optional prebuilt resolution plan for [`Self::energies`].
804    ///
805    /// When a caller (typically spatial dispatch) builds the plan
806    /// once for a shared grid, passing it here lets every per-pixel
807    /// `evaluate()` / `analytical_jacobian()` call reuse the hoisted
808    /// TOF / kernel-interp / bracket work.  `None` ⇒ per-call
809    /// broadening (same output as pre-plan main).
810    resolution_plan: Option<Arc<ResolutionPlan>>,
811    /// Optional sparse empirical cubature plan.
812    ///
813    /// See [`PrecomputedTransmissionModel::sparse_cubature_plan`]
814    /// for the dispatch contract.  In this per-pixel model the
815    /// cubature is additionally constrained: if `temperature_index`
816    /// is `Some` or the temperature changes between evaluate calls,
817    /// the σ the cubature was built against becomes stale so the
818    /// dispatch silently falls back.
819    sparse_cubature_plan: Option<Arc<SparseEmpiricalCubaturePlan>>,
820    /// Optional scalar (k = 1) surrogate plan.
821    /// Parallel to `sparse_cubature_plan` but dispatches only for
822    /// `n_density_params == 1`.
823    sparse_scalar_plan: Option<Arc<ScalarSurrogatePlan>>,
824}
825
826impl TransmissionFitModel {
827    /// Create a validated `TransmissionFitModel`.
828    ///
829    /// When `external_base_xs` is `Some`, uses those precomputed unbroadened
830    /// cross-sections instead of computing them (expensive Reich-Moore).
831    /// `spatial_map` precomputes once for all pixels and passes them here.
832    ///
833    /// # Errors
834    /// Returns `FittingError::InvalidConfig` if `temperature_index` overlaps
835    /// with `density_indices`, or if `external_base_xs` has a mismatched shape.
836    pub fn new(
837        energies: Vec<f64>,
838        resonance_data: Vec<ResonanceData>,
839        temperature_k: f64,
840        instrument: Option<Arc<InstrumentParams>>,
841        density_mapping: (Vec<usize>, Vec<f64>),
842        temperature_index: Option<usize>,
843        external_base_xs: Option<Arc<Vec<Vec<f64>>>>,
844    ) -> Result<Self, FittingError> {
845        let (density_indices, density_ratios) = density_mapping;
846        if density_indices.len() != resonance_data.len() {
847            return Err(FittingError::InvalidConfig(format!(
848                "density_indices has {} entries but resonance_data has {}",
849                density_indices.len(),
850                resonance_data.len(),
851            )));
852        }
853        if density_ratios.len() != resonance_data.len() {
854            return Err(FittingError::InvalidConfig(format!(
855                "density_ratios has {} entries but resonance_data has {}",
856                density_ratios.len(),
857                resonance_data.len(),
858            )));
859        }
860        if let Some(ti) = temperature_index
861            && density_indices.contains(&ti)
862        {
863            return Err(FittingError::InvalidConfig(
864                "temperature_index must not overlap with density_indices".into(),
865            ));
866        }
867        // Validate external base XS shape before accepting.
868        if let Some(ref xs) = external_base_xs {
869            if xs.len() != resonance_data.len() {
870                return Err(FittingError::InvalidConfig(format!(
871                    "external_base_xs has {} isotopes but resonance_data has {}",
872                    xs.len(),
873                    resonance_data.len(),
874                )));
875            }
876            for (i, row) in xs.iter().enumerate() {
877                if row.len() != energies.len() {
878                    return Err(FittingError::InvalidConfig(format!(
879                        "external_base_xs[{i}] has {} energies but expected {}",
880                        row.len(),
881                        energies.len(),
882                    )));
883                }
884            }
885        }
886        let base_xs = match external_base_xs {
887            Some(xs) => Some(xs),
888            None if temperature_index.is_some() => Some(Arc::new(
889                transmission::unbroadened_cross_sections(&energies, &resonance_data, None)
890                    .map_err(|e| {
891                        FittingError::InvalidConfig(format!(
892                            "failed to compute unbroadened cross-sections: {e}"
893                        ))
894                    })?,
895            )),
896            None => None,
897        };
898        Ok(Self {
899            energies,
900            resonance_data,
901            temperature_k,
902            instrument,
903            density_indices,
904            density_ratios,
905            temperature_index,
906            base_xs,
907            cached_broadened_xs: RefCell::new(None),
908            cached_dxs_dt: RefCell::new(None),
909            cached_work_layout: RefCell::new(None),
910            cached_temperature: Cell::new(f64::NAN),
911            resolution_plan: None,
912            sparse_cubature_plan: None,
913            sparse_scalar_plan: None,
914        })
915    }
916
917    /// Attach a prebuilt resolution plan for the model's energy grid.
918    ///
919    /// Safe to call before any `evaluate()`.  Caller contract:
920    /// `plan.target_energies() == energies` — violating this will
921    /// fail on the first broadening call, either via a length
922    /// mismatch or, for a different same-length grid,
923    /// `ResolutionError::PlanGridMismatch`.
924    #[must_use]
925    pub fn with_resolution_plan(mut self, plan: Option<Arc<ResolutionPlan>>) -> Self {
926        self.resolution_plan = plan;
927        self
928    }
929
930    /// Attach a prebuilt sparse empirical cubature plan.  See
931    /// [`PrecomputedTransmissionModel::sparse_cubature_plan`] for the
932    /// dispatch conditions.
933    #[must_use]
934    pub fn with_sparse_cubature_plan(
935        mut self,
936        plan: Option<Arc<SparseEmpiricalCubaturePlan>>,
937    ) -> Self {
938        self.sparse_cubature_plan = plan;
939        self
940    }
941
942    /// Attach a prebuilt scalar (k = 1) surrogate plan.  See
943    /// [`PrecomputedTransmissionModel::sparse_scalar_plan`] for the
944    /// dispatch conditions.
945    #[must_use]
946    pub fn with_sparse_scalar_plan(mut self, plan: Option<Arc<ScalarSurrogatePlan>>) -> Self {
947        self.sparse_scalar_plan = plan;
948        self
949    }
950}
951
952impl FitModel for TransmissionFitModel {
953    fn evaluate(&self, params: &[f64]) -> Result<Vec<f64>, FittingError> {
954        debug_assert!(
955            self.density_indices.iter().all(|&i| i < params.len()),
956            "density_indices out of bounds for params (len={})",
957            params.len(),
958        );
959        debug_assert!(
960            self.temperature_index.is_none_or(|i| i < params.len()),
961            "temperature_index out of bounds for params (len={})",
962            params.len(),
963        );
964
965        // Cubature fast path: plan present, resolution on, no
966        // temperature fit (σ the cubature was built against must not
967        // change at runtime).  k=1 grouped case and per-isotope T-fit
968        // falls through to the exact path.
969        if let (Some(cubature), Some(inst)) = (&self.sparse_cubature_plan, &self.instrument)
970            && self.temperature_index.is_none()
971        {
972            let params_indices = density_param_indices(&self.density_indices);
973            if cubature_eligible(
974                cubature,
975                &self.energies,
976                &inst.resolution,
977                self.resolution_plan.as_deref(),
978                params_indices.len(),
979            ) {
980                // Caller contract: for grouped fits the cubature
981                // was built with σ already aggregated by ratios
982                // (`σ_group_j = Σ_{i ∈ group_j} ratio_i · σ_i`),
983                // so the online `forward(n)` receives only the
984                // per-group density vector and multiplies by the
985                // pre-aggregated atoms internally.
986                let n: Vec<f64> = params_indices.iter().map(|&i| params[i]).collect();
987                if density_within_box(cubature, &n) {
988                    return Ok(cubature.forward(&n));
989                }
990                // Density escaped training box → fall through.
991            }
992        }
993
994        // Scalar (k = 1) surrogate fast path was removed from this
995        // model: `TransmissionFitModel`'s on-the-fly σ compute
996        // couldn't be cheaply fingerprint-checked against the
997        // plan's σ, leaving a same-grid stale-plan correctness
998        // hole.  Production spatial dispatch attaches scalar plans
999        // to [`PrecomputedTransmissionModel`] (via
1000        // `UnifiedFitConfig::with_precomputed_cross_sections` +
1001        // `with_precomputed_sparse_scalar_plan`), which DOES
1002        // enforce σ-fingerprint + Arc::ptr_eq guards.  The
1003        // `sparse_scalar_plan` field and setter remain here for
1004        // API consistency with `PrecomputedTransmissionModel`, but
1005        // this model will always fall through to the exact path.
1006
1007        let temperature_k = match self.temperature_index {
1008            Some(idx) => params[idx],
1009            None => self.temperature_k,
1010        };
1011
1012        if let Some(ref base_xs) = self.base_xs {
1013            // Fast path: reuse cached unbroadened XS, only redo Doppler + Beer-Lambert.
1014            // Validate temperature (same rules as SampleParams::new in the slow path)
1015            // so the optimizer can't silently evaluate an unphysical model.
1016            if !temperature_k.is_finite() || temperature_k < 0.0 {
1017                return Err(FittingError::EvaluationFailed(format!(
1018                    "Invalid temperature: {temperature_k} K (must be finite and non-negative)"
1019                )));
1020            }
1021
1022            // Compute broadened XS on the WORKING grid (or reuse cache if
1023            // temperature unchanged).  Caching avoids redundant Doppler
1024            // broadening on rejected LM steps (same T, different lambda) and
1025            // enables analytical_jacobian() to read the broadened σ for the
1026            // density columns AND to build the inner derivative on the same
1027            // working grid.
1028            //
1029            // Issue #608: Doppler + Beer-Lambert + resolution all run on the
1030            // working grid (auxiliary extended grid when Gaussian resolution is
1031            // active), with the data points extracted LAST — matching
1032            // forward_model.  The previous cached path collapsed σ to the
1033            // coarse data grid before resolution, degrading the convolution.
1034            //
1035            // Derivative ∂σ/∂T is computed on-demand in analytical_jacobian(),
1036            // NOT here — evaluate() is called many times during line search
1037            // trials, and the derivative overhead would dominate.
1038            let (broadened_xs, layout) = if (temperature_k - self.cached_temperature.get()).abs()
1039                < 1e-15
1040                && self.cached_broadened_xs.borrow().is_some()
1041            {
1042                (
1043                    Rc::clone(self.cached_broadened_xs.borrow().as_ref().unwrap()),
1044                    Rc::clone(self.cached_work_layout.borrow().as_ref().unwrap()),
1045                )
1046            } else {
1047                let working = transmission::broadened_cross_sections_from_base_on_working_grid(
1048                    &self.energies,
1049                    base_xs,
1050                    &self.resonance_data,
1051                    temperature_k,
1052                    self.instrument.as_deref(),
1053                )
1054                .map_err(|e| FittingError::EvaluationFailed(e.to_string()))?;
1055                let xs = Rc::new(working.sigma);
1056                let layout = Rc::new(working.layout);
1057                *self.cached_broadened_xs.borrow_mut() = Some(Rc::clone(&xs));
1058                *self.cached_work_layout.borrow_mut() = Some(Rc::clone(&layout));
1059                // Invalidate derivative cache — temperature changed, old ∂σ/∂T stale.
1060                *self.cached_dxs_dt.borrow_mut() = None;
1061                self.cached_temperature.set(temperature_k);
1062                (xs, layout)
1063            };
1064
1065            // Beer-Lambert on the working grid: T(E) = exp(-Σᵢ nᵢ · rᵢ · σᵢ(E))
1066            // where rᵢ is the fractional ratio (1.0 for ungrouped isotopes).
1067            let work_len = layout.energies.len();
1068            let mut neg_opt = vec![0.0f64; work_len];
1069            for (i, xs) in broadened_xs.iter().enumerate() {
1070                let density = params[self.density_indices[i]];
1071                let ratio = self.density_ratios[i];
1072                for (j, &sigma) in xs.iter().enumerate() {
1073                    neg_opt[j] -= density * ratio * sigma;
1074                }
1075            }
1076            let transmission: Vec<f64> = neg_opt.iter().map(|&d| d.exp()).collect();
1077
1078            // Issue #442: apply resolution broadening to the total transmission
1079            // AFTER Beer-Lambert, on the working grid; then extract the data
1080            // points last (issue #608).  For Gaussian resolution `resolution_plan`
1081            // is `None` (the planned path is tabulated-only) and broadening runs
1082            // on `layout.energies`; for tabulated resolution the working grid IS
1083            // the data grid so the data-grid plan still matches.
1084            if let Some(ref inst) = self.instrument {
1085                let t_broadened = resolution::apply_resolution_with_plan(
1086                    self.resolution_plan.as_deref(),
1087                    &layout.energies,
1088                    &transmission,
1089                    &inst.resolution,
1090                )
1091                .map_err(|e| {
1092                    FittingError::EvaluationFailed(format!("resolution broadening: {e}"))
1093                })?;
1094                Ok(layout.extract(&t_broadened))
1095            } else {
1096                Ok(layout.extract(&transmission))
1097            }
1098        } else {
1099            // Original path: full forward model (no temperature fitting).
1100            // Apply ratio weights: effective density = params[idx] * ratio.
1101            let isotopes: Vec<(ResonanceData, f64)> = self
1102                .resonance_data
1103                .iter()
1104                .enumerate()
1105                .map(|(i, rd)| {
1106                    (
1107                        rd.clone(),
1108                        params[self.density_indices[i]] * self.density_ratios[i],
1109                    )
1110                })
1111                .collect();
1112
1113            let sample = SampleParams::new(temperature_k, isotopes)
1114                .map_err(|e| FittingError::EvaluationFailed(e.to_string()))?;
1115
1116            transmission::forward_model(&self.energies, &sample, self.instrument.as_deref())
1117                .map_err(|e| FittingError::EvaluationFailed(e.to_string()))
1118        }
1119    }
1120
1121    /// Analytical Jacobian for the transmission model with temperature fitting.
1122    ///
1123    /// When `base_xs` is available (temperature fitting path):
1124    /// - **Density columns**: `∂T/∂nᵢ = -σᵢ(E)·T(E)` using cached broadened XS
1125    ///   from the most recent `evaluate()` call.  Same formula as
1126    ///   `PrecomputedTransmissionModel`, zero extra broadening calls.
1127    /// - **Temperature column**: analytical chain rule via on-demand `∂σ/∂T`.
1128    ///   `∂T/∂T_temp = -T(E) · Σᵢ nᵢ·rᵢ·∂σᵢ/∂T`.  The derivative is
1129    ///   computed once per temperature via
1130    ///   `broadened_cross_sections_with_analytical_derivative_from_base()`
1131    ///   and cached until temperature changes.  Costs one broadening call
1132    ///   per Jacobian (same as the old FD approach, but exact).
1133    ///
1134    /// Returns `None` for the no-base_xs path (full forward model), which
1135    /// falls back to finite-difference in the LM solver.
1136    /// Analytical Jacobian for density and temperature fitting.
1137    ///
1138    /// Without resolution:
1139    ///   ∂T/∂N_g = -(Σ_{i∈g} rᵢ σᵢ) · T
1140    ///   ∂T/∂Temp = -T · Σᵢ nᵢ rᵢ ∂σᵢ/∂T
1141    ///
1142    /// With resolution (R is a linear operator):
1143    ///   ∂T_obs/∂N_g = R\[-(Σ_{i∈g} rᵢ σᵢ) · T\]
1144    ///   ∂T_obs/∂Temp = R\[-T · Σᵢ nᵢ rᵢ ∂σᵢ/∂T\]
1145    ///
1146    /// Returns `None` only when `base_xs` is not available (full forward
1147    /// model path falls back to FD) or when the temperature cache is stale.
1148    fn analytical_jacobian(
1149        &self,
1150        params: &[f64],
1151        free_param_indices: &[usize],
1152        y_current: &[f64],
1153    ) -> Option<FlatMatrix> {
1154        // Cubature fast path — same eligibility as `evaluate()` plus
1155        // the requirement that every free param is a density param.
1156        if let (Some(cubature), Some(inst)) = (&self.sparse_cubature_plan, &self.instrument)
1157            && self.temperature_index.is_none()
1158        {
1159            let params_indices = density_param_indices(&self.density_indices);
1160            if cubature_eligible(
1161                cubature,
1162                &self.energies,
1163                &inst.resolution,
1164                self.resolution_plan.as_deref(),
1165                params_indices.len(),
1166            ) {
1167                let col_map: Option<Vec<usize>> = free_param_indices
1168                    .iter()
1169                    .map(|&fp| params_indices.iter().position(|&i| i == fp))
1170                    .collect();
1171                if let Some(col_map) = col_map {
1172                    let n: Vec<f64> = params_indices.iter().map(|&i| params[i]).collect();
1173                    if density_within_box(cubature, &n) {
1174                        // In-box: take the cubature Jacobian fast
1175                        // path.  Out-of-box falls through to the
1176                        // exact analytical Jacobian below.
1177                        let (_t, jac_flat) = cubature.forward_and_jacobian(&n);
1178                        let k = params_indices.len();
1179                        let n_e = self.energies.len();
1180                        let mut jacobian = FlatMatrix::zeros(n_e, free_param_indices.len());
1181                        for (col, &ell) in col_map.iter().enumerate() {
1182                            for i in 0..n_e {
1183                                *jacobian.get_mut(i, col) = jac_flat[i * k + ell];
1184                            }
1185                        }
1186                        return Some(jacobian);
1187                    }
1188                }
1189            }
1190        }
1191
1192        // Scalar (k = 1) surrogate Jacobian fast path removed —
1193        // see the docstring at the corresponding
1194        // site in `TransmissionFitModel::evaluate()` above.
1195
1196        // Only provide analytical Jacobian when base_xs is available
1197        // (temperature-fitting fast path with cached broadened XS).
1198        let _base_xs_guard = self.base_xs.as_ref()?;
1199        let cached_xs = self.cached_broadened_xs.borrow();
1200        let broadened_xs = cached_xs.as_ref()?;
1201        // Working-grid layout matching the cached σ (issue #608).  Inner
1202        // derivatives are formed on this grid, resolution-broadened there, and
1203        // the data points are extracted LAST.
1204        let cached_layout = self.cached_work_layout.borrow();
1205        let layout = cached_layout.as_ref()?;
1206
1207        // Guard: verify the cache matches the current parameter temperature.
1208        if let Some(ti) = self.temperature_index {
1209            let param_temp = params[ti];
1210            if (param_temp - self.cached_temperature.get()).abs() > 1e-15 {
1211                return None;
1212            }
1213        }
1214
1215        let n_e = y_current.len();
1216        let work_len = layout.energies.len();
1217        let n_free = free_param_indices.len();
1218        let mut jacobian = FlatMatrix::zeros(n_e, n_free);
1219
1220        let temp_col = self
1221            .temperature_index
1222            .and_then(|ti| free_param_indices.iter().position(|&fp| fp == ti));
1223
1224        // The UNRESOLVED transmission T(E) on the WORKING grid, used to form
1225        // inner derivatives before resolution.  Issue #608: with resolution,
1226        // y_current is T_obs = R[T] on the DATA grid — not usable as the inner
1227        // T on the working grid — so recompute T from the cached working-grid
1228        // σ.  Without resolution the working grid is the data grid (identity
1229        // layout) and y_current IS T, so reuse it to stay bit-identical.
1230        let t_unresolved: Option<Vec<f64>> = if self.instrument.is_some() {
1231            let mut neg_opt = vec![0.0f64; work_len];
1232            for (iso, xs) in broadened_xs.iter().enumerate() {
1233                let density = params[self.density_indices[iso]];
1234                let ratio = self.density_ratios[iso];
1235                for (j, &sigma) in xs.iter().enumerate() {
1236                    neg_opt[j] -= density * ratio * sigma;
1237                }
1238            }
1239            Some(neg_opt.iter().map(|&d| d.exp()).collect())
1240        } else {
1241            None
1242        };
1243        // T(E) on the working grid for the inner derivatives.
1244        let t_for_deriv: &[f64] = t_unresolved.as_deref().unwrap_or(y_current);
1245
1246        // ── Density columns: ∂T/∂N_g or ∂T_obs/∂N_g ──
1247        for (col, &fp_idx) in free_param_indices.iter().enumerate() {
1248            if temp_col == Some(col) {
1249                continue;
1250            }
1251            let mut sigma_sum = vec![0.0f64; work_len];
1252            for (iso, &di) in self.density_indices.iter().enumerate() {
1253                if di == fp_idx {
1254                    let ratio = self.density_ratios[iso];
1255                    for (j, &sigma) in broadened_xs[iso].iter().enumerate() {
1256                        sigma_sum[j] += ratio * sigma;
1257                    }
1258                }
1259            }
1260            // Inner derivative on the working grid: -σ_sum · T_unresolved.
1261            let inner: Vec<f64> = (0..work_len)
1262                .map(|i| -sigma_sum[i] * t_for_deriv[i])
1263                .collect();
1264
1265            if let Some(ref inst) = self.instrument {
1266                // ∂T_obs/∂N_g = extract(R[inner]) — resolution on the working
1267                // grid, data points extracted last (issue #608).
1268                let resolved = resolution::apply_resolution_with_plan(
1269                    self.resolution_plan.as_deref(),
1270                    &layout.energies,
1271                    &inner,
1272                    &inst.resolution,
1273                )
1274                .ok()?;
1275                let resolved = layout.extract(&resolved);
1276                for (i, &val) in resolved.iter().enumerate() {
1277                    *jacobian.get_mut(i, col) = val;
1278                }
1279            } else {
1280                // No resolution → identity layout, inner is already data grid.
1281                for (i, &val) in inner.iter().enumerate() {
1282                    *jacobian.get_mut(i, col) = val;
1283                }
1284            }
1285        }
1286
1287        // ── Temperature column: ∂T/∂Temp or ∂T_obs/∂Temp ──
1288        if let Some(col) = temp_col {
1289            // Compute ∂σ/∂T (on the working grid) on-demand if not cached.
1290            {
1291                let needs_compute = self.cached_dxs_dt.borrow().as_ref().is_none();
1292                if needs_compute {
1293                    let base_xs = self.base_xs.as_ref()?;
1294                    let temperature_k = self.cached_temperature.get();
1295                    let working =
1296                        transmission::broadened_cross_sections_with_analytical_derivative_from_base_on_working_grid(
1297                            &self.energies,
1298                            base_xs,
1299                            &self.resonance_data,
1300                            temperature_k,
1301                            self.instrument.as_deref(),
1302                        )
1303                        .ok()?;
1304                    *self.cached_dxs_dt.borrow_mut() = Some(Rc::new(working.dsigma_dt));
1305                }
1306            }
1307            let cached_dxs = self.cached_dxs_dt.borrow();
1308            let dxs_dt = cached_dxs.as_ref()?;
1309
1310            // Inner derivative on the working grid: -T · Σᵢ nᵢ rᵢ ∂σᵢ/∂T.
1311            let inner: Vec<f64> = (0..work_len)
1312                .map(|i| {
1313                    let mut sum_n_dsigma = 0.0f64;
1314                    for (iso, dxs) in dxs_dt.iter().enumerate() {
1315                        let density = params[self.density_indices[iso]];
1316                        let ratio = self.density_ratios[iso];
1317                        sum_n_dsigma += density * ratio * dxs[i];
1318                    }
1319                    -t_for_deriv[i] * sum_n_dsigma
1320                })
1321                .collect();
1322
1323            if let Some(ref inst) = self.instrument {
1324                let resolved = resolution::apply_resolution_with_plan(
1325                    self.resolution_plan.as_deref(),
1326                    &layout.energies,
1327                    &inner,
1328                    &inst.resolution,
1329                )
1330                .ok()?;
1331                let resolved = layout.extract(&resolved);
1332                for (i, &val) in resolved.iter().enumerate() {
1333                    *jacobian.get_mut(i, col) = val;
1334                }
1335            } else {
1336                for (i, &val) in inner.iter().enumerate() {
1337                    *jacobian.get_mut(i, col) = val;
1338                }
1339            }
1340        }
1341
1342        Some(jacobian)
1343    }
1344}
1345
1346/// Wraps a transmission model with SAMMY-style normalization and background.
1347///
1348/// T_out(E) = Anorm × T_inner(E) + BackA + BackB / √E + BackC × √E
1349///          + BackD × exp(−BackF / √E)
1350///
1351/// The normalization and background parameters are additional entries in the
1352/// parameter vector, appended after the density (and optional temperature)
1353/// parameters of the inner model.
1354///
1355/// The exponential tail (BackD, BackF) is optional.  When
1356/// `back_d_index` and `back_f_index` are `None`, the model reduces to
1357/// the 4-parameter form.
1358///
1359/// ## SAMMY Reference
1360/// SAMMY manual Sec III.E.2 — NORMAlization and BACKGround cards.
1361/// SAMMY fits up to 6 background terms; we implement all 6:
1362///   Anorm, constant BackA, 1/√E term BackB, √E term BackC,
1363///   exponential amplitude BackD, exponential decay BackF.
1364pub struct NormalizedTransmissionModel<M: FitModel> {
1365    /// The inner (pure Beer-Lambert) transmission model.
1366    inner: M,
1367    /// Precomputed √E for each energy bin.  Computed once in `new()`.
1368    sqrt_energies: Vec<f64>,
1369    /// Precomputed 1/√E for each energy bin.  Computed once in `new()`.
1370    inv_sqrt_energies: Vec<f64>,
1371    /// Index of the Anorm parameter in the full parameter vector.
1372    anorm_index: usize,
1373    /// Index of the BackA (constant background) parameter.
1374    back_a_index: usize,
1375    /// Index of the BackB (1/√E background) parameter.
1376    back_b_index: usize,
1377    /// Index of the BackC (√E background) parameter.
1378    back_c_index: usize,
1379    /// Index of BackD (exponential amplitude) in the parameter vector.
1380    /// `None` disables the exponential tail term.
1381    back_d_index: Option<usize>,
1382    /// Index of BackF (exponential decay constant) in the parameter vector.
1383    /// `None` disables the exponential tail term.
1384    back_f_index: Option<usize>,
1385}
1386
1387impl<M: FitModel> NormalizedTransmissionModel<M> {
1388    /// Create a new normalized transmission model (4-parameter, no exponential tail).
1389    ///
1390    /// # Arguments
1391    /// * `inner` — The inner transmission model (Beer-Lambert).
1392    /// * `energies` — Energy grid in eV (must be positive).
1393    /// * `anorm_index` — Index of Anorm in the parameter vector.
1394    /// * `back_a_index` — Index of BackA in the parameter vector.
1395    /// * `back_b_index` — Index of BackB in the parameter vector.
1396    /// * `back_c_index` — Index of BackC in the parameter vector.
1397    pub fn new(
1398        inner: M,
1399        energies: &[f64],
1400        anorm_index: usize,
1401        back_a_index: usize,
1402        back_b_index: usize,
1403        back_c_index: usize,
1404    ) -> Self {
1405        let sqrt_energies: Vec<f64> = energies.iter().map(|&e| e.sqrt()).collect();
1406        let inv_sqrt_energies: Vec<f64> = sqrt_energies
1407            .iter()
1408            .map(|&se| if se > 0.0 { 1.0 / se } else { 0.0 })
1409            .collect();
1410        Self {
1411            inner,
1412            sqrt_energies,
1413            inv_sqrt_energies,
1414            anorm_index,
1415            back_a_index,
1416            back_b_index,
1417            back_c_index,
1418            back_d_index: None,
1419            back_f_index: None,
1420        }
1421    }
1422
1423    /// Create a normalized transmission model with the SAMMY exponential tail.
1424    ///
1425    /// Adds BackD × exp(−BackF / √E) to the 4-parameter background model.
1426    ///
1427    /// # Arguments
1428    /// * `back_d_index` — Index of BackD (exponential amplitude) in the parameter vector.
1429    /// * `back_f_index` — Index of BackF (exponential decay constant) in the parameter vector.
1430    #[allow(clippy::too_many_arguments)]
1431    pub fn new_with_exponential(
1432        inner: M,
1433        energies: &[f64],
1434        anorm_index: usize,
1435        back_a_index: usize,
1436        back_b_index: usize,
1437        back_c_index: usize,
1438        back_d_index: usize,
1439        back_f_index: usize,
1440    ) -> Self {
1441        let sqrt_energies: Vec<f64> = energies.iter().map(|&e| e.sqrt()).collect();
1442        let inv_sqrt_energies: Vec<f64> = sqrt_energies
1443            .iter()
1444            .map(|&se| if se > 0.0 { 1.0 / se } else { 0.0 })
1445            .collect();
1446        Self {
1447            inner,
1448            sqrt_energies,
1449            inv_sqrt_energies,
1450            anorm_index,
1451            back_a_index,
1452            back_b_index,
1453            back_c_index,
1454            back_d_index: Some(back_d_index),
1455            back_f_index: Some(back_f_index),
1456        }
1457    }
1458}
1459
1460impl<M: FitModel> FitModel for NormalizedTransmissionModel<M> {
1461    fn evaluate(&self, params: &[f64]) -> Result<Vec<f64>, FittingError> {
1462        let t_inner = self.inner.evaluate(params)?;
1463        let anorm = params[self.anorm_index];
1464        let back_a = params[self.back_a_index];
1465        let back_b = params[self.back_b_index];
1466        let back_c = params[self.back_c_index];
1467
1468        // Optional exponential tail: BackD × exp(−BackF / √E)
1469        let (back_d, back_f) = match (self.back_d_index, self.back_f_index) {
1470            (Some(di), Some(fi)) => (params[di], params[fi]),
1471            _ => (0.0, 0.0),
1472        };
1473        let has_exp = self.back_d_index.is_some();
1474
1475        let result: Vec<f64> = t_inner
1476            .iter()
1477            .enumerate()
1478            .map(|(i, &t)| {
1479                let mut val = anorm * t
1480                    + back_a
1481                    + back_b * self.inv_sqrt_energies[i]
1482                    + back_c * self.sqrt_energies[i];
1483                if has_exp {
1484                    val += back_d * (-back_f * self.inv_sqrt_energies[i]).exp();
1485                }
1486                val
1487            })
1488            .collect();
1489        Ok(result)
1490    }
1491
1492    /// Analytical Jacobian for the normalized transmission model.
1493    ///
1494    /// For each free parameter:
1495    /// - If it belongs to the inner model (density or temperature):
1496    ///   ∂T_out/∂p = Anorm × ∂T_inner/∂p  (inner Jacobian scaled by Anorm)
1497    /// - ∂T_out/∂Anorm  = T_inner(E)
1498    /// - ∂T_out/∂BackA  = 1
1499    /// - ∂T_out/∂BackB  = 1/√E
1500    /// - ∂T_out/∂BackC  = √E
1501    /// - ∂T_out/∂BackD  = exp(−BackF / √E)
1502    /// - ∂T_out/∂BackF  = −BackD × exp(−BackF / √E) / √E
1503    fn analytical_jacobian(
1504        &self,
1505        params: &[f64],
1506        free_param_indices: &[usize],
1507        y_current: &[f64],
1508    ) -> Option<FlatMatrix> {
1509        let n_e = y_current.len();
1510        let n_free = free_param_indices.len();
1511
1512        // Compute T_inner for Anorm column and for scaling inner Jacobian.
1513        // T_inner = (T_out - BackA - BackB/√E - BackC×√E) / Anorm
1514        // But to avoid numerical issues, recompute from the inner model.
1515        let t_inner = self.inner.evaluate(params).ok()?;
1516
1517        let anorm = params[self.anorm_index];
1518
1519        // Identify which free params are background params vs inner params.
1520        let mut bg_indices_set = vec![
1521            self.anorm_index,
1522            self.back_a_index,
1523            self.back_b_index,
1524            self.back_c_index,
1525        ];
1526        if let Some(di) = self.back_d_index {
1527            bg_indices_set.push(di);
1528        }
1529        if let Some(fi) = self.back_f_index {
1530            bg_indices_set.push(fi);
1531        }
1532
1533        // Collect inner model's free param indices (those not in bg_indices).
1534        let inner_free_indices: Vec<usize> = free_param_indices
1535            .iter()
1536            .copied()
1537            .filter(|idx| !bg_indices_set.contains(idx))
1538            .collect();
1539
1540        // Get inner Jacobian if there are inner free params.
1541        // y_current for the inner model is t_inner, not the outer y_current.
1542        let inner_jac = if !inner_free_indices.is_empty() {
1543            self.inner
1544                .analytical_jacobian(params, &inner_free_indices, &t_inner)
1545        } else {
1546            None
1547        };
1548
1549        // Precompute exp(−BackF / √E) for the exponential tail columns.
1550        let exp_terms: Vec<f64> =
1551            if let (Some(di), Some(fi)) = (self.back_d_index, self.back_f_index) {
1552                let _back_d = params[di];
1553                let back_f = params[fi];
1554                self.inv_sqrt_energies
1555                    .iter()
1556                    .map(|&inv_se| (-back_f * inv_se).exp())
1557                    .collect()
1558            } else {
1559                vec![]
1560            };
1561
1562        let mut jacobian = FlatMatrix::zeros(n_e, n_free);
1563
1564        // Map inner free param index → column in inner Jacobian.
1565        let mut inner_col_map = std::collections::HashMap::new();
1566        for (col, &idx) in inner_free_indices.iter().enumerate() {
1567            inner_col_map.insert(idx, col);
1568        }
1569
1570        for (col, &fp_idx) in free_param_indices.iter().enumerate() {
1571            if fp_idx == self.anorm_index {
1572                // ∂T_out/∂Anorm = T_inner(E)
1573                for (i, &ti) in t_inner.iter().enumerate() {
1574                    *jacobian.get_mut(i, col) = ti;
1575                }
1576            } else if fp_idx == self.back_a_index {
1577                // ∂T_out/∂BackA = 1
1578                for i in 0..n_e {
1579                    *jacobian.get_mut(i, col) = 1.0;
1580                }
1581            } else if fp_idx == self.back_b_index {
1582                // ∂T_out/∂BackB = 1/√E
1583                for (i, &inv_se) in self.inv_sqrt_energies.iter().enumerate() {
1584                    *jacobian.get_mut(i, col) = inv_se;
1585                }
1586            } else if fp_idx == self.back_c_index {
1587                // ∂T_out/∂BackC = √E
1588                for (i, &se) in self.sqrt_energies.iter().enumerate() {
1589                    *jacobian.get_mut(i, col) = se;
1590                }
1591            } else if self.back_d_index == Some(fp_idx) {
1592                // ∂T_out/∂BackD = exp(−BackF / √E)
1593                for (i, &et) in exp_terms.iter().enumerate() {
1594                    *jacobian.get_mut(i, col) = et;
1595                }
1596            } else if self.back_f_index == Some(fp_idx) {
1597                // ∂T_out/∂BackF = −BackD × exp(−BackF / √E) / √E
1598                let back_d = params[self.back_d_index.unwrap()];
1599                for (i, (&et, &inv_se)) in exp_terms
1600                    .iter()
1601                    .zip(self.inv_sqrt_energies.iter())
1602                    .enumerate()
1603                {
1604                    *jacobian.get_mut(i, col) = -back_d * et * inv_se;
1605                }
1606            } else if let Some(&inner_col) = inner_col_map.get(&fp_idx) {
1607                // Inner model parameter: ∂T_out/∂p = Anorm × ∂T_inner/∂p
1608                if let Some(ref jac) = inner_jac {
1609                    for i in 0..n_e {
1610                        *jacobian.get_mut(i, col) = anorm * jac.get(i, inner_col);
1611                    }
1612                } else {
1613                    // Inner model did not provide analytical Jacobian —
1614                    // fall back to finite-difference for the whole thing.
1615                    return None;
1616                }
1617            } else {
1618                // Unknown parameter — should not happen, but fall back to FD.
1619                return None;
1620            }
1621        }
1622
1623        Some(jacobian)
1624    }
1625}
1626
1627// ── Energy-scale transmission model (SAMMY TZERO equivalent) ─────────────
1628
1629/// Transmission model with energy-scale calibration parameters (t₀, L_scale).
1630///
1631/// Carries per-isotope resonance data (NOT a precomputed σ grid) and rebuilds
1632/// the TRUE cross-section at the corrected energies on each evaluation
1633/// (issue #608), matching `forward_model`:
1634///   1. Convert nominal energy → TOF: `t = TOF_FACTOR * L / √E_nom`
1635///   2. Apply calibration: `t_corr = t - t₀`,
1636///      `E_corr = (TOF_FACTOR * L * L_scale / t_corr)²`
1637///   3. Evaluate σ(E_corr) directly via `reich_moore` + Doppler on a working
1638///      grid built from `E_corr` (auxiliary extended grid under Gaussian
1639///      resolution; `E_corr` itself for tabulated / no resolution) — NOT
1640///      interpolation of a fixed σ grid, which clamps at the auxiliary
1641///      boundary and drops resonance fine-structure.
1642///   4. Beer-Lambert + resolution on the working grid, then extract the data
1643///      points last.
1644///
1645/// This is equivalent to SAMMY's TZERO parameters.
1646///
1647/// The Jacobian for t₀ and L_scale defaults to **partial-GAL** since
1648/// issue #489: central FD on `t0` only (2 evals) plus an inline rank-1
1649/// derivation of the `L_scale` column. The previous central-FD-on-both
1650/// (4-eval) behaviour is reachable via `with_jacobian_method`,
1651/// `NEREIDS_TZERO_JACOBIAN=fd2`, or `tzero_jacobian="fd2"` Python kwarg.
1652/// See [`EnergyScaleJacobianMethod`] for full method documentation.
1653pub struct EnergyScaleTransmissionModel {
1654    /// Resonance parameters per isotope.  Issue #608: the energy-scale model
1655    /// evaluates the TRUE cross-section at the corrected energies (matching
1656    /// `forward_model`) instead of interpolating a precomputed σ grid, so it
1657    /// carries resonance data and rebuilds σ on the corrected working grid each
1658    /// `evaluate`.  This is the only way to reproduce SAMMY's σ(E_corr) under
1659    /// the energy-scale shift with full boundary + resonance-fine-structure
1660    /// fidelity; interpolating a fixed precomputed σ cannot (it clamps at the
1661    /// auxiliary boundary and misses fine-structure).
1662    resonance_data: Arc<Vec<ResonanceData>>,
1663    /// Density parameter index per isotope (same convention as
1664    /// `PrecomputedTransmissionModel`).
1665    density_indices: Arc<Vec<usize>>,
1666    /// Fractional ratio per isotope (1.0 when ungrouped).  Per-isotope
1667    /// thickness is `params[density_indices[i]] * density_ratios[i]`.
1668    density_ratios: Arc<Vec<f64>>,
1669    /// Sample temperature (K) for Doppler broadening at the corrected energies.
1670    temperature_k: f64,
1671    /// Nominal energy grid (eV, ascending).
1672    nominal_energies: Vec<f64>,
1673    /// Flight path length in meters (used for TOF↔energy conversion).
1674    flight_path_m: f64,
1675    /// TOF factor: sqrt(m_n / (2 * eV)) in μs·√eV/m.
1676    tof_factor: f64,
1677    /// Index of t₀ (μs) in the parameter vector.
1678    t0_index: usize,
1679    /// Index of L_scale (dimensionless) in the parameter vector.
1680    l_scale_index: usize,
1681    /// Instrument resolution parameters (applied after Beer-Lambert).
1682    instrument: Option<Arc<transmission::InstrumentParams>>,
1683    /// Plan cache keyed on `(t0_bits, l_scale_bits)`.  Within one KL
1684    /// outer iteration (deviance + gradient + Fisher all at the same
1685    /// `params`) `evaluate_at` is called 3× at identical `(t0, L)`;
1686    /// the density-column path of `analytical_jacobian` wants a plan
1687    /// at that same `(t0, L)` too — that's 4 cache hits per outer
1688    /// iter on KL+periso+TZERO.  Finite-difference probes land at a
1689    /// different `(t0, L)` bit-pattern from the accepted probe and
1690    /// are routed through `evaluate_at_with_cache(..., false)` so
1691    /// they stay on the non-plan broadening path — no plan is built
1692    /// or inserted for FD probes, so they neither miss nor pollute
1693    /// the cache.
1694    ///
1695    /// **Capacity 2** (FIFO on miss): this survives LM backtracking,
1696    /// where a proposed-but-rejected trial step evaluates at a new
1697    /// `(t0, L)` key and would otherwise evict the accepted-step
1698    /// plan.  With capacity 2, the accepted plan stays resident
1699    /// alongside the trial plan; if the trial is rejected, the next
1700    /// iteration's evaluate at the accepted `(t0, L)` still hits.
1701    /// Only when a genuine new accepted step lands do we start
1702    /// aging the oldest entry out (#483 A1).
1703    ///
1704    /// `RefCell` is safe: `TransmissionFitModel`-family models are
1705    /// rebuilt per-pixel and never shared across rayon workers.
1706    cached_plans: RefCell<CachedPlanRing>,
1707    /// Capacity-1 cache of the working-grid σ keyed on `(t0_bits, L_scale_bits)`
1708    /// (issue #608 perf): a base-point `evaluate` + the Jacobian's density
1709    /// columns at the same probe reuse one reich_moore+Doppler build instead of
1710    /// rebuilding it twice.  `RefCell` is safe — the model is rebuilt per pixel
1711    /// and never shared across threads.
1712    cached_work_xs: RefCell<CachedWorkXs>,
1713    /// Method for the t0 / L_scale Jacobian columns. Initialised from
1714    /// [`EnergyScaleJacobianMethod::from_env`] in [`Self::new`], which
1715    /// defaults to `PartialGal` since issue #489 (and respects the
1716    /// `NEREIDS_TZERO_JACOBIAN` env var as a global override). Can be
1717    /// overridden per-instance via [`Self::with_jacobian_method`].
1718    jacobian_method: EnergyScaleJacobianMethod,
1719}
1720
1721/// Capacity-1 working-grid σ cache entry, keyed on `(t0_bits, l_scale_bits)`.
1722/// Named alias to keep the field type within clippy's `type_complexity` budget
1723/// (issue #608).
1724type CachedWorkXs = Option<((u64, u64), Rc<transmission::WorkingGridXs>)>;
1725
1726/// One `(t0_bits, l_scale_bits)` → `ResolutionPlan` entry.  Named
1727/// struct to keep the cache field type within clippy's
1728/// `type_complexity` budget.
1729#[derive(Debug, Clone)]
1730struct CachedPlanEntry {
1731    key: (u64, u64),
1732    plan: Arc<ResolutionPlan>,
1733}
1734
1735/// Capacity-2 FIFO ring of plan entries.  Two entries suffice to
1736/// survive a single-trial LM backtrack (accepted + trial); deeper
1737/// backtracking chains still lose the accepted plan eventually, but
1738/// those are rare in production and cheaper to miss than the default
1739/// non-plan path.  Issue #483 A1.
1740#[derive(Debug, Default)]
1741struct CachedPlanRing {
1742    /// Slot 0 is the most-recently-inserted entry; slot 1 is the
1743    /// previous entry.  Lookup checks both; insert shifts 0 → 1 and
1744    /// places the new entry at 0.
1745    slots: [Option<CachedPlanEntry>; 2],
1746}
1747
1748impl CachedPlanRing {
1749    fn lookup(&self, key: (u64, u64)) -> Option<Arc<ResolutionPlan>> {
1750        for slot in &self.slots {
1751            if let Some(entry) = slot
1752                && entry.key == key
1753            {
1754                return Some(Arc::clone(&entry.plan));
1755            }
1756        }
1757        None
1758    }
1759
1760    fn insert(&mut self, entry: CachedPlanEntry) {
1761        // Shift oldest out, newest to slot 0.
1762        self.slots[1] = self.slots[0].take();
1763        self.slots[0] = Some(entry);
1764    }
1765}
1766
1767/// Method for computing the t0 / L_scale columns of the
1768/// `EnergyScaleTransmissionModel` Jacobian.
1769///
1770/// - `PartialGal` (default since issue #489): central FD on `t0` only
1771///   (2 evaluations); derive `L_scale` column inline via the rank-1
1772///   identity `J[:, L_scale] = ((tof - t0) / L_scale) * J[:, t0]` per
1773///   energy bin. Halves the FD probe count on workloads where both
1774///   calibration parameters are free.
1775///
1776///   **Correctness regime**: exact in the no-resolution limit and the
1777///   narrow-kernel limit. With a non-trivial resolution operator `R`,
1778///   the rank-1 simplification additionally assumes per-bin uniformity
1779///   of `(tof - t0) / L_scale` over the kernel support — necessary
1780///   because `R` mixes source bins whose ratios differ. `broaden_presorted`
1781///   uses `self.flight_path_m` (not the model's `L_nominal * L_scale`) so
1782///   tabulated kernels satisfy the structural factorisation through
1783///   `e_corr`, but the per-bin homogeneity assumption is empirical.
1784///   On real VENUS Hf 120-min KL+per-iso+TZERO 4×4 the approximation is
1785///   tight enough that 15/16 pixels converge within 0.1·σ_Fisher of FD2;
1786///   median wall-time speedup 1.28× over FD2.
1787/// - `FiniteDifference`: central FD on the full inner forward chain,
1788///   4 forward evaluations per Jacobian (h_t0=1e-4, h_ls=1e-7).
1789///   The pre-#489 production default; reachable via
1790///   `NEREIDS_TZERO_JACOBIAN=fd2` env var or `tzero_jacobian="fd2"`
1791///   Python kwarg.
1792#[derive(Debug, Clone, Copy, PartialEq, Eq)]
1793pub enum EnergyScaleJacobianMethod {
1794    FiniteDifference,
1795    PartialGal,
1796}
1797
1798impl EnergyScaleJacobianMethod {
1799    /// Resolve the default Jacobian method from the
1800    /// `NEREIDS_TZERO_JACOBIAN` env var.
1801    ///
1802    /// The env var is read **once per process** via a `OnceLock`. Per
1803    /// `EnergyScaleTransmissionModel::new` is hot under
1804    /// `spatial_map_typed` (one model per pixel; 262 144 calls per
1805    /// 512×512 map), so `std::env::var` would otherwise be a syscall
1806    /// hot spot. Tests that need to swap the default must use
1807    /// `EnergyScaleTransmissionModel::with_jacobian_method` (which
1808    /// bypasses the cache); changing the env var mid-process has no
1809    /// effect.
1810    ///
1811    /// An unrecognized or removed value (e.g. the `"chain"` method dropped in
1812    /// #608) emits a one-time `eprintln` warning and falls back to the
1813    /// `PartialGal` default rather than being silently masked.  It does not
1814    /// panic — `new` is a hot, infallible, per-pixel constructor across the
1815    /// PyO3 boundary; the Python `tzero_jacobian=` kwarg is the strict
1816    /// (hard-erroring) override path.
1817    fn from_env() -> Self {
1818        use std::sync::OnceLock;
1819        static CACHED: OnceLock<EnergyScaleJacobianMethod> = OnceLock::new();
1820        *CACHED.get_or_init(Self::resolve_env_uncached)
1821    }
1822
1823    fn resolve_env_uncached() -> Self {
1824        let Ok(v) = std::env::var("NEREIDS_TZERO_JACOBIAN") else {
1825            // Unset → the documented #489 default, silently.
1826            return Self::PartialGal;
1827        };
1828        if v.eq_ignore_ascii_case("fd2")
1829            || v.eq_ignore_ascii_case("finite-difference")
1830            || v.eq_ignore_ascii_case("finite_difference")
1831        {
1832            Self::FiniteDifference
1833        } else if v.eq_ignore_ascii_case("partial-gal") || v.eq_ignore_ascii_case("partial_gal") {
1834            Self::PartialGal
1835        } else {
1836            // Set to an unrecognized / removed method name.  The legacy
1837            // `"chain"` / `"frozen-r"` / `"frozen_r"` FrozenResolutionChainRule
1838            // method was removed in #608 (it interpolated a precomputed σ on the
1839            // data grid, incompatible with the true-σ aux-grid `evaluate`;
1840            // FD / PartialGal of the corrected evaluate is the exact
1841            // replacement).  The Python `tzero_jacobian=` kwarg HARD-ERRORS on
1842            // these names (bindings/python `parse_tzero_jacobian`); `from_env` is
1843            // an infallible, process-cached, per-pixel constructor path that must
1844            // not panic across the PyO3 boundary (cf. the #608 `working_xs`
1845            // Err-not-panic guard), so it cannot itself return an error.  Warn
1846            // loudly (once, via the `OnceLock` in `from_env`) so the override is
1847            // NOT silently masked, then fall back to the PartialGal default —
1848            // matching the kwarg in *surfacing* the bad value while staying
1849            // non-fatal on this hot, infallible path.
1850            eprintln!(
1851                "warning: NEREIDS_TZERO_JACOBIAN=\"{v}\" is not a recognized \
1852                 Jacobian method (\"chain\" / \"frozen-r\" were removed in #608); \
1853                 using the default \"partial-gal\". Valid values: \"fd2\", \
1854                 \"partial-gal\"."
1855            );
1856            Self::PartialGal
1857        }
1858    }
1859}
1860
1861impl EnergyScaleTransmissionModel {
1862    /// Create a new energy-scale transmission model.
1863    ///
1864    /// # Arguments
1865    /// * `resonance_data` — Resonance parameters per isotope; σ is evaluated at
1866    ///   the corrected energies via `reich_moore` + Doppler (issue #608).
1867    /// * `density_indices` — Maps isotope index → density parameter index.
1868    /// * `density_ratios` — Fractional ratio per isotope (1.0 when ungrouped).
1869    /// * `temperature_k` — Sample temperature (K) for Doppler broadening.
1870    /// * `nominal_energies` — Energy grid in eV (ascending).
1871    /// * `flight_path_m` — Nominal flight path in meters.
1872    /// * `t0_index` — Index of t₀ parameter.
1873    /// * `l_scale_index` — Index of L_scale parameter.
1874    /// * `instrument` — Optional resolution function.
1875    #[allow(clippy::too_many_arguments)]
1876    pub fn new(
1877        resonance_data: Arc<Vec<ResonanceData>>,
1878        density_indices: Arc<Vec<usize>>,
1879        density_ratios: Arc<Vec<f64>>,
1880        temperature_k: f64,
1881        nominal_energies: Vec<f64>,
1882        flight_path_m: f64,
1883        t0_index: usize,
1884        l_scale_index: usize,
1885        instrument: Option<Arc<transmission::InstrumentParams>>,
1886    ) -> Self {
1887        // TOF_FACTOR = sqrt(m_n / (2 · eV)) · 1e6 [μs·√eV/m].
1888        // Use the CODATA 2018 values from nereids-core::constants so that
1889        // this model, calibration.rs, and core::tof_to_energy all agree to
1890        // machine precision (previously the inline approximations differed
1891        // by ~5e-5 relative, enough to visibly shift sharp resonances).
1892        let tof_factor = (0.5 * NEUTRON_MASS_KG / EV_TO_JOULES).sqrt() * 1.0e6;
1893        Self {
1894            resonance_data,
1895            density_indices,
1896            density_ratios,
1897            temperature_k,
1898            nominal_energies,
1899            flight_path_m,
1900            tof_factor,
1901            t0_index,
1902            l_scale_index,
1903            instrument,
1904            cached_plans: RefCell::new(CachedPlanRing::default()),
1905            cached_work_xs: RefCell::new(None),
1906            jacobian_method: EnergyScaleJacobianMethod::from_env(),
1907        }
1908    }
1909
1910    /// Override the t0 / L_scale Jacobian method for this model instance.
1911    /// Bypasses the `NEREIDS_TZERO_JACOBIAN` env var.
1912    #[must_use]
1913    pub fn with_jacobian_method(mut self, method: EnergyScaleJacobianMethod) -> Self {
1914        self.jacobian_method = method;
1915        self
1916    }
1917
1918    /// Build or reuse the broadening plan for the current `(t0, L_scale)`
1919    /// probe.  Capacity-2 FIFO ring keyed on raw `f64` bits, matching
1920    /// the invariant that `corrected_energies(t0, L)` is a pure
1921    /// function of `(t0_bits, L_bits)` and `self.nominal_energies`
1922    /// (fixed for the model's lifetime).
1923    ///
1924    /// Capacity 2 survives one LM backtrack rejection: the previous
1925    /// (accepted) entry stays in slot 1 while the trial-step entry
1926    /// occupies slot 0, so a rejection followed by an evaluate at the
1927    /// restored accepted `(t0, L)` still hits (#483 A1).
1928    ///
1929    /// Returns `None` for Gaussian resolution (no plan representation)
1930    /// or when the `build_resolution_plan` call fails (unsorted grid) —
1931    /// both cases transparently fall back to the non-plan
1932    /// `apply_resolution` path via `apply_resolution_with_plan(None, …)`.
1933    ///
1934    /// `working_energies` is the broadening grid the plan is built on — the
1935    /// model's WORKING grid (`work.layout.energies`), which every caller passes
1936    /// post-#608.  For tabulated resolution (the only case that builds a plan)
1937    /// the working grid IS the corrected data grid; for Gaussian it is the
1938    /// auxiliary extended grid, but that path returns `None` above before the
1939    /// grid is used.
1940    fn cached_resolution_plan(
1941        &self,
1942        t0_us: f64,
1943        l_scale: f64,
1944        working_energies: &[f64],
1945    ) -> Option<Arc<ResolutionPlan>> {
1946        let inst = self.instrument.as_ref()?;
1947        // Match on a reference to `inst.resolution` defensively so the
1948        // check never attempts to move a non-`Copy` `ResolutionFunction`
1949        // out of a shared `Arc<InstrumentParams>`.
1950        if !matches!(
1951            &inst.resolution,
1952            nereids_physics::resolution::ResolutionFunction::Tabulated(_)
1953        ) {
1954            // Gaussian resolution has no plan; nothing to cache.
1955            return None;
1956        }
1957        let key = (t0_us.to_bits(), l_scale.to_bits());
1958        if let Some(plan) = self.cached_plans.borrow().lookup(key) {
1959            return Some(plan);
1960        }
1961        // Miss: build, insert, return.
1962        let plan = resolution::build_resolution_plan(working_energies, &inst.resolution)
1963            .ok()
1964            .flatten()?;
1965        let arc = Arc::new(plan);
1966        self.cached_plans.borrow_mut().insert(CachedPlanEntry {
1967            key,
1968            plan: Arc::clone(&arc),
1969        });
1970        Some(arc)
1971    }
1972
1973    /// Compute the corrected energy grid for given (t₀, L_scale).
1974    ///
1975    /// **Physical bound on `t0_us`.**  The corrected TOF is `tof - t0_us`,
1976    /// where `tof = tof_factor · L / √E_nom`.  For the corrected grid to
1977    /// remain physical, `tof_corr > 0` must hold for every bin — i.e.
1978    /// `t0_us < min_i(tof_i) = tof_factor · L / √(max E_nom)`.  The
1979    /// `EnergyScaleTransmissionModel` pipeline registers `t0_us` with
1980    /// bounds of ±10 μs, which safely satisfies this invariant for VENUS
1981    /// (L = 25 m, E ≤ 200 eV gives `min_tof ≈ 17.7 μs`).
1982    ///
1983    /// As a defensive measure — if a caller ever invokes this function
1984    /// with a `t0_us` that would push any bin's `tof_corr` below zero —
1985    /// we clamp `t0_us` to just under `min_tof` so the corrected grid
1986    /// stays monotone and physical.  This is a safety net; the expected
1987    /// path is that the optimizer's parameter bounds keep `t0_us` well
1988    /// below the clamp threshold.
1989    fn corrected_energies(&self, t0_us: f64, l_scale: f64) -> Vec<f64> {
1990        if self.nominal_energies.is_empty() {
1991            return Vec::new();
1992        }
1993        let l_eff = self.flight_path_m * l_scale;
1994        // min(tof) over the grid = tof_factor * L / sqrt(max E_nom).
1995        let min_tof = self
1996            .nominal_energies
1997            .iter()
1998            .copied()
1999            .fold(f64::INFINITY, |acc, e| {
2000                acc.min(self.tof_factor * self.flight_path_m / e.sqrt())
2001            });
2002        let t0_limit = min_tof * (1.0 - 1.0e-12);
2003        let t0_clamped = t0_us.min(t0_limit);
2004        self.nominal_energies
2005            .iter()
2006            .map(|&e_nom| {
2007                let tof = self.tof_factor * self.flight_path_m / e_nom.sqrt();
2008                let tof_corr = tof - t0_clamped;
2009                (self.tof_factor * l_eff / tof_corr).powi(2)
2010            })
2011            .collect()
2012    }
2013
2014    /// Doppler-broadened TRUE σ per isotope on the working grid built from the
2015    /// corrected energies, plus the data-grid layout (issue #608).
2016    ///
2017    /// Mirrors `forward_model`: builds the auxiliary extended grid on `e_corr`
2018    /// WITH the model's resonance data (boundary extension + resonance
2019    /// fine-structure), evaluates σ via `reich_moore` at those energies, and
2020    /// Doppler-broadens there.  The corrected grid is re-derived per
2021    /// `(t0, L_scale)` probe, so the working grid + σ are rebuilt each call —
2022    /// the only way to reproduce SAMMY's σ(E_corr) under the energy-scale shift
2023    /// (boundary + fine-structure fidelity).  For tabulated / no resolution the
2024    /// working grid is `e_corr` itself with an identity layout.
2025    fn working_xs(&self, e_corr: &[f64]) -> Result<transmission::WorkingGridXs, FittingError> {
2026        // Issue #608: a degenerate calibration can drive corrected energies to
2027        // 0 (l_scale → 0) or non-finite (l_scale → ∞).  `reich_moore` asserts
2028        // positive finite energy (an always-on `assert!`), so without this guard
2029        // such inputs PANIC inside `broadened_cross_sections_on_working_grid` —
2030        // a process abort across the PyO3 boundary.  Return a graceful Err so the
2031        // LM/KL/Python callers see a failed evaluate instead of a panic.
2032        //
2033        // BEHAVIOR CHANGE vs pre-#608: the old model interpolated a precomputed σ
2034        // and CLAMPED a degenerate corrected energy to the grid edge, continuing
2035        // the fit with a (finite but unphysical) value; the true-σ model instead
2036        // FAILS the evaluate rather than fabricating σ at a non-positive energy.
2037        // Reachable only by a degenerate calibration, which production keeps out
2038        // of reach: `validate_energy_scale_params` rejects `l_scale_init <= 0` at
2039        // setup and `corrected_energies` clamps `t0` below the min TOF, so a real
2040        // fit never drives `e_corr` to 0 / ∞; this guard is the runtime backstop.
2041        if let Some(&bad) = e_corr.iter().find(|&&e| !e.is_finite() || e <= 0.0) {
2042            return Err(FittingError::EvaluationFailed(format!(
2043                "energy-scale corrected energy is non-positive or non-finite ({bad}); \
2044                 t0 / L_scale give a degenerate calibration"
2045            )));
2046        }
2047        transmission::broadened_cross_sections_on_working_grid(
2048            e_corr,
2049            &self.resonance_data,
2050            self.temperature_k,
2051            self.instrument.as_deref(),
2052            None,
2053        )
2054        .map_err(|e| FittingError::EvaluationFailed(e.to_string()))
2055    }
2056
2057    /// Working-grid σ for the current probe, cached (capacity 1, keyed on
2058    /// `(t0, L_scale)` bits) so a base-point `evaluate` and the Jacobian's
2059    /// density columns at the SAME probe share one reich_moore+Doppler build
2060    /// instead of rebuilding it twice (issue #608 perf).  FD probes at
2061    /// perturbed `(t0, L_scale)` miss and rebuild, as required.
2062    fn working_xs_for(
2063        &self,
2064        params: &[f64],
2065        e_corr: &[f64],
2066    ) -> Result<Rc<transmission::WorkingGridXs>, FittingError> {
2067        let key = (
2068            params[self.t0_index].to_bits(),
2069            params[self.l_scale_index].to_bits(),
2070        );
2071        let hit = self
2072            .cached_work_xs
2073            .borrow()
2074            .as_ref()
2075            .and_then(|(k, xs)| (*k == key).then(|| Rc::clone(xs)));
2076        if let Some(xs) = hit {
2077            return Ok(xs);
2078        }
2079        let xs = Rc::new(self.working_xs(e_corr)?);
2080        *self.cached_work_xs.borrow_mut() = Some((key, Rc::clone(&xs)));
2081        Ok(xs)
2082    }
2083
2084    /// Evaluate transmission at given parameters (densities + t0 + l_scale).
2085    ///
2086    /// When `use_plan_cache` is `true`, the struct-level `(t0, L_scale)`-
2087    /// keyed plan cache is consulted and populated — appropriate for
2088    /// evaluate calls that will be followed by more work at the SAME
2089    /// probe (e.g. `FitModel::evaluate` + `analytical_jacobian` density
2090    /// cols within one KL outer iter).  When `false`, broadening goes
2091    /// through the non-plan path unchanged — appropriate for the
2092    /// one-shot LM FD probes at `(t0 ± h, L)` / `(t0, L ± h)` where
2093    /// a plan build has no reuse to amortize.  Issue #483 A1.
2094    fn evaluate_at_with_cache(
2095        &self,
2096        params: &[f64],
2097        e_corr: &[f64],
2098        use_plan_cache: bool,
2099    ) -> Result<Vec<f64>, FittingError> {
2100        // Issue #608: evaluate the TRUE σ at the corrected energies on the
2101        // working grid (auxiliary extended grid for Gaussian resolution; the
2102        // data grid for tabulated / no resolution) — reich_moore + Doppler on
2103        // the working grid, Beer-Lambert, resolution, extract the data points
2104        // last — exactly as `forward_model` does.  This replaces interpolating
2105        // a precomputed σ, which clamped at the auxiliary boundary and dropped
2106        // resonance fine-structure (a forward_model-fidelity gap; #608).
2107        let work = self.working_xs_for(params, e_corr)?;
2108        let work_e = &work.layout.energies;
2109
2110        // Beer-Lambert on the working grid: T = exp(-Σᵢ nᵢ·rᵢ·σᵢ(E)), where rᵢ
2111        // is the fractional ratio (1.0 for ungrouped isotopes).  No density > 0
2112        // guard — exp(−n·σ) is well-defined for negative n, matching
2113        // PrecomputedTransmissionModel (issue #109.1).
2114        let mut neg_opt = vec![0.0f64; work_e.len()];
2115        for (iso, xs) in work.sigma.iter().enumerate() {
2116            let density = params[self.density_indices[iso]];
2117            let ratio = self.density_ratios[iso];
2118            for (j, &sigma) in xs.iter().enumerate() {
2119                neg_opt[j] -= density * ratio * sigma;
2120            }
2121        }
2122        let t_unbroadened: Vec<f64> = neg_opt.iter().map(|&d| d.exp()).collect();
2123
2124        let Some(inst) = self.instrument.as_ref() else {
2125            // No resolution: the working grid IS the data grid (identity
2126            // layout), so `extract` is a no-op clone.
2127            return Ok(work.layout.extract(&t_unbroadened));
2128        };
2129
2130        // Resolution on the working grid, then extract the data points last
2131        // (issue #442 + #608).  For tabulated resolution the working grid IS
2132        // `e_corr`, so the `(t0, L_scale)`-keyed plan (built on `e_corr`) still
2133        // matches; for Gaussian the plan is `None` and broadening runs on the
2134        // auxiliary grid via `apply_resolution`.
2135        let plan = if use_plan_cache {
2136            let t0 = params[self.t0_index];
2137            let l_scale = params[self.l_scale_index];
2138            self.cached_resolution_plan(t0, l_scale, work_e)
2139        } else {
2140            None
2141        };
2142        let t_broadened = resolution::apply_resolution_with_plan(
2143            plan.as_deref(),
2144            work_e,
2145            &t_unbroadened,
2146            &inst.resolution,
2147        )
2148        .map_err(|e| FittingError::EvaluationFailed(format!("resolution broadening: {e}")))?;
2149        Ok(work.layout.extract(&t_broadened))
2150    }
2151}
2152
2153impl FitModel for EnergyScaleTransmissionModel {
2154    fn evaluate(&self, params: &[f64]) -> Result<Vec<f64>, FittingError> {
2155        let t0 = params[self.t0_index];
2156        let l_scale = params[self.l_scale_index];
2157        let e_corr = self.corrected_energies(t0, l_scale);
2158        // Public `evaluate` uses the plan cache: downstream the
2159        // Jacobian (+ joint-Poisson's gradient + Fisher) will re-call
2160        // `evaluate` at the SAME `(t0, L_scale)` before the next LM
2161        // step, and the density-col path of `analytical_jacobian`
2162        // also wants a plan at this probe — all of those hit the
2163        // cache.  LM's own FD probes — one-coordinate-at-a-time
2164        // central differences at `(t0 ± h, L_scale)` or
2165        // `(t0, L_scale ± h)` — go through a dedicated non-cache
2166        // path in `analytical_jacobian` below, so they don't add
2167        // plan-build overhead.  Issue #483 A1.
2168        self.evaluate_at_with_cache(params, &e_corr, true)
2169    }
2170
2171    /// Jacobian: analytical for density parameters, finite-difference for t₀ and L_scale.
2172    fn analytical_jacobian(
2173        &self,
2174        params: &[f64],
2175        free_param_indices: &[usize],
2176        _y_current: &[f64],
2177    ) -> Option<FlatMatrix> {
2178        let n_e = self.nominal_energies.len();
2179        let n_free = free_param_indices.len();
2180        let mut jacobian = FlatMatrix::zeros(n_e, n_free);
2181
2182        let t0 = params[self.t0_index];
2183        let l_scale = params[self.l_scale_index];
2184        let e_corr = self.corrected_energies(t0, l_scale);
2185        let energy_scale_method = self.jacobian_method;
2186        let t0_free_pos = free_param_indices
2187            .iter()
2188            .position(|&idx| idx == self.t0_index);
2189        let l_scale_free_pos = free_param_indices
2190            .iter()
2191            .position(|&idx| idx == self.l_scale_index);
2192        // Partial-GAL t0 FD pair (precomputed once; the L_scale column
2193        // is derived from this column inline below). Skipped when:
2194        // - method is not PartialGal, OR
2195        // - either t0 or L_scale is fixed (the rank-1 derivation needs
2196        //   both columns paired), OR
2197        // - `t0 + h` would land at or above the `corrected_energies`
2198        //   clamp (`min_tof * (1 - 1e-12)`). At the clamp, both `±h`
2199        //   probes collapse to the same clamped value: the t0 FD column
2200        //   becomes ~0, and the rank-1 L_scale column would also be ~0
2201        //   even though `corrected_energies` does NOT clamp on
2202        //   `L_scale`. Falling through here lets the standard
2203        //   per-coordinate FD path below compute the L_scale column
2204        //   correctly. Issue #489.
2205        let partial_gal_t0_column = if energy_scale_method == EnergyScaleJacobianMethod::PartialGal
2206            && t0_free_pos.is_some()
2207            && l_scale_free_pos.is_some()
2208        {
2209            let h = 1e-4;
2210            let min_tof_us = self
2211                .nominal_energies
2212                .iter()
2213                .map(|&e| self.tof_factor * self.flight_path_m / e.sqrt())
2214                .fold(f64::INFINITY, f64::min);
2215            let t0_limit = min_tof_us * (1.0 - 1.0e-12);
2216            // Need (t0 + h) strictly below the clamp so the +h probe
2217            // returns a distinct corrected grid; otherwise fall through.
2218            if t0 + h >= t0_limit {
2219                None
2220            } else {
2221                let mut p_plus = params.to_vec();
2222                let mut p_minus = params.to_vec();
2223                p_plus[self.t0_index] += h;
2224                p_minus[self.t0_index] -= h;
2225                let e_corr_plus =
2226                    self.corrected_energies(p_plus[self.t0_index], p_plus[self.l_scale_index]);
2227                let e_corr_minus =
2228                    self.corrected_energies(p_minus[self.t0_index], p_minus[self.l_scale_index]);
2229                let y_plus = match self.evaluate_at_with_cache(&p_plus, &e_corr_plus, false) {
2230                    Ok(v) => v,
2231                    Err(_) => return None,
2232                };
2233                let y_minus = match self.evaluate_at_with_cache(&p_minus, &e_corr_minus, false) {
2234                    Ok(v) => v,
2235                    Err(_) => return None,
2236                };
2237                // Per-cell finiteness check.  Without it a NaN in
2238                // `y_plus[i]` / `y_minus[i]` propagates into both the
2239                // t0 column AND the L_scale column derived from it via
2240                // the rank-1 reconstruction at `scale * partial_t0_col[i]`
2241                // (~line 2280), poisoning the post-convergence
2242                // covariance the same way lm.rs `compute_jacobian` was
2243                // vulnerable.  Mirror that fix: zero the entry rather
2244                // than dropping the column — masked rows (NaN by design
2245                // in some test contracts) get skipped downstream by the
2246                // active-mask row-skip in the LM normal-equation
2247                // assembly, so a 0 in a masked row is benign.
2248                let mut col = vec![0.0f64; n_e];
2249                for i in 0..n_e {
2250                    let a = y_plus[i];
2251                    let b = y_minus[i];
2252                    if a.is_finite() && b.is_finite() {
2253                        col[i] = (a - b) / (2.0 * h);
2254                    }
2255                    // else: leave col[i] at 0.0; downstream L_scale
2256                    // reconstruction `scale * 0 == 0` is consistent.
2257                }
2258                Some(col)
2259            }
2260        } else {
2261            None
2262        };
2263
2264        // Issue #608: density columns are formed on the WORKING grid (auxiliary
2265        // extended grid for Gaussian resolution; `e_corr` for tabulated / no
2266        // resolution) from the TRUE σ at the corrected energies (reich_moore +
2267        // Doppler), resolution-broadened there, and the data points extracted
2268        // last — matching `forward_model` and `evaluate`.
2269        let work = match self.working_xs_for(params, &e_corr) {
2270            Ok(w) => w,
2271            Err(_) => return None,
2272        };
2273        let work_layout = &work.layout;
2274        let work_e = &work_layout.energies;
2275
2276        // Unresolved T on the WORKING grid: T = exp(-Σᵢ nᵢ·rᵢ·σᵢ).
2277        let mut neg_opt = vec![0.0f64; work_e.len()];
2278        for (iso, xs) in work.sigma.iter().enumerate() {
2279            let density = params[self.density_indices[iso]];
2280            let ratio = self.density_ratios[iso];
2281            for (j, &sigma) in xs.iter().enumerate() {
2282                neg_opt[j] -= density * ratio * sigma;
2283            }
2284        }
2285        let t_unresolved: Vec<f64> = neg_opt.iter().map(|&d| d.exp()).collect();
2286
2287        // Density-column plan: Issue #483 A1 routes through the
2288        // struct-level `(t0, L_scale)`-keyed cache.  Built on the working grid
2289        // (== `e_corr` for tabulated, where the plan is meaningful; `None` for
2290        // Gaussian).  When `self.evaluate(params)` ran earlier in the same KL
2291        // outer iteration the cache was already populated at the current
2292        // `(t0, L_scale)` and this lookup is a cheap Arc clone.
2293        //
2294        // An earlier `n_density_cols >= 2` gate is
2295        // dropped here: the cache makes the plan build a one-shot
2296        // cost amortized across every evaluate at `(t0, L_scale)` in
2297        // the surrounding KL iteration, so even the N_density = 1
2298        // case (A.1 / KL+grouped+TZERO) now benefits from plan
2299        // reuse across 3 evaluates + 2 jacobians per outer iter.
2300        // The non-tabulated / build-failure branches still return
2301        // `None` → `apply_resolution_with_plan(None, …)` forwards
2302        // byte-identically to `apply_resolution`.
2303        let density_plan = self.cached_resolution_plan(t0, l_scale, work_e);
2304
2305        for (col, &fp_idx) in free_param_indices.iter().enumerate() {
2306            if fp_idx == self.t0_index || fp_idx == self.l_scale_index {
2307                // partial-GAL: when both t0 and L_scale are free, the t0
2308                // column comes from a single pre-computed FD pair (above),
2309                // and the L_scale column is the per-bin rank-1 derivation
2310                //   J[:, L_scale]_i = ((tof_i - t0_clamped) / L_scale) * J[:, t0]_i.
2311                //
2312                // The structural factorisation through `e_corr` holds
2313                // when `R` depends on `(t0, L_scale)` only through
2314                // `e_corr` — `broaden_presorted` uses `self.flight_path_m`
2315                // (not the model's `L_nominal * L_scale`) for
2316                // `tof_center` / `e_prime`, so tabulated kernels satisfy
2317                // it. The per-bin rank-1 simplification additionally
2318                // assumes per-bin homogeneity of `(tof - t0) / L_scale`
2319                // across the kernel support; see the
2320                // `EnergyScaleJacobianMethod` doc for the empirical
2321                // characterisation. When only one of t0 / L_scale is
2322                // free, we fall through to the standard FD path below.
2323                if let Some(partial_t0_col) = &partial_gal_t0_column {
2324                    if fp_idx == self.t0_index {
2325                        for (i, &val) in partial_t0_col.iter().enumerate() {
2326                            *jacobian.get_mut(i, col) = val;
2327                        }
2328                        continue;
2329                    }
2330                    if fp_idx == self.l_scale_index {
2331                        let l_scale = params[self.l_scale_index];
2332                        // Issue #500: at `l_scale ≈ 0` the rank-1 factor
2333                        // `(tof - t0_clamped) / l_scale` blows up and
2334                        // produces NaN columns when combined with the
2335                        // FD-based t0 reference (which goes to ~0 at the
2336                        // same boundary).  Skip the partial-GAL path
2337                        // and fall through to the per-coordinate FD
2338                        // section below — mirrors the t0 clamp-boundary
2339                        // fallthrough (when
2340                        // `partial_gal_t0_column` is `None`, the entire
2341                        // partial-GAL block is skipped).  Production
2342                        // L_scale bounds are typically `[0.99, 1.01]`,
2343                        // so this guard fires only at API edge cases.
2344                        if l_scale.abs() >= L_SCALE_EPSILON {
2345                            let t0 = params[self.t0_index];
2346                            // Match the `corrected_energies` t0 clamp so the
2347                            // (tof - t0) factor in the rank-1 derivation
2348                            // agrees with the production forward at the
2349                            // clamp boundary.
2350                            let min_tof_us = self
2351                                .nominal_energies
2352                                .iter()
2353                                .map(|&e| self.tof_factor * self.flight_path_m / e.sqrt())
2354                                .fold(f64::INFINITY, f64::min);
2355                            let t0_clamped = t0.min(min_tof_us * (1.0 - 1.0e-12));
2356                            for (i, &e_nom) in self.nominal_energies.iter().enumerate() {
2357                                let tof_i = self.tof_factor * self.flight_path_m / e_nom.sqrt();
2358                                let scale = (tof_i - t0_clamped) / l_scale;
2359                                *jacobian.get_mut(i, col) = scale * partial_t0_col[i];
2360                            }
2361                            continue;
2362                        }
2363                        // l_scale ≈ 0: do NOT `continue`; flow falls
2364                        // through to the FD path below for this column.
2365                    }
2366                }
2367                // Finite difference for energy-scale parameters.
2368                //
2369                // Central-difference probes perturb one coordinate at
2370                // a time: `(t0 ± h, L_scale)` when differentiating in
2371                // `t0`, or `(t0, L_scale ± h)` when differentiating
2372                // in `L_scale`.  Each perturbed point is a distinct
2373                // `(t0, L_scale)` key that would miss the struct
2374                // plan cache, and building a plan for the probe has
2375                // no reuse to amortize.  Route them through
2376                // `evaluate_at_with_cache(..., false)` so they stay
2377                // on the original non-plan `apply_resolution` path.
2378                // The public `FitModel::evaluate` path continues to
2379                // use the cache for the many-uses-per-probe callers
2380                // (KL solver's deviance + gradient + Fisher at the
2381                // current probe).  Issue #483 A1.
2382                let h = if fp_idx == self.t0_index { 1e-4 } else { 1e-7 };
2383                let mut p_plus = params.to_vec();
2384                let mut p_minus = params.to_vec();
2385                p_plus[fp_idx] += h;
2386                p_minus[fp_idx] -= h;
2387                let t0_plus = p_plus[self.t0_index];
2388                let l_plus = p_plus[self.l_scale_index];
2389                let t0_minus = p_minus[self.t0_index];
2390                let l_minus = p_minus[self.l_scale_index];
2391                let e_corr_plus = self.corrected_energies(t0_plus, l_plus);
2392                let e_corr_minus = self.corrected_energies(t0_minus, l_minus);
2393                let y_plus = match self.evaluate_at_with_cache(&p_plus, &e_corr_plus, false) {
2394                    Ok(v) => v,
2395                    Err(_) => return None,
2396                };
2397                let y_minus = match self.evaluate_at_with_cache(&p_minus, &e_corr_minus, false) {
2398                    Ok(v) => v,
2399                    Err(_) => return None,
2400                };
2401                // Per-cell finiteness check — mirrors the lm.rs
2402                // `compute_jacobian` FD path.  A NaN in the perturbed
2403                // model at an active row would otherwise feed NaN
2404                // through the post-convergence covariance; per-cell
2405                // skip leaves masked-row NaN benign (the LM normal-
2406                // equation assembly already row-skips those).
2407                for i in 0..n_e {
2408                    let a = y_plus[i];
2409                    let b = y_minus[i];
2410                    if a.is_finite() && b.is_finite() {
2411                        *jacobian.get_mut(i, col) = (a - b) / (2.0 * h);
2412                    }
2413                    // else: leave at zero-default.
2414                }
2415            } else {
2416                // Density parameter: analytical derivative on the WORKING grid
2417                // (issue #608) from the TRUE σ, resolution-broadened there,
2418                // data points extracted last.
2419                // ∂T/∂n_g = extract(R[-(Σ_{iso∈g} rᵢ·σ_iso(E)) · T_unresolved(E)])
2420                let mut sigma_sum = vec![0.0f64; work_e.len()];
2421                for (iso, &di) in self.density_indices.iter().enumerate() {
2422                    if di == fp_idx {
2423                        let ratio = self.density_ratios[iso];
2424                        for (j, &sigma) in work.sigma[iso].iter().enumerate() {
2425                            sigma_sum[j] += ratio * sigma;
2426                        }
2427                    }
2428                }
2429                let inner_deriv: Vec<f64> = (0..work_e.len())
2430                    .map(|i| -sigma_sum[i] * t_unresolved[i])
2431                    .collect();
2432
2433                // Apply resolution to derivative if enabled.
2434                //
2435                // When `density_plan` is `Some` (tabulated resolution
2436                // + populated cache) we hit the struct-level
2437                // `(t0, L_scale)`-keyed plan (built on the working grid, which
2438                // equals `e_corr` for tabulated).  When `None` (Gaussian
2439                // resolution or build failure),
2440                // `apply_resolution_with_plan(None, …)` transparently
2441                // forwards to `apply_resolution` — bit-exact with
2442                // the pre-cache path.  Issue #483 A1.
2443                if let Some(inst) = &self.instrument {
2444                    let resolved_deriv = match resolution::apply_resolution_with_plan(
2445                        density_plan.as_deref(),
2446                        work_e,
2447                        &inner_deriv,
2448                        &inst.resolution,
2449                    ) {
2450                        Ok(v) => v,
2451                        Err(_) => return None,
2452                    };
2453                    let resolved_deriv = work_layout.extract(&resolved_deriv);
2454                    for (i, &val) in resolved_deriv.iter().enumerate() {
2455                        *jacobian.get_mut(i, col) = val;
2456                    }
2457                } else {
2458                    // No resolution → identity layout, inner is already data grid.
2459                    for (i, &val) in inner_deriv.iter().enumerate() {
2460                        *jacobian.get_mut(i, col) = val;
2461                    }
2462                }
2463            }
2464        }
2465
2466        Some(jacobian)
2467    }
2468}
2469
2470// ── ForwardModel implementations (Phase 1) ──────────────────────────────
2471//
2472// Each implementation delegates to the existing FitModel logic.
2473// `predict` == `evaluate`, `jacobian` converts FlatMatrix → Vec<Vec<f64>>.
2474
2475use crate::forward_model::ForwardModel;
2476
2477impl ForwardModel for PrecomputedTransmissionModel {
2478    fn predict(&self, params: &[f64]) -> Result<Vec<f64>, FittingError> {
2479        self.evaluate(params)
2480    }
2481
2482    fn jacobian(
2483        &self,
2484        params: &[f64],
2485        free_param_indices: &[usize],
2486        y_current: &[f64],
2487    ) -> Option<Vec<Vec<f64>>> {
2488        let fm = self.analytical_jacobian(params, free_param_indices, y_current)?;
2489        Some(flat_matrix_to_vecs(&fm, free_param_indices.len()))
2490    }
2491
2492    fn n_data(&self) -> usize {
2493        // Issue #608: when a Gaussian working-grid layout is attached,
2494        // `cross_sections` lives on the (longer) working grid, but the number of
2495        // DATA points the model predicts is the layout's data-index count.
2496        // Without a layout the working grid IS the data grid.
2497        if let Some(layout) = &self.work_layout {
2498            layout.data_indices.len()
2499        } else if self.cross_sections.is_empty() {
2500            0
2501        } else {
2502            self.cross_sections[0].len()
2503        }
2504    }
2505
2506    fn n_params(&self) -> usize {
2507        // Max index in density_indices + 1
2508        self.density_indices
2509            .iter()
2510            .copied()
2511            .max()
2512            .map_or(0, |m| m + 1)
2513    }
2514}
2515
2516impl ForwardModel for TransmissionFitModel {
2517    fn predict(&self, params: &[f64]) -> Result<Vec<f64>, FittingError> {
2518        self.evaluate(params)
2519    }
2520
2521    fn jacobian(
2522        &self,
2523        params: &[f64],
2524        free_param_indices: &[usize],
2525        y_current: &[f64],
2526    ) -> Option<Vec<Vec<f64>>> {
2527        let fm = self.analytical_jacobian(params, free_param_indices, y_current)?;
2528        Some(flat_matrix_to_vecs(&fm, free_param_indices.len()))
2529    }
2530
2531    fn n_data(&self) -> usize {
2532        self.energies.len()
2533    }
2534
2535    fn n_params(&self) -> usize {
2536        let max_density = self.density_indices.iter().copied().max().unwrap_or(0);
2537        let max_temp = self.temperature_index.unwrap_or(0);
2538        max_density.max(max_temp) + 1
2539    }
2540}
2541
2542impl<M: FitModel> ForwardModel for NormalizedTransmissionModel<M> {
2543    fn predict(&self, params: &[f64]) -> Result<Vec<f64>, FittingError> {
2544        self.evaluate(params)
2545    }
2546
2547    fn jacobian(
2548        &self,
2549        params: &[f64],
2550        free_param_indices: &[usize],
2551        y_current: &[f64],
2552    ) -> Option<Vec<Vec<f64>>> {
2553        let fm = self.analytical_jacobian(params, free_param_indices, y_current)?;
2554        Some(flat_matrix_to_vecs(&fm, free_param_indices.len()))
2555    }
2556
2557    fn n_data(&self) -> usize {
2558        self.sqrt_energies.len()
2559    }
2560
2561    fn n_params(&self) -> usize {
2562        // The background indices are the highest parameter indices.
2563        let mut max_idx = self
2564            .anorm_index
2565            .max(self.back_a_index)
2566            .max(self.back_b_index)
2567            .max(self.back_c_index);
2568        if let Some(di) = self.back_d_index {
2569            max_idx = max_idx.max(di);
2570        }
2571        if let Some(fi) = self.back_f_index {
2572            max_idx = max_idx.max(fi);
2573        }
2574        max_idx + 1
2575    }
2576}
2577
2578impl ForwardModel for EnergyScaleTransmissionModel {
2579    fn predict(&self, params: &[f64]) -> Result<Vec<f64>, FittingError> {
2580        self.evaluate(params)
2581    }
2582
2583    fn jacobian(
2584        &self,
2585        params: &[f64],
2586        free_param_indices: &[usize],
2587        y_current: &[f64],
2588    ) -> Option<Vec<Vec<f64>>> {
2589        let fm = self.analytical_jacobian(params, free_param_indices, y_current)?;
2590        Some(flat_matrix_to_vecs(&fm, free_param_indices.len()))
2591    }
2592
2593    fn n_data(&self) -> usize {
2594        self.nominal_energies.len()
2595    }
2596
2597    fn n_params(&self) -> usize {
2598        self.t0_index.max(self.l_scale_index) + 1
2599    }
2600}
2601
2602/// Convert a `FlatMatrix` (row-major) to `Vec<Vec<f64>>` (column-major).
2603///
2604/// Returns `cols` vectors, each of length `fm.nrows()`.
2605fn flat_matrix_to_vecs(fm: &FlatMatrix, cols: usize) -> Vec<Vec<f64>> {
2606    let nrows = fm.nrows;
2607    (0..cols)
2608        .map(|j| (0..nrows).map(|i| fm.get(i, j)).collect())
2609        .collect()
2610}
2611
2612#[cfg(test)]
2613mod tests {
2614    use super::*;
2615    use crate::lm::{self, FitModel, LmConfig};
2616    use crate::parameters::{FitParameter, ParameterSet};
2617    use nereids_core::types::Isotope;
2618    use nereids_endf::resonance::test_support::u238_single_resonance;
2619    use nereids_endf::resonance::{LGroup, Resonance, ResonanceFormalism, ResonanceRange};
2620
2621    /// ∞-norm of the residual between two equal-length spectra.
2622    /// (Issue #608 aux-grid regression-test helper.)
2623    fn max_abs_diff(a: &[f64], b: &[f64]) -> f64 {
2624        a.iter()
2625            .zip(b.iter())
2626            .map(|(x, y)| (x - y).abs())
2627            .fold(0.0f64, f64::max)
2628    }
2629
2630    /// ∞-norm (max |value|) of a spectrum — a scale for relative thresholds.
2631    fn max_abs(a: &[f64]) -> f64 {
2632        a.iter().map(|x| x.abs()).fold(0.0f64, f64::max)
2633    }
2634
2635    // ── PrecomputedTransmissionModel ─────────────────────────────────────────
2636
2637    /// Verify Beer-Lambert: T(E) = exp(-Σᵢ nᵢ·σᵢ(E)).
2638    #[test]
2639    fn precomputed_evaluate_matches_beer_lambert() {
2640        let model = make_precomputed(
2641            vec![
2642                vec![1.0, 2.0, 3.0], // isotope 0
2643                vec![0.5, 0.5, 0.5], // isotope 1
2644            ],
2645            vec![0, 1],
2646        );
2647
2648        let params = [0.2f64, 0.4f64];
2649        let y = model.evaluate(&params).unwrap();
2650
2651        let expected: Vec<f64> = (0..3)
2652            .map(|i| {
2653                let s0 = [1.0, 2.0, 3.0][i];
2654                let s1 = [0.5, 0.5, 0.5][i];
2655                (-params[0] * s0 - params[1] * s1).exp()
2656            })
2657            .collect();
2658
2659        assert_eq!(y.len(), 3);
2660        for (yi, ei) in y.iter().zip(expected.iter()) {
2661            assert!(
2662                (yi - ei).abs() < 1e-12,
2663                "evaluate mismatch: got {yi}, expected {ei}"
2664            );
2665        }
2666    }
2667
2668    /// Analytical Jacobian ∂T/∂nᵢ = -σᵢ(E)·T(E) must match central-difference FD.
2669    #[test]
2670    fn precomputed_analytical_jacobian_matches_finite_difference() {
2671        let model = make_precomputed(
2672            vec![
2673                vec![1.0, 2.0, 3.0], // isotope 0
2674                vec![0.5, 0.5, 0.5], // isotope 1
2675            ],
2676            vec![0, 1],
2677        );
2678
2679        let params = [0.2f64, 0.4f64];
2680        let y = model.evaluate(&params).unwrap();
2681        let free = vec![0usize, 1usize];
2682
2683        let jac = model
2684            .analytical_jacobian(&params, &free, &y)
2685            .expect("analytical_jacobian should return Some(_)");
2686
2687        assert_eq!(jac.nrows, 3); // n_energies
2688        assert_eq!(jac.ncols, 2); // n_free_params
2689
2690        // Central-difference reference.
2691        let h = 1e-6f64;
2692        for (col, &p_idx) in free.iter().enumerate() {
2693            let mut p_plus = params;
2694            let mut p_minus = params;
2695            p_plus[p_idx] += h;
2696            p_minus[p_idx] -= h;
2697
2698            let y_plus = model.evaluate(&p_plus).unwrap();
2699            let y_minus = model.evaluate(&p_minus).unwrap();
2700
2701            for i in 0..3 {
2702                let fd = (y_plus[i] - y_minus[i]) / (2.0 * h);
2703                let ana = jac.get(i, col);
2704                assert!(
2705                    (fd - ana).abs() < 1e-6,
2706                    "Jacobian mismatch (row {i}, col {col}): FD={fd:.8}, analytical={ana:.8}"
2707                );
2708            }
2709        }
2710    }
2711
2712    /// When two isotopes share a density parameter, the Jacobian column must
2713    /// equal -T(E) * (σ₀(E) + σ₁(E)), not just the first isotope's σ.
2714    #[test]
2715    fn precomputed_jacobian_tied_parameters_sums_both_isotopes() {
2716        // Two isotopes mapped to the same density parameter (index 0).
2717        let model = make_precomputed(
2718            vec![
2719                vec![1.0, 2.0, 3.0], // isotope 0
2720                vec![0.5, 1.0, 1.5], // isotope 1 — tied to same param
2721            ],
2722            vec![0, 0], // both isotopes share param[0]
2723        );
2724
2725        let params = [0.1f64];
2726        let y = model.evaluate(&params).unwrap();
2727        let free = vec![0usize];
2728
2729        let jac = model
2730            .analytical_jacobian(&params, &free, &y)
2731            .expect("analytical_jacobian should return Some(_)");
2732
2733        // Expected: ∂T/∂n = -T(E) * (σ₀(E) + σ₁(E))
2734        for i in 0..3 {
2735            let sigma_sum = [1.0, 2.0, 3.0][i] + [0.5, 1.0, 1.5][i];
2736            let expected = -y[i] * sigma_sum;
2737            assert!(
2738                (jac.get(i, 0) - expected).abs() < 1e-12,
2739                "Tied Jacobian mismatch at E[{i}]: got {}, expected {expected}",
2740                jac.get(i, 0)
2741            );
2742        }
2743    }
2744
2745    // ── TransmissionFitModel ─────────────────────────────────────────────────
2746
2747    #[test]
2748    fn test_recover_single_isotope_thickness() {
2749        let data = u238_single_resonance();
2750        let true_thickness = 0.0005;
2751
2752        // Generate synthetic data
2753        let energies: Vec<f64> = (0..201).map(|i| 1.0 + (i as f64) * 0.05).collect();
2754
2755        let model = TransmissionFitModel::new(
2756            energies.clone(),
2757            vec![data],
2758            0.0,
2759            None,
2760            (vec![0], vec![1.0]),
2761            None,
2762            None,
2763        )
2764        .unwrap();
2765
2766        let y_obs = model.evaluate(&[true_thickness]).unwrap();
2767        let sigma = vec![0.01; y_obs.len()]; // 1% uncertainty
2768
2769        let mut params = ParameterSet::new(vec![
2770            FitParameter::non_negative("thickness", 0.001), // initial guess 2× off
2771        ]);
2772
2773        let result =
2774            lm::levenberg_marquardt(&model, &y_obs, &sigma, &mut params, &LmConfig::default())
2775                .unwrap();
2776
2777        assert!(result.converged, "Fit did not converge");
2778        let fitted = result.params[0];
2779        assert!(
2780            (fitted - true_thickness).abs() / true_thickness < 0.01,
2781            "Fitted thickness = {}, true = {}, error = {:.1}%",
2782            fitted,
2783            true_thickness,
2784            (fitted - true_thickness).abs() / true_thickness * 100.0,
2785        );
2786    }
2787
2788    #[test]
2789    fn test_recover_two_isotope_thicknesses() {
2790        let u238 = u238_single_resonance();
2791
2792        // Second isotope with resonance at 20 eV
2793        let other = ResonanceData {
2794            isotope: Isotope::new(1, 10).unwrap(),
2795            za: 1010,
2796            awr: 10.0,
2797            ranges: vec![ResonanceRange {
2798                energy_low: 0.0,
2799                energy_high: 100.0,
2800                resolved: true,
2801                formalism: ResonanceFormalism::ReichMoore,
2802                target_spin: 0.0,
2803                scattering_radius: 5.0,
2804                naps: 1,
2805                l_groups: vec![LGroup {
2806                    l: 0,
2807                    awr: 10.0,
2808                    apl: 5.0,
2809                    qx: 0.0,
2810                    lrx: 0,
2811                    resonances: vec![Resonance {
2812                        energy: 20.0,
2813                        j: 0.5,
2814                        gn: 0.1,
2815                        gg: 0.05,
2816                        gfa: 0.0,
2817                        gfb: 0.0,
2818                    }],
2819                }],
2820                rml: None,
2821                urr: None,
2822                ap_table: None,
2823                r_external: vec![],
2824            }],
2825        };
2826
2827        let true_t1 = 0.0003;
2828        let true_t2 = 0.0001;
2829
2830        let energies: Vec<f64> = (0..301).map(|i| 1.0 + (i as f64) * 0.1).collect();
2831
2832        let model = TransmissionFitModel::new(
2833            energies.clone(),
2834            vec![u238, other],
2835            0.0,
2836            None,
2837            (vec![0, 1], vec![1.0, 1.0]),
2838            None,
2839            None,
2840        )
2841        .unwrap();
2842
2843        let y_obs = model.evaluate(&[true_t1, true_t2]).unwrap();
2844        let sigma = vec![0.01; y_obs.len()];
2845
2846        let mut params = ParameterSet::new(vec![
2847            FitParameter::non_negative("U-238 thickness", 0.001),
2848            FitParameter::non_negative("Other thickness", 0.001),
2849        ]);
2850
2851        let result =
2852            lm::levenberg_marquardt(&model, &y_obs, &sigma, &mut params, &LmConfig::default())
2853                .unwrap();
2854
2855        assert!(
2856            result.converged,
2857            "Fit did not converge after {} iterations",
2858            result.iterations
2859        );
2860
2861        let (fit_t1, fit_t2) = (result.params[0], result.params[1]);
2862        assert!(
2863            (fit_t1 - true_t1).abs() / true_t1 < 0.05,
2864            "U-238: fitted={}, true={}, error={:.1}%",
2865            fit_t1,
2866            true_t1,
2867            (fit_t1 - true_t1).abs() / true_t1 * 100.0,
2868        );
2869        assert!(
2870            (fit_t2 - true_t2).abs() / true_t2 < 0.05,
2871            "Other: fitted={}, true={}, error={:.1}%",
2872            fit_t2,
2873            true_t2,
2874            (fit_t2 - true_t2).abs() / true_t2 * 100.0,
2875        );
2876    }
2877
2878    // ── Temperature fitting ──────────────────────────────────────────────────
2879
2880    /// Verify that temperature_index makes evaluate() read T from the
2881    /// parameter vector instead of the fixed `temperature_k` field.
2882    #[test]
2883    fn temperature_index_overrides_fixed_temperature() {
2884        let data = u238_single_resonance();
2885        let energies: Vec<f64> = (0..201).map(|i| 1.0 + (i as f64) * 0.05).collect();
2886
2887        // Model with fixed temperature = 0 K but temperature_index pointing
2888        // to params[1].
2889        let model = TransmissionFitModel::new(
2890            energies.clone(),
2891            vec![data.clone()],
2892            0.0,
2893            None,
2894            (vec![0], vec![1.0]),
2895            Some(1),
2896            None,
2897        )
2898        .unwrap();
2899
2900        // Model with fixed temperature = 300 K (no temperature_index).
2901        let model_fixed = TransmissionFitModel::new(
2902            energies.clone(),
2903            vec![data],
2904            300.0,
2905            None,
2906            (vec![0], vec![1.0]),
2907            None,
2908            None,
2909        )
2910        .unwrap();
2911
2912        let density = 0.0005;
2913        let y_via_index = model.evaluate(&[density, 300.0]).unwrap();
2914        let y_via_fixed = model_fixed.evaluate(&[density]).unwrap();
2915
2916        for (a, b) in y_via_index.iter().zip(y_via_fixed.iter()) {
2917            assert!(
2918                (a - b).abs() < 1e-12,
2919                "temperature_index path disagrees with fixed path: {} vs {}",
2920                a,
2921                b
2922            );
2923        }
2924    }
2925
2926    /// Recover temperature from Doppler-broadened synthetic data.
2927    ///
2928    /// Generates transmission at T_true with known density, then fits both
2929    /// density and temperature simultaneously.
2930    #[test]
2931    fn test_recover_temperature() {
2932        let data = u238_single_resonance();
2933        let true_density = 0.0005;
2934        let true_temp = 300.0; // K
2935
2936        // Energy grid around the 6.674 eV resonance.
2937        let energies: Vec<f64> = (0..401).map(|i| 4.0 + (i as f64) * 0.025).collect();
2938
2939        // Generate synthetic data at the true temperature.
2940        let model = TransmissionFitModel::new(
2941            energies.clone(),
2942            vec![data],
2943            0.0, // ignored — temperature_index is set
2944            None,
2945            (vec![0], vec![1.0]),
2946            Some(1), // params[1] = temperature
2947            None,
2948        )
2949        .unwrap();
2950
2951        let mut y_obs = model.evaluate(&[true_density, true_temp]).unwrap();
2952        // Add tiny deterministic noise so reduced_chi2 stays positive.
2953        // Without noise, the analytical Jacobian converges to exact parameters,
2954        // yielding chi2 ≈ 0, which makes covariance ≈ 0 and uncertainty NaN.
2955        for (i, y) in y_obs.iter_mut().enumerate() {
2956            *y *= 1.0 + 1e-5 * ((i % 7) as f64 - 3.0);
2957        }
2958        let sigma = vec![0.005; y_obs.len()];
2959
2960        // Fit with initial guesses offset from truth.
2961        let mut params = ParameterSet::new(vec![
2962            FitParameter::non_negative("density", 0.001),
2963            FitParameter {
2964                name: "temperature_k".into(),
2965                value: 200.0, // initial guess 100 K off
2966                lower: 1.0,
2967                upper: 2000.0,
2968                fixed: false,
2969            },
2970        ]);
2971
2972        let config = LmConfig {
2973            max_iter: 200,
2974            ..LmConfig::default()
2975        };
2976
2977        let result = lm::levenberg_marquardt(&model, &y_obs, &sigma, &mut params, &config).unwrap();
2978
2979        assert!(
2980            result.converged,
2981            "Temperature fit did not converge after {} iterations",
2982            result.iterations
2983        );
2984
2985        let fit_density = result.params[0];
2986        let fit_temp = result.params[1];
2987
2988        // Tiny deterministic noise (max ±3e-5): optimizer should converge to within 0.1%.
2989        assert!(
2990            (fit_density - true_density).abs() / true_density < 0.001,
2991            "Density: fitted={}, true={}, error={:.1}%",
2992            fit_density,
2993            true_density,
2994            (fit_density - true_density).abs() / true_density * 100.0,
2995        );
2996        assert!(
2997            (fit_temp - true_temp).abs() / true_temp < 0.001,
2998            "Temperature: fitted={:.1} K, true={:.1} K, error={:.1}%",
2999            fit_temp,
3000            true_temp,
3001            (fit_temp - true_temp).abs() / true_temp * 100.0,
3002        );
3003
3004        // Verify uncertainty is reported.
3005        let unc = result
3006            .uncertainties
3007            .expect("uncertainties should be available");
3008        assert!(
3009            unc.len() == 2,
3010            "expected 2 uncertainties, got {}",
3011            unc.len()
3012        );
3013        assert!(
3014            unc[1] > 0.0 && unc[1].is_finite(),
3015            "temperature uncertainty should be positive and finite, got {}",
3016            unc[1]
3017        );
3018    }
3019
3020    /// Analytical Jacobian for TransmissionFitModel (with temperature) must
3021    /// agree with central-difference finite-difference Jacobian.
3022    ///
3023    /// This validates both the density columns (∂T/∂nᵢ = -σᵢ·T) and the
3024    /// temperature column (forward FD at T+dT).
3025    #[test]
3026    fn transmission_fit_model_analytical_jacobian_matches_fd() {
3027        let data = u238_single_resonance();
3028        let energies: Vec<f64> = (0..201).map(|i| 1.0 + (i as f64) * 0.05).collect();
3029
3030        let model = TransmissionFitModel::new(
3031            energies,
3032            vec![data],
3033            0.0,
3034            None,
3035            (vec![0], vec![1.0]),
3036            Some(1), // params[1] = temperature
3037            None,
3038        )
3039        .unwrap();
3040
3041        let params = [0.0005f64, 300.0f64]; // density, temperature
3042        let y = model.evaluate(&params).unwrap();
3043        let free = vec![0usize, 1usize];
3044
3045        let jac = model
3046            .analytical_jacobian(&params, &free, &y)
3047            .expect("analytical_jacobian should return Some(_)");
3048
3049        assert_eq!(jac.nrows, y.len());
3050        assert_eq!(jac.ncols, 2);
3051
3052        // Central-difference reference.
3053        let h = 1e-6f64;
3054        for (col, &p_idx) in free.iter().enumerate() {
3055            let mut p_plus = params;
3056            let mut p_minus = params;
3057            p_plus[p_idx] += h * (1.0 + params[p_idx].abs());
3058            p_minus[p_idx] -= h * (1.0 + params[p_idx].abs());
3059
3060            let y_plus = model.evaluate(&p_plus).unwrap();
3061            let y_minus = model.evaluate(&p_minus).unwrap();
3062
3063            let actual_2h = p_plus[p_idx] - p_minus[p_idx];
3064            for i in 0..y.len() {
3065                let fd = (y_plus[i] - y_minus[i]) / actual_2h;
3066                let ana = jac.get(i, col);
3067                let err = (fd - ana).abs();
3068                // Use a meaningful floor: when both FD and analytical values
3069                // are below 1e-10, relative error comparisons are dominated
3070                // by floating-point noise and are not physically meaningful.
3071                //
3072                // The floor was raised from 1e-15 to 1e-10 alongside the
3073                // B=S_l boundary condition fix in the Reich-Moore U-matrix.
3074                // That fix shifted near-zero cross-section values from
3075                // O(1e-15) to O(1e-10), making the old floor too tight for
3076                // floating-point comparison at those magnitudes.
3077                let scale = fd.abs().max(ana.abs()).max(1e-10);
3078                assert!(
3079                    err / scale < 0.01,
3080                    "Jacobian mismatch (row {i}, col {col}): FD={fd:.8}, analytical={ana:.8}, \
3081                     rel_err={:.4}",
3082                    err / scale,
3083                );
3084            }
3085        }
3086    }
3087
3088    /// Verify that the broadened-XS cache avoids redundant recomputation.
3089    /// Calling evaluate() twice with the same temperature should produce
3090    /// identical results and reuse the cache.
3091    #[test]
3092    fn transmission_fit_model_cache_reuse() {
3093        let data = u238_single_resonance();
3094        let energies: Vec<f64> = (0..201).map(|i| 1.0 + (i as f64) * 0.05).collect();
3095
3096        let model = TransmissionFitModel::new(
3097            energies,
3098            vec![data],
3099            0.0,
3100            None,
3101            (vec![0], vec![1.0]),
3102            Some(1),
3103            None,
3104        )
3105        .unwrap();
3106
3107        // First call populates the cache.
3108        let y1 = model.evaluate(&[0.0005, 300.0]).unwrap();
3109        assert!(model.cached_broadened_xs.borrow().is_some());
3110        assert!((model.cached_temperature.get() - 300.0).abs() < 1e-15);
3111
3112        // Second call with same temperature but different density should
3113        // reuse cached broadened XS (no rebroadening).
3114        let y2 = model.evaluate(&[0.001, 300.0]).unwrap();
3115        assert!((model.cached_temperature.get() - 300.0).abs() < 1e-15);
3116
3117        // Results must differ (different density) but cache temperature unchanged.
3118        assert!(
3119            (y1[100] - y2[100]).abs() > 1e-10,
3120            "different densities should produce different transmission"
3121        );
3122
3123        // Change temperature — cache should update.
3124        let _y3 = model.evaluate(&[0.0005, 600.0]).unwrap();
3125        assert!((model.cached_temperature.get() - 600.0).abs() < 1e-15);
3126    }
3127
3128    // ── NormalizedTransmissionModel ─────────────────────────────────────────
3129
3130    /// Helper: make a PrecomputedTransmissionModel with given cross-sections
3131    /// and no resolution (Beer-Lambert only).
3132    fn make_precomputed(
3133        xs: Vec<Vec<f64>>,
3134        density_indices: Vec<usize>,
3135    ) -> PrecomputedTransmissionModel {
3136        PrecomputedTransmissionModel {
3137            cross_sections: Arc::new(xs),
3138            density_indices: Arc::new(density_indices),
3139            energies: None,
3140            instrument: None,
3141            resolution_plan: None,
3142            sparse_cubature_plan: None,
3143            sparse_scalar_plan: None,
3144            work_layout: None,
3145        }
3146    }
3147
3148    // ── Cubature dispatch tests ─────────────────────────────────────────
3149
3150    /// Helper: build a synthetic resolution kernel + plan + matrix.
3151    /// CI-hermetic (no PLEIADES fixture) using the same synthetic-
3152    /// overlap-plan pattern as the surrogate module's tests.
3153    fn synthetic_resolution_setup(
3154        n_grid: usize,
3155        half_kernel: usize,
3156    ) -> (
3157        Vec<f64>,
3158        Arc<ResolutionPlan>,
3159        nereids_physics::resolution::ResolutionMatrix,
3160    ) {
3161        assert!(n_grid > 2 * half_kernel);
3162        let energies: Vec<f64> = (0..n_grid).map(|i| 10.0 + i as f64).collect();
3163        let mut starts: Vec<u32> = vec![0];
3164        let mut lo_idx: Vec<u32> = Vec::new();
3165        let mut frac_arr: Vec<f64> = Vec::new();
3166        let mut weight_arr: Vec<f64> = Vec::new();
3167        let mut norm: Vec<f64> = Vec::with_capacity(n_grid);
3168        for i in 0..n_grid {
3169            let lo_min = i.saturating_sub(half_kernel);
3170            let lo_max = (i + half_kernel).min(n_grid - 2);
3171            let mut row_norm = 0.0_f64;
3172            for lo in lo_min..=lo_max {
3173                let d = (lo as i64 - i as i64).abs() as f64;
3174                let w = 1.0 - d / (half_kernel as f64 + 1.0);
3175                lo_idx.push(lo as u32);
3176                frac_arr.push(0.5);
3177                weight_arr.push(w);
3178                row_norm += w;
3179            }
3180            norm.push(row_norm);
3181            starts.push(lo_idx.len() as u32);
3182        }
3183        let plan = nereids_physics::resolution::test_support::plan_from_raw_parts(
3184            energies.clone(),
3185            starts,
3186            lo_idx,
3187            frac_arr,
3188            weight_arr,
3189            norm,
3190        );
3191        let matrix = plan.compile_to_matrix();
3192        (energies, Arc::new(plan), matrix)
3193    }
3194
3195    /// Helper: build a k-isotope synthetic σ stack.
3196    fn synthetic_sigmas(n_grid: usize, k: usize) -> Vec<Vec<f64>> {
3197        let mut out = Vec::with_capacity(k);
3198        for j in 0..k {
3199            let center = 10.0 + (j as f64 + 1.0) * (n_grid as f64) / (k as f64 + 1.0);
3200            let width = 3.0;
3201            out.push(
3202                (0..n_grid)
3203                    .map(|ell| {
3204                        let e = 10.0 + ell as f64;
3205                        100.0 * (-((e - center).powi(2)) / (width * width)).exp() + 5.0
3206                    })
3207                    .collect(),
3208            );
3209        }
3210        out
3211    }
3212
3213    /// Helper: build a sparse cubature plan against a known
3214    /// (matrix, σ stack) pair, with the canonical design-study training
3215    /// rule.
3216    fn build_cubature(
3217        matrix: &nereids_physics::resolution::ResolutionMatrix,
3218        sigmas: &[Vec<f64>],
3219        train_max: Vec<f64>,
3220    ) -> Arc<SparseEmpiricalCubaturePlan> {
3221        let k = sigmas.len();
3222        let n_rows = matrix.len();
3223        let mut flat = Vec::with_capacity(k * n_rows);
3224        for row in sigmas {
3225            flat.extend_from_slice(row);
3226        }
3227        let training = SparseEmpiricalCubaturePlan::default_training_points(&train_max);
3228        let anchor = SparseEmpiricalCubaturePlan::default_jacobian_anchor(&train_max);
3229        Arc::new(
3230            SparseEmpiricalCubaturePlan::build(matrix, &flat, k, &training, &anchor)
3231                .expect("synthetic cubature build"),
3232        )
3233    }
3234
3235    /// Build an `InstrumentParams` wrapping a trivial delta-like
3236    /// tabulated resolution (single ref energy, δ-kernel).  Used
3237    /// only because the dispatch guards check `instrument.is_some()`
3238    /// AND require `ResolutionFunction::Tabulated(_)`.  The actual
3239    /// broadening wouldn't fire on the cubature path regardless
3240    /// (cubature folds `apply_resolution*` into its atom sweep).
3241    fn make_trivial_instrument() -> Arc<InstrumentParams> {
3242        use nereids_physics::resolution::ResolutionFunction;
3243        // Tabulated resolution required for cubature-dispatch tests:
3244        // the eligibility guard refuses the dispatch when the active
3245        // instrument resolution isn't `ResolutionFunction::Tabulated`.
3246        // The test_support helper builds a minimal delta-like kernel;
3247        // the broadening never actually runs on the cubature path
3248        // (cubature.forward replaces apply_resolution entirely).
3249        let tab =
3250            Arc::new(nereids_physics::resolution::test_support::trivial_tabulated_resolution(25.0));
3251        let res_fn = ResolutionFunction::Tabulated(tab);
3252        Arc::new(InstrumentParams { resolution: res_fn })
3253    }
3254
3255    #[test]
3256    fn precomputed_cubature_dispatches_at_k2_matching_k() {
3257        // k = 2 with an installed cubature plan: evaluate should
3258        // return the cubature's forward output (which differs from
3259        // the exact `exp(-Σ n σ) + apply_r` path ONLY at held-out
3260        // densities; at training densities the LP pins them exactly).
3261        let n_grid = 40_usize;
3262        let (energies, plan, matrix) = synthetic_resolution_setup(n_grid, 4);
3263        let sigmas = synthetic_sigmas(n_grid, 2);
3264        let train_max = vec![1e-4_f64, 1e-4];
3265        let cubature = build_cubature(&matrix, &sigmas, train_max.clone());
3266
3267        // Build the model with cubature installed.  The resolution
3268        // plan MUST also be installed for the cubature dispatch to
3269        // fire — without it, `cubature_eligible` refuses the plan
3270        // on the grounds that the cubature would be silently
3271        // bypassing an unknown resolution operator.
3272        let mut model = PrecomputedTransmissionModel {
3273            cross_sections: Arc::new(sigmas.clone()),
3274            density_indices: Arc::new(vec![0, 1]),
3275            energies: Some(Arc::new(energies.clone())),
3276            instrument: Some(make_trivial_instrument()),
3277            resolution_plan: Some(Arc::clone(&plan)),
3278            sparse_cubature_plan: Some(Arc::clone(&cubature)),
3279            sparse_scalar_plan: None,
3280            work_layout: None,
3281        };
3282
3283        // Evaluate at a training density: cubature ≡ exact to LP
3284        // tolerance.
3285        let n = [0.25 * train_max[0], 0.25 * train_max[1]];
3286        let t_cubature = model.evaluate(&n).unwrap();
3287
3288        // Disable cubature → exact cannot match bit-for-bit (different
3289        // summation order).  But we can compute the cubature output
3290        // directly and confirm it equals what `evaluate()` returned.
3291        model.sparse_cubature_plan = None;
3292        let t_exact_via_r = {
3293            // exp(-Σ n σ) then apply_r
3294            let n_grid_local = n_grid;
3295            let mut neg_opt = vec![0.0_f64; n_grid_local];
3296            for (j, &nj) in n.iter().enumerate() {
3297                for (ell, &sig) in sigmas[j].iter().enumerate() {
3298                    neg_opt[ell] -= nj * sig;
3299                }
3300            }
3301            let t_un: Vec<f64> = neg_opt.iter().map(|&d| d.exp()).collect();
3302            nereids_physics::resolution::apply_r(&matrix, &t_un)
3303        };
3304        let t_cubature_direct = cubature.forward(&n);
3305
3306        // Sanity: cubature direct output matches what evaluate() returned.
3307        for (a, b) in t_cubature.iter().zip(t_cubature_direct.iter()) {
3308            assert!((a - b).abs() < 1e-14);
3309        }
3310        // Cubature vs exact at training density: LP-pinned equivalence.
3311        let max_err = t_cubature
3312            .iter()
3313            .zip(t_exact_via_r.iter())
3314            .map(|(a, b)| {
3315                let denom = a.abs().max(b.abs()).max(1e-12);
3316                (a - b).abs() / denom
3317            })
3318            .fold(0.0_f64, f64::max);
3319        assert!(
3320            max_err < 1e-9,
3321            "at training density, cubature vs exact max err = {max_err:.3e}",
3322        );
3323    }
3324
3325    #[test]
3326    fn precomputed_cubature_falls_back_at_k1() {
3327        // k = 1 with a k=2 cubature → cubature_eligible returns false
3328        // (plan.k mismatch with n_density_params), dispatch MUST
3329        // fall back to the exact `exp(-n σ) + apply_resolution`
3330        // path.  We prove fallback via byte-identity: constructing a
3331        // second model WITHOUT the cubature plan must produce
3332        // exactly the same output as the first model WITH the
3333        // ineligible plan.  A false-positive dispatch would violate
3334        // this invariant because the k=2 cubature's atoms live in
3335        // ℝ² and `cubature.forward([n])` would panic on the
3336        // input-length check in `SparseEmpiricalCubaturePlan::forward`
3337        // — OR, worse, if the guard check accidentally accepted a
3338        // k=2 plan for a k=1 model the output would numerically
3339        // differ from straight Beer-Lambert by more than
3340        // floating-point noise.
3341        let n_grid = 40_usize;
3342        // `plan` intentionally unused here: this test wants both
3343        // model variants in the no-dispatch state (no cubature can
3344        // fire because k=1 vs cubature.k=2), so installing a
3345        // resolution plan would add work without changing the
3346        // tested invariant.
3347        let (energies, _plan, matrix) = synthetic_resolution_setup(n_grid, 4);
3348        let sigmas_k2 = synthetic_sigmas(n_grid, 2);
3349        let cubature_k2 = build_cubature(&matrix, &sigmas_k2, vec![1e-4_f64, 1e-4]);
3350
3351        // Model has k = 1 (only one isotope in cross_sections), but a
3352        // k = 2 cubature is installed → must fall back.
3353        let sigmas_k1 = synthetic_sigmas(n_grid, 1);
3354        let model_with_plan = PrecomputedTransmissionModel {
3355            cross_sections: Arc::new(sigmas_k1.clone()),
3356            density_indices: Arc::new(vec![0]),
3357            energies: Some(Arc::new(energies.clone())),
3358            instrument: Some(make_trivial_instrument()),
3359            resolution_plan: None,
3360            sparse_cubature_plan: Some(Arc::clone(&cubature_k2)),
3361            sparse_scalar_plan: None,
3362            work_layout: None,
3363        };
3364        let model_without_plan = PrecomputedTransmissionModel {
3365            cross_sections: Arc::new(sigmas_k1.clone()),
3366            density_indices: Arc::new(vec![0]),
3367            energies: Some(Arc::new(energies.clone())),
3368            instrument: Some(make_trivial_instrument()),
3369            resolution_plan: None,
3370            sparse_cubature_plan: None,
3371            sparse_scalar_plan: None,
3372            work_layout: None,
3373        };
3374
3375        let n = [1e-4_f64];
3376        let t_with = model_with_plan.evaluate(&n).unwrap();
3377        let t_without = model_without_plan.evaluate(&n).unwrap();
3378        assert_eq!(t_with.len(), n_grid);
3379        assert_eq!(t_without.len(), n_grid);
3380        // Byte identity: ineligible-plan dispatch MUST equal
3381        // no-plan dispatch exactly.
3382        for (a, b) in t_with.iter().zip(t_without.iter()) {
3383            assert_eq!(
3384                a.to_bits(),
3385                b.to_bits(),
3386                "fallback path must be byte-identical to the no-plan path; \
3387                 otherwise the k=2 cubature is silently firing on a k=1 model",
3388            );
3389        }
3390    }
3391
3392    #[test]
3393    fn precomputed_cubature_no_plan_means_exact_path() {
3394        // No cubature installed → byte-identical to the
3395        // pre-cubature-dispatch path.  This is the regression guard:
3396        // the dispatch addition
3397        // must not change the default forward path.
3398        let n_grid = 40_usize;
3399        let (_energies, _plan, _matrix) = synthetic_resolution_setup(n_grid, 4);
3400        let sigmas = synthetic_sigmas(n_grid, 2);
3401
3402        let model = PrecomputedTransmissionModel {
3403            cross_sections: Arc::new(sigmas.clone()),
3404            density_indices: Arc::new(vec![0, 1]),
3405            energies: None,
3406            instrument: None,
3407            resolution_plan: None,
3408            sparse_cubature_plan: None,
3409            sparse_scalar_plan: None,
3410            work_layout: None,
3411        };
3412
3413        let n = [1e-4_f64, 1e-4];
3414        let t = model.evaluate(&n).unwrap();
3415        // Exact Beer-Lambert: T = exp(-Σ n σ).
3416        for (ell, &t_val) in t.iter().enumerate() {
3417            let tau: f64 = sigmas
3418                .iter()
3419                .zip(n.iter())
3420                .map(|(s, &ni)| ni * s[ell])
3421                .sum();
3422            let expected = (-tau).exp();
3423            assert!(
3424                (t_val - expected).abs() < 1e-14,
3425                "at ell={ell}: got {t_val}, expected {expected}",
3426            );
3427        }
3428    }
3429
3430    #[test]
3431    fn precomputed_cubature_jacobian_matches_forward_derivative() {
3432        // Cubature Jacobian columns should equal the per-isotope
3433        // derivatives of the cubature forward output at the anchor.
3434        let n_grid = 40_usize;
3435        let (energies, plan, matrix) = synthetic_resolution_setup(n_grid, 4);
3436        let sigmas = synthetic_sigmas(n_grid, 2);
3437        let train_max = vec![1e-4_f64, 1e-4];
3438        let cubature = build_cubature(&matrix, &sigmas, train_max.clone());
3439
3440        let model = PrecomputedTransmissionModel {
3441            cross_sections: Arc::new(sigmas),
3442            density_indices: Arc::new(vec![0, 1]),
3443            energies: Some(Arc::new(energies)),
3444            instrument: Some(make_trivial_instrument()),
3445            resolution_plan: Some(Arc::clone(&plan)),
3446            sparse_cubature_plan: Some(Arc::clone(&cubature)),
3447            sparse_scalar_plan: None,
3448            work_layout: None,
3449        };
3450
3451        // Use anchor density: LP pins Jacobian exactly here.
3452        let anchor = SparseEmpiricalCubaturePlan::default_jacobian_anchor(&train_max);
3453        let y_curr = model.evaluate(&anchor).unwrap();
3454        let jac = model
3455            .analytical_jacobian(&anchor, &[0, 1], &y_curr)
3456            .expect("cubature Jacobian path");
3457
3458        // Cross-check: cubature.forward_and_jacobian should give the
3459        // same J.
3460        let (_t_ref, jac_flat_ref) = cubature.forward_and_jacobian(&anchor);
3461        for i in 0..n_grid {
3462            for col in 0..2 {
3463                let from_model = jac.get(i, col);
3464                let from_cubature = jac_flat_ref[i * 2 + col];
3465                assert!(
3466                    (from_model - from_cubature).abs() < 1e-14,
3467                    "row {i} col {col}: model = {from_model}, cubature = {from_cubature}",
3468                );
3469            }
3470        }
3471    }
3472
3473    // ── TransmissionFitModel cubature dispatch tests ──────────────────
3474    //
3475    // The per-pixel `TransmissionFitModel` fires the cubature path
3476    // with extra guards (`temperature_index.is_none()` for σ stack
3477    // stability).  These tests exercise BOTH `evaluate()` and
3478    // `analytical_jacobian()` directly on `TransmissionFitModel`,
3479    // not the precomputed variant.
3480
3481    /// Build a minimal `TransmissionFitModel` with a single trivial
3482    /// resonance per isotope + the synthetic σ used for the
3483    /// Precomputed tests, so the cubature dispatch condition can
3484    /// trigger without loading full ENDF data.
3485    fn make_trivial_fit_model(energies: Vec<f64>, k: usize) -> TransmissionFitModel {
3486        // Build k synthetic Isotope / ResonanceData pairs — the fit
3487        // model doesn't actually consult them when the cubature
3488        // dispatch fires (cubature.forward replaces `exp(-Σ n σ) +
3489        // apply_resolution`).  But the constructor still validates
3490        // the count.
3491        // Minimal ResonanceData — the cubature dispatch fires
3492        // before any ENDF-derived code runs, so `ranges` can be
3493        // empty.  When the dispatch falls through (tests that check
3494        // the exact path), we don't exercise cross_sections from
3495        // these resonance_data either; the model uses
3496        // `precomputed_cross_sections` / `base_xs`.
3497        let resonance_data: Vec<ResonanceData> = (0..k)
3498            .map(|j| {
3499                let iso = Isotope::new(40 + j as u32, 96 + j as u32).unwrap();
3500                ResonanceData {
3501                    isotope: iso,
3502                    za: ((40 + j) * 1000 + (96 + j)) as u32,
3503                    awr: 96.0 + j as f64,
3504                    ranges: vec![],
3505                }
3506            })
3507            .collect();
3508
3509        TransmissionFitModel::new(
3510            energies,
3511            resonance_data,
3512            293.6,
3513            Some(make_trivial_instrument()),
3514            ((0..k).collect(), vec![1.0; k]),
3515            None,
3516            None,
3517        )
3518        .expect("TransmissionFitModel::new")
3519    }
3520
3521    #[test]
3522    fn fit_model_cubature_dispatches_at_anchor() {
3523        // Build a k = 2 cubature and a TransmissionFitModel whose
3524        // density_indices / ratios map directly (identity) onto it.
3525        // `evaluate()` at the anchor density MUST equal
3526        // `cubature.forward(anchor)` exactly — the LP equality
3527        // constraint pins it.
3528        let n_grid = 40_usize;
3529        let (energies, plan, matrix) = synthetic_resolution_setup(n_grid, 4);
3530        let sigmas = synthetic_sigmas(n_grid, 2);
3531        let train_max = vec![1e-4_f64, 1e-4];
3532        let cubature = build_cubature(&matrix, &sigmas, train_max.clone());
3533
3534        // Install BOTH the resolution plan and the cubature plan:
3535        // the eligibility guard requires `resolution_plan.is_some()`
3536        // so the cubature doesn't silently bypass an unknown
3537        // resolution operator.
3538        let model = make_trivial_fit_model(energies.clone(), 2)
3539            .with_resolution_plan(Some(Arc::clone(&plan)))
3540            .with_sparse_cubature_plan(Some(cubature.clone()));
3541
3542        // Evaluate at a training density (LP pins exactly) → model
3543        // output equals cubature output.
3544        let n = [0.25 * train_max[0], 0.25 * train_max[1]];
3545        let t_model = model.evaluate(&n).unwrap();
3546        let t_cub = cubature.forward(&n);
3547        assert_eq!(t_model.len(), n_grid);
3548        for (a, b) in t_model.iter().zip(t_cub.iter()) {
3549            assert_eq!(
3550                a.to_bits(),
3551                b.to_bits(),
3552                "TransmissionFitModel cubature dispatch must return cubature.forward() byte-exact at the LP-pinned anchor",
3553            );
3554        }
3555    }
3556
3557    #[test]
3558    fn fit_model_cubature_jacobian_matches_cubature_output() {
3559        // Same pattern as the Precomputed Jacobian test but on
3560        // TransmissionFitModel.  analytical_jacobian at the anchor
3561        // density must return exactly `cubature.forward_and_jacobian(n)`'s
3562        // J matrix.
3563        let n_grid = 40_usize;
3564        let (energies, plan, matrix) = synthetic_resolution_setup(n_grid, 4);
3565        let sigmas = synthetic_sigmas(n_grid, 2);
3566        let train_max = vec![1e-4_f64, 1e-4];
3567        let cubature = build_cubature(&matrix, &sigmas, train_max.clone());
3568
3569        let model = make_trivial_fit_model(energies, 2)
3570            .with_resolution_plan(Some(Arc::clone(&plan)))
3571            .with_sparse_cubature_plan(Some(cubature.clone()));
3572
3573        let anchor = SparseEmpiricalCubaturePlan::default_jacobian_anchor(&train_max);
3574        let y_curr = model.evaluate(&anchor).unwrap();
3575        let jac = model
3576            .analytical_jacobian(&anchor, &[0, 1], &y_curr)
3577            .expect("cubature Jacobian path on TransmissionFitModel");
3578        let (_t_ref, jac_flat_ref) = cubature.forward_and_jacobian(&anchor);
3579        for i in 0..n_grid {
3580            for col in 0..2 {
3581                let from_model = jac.get(i, col);
3582                let from_cubature = jac_flat_ref[i * 2 + col];
3583                assert_eq!(
3584                    from_model.to_bits(),
3585                    from_cubature.to_bits(),
3586                    "row {i} col {col}: TransmissionFitModel must return cubature J byte-exact",
3587                );
3588            }
3589        }
3590    }
3591
3592    #[test]
3593    fn fit_model_cubature_falls_back_on_grid_mismatch() {
3594        // Build a cubature on one grid, install it on a model with a
3595        // DIFFERENT same-length grid.  Dispatch must refuse the plan
3596        // via the new `to_bits()` grid-identity check and produce
3597        // byte-identical output to the no-plan model (exact path).
3598        let n_grid = 40_usize;
3599        let (energies_a, _plan, matrix) = synthetic_resolution_setup(n_grid, 4);
3600        let sigmas = synthetic_sigmas(n_grid, 2);
3601        let train_max = vec![1e-4_f64, 1e-4];
3602        let cubature = build_cubature(&matrix, &sigmas, train_max);
3603
3604        // A different same-length grid (shifted by 1 eV).
3605        let energies_b: Vec<f64> = energies_a.iter().map(|&e| e + 1.0).collect();
3606
3607        let model_with_stale_plan =
3608            make_trivial_fit_model(energies_b.clone(), 2).with_sparse_cubature_plan(Some(cubature));
3609        let model_without_plan = make_trivial_fit_model(energies_b, 2);
3610
3611        let n = [1e-5_f64, 1e-5];
3612        let t_stale = model_with_stale_plan.evaluate(&n).unwrap();
3613        let t_exact = model_without_plan.evaluate(&n).unwrap();
3614        for (a, b) in t_stale.iter().zip(t_exact.iter()) {
3615            assert_eq!(
3616                a.to_bits(),
3617                b.to_bits(),
3618                "stale-grid cubature plan MUST NOT fire; evaluate() must match no-plan byte-exactly",
3619            );
3620        }
3621    }
3622
3623    #[test]
3624    fn fit_model_cubature_falls_back_when_density_escapes_box() {
3625        // Build cubature with train_max = [1e-4, 1e-4], install
3626        // the density_box, then call evaluate() with a density
3627        // WELL beyond the 1.5× tolerance.  Dispatch must fall back
3628        // to the exact path rather than silently extrapolate the
3629        // surrogate outside its trained region.
3630        let n_grid = 40_usize;
3631        let (energies, plan, matrix) = synthetic_resolution_setup(n_grid, 4);
3632        let sigmas = synthetic_sigmas(n_grid, 2);
3633        let train_max = vec![1e-4_f64, 1e-4];
3634
3635        // Build cubature AND attach the density_box.
3636        let cubature = {
3637            let flat: Vec<f64> = sigmas.iter().flat_map(|s| s.iter().copied()).collect();
3638            let training = SparseEmpiricalCubaturePlan::default_training_points(&train_max);
3639            let anchor = SparseEmpiricalCubaturePlan::default_jacobian_anchor(&train_max);
3640            Arc::new(
3641                SparseEmpiricalCubaturePlan::build(&matrix, &flat, 2, &training, &anchor)
3642                    .expect("build")
3643                    .with_density_box(train_max.clone()),
3644            )
3645        };
3646
3647        let model_with = make_trivial_fit_model(energies.clone(), 2)
3648            .with_resolution_plan(Some(Arc::clone(&plan)))
3649            .with_sparse_cubature_plan(Some(Arc::clone(&cubature)));
3650        let model_without =
3651            make_trivial_fit_model(energies, 2).with_resolution_plan(Some(Arc::clone(&plan)));
3652
3653        // Escape: 5× the training max → well outside the 1.5× tolerance.
3654        let n_escape = [5.0 * train_max[0], 5.0 * train_max[1]];
3655        let t_with = model_with.evaluate(&n_escape).unwrap();
3656        let t_without = model_without.evaluate(&n_escape).unwrap();
3657        // If the guard fired correctly, the cubature-installed
3658        // model falls back to the exact path and produces the same
3659        // output as the no-plan model — byte-identical.
3660        for (a, b) in t_with.iter().zip(t_without.iter()) {
3661            assert_eq!(
3662                a.to_bits(),
3663                b.to_bits(),
3664                "density-box escape guard MUST fall back to exact path byte-identically",
3665            );
3666        }
3667    }
3668
3669    #[test]
3670    fn fit_model_cubature_dispatches_without_resolution_plan_attached() {
3671        // Single-spectrum regression: callers of the non-spatial
3672        // `fit_spectrum_typed` / `build_transmission_model` path
3673        // attach a cubature via
3674        // `UnifiedFitConfig::with_precomputed_sparse_cubature_plan`
3675        // but typically don't also pre-build a `ResolutionPlan` (the
3676        // per-call `apply_resolution` broaden path is used
3677        // otherwise).  The cubature fast path MUST still fire — a
3678        // prior `resolution_plan.is_some()` requirement
3679        // made the new API inert on this surface.
3680        let n_grid = 40_usize;
3681        let (energies, _plan, matrix) = synthetic_resolution_setup(n_grid, 4);
3682        let sigmas = synthetic_sigmas(n_grid, 2);
3683        let train_max = vec![1e-4_f64, 1e-4];
3684        let cubature = build_cubature(&matrix, &sigmas, train_max.clone());
3685
3686        // Intentionally NOT installing a resolution plan.  The
3687        // instrument's tabulated resolution is enough.
3688        let model = make_trivial_fit_model(energies.clone(), 2)
3689            .with_sparse_cubature_plan(Some(Arc::clone(&cubature)));
3690
3691        let n = [0.25 * train_max[0], 0.25 * train_max[1]];
3692        let t_model = model.evaluate(&n).unwrap();
3693        let t_cub = cubature.forward(&n);
3694        for (a, b) in t_model.iter().zip(t_cub.iter()) {
3695            assert_eq!(
3696                a.to_bits(),
3697                b.to_bits(),
3698                "cubature dispatch must fire on single-spectrum path without a separate ResolutionPlan attached",
3699            );
3700        }
3701    }
3702
3703    // ── Scalar (k = 1) dispatch-guard tests ───────────────────────────
3704    //
3705    // The cubature tests above cover
3706    // the k ≥ 2 path; the scalar path is a separate surrogate with its
3707    // own eligibility guard (`scalar_eligible`) and its own
3708    // density-box guard (`scalar_density_within_box`).  These tests
3709    // exercise the scalar-specific guards: k=1-only, grid-identity
3710    // via `to_bits()`, tabulated-only instrument resolution,
3711    // density-box escape, and that the pure no-plan path remains
3712    // byte-identical to the pre-surrogate path.
3713
3714    /// Helper: build a synthetic scalar (k = 1) Chebyshev plan on
3715    /// the same grid as the cubature helpers.  Takes an
3716    /// `Arc<ResolutionPlan>` so tests can share the same Arc
3717    /// pointer with the model's `resolution_plan` (required by the
3718    /// `Arc::ptr_eq` dispatch guard).
3719    fn build_scalar_plan(
3720        res_plan: Arc<ResolutionPlan>,
3721        sigma_k1: &[f64],
3722        n_max: f64,
3723    ) -> Arc<ScalarSurrogatePlan> {
3724        Arc::new(
3725            nereids_physics::surrogate::ScalarChebyshevPlan::build(res_plan, sigma_k1, n_max, 16)
3726                .expect("synthetic scalar Chebyshev build"),
3727        )
3728    }
3729
3730    /// Helper: build a `PrecomputedTransmissionModel` with the
3731    /// caller-chosen σ / k / resolution-plan / scalar-plan state.
3732    /// Mirrors `make_trivial_fit_model` but targets the model that
3733    /// actually dispatches scalar in production (spatial routes
3734    /// scalar-eligible k=1 through `PrecomputedTransmissionModel`).
3735    fn make_precomp_for_scalar(
3736        energies: Vec<f64>,
3737        sigmas: Vec<Vec<f64>>,
3738        density_indices: Vec<usize>,
3739        resolution_plan: Option<Arc<ResolutionPlan>>,
3740        scalar_plan: Option<Arc<ScalarSurrogatePlan>>,
3741    ) -> PrecomputedTransmissionModel {
3742        PrecomputedTransmissionModel {
3743            cross_sections: Arc::new(sigmas),
3744            density_indices: Arc::new(density_indices),
3745            energies: Some(Arc::new(energies)),
3746            instrument: Some(make_trivial_instrument()),
3747            resolution_plan,
3748            sparse_cubature_plan: None,
3749            sparse_scalar_plan: scalar_plan,
3750            work_layout: None,
3751        }
3752    }
3753
3754    #[test]
3755    fn precomputed_scalar_dispatches_at_k1() {
3756        // k = 1 with both the scalar plan and the resolution plan
3757        // installed (same Arc) and σ matching the plan's
3758        // fingerprint: evaluate() must return the scalar plan's
3759        // forward output byte-exact.
3760        let n_grid = 40_usize;
3761        let (energies, res_plan, _matrix) = synthetic_resolution_setup(n_grid, 4);
3762        let sigmas = synthetic_sigmas(n_grid, 1);
3763        let n_max = 2.0 * 1e-4_f64;
3764        let scalar = build_scalar_plan(Arc::clone(&res_plan), &sigmas[0], n_max);
3765
3766        let model = make_precomp_for_scalar(
3767            energies,
3768            sigmas,
3769            vec![0],
3770            Some(Arc::clone(&res_plan)),
3771            Some(Arc::clone(&scalar)),
3772        );
3773
3774        let n = [0.5 * n_max];
3775        let t_model = model.evaluate(&n).unwrap();
3776        let t_scalar = scalar.forward_scalar(n[0]);
3777        assert_eq!(t_model.len(), n_grid);
3778        for (a, b) in t_model.iter().zip(t_scalar.iter()) {
3779            assert_eq!(
3780                a.to_bits(),
3781                b.to_bits(),
3782                "scalar dispatch must return forward_scalar() byte-exact",
3783            );
3784        }
3785    }
3786
3787    #[test]
3788    fn precomputed_scalar_jacobian_matches_derivative() {
3789        let n_grid = 40_usize;
3790        let (energies, res_plan, _matrix) = synthetic_resolution_setup(n_grid, 4);
3791        let sigmas = synthetic_sigmas(n_grid, 1);
3792        let n_max = 2.0 * 1e-4_f64;
3793        let scalar = build_scalar_plan(Arc::clone(&res_plan), &sigmas[0], n_max);
3794
3795        let model = make_precomp_for_scalar(
3796            energies,
3797            sigmas,
3798            vec![0],
3799            Some(Arc::clone(&res_plan)),
3800            Some(Arc::clone(&scalar)),
3801        );
3802
3803        let n = [0.5 * n_max];
3804        let y_curr = model.evaluate(&n).unwrap();
3805        let jac = model
3806            .analytical_jacobian(&n, &[0], &y_curr)
3807            .expect("scalar Jacobian path");
3808        let (_t_ref, dt_ref) = scalar.forward_and_derivative_scalar(n[0]);
3809        assert_eq!(jac.ncols, 1);
3810        assert_eq!(jac.nrows, n_grid);
3811        for (i, &dt_i) in dt_ref.iter().enumerate().take(n_grid) {
3812            assert_eq!(
3813                jac.get(i, 0).to_bits(),
3814                dt_i.to_bits(),
3815                "row {i}: scalar dT/dn must be byte-exact",
3816            );
3817        }
3818    }
3819
3820    #[test]
3821    fn precomputed_scalar_falls_back_at_k2() {
3822        // k = 2 with a scalar plan installed → `scalar_eligible`
3823        // rejects `cross_sections.len() == 1` guard (k=2 model has
3824        // 2 σ rows).  Dispatch falls back.
3825        let n_grid = 40_usize;
3826        let (energies, res_plan, _matrix) = synthetic_resolution_setup(n_grid, 4);
3827        let sigmas_k2 = synthetic_sigmas(n_grid, 2);
3828        let sigma_k1 = synthetic_sigmas(n_grid, 1).remove(0);
3829        let n_max = 2.0 * 1e-4_f64;
3830        let scalar = build_scalar_plan(Arc::clone(&res_plan), &sigma_k1, n_max);
3831
3832        let model_with = make_precomp_for_scalar(
3833            energies.clone(),
3834            sigmas_k2.clone(),
3835            vec![0, 1],
3836            Some(Arc::clone(&res_plan)),
3837            Some(scalar),
3838        );
3839        let model_without = make_precomp_for_scalar(
3840            energies,
3841            sigmas_k2,
3842            vec![0, 1],
3843            Some(Arc::clone(&res_plan)),
3844            None,
3845        );
3846        let n = [1e-4_f64, 2e-4];
3847        let t_with = model_with.evaluate(&n).unwrap();
3848        let t_without = model_without.evaluate(&n).unwrap();
3849        for (a, b) in t_with.iter().zip(t_without.iter()) {
3850            assert_eq!(
3851                a.to_bits(),
3852                b.to_bits(),
3853                "scalar plan must refuse k=2 dispatch → byte-identical fallback",
3854            );
3855        }
3856    }
3857
3858    #[test]
3859    fn precomputed_scalar_falls_back_on_stale_resolution_plan() {
3860        // Same-grid
3861        // DIFFERENT-kernel ResolutionPlan swap must not silently
3862        // dispatch.  The `Arc::ptr_eq` guard on the scalar plan's
3863        // stored source plan is the O(1) check that closes this.
3864        let n_grid = 40_usize;
3865        let (energies, res_plan_a, _matrix) = synthetic_resolution_setup(n_grid, 4);
3866        let sigmas = synthetic_sigmas(n_grid, 1);
3867        let n_max = 2.0 * 1e-4_f64;
3868        let scalar = build_scalar_plan(Arc::clone(&res_plan_a), &sigmas[0], n_max);
3869
3870        // Build a DIFFERENT ResolutionPlan on the same grid (wider
3871        // kernel) and attach it to the model.  Even though the
3872        // grid matches bit-for-bit, the scalar plan was built from
3873        // res_plan_a and its `source_resolution_plan` Arc differs
3874        // from res_plan_b → dispatch refuses.
3875        let (_e_b, res_plan_b, _matrix_b) = synthetic_resolution_setup(n_grid, 6);
3876        let model_stale = make_precomp_for_scalar(
3877            energies.clone(),
3878            sigmas.clone(),
3879            vec![0],
3880            Some(Arc::clone(&res_plan_b)),
3881            Some(Arc::clone(&scalar)),
3882        );
3883        let model_noplan =
3884            make_precomp_for_scalar(energies, sigmas, vec![0], Some(res_plan_b), None);
3885        let n = [0.25 * n_max];
3886        let t_stale = model_stale.evaluate(&n).unwrap();
3887        let t_exact = model_noplan.evaluate(&n).unwrap();
3888        for (a, b) in t_stale.iter().zip(t_exact.iter()) {
3889            assert_eq!(
3890                a.to_bits(),
3891                b.to_bits(),
3892                "scalar plan with non-ptr_eq source_resolution_plan MUST NOT fire",
3893            );
3894        }
3895    }
3896
3897    #[test]
3898    fn precomputed_scalar_falls_back_on_stale_sigma() {
3899        // Plan built
3900        // from σ_A, attached to a model whose cross_sections[0] is
3901        // σ_B on the same grid with the same resolution plan →
3902        // σ-fingerprint mismatch forces fallback.
3903        let n_grid = 40_usize;
3904        let (energies, res_plan, _matrix) = synthetic_resolution_setup(n_grid, 4);
3905        let sigma_a = synthetic_sigmas(n_grid, 1);
3906        // σ_B: flip one element of σ_A so the fingerprint differs
3907        // but the shape / magnitude is plausible.
3908        let mut sigma_b = sigma_a.clone();
3909        sigma_b[0][n_grid / 2] += 1.0; // tiny perturbation → different fingerprint
3910        let n_max = 2.0 * 1e-4_f64;
3911        let scalar = build_scalar_plan(Arc::clone(&res_plan), &sigma_a[0], n_max);
3912
3913        let model_stale = make_precomp_for_scalar(
3914            energies.clone(),
3915            sigma_b.clone(),
3916            vec![0],
3917            Some(Arc::clone(&res_plan)),
3918            Some(scalar),
3919        );
3920        let model_noplan =
3921            make_precomp_for_scalar(energies, sigma_b, vec![0], Some(res_plan), None);
3922        let n = [0.25 * n_max];
3923        let t_stale = model_stale.evaluate(&n).unwrap();
3924        let t_exact = model_noplan.evaluate(&n).unwrap();
3925        for (a, b) in t_stale.iter().zip(t_exact.iter()) {
3926            assert_eq!(
3927                a.to_bits(),
3928                b.to_bits(),
3929                "σ-fingerprint mismatch MUST force fallback → byte-identical to no-plan",
3930            );
3931        }
3932    }
3933
3934    #[test]
3935    fn precomputed_scalar_falls_back_when_density_escapes_box() {
3936        let n_grid = 40_usize;
3937        let (energies, res_plan, _matrix) = synthetic_resolution_setup(n_grid, 4);
3938        let sigmas = synthetic_sigmas(n_grid, 1);
3939        let n_max = 2.0 * 1e-4_f64;
3940        let scalar = build_scalar_plan(Arc::clone(&res_plan), &sigmas[0], n_max);
3941
3942        let model_with = make_precomp_for_scalar(
3943            energies.clone(),
3944            sigmas.clone(),
3945            vec![0],
3946            Some(Arc::clone(&res_plan)),
3947            Some(Arc::clone(&scalar)),
3948        );
3949        let model_without =
3950            make_precomp_for_scalar(energies, sigmas, vec![0], Some(Arc::clone(&res_plan)), None);
3951        let n_escape = [2.0 * n_max];
3952        let t_with = model_with.evaluate(&n_escape).unwrap();
3953        let t_without = model_without.evaluate(&n_escape).unwrap();
3954        for (a, b) in t_with.iter().zip(t_without.iter()) {
3955            assert_eq!(
3956                a.to_bits(),
3957                b.to_bits(),
3958                "density-box escape guard must fall back byte-identically",
3959            );
3960        }
3961    }
3962
3963    #[test]
3964    fn precomputed_scalar_rejects_nonfinite_density() {
3965        let n_grid = 40_usize;
3966        let (energies, res_plan, _matrix) = synthetic_resolution_setup(n_grid, 4);
3967        let sigmas = synthetic_sigmas(n_grid, 1);
3968        let n_max = 2.0 * 1e-4_f64;
3969        let scalar = build_scalar_plan(Arc::clone(&res_plan), &sigmas[0], n_max);
3970
3971        let model_with = make_precomp_for_scalar(
3972            energies.clone(),
3973            sigmas.clone(),
3974            vec![0],
3975            Some(Arc::clone(&res_plan)),
3976            Some(Arc::clone(&scalar)),
3977        );
3978        let model_without =
3979            make_precomp_for_scalar(energies, sigmas, vec![0], Some(Arc::clone(&res_plan)), None);
3980        for bad_n in [f64::NAN, f64::INFINITY, -1e-6_f64] {
3981            let n = [bad_n];
3982            let t_with = model_with.evaluate(&n).unwrap();
3983            let t_without = model_without.evaluate(&n).unwrap();
3984            for (i, (a, b)) in t_with.iter().zip(t_without.iter()).enumerate() {
3985                assert_eq!(
3986                    a.to_bits(),
3987                    b.to_bits(),
3988                    "n = {bad_n}: scalar guard must fall back byte-exactly; row {i}",
3989                );
3990            }
3991        }
3992    }
3993
3994    #[test]
3995    fn scalar_density_within_box_direct_guard() {
3996        // Unit-test the scalar_density_within_box helper directly
3997        // without going through the model dispatch.  Chebyshev is a
3998        // polynomial interpolant that diverges exponentially outside
3999        // `[0, n_max]` — measured: 73 % rel err
4000        // at `1.5 × n_max`.  The guard is therefore **strict**
4001        // `n ≤ train_max`, not the cubature's 1.5× tolerance.
4002        let n_grid = 16_usize;
4003        let (_energies, res_plan, _matrix) = synthetic_resolution_setup(n_grid, 2);
4004        let sigmas = synthetic_sigmas(n_grid, 1);
4005        let n_max = 1e-4_f64;
4006        let plan =
4007            nereids_physics::surrogate::ScalarChebyshevPlan::build(res_plan, &sigmas[0], n_max, 16)
4008                .expect("build");
4009
4010        // Inside the box: accepted.
4011        assert!(scalar_density_within_box(&plan, 0.0));
4012        assert!(scalar_density_within_box(&plan, 0.5 * n_max));
4013        assert!(scalar_density_within_box(&plan, n_max));
4014        // Any positive excursion past the box is rejected (no
4015        // 1.5× tolerance).
4016        assert!(!scalar_density_within_box(
4017            &plan,
4018            n_max * (1.0 + f64::EPSILON)
4019        ));
4020        assert!(!scalar_density_within_box(&plan, 1.01 * n_max));
4021        assert!(!scalar_density_within_box(&plan, 1.5 * n_max));
4022        assert!(!scalar_density_within_box(&plan, 2.0 * n_max));
4023        // Non-finite and negative must be rejected.
4024        assert!(!scalar_density_within_box(&plan, f64::NAN));
4025        assert!(!scalar_density_within_box(&plan, f64::INFINITY));
4026        assert!(!scalar_density_within_box(&plan, f64::NEG_INFINITY));
4027        assert!(!scalar_density_within_box(&plan, -1e-9));
4028    }
4029
4030    #[test]
4031    fn density_param_indices_sorted_by_value() {
4032        // First-appearance order would swap columns for non-
4033        // monotonic group layouts like [1, 0, 1].  Sorted-by-value
4034        // keeps dispatch aligned with the cubature's σ-stack
4035        // indexing (`sigmas[j * n_rows + ℓ]` = σ for density param
4036        // j).
4037        assert_eq!(density_param_indices(&[0, 0, 0]), vec![0]);
4038        assert_eq!(density_param_indices(&[0, 1, 2, 3]), vec![0, 1, 2, 3]);
4039        assert_eq!(density_param_indices(&[1, 0, 1]), vec![0, 1]);
4040        assert_eq!(density_param_indices(&[3, 1, 2, 0, 2]), vec![0, 1, 2, 3]);
4041    }
4042
4043    /// Verify that NormalizedTransmissionModel with identity normalization
4044    /// (Anorm=1, all background=0) gives the same result as the inner model.
4045    #[test]
4046    fn normalized_identity_matches_inner() {
4047        let xs = vec![
4048            vec![1.0, 2.0, 3.0], // isotope 0
4049            vec![0.5, 0.5, 0.5], // isotope 1
4050        ];
4051        let inner_ref = make_precomputed(xs.clone(), vec![0, 1]);
4052        let inner_wrap = make_precomputed(xs, vec![0, 1]);
4053
4054        let energies = [4.0, 9.0, 16.0];
4055        // params: [density0, density1, Anorm, BackA, BackB, BackC]
4056        let model = NormalizedTransmissionModel::new(inner_wrap, &energies, 2, 3, 4, 5);
4057
4058        let params = [0.2, 0.4, 1.0, 0.0, 0.0, 0.0];
4059        let y_norm = model.evaluate(&params).unwrap();
4060        let y_inner = inner_ref.evaluate(&params).unwrap();
4061
4062        for (a, b) in y_norm.iter().zip(y_inner.iter()) {
4063            assert!(
4064                (a - b).abs() < 1e-12,
4065                "identity normalization should match inner: {} vs {}",
4066                a,
4067                b
4068            );
4069        }
4070    }
4071
4072    /// Verify the normalization formula:
4073    /// T_out = Anorm * T_inner + BackA + BackB/sqrt(E) + BackC*sqrt(E)
4074    #[test]
4075    fn normalized_formula_correct() {
4076        let xs = vec![vec![1.0, 2.0, 3.0]];
4077        let inner_ref = make_precomputed(xs.clone(), vec![0]);
4078        let inner_wrap = make_precomputed(xs, vec![0]);
4079
4080        let energies = [4.0, 9.0, 16.0]; // sqrt = [2, 3, 4]
4081        let model = NormalizedTransmissionModel::new(inner_wrap, &energies, 1, 2, 3, 4);
4082
4083        // params: [density, Anorm, BackA, BackB, BackC]
4084        let anorm = 0.95;
4085        let back_a = 0.01;
4086        let back_b = 0.02;
4087        let back_c = 0.005;
4088        let density = 0.3;
4089        let params = [density, anorm, back_a, back_b, back_c];
4090
4091        let y = model.evaluate(&params).unwrap();
4092        let t_inner = inner_ref.evaluate(&params).unwrap();
4093
4094        for (i, (&yi, &ti)) in y.iter().zip(t_inner.iter()).enumerate() {
4095            let sqrt_e = energies[i].sqrt();
4096            let expected = anorm * ti + back_a + back_b / sqrt_e + back_c * sqrt_e;
4097            assert!(
4098                (yi - expected).abs() < 1e-12,
4099                "E[{i}]: got {yi}, expected {expected}"
4100            );
4101        }
4102    }
4103
4104    /// Analytical Jacobian of NormalizedTransmissionModel must match
4105    /// central-difference finite-difference.
4106    #[test]
4107    fn normalized_analytical_jacobian_matches_fd() {
4108        let xs = vec![
4109            vec![1.0, 2.0, 3.0], // isotope 0
4110            vec![0.5, 0.5, 0.5], // isotope 1
4111        ];
4112        let inner = make_precomputed(xs, vec![0, 1]);
4113
4114        let energies = [4.0, 9.0, 16.0];
4115        // params: [density0, density1, Anorm, BackA, BackB, BackC]
4116        let model = NormalizedTransmissionModel::new(inner, &energies, 2, 3, 4, 5);
4117
4118        let params = [0.2, 0.4, 0.95, 0.01, 0.02, 0.005];
4119        let y = model.evaluate(&params).unwrap();
4120        let free: Vec<usize> = (0..6).collect();
4121
4122        let jac = model
4123            .analytical_jacobian(&params, &free, &y)
4124            .expect("analytical_jacobian should return Some");
4125
4126        assert_eq!(jac.nrows, 3);
4127        assert_eq!(jac.ncols, 6);
4128
4129        // Central-difference reference
4130        let h = 1e-7;
4131        for (col, &p_idx) in free.iter().enumerate() {
4132            let mut p_plus = params;
4133            let mut p_minus = params;
4134            p_plus[p_idx] += h;
4135            p_minus[p_idx] -= h;
4136
4137            let y_plus = model.evaluate(&p_plus).unwrap();
4138            let y_minus = model.evaluate(&p_minus).unwrap();
4139
4140            for i in 0..3 {
4141                let fd = (y_plus[i] - y_minus[i]) / (2.0 * h);
4142                let ana = jac.get(i, col);
4143                let err = (fd - ana).abs();
4144                let scale = fd.abs().max(ana.abs()).max(1e-10);
4145                assert!(
4146                    err / scale < 1e-4,
4147                    "Jacobian mismatch (row {i}, col {col}): FD={fd:.8}, analytical={ana:.8}, \
4148                     rel_err={:.6}",
4149                    err / scale,
4150                );
4151            }
4152        }
4153    }
4154
4155    /// Verify that when some background params are fixed (not in
4156    /// free_param_indices), the Jacobian columns are correct.
4157    #[test]
4158    fn normalized_jacobian_partial_free() {
4159        let xs = vec![vec![1.0, 2.0, 3.0]];
4160        let inner = make_precomputed(xs, vec![0]);
4161
4162        let energies = [4.0, 9.0, 16.0];
4163        let model = NormalizedTransmissionModel::new(inner, &energies, 1, 2, 3, 4);
4164
4165        // params: [density, Anorm, BackA, BackB, BackC]
4166        let params = [0.3, 0.95, 0.01, 0.0, 0.0];
4167        let y = model.evaluate(&params).unwrap();
4168        // Only density and Anorm are free
4169        let free = vec![0usize, 1usize];
4170
4171        let jac = model
4172            .analytical_jacobian(&params, &free, &y)
4173            .expect("should return Some for partial free");
4174
4175        assert_eq!(jac.nrows, 3);
4176        assert_eq!(jac.ncols, 2);
4177
4178        // Central-difference reference
4179        let h = 1e-7;
4180        for (col, &p_idx) in free.iter().enumerate() {
4181            let mut p_plus = params;
4182            let mut p_minus = params;
4183            p_plus[p_idx] += h;
4184            p_minus[p_idx] -= h;
4185
4186            let y_plus = model.evaluate(&p_plus).unwrap();
4187            let y_minus = model.evaluate(&p_minus).unwrap();
4188
4189            for i in 0..3 {
4190                let fd = (y_plus[i] - y_minus[i]) / (2.0 * h);
4191                let ana = jac.get(i, col);
4192                let err = (fd - ana).abs();
4193                let scale = fd.abs().max(ana.abs()).max(1e-10);
4194                assert!(
4195                    err / scale < 1e-4,
4196                    "Jacobian mismatch (row {i}, col {col}): FD={fd:.8}, analytical={ana:.8}"
4197                );
4198            }
4199        }
4200    }
4201
4202    /// End-to-end: fit recovers known Anorm + BackA from synthetic data.
4203    #[test]
4204    fn normalized_fit_recovers_anorm_and_backa() {
4205        let xs = vec![vec![1.0, 2.0, 3.0, 2.0, 1.5]];
4206        let inner = make_precomputed(xs, vec![0]);
4207
4208        let energies = [4.0, 9.0, 16.0, 25.0, 36.0];
4209        let model = NormalizedTransmissionModel::new(inner, &energies, 1, 2, 3, 4);
4210
4211        // True parameters
4212        let true_density = 0.2;
4213        let true_anorm = 0.95;
4214        let true_back_a = 0.02;
4215        let true_params = [true_density, true_anorm, true_back_a, 0.0, 0.0];
4216
4217        let y_obs = model.evaluate(&true_params).unwrap();
4218        let sigma = vec![0.001; y_obs.len()];
4219
4220        // Initial guesses offset from truth
4221        let mut params = ParameterSet::new(vec![
4222            FitParameter::non_negative("density", 0.1),
4223            FitParameter {
4224                name: "anorm".into(),
4225                value: 1.0,
4226                lower: 0.5,
4227                upper: 1.5,
4228                fixed: false,
4229            },
4230            FitParameter::unbounded("back_a", 0.0),
4231            FitParameter::fixed("back_b", 0.0),
4232            FitParameter::fixed("back_c", 0.0),
4233        ]);
4234
4235        let config = LmConfig {
4236            max_iter: 200,
4237            ..LmConfig::default()
4238        };
4239
4240        let result = lm::levenberg_marquardt(&model, &y_obs, &sigma, &mut params, &config).unwrap();
4241
4242        assert!(result.converged, "Fit should converge");
4243
4244        let fit_density = result.params[0];
4245        let fit_anorm = result.params[1];
4246        let fit_back_a = result.params[2];
4247
4248        assert!(
4249            (fit_density - true_density).abs() / true_density < 0.01,
4250            "density: fitted={fit_density}, true={true_density}"
4251        );
4252        assert!(
4253            (fit_anorm - true_anorm).abs() / true_anorm < 0.01,
4254            "anorm: fitted={fit_anorm}, true={true_anorm}"
4255        );
4256        assert!(
4257            (fit_back_a - true_back_a).abs() < 0.001,
4258            "back_a: fitted={fit_back_a}, true={true_back_a}"
4259        );
4260    }
4261
4262    // ── Phase 1: ForwardModel tests ──
4263
4264    #[test]
4265    fn forward_model_predict_equals_fit_model_evaluate_precomputed() {
4266        use crate::forward_model::ForwardModel;
4267        let xs = vec![vec![1.0, 2.0, 3.0, 2.0, 1.5]];
4268        let model = make_precomputed(xs, vec![0]);
4269        let params = [0.001];
4270        let fm_result = model.evaluate(&params).unwrap();
4271        let fwd_result = model.predict(&params).unwrap();
4272        assert_eq!(fm_result, fwd_result);
4273        assert_eq!(model.n_data(), 5);
4274        assert_eq!(model.n_params(), 1);
4275    }
4276
4277    #[test]
4278    fn forward_model_predict_equals_fit_model_evaluate_normalized() {
4279        use crate::forward_model::ForwardModel;
4280        let xs = vec![vec![1.0, 2.0, 3.0, 2.0, 1.5]];
4281        let inner = make_precomputed(xs, vec![0]);
4282        let energies = [4.0, 9.0, 16.0, 25.0, 36.0];
4283        let model = NormalizedTransmissionModel::new(inner, &energies, 1, 2, 3, 4);
4284        let params = [0.001, 0.95, 0.01, 0.0, 0.0];
4285        let fm_result = model.evaluate(&params).unwrap();
4286        let fwd_result = model.predict(&params).unwrap();
4287        assert_eq!(fm_result, fwd_result);
4288        assert_eq!(model.n_data(), 5);
4289        assert_eq!(model.n_params(), 5);
4290    }
4291
4292    #[test]
4293    fn forward_model_jacobian_columns_match_precomputed() {
4294        use crate::forward_model::ForwardModel;
4295        let xs = vec![vec![1.0, 2.0, 3.0], vec![0.5, 1.5, 2.5]];
4296        let model = make_precomputed(xs, vec![0, 1]);
4297        let params = [0.001, 0.002];
4298        let y = model.predict(&params).unwrap();
4299        let free_indices = vec![0, 1];
4300        let jac = model
4301            .jacobian(&params, &free_indices, &y)
4302            .expect("analytical jacobian should be available");
4303        assert_eq!(jac.len(), 2); // 2 columns (one per free param)
4304        assert_eq!(jac[0].len(), 3); // 3 rows (one per energy bin)
4305    }
4306
4307    // ── Issue #442 Step 3 regression tests ─────────────────────────────────
4308
4309    /// Issue #442: PrecomputedTransmissionModel with resolution must match
4310    /// forward_model() with resolution for the same single-isotope sample.
4311    #[test]
4312    fn precomputed_with_resolution_matches_forward_model() {
4313        use nereids_physics::resolution::ResolutionFunction;
4314
4315        let data = u238_single_resonance();
4316        let thickness = 0.0005;
4317        let temperature = 300.0;
4318        let energies: Vec<f64> = (0..401).map(|i| 4.0 + (i as f64) * 0.015).collect();
4319
4320        let inst = Arc::new(InstrumentParams {
4321            resolution: ResolutionFunction::Gaussian(
4322                nereids_physics::resolution::ResolutionParams::new(25.0, 0.5, 0.005, 0.0).unwrap(),
4323            ),
4324        });
4325
4326        // Reference: forward_model() (already fixed in Step 1).
4327        let sample = SampleParams::new(temperature, vec![(data.clone(), thickness)]).unwrap();
4328        let t_forward = transmission::forward_model(&energies, &sample, Some(&inst)).unwrap();
4329
4330        // Precomputed path: Doppler-only XS → PrecomputedTransmissionModel.
4331        let xs = transmission::broadened_cross_sections(
4332            &energies,
4333            std::slice::from_ref(&data),
4334            temperature,
4335            Some(&inst), // aux grid for Doppler accuracy
4336            None,
4337        )
4338        .unwrap();
4339        let model = PrecomputedTransmissionModel {
4340            cross_sections: Arc::new(xs),
4341            density_indices: Arc::new(vec![0]),
4342            energies: Some(Arc::new(energies.clone())),
4343            instrument: Some(Arc::clone(&inst)),
4344            resolution_plan: None,
4345            sparse_cubature_plan: None,
4346            sparse_scalar_plan: None,
4347            work_layout: None,
4348        };
4349        let t_precomputed = model.evaluate(&[thickness]).unwrap();
4350
4351        // Both should agree closely on the interior grid.
4352        // Small differences are expected from extended-grid Doppler
4353        // in forward_model vs data-grid Doppler in broadened_cross_sections.
4354        let interior = 20..energies.len() - 20;
4355        let mut max_err = 0.0f64;
4356        for i in interior {
4357            let err = (t_forward[i] - t_precomputed[i]).abs();
4358            max_err = max_err.max(err);
4359        }
4360        assert!(
4361            max_err < 0.02,
4362            "PrecomputedTransmissionModel with resolution should match \
4363             forward_model.  Max error = {max_err}"
4364        );
4365    }
4366
4367    /// Issue #442: PrecomputedTransmissionModel without resolution must
4368    /// behave identically to the pre-fix version (pure Beer-Lambert).
4369    #[test]
4370    fn precomputed_without_resolution_unchanged() {
4371        let model_no_res = make_precomputed(
4372            vec![vec![100.0, 200.0, 50.0]], // one isotope
4373            vec![0],
4374        );
4375        let params = [0.001f64]; // density
4376        let t = model_no_res.evaluate(&params).unwrap();
4377
4378        // Expected: pure Beer-Lambert.
4379        let expected: Vec<f64> = [100.0, 200.0, 50.0]
4380            .iter()
4381            .map(|&sigma| (-params[0] * sigma).exp())
4382            .collect();
4383
4384        for (i, (&ti, &ei)) in t.iter().zip(expected.iter()).enumerate() {
4385            assert!(
4386                (ti - ei).abs() < 1e-14,
4387                "No-resolution mismatch at bin {i}: got {ti}, expected {ei}"
4388            );
4389        }
4390
4391        // Analytical Jacobian should still be available when instrument is None.
4392        let y = model_no_res.evaluate(&params).unwrap();
4393        assert!(
4394            model_no_res
4395                .analytical_jacobian(&params, &[0], &y)
4396                .is_some(),
4397            "Analytical Jacobian must be available when instrument is None"
4398        );
4399    }
4400
4401    /// PrecomputedTransmissionModel with resolution: analytical Jacobian
4402    /// exists and density derivative matches finite difference.
4403    #[test]
4404    fn precomputed_jacobian_with_resolution_matches_fd() {
4405        use nereids_physics::resolution::ResolutionFunction;
4406
4407        let data = u238_single_resonance();
4408        let temperature = 300.0;
4409        let energies: Vec<f64> = (0..201).map(|i| 4.0 + (i as f64) * 0.025).collect();
4410        let inst = Arc::new(InstrumentParams {
4411            resolution: ResolutionFunction::Gaussian(
4412                nereids_physics::resolution::ResolutionParams::new(25.0, 0.5, 0.005, 0.0).unwrap(),
4413            ),
4414        });
4415
4416        let xs = transmission::broadened_cross_sections(
4417            &energies,
4418            std::slice::from_ref(&data),
4419            temperature,
4420            Some(&inst),
4421            None,
4422        )
4423        .unwrap();
4424        let model = PrecomputedTransmissionModel {
4425            cross_sections: Arc::new(xs),
4426            density_indices: Arc::new(vec![0]),
4427            energies: Some(Arc::new(energies.clone())),
4428            instrument: Some(Arc::clone(&inst)),
4429            resolution_plan: None,
4430            sparse_cubature_plan: None,
4431            sparse_scalar_plan: None,
4432            work_layout: None,
4433        };
4434
4435        let params = [0.0005f64];
4436        let y = model.evaluate(&params).unwrap();
4437
4438        let jac = model
4439            .analytical_jacobian(&params, &[0], &y)
4440            .expect("analytical Jacobian must be available with resolution");
4441
4442        // Finite-difference reference.
4443        let h = 1e-7;
4444        let y_plus = model.evaluate(&[params[0] + h]).unwrap();
4445        let y_minus = model.evaluate(&[params[0] - h]).unwrap();
4446
4447        let interior = 20..energies.len() - 20;
4448        let mut max_rel_err = 0.0f64;
4449        for i in interior {
4450            let fd = (y_plus[i] - y_minus[i]) / (2.0 * h);
4451            let ana = jac.get(i, 0);
4452            let denom = fd.abs().max(ana.abs()).max(1e-30);
4453            max_rel_err = max_rel_err.max((ana - fd).abs() / denom);
4454        }
4455        assert!(
4456            max_rel_err < 0.01,
4457            "PrecomputedTM analytical Jacobian with resolution vs FD: \
4458             max relative error = {max_rel_err}"
4459        );
4460    }
4461
4462    /// PrecomputedTransmissionModel with resolution + shared density param:
4463    /// grouped isotope Jacobian matches FD.
4464    #[test]
4465    fn precomputed_jacobian_grouped_with_resolution_matches_fd() {
4466        use nereids_physics::resolution::ResolutionFunction;
4467
4468        let energies: Vec<f64> = (0..100).map(|i| 1.0 + i as f64 * 0.1).collect();
4469        let inst = Arc::new(InstrumentParams {
4470            resolution: ResolutionFunction::Gaussian(
4471                nereids_physics::resolution::ResolutionParams::new(25.0, 0.5, 0.005, 0.0).unwrap(),
4472            ),
4473        });
4474        // Two isotopes sharing one density parameter.
4475        let xs = vec![vec![10.0; 100], vec![5.0; 100]];
4476        let model = PrecomputedTransmissionModel {
4477            cross_sections: Arc::new(xs),
4478            density_indices: Arc::new(vec![0, 0]), // both share param[0]
4479            energies: Some(Arc::new(energies.clone())),
4480            instrument: Some(Arc::clone(&inst)),
4481            resolution_plan: None,
4482            sparse_cubature_plan: None,
4483            sparse_scalar_plan: None,
4484            work_layout: None,
4485        };
4486
4487        let params = [0.001f64];
4488        let y = model.evaluate(&params).unwrap();
4489        let jac = model
4490            .analytical_jacobian(&params, &[0], &y)
4491            .expect("analytical Jacobian must be available");
4492
4493        let h = 1e-7;
4494        let y_plus = model.evaluate(&[params[0] + h]).unwrap();
4495        let y_minus = model.evaluate(&[params[0] - h]).unwrap();
4496
4497        let mut max_rel_err = 0.0f64;
4498        for i in 10..energies.len() - 10 {
4499            let fd = (y_plus[i] - y_minus[i]) / (2.0 * h);
4500            let ana = jac.get(i, 0);
4501            let denom = fd.abs().max(ana.abs()).max(1e-30);
4502            max_rel_err = max_rel_err.max((ana - fd).abs() / denom);
4503        }
4504        assert!(
4505            max_rel_err < 0.01,
4506            "Grouped PrecomputedTM analytical Jacobian with resolution vs FD: \
4507             max relative error = {max_rel_err}"
4508        );
4509    }
4510
4511    // ── TransmissionFitModel Jacobian with resolution ──────────────────────
4512
4513    /// TransmissionFitModel with resolution: analytical Jacobian exists and
4514    /// density + temperature columns match finite difference.
4515    #[test]
4516    fn transmission_fit_model_jacobian_with_resolution_matches_fd() {
4517        use nereids_physics::resolution::ResolutionFunction;
4518
4519        let data = u238_single_resonance();
4520        let energies: Vec<f64> = (0..201).map(|i| 4.0 + (i as f64) * 0.025).collect();
4521        let inst = Arc::new(InstrumentParams {
4522            resolution: ResolutionFunction::Gaussian(
4523                nereids_physics::resolution::ResolutionParams::new(25.0, 0.5, 0.005, 0.0).unwrap(),
4524            ),
4525        });
4526
4527        let model = TransmissionFitModel::new(
4528            energies.clone(),
4529            vec![data],
4530            300.0,
4531            Some(inst),
4532            (vec![0], vec![1.0]),
4533            Some(1), // temperature_index = 1
4534            None,
4535        )
4536        .unwrap();
4537
4538        let params = [0.0005f64, 300.0];
4539        let y = model.evaluate(&params).unwrap();
4540        let free = vec![0usize, 1usize];
4541
4542        let jac = model
4543            .analytical_jacobian(&params, &free, &y)
4544            .expect("analytical Jacobian must be available with resolution");
4545
4546        // FD for each free param.
4547        let h_density = 1e-7;
4548        let h_temp = 0.01; // temperature needs larger step
4549
4550        for (col, (&fp_idx, &h)) in free.iter().zip([h_density, h_temp].iter()).enumerate() {
4551            let mut p_plus = params;
4552            let mut p_minus = params;
4553            p_plus[fp_idx] += h;
4554            p_minus[fp_idx] -= h;
4555            let y_plus = model.evaluate(&p_plus).unwrap();
4556            let y_minus = model.evaluate(&p_minus).unwrap();
4557
4558            let interior = 20..energies.len() - 20;
4559            let mut max_rel_err = 0.0f64;
4560            for i in interior {
4561                let fd = (y_plus[i] - y_minus[i]) / (2.0 * h);
4562                let ana = jac.get(i, col);
4563                let denom = fd.abs().max(ana.abs()).max(1e-30);
4564                max_rel_err = max_rel_err.max((ana - fd).abs() / denom);
4565            }
4566            let label = if col == 0 { "density" } else { "temperature" };
4567            assert!(
4568                max_rel_err < 0.05,
4569                "TransmissionFitModel {label} column with resolution vs FD: \
4570                 max relative error = {max_rel_err}"
4571            );
4572        }
4573    }
4574
4575    /// TransmissionFitModel without resolution: analytical Jacobian still
4576    /// available and unchanged.
4577    #[test]
4578    fn transmission_fit_model_jacobian_available_without_resolution() {
4579        let data = u238_single_resonance();
4580        let energies: Vec<f64> = (0..101).map(|i| 4.0 + (i as f64) * 0.05).collect();
4581
4582        let model = TransmissionFitModel::new(
4583            energies,
4584            vec![data],
4585            300.0,
4586            None,
4587            (vec![0], vec![1.0]),
4588            Some(1),
4589            None,
4590        )
4591        .unwrap();
4592
4593        let params = [0.0005, 300.0];
4594        let y = model.evaluate(&params).unwrap();
4595
4596        assert!(
4597            model.analytical_jacobian(&params, &[0, 1], &y).is_some(),
4598            "TransmissionFitModel analytical Jacobian must be available \
4599             when resolution is disabled"
4600        );
4601    }
4602
4603    // ── Issue #442: TransmissionFitModel temperature-path resolution fix ───
4604
4605    /// TransmissionFitModel::evaluate() with fit_temperature=true and
4606    /// resolution enabled must match forward_model() for the same sample.
4607    #[test]
4608    fn transmission_fit_model_temp_path_with_resolution_matches_forward_model() {
4609        use nereids_physics::resolution::ResolutionFunction;
4610
4611        let data = u238_single_resonance();
4612        let thickness = 0.0005;
4613        let temperature = 300.0;
4614        let energies: Vec<f64> = (0..401).map(|i| 4.0 + (i as f64) * 0.015).collect();
4615
4616        let inst = Arc::new(InstrumentParams {
4617            resolution: ResolutionFunction::Gaussian(
4618                nereids_physics::resolution::ResolutionParams::new(25.0, 0.5, 0.005, 0.0).unwrap(),
4619            ),
4620        });
4621
4622        // Reference: forward_model() (corrected in Step 1).
4623        let sample = SampleParams::new(temperature, vec![(data.clone(), thickness)]).unwrap();
4624        let t_ref = transmission::forward_model(&energies, &sample, Some(&inst)).unwrap();
4625
4626        // Temperature-fitting path through TransmissionFitModel.
4627        let model = TransmissionFitModel::new(
4628            energies.clone(),
4629            vec![data],
4630            temperature,
4631            Some(Arc::clone(&inst)),
4632            (vec![0], vec![1.0]),
4633            Some(1), // temperature_index
4634            None,
4635        )
4636        .unwrap();
4637
4638        // params = [density, temperature]
4639        let t_model = model.evaluate(&[thickness, temperature]).unwrap();
4640
4641        // Compare on interior (skip boundary effects from extended grid
4642        // differences between forward_model and broadened_cross_sections_from_base).
4643        let interior = 20..energies.len() - 20;
4644        let mut max_err = 0.0f64;
4645        for i in interior {
4646            max_err = max_err.max((t_ref[i] - t_model[i]).abs());
4647        }
4648        assert!(
4649            max_err < 0.02,
4650            "TransmissionFitModel temperature path with resolution should match \
4651             forward_model.  Max error = {max_err}"
4652        );
4653    }
4654
4655    /// TransmissionFitModel temperature path without resolution must be
4656    /// unchanged (pure Doppler + Beer-Lambert).
4657    #[test]
4658    fn transmission_fit_model_temp_path_no_resolution_unchanged() {
4659        let data = u238_single_resonance();
4660        let thickness = 0.0005;
4661        let temperature = 300.0;
4662        let energies: Vec<f64> = (0..201).map(|i| 4.0 + (i as f64) * 0.025).collect();
4663
4664        // Reference: forward_model without resolution.
4665        let sample = SampleParams::new(temperature, vec![(data.clone(), thickness)]).unwrap();
4666        let t_ref = transmission::forward_model(&energies, &sample, None).unwrap();
4667
4668        // TransmissionFitModel, no resolution.
4669        let model = TransmissionFitModel::new(
4670            energies.clone(),
4671            vec![data],
4672            temperature,
4673            None,
4674            (vec![0], vec![1.0]),
4675            Some(1),
4676            None,
4677        )
4678        .unwrap();
4679
4680        let t_model = model.evaluate(&[thickness, temperature]).unwrap();
4681
4682        for (i, (&r, &m)) in t_ref.iter().zip(t_model.iter()).enumerate() {
4683            assert!(
4684                (r - m).abs() < 1e-12,
4685                "No-resolution mismatch at E[{i}]={}: ref={r}, model={m}",
4686                energies[i]
4687            );
4688        }
4689    }
4690
4691    // ── Issue #608: LM-fit resolution must use the auxiliary grid ────────────
4692    //
4693    // The pre-#608 cached / precomputed / energy-scale paths applied resolution
4694    // broadening on the COARSE data grid, unlike `forward_model`, which broadens
4695    // on the auxiliary extended grid and extracts the data points last.  The
4696    // tests below pin every fixed path to `forward_model` — an INDEPENDENT
4697    // oracle: it computes σ inline (`reich_moore::cross_sections_at_energy`) and
4698    // never calls the `broadened_cross_sections` family this fix touches — to
4699    // MACHINE PRECISION over the FULL grid, including the boundary points the
4700    // earlier #442 tests (tol 2e-2, interior-only) excluded.  Each test verifies
4701    // the kernel actually broadens the spectrum (a non-vacuity pre-check, so a
4702    // shared-primitive oracle cannot pass vacuously) and, where it can construct the
4703    // old path, shows the old coarse-grid result differed materially — proving
4704    // the fix is a real correction, not a no-op.  Jacobian columns are checked
4705    // against central finite differences of the (now aux-correct) `evaluate`.
4706    //
4707    // SCOPE of the 1e-9 bound: these tests pin GRID FIDELITY — that
4708    // each fixed path builds the same auxiliary grid + layout as `forward_model`
4709    // and extracts the data points identically.  The resolution KERNEL primitive
4710    // itself (`apply_resolution_*`, `build_aux_grid`, `doppler::doppler_broaden`)
4711    // is SHARED with the oracle, so a kernel error common to both would pass
4712    // here; the kernel's physics is validated independently against SAMMY in
4713    // `nereids-physics` (`resolution.rs`, `samtry_validation.rs`).  The
4714    // non-vacuity `‖kernel − none‖` guards keep this shared-primitive oracle
4715    // non-circular for what it asserts (the #608 grid wiring).
4716
4717    /// Issue #608: the spatial production path (`PrecomputedTransmissionModel`)
4718    /// must broaden resolution on the auxiliary grid, matching `forward_model`.
4719    #[test]
4720    fn issue_608_precomputed_aux_grid_resolution_matches_forward_model() {
4721        use nereids_physics::resolution::ResolutionFunction;
4722
4723        let data = u238_single_resonance();
4724        let thickness = 0.0005;
4725        let temperature = 300.0;
4726        let energies: Vec<f64> = (0..401).map(|i| 4.0 + (i as f64) * 0.015).collect();
4727        let inst = Arc::new(InstrumentParams {
4728            resolution: ResolutionFunction::Gaussian(
4729                nereids_physics::resolution::ResolutionParams::new(25.0, 0.5, 0.005, 0.0).unwrap(),
4730            ),
4731        });
4732
4733        // Independent oracle (computes σ inline; broadens on the aux grid).
4734        let sample = SampleParams::new(temperature, vec![(data.clone(), thickness)]).unwrap();
4735        let t_ref = transmission::forward_model(&energies, &sample, Some(&inst)).unwrap();
4736
4737        // Non-vacuity: the kernel must actually broaden the spectrum, else
4738        // aux-grid vs data-grid broadening would be indistinguishable.
4739        let t_nores = transmission::forward_model(&energies, &sample, None).unwrap();
4740        let broaden = max_abs_diff(&t_ref, &t_nores);
4741        assert!(
4742            broaden > 1e-3 * max_abs(&t_nores),
4743            "resolution kernel must broaden the spectrum non-trivially (got {broaden:.3e})"
4744        );
4745
4746        // FIXED path: working-grid σ + layout, exactly as `spatial_map_typed` builds it.
4747        let working = transmission::broadened_cross_sections_on_working_grid(
4748            &energies,
4749            std::slice::from_ref(&data),
4750            temperature,
4751            Some(&inst),
4752            None,
4753        )
4754        .unwrap();
4755        assert!(
4756            !working.layout.is_identity(),
4757            "Gaussian resolution must build a non-identity auxiliary grid — else \
4758             this test does not exercise the #608 fix"
4759        );
4760        let model_fixed = PrecomputedTransmissionModel {
4761            cross_sections: Arc::new(working.sigma),
4762            density_indices: Arc::new(vec![0]),
4763            energies: Some(Arc::new(energies.clone())),
4764            instrument: Some(Arc::clone(&inst)),
4765            resolution_plan: None,
4766            sparse_cubature_plan: None,
4767            sparse_scalar_plan: None,
4768            work_layout: Some(Arc::new(working.layout)),
4769        };
4770        let t_fixed = model_fixed.evaluate(&[thickness]).unwrap();
4771
4772        // OLD path: data-grid σ, no layout — broadens on the coarse data grid
4773        // (the configuration the pre-#608 spatial pipeline produced).
4774        let xs_data = transmission::broadened_cross_sections(
4775            &energies,
4776            std::slice::from_ref(&data),
4777            temperature,
4778            Some(&inst),
4779            None,
4780        )
4781        .unwrap();
4782        let model_old = PrecomputedTransmissionModel {
4783            cross_sections: Arc::new(xs_data),
4784            density_indices: Arc::new(vec![0]),
4785            energies: Some(Arc::new(energies.clone())),
4786            instrument: Some(Arc::clone(&inst)),
4787            resolution_plan: None,
4788            sparse_cubature_plan: None,
4789            sparse_scalar_plan: None,
4790            work_layout: None,
4791        };
4792        let t_old = model_old.evaluate(&[thickness]).unwrap();
4793
4794        let err_fixed = max_abs_diff(&t_fixed, &t_ref);
4795        let err_old = max_abs_diff(&t_old, &t_ref);
4796
4797        assert!(
4798            err_fixed < 1e-9,
4799            "aux-grid PrecomputedTransmissionModel must match forward_model to \
4800             machine precision over the full grid (got {err_fixed:.3e})"
4801        );
4802        assert!(
4803            err_old > 1e-4 && err_old > 1e4 * err_fixed.max(1e-15),
4804            "old coarse-grid path should differ from forward_model far more than \
4805             the fixed path (old={err_old:.3e}, fixed={err_fixed:.3e})"
4806        );
4807    }
4808
4809    /// Issue #608: `PrecomputedTransmissionModel::analytical_jacobian` forms the
4810    /// inner derivative on the auxiliary grid; it must match central finite
4811    /// differences of the (aux-correct) `evaluate`.
4812    #[test]
4813    fn issue_608_precomputed_aux_grid_jacobian_matches_fd() {
4814        use nereids_physics::resolution::ResolutionFunction;
4815
4816        let data = u238_single_resonance();
4817        let thickness = 0.0005;
4818        let temperature = 300.0;
4819        let energies: Vec<f64> = (0..401).map(|i| 4.0 + (i as f64) * 0.015).collect();
4820        let inst = Arc::new(InstrumentParams {
4821            resolution: ResolutionFunction::Gaussian(
4822                nereids_physics::resolution::ResolutionParams::new(25.0, 0.5, 0.005, 0.0).unwrap(),
4823            ),
4824        });
4825        let working = transmission::broadened_cross_sections_on_working_grid(
4826            &energies,
4827            std::slice::from_ref(&data),
4828            temperature,
4829            Some(&inst),
4830            None,
4831        )
4832        .unwrap();
4833        let model = PrecomputedTransmissionModel {
4834            cross_sections: Arc::new(working.sigma),
4835            density_indices: Arc::new(vec![0]),
4836            energies: Some(Arc::new(energies.clone())),
4837            instrument: Some(Arc::clone(&inst)),
4838            resolution_plan: None,
4839            sparse_cubature_plan: None,
4840            sparse_scalar_plan: None,
4841            work_layout: Some(Arc::new(working.layout)),
4842        };
4843
4844        let params = [thickness];
4845        let free = [0usize];
4846        let y0 = model.evaluate(&params).unwrap();
4847        let jac = model
4848            .analytical_jacobian(&params, &free, &y0)
4849            .expect("analytical jacobian must be available with resolution + aux grid");
4850
4851        let h = 1e-7;
4852        let mut pp = params;
4853        let mut pm = params;
4854        pp[0] += h;
4855        pm[0] -= h;
4856        let yp = model.evaluate(&pp).unwrap();
4857        let ym = model.evaluate(&pm).unwrap();
4858
4859        let mut scale = 0.0f64;
4860        let mut max_err = 0.0f64;
4861        for i in 0..y0.len() {
4862            let fd = (yp[i] - ym[i]) / (2.0 * h);
4863            let an = jac.get(i, 0);
4864            scale = scale.max(an.abs());
4865            max_err = max_err.max((fd - an).abs());
4866        }
4867        let rel = max_err / scale.max(1e-30);
4868        assert!(
4869            rel < 1e-6,
4870            "analytical density Jacobian must match central FD (rel err {rel:.3e})"
4871        );
4872    }
4873
4874    /// Issue #608: `TransmissionFitModel`'s cached temperature-fit `evaluate`
4875    /// must broaden on the auxiliary grid, matching `forward_model` to machine
4876    /// precision over the full grid (the #442 test tolerated 2e-2, interior-only).
4877    #[test]
4878    fn issue_608_transmission_fit_temp_path_aux_grid_matches_forward_model() {
4879        use nereids_physics::resolution::ResolutionFunction;
4880
4881        let data = u238_single_resonance();
4882        let thickness = 0.0005;
4883        let temperature = 300.0;
4884        let energies: Vec<f64> = (0..401).map(|i| 4.0 + (i as f64) * 0.015).collect();
4885        let inst = Arc::new(InstrumentParams {
4886            resolution: ResolutionFunction::Gaussian(
4887                nereids_physics::resolution::ResolutionParams::new(25.0, 0.5, 0.005, 0.0).unwrap(),
4888            ),
4889        });
4890
4891        let sample = SampleParams::new(temperature, vec![(data.clone(), thickness)]).unwrap();
4892        let t_ref = transmission::forward_model(&energies, &sample, Some(&inst)).unwrap();
4893        let t_nores = transmission::forward_model(&energies, &sample, None).unwrap();
4894        let broaden = max_abs_diff(&t_ref, &t_nores);
4895        assert!(
4896            broaden > 1e-3 * max_abs(&t_nores),
4897            "resolution kernel must broaden the spectrum non-trivially (got {broaden:.3e})"
4898        );
4899
4900        let model = TransmissionFitModel::new(
4901            energies.clone(),
4902            vec![data],
4903            temperature,
4904            Some(Arc::clone(&inst)),
4905            (vec![0], vec![1.0]),
4906            Some(1), // temperature_index → exercises the cached temperature path
4907            None,
4908        )
4909        .unwrap();
4910        let t_model = model.evaluate(&[thickness, temperature]).unwrap();
4911
4912        let err = max_abs_diff(&t_model, &t_ref);
4913        assert!(
4914            err < 1e-9,
4915            "aux-grid TransmissionFitModel temperature path must match \
4916             forward_model over the full grid (got {err:.3e})"
4917        );
4918    }
4919
4920    /// Issue #608: `TransmissionFitModel::analytical_jacobian` (cached temp path)
4921    /// forms density and temperature inner derivatives on the auxiliary grid;
4922    /// both columns must match central finite differences of `evaluate`.
4923    #[test]
4924    fn issue_608_transmission_fit_temp_path_jacobian_matches_fd() {
4925        use nereids_physics::resolution::ResolutionFunction;
4926
4927        let data = u238_single_resonance();
4928        let thickness = 0.0005;
4929        let temperature = 300.0;
4930        let energies: Vec<f64> = (0..401).map(|i| 4.0 + (i as f64) * 0.015).collect();
4931        let inst = Arc::new(InstrumentParams {
4932            resolution: ResolutionFunction::Gaussian(
4933                nereids_physics::resolution::ResolutionParams::new(25.0, 0.5, 0.005, 0.0).unwrap(),
4934            ),
4935        });
4936
4937        let model = TransmissionFitModel::new(
4938            energies.clone(),
4939            vec![data],
4940            temperature,
4941            Some(Arc::clone(&inst)),
4942            (vec![0], vec![1.0]),
4943            Some(1),
4944            None,
4945        )
4946        .unwrap();
4947
4948        let params = [thickness, temperature];
4949        // evaluate() populates the broadened-σ cache at these params; the
4950        // analytical jacobian reads that cache, so compute it BEFORE any FD
4951        // perturbation mutates the cache.
4952        let y0 = model.evaluate(&params).unwrap();
4953        let free = [0usize, 1usize];
4954        let jac = model
4955            .analytical_jacobian(&params, &free, &y0)
4956            .expect("analytical jacobian must be available with resolution + aux grid");
4957
4958        // Per-parameter central-FD step (absolute): density ~5e-4, temperature 300 K.
4959        let steps = [1e-7, 1e-2];
4960        for (col, &p_idx) in free.iter().enumerate() {
4961            let h = steps[col];
4962            let mut pp = params;
4963            let mut pm = params;
4964            pp[p_idx] += h;
4965            pm[p_idx] -= h;
4966            let yp = model.evaluate(&pp).unwrap();
4967            let ym = model.evaluate(&pm).unwrap();
4968            let mut scale = 0.0f64;
4969            let mut max_err = 0.0f64;
4970            for i in 0..y0.len() {
4971                let fd = (yp[i] - ym[i]) / (2.0 * h);
4972                let an = jac.get(i, col);
4973                scale = scale.max(an.abs());
4974                max_err = max_err.max((fd - an).abs());
4975            }
4976            let rel = max_err / scale.max(1e-30);
4977            assert!(
4978                rel < 1e-5,
4979                "analytical Jacobian column {col} must match central FD (rel err {rel:.3e})"
4980            );
4981        }
4982    }
4983
4984    /// Issue #608: EnergyScale must evaluate the TRUE σ at the
4985    /// corrected energies on the auxiliary grid — INCLUDING the boundary
4986    /// extension points — exactly like `forward_model`, not clamp a precomputed
4987    /// σ.  With the U-238 resonance near the grid EDGE (where the pre-fix clamp
4988    /// deviated most) and Gaussian resolution active, EnergyScale at identity
4989    /// calibration must match `forward_model` — an independent oracle that
4990    /// evaluates σ inline — to machine precision over the FULL grid.  This is the
4991    /// non-circular replacement for the previous flat-σ/clamp-oracle test (which
4992    /// could not detect the boundary deviation).
4993    #[test]
4994    fn issue_608_energy_scale_aux_grid_true_sigma_matches_forward_model() {
4995        use nereids_physics::resolution::ResolutionFunction;
4996
4997        let data = u238_single_resonance();
4998        let density = 0.01;
4999        // Grid placing the U-238 resonance (~6.67 eV) near the UPPER edge, so σ
5000        // is strongly non-flat at the boundary — exactly where clamping (the
5001        // pre-#608 behaviour) deviated from true physics.
5002        let energies: Vec<f64> = (0..121).map(|i| 5.0 + (i as f64) * 0.015).collect();
5003        let inst = Arc::new(InstrumentParams {
5004            resolution: ResolutionFunction::Gaussian(
5005                nereids_physics::resolution::ResolutionParams::new(25.0, 0.5, 0.005, 0.0).unwrap(),
5006            ),
5007        });
5008
5009        let model = make_energy_scale_u238(energies.clone(), Some(Arc::clone(&inst)));
5010        let t_es = model.evaluate(&[density, 0.0, 1.0]).unwrap();
5011
5012        // Independent oracle: forward_model evaluates σ inline on the aux grid.
5013        let sample = SampleParams::new(300.0, vec![(data, density)]).unwrap();
5014        let t_ref = transmission::forward_model(&energies, &sample, Some(&inst)).unwrap();
5015
5016        // Non-vacuity: the resolution kernel must broaden the spectrum.
5017        let t_nores = transmission::forward_model(&energies, &sample, None).unwrap();
5018        let broaden = max_abs_diff(&t_ref, &t_nores);
5019        assert!(
5020            broaden > 1e-3 * max_abs(&t_nores),
5021            "resolution kernel must broaden the spectrum non-trivially (got {broaden:.3e})"
5022        );
5023
5024        // True-σ aux-grid EnergyScale matches forward_model over the FULL grid,
5025        // including the resonance-near-edge boundary where the old clamp failed.
5026        let err = max_abs_diff(&t_es, &t_ref);
5027        assert!(
5028            err < 1e-9,
5029            "EnergyScale identity-calibration evaluate must match forward_model to \
5030             machine precision over the full grid (got {err:.3e})"
5031        );
5032    }
5033
5034    /// Issue #608: the GROUPED energy-scale path — multiple isotopes mapped
5035    /// to ONE density parameter with non-unity ratios — is reachable in
5036    /// production (`with_groups` + `fit_energy_scale`) but was exercised by no
5037    /// test; every other energy-scale test used a single isotope
5038    /// (`density_indices=[0]`, ratio 1.0).  Build two DISTINCT isotopes sharing
5039    /// density param 0 with ratios (0.7, 0.3) and verify the per-member
5040    /// Beer-Lambert accumulation (`Σᵢ n·ratioᵢ·σᵢ`) matches `forward_model` with
5041    /// per-isotope effective densities — plus an FD check on the single shared
5042    /// density column.
5043    #[test]
5044    fn issue_608_energy_scale_grouped_density_matches_forward_model() {
5045        use nereids_endf::resonance::test_support::synthetic_single_resonance;
5046        use nereids_physics::resolution::ResolutionFunction;
5047
5048        let iso0 = u238_single_resonance(); // resonance @ ~6.674 eV
5049        let iso1 = synthetic_single_resonance(72, 178, 176.0, 7.5); // distinct @ 7.5 eV
5050        let density = 0.01_f64;
5051        let ratios = [0.7_f64, 0.3_f64];
5052        // Grid overlapping BOTH resonances so σ0 ≠ σ1 (a swapped ratio / wrong
5053        // index shifts T detectably — proven by the swap guard below).
5054        let energies: Vec<f64> = (0..201).map(|i| 5.0 + (i as f64) * 0.02).collect();
5055        let inst = Arc::new(InstrumentParams {
5056            resolution: ResolutionFunction::Gaussian(
5057                nereids_physics::resolution::ResolutionParams::new(25.0, 0.5, 0.005, 0.0).unwrap(),
5058            ),
5059        });
5060
5061        let model = EnergyScaleTransmissionModel::new(
5062            Arc::new(vec![iso0.clone(), iso1.clone()]),
5063            Arc::new(vec![0, 0]), // both isotopes → density param 0 (grouped)
5064            Arc::new(vec![ratios[0], ratios[1]]),
5065            300.0,
5066            energies.clone(),
5067            25.0,
5068            1, // t0 index
5069            2, // l_scale index
5070            Some(Arc::clone(&inst)),
5071        );
5072        let params = [density, 0.0, 1.0]; // identity calibration (t0=0, l_scale=1)
5073        let t_es = model.evaluate(&params).unwrap();
5074
5075        // Independent oracle: forward_model with per-isotope effective areal
5076        // densities n·ratioᵢ.  Beer-Lambert is additive over isotopes, so the
5077        // grouped model (one density param × per-iso ratio) must equal a two-
5078        // isotope sample with densities (n·0.7, n·0.3).
5079        let sample = SampleParams::new(
5080            300.0,
5081            vec![
5082                (iso0.clone(), density * ratios[0]),
5083                (iso1.clone(), density * ratios[1]),
5084            ],
5085        )
5086        .unwrap();
5087        let t_ref = transmission::forward_model(&energies, &sample, Some(&inst)).unwrap();
5088
5089        // Non-vacuity: the kernel must broaden the grouped spectrum.
5090        let t_nores = transmission::forward_model(&energies, &sample, None).unwrap();
5091        assert!(
5092            max_abs_diff(&t_ref, &t_nores) > 1e-3 * max_abs(&t_nores),
5093            "resolution kernel must broaden the grouped spectrum non-trivially"
5094        );
5095
5096        // Discrimination: swapping the two ratios MUST change T (proves σ0 ≠ σ1
5097        // over the grid, so the match assertion below is sensitive to a ratio /
5098        // index mix-up in the per-member accumulation — i.e. non-vacuous).
5099        let model_swapped = EnergyScaleTransmissionModel::new(
5100            Arc::new(vec![iso0.clone(), iso1.clone()]),
5101            Arc::new(vec![0, 0]),
5102            Arc::new(vec![ratios[1], ratios[0]]), // swapped
5103            300.0,
5104            energies.clone(),
5105            25.0,
5106            1,
5107            2,
5108            Some(Arc::clone(&inst)),
5109        );
5110        let t_swapped = model_swapped.evaluate(&params).unwrap();
5111        assert!(
5112            max_abs_diff(&t_es, &t_swapped) > 1e-4,
5113            "swapping the two density ratios must change T (else the test could \
5114             not distinguish the ratio→isotope assignment)"
5115        );
5116
5117        // Grouped evaluate matches the independent oracle to machine precision.
5118        let err = max_abs_diff(&t_es, &t_ref);
5119        assert!(
5120            err < 1e-9,
5121            "grouped EnergyScale (2 isotopes → 1 density param, ratios {ratios:?}) \
5122             must match forward_model with per-isotope effective densities to \
5123             machine precision (got {err:.3e})"
5124        );
5125
5126        // FD check on the single shared density column: ∂T/∂n accumulates
5127        // ratioᵢ·σᵢ over BOTH grouped isotopes.
5128        let free = vec![0usize];
5129        let jac = model
5130            .analytical_jacobian(&params, &free, &t_es)
5131            .expect("Jacobian should be available");
5132        let h = 1e-7;
5133        let mut pp = params;
5134        let mut pm = params;
5135        pp[0] += h;
5136        pm[0] -= h;
5137        let yp = model.evaluate(&pp).unwrap();
5138        let ym = model.evaluate(&pm).unwrap();
5139        for row in 0..energies.len() {
5140            let fd = (yp[row] - ym[row]) / (2.0 * h);
5141            let anal = jac.get(row, 0);
5142            let abs_err = (anal - fd).abs();
5143            let rel_err = abs_err / fd.abs().max(1e-15);
5144            assert!(
5145                rel_err < 1e-3 || abs_err < 1e-8,
5146                "grouped density col bin {row}: anal={anal:.6e} fd={fd:.6e} rel={rel_err:.2e}"
5147            );
5148        }
5149    }
5150
5151    /// Resolution-enabled temperature path must produce measurably different
5152    /// results from the unresolved path (verifies resolution is being applied).
5153    #[test]
5154    fn transmission_fit_model_temp_path_resolution_makes_difference() {
5155        use nereids_physics::resolution::ResolutionFunction;
5156
5157        let data = u238_single_resonance();
5158        let thickness = 0.0005;
5159        let temperature = 300.0;
5160        let energies: Vec<f64> = (0..401).map(|i| 4.0 + (i as f64) * 0.015).collect();
5161
5162        let inst = Arc::new(InstrumentParams {
5163            resolution: ResolutionFunction::Gaussian(
5164                nereids_physics::resolution::ResolutionParams::new(25.0, 0.5, 0.005, 0.0).unwrap(),
5165            ),
5166        });
5167
5168        // With resolution.
5169        let model_res = TransmissionFitModel::new(
5170            energies.clone(),
5171            vec![data.clone()],
5172            temperature,
5173            Some(inst),
5174            (vec![0], vec![1.0]),
5175            Some(1),
5176            None,
5177        )
5178        .unwrap();
5179        let t_res = model_res.evaluate(&[thickness, temperature]).unwrap();
5180
5181        // Without resolution.
5182        let model_no = TransmissionFitModel::new(
5183            energies.clone(),
5184            vec![data],
5185            temperature,
5186            None,
5187            (vec![0], vec![1.0]),
5188            Some(1),
5189            None,
5190        )
5191        .unwrap();
5192        let t_no = model_no.evaluate(&[thickness, temperature]).unwrap();
5193
5194        let interior = 20..energies.len() - 20;
5195        let max_diff: f64 = interior
5196            .map(|i| (t_res[i] - t_no[i]).abs())
5197            .fold(0.0f64, f64::max);
5198        assert!(
5199            max_diff > 1e-4,
5200            "Resolution should make a measurable difference in the temperature \
5201             path, but max diff = {max_diff}"
5202        );
5203    }
5204
5205    // ── Exponential background (BackD, BackF) tests ──
5206
5207    /// Verify that new_with_exponential evaluate() matches the formula:
5208    /// T_out = Anorm*T_inner + BackA + BackB/√E + BackC*√E + BackD*exp(-BackF/√E)
5209    #[test]
5210    fn exponential_evaluate_formula_correct() {
5211        let xs = vec![vec![1.0, 2.0, 3.0]];
5212        let inner = make_precomputed(xs, vec![0]);
5213        let energies = [4.0, 9.0, 25.0]; // sqrt = [2, 3, 5]
5214
5215        let model =
5216            NormalizedTransmissionModel::new_with_exponential(inner, &energies, 1, 2, 3, 4, 5, 6);
5217
5218        // params: [density, anorm, back_a, back_b, back_c, back_d, back_f]
5219        let density = 0.1;
5220        let anorm = 1.02;
5221        let back_a = 0.01;
5222        let back_b = 0.005;
5223        let back_c = 0.002;
5224        let back_d = 0.05;
5225        let back_f = 3.0;
5226        let params = [density, anorm, back_a, back_b, back_c, back_d, back_f];
5227
5228        let y = model.evaluate(&params).unwrap();
5229
5230        // Manually compute expected
5231        let xs_vals = [1.0, 2.0, 3.0];
5232        let sqrt_e = [2.0, 3.0, 5.0];
5233        for i in 0..3 {
5234            let t_inner = (-density * xs_vals[i]).exp();
5235            let expected = anorm * t_inner
5236                + back_a
5237                + back_b / sqrt_e[i]
5238                + back_c * sqrt_e[i]
5239                + back_d * (-back_f / sqrt_e[i]).exp();
5240            assert!(
5241                (y[i] - expected).abs() < 1e-12,
5242                "bin {i}: got {}, expected {expected}",
5243                y[i]
5244            );
5245        }
5246    }
5247
5248    /// Analytical Jacobian for BackD and BackF columns must match central FD.
5249    #[test]
5250    fn exponential_jacobian_matches_finite_difference() {
5251        let xs = vec![vec![1.0, 2.0, 3.0, 0.5, 1.5]];
5252        let inner = make_precomputed(xs, vec![0]);
5253        let energies = [0.1, 1.0, 4.0, 25.0, 100.0]; // span 0.1–100 eV
5254
5255        let model =
5256            NormalizedTransmissionModel::new_with_exponential(inner, &energies, 1, 2, 3, 4, 5, 6);
5257
5258        // params: [density, anorm, back_a, back_b, back_c, back_d, back_f]
5259        let params = [0.1, 1.02, 0.01, 0.005, 0.002, 0.05, 3.0];
5260        let y = model.evaluate(&params).unwrap();
5261        let free_indices: Vec<usize> = (0..7).collect();
5262
5263        let jac = model
5264            .analytical_jacobian(&params, &free_indices, &y)
5265            .expect("analytical Jacobian should be available");
5266
5267        // Central finite difference for all parameters
5268        let h = 1e-7;
5269        for (col, &pidx) in free_indices.iter().enumerate() {
5270            let mut p_plus = params.to_vec();
5271            let mut p_minus = params.to_vec();
5272            p_plus[pidx] += h;
5273            p_minus[pidx] -= h;
5274            let y_plus = model.evaluate(&p_plus).unwrap();
5275            let y_minus = model.evaluate(&p_minus).unwrap();
5276
5277            for row in 0..energies.len() {
5278                let fd = (y_plus[row] - y_minus[row]) / (2.0 * h);
5279                let anal = jac.get(row, col);
5280                let abs_err = (anal - fd).abs();
5281                let rel_err = abs_err / fd.abs().max(1e-15);
5282                assert!(
5283                    rel_err < 1e-5 || abs_err < 1e-10,
5284                    "param {pidx} (col {col}), bin {row}: analytical={anal:.10e}, \
5285                     fd={fd:.10e}, rel_err={rel_err:.2e}"
5286                );
5287            }
5288        }
5289    }
5290
5291    /// Round-trip: fit recovers all 6 background + density from noiseless data.
5292    #[test]
5293    fn exponential_fit_recovers_all_params() {
5294        let xs = vec![vec![1.0, 2.0, 3.0, 2.0, 1.5, 0.8, 1.2, 2.5]];
5295        let inner = make_precomputed(xs, vec![0]);
5296        let energies = [0.5, 1.0, 4.0, 9.0, 16.0, 25.0, 36.0, 64.0];
5297
5298        let model =
5299            NormalizedTransmissionModel::new_with_exponential(inner, &energies, 1, 2, 3, 4, 5, 6);
5300
5301        // True parameters
5302        let true_density = 0.15;
5303        let true_anorm = 1.02;
5304        let true_back_a = 0.01;
5305        let true_back_b = 0.005;
5306        let true_back_c = 0.002;
5307        let true_back_d = 0.03;
5308        let true_back_f = 2.0;
5309        let true_params = [
5310            true_density,
5311            true_anorm,
5312            true_back_a,
5313            true_back_b,
5314            true_back_c,
5315            true_back_d,
5316            true_back_f,
5317        ];
5318
5319        let y_obs = model.evaluate(&true_params).unwrap();
5320        let sigma = vec![0.001; y_obs.len()];
5321
5322        let mut params = ParameterSet::new(vec![
5323            FitParameter::non_negative("density", 0.1),
5324            FitParameter {
5325                name: "anorm".into(),
5326                value: 1.0,
5327                lower: 0.5,
5328                upper: 1.5,
5329                fixed: false,
5330            },
5331            FitParameter {
5332                name: "back_a".into(),
5333                value: 0.0,
5334                lower: -0.5,
5335                upper: 0.5,
5336                fixed: false,
5337            },
5338            FitParameter {
5339                name: "back_b".into(),
5340                value: 0.0,
5341                lower: -0.5,
5342                upper: 0.5,
5343                fixed: false,
5344            },
5345            FitParameter {
5346                name: "back_c".into(),
5347                value: 0.0,
5348                lower: -0.5,
5349                upper: 0.5,
5350                fixed: false,
5351            },
5352            FitParameter {
5353                name: "back_d".into(),
5354                value: 0.01,
5355                lower: 0.0,
5356                upper: 1.0,
5357                fixed: false,
5358            },
5359            FitParameter {
5360                name: "back_f".into(),
5361                value: 1.0,
5362                lower: 0.0,
5363                upper: 100.0,
5364                fixed: false,
5365            },
5366        ]);
5367
5368        let config = LmConfig {
5369            max_iter: 500,
5370            ..LmConfig::default()
5371        };
5372
5373        let result = lm::levenberg_marquardt(&model, &y_obs, &sigma, &mut params, &config).unwrap();
5374
5375        assert!(result.converged, "Fit should converge");
5376
5377        let fitted = &result.params;
5378        let check = |name, fitted_val: f64, true_val: f64, tol: f64| {
5379            let err = (fitted_val - true_val).abs();
5380            let rel = err / true_val.abs().max(1e-10);
5381            assert!(
5382                rel < tol || err < 1e-6,
5383                "{name}: fitted={fitted_val:.6}, true={true_val:.6}, rel_err={rel:.4}"
5384            );
5385        };
5386
5387        check("density", fitted[0], true_density, 0.10);
5388        check("anorm", fitted[1], true_anorm, 0.10);
5389        check("back_a", fitted[2], true_back_a, 0.10);
5390        check("back_b", fitted[3], true_back_b, 0.10);
5391        check("back_c", fitted[4], true_back_c, 0.10);
5392        check("back_d", fitted[5], true_back_d, 0.10);
5393        check("back_f", fitted[6], true_back_f, 0.10);
5394    }
5395
5396    // ── EnergyScaleTransmissionModel tests ──
5397
5398    /// Verify that corrected_energies shifts the grid correctly.
5399    /// Build a single-isotope (U-238) EnergyScale model for tests: density at
5400    /// param 0, t0 at param 1, l_scale at param 2.  Issue #608: σ is evaluated
5401    /// from the resonance at the corrected energies (matching forward_model), so
5402    /// test grids should overlap the U-238 resonance (~6.67 eV) for non-trivial
5403    /// σ.  Temperature 300 K, flight path 25 m.
5404    fn make_energy_scale_u238(
5405        energies: Vec<f64>,
5406        instrument: Option<Arc<InstrumentParams>>,
5407    ) -> EnergyScaleTransmissionModel {
5408        EnergyScaleTransmissionModel::new(
5409            Arc::new(vec![u238_single_resonance()]),
5410            Arc::new(vec![0]),
5411            Arc::new(vec![1.0]),
5412            300.0,
5413            energies,
5414            25.0,
5415            1,
5416            2,
5417            instrument,
5418        )
5419    }
5420
5421    #[test]
5422    fn energy_scale_corrected_energies() {
5423        let energies = vec![10.0, 20.0, 50.0, 100.0, 200.0];
5424        let model = make_energy_scale_u238(energies.clone(), None);
5425
5426        // With t0=0, l_scale=1: corrected energies should equal nominal
5427        let e_corr = model.corrected_energies(0.0, 1.0);
5428        for (i, (&nom, &corr)) in energies.iter().zip(e_corr.iter()).enumerate() {
5429            assert!(
5430                (nom - corr).abs() / nom < 1e-10,
5431                "bin {i}: nominal={nom}, corrected={corr}"
5432            );
5433        }
5434
5435        // With l_scale > 1: all corrected energies should increase
5436        let e_corr_ls = model.corrected_energies(0.0, 1.005);
5437        for (i, (&nom, &corr)) in energies.iter().zip(e_corr_ls.iter()).enumerate() {
5438            assert!(
5439                corr > nom,
5440                "bin {i}: l_scale=1.005 should increase energy, got nom={nom}, corr={corr}"
5441            );
5442        }
5443
5444        // With t0 > 0: energies should increase (shorter effective TOF)
5445        let e_corr_t0 = model.corrected_energies(1.0, 1.0);
5446        for (i, (&nom, &corr)) in energies.iter().zip(e_corr_t0.iter()).enumerate() {
5447            assert!(
5448                corr > nom,
5449                "bin {i}: t0=1.0 should increase energy, got nom={nom}, corr={corr}"
5450            );
5451        }
5452    }
5453
5454    /// Issue #608: at identity calibration (t0=0, l_scale=1) the corrected grid
5455    /// equals the nominal grid, so EnergyScale must evaluate the SAME true σ as
5456    /// `forward_model` — the independent oracle — to machine precision.
5457    #[test]
5458    fn energy_scale_evaluate_identity() {
5459        let energies: Vec<f64> = (0..201).map(|i| 4.0 + (i as f64) * 0.03).collect();
5460        let density = 0.01;
5461        let model_es = make_energy_scale_u238(energies.clone(), None);
5462        let y_es = model_es.evaluate(&[density, 0.0, 1.0]).unwrap();
5463
5464        let sample = SampleParams::new(300.0, vec![(u238_single_resonance(), density)]).unwrap();
5465        let y_ref = transmission::forward_model(&energies, &sample, None).unwrap();
5466
5467        for (i, (&a, &b)) in y_es.iter().zip(y_ref.iter()).enumerate() {
5468            assert!(
5469                (a - b).abs() < 1e-10,
5470                "bin {i}: energy_scale={a}, forward_model={b}"
5471            );
5472        }
5473    }
5474
5475    /// Jacobian for energy-scale model: density columns must match FD.
5476    #[test]
5477    fn energy_scale_jacobian_density_matches_fd() {
5478        let energies: Vec<f64> = (0..101).map(|i| 4.0 + (i as f64) * 0.06).collect();
5479        let model = make_energy_scale_u238(energies.clone(), None);
5480
5481        let params = [0.01, 0.5, 1.002]; // density, t0, l_scale
5482        let y = model.evaluate(&params).unwrap();
5483        // Density column only (matching this test's name).  The energy-scale
5484        // (t0 / L_scale) columns are FD-based and method-dependent; they are
5485        // covered against a matching-h FD2 reference by the partial_gal_* tests.
5486        // Comparing them to a different-h FD here would be apples-to-oranges,
5487        // especially on the sharp U-238 resonance (#608 migration
5488        // to true-σ resonance data).
5489        let free = vec![0];
5490        let jac = model
5491            .analytical_jacobian(&params, &free, &y)
5492            .expect("Jacobian should be available");
5493
5494        let h = 1e-7;
5495        for (col, &pidx) in free.iter().enumerate() {
5496            let mut pp = params.to_vec();
5497            let mut pm = params.to_vec();
5498            pp[pidx] += h;
5499            pm[pidx] -= h;
5500            let yp = model.evaluate(&pp).unwrap();
5501            let ym = model.evaluate(&pm).unwrap();
5502            for row in 0..energies.len() {
5503                let fd = (yp[row] - ym[row]) / (2.0 * h);
5504                let anal = jac.get(row, col);
5505                let abs_err = (anal - fd).abs();
5506                let rel_err = abs_err / fd.abs().max(1e-15);
5507                assert!(
5508                    rel_err < 1e-3 || abs_err < 1e-8,
5509                    "param {pidx} col {col} bin {row}: anal={anal:.6e} fd={fd:.6e} rel={rel_err:.2e}"
5510                );
5511            }
5512        }
5513    }
5514
5515    /// LM fit with energy-scale model recovers l_scale from shifted data.
5516    ///
5517    /// Uses a sharp Breit-Wigner-like resonance on a dense grid so the
5518    /// energy shift is unambiguous.  Only l_scale is varied (t0 fixed
5519    /// at 0) to avoid degenerate local minima.
5520    #[test]
5521    fn energy_scale_fit_recovers_l_scale() {
5522        // Dense grid over the sharp U-238 resonance (~6.67 eV) so the energy
5523        // shift is unambiguous.  Only l_scale is varied (t0 fixed at 0).
5524        let energies: Vec<f64> = (0..200).map(|i| 4.0 + (i as f64) * 0.03).collect();
5525
5526        let true_density = 0.001;
5527        let true_ls = 1.003;
5528
5529        let model = make_energy_scale_u238(energies, None);
5530        let true_params = [true_density, 0.0, true_ls];
5531        let y_obs = model.evaluate(&true_params).unwrap();
5532        let sigma = vec![0.001; y_obs.len()];
5533
5534        let mut params = ParameterSet::new(vec![
5535            FitParameter::non_negative("density", 0.0005),
5536            FitParameter::fixed("t0", 0.0),
5537            FitParameter {
5538                name: "l_scale".into(),
5539                value: 1.0,
5540                lower: 0.99,
5541                upper: 1.01,
5542                fixed: false,
5543            },
5544        ]);
5545
5546        let config = LmConfig {
5547            max_iter: 200,
5548            ..LmConfig::default()
5549        };
5550
5551        let result = lm::levenberg_marquardt(&model, &y_obs, &sigma, &mut params, &config).unwrap();
5552
5553        assert!(result.converged, "Fit should converge");
5554        let f = &result.params;
5555        assert!(
5556            (f[0] - true_density).abs() / true_density < 0.05,
5557            "density: fitted={}, true={true_density}",
5558            f[0]
5559        );
5560        assert!(
5561            (f[2] - true_ls).abs() < 0.001,
5562            "l_scale: fitted={}, true={true_ls}",
5563            f[2]
5564        );
5565    }
5566
5567    /// Partial-GAL Jacobian with NO resolution should match FD2 to f64
5568    /// roundoff: in this regime the rank-1 identity
5569    /// `J[:, L_scale] = ((tof - t0) / L_scale) * J[:, t0]` is exact (the
5570    /// forward chain factorises through `e_corr` only, with no
5571    /// resolution operator to introduce additional `(t0, L_scale)`
5572    /// dependence).  Issue #489.
5573    #[test]
5574    fn partial_gal_no_resolution_matches_fd2() {
5575        let energies: Vec<f64> = (0..101).map(|i| 4.0 + (i as f64) * 0.06).collect();
5576        // Pin both reference and alt models explicitly via
5577        // `with_jacobian_method` so the test is independent of the
5578        // process-global `NEREIDS_TZERO_JACOBIAN` env var.  Without
5579        // pinning, the post-#489 default of `PartialGal` would make
5580        // the "FD2 reference" actually run partial-GAL (vacuous
5581        // self-comparison).
5582        let mut model = make_energy_scale_u238(energies.clone(), None)
5583            .with_jacobian_method(EnergyScaleJacobianMethod::FiniteDifference);
5584
5585        let params = [0.001, 0.05, 1.002]; // density, t0, l_scale
5586        let free = vec![0, 1, 2];
5587
5588        // FD2 reference Jacobian (explicitly pinned above).
5589        let jac_fd2 = model
5590            .analytical_jacobian(&params, &free, &model.evaluate(&params).unwrap())
5591            .expect("FD2 Jacobian should be available");
5592
5593        // Partial-GAL Jacobian.
5594        model = model.with_jacobian_method(EnergyScaleJacobianMethod::PartialGal);
5595        let jac_pg = model
5596            .analytical_jacobian(&params, &free, &model.evaluate(&params).unwrap())
5597            .expect("partial-GAL Jacobian should be available");
5598
5599        // Density column: identical (analytical, not affected by method).
5600        for i in 0..energies.len() {
5601            let fd2 = jac_fd2.get(i, 0);
5602            let pg = jac_pg.get(i, 0);
5603            assert!(
5604                (fd2 - pg).abs() < 1e-15,
5605                "density bin {i}: fd2={fd2:.6e} pg={pg:.6e}"
5606            );
5607        }
5608        // t0 column: identical (both methods use the same FD pair when
5609        // both t0 and L_scale are free; partial-GAL just hoists it out
5610        // of the loop).
5611        for i in 0..energies.len() {
5612            let fd2 = jac_fd2.get(i, 1);
5613            let pg = jac_pg.get(i, 1);
5614            assert!(
5615                (fd2 - pg).abs() < 1e-15,
5616                "t0 bin {i}: fd2={fd2:.6e} pg={pg:.6e}"
5617            );
5618        }
5619        // L_scale column: the rank-1 derivation is analytically exact without
5620        // resolution.  The only residual vs FD2 is the difference in central-FD
5621        // truncation — PartialGal's L_scale inherits the t0 step (h=1e-4), FD2
5622        // takes a direct L_scale step (h=1e-7).  On the sharp U-238 resonance
5623        // that truncation dominates small-derivative TAIL bins (per-bin rel can
5624        // hit a few % there while contributing negligibly to the spectrum), so
5625        // compare the aggregate relative L₂ — the same robust metric the
5626        // with-resolution sister test uses.  Measured ~8.0e-3 here; the bound
5627        // (2.5e-2) gives ~3× headroom yet is far below the O(1) a broken rank-1
5628        // identity would produce.
5629        let mut num_sq = 0.0_f64;
5630        let mut den_sq = 0.0_f64;
5631        for i in 0..energies.len() {
5632            let fd2 = jac_fd2.get(i, 2);
5633            let pg = jac_pg.get(i, 2);
5634            let diff = pg - fd2;
5635            num_sq += diff * diff;
5636            den_sq += fd2 * fd2;
5637        }
5638        let rel_l2 = (num_sq / den_sq.max(1e-30)).sqrt();
5639        assert!(
5640            rel_l2 < 2.5e-2,
5641            "L_scale rank-1 vs FD2 rel L₂ = {rel_l2:.3e} (expected ≪ 1 without \
5642             resolution — the rank-1 identity is exact up to FD truncation)"
5643        );
5644    }
5645
5646    /// When only L_scale is free (t0 fixed), partial-GAL falls through
5647    /// to standard FD: there is no t0 column to derive L_scale from,
5648    /// so the per-coordinate FD path must still be used. Verifies the
5649    /// dispatch logic correctly handles this case.
5650    #[test]
5651    fn partial_gal_l_scale_only_falls_through_to_fd() {
5652        let energies: Vec<f64> = (0..101).map(|i| 4.0 + (i as f64) * 0.06).collect();
5653        let model = make_energy_scale_u238(energies.clone(), None)
5654            .with_jacobian_method(EnergyScaleJacobianMethod::PartialGal);
5655
5656        let params = [0.001, 0.0, 1.002];
5657        let free = vec![0, 2]; // density + L_scale (no t0)
5658        let y = model.evaluate(&params).unwrap();
5659        let jac = model
5660            .analytical_jacobian(&params, &free, &y)
5661            .expect("Jacobian should be available even when t0 not free");
5662
5663        // L_scale column should match a manual central FD reference.
5664        let h = 1e-7;
5665        let mut pp = params.to_vec();
5666        let mut pm = params.to_vec();
5667        pp[2] += h;
5668        pm[2] -= h;
5669        let yp = model.evaluate(&pp).unwrap();
5670        let ym = model.evaluate(&pm).unwrap();
5671        for i in 0..energies.len() {
5672            let fd = (yp[i] - ym[i]) / (2.0 * h);
5673            let anal = jac.get(i, 1);
5674            let abs_err = (anal - fd).abs();
5675            let rel_err = abs_err / fd.abs().max(1e-15);
5676            assert!(
5677                rel_err < 1e-3 || abs_err < 1e-8,
5678                "L_scale bin {i}: anal={anal:.6e} fd={fd:.6e} rel={rel_err:.2e}"
5679            );
5680        }
5681    }
5682
5683    /// Regression for #500: at `l_scale ≈ 0` the partial-GAL rank-1
5684    /// derivation `(tof - t0_clamped) / l_scale` divides by zero,
5685    /// producing a NaN L_scale Jacobian column.  After the fix, the
5686    /// L_scale column falls through to the per-coordinate FD path —
5687    /// every Jacobian entry must be finite, and the L_scale column
5688    /// must agree with the FD2 reference (which uses the same
5689    /// per-coordinate FD).
5690    ///
5691    /// Setup mirrors `partial_gal_no_resolution_matches_fd2` so the
5692    /// FD-tolerance comparison against FD2 stays apples-to-apples.
5693    #[test]
5694    fn partial_gal_l_scale_zero_falls_through_to_finite_jacobian() {
5695        let energies: Vec<f64> = (0..101).map(|i| 4.0 + (i as f64) * 0.06).collect();
5696        let mut model = make_energy_scale_u238(energies.clone(), None)
5697            .with_jacobian_method(EnergyScaleJacobianMethod::FiniteDifference);
5698
5699        // l_scale = 0.0 — well below `L_SCALE_EPSILON = 1e-12` — so the
5700        // partial-GAL guard fires and falls through to FD.
5701        //
5702        // **Active code path (regression target):** the test inputs
5703        // are chosen so the new `L_SCALE_EPSILON` guard is what fires,
5704        // *not* the older `t0 + h >= t0_limit` precompute fallthrough.
5705        // Specifically:
5706        //
5707        //   - `min_tof_us = tof_factor * 25.0 / sqrt(max_E ≈ 10.0) ≈ 5.7e2 µs`
5708        //   - `t0 + h = 0.05 + 1e-4 = 0.0501 µs ≪ min_tof * (1 - 1e-12)`
5709        //
5710        // So `partial_gal_t0_column = Some(...)` (not `None`), the
5711        // partial-GAL block at line ~2208 enters, and the L_scale
5712        // branch reaches the new `l_scale.abs() < L_SCALE_EPSILON`
5713        // guard.  Pre-fix, the inner `(tof_i - t0_clamped) / 0.0 =
5714        // ±inf` then `inf * 0 = NaN` would poison the column.  If a
5715        // future refactor changes these inputs, verify that
5716        // `partial_gal_t0_column.is_some()` still holds for this test
5717        // — otherwise the regression target shifts to a different
5718        // code path.
5719        let params = [0.001, 0.05, 1e-13]; // density, t0, l_scale ≈ 0 (< L_SCALE_EPSILON)
5720        let free = vec![0, 1, 2];
5721
5722        // FD2 reference Jacobian.  FD2 computes each column via its
5723        // own per-coordinate FD pair, so it produces well-defined
5724        // finite values at l_scale = 0 (no division by l_scale in the
5725        // FD2 path).
5726        let jac_fd2 = model
5727            .analytical_jacobian(&params, &free, &model.evaluate(&params).unwrap())
5728            .expect("FD2 Jacobian should be available at l_scale = 0");
5729
5730        // Partial-GAL Jacobian — with the #500 guard, L_scale column
5731        // falls through to the same per-coordinate FD path.
5732        model = model.with_jacobian_method(EnergyScaleJacobianMethod::PartialGal);
5733        let jac_pg = model
5734            .analytical_jacobian(&params, &free, &model.evaluate(&params).unwrap())
5735            .expect("partial-GAL Jacobian should be available at l_scale = 0 (fallthrough to FD)");
5736
5737        // Primary regression: every entry finite.  Pre-fix the L_scale
5738        // column would be NaN from the `1 / l_scale` division.
5739        for i in 0..jac_pg.nrows {
5740            for c in 0..jac_pg.ncols {
5741                let v = jac_pg.get(i, c);
5742                assert!(
5743                    v.is_finite(),
5744                    "partial-GAL Jacobian must be finite at l_scale = 0; \
5745                     got non-finite at ({i},{c}) = {v}"
5746                );
5747            }
5748        }
5749
5750        // Bit-equivalent to FD2 across every column — confirms the
5751        // L_scale fallthrough lands on the same FD code path FD2 uses,
5752        // and the density / t0 columns are unchanged by the guard.
5753        for c in 0..jac_pg.ncols {
5754            for i in 0..jac_pg.nrows {
5755                let fd2 = jac_fd2.get(i, c);
5756                let pg = jac_pg.get(i, c);
5757                let abs_err = (fd2 - pg).abs();
5758                let rel_err = abs_err / fd2.abs().max(1e-15);
5759                assert!(
5760                    rel_err < 1e-3 || abs_err < 1e-8,
5761                    "partial-GAL must match FD2 at l_scale = 0; \
5762                     col {c} bin {i}: fd2={fd2:.6e} pg={pg:.6e} rel={rel_err:.2e}"
5763                );
5764            }
5765        }
5766    }
5767
5768    /// In-tree regression for the partial-GAL rank-1 approximation in
5769    /// the presence of a non-trivial resolution kernel.  Issue #499.
5770    ///
5771    /// **Motivation.**  The empirical bound supporting the post-#489
5772    /// default flip to `PartialGal` was measured on real VENUS Hf
5773    /// 120-min KL+per-iso+TZERO 4×4 data: 15 of 16 fitted pixels landed
5774    /// within 0.1·σ_Fisher of the FD2 reference for the L_scale
5775    /// column.  That measurement was made against the production
5776    /// USR/FTS tabulated resolution kernel, which ORNL release policy
5777    /// keeps out of the repository — so it cannot ship as an in-tree
5778    /// fixture.  Without an in-tree analogue, a future refactor could
5779    /// silently regress the rank-1 bound on real workloads.
5780    ///
5781    /// **Synthetic stand-in.**  This test exercises the same code path
5782    /// with a sharp Gaussian "resonance" cross-section convolved by a
5783    /// Gaussian resolution kernel — a deliberately rough stand-in for
5784    /// the SAMMY-format tabulated VENUS USR/FTS kernel.  The Gaussian
5785    /// kernel is *not* a fidelity replacement; it is the simplest
5786    /// non-trivial resolution operator that activates the partial-GAL
5787    /// resolution-bearing branch without introducing a binary fixture.
5788    ///
5789    /// **Tolerance.**  Density and t0 columns retain the same tight
5790    /// bound as the no-resolution test (the resolution operator does
5791    /// not couple into those columns differently).  The L_scale column
5792    /// is checked via relative L₂ norm against the FD2 reference with
5793    /// tolerance `PARTIAL_GAL_REL_L2_TOLERANCE = 1.5e-5`.  On the U-238
5794    /// resonance grid below (kernel sized so it spans several bins and
5795    /// meaningfully broadens the resonance) the measured relative L₂ is
5796    /// `~4.3e-6` — the tolerance gives roughly 3× headroom over the current
5797    /// measurement, tight enough to catch a non-trivial regression
5798    /// of the rank-1 simplification while loose enough to absorb
5799    /// FD-truncation noise.  An upstream pre-check (see below)
5800    /// asserts the kernel itself is non-trivial so a future tweak
5801    /// to grid spacing or kernel parameters cannot silently degrade
5802    /// this back into a vacuous "no-resolution-in-disguise" test
5803    /// (a regression that has occurred before).  The
5804    /// measured relative L₂ surfaces in the assert message if the
5805    /// bound is ever exceeded so future tightening is straightforward.
5806    #[test]
5807    fn partial_gal_with_resolution_matches_fd2() {
5808        use nereids_physics::resolution::{ResolutionFunction, ResolutionParams};
5809
5810        // Tolerance for relative L₂ error on the L_scale column.
5811        // Measured rel L₂ on this synthetic grid is ~9.3e-4; 3e-3
5812        // gives ~3× headroom — tight enough to catch a non-trivial
5813        // regression of the rank-1 simplification under resolution,
5814        // loose enough to absorb FD truncation noise.  See rustdoc
5815        // above for why this bound is conservative rather than the
5816        // tighter empirical 0.1·σ_Fisher seen on real workloads.
5817        const PARTIAL_GAL_REL_L2_TOLERANCE: f64 = 1.5e-5;
5818
5819        // Dense grid over the sharp U-238 resonance (~6.67 eV) so the σ feature
5820        // is well resolved and the resolution kernel meaningfully broadens it.
5821        let energies: Vec<f64> = (0..101).map(|i| 4.0 + (i as f64) * 0.06).collect();
5822
5823        // Gaussian resolution kernel sized to be NON-TRIVIAL on this grid (it
5824        // broadens the U-238 resonance by ~1%, verified by the pre-check below).
5825        // A kernel-too-narrow vacuous-test regression has occurred before;
5826        // the pre-check guards against re-introducing it.
5827        let instrument = Some(Arc::new(InstrumentParams {
5828            resolution: ResolutionFunction::Gaussian(
5829                ResolutionParams::new(25.0, 0.5, 0.005, 0.0).unwrap(),
5830            ),
5831        }));
5832
5833        // Pin FD2 first so the comparison is independent of the
5834        // process-global `NEREIDS_TZERO_JACOBIAN` env var, matching
5835        // the pattern used by `partial_gal_no_resolution_matches_fd2`.
5836        let mut model = make_energy_scale_u238(energies.clone(), instrument)
5837            .with_jacobian_method(EnergyScaleJacobianMethod::FiniteDifference);
5838
5839        let params = [0.001, 0.05, 1.002]; // density, t0, l_scale
5840        let free = vec![0, 1, 2];
5841
5842        // Pre-check: confirm the resolution kernel actually broadens the
5843        // spectrum on this grid, so the comparison is not a vacuous
5844        // no-resolution-in-disguise test.
5845        let model_no_resolution = make_energy_scale_u238(energies.clone(), None)
5846            .with_jacobian_method(EnergyScaleJacobianMethod::FiniteDifference);
5847        let t_no_res = model_no_resolution.evaluate(&params).unwrap();
5848        let t_with_res = model.evaluate(&params).unwrap();
5849        let diff_inf = t_no_res
5850            .iter()
5851            .zip(t_with_res.iter())
5852            .map(|(a, b)| (a - b).abs())
5853            .fold(0.0_f64, f64::max);
5854        let t_inf = t_no_res.iter().map(|x| x.abs()).fold(0.0_f64, f64::max);
5855        assert!(
5856            diff_inf > 1e-3 * t_inf,
5857            "resolution kernel must broaden the spectrum nontrivially \
5858             (got ||T_kernel - T_none||_∞ = {diff_inf:.3e}, ||T_none||_∞ = {t_inf:.3e}, \
5859             ratio = {ratio:.3e}); widen the kernel or sharpen the resonance",
5860            ratio = diff_inf / t_inf.max(1e-30),
5861        );
5862
5863        // FD2 reference Jacobian (explicitly pinned above).
5864        let jac_fd2 = model
5865            .analytical_jacobian(&params, &free, &model.evaluate(&params).unwrap())
5866            .expect("FD2 Jacobian should be available with resolution kernel");
5867
5868        // Flip to partial-GAL.
5869        model = model.with_jacobian_method(EnergyScaleJacobianMethod::PartialGal);
5870        let jac_pg = model
5871            .analytical_jacobian(&params, &free, &model.evaluate(&params).unwrap())
5872            .expect("partial-GAL Jacobian should be available with resolution kernel");
5873
5874        // Density column: tight bound — resolution doesn't change the
5875        // density derivative path.
5876        for i in 0..energies.len() {
5877            let fd2 = jac_fd2.get(i, 0);
5878            let pg = jac_pg.get(i, 0);
5879            let abs_err = (fd2 - pg).abs();
5880            let rel_err = abs_err / fd2.abs().max(1e-15);
5881            assert!(
5882                rel_err < 1e-3 || abs_err < 1e-8,
5883                "density bin {i}: fd2={fd2:.6e} pg={pg:.6e} rel={rel_err:.2e}"
5884            );
5885        }
5886
5887        // t0 column: tight bound — both methods use the same FD pair
5888        // on t0 (partial-GAL just hoists it out of the per-coord loop).
5889        for i in 0..energies.len() {
5890            let fd2 = jac_fd2.get(i, 1);
5891            let pg = jac_pg.get(i, 1);
5892            let abs_err = (fd2 - pg).abs();
5893            let rel_err = abs_err / fd2.abs().max(1e-15);
5894            assert!(
5895                rel_err < 1e-3 || abs_err < 1e-8,
5896                "t0 bin {i}: fd2={fd2:.6e} pg={pg:.6e} rel={rel_err:.2e}"
5897            );
5898        }
5899
5900        // L_scale column: relative L₂ norm bound.  In the presence of
5901        // a non-trivial resolution kernel the rank-1 identity is no
5902        // longer exact; the resolution operator introduces an
5903        // additional (t0, L_scale)-dependence that partial-GAL
5904        // approximates as zero.  The bound captures the residual.
5905        let mut num_sq = 0.0_f64;
5906        let mut den_sq = 0.0_f64;
5907        for i in 0..energies.len() {
5908            let fd2 = jac_fd2.get(i, 2);
5909            let pg = jac_pg.get(i, 2);
5910            let diff = pg - fd2;
5911            num_sq += diff * diff;
5912            den_sq += fd2 * fd2;
5913        }
5914        let rel_l2 = (num_sq / den_sq.max(1e-30)).sqrt();
5915        assert!(
5916            rel_l2 < PARTIAL_GAL_REL_L2_TOLERANCE,
5917            "L_scale rel L₂ = {rel_l2:.4e} exceeds tolerance {tol:.4e}; \
5918             tighten or loosen `PARTIAL_GAL_REL_L2_TOLERANCE` (see rustdoc)",
5919            tol = PARTIAL_GAL_REL_L2_TOLERANCE,
5920        );
5921    }
5922}