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