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