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