use std::cell::{Cell, RefCell};
use std::rc::Rc;
use std::sync::Arc;
use nereids_core::constants::{EV_TO_JOULES, NEUTRON_MASS_KG};
use nereids_endf::resonance::ResonanceData;
use nereids_physics::resolution::{self, ResolutionFunction, ResolutionPlan};
use nereids_physics::surrogate::{ScalarSurrogatePlan, SparseEmpiricalCubaturePlan};
use nereids_physics::transmission::{self, InstrumentParams, SampleParams};
use crate::error::FittingError;
use crate::lm::{FitModel, FlatMatrix};
const L_SCALE_EPSILON: f64 = 1.0e-12;
pub struct PrecomputedTransmissionModel {
pub cross_sections: Arc<Vec<Vec<f64>>>,
pub density_indices: Arc<Vec<usize>>,
pub energies: Option<Arc<Vec<f64>>>,
pub instrument: Option<Arc<InstrumentParams>>,
pub resolution_plan: Option<Arc<ResolutionPlan>>,
pub sparse_cubature_plan: Option<Arc<SparseEmpiricalCubaturePlan>>,
pub sparse_scalar_plan: Option<Arc<ScalarSurrogatePlan>>,
pub work_layout: Option<Arc<transmission::WorkingGridLayout>>,
}
fn density_param_indices(density_indices: &[usize]) -> Vec<usize> {
let mut seen: Vec<usize> = density_indices.to_vec();
seen.sort_unstable();
seen.dedup();
seen
}
fn scalar_eligible(
plan: &ScalarSurrogatePlan,
energies: &[f64],
instrument_resolution: &ResolutionFunction,
resolution_plan: Option<&Arc<ResolutionPlan>>,
sigma_row: &[f64],
n_density_params: usize,
) -> bool {
if n_density_params != 1 {
return false;
}
if plan.len() != energies.len() {
return false;
}
if !matches!(instrument_resolution, ResolutionFunction::Tabulated(_)) {
return false;
}
let plan_grid = plan.target_energies();
for (e_cur, e_plan) in energies.iter().zip(plan_grid) {
if e_cur.to_bits() != e_plan.to_bits() {
return false;
}
}
let Some(model_plan) = resolution_plan else {
return false;
};
if !Arc::ptr_eq(model_plan, plan.source_resolution_plan()) {
return false;
}
if model_plan.target_energies().len() != energies.len() {
return false;
}
for (e_cur, e_res) in energies.iter().zip(model_plan.target_energies()) {
if e_cur.to_bits() != e_res.to_bits() {
return false;
}
}
if nereids_physics::surrogate::fingerprint_f64_slice(sigma_row) != plan.sigma_fingerprint() {
return false;
}
true
}
fn scalar_density_within_box(plan: &ScalarSurrogatePlan, n: f64) -> bool {
let Some(train_max) = plan.density_box() else {
return true;
};
if !n.is_finite() || n < 0.0 {
return false;
}
n <= train_max
}
fn density_within_box(plan: &SparseEmpiricalCubaturePlan, n: &[f64]) -> bool {
let Some(train_max) = plan.density_box() else {
return true;
};
if train_max.len() != n.len() {
return false;
}
const TOLERANCE: f64 = 1.5; for (&n_i, &max_i) in n.iter().zip(train_max) {
if !n_i.is_finite() || n_i < 0.0 {
return false;
}
if n_i > max_i * TOLERANCE {
return false;
}
}
true
}
fn cubature_eligible(
plan: &SparseEmpiricalCubaturePlan,
energies: &[f64],
instrument_resolution: &ResolutionFunction,
resolution_plan: Option<&ResolutionPlan>,
n_density_params: usize,
) -> bool {
if n_density_params < 2 {
return false;
}
if plan.k() != n_density_params {
return false;
}
if plan.len() != energies.len() {
return false;
}
if !matches!(instrument_resolution, ResolutionFunction::Tabulated(_)) {
return false;
}
let cub_grid = plan.target_energies();
for (e_cur, e_cub) in energies.iter().zip(cub_grid) {
if e_cur.to_bits() != e_cub.to_bits() {
return false;
}
}
if let Some(res_plan) = resolution_plan {
if res_plan.target_energies().len() != energies.len() {
return false;
}
let res_grid = res_plan.target_energies();
for (e_cur, e_res) in energies.iter().zip(res_grid) {
if e_cur.to_bits() != e_res.to_bits() {
return false;
}
}
}
true
}
impl PrecomputedTransmissionModel {
fn work_energies(&self) -> Option<&[f64]> {
match (&self.work_layout, &self.energies) {
(Some(layout), _) => Some(layout.energies.as_slice()),
(None, Some(energies)) => Some(energies.as_slice()),
(None, None) => None,
}
}
fn extract_data_points(&self, working: &[f64]) -> Vec<f64> {
match &self.work_layout {
Some(layout) => layout.extract(working),
None => working.to_vec(),
}
}
}
impl FitModel for PrecomputedTransmissionModel {
fn evaluate(&self, params: &[f64]) -> Result<Vec<f64>, FittingError> {
if self.cross_sections.is_empty() {
return Err(FittingError::InvalidConfig(
"PrecomputedTransmissionModel.cross_sections must not be empty".into(),
));
}
let n_e = self.cross_sections[0].len();
if let (Some(cubature), Some(inst), Some(energies)) =
(&self.sparse_cubature_plan, &self.instrument, &self.energies)
{
let params_indices = density_param_indices(&self.density_indices);
if cubature_eligible(
cubature,
energies,
&inst.resolution,
self.resolution_plan.as_deref(),
params_indices.len(),
) {
let n: Vec<f64> = params_indices.iter().map(|&i| params[i]).collect();
if density_within_box(cubature, &n) {
return Ok(cubature.forward(&n));
}
}
}
if let (Some(scalar), Some(inst), Some(energies)) =
(&self.sparse_scalar_plan, &self.instrument, &self.energies)
{
let params_indices = density_param_indices(&self.density_indices);
if self.cross_sections.len() == 1
&& self.density_indices.len() == 1
&& self.density_indices[0] == params_indices[0]
&& scalar_eligible(
scalar,
energies,
&inst.resolution,
self.resolution_plan.as_ref(),
&self.cross_sections[0],
params_indices.len(),
)
{
let n = params[params_indices[0]];
if scalar_density_within_box(scalar, n) {
return Ok(scalar.forward_scalar(n));
}
}
}
let mut neg_opt = vec![0.0f64; n_e];
for (i, xs) in self.cross_sections.iter().enumerate() {
let density = params[self.density_indices[i]];
for (j, &sigma) in xs.iter().enumerate() {
neg_opt[j] -= density * sigma;
}
}
let transmission: Vec<f64> = neg_opt.iter().map(|&d| d.exp()).collect();
if let (Some(inst), Some(work_energies)) = (&self.instrument, self.work_energies()) {
let t_broadened = resolution::apply_resolution_with_plan(
self.resolution_plan.as_deref(),
work_energies,
&transmission,
&inst.resolution,
)
.map_err(|e| FittingError::EvaluationFailed(format!("resolution broadening: {e}")))?;
Ok(self.extract_data_points(&t_broadened))
} else {
Ok(self.extract_data_points(&transmission))
}
}
fn analytical_jacobian(
&self,
params: &[f64],
free_param_indices: &[usize],
y_current: &[f64],
) -> Option<FlatMatrix> {
let n_e = if self.cross_sections.is_empty() {
return None;
} else {
self.cross_sections[0].len()
};
let n_free = free_param_indices.len();
if let (Some(cubature), Some(inst), Some(energies)) =
(&self.sparse_cubature_plan, &self.instrument, &self.energies)
{
let params_indices = density_param_indices(&self.density_indices);
if cubature_eligible(
cubature,
energies,
&inst.resolution,
self.resolution_plan.as_deref(),
params_indices.len(),
) {
let col_map: Option<Vec<usize>> = free_param_indices
.iter()
.map(|&fp| params_indices.iter().position(|&i| i == fp))
.collect();
if let Some(col_map) = col_map {
let n: Vec<f64> = params_indices.iter().map(|&i| params[i]).collect();
if density_within_box(cubature, &n) {
let (_t, jac_flat) = cubature.forward_and_jacobian(&n);
let k = params_indices.len();
let mut jacobian = FlatMatrix::zeros(n_e, n_free);
for (col, &ell) in col_map.iter().enumerate() {
for i in 0..n_e {
*jacobian.get_mut(i, col) = jac_flat[i * k + ell];
}
}
return Some(jacobian);
}
}
}
}
if let (Some(scalar), Some(inst), Some(energies)) =
(&self.sparse_scalar_plan, &self.instrument, &self.energies)
{
let params_indices = density_param_indices(&self.density_indices);
if self.cross_sections.len() == 1
&& self.density_indices.len() == 1
&& self.density_indices[0] == params_indices[0]
&& scalar_eligible(
scalar,
energies,
&inst.resolution,
self.resolution_plan.as_ref(),
&self.cross_sections[0],
params_indices.len(),
)
&& free_param_indices.len() == 1
&& free_param_indices[0] == params_indices[0]
{
let n = params[params_indices[0]];
if scalar_density_within_box(scalar, n) {
let (_t, dt) = scalar.forward_and_derivative_scalar(n);
let mut jacobian = FlatMatrix::zeros(n_e, 1);
for (i, &v) in dt.iter().enumerate() {
*jacobian.get_mut(i, 0) = v;
}
return Some(jacobian);
}
}
}
let fp_xs_sums: Vec<Vec<f64>> = free_param_indices
.iter()
.map(|&fp_idx| {
let mut sum = vec![0.0f64; n_e];
for (iso, &di) in self.density_indices.iter().enumerate() {
if di == fp_idx {
for (j, &sigma) in self.cross_sections[iso].iter().enumerate() {
sum[j] += sigma;
}
}
}
sum
})
.collect();
let n_data = y_current.len();
if let (Some(inst), Some(work_energies)) = (&self.instrument, self.work_energies()) {
let mut neg_opt = vec![0.0f64; n_e];
for (i, xs) in self.cross_sections.iter().enumerate() {
let density = params[self.density_indices[i]];
for (j, &sigma) in xs.iter().enumerate() {
neg_opt[j] -= density * sigma;
}
}
let t_unresolved: Vec<f64> = neg_opt.iter().map(|&d| d.exp()).collect();
let mut jacobian = FlatMatrix::zeros(n_data, n_free);
for (col, xs_sum) in fp_xs_sums.iter().enumerate() {
let inner_deriv: Vec<f64> =
(0..n_e).map(|i| -xs_sum[i] * t_unresolved[i]).collect();
let resolved_deriv = resolution::apply_resolution_with_plan(
self.resolution_plan.as_deref(),
work_energies,
&inner_deriv,
&inst.resolution,
)
.ok()?;
let resolved_deriv = self.extract_data_points(&resolved_deriv);
for (i, &val) in resolved_deriv.iter().enumerate() {
*jacobian.get_mut(i, col) = val;
}
}
Some(jacobian)
} else {
let mut jacobian = FlatMatrix::zeros(n_data, n_free);
for i in 0..n_data {
for (j, xs_sum) in fp_xs_sums.iter().enumerate() {
*jacobian.get_mut(i, j) = -xs_sum[i] * y_current[i];
}
}
Some(jacobian)
}
}
}
pub struct TransmissionFitModel {
energies: Vec<f64>,
resonance_data: Vec<ResonanceData>,
temperature_k: f64,
instrument: Option<Arc<InstrumentParams>>,
density_indices: Vec<usize>,
density_ratios: Vec<f64>,
temperature_index: Option<usize>,
base_xs: Option<Arc<Vec<Vec<f64>>>>,
cached_broadened_xs: RefCell<Option<Rc<Vec<Vec<f64>>>>>,
cached_dxs_dt: RefCell<Option<Rc<Vec<Vec<f64>>>>>,
cached_work_layout: RefCell<Option<Rc<transmission::WorkingGridLayout>>>,
cached_temperature: Cell<f64>,
resolution_plan: Option<Arc<ResolutionPlan>>,
sparse_cubature_plan: Option<Arc<SparseEmpiricalCubaturePlan>>,
sparse_scalar_plan: Option<Arc<ScalarSurrogatePlan>>,
}
impl TransmissionFitModel {
pub fn new(
energies: Vec<f64>,
resonance_data: Vec<ResonanceData>,
temperature_k: f64,
instrument: Option<Arc<InstrumentParams>>,
density_mapping: (Vec<usize>, Vec<f64>),
temperature_index: Option<usize>,
external_base_xs: Option<Arc<Vec<Vec<f64>>>>,
) -> Result<Self, FittingError> {
let (density_indices, density_ratios) = density_mapping;
if density_indices.len() != resonance_data.len() {
return Err(FittingError::InvalidConfig(format!(
"density_indices has {} entries but resonance_data has {}",
density_indices.len(),
resonance_data.len(),
)));
}
if density_ratios.len() != resonance_data.len() {
return Err(FittingError::InvalidConfig(format!(
"density_ratios has {} entries but resonance_data has {}",
density_ratios.len(),
resonance_data.len(),
)));
}
if let Some(ti) = temperature_index
&& density_indices.contains(&ti)
{
return Err(FittingError::InvalidConfig(
"temperature_index must not overlap with density_indices".into(),
));
}
if let Some(ref xs) = external_base_xs {
if xs.len() != resonance_data.len() {
return Err(FittingError::InvalidConfig(format!(
"external_base_xs has {} isotopes but resonance_data has {}",
xs.len(),
resonance_data.len(),
)));
}
for (i, row) in xs.iter().enumerate() {
if row.len() != energies.len() {
return Err(FittingError::InvalidConfig(format!(
"external_base_xs[{i}] has {} energies but expected {}",
row.len(),
energies.len(),
)));
}
}
}
let base_xs = match external_base_xs {
Some(xs) => Some(xs),
None if temperature_index.is_some() => Some(Arc::new(
transmission::unbroadened_cross_sections(&energies, &resonance_data, None)
.map_err(|e| {
FittingError::InvalidConfig(format!(
"failed to compute unbroadened cross-sections: {e}"
))
})?,
)),
None => None,
};
Ok(Self {
energies,
resonance_data,
temperature_k,
instrument,
density_indices,
density_ratios,
temperature_index,
base_xs,
cached_broadened_xs: RefCell::new(None),
cached_dxs_dt: RefCell::new(None),
cached_work_layout: RefCell::new(None),
cached_temperature: Cell::new(f64::NAN),
resolution_plan: None,
sparse_cubature_plan: None,
sparse_scalar_plan: None,
})
}
#[must_use]
pub fn with_resolution_plan(mut self, plan: Option<Arc<ResolutionPlan>>) -> Self {
self.resolution_plan = plan;
self
}
#[must_use]
pub fn with_sparse_cubature_plan(
mut self,
plan: Option<Arc<SparseEmpiricalCubaturePlan>>,
) -> Self {
self.sparse_cubature_plan = plan;
self
}
#[must_use]
pub fn with_sparse_scalar_plan(mut self, plan: Option<Arc<ScalarSurrogatePlan>>) -> Self {
self.sparse_scalar_plan = plan;
self
}
}
impl FitModel for TransmissionFitModel {
fn evaluate(&self, params: &[f64]) -> Result<Vec<f64>, FittingError> {
debug_assert!(
self.density_indices.iter().all(|&i| i < params.len()),
"density_indices out of bounds for params (len={})",
params.len(),
);
debug_assert!(
self.temperature_index.is_none_or(|i| i < params.len()),
"temperature_index out of bounds for params (len={})",
params.len(),
);
if let (Some(cubature), Some(inst)) = (&self.sparse_cubature_plan, &self.instrument)
&& self.temperature_index.is_none()
{
let params_indices = density_param_indices(&self.density_indices);
if cubature_eligible(
cubature,
&self.energies,
&inst.resolution,
self.resolution_plan.as_deref(),
params_indices.len(),
) {
let n: Vec<f64> = params_indices.iter().map(|&i| params[i]).collect();
if density_within_box(cubature, &n) {
return Ok(cubature.forward(&n));
}
}
}
let temperature_k = match self.temperature_index {
Some(idx) => params[idx],
None => self.temperature_k,
};
if let Some(ref base_xs) = self.base_xs {
if !temperature_k.is_finite() || temperature_k < 0.0 {
return Err(FittingError::EvaluationFailed(format!(
"Invalid temperature: {temperature_k} K (must be finite and non-negative)"
)));
}
let (broadened_xs, layout) = if (temperature_k - self.cached_temperature.get()).abs()
< 1e-15
&& self.cached_broadened_xs.borrow().is_some()
{
(
Rc::clone(self.cached_broadened_xs.borrow().as_ref().unwrap()),
Rc::clone(self.cached_work_layout.borrow().as_ref().unwrap()),
)
} else {
let working = transmission::broadened_cross_sections_from_base_on_working_grid(
&self.energies,
base_xs,
&self.resonance_data,
temperature_k,
self.instrument.as_deref(),
)
.map_err(|e| FittingError::EvaluationFailed(e.to_string()))?;
let xs = Rc::new(working.sigma);
let layout = Rc::new(working.layout);
*self.cached_broadened_xs.borrow_mut() = Some(Rc::clone(&xs));
*self.cached_work_layout.borrow_mut() = Some(Rc::clone(&layout));
*self.cached_dxs_dt.borrow_mut() = None;
self.cached_temperature.set(temperature_k);
(xs, layout)
};
let work_len = layout.energies.len();
let mut neg_opt = vec![0.0f64; work_len];
for (i, xs) in broadened_xs.iter().enumerate() {
let density = params[self.density_indices[i]];
let ratio = self.density_ratios[i];
for (j, &sigma) in xs.iter().enumerate() {
neg_opt[j] -= density * ratio * sigma;
}
}
let transmission: Vec<f64> = neg_opt.iter().map(|&d| d.exp()).collect();
if let Some(ref inst) = self.instrument {
let t_broadened = resolution::apply_resolution_with_plan(
self.resolution_plan.as_deref(),
&layout.energies,
&transmission,
&inst.resolution,
)
.map_err(|e| {
FittingError::EvaluationFailed(format!("resolution broadening: {e}"))
})?;
Ok(layout.extract(&t_broadened))
} else {
Ok(layout.extract(&transmission))
}
} else {
let isotopes: Vec<(ResonanceData, f64)> = self
.resonance_data
.iter()
.enumerate()
.map(|(i, rd)| {
(
rd.clone(),
params[self.density_indices[i]] * self.density_ratios[i],
)
})
.collect();
let sample = SampleParams::new(temperature_k, isotopes)
.map_err(|e| FittingError::EvaluationFailed(e.to_string()))?;
transmission::forward_model(&self.energies, &sample, self.instrument.as_deref())
.map_err(|e| FittingError::EvaluationFailed(e.to_string()))
}
}
fn analytical_jacobian(
&self,
params: &[f64],
free_param_indices: &[usize],
y_current: &[f64],
) -> Option<FlatMatrix> {
if let (Some(cubature), Some(inst)) = (&self.sparse_cubature_plan, &self.instrument)
&& self.temperature_index.is_none()
{
let params_indices = density_param_indices(&self.density_indices);
if cubature_eligible(
cubature,
&self.energies,
&inst.resolution,
self.resolution_plan.as_deref(),
params_indices.len(),
) {
let col_map: Option<Vec<usize>> = free_param_indices
.iter()
.map(|&fp| params_indices.iter().position(|&i| i == fp))
.collect();
if let Some(col_map) = col_map {
let n: Vec<f64> = params_indices.iter().map(|&i| params[i]).collect();
if density_within_box(cubature, &n) {
let (_t, jac_flat) = cubature.forward_and_jacobian(&n);
let k = params_indices.len();
let n_e = self.energies.len();
let mut jacobian = FlatMatrix::zeros(n_e, free_param_indices.len());
for (col, &ell) in col_map.iter().enumerate() {
for i in 0..n_e {
*jacobian.get_mut(i, col) = jac_flat[i * k + ell];
}
}
return Some(jacobian);
}
}
}
}
let _base_xs_guard = self.base_xs.as_ref()?;
let cached_xs = self.cached_broadened_xs.borrow();
let broadened_xs = cached_xs.as_ref()?;
let cached_layout = self.cached_work_layout.borrow();
let layout = cached_layout.as_ref()?;
if let Some(ti) = self.temperature_index {
let param_temp = params[ti];
if (param_temp - self.cached_temperature.get()).abs() > 1e-15 {
return None;
}
}
let n_e = y_current.len();
let work_len = layout.energies.len();
let n_free = free_param_indices.len();
let mut jacobian = FlatMatrix::zeros(n_e, n_free);
let temp_col = self
.temperature_index
.and_then(|ti| free_param_indices.iter().position(|&fp| fp == ti));
let t_unresolved: Option<Vec<f64>> = if self.instrument.is_some() {
let mut neg_opt = vec![0.0f64; work_len];
for (iso, xs) in broadened_xs.iter().enumerate() {
let density = params[self.density_indices[iso]];
let ratio = self.density_ratios[iso];
for (j, &sigma) in xs.iter().enumerate() {
neg_opt[j] -= density * ratio * sigma;
}
}
Some(neg_opt.iter().map(|&d| d.exp()).collect())
} else {
None
};
let t_for_deriv: &[f64] = t_unresolved.as_deref().unwrap_or(y_current);
for (col, &fp_idx) in free_param_indices.iter().enumerate() {
if temp_col == Some(col) {
continue;
}
let mut sigma_sum = vec![0.0f64; work_len];
for (iso, &di) in self.density_indices.iter().enumerate() {
if di == fp_idx {
let ratio = self.density_ratios[iso];
for (j, &sigma) in broadened_xs[iso].iter().enumerate() {
sigma_sum[j] += ratio * sigma;
}
}
}
let inner: Vec<f64> = (0..work_len)
.map(|i| -sigma_sum[i] * t_for_deriv[i])
.collect();
if let Some(ref inst) = self.instrument {
let resolved = resolution::apply_resolution_with_plan(
self.resolution_plan.as_deref(),
&layout.energies,
&inner,
&inst.resolution,
)
.ok()?;
let resolved = layout.extract(&resolved);
for (i, &val) in resolved.iter().enumerate() {
*jacobian.get_mut(i, col) = val;
}
} else {
for (i, &val) in inner.iter().enumerate() {
*jacobian.get_mut(i, col) = val;
}
}
}
if let Some(col) = temp_col {
{
let needs_compute = self.cached_dxs_dt.borrow().as_ref().is_none();
if needs_compute {
let base_xs = self.base_xs.as_ref()?;
let temperature_k = self.cached_temperature.get();
let working =
transmission::broadened_cross_sections_with_analytical_derivative_from_base_on_working_grid(
&self.energies,
base_xs,
&self.resonance_data,
temperature_k,
self.instrument.as_deref(),
)
.ok()?;
*self.cached_dxs_dt.borrow_mut() = Some(Rc::new(working.dsigma_dt));
}
}
let cached_dxs = self.cached_dxs_dt.borrow();
let dxs_dt = cached_dxs.as_ref()?;
let inner: Vec<f64> = (0..work_len)
.map(|i| {
let mut sum_n_dsigma = 0.0f64;
for (iso, dxs) in dxs_dt.iter().enumerate() {
let density = params[self.density_indices[iso]];
let ratio = self.density_ratios[iso];
sum_n_dsigma += density * ratio * dxs[i];
}
-t_for_deriv[i] * sum_n_dsigma
})
.collect();
if let Some(ref inst) = self.instrument {
let resolved = resolution::apply_resolution_with_plan(
self.resolution_plan.as_deref(),
&layout.energies,
&inner,
&inst.resolution,
)
.ok()?;
let resolved = layout.extract(&resolved);
for (i, &val) in resolved.iter().enumerate() {
*jacobian.get_mut(i, col) = val;
}
} else {
for (i, &val) in inner.iter().enumerate() {
*jacobian.get_mut(i, col) = val;
}
}
}
Some(jacobian)
}
}
pub struct NormalizedTransmissionModel<M: FitModel> {
inner: M,
sqrt_energies: Vec<f64>,
inv_sqrt_energies: Vec<f64>,
anorm_index: usize,
back_a_index: usize,
back_b_index: usize,
back_c_index: usize,
back_d_index: Option<usize>,
back_f_index: Option<usize>,
}
impl<M: FitModel> NormalizedTransmissionModel<M> {
pub fn new(
inner: M,
energies: &[f64],
anorm_index: usize,
back_a_index: usize,
back_b_index: usize,
back_c_index: usize,
) -> Self {
let sqrt_energies: Vec<f64> = energies.iter().map(|&e| e.sqrt()).collect();
let inv_sqrt_energies: Vec<f64> = sqrt_energies
.iter()
.map(|&se| if se > 0.0 { 1.0 / se } else { 0.0 })
.collect();
Self {
inner,
sqrt_energies,
inv_sqrt_energies,
anorm_index,
back_a_index,
back_b_index,
back_c_index,
back_d_index: None,
back_f_index: None,
}
}
#[allow(clippy::too_many_arguments)]
pub fn new_with_exponential(
inner: M,
energies: &[f64],
anorm_index: usize,
back_a_index: usize,
back_b_index: usize,
back_c_index: usize,
back_d_index: usize,
back_f_index: usize,
) -> Self {
let sqrt_energies: Vec<f64> = energies.iter().map(|&e| e.sqrt()).collect();
let inv_sqrt_energies: Vec<f64> = sqrt_energies
.iter()
.map(|&se| if se > 0.0 { 1.0 / se } else { 0.0 })
.collect();
Self {
inner,
sqrt_energies,
inv_sqrt_energies,
anorm_index,
back_a_index,
back_b_index,
back_c_index,
back_d_index: Some(back_d_index),
back_f_index: Some(back_f_index),
}
}
}
impl<M: FitModel> FitModel for NormalizedTransmissionModel<M> {
fn evaluate(&self, params: &[f64]) -> Result<Vec<f64>, FittingError> {
let t_inner = self.inner.evaluate(params)?;
let anorm = params[self.anorm_index];
let back_a = params[self.back_a_index];
let back_b = params[self.back_b_index];
let back_c = params[self.back_c_index];
let (back_d, back_f) = match (self.back_d_index, self.back_f_index) {
(Some(di), Some(fi)) => (params[di], params[fi]),
_ => (0.0, 0.0),
};
let has_exp = self.back_d_index.is_some();
let result: Vec<f64> = t_inner
.iter()
.enumerate()
.map(|(i, &t)| {
let mut val = anorm * t
+ back_a
+ back_b * self.inv_sqrt_energies[i]
+ back_c * self.sqrt_energies[i];
if has_exp {
val += back_d * (-back_f * self.inv_sqrt_energies[i]).exp();
}
val
})
.collect();
Ok(result)
}
fn analytical_jacobian(
&self,
params: &[f64],
free_param_indices: &[usize],
y_current: &[f64],
) -> Option<FlatMatrix> {
let n_e = y_current.len();
let n_free = free_param_indices.len();
let t_inner = self.inner.evaluate(params).ok()?;
let anorm = params[self.anorm_index];
let mut bg_indices_set = vec![
self.anorm_index,
self.back_a_index,
self.back_b_index,
self.back_c_index,
];
if let Some(di) = self.back_d_index {
bg_indices_set.push(di);
}
if let Some(fi) = self.back_f_index {
bg_indices_set.push(fi);
}
let inner_free_indices: Vec<usize> = free_param_indices
.iter()
.copied()
.filter(|idx| !bg_indices_set.contains(idx))
.collect();
let inner_jac = if !inner_free_indices.is_empty() {
self.inner
.analytical_jacobian(params, &inner_free_indices, &t_inner)
} else {
None
};
let exp_terms: Vec<f64> =
if let (Some(di), Some(fi)) = (self.back_d_index, self.back_f_index) {
let _back_d = params[di];
let back_f = params[fi];
self.inv_sqrt_energies
.iter()
.map(|&inv_se| (-back_f * inv_se).exp())
.collect()
} else {
vec![]
};
let mut jacobian = FlatMatrix::zeros(n_e, n_free);
let mut inner_col_map = std::collections::HashMap::new();
for (col, &idx) in inner_free_indices.iter().enumerate() {
inner_col_map.insert(idx, col);
}
for (col, &fp_idx) in free_param_indices.iter().enumerate() {
if fp_idx == self.anorm_index {
for (i, &ti) in t_inner.iter().enumerate() {
*jacobian.get_mut(i, col) = ti;
}
} else if fp_idx == self.back_a_index {
for i in 0..n_e {
*jacobian.get_mut(i, col) = 1.0;
}
} else if fp_idx == self.back_b_index {
for (i, &inv_se) in self.inv_sqrt_energies.iter().enumerate() {
*jacobian.get_mut(i, col) = inv_se;
}
} else if fp_idx == self.back_c_index {
for (i, &se) in self.sqrt_energies.iter().enumerate() {
*jacobian.get_mut(i, col) = se;
}
} else if self.back_d_index == Some(fp_idx) {
for (i, &et) in exp_terms.iter().enumerate() {
*jacobian.get_mut(i, col) = et;
}
} else if self.back_f_index == Some(fp_idx) {
let back_d = params[self.back_d_index.unwrap()];
for (i, (&et, &inv_se)) in exp_terms
.iter()
.zip(self.inv_sqrt_energies.iter())
.enumerate()
{
*jacobian.get_mut(i, col) = -back_d * et * inv_se;
}
} else if let Some(&inner_col) = inner_col_map.get(&fp_idx) {
if let Some(ref jac) = inner_jac {
for i in 0..n_e {
*jacobian.get_mut(i, col) = anorm * jac.get(i, inner_col);
}
} else {
return None;
}
} else {
return None;
}
}
Some(jacobian)
}
}
pub struct EnergyScaleTransmissionModel {
resonance_data: Arc<Vec<ResonanceData>>,
density_indices: Arc<Vec<usize>>,
density_ratios: Arc<Vec<f64>>,
temperature_k: f64,
nominal_energies: Vec<f64>,
flight_path_m: f64,
tof_factor: f64,
t0_index: usize,
l_scale_index: usize,
instrument: Option<Arc<transmission::InstrumentParams>>,
cached_plans: RefCell<CachedPlanRing>,
cached_work_xs: RefCell<CachedWorkXs>,
jacobian_method: EnergyScaleJacobianMethod,
}
type CachedWorkXs = Option<((u64, u64), Rc<transmission::WorkingGridXs>)>;
#[derive(Debug, Clone)]
struct CachedPlanEntry {
key: (u64, u64),
plan: Arc<ResolutionPlan>,
}
#[derive(Debug, Default)]
struct CachedPlanRing {
slots: [Option<CachedPlanEntry>; 2],
}
impl CachedPlanRing {
fn lookup(&self, key: (u64, u64)) -> Option<Arc<ResolutionPlan>> {
for slot in &self.slots {
if let Some(entry) = slot
&& entry.key == key
{
return Some(Arc::clone(&entry.plan));
}
}
None
}
fn insert(&mut self, entry: CachedPlanEntry) {
self.slots[1] = self.slots[0].take();
self.slots[0] = Some(entry);
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum EnergyScaleJacobianMethod {
FiniteDifference,
PartialGal,
}
impl EnergyScaleJacobianMethod {
fn from_env() -> Self {
use std::sync::OnceLock;
static CACHED: OnceLock<EnergyScaleJacobianMethod> = OnceLock::new();
*CACHED.get_or_init(Self::resolve_env_uncached)
}
fn resolve_env_uncached() -> Self {
let Ok(v) = std::env::var("NEREIDS_TZERO_JACOBIAN") else {
return Self::PartialGal;
};
if v.eq_ignore_ascii_case("fd2")
|| v.eq_ignore_ascii_case("finite-difference")
|| v.eq_ignore_ascii_case("finite_difference")
{
Self::FiniteDifference
} else if v.eq_ignore_ascii_case("partial-gal") || v.eq_ignore_ascii_case("partial_gal") {
Self::PartialGal
} else {
eprintln!(
"warning: NEREIDS_TZERO_JACOBIAN=\"{v}\" is not a recognized \
Jacobian method (\"chain\" / \"frozen-r\" were removed in #608); \
using the default \"partial-gal\". Valid values: \"fd2\", \
\"partial-gal\"."
);
Self::PartialGal
}
}
}
impl EnergyScaleTransmissionModel {
#[allow(clippy::too_many_arguments)]
pub fn new(
resonance_data: Arc<Vec<ResonanceData>>,
density_indices: Arc<Vec<usize>>,
density_ratios: Arc<Vec<f64>>,
temperature_k: f64,
nominal_energies: Vec<f64>,
flight_path_m: f64,
t0_index: usize,
l_scale_index: usize,
instrument: Option<Arc<transmission::InstrumentParams>>,
) -> Self {
let tof_factor = (0.5 * NEUTRON_MASS_KG / EV_TO_JOULES).sqrt() * 1.0e6;
Self {
resonance_data,
density_indices,
density_ratios,
temperature_k,
nominal_energies,
flight_path_m,
tof_factor,
t0_index,
l_scale_index,
instrument,
cached_plans: RefCell::new(CachedPlanRing::default()),
cached_work_xs: RefCell::new(None),
jacobian_method: EnergyScaleJacobianMethod::from_env(),
}
}
#[must_use]
pub fn with_jacobian_method(mut self, method: EnergyScaleJacobianMethod) -> Self {
self.jacobian_method = method;
self
}
fn cached_resolution_plan(
&self,
t0_us: f64,
l_scale: f64,
working_energies: &[f64],
) -> Option<Arc<ResolutionPlan>> {
let inst = self.instrument.as_ref()?;
if !matches!(
&inst.resolution,
nereids_physics::resolution::ResolutionFunction::Tabulated(_)
) {
return None;
}
let key = (t0_us.to_bits(), l_scale.to_bits());
if let Some(plan) = self.cached_plans.borrow().lookup(key) {
return Some(plan);
}
let plan = resolution::build_resolution_plan(working_energies, &inst.resolution)
.ok()
.flatten()?;
let arc = Arc::new(plan);
self.cached_plans.borrow_mut().insert(CachedPlanEntry {
key,
plan: Arc::clone(&arc),
});
Some(arc)
}
fn corrected_energies(&self, t0_us: f64, l_scale: f64) -> Vec<f64> {
if self.nominal_energies.is_empty() {
return Vec::new();
}
let l_eff = self.flight_path_m * l_scale;
let min_tof = self
.nominal_energies
.iter()
.copied()
.fold(f64::INFINITY, |acc, e| {
acc.min(self.tof_factor * self.flight_path_m / e.sqrt())
});
let t0_limit = min_tof * (1.0 - 1.0e-12);
let t0_clamped = t0_us.min(t0_limit);
self.nominal_energies
.iter()
.map(|&e_nom| {
let tof = self.tof_factor * self.flight_path_m / e_nom.sqrt();
let tof_corr = tof - t0_clamped;
(self.tof_factor * l_eff / tof_corr).powi(2)
})
.collect()
}
fn working_xs(&self, e_corr: &[f64]) -> Result<transmission::WorkingGridXs, FittingError> {
if let Some(&bad) = e_corr.iter().find(|&&e| !e.is_finite() || e <= 0.0) {
return Err(FittingError::EvaluationFailed(format!(
"energy-scale corrected energy is non-positive or non-finite ({bad}); \
t0 / L_scale give a degenerate calibration"
)));
}
transmission::broadened_cross_sections_on_working_grid(
e_corr,
&self.resonance_data,
self.temperature_k,
self.instrument.as_deref(),
None,
)
.map_err(|e| FittingError::EvaluationFailed(e.to_string()))
}
fn working_xs_for(
&self,
params: &[f64],
e_corr: &[f64],
) -> Result<Rc<transmission::WorkingGridXs>, FittingError> {
let key = (
params[self.t0_index].to_bits(),
params[self.l_scale_index].to_bits(),
);
let hit = self
.cached_work_xs
.borrow()
.as_ref()
.and_then(|(k, xs)| (*k == key).then(|| Rc::clone(xs)));
if let Some(xs) = hit {
return Ok(xs);
}
let xs = Rc::new(self.working_xs(e_corr)?);
*self.cached_work_xs.borrow_mut() = Some((key, Rc::clone(&xs)));
Ok(xs)
}
fn evaluate_at_with_cache(
&self,
params: &[f64],
e_corr: &[f64],
use_plan_cache: bool,
) -> Result<Vec<f64>, FittingError> {
let work = self.working_xs_for(params, e_corr)?;
let work_e = &work.layout.energies;
let mut neg_opt = vec![0.0f64; work_e.len()];
for (iso, xs) in work.sigma.iter().enumerate() {
let density = params[self.density_indices[iso]];
let ratio = self.density_ratios[iso];
for (j, &sigma) in xs.iter().enumerate() {
neg_opt[j] -= density * ratio * sigma;
}
}
let t_unbroadened: Vec<f64> = neg_opt.iter().map(|&d| d.exp()).collect();
let Some(inst) = self.instrument.as_ref() else {
return Ok(work.layout.extract(&t_unbroadened));
};
let plan = if use_plan_cache {
let t0 = params[self.t0_index];
let l_scale = params[self.l_scale_index];
self.cached_resolution_plan(t0, l_scale, work_e)
} else {
None
};
let t_broadened = resolution::apply_resolution_with_plan(
plan.as_deref(),
work_e,
&t_unbroadened,
&inst.resolution,
)
.map_err(|e| FittingError::EvaluationFailed(format!("resolution broadening: {e}")))?;
Ok(work.layout.extract(&t_broadened))
}
}
impl FitModel for EnergyScaleTransmissionModel {
fn evaluate(&self, params: &[f64]) -> Result<Vec<f64>, FittingError> {
let t0 = params[self.t0_index];
let l_scale = params[self.l_scale_index];
let e_corr = self.corrected_energies(t0, l_scale);
self.evaluate_at_with_cache(params, &e_corr, true)
}
fn analytical_jacobian(
&self,
params: &[f64],
free_param_indices: &[usize],
_y_current: &[f64],
) -> Option<FlatMatrix> {
let n_e = self.nominal_energies.len();
let n_free = free_param_indices.len();
let mut jacobian = FlatMatrix::zeros(n_e, n_free);
let t0 = params[self.t0_index];
let l_scale = params[self.l_scale_index];
let e_corr = self.corrected_energies(t0, l_scale);
let energy_scale_method = self.jacobian_method;
let t0_free_pos = free_param_indices
.iter()
.position(|&idx| idx == self.t0_index);
let l_scale_free_pos = free_param_indices
.iter()
.position(|&idx| idx == self.l_scale_index);
let partial_gal_t0_column = if energy_scale_method == EnergyScaleJacobianMethod::PartialGal
&& t0_free_pos.is_some()
&& l_scale_free_pos.is_some()
{
let h = 1e-4;
let min_tof_us = self
.nominal_energies
.iter()
.map(|&e| self.tof_factor * self.flight_path_m / e.sqrt())
.fold(f64::INFINITY, f64::min);
let t0_limit = min_tof_us * (1.0 - 1.0e-12);
if t0 + h >= t0_limit {
None
} else {
let mut p_plus = params.to_vec();
let mut p_minus = params.to_vec();
p_plus[self.t0_index] += h;
p_minus[self.t0_index] -= h;
let e_corr_plus =
self.corrected_energies(p_plus[self.t0_index], p_plus[self.l_scale_index]);
let e_corr_minus =
self.corrected_energies(p_minus[self.t0_index], p_minus[self.l_scale_index]);
let y_plus = match self.evaluate_at_with_cache(&p_plus, &e_corr_plus, false) {
Ok(v) => v,
Err(_) => return None,
};
let y_minus = match self.evaluate_at_with_cache(&p_minus, &e_corr_minus, false) {
Ok(v) => v,
Err(_) => return None,
};
let mut col = vec![0.0f64; n_e];
for i in 0..n_e {
let a = y_plus[i];
let b = y_minus[i];
if a.is_finite() && b.is_finite() {
col[i] = (a - b) / (2.0 * h);
}
}
Some(col)
}
} else {
None
};
let work = match self.working_xs_for(params, &e_corr) {
Ok(w) => w,
Err(_) => return None,
};
let work_layout = &work.layout;
let work_e = &work_layout.energies;
let mut neg_opt = vec![0.0f64; work_e.len()];
for (iso, xs) in work.sigma.iter().enumerate() {
let density = params[self.density_indices[iso]];
let ratio = self.density_ratios[iso];
for (j, &sigma) in xs.iter().enumerate() {
neg_opt[j] -= density * ratio * sigma;
}
}
let t_unresolved: Vec<f64> = neg_opt.iter().map(|&d| d.exp()).collect();
let density_plan = self.cached_resolution_plan(t0, l_scale, work_e);
for (col, &fp_idx) in free_param_indices.iter().enumerate() {
if fp_idx == self.t0_index || fp_idx == self.l_scale_index {
if let Some(partial_t0_col) = &partial_gal_t0_column {
if fp_idx == self.t0_index {
for (i, &val) in partial_t0_col.iter().enumerate() {
*jacobian.get_mut(i, col) = val;
}
continue;
}
if fp_idx == self.l_scale_index {
let l_scale = params[self.l_scale_index];
if l_scale.abs() >= L_SCALE_EPSILON {
let t0 = params[self.t0_index];
let min_tof_us = self
.nominal_energies
.iter()
.map(|&e| self.tof_factor * self.flight_path_m / e.sqrt())
.fold(f64::INFINITY, f64::min);
let t0_clamped = t0.min(min_tof_us * (1.0 - 1.0e-12));
for (i, &e_nom) in self.nominal_energies.iter().enumerate() {
let tof_i = self.tof_factor * self.flight_path_m / e_nom.sqrt();
let scale = (tof_i - t0_clamped) / l_scale;
*jacobian.get_mut(i, col) = scale * partial_t0_col[i];
}
continue;
}
}
}
let h = if fp_idx == self.t0_index { 1e-4 } else { 1e-7 };
let mut p_plus = params.to_vec();
let mut p_minus = params.to_vec();
p_plus[fp_idx] += h;
p_minus[fp_idx] -= h;
let t0_plus = p_plus[self.t0_index];
let l_plus = p_plus[self.l_scale_index];
let t0_minus = p_minus[self.t0_index];
let l_minus = p_minus[self.l_scale_index];
let e_corr_plus = self.corrected_energies(t0_plus, l_plus);
let e_corr_minus = self.corrected_energies(t0_minus, l_minus);
let y_plus = match self.evaluate_at_with_cache(&p_plus, &e_corr_plus, false) {
Ok(v) => v,
Err(_) => return None,
};
let y_minus = match self.evaluate_at_with_cache(&p_minus, &e_corr_minus, false) {
Ok(v) => v,
Err(_) => return None,
};
for i in 0..n_e {
let a = y_plus[i];
let b = y_minus[i];
if a.is_finite() && b.is_finite() {
*jacobian.get_mut(i, col) = (a - b) / (2.0 * h);
}
}
} else {
let mut sigma_sum = vec![0.0f64; work_e.len()];
for (iso, &di) in self.density_indices.iter().enumerate() {
if di == fp_idx {
let ratio = self.density_ratios[iso];
for (j, &sigma) in work.sigma[iso].iter().enumerate() {
sigma_sum[j] += ratio * sigma;
}
}
}
let inner_deriv: Vec<f64> = (0..work_e.len())
.map(|i| -sigma_sum[i] * t_unresolved[i])
.collect();
if let Some(inst) = &self.instrument {
let resolved_deriv = match resolution::apply_resolution_with_plan(
density_plan.as_deref(),
work_e,
&inner_deriv,
&inst.resolution,
) {
Ok(v) => v,
Err(_) => return None,
};
let resolved_deriv = work_layout.extract(&resolved_deriv);
for (i, &val) in resolved_deriv.iter().enumerate() {
*jacobian.get_mut(i, col) = val;
}
} else {
for (i, &val) in inner_deriv.iter().enumerate() {
*jacobian.get_mut(i, col) = val;
}
}
}
}
Some(jacobian)
}
}
use crate::forward_model::ForwardModel;
impl ForwardModel for PrecomputedTransmissionModel {
fn predict(&self, params: &[f64]) -> Result<Vec<f64>, FittingError> {
self.evaluate(params)
}
fn jacobian(
&self,
params: &[f64],
free_param_indices: &[usize],
y_current: &[f64],
) -> Option<Vec<Vec<f64>>> {
let fm = self.analytical_jacobian(params, free_param_indices, y_current)?;
Some(flat_matrix_to_vecs(&fm, free_param_indices.len()))
}
fn n_data(&self) -> usize {
if let Some(layout) = &self.work_layout {
layout.data_indices.len()
} else if self.cross_sections.is_empty() {
0
} else {
self.cross_sections[0].len()
}
}
fn n_params(&self) -> usize {
self.density_indices
.iter()
.copied()
.max()
.map_or(0, |m| m + 1)
}
}
impl ForwardModel for TransmissionFitModel {
fn predict(&self, params: &[f64]) -> Result<Vec<f64>, FittingError> {
self.evaluate(params)
}
fn jacobian(
&self,
params: &[f64],
free_param_indices: &[usize],
y_current: &[f64],
) -> Option<Vec<Vec<f64>>> {
let fm = self.analytical_jacobian(params, free_param_indices, y_current)?;
Some(flat_matrix_to_vecs(&fm, free_param_indices.len()))
}
fn n_data(&self) -> usize {
self.energies.len()
}
fn n_params(&self) -> usize {
let max_density = self.density_indices.iter().copied().max().unwrap_or(0);
let max_temp = self.temperature_index.unwrap_or(0);
max_density.max(max_temp) + 1
}
}
impl<M: FitModel> ForwardModel for NormalizedTransmissionModel<M> {
fn predict(&self, params: &[f64]) -> Result<Vec<f64>, FittingError> {
self.evaluate(params)
}
fn jacobian(
&self,
params: &[f64],
free_param_indices: &[usize],
y_current: &[f64],
) -> Option<Vec<Vec<f64>>> {
let fm = self.analytical_jacobian(params, free_param_indices, y_current)?;
Some(flat_matrix_to_vecs(&fm, free_param_indices.len()))
}
fn n_data(&self) -> usize {
self.sqrt_energies.len()
}
fn n_params(&self) -> usize {
let mut max_idx = self
.anorm_index
.max(self.back_a_index)
.max(self.back_b_index)
.max(self.back_c_index);
if let Some(di) = self.back_d_index {
max_idx = max_idx.max(di);
}
if let Some(fi) = self.back_f_index {
max_idx = max_idx.max(fi);
}
max_idx + 1
}
}
impl ForwardModel for EnergyScaleTransmissionModel {
fn predict(&self, params: &[f64]) -> Result<Vec<f64>, FittingError> {
self.evaluate(params)
}
fn jacobian(
&self,
params: &[f64],
free_param_indices: &[usize],
y_current: &[f64],
) -> Option<Vec<Vec<f64>>> {
let fm = self.analytical_jacobian(params, free_param_indices, y_current)?;
Some(flat_matrix_to_vecs(&fm, free_param_indices.len()))
}
fn n_data(&self) -> usize {
self.nominal_energies.len()
}
fn n_params(&self) -> usize {
self.t0_index.max(self.l_scale_index) + 1
}
}
fn flat_matrix_to_vecs(fm: &FlatMatrix, cols: usize) -> Vec<Vec<f64>> {
let nrows = fm.nrows;
(0..cols)
.map(|j| (0..nrows).map(|i| fm.get(i, j)).collect())
.collect()
}
#[cfg(test)]
mod tests {
use super::*;
use crate::lm::{self, FitModel, LmConfig};
use crate::parameters::{FitParameter, ParameterSet};
use nereids_core::types::Isotope;
use nereids_endf::resonance::test_support::u238_single_resonance;
use nereids_endf::resonance::{LGroup, Resonance, ResonanceFormalism, ResonanceRange};
fn max_abs_diff(a: &[f64], b: &[f64]) -> f64 {
a.iter()
.zip(b.iter())
.map(|(x, y)| (x - y).abs())
.fold(0.0f64, f64::max)
}
fn max_abs(a: &[f64]) -> f64 {
a.iter().map(|x| x.abs()).fold(0.0f64, f64::max)
}
#[test]
fn precomputed_evaluate_matches_beer_lambert() {
let model = make_precomputed(
vec![
vec![1.0, 2.0, 3.0], vec![0.5, 0.5, 0.5], ],
vec![0, 1],
);
let params = [0.2f64, 0.4f64];
let y = model.evaluate(¶ms).unwrap();
let expected: Vec<f64> = (0..3)
.map(|i| {
let s0 = [1.0, 2.0, 3.0][i];
let s1 = [0.5, 0.5, 0.5][i];
(-params[0] * s0 - params[1] * s1).exp()
})
.collect();
assert_eq!(y.len(), 3);
for (yi, ei) in y.iter().zip(expected.iter()) {
assert!(
(yi - ei).abs() < 1e-12,
"evaluate mismatch: got {yi}, expected {ei}"
);
}
}
#[test]
fn precomputed_analytical_jacobian_matches_finite_difference() {
let model = make_precomputed(
vec![
vec![1.0, 2.0, 3.0], vec![0.5, 0.5, 0.5], ],
vec![0, 1],
);
let params = [0.2f64, 0.4f64];
let y = model.evaluate(¶ms).unwrap();
let free = vec![0usize, 1usize];
let jac = model
.analytical_jacobian(¶ms, &free, &y)
.expect("analytical_jacobian should return Some(_)");
assert_eq!(jac.nrows, 3); assert_eq!(jac.ncols, 2);
let h = 1e-6f64;
for (col, &p_idx) in free.iter().enumerate() {
let mut p_plus = params;
let mut p_minus = params;
p_plus[p_idx] += h;
p_minus[p_idx] -= h;
let y_plus = model.evaluate(&p_plus).unwrap();
let y_minus = model.evaluate(&p_minus).unwrap();
for i in 0..3 {
let fd = (y_plus[i] - y_minus[i]) / (2.0 * h);
let ana = jac.get(i, col);
assert!(
(fd - ana).abs() < 1e-6,
"Jacobian mismatch (row {i}, col {col}): FD={fd:.8}, analytical={ana:.8}"
);
}
}
}
#[test]
fn precomputed_jacobian_tied_parameters_sums_both_isotopes() {
let model = make_precomputed(
vec![
vec![1.0, 2.0, 3.0], vec![0.5, 1.0, 1.5], ],
vec![0, 0], );
let params = [0.1f64];
let y = model.evaluate(¶ms).unwrap();
let free = vec![0usize];
let jac = model
.analytical_jacobian(¶ms, &free, &y)
.expect("analytical_jacobian should return Some(_)");
for i in 0..3 {
let sigma_sum = [1.0, 2.0, 3.0][i] + [0.5, 1.0, 1.5][i];
let expected = -y[i] * sigma_sum;
assert!(
(jac.get(i, 0) - expected).abs() < 1e-12,
"Tied Jacobian mismatch at E[{i}]: got {}, expected {expected}",
jac.get(i, 0)
);
}
}
#[test]
fn test_recover_single_isotope_thickness() {
let data = u238_single_resonance();
let true_thickness = 0.0005;
let energies: Vec<f64> = (0..201).map(|i| 1.0 + (i as f64) * 0.05).collect();
let model = TransmissionFitModel::new(
energies.clone(),
vec![data],
0.0,
None,
(vec![0], vec![1.0]),
None,
None,
)
.unwrap();
let y_obs = model.evaluate(&[true_thickness]).unwrap();
let sigma = vec![0.01; y_obs.len()];
let mut params = ParameterSet::new(vec![
FitParameter::non_negative("thickness", 0.001), ]);
let result =
lm::levenberg_marquardt(&model, &y_obs, &sigma, &mut params, &LmConfig::default())
.unwrap();
assert!(result.converged, "Fit did not converge");
let fitted = result.params[0];
assert!(
(fitted - true_thickness).abs() / true_thickness < 0.01,
"Fitted thickness = {}, true = {}, error = {:.1}%",
fitted,
true_thickness,
(fitted - true_thickness).abs() / true_thickness * 100.0,
);
}
#[test]
fn test_recover_two_isotope_thicknesses() {
let u238 = u238_single_resonance();
let other = ResonanceData {
isotope: Isotope::new(1, 10).unwrap(),
za: 1010,
awr: 10.0,
ranges: vec![ResonanceRange {
energy_low: 0.0,
energy_high: 100.0,
resolved: true,
formalism: ResonanceFormalism::ReichMoore,
target_spin: 0.0,
scattering_radius: 5.0,
naps: 1,
l_groups: vec![LGroup {
l: 0,
awr: 10.0,
apl: 5.0,
qx: 0.0,
lrx: 0,
resonances: vec![Resonance {
energy: 20.0,
j: 0.5,
gn: 0.1,
gg: 0.05,
gfa: 0.0,
gfb: 0.0,
}],
}],
rml: None,
urr: None,
ap_table: None,
r_external: vec![],
}],
};
let true_t1 = 0.0003;
let true_t2 = 0.0001;
let energies: Vec<f64> = (0..301).map(|i| 1.0 + (i as f64) * 0.1).collect();
let model = TransmissionFitModel::new(
energies.clone(),
vec![u238, other],
0.0,
None,
(vec![0, 1], vec![1.0, 1.0]),
None,
None,
)
.unwrap();
let y_obs = model.evaluate(&[true_t1, true_t2]).unwrap();
let sigma = vec![0.01; y_obs.len()];
let mut params = ParameterSet::new(vec![
FitParameter::non_negative("U-238 thickness", 0.001),
FitParameter::non_negative("Other thickness", 0.001),
]);
let result =
lm::levenberg_marquardt(&model, &y_obs, &sigma, &mut params, &LmConfig::default())
.unwrap();
assert!(
result.converged,
"Fit did not converge after {} iterations",
result.iterations
);
let (fit_t1, fit_t2) = (result.params[0], result.params[1]);
assert!(
(fit_t1 - true_t1).abs() / true_t1 < 0.05,
"U-238: fitted={}, true={}, error={:.1}%",
fit_t1,
true_t1,
(fit_t1 - true_t1).abs() / true_t1 * 100.0,
);
assert!(
(fit_t2 - true_t2).abs() / true_t2 < 0.05,
"Other: fitted={}, true={}, error={:.1}%",
fit_t2,
true_t2,
(fit_t2 - true_t2).abs() / true_t2 * 100.0,
);
}
#[test]
fn temperature_index_overrides_fixed_temperature() {
let data = u238_single_resonance();
let energies: Vec<f64> = (0..201).map(|i| 1.0 + (i as f64) * 0.05).collect();
let model = TransmissionFitModel::new(
energies.clone(),
vec![data.clone()],
0.0,
None,
(vec![0], vec![1.0]),
Some(1),
None,
)
.unwrap();
let model_fixed = TransmissionFitModel::new(
energies.clone(),
vec![data],
300.0,
None,
(vec![0], vec![1.0]),
None,
None,
)
.unwrap();
let density = 0.0005;
let y_via_index = model.evaluate(&[density, 300.0]).unwrap();
let y_via_fixed = model_fixed.evaluate(&[density]).unwrap();
for (a, b) in y_via_index.iter().zip(y_via_fixed.iter()) {
assert!(
(a - b).abs() < 1e-12,
"temperature_index path disagrees with fixed path: {} vs {}",
a,
b
);
}
}
#[test]
fn test_recover_temperature() {
let data = u238_single_resonance();
let true_density = 0.0005;
let true_temp = 300.0;
let energies: Vec<f64> = (0..401).map(|i| 4.0 + (i as f64) * 0.025).collect();
let model = TransmissionFitModel::new(
energies.clone(),
vec![data],
0.0, None,
(vec![0], vec![1.0]),
Some(1), None,
)
.unwrap();
let mut y_obs = model.evaluate(&[true_density, true_temp]).unwrap();
for (i, y) in y_obs.iter_mut().enumerate() {
*y *= 1.0 + 1e-5 * ((i % 7) as f64 - 3.0);
}
let sigma = vec![0.005; y_obs.len()];
let mut params = ParameterSet::new(vec![
FitParameter::non_negative("density", 0.001),
FitParameter {
name: "temperature_k".into(),
value: 200.0, lower: 1.0,
upper: 2000.0,
fixed: false,
},
]);
let config = LmConfig {
max_iter: 200,
..LmConfig::default()
};
let result = lm::levenberg_marquardt(&model, &y_obs, &sigma, &mut params, &config).unwrap();
assert!(
result.converged,
"Temperature fit did not converge after {} iterations",
result.iterations
);
let fit_density = result.params[0];
let fit_temp = result.params[1];
assert!(
(fit_density - true_density).abs() / true_density < 0.001,
"Density: fitted={}, true={}, error={:.1}%",
fit_density,
true_density,
(fit_density - true_density).abs() / true_density * 100.0,
);
assert!(
(fit_temp - true_temp).abs() / true_temp < 0.001,
"Temperature: fitted={:.1} K, true={:.1} K, error={:.1}%",
fit_temp,
true_temp,
(fit_temp - true_temp).abs() / true_temp * 100.0,
);
let unc = result
.uncertainties
.expect("uncertainties should be available");
assert!(
unc.len() == 2,
"expected 2 uncertainties, got {}",
unc.len()
);
assert!(
unc[1] > 0.0 && unc[1].is_finite(),
"temperature uncertainty should be positive and finite, got {}",
unc[1]
);
}
#[test]
fn transmission_fit_model_analytical_jacobian_matches_fd() {
let data = u238_single_resonance();
let energies: Vec<f64> = (0..201).map(|i| 1.0 + (i as f64) * 0.05).collect();
let model = TransmissionFitModel::new(
energies,
vec![data],
0.0,
None,
(vec![0], vec![1.0]),
Some(1), None,
)
.unwrap();
let params = [0.0005f64, 300.0f64]; let y = model.evaluate(¶ms).unwrap();
let free = vec![0usize, 1usize];
let jac = model
.analytical_jacobian(¶ms, &free, &y)
.expect("analytical_jacobian should return Some(_)");
assert_eq!(jac.nrows, y.len());
assert_eq!(jac.ncols, 2);
let h = 1e-6f64;
for (col, &p_idx) in free.iter().enumerate() {
let mut p_plus = params;
let mut p_minus = params;
p_plus[p_idx] += h * (1.0 + params[p_idx].abs());
p_minus[p_idx] -= h * (1.0 + params[p_idx].abs());
let y_plus = model.evaluate(&p_plus).unwrap();
let y_minus = model.evaluate(&p_minus).unwrap();
let actual_2h = p_plus[p_idx] - p_minus[p_idx];
for i in 0..y.len() {
let fd = (y_plus[i] - y_minus[i]) / actual_2h;
let ana = jac.get(i, col);
let err = (fd - ana).abs();
let scale = fd.abs().max(ana.abs()).max(1e-10);
assert!(
err / scale < 0.01,
"Jacobian mismatch (row {i}, col {col}): FD={fd:.8}, analytical={ana:.8}, \
rel_err={:.4}",
err / scale,
);
}
}
}
#[test]
fn transmission_fit_model_cache_reuse() {
let data = u238_single_resonance();
let energies: Vec<f64> = (0..201).map(|i| 1.0 + (i as f64) * 0.05).collect();
let model = TransmissionFitModel::new(
energies,
vec![data],
0.0,
None,
(vec![0], vec![1.0]),
Some(1),
None,
)
.unwrap();
let y1 = model.evaluate(&[0.0005, 300.0]).unwrap();
assert!(model.cached_broadened_xs.borrow().is_some());
assert!((model.cached_temperature.get() - 300.0).abs() < 1e-15);
let y2 = model.evaluate(&[0.001, 300.0]).unwrap();
assert!((model.cached_temperature.get() - 300.0).abs() < 1e-15);
assert!(
(y1[100] - y2[100]).abs() > 1e-10,
"different densities should produce different transmission"
);
let _y3 = model.evaluate(&[0.0005, 600.0]).unwrap();
assert!((model.cached_temperature.get() - 600.0).abs() < 1e-15);
}
fn make_precomputed(
xs: Vec<Vec<f64>>,
density_indices: Vec<usize>,
) -> PrecomputedTransmissionModel {
PrecomputedTransmissionModel {
cross_sections: Arc::new(xs),
density_indices: Arc::new(density_indices),
energies: None,
instrument: None,
resolution_plan: None,
sparse_cubature_plan: None,
sparse_scalar_plan: None,
work_layout: None,
}
}
fn synthetic_resolution_setup(
n_grid: usize,
half_kernel: usize,
) -> (
Vec<f64>,
Arc<ResolutionPlan>,
nereids_physics::resolution::ResolutionMatrix,
) {
assert!(n_grid > 2 * half_kernel);
let energies: Vec<f64> = (0..n_grid).map(|i| 10.0 + i as f64).collect();
let mut starts: Vec<u32> = vec![0];
let mut lo_idx: Vec<u32> = Vec::new();
let mut frac_arr: Vec<f64> = Vec::new();
let mut weight_arr: Vec<f64> = Vec::new();
let mut norm: Vec<f64> = Vec::with_capacity(n_grid);
for i in 0..n_grid {
let lo_min = i.saturating_sub(half_kernel);
let lo_max = (i + half_kernel).min(n_grid - 2);
let mut row_norm = 0.0_f64;
for lo in lo_min..=lo_max {
let d = (lo as i64 - i as i64).abs() as f64;
let w = 1.0 - d / (half_kernel as f64 + 1.0);
lo_idx.push(lo as u32);
frac_arr.push(0.5);
weight_arr.push(w);
row_norm += w;
}
norm.push(row_norm);
starts.push(lo_idx.len() as u32);
}
let plan = nereids_physics::resolution::test_support::plan_from_raw_parts(
energies.clone(),
starts,
lo_idx,
frac_arr,
weight_arr,
norm,
);
let matrix = plan.compile_to_matrix();
(energies, Arc::new(plan), matrix)
}
fn synthetic_sigmas(n_grid: usize, k: usize) -> Vec<Vec<f64>> {
let mut out = Vec::with_capacity(k);
for j in 0..k {
let center = 10.0 + (j as f64 + 1.0) * (n_grid as f64) / (k as f64 + 1.0);
let width = 3.0;
out.push(
(0..n_grid)
.map(|ell| {
let e = 10.0 + ell as f64;
100.0 * (-((e - center).powi(2)) / (width * width)).exp() + 5.0
})
.collect(),
);
}
out
}
fn build_cubature(
matrix: &nereids_physics::resolution::ResolutionMatrix,
sigmas: &[Vec<f64>],
train_max: Vec<f64>,
) -> Arc<SparseEmpiricalCubaturePlan> {
let k = sigmas.len();
let n_rows = matrix.len();
let mut flat = Vec::with_capacity(k * n_rows);
for row in sigmas {
flat.extend_from_slice(row);
}
let training = SparseEmpiricalCubaturePlan::default_training_points(&train_max);
let anchor = SparseEmpiricalCubaturePlan::default_jacobian_anchor(&train_max);
Arc::new(
SparseEmpiricalCubaturePlan::build(matrix, &flat, k, &training, &anchor)
.expect("synthetic cubature build"),
)
}
fn make_trivial_instrument() -> Arc<InstrumentParams> {
use nereids_physics::resolution::ResolutionFunction;
let tab =
Arc::new(nereids_physics::resolution::test_support::trivial_tabulated_resolution(25.0));
let res_fn = ResolutionFunction::Tabulated(tab);
Arc::new(InstrumentParams { resolution: res_fn })
}
#[test]
fn precomputed_cubature_dispatches_at_k2_matching_k() {
let n_grid = 40_usize;
let (energies, plan, matrix) = synthetic_resolution_setup(n_grid, 4);
let sigmas = synthetic_sigmas(n_grid, 2);
let train_max = vec![1e-4_f64, 1e-4];
let cubature = build_cubature(&matrix, &sigmas, train_max.clone());
let mut model = PrecomputedTransmissionModel {
cross_sections: Arc::new(sigmas.clone()),
density_indices: Arc::new(vec![0, 1]),
energies: Some(Arc::new(energies.clone())),
instrument: Some(make_trivial_instrument()),
resolution_plan: Some(Arc::clone(&plan)),
sparse_cubature_plan: Some(Arc::clone(&cubature)),
sparse_scalar_plan: None,
work_layout: None,
};
let n = [0.25 * train_max[0], 0.25 * train_max[1]];
let t_cubature = model.evaluate(&n).unwrap();
model.sparse_cubature_plan = None;
let t_exact_via_r = {
let n_grid_local = n_grid;
let mut neg_opt = vec![0.0_f64; n_grid_local];
for (j, &nj) in n.iter().enumerate() {
for (ell, &sig) in sigmas[j].iter().enumerate() {
neg_opt[ell] -= nj * sig;
}
}
let t_un: Vec<f64> = neg_opt.iter().map(|&d| d.exp()).collect();
nereids_physics::resolution::apply_r(&matrix, &t_un)
};
let t_cubature_direct = cubature.forward(&n);
for (a, b) in t_cubature.iter().zip(t_cubature_direct.iter()) {
assert!((a - b).abs() < 1e-14);
}
let max_err = t_cubature
.iter()
.zip(t_exact_via_r.iter())
.map(|(a, b)| {
let denom = a.abs().max(b.abs()).max(1e-12);
(a - b).abs() / denom
})
.fold(0.0_f64, f64::max);
assert!(
max_err < 1e-9,
"at training density, cubature vs exact max err = {max_err:.3e}",
);
}
#[test]
fn precomputed_cubature_falls_back_at_k1() {
let n_grid = 40_usize;
let (energies, _plan, matrix) = synthetic_resolution_setup(n_grid, 4);
let sigmas_k2 = synthetic_sigmas(n_grid, 2);
let cubature_k2 = build_cubature(&matrix, &sigmas_k2, vec![1e-4_f64, 1e-4]);
let sigmas_k1 = synthetic_sigmas(n_grid, 1);
let model_with_plan = PrecomputedTransmissionModel {
cross_sections: Arc::new(sigmas_k1.clone()),
density_indices: Arc::new(vec![0]),
energies: Some(Arc::new(energies.clone())),
instrument: Some(make_trivial_instrument()),
resolution_plan: None,
sparse_cubature_plan: Some(Arc::clone(&cubature_k2)),
sparse_scalar_plan: None,
work_layout: None,
};
let model_without_plan = PrecomputedTransmissionModel {
cross_sections: Arc::new(sigmas_k1.clone()),
density_indices: Arc::new(vec![0]),
energies: Some(Arc::new(energies.clone())),
instrument: Some(make_trivial_instrument()),
resolution_plan: None,
sparse_cubature_plan: None,
sparse_scalar_plan: None,
work_layout: None,
};
let n = [1e-4_f64];
let t_with = model_with_plan.evaluate(&n).unwrap();
let t_without = model_without_plan.evaluate(&n).unwrap();
assert_eq!(t_with.len(), n_grid);
assert_eq!(t_without.len(), n_grid);
for (a, b) in t_with.iter().zip(t_without.iter()) {
assert_eq!(
a.to_bits(),
b.to_bits(),
"fallback path must be byte-identical to the no-plan path; \
otherwise the k=2 cubature is silently firing on a k=1 model",
);
}
}
#[test]
fn precomputed_cubature_no_plan_means_exact_path() {
let n_grid = 40_usize;
let (_energies, _plan, _matrix) = synthetic_resolution_setup(n_grid, 4);
let sigmas = synthetic_sigmas(n_grid, 2);
let model = PrecomputedTransmissionModel {
cross_sections: Arc::new(sigmas.clone()),
density_indices: Arc::new(vec![0, 1]),
energies: None,
instrument: None,
resolution_plan: None,
sparse_cubature_plan: None,
sparse_scalar_plan: None,
work_layout: None,
};
let n = [1e-4_f64, 1e-4];
let t = model.evaluate(&n).unwrap();
for (ell, &t_val) in t.iter().enumerate() {
let tau: f64 = sigmas
.iter()
.zip(n.iter())
.map(|(s, &ni)| ni * s[ell])
.sum();
let expected = (-tau).exp();
assert!(
(t_val - expected).abs() < 1e-14,
"at ell={ell}: got {t_val}, expected {expected}",
);
}
}
#[test]
fn precomputed_cubature_jacobian_matches_forward_derivative() {
let n_grid = 40_usize;
let (energies, plan, matrix) = synthetic_resolution_setup(n_grid, 4);
let sigmas = synthetic_sigmas(n_grid, 2);
let train_max = vec![1e-4_f64, 1e-4];
let cubature = build_cubature(&matrix, &sigmas, train_max.clone());
let model = PrecomputedTransmissionModel {
cross_sections: Arc::new(sigmas),
density_indices: Arc::new(vec![0, 1]),
energies: Some(Arc::new(energies)),
instrument: Some(make_trivial_instrument()),
resolution_plan: Some(Arc::clone(&plan)),
sparse_cubature_plan: Some(Arc::clone(&cubature)),
sparse_scalar_plan: None,
work_layout: None,
};
let anchor = SparseEmpiricalCubaturePlan::default_jacobian_anchor(&train_max);
let y_curr = model.evaluate(&anchor).unwrap();
let jac = model
.analytical_jacobian(&anchor, &[0, 1], &y_curr)
.expect("cubature Jacobian path");
let (_t_ref, jac_flat_ref) = cubature.forward_and_jacobian(&anchor);
for i in 0..n_grid {
for col in 0..2 {
let from_model = jac.get(i, col);
let from_cubature = jac_flat_ref[i * 2 + col];
assert!(
(from_model - from_cubature).abs() < 1e-14,
"row {i} col {col}: model = {from_model}, cubature = {from_cubature}",
);
}
}
}
fn make_trivial_fit_model(energies: Vec<f64>, k: usize) -> TransmissionFitModel {
let resonance_data: Vec<ResonanceData> = (0..k)
.map(|j| {
let iso = Isotope::new(40 + j as u32, 96 + j as u32).unwrap();
ResonanceData {
isotope: iso,
za: ((40 + j) * 1000 + (96 + j)) as u32,
awr: 96.0 + j as f64,
ranges: vec![],
}
})
.collect();
TransmissionFitModel::new(
energies,
resonance_data,
293.6,
Some(make_trivial_instrument()),
((0..k).collect(), vec![1.0; k]),
None,
None,
)
.expect("TransmissionFitModel::new")
}
#[test]
fn fit_model_cubature_dispatches_at_anchor() {
let n_grid = 40_usize;
let (energies, plan, matrix) = synthetic_resolution_setup(n_grid, 4);
let sigmas = synthetic_sigmas(n_grid, 2);
let train_max = vec![1e-4_f64, 1e-4];
let cubature = build_cubature(&matrix, &sigmas, train_max.clone());
let model = make_trivial_fit_model(energies.clone(), 2)
.with_resolution_plan(Some(Arc::clone(&plan)))
.with_sparse_cubature_plan(Some(cubature.clone()));
let n = [0.25 * train_max[0], 0.25 * train_max[1]];
let t_model = model.evaluate(&n).unwrap();
let t_cub = cubature.forward(&n);
assert_eq!(t_model.len(), n_grid);
for (a, b) in t_model.iter().zip(t_cub.iter()) {
assert_eq!(
a.to_bits(),
b.to_bits(),
"TransmissionFitModel cubature dispatch must return cubature.forward() byte-exact at the LP-pinned anchor",
);
}
}
#[test]
fn fit_model_cubature_jacobian_matches_cubature_output() {
let n_grid = 40_usize;
let (energies, plan, matrix) = synthetic_resolution_setup(n_grid, 4);
let sigmas = synthetic_sigmas(n_grid, 2);
let train_max = vec![1e-4_f64, 1e-4];
let cubature = build_cubature(&matrix, &sigmas, train_max.clone());
let model = make_trivial_fit_model(energies, 2)
.with_resolution_plan(Some(Arc::clone(&plan)))
.with_sparse_cubature_plan(Some(cubature.clone()));
let anchor = SparseEmpiricalCubaturePlan::default_jacobian_anchor(&train_max);
let y_curr = model.evaluate(&anchor).unwrap();
let jac = model
.analytical_jacobian(&anchor, &[0, 1], &y_curr)
.expect("cubature Jacobian path on TransmissionFitModel");
let (_t_ref, jac_flat_ref) = cubature.forward_and_jacobian(&anchor);
for i in 0..n_grid {
for col in 0..2 {
let from_model = jac.get(i, col);
let from_cubature = jac_flat_ref[i * 2 + col];
assert_eq!(
from_model.to_bits(),
from_cubature.to_bits(),
"row {i} col {col}: TransmissionFitModel must return cubature J byte-exact",
);
}
}
}
#[test]
fn fit_model_cubature_falls_back_on_grid_mismatch() {
let n_grid = 40_usize;
let (energies_a, _plan, matrix) = synthetic_resolution_setup(n_grid, 4);
let sigmas = synthetic_sigmas(n_grid, 2);
let train_max = vec![1e-4_f64, 1e-4];
let cubature = build_cubature(&matrix, &sigmas, train_max);
let energies_b: Vec<f64> = energies_a.iter().map(|&e| e + 1.0).collect();
let model_with_stale_plan =
make_trivial_fit_model(energies_b.clone(), 2).with_sparse_cubature_plan(Some(cubature));
let model_without_plan = make_trivial_fit_model(energies_b, 2);
let n = [1e-5_f64, 1e-5];
let t_stale = model_with_stale_plan.evaluate(&n).unwrap();
let t_exact = model_without_plan.evaluate(&n).unwrap();
for (a, b) in t_stale.iter().zip(t_exact.iter()) {
assert_eq!(
a.to_bits(),
b.to_bits(),
"stale-grid cubature plan MUST NOT fire; evaluate() must match no-plan byte-exactly",
);
}
}
#[test]
fn fit_model_cubature_falls_back_when_density_escapes_box() {
let n_grid = 40_usize;
let (energies, plan, matrix) = synthetic_resolution_setup(n_grid, 4);
let sigmas = synthetic_sigmas(n_grid, 2);
let train_max = vec![1e-4_f64, 1e-4];
let cubature = {
let flat: Vec<f64> = sigmas.iter().flat_map(|s| s.iter().copied()).collect();
let training = SparseEmpiricalCubaturePlan::default_training_points(&train_max);
let anchor = SparseEmpiricalCubaturePlan::default_jacobian_anchor(&train_max);
Arc::new(
SparseEmpiricalCubaturePlan::build(&matrix, &flat, 2, &training, &anchor)
.expect("build")
.with_density_box(train_max.clone()),
)
};
let model_with = make_trivial_fit_model(energies.clone(), 2)
.with_resolution_plan(Some(Arc::clone(&plan)))
.with_sparse_cubature_plan(Some(Arc::clone(&cubature)));
let model_without =
make_trivial_fit_model(energies, 2).with_resolution_plan(Some(Arc::clone(&plan)));
let n_escape = [5.0 * train_max[0], 5.0 * train_max[1]];
let t_with = model_with.evaluate(&n_escape).unwrap();
let t_without = model_without.evaluate(&n_escape).unwrap();
for (a, b) in t_with.iter().zip(t_without.iter()) {
assert_eq!(
a.to_bits(),
b.to_bits(),
"density-box escape guard MUST fall back to exact path byte-identically",
);
}
}
#[test]
fn fit_model_cubature_dispatches_without_resolution_plan_attached() {
let n_grid = 40_usize;
let (energies, _plan, matrix) = synthetic_resolution_setup(n_grid, 4);
let sigmas = synthetic_sigmas(n_grid, 2);
let train_max = vec![1e-4_f64, 1e-4];
let cubature = build_cubature(&matrix, &sigmas, train_max.clone());
let model = make_trivial_fit_model(energies.clone(), 2)
.with_sparse_cubature_plan(Some(Arc::clone(&cubature)));
let n = [0.25 * train_max[0], 0.25 * train_max[1]];
let t_model = model.evaluate(&n).unwrap();
let t_cub = cubature.forward(&n);
for (a, b) in t_model.iter().zip(t_cub.iter()) {
assert_eq!(
a.to_bits(),
b.to_bits(),
"cubature dispatch must fire on single-spectrum path without a separate ResolutionPlan attached",
);
}
}
fn build_scalar_plan(
res_plan: Arc<ResolutionPlan>,
sigma_k1: &[f64],
n_max: f64,
) -> Arc<ScalarSurrogatePlan> {
Arc::new(
nereids_physics::surrogate::ScalarChebyshevPlan::build(res_plan, sigma_k1, n_max, 16)
.expect("synthetic scalar Chebyshev build"),
)
}
fn make_precomp_for_scalar(
energies: Vec<f64>,
sigmas: Vec<Vec<f64>>,
density_indices: Vec<usize>,
resolution_plan: Option<Arc<ResolutionPlan>>,
scalar_plan: Option<Arc<ScalarSurrogatePlan>>,
) -> PrecomputedTransmissionModel {
PrecomputedTransmissionModel {
cross_sections: Arc::new(sigmas),
density_indices: Arc::new(density_indices),
energies: Some(Arc::new(energies)),
instrument: Some(make_trivial_instrument()),
resolution_plan,
sparse_cubature_plan: None,
sparse_scalar_plan: scalar_plan,
work_layout: None,
}
}
#[test]
fn precomputed_scalar_dispatches_at_k1() {
let n_grid = 40_usize;
let (energies, res_plan, _matrix) = synthetic_resolution_setup(n_grid, 4);
let sigmas = synthetic_sigmas(n_grid, 1);
let n_max = 2.0 * 1e-4_f64;
let scalar = build_scalar_plan(Arc::clone(&res_plan), &sigmas[0], n_max);
let model = make_precomp_for_scalar(
energies,
sigmas,
vec![0],
Some(Arc::clone(&res_plan)),
Some(Arc::clone(&scalar)),
);
let n = [0.5 * n_max];
let t_model = model.evaluate(&n).unwrap();
let t_scalar = scalar.forward_scalar(n[0]);
assert_eq!(t_model.len(), n_grid);
for (a, b) in t_model.iter().zip(t_scalar.iter()) {
assert_eq!(
a.to_bits(),
b.to_bits(),
"scalar dispatch must return forward_scalar() byte-exact",
);
}
}
#[test]
fn precomputed_scalar_jacobian_matches_derivative() {
let n_grid = 40_usize;
let (energies, res_plan, _matrix) = synthetic_resolution_setup(n_grid, 4);
let sigmas = synthetic_sigmas(n_grid, 1);
let n_max = 2.0 * 1e-4_f64;
let scalar = build_scalar_plan(Arc::clone(&res_plan), &sigmas[0], n_max);
let model = make_precomp_for_scalar(
energies,
sigmas,
vec![0],
Some(Arc::clone(&res_plan)),
Some(Arc::clone(&scalar)),
);
let n = [0.5 * n_max];
let y_curr = model.evaluate(&n).unwrap();
let jac = model
.analytical_jacobian(&n, &[0], &y_curr)
.expect("scalar Jacobian path");
let (_t_ref, dt_ref) = scalar.forward_and_derivative_scalar(n[0]);
assert_eq!(jac.ncols, 1);
assert_eq!(jac.nrows, n_grid);
for (i, &dt_i) in dt_ref.iter().enumerate().take(n_grid) {
assert_eq!(
jac.get(i, 0).to_bits(),
dt_i.to_bits(),
"row {i}: scalar dT/dn must be byte-exact",
);
}
}
#[test]
fn precomputed_scalar_falls_back_at_k2() {
let n_grid = 40_usize;
let (energies, res_plan, _matrix) = synthetic_resolution_setup(n_grid, 4);
let sigmas_k2 = synthetic_sigmas(n_grid, 2);
let sigma_k1 = synthetic_sigmas(n_grid, 1).remove(0);
let n_max = 2.0 * 1e-4_f64;
let scalar = build_scalar_plan(Arc::clone(&res_plan), &sigma_k1, n_max);
let model_with = make_precomp_for_scalar(
energies.clone(),
sigmas_k2.clone(),
vec![0, 1],
Some(Arc::clone(&res_plan)),
Some(scalar),
);
let model_without = make_precomp_for_scalar(
energies,
sigmas_k2,
vec![0, 1],
Some(Arc::clone(&res_plan)),
None,
);
let n = [1e-4_f64, 2e-4];
let t_with = model_with.evaluate(&n).unwrap();
let t_without = model_without.evaluate(&n).unwrap();
for (a, b) in t_with.iter().zip(t_without.iter()) {
assert_eq!(
a.to_bits(),
b.to_bits(),
"scalar plan must refuse k=2 dispatch → byte-identical fallback",
);
}
}
#[test]
fn precomputed_scalar_falls_back_on_stale_resolution_plan() {
let n_grid = 40_usize;
let (energies, res_plan_a, _matrix) = synthetic_resolution_setup(n_grid, 4);
let sigmas = synthetic_sigmas(n_grid, 1);
let n_max = 2.0 * 1e-4_f64;
let scalar = build_scalar_plan(Arc::clone(&res_plan_a), &sigmas[0], n_max);
let (_e_b, res_plan_b, _matrix_b) = synthetic_resolution_setup(n_grid, 6);
let model_stale = make_precomp_for_scalar(
energies.clone(),
sigmas.clone(),
vec![0],
Some(Arc::clone(&res_plan_b)),
Some(Arc::clone(&scalar)),
);
let model_noplan =
make_precomp_for_scalar(energies, sigmas, vec![0], Some(res_plan_b), None);
let n = [0.25 * n_max];
let t_stale = model_stale.evaluate(&n).unwrap();
let t_exact = model_noplan.evaluate(&n).unwrap();
for (a, b) in t_stale.iter().zip(t_exact.iter()) {
assert_eq!(
a.to_bits(),
b.to_bits(),
"scalar plan with non-ptr_eq source_resolution_plan MUST NOT fire",
);
}
}
#[test]
fn precomputed_scalar_falls_back_on_stale_sigma() {
let n_grid = 40_usize;
let (energies, res_plan, _matrix) = synthetic_resolution_setup(n_grid, 4);
let sigma_a = synthetic_sigmas(n_grid, 1);
let mut sigma_b = sigma_a.clone();
sigma_b[0][n_grid / 2] += 1.0; let n_max = 2.0 * 1e-4_f64;
let scalar = build_scalar_plan(Arc::clone(&res_plan), &sigma_a[0], n_max);
let model_stale = make_precomp_for_scalar(
energies.clone(),
sigma_b.clone(),
vec![0],
Some(Arc::clone(&res_plan)),
Some(scalar),
);
let model_noplan =
make_precomp_for_scalar(energies, sigma_b, vec![0], Some(res_plan), None);
let n = [0.25 * n_max];
let t_stale = model_stale.evaluate(&n).unwrap();
let t_exact = model_noplan.evaluate(&n).unwrap();
for (a, b) in t_stale.iter().zip(t_exact.iter()) {
assert_eq!(
a.to_bits(),
b.to_bits(),
"σ-fingerprint mismatch MUST force fallback → byte-identical to no-plan",
);
}
}
#[test]
fn precomputed_scalar_falls_back_when_density_escapes_box() {
let n_grid = 40_usize;
let (energies, res_plan, _matrix) = synthetic_resolution_setup(n_grid, 4);
let sigmas = synthetic_sigmas(n_grid, 1);
let n_max = 2.0 * 1e-4_f64;
let scalar = build_scalar_plan(Arc::clone(&res_plan), &sigmas[0], n_max);
let model_with = make_precomp_for_scalar(
energies.clone(),
sigmas.clone(),
vec![0],
Some(Arc::clone(&res_plan)),
Some(Arc::clone(&scalar)),
);
let model_without =
make_precomp_for_scalar(energies, sigmas, vec![0], Some(Arc::clone(&res_plan)), None);
let n_escape = [2.0 * n_max];
let t_with = model_with.evaluate(&n_escape).unwrap();
let t_without = model_without.evaluate(&n_escape).unwrap();
for (a, b) in t_with.iter().zip(t_without.iter()) {
assert_eq!(
a.to_bits(),
b.to_bits(),
"density-box escape guard must fall back byte-identically",
);
}
}
#[test]
fn precomputed_scalar_rejects_nonfinite_density() {
let n_grid = 40_usize;
let (energies, res_plan, _matrix) = synthetic_resolution_setup(n_grid, 4);
let sigmas = synthetic_sigmas(n_grid, 1);
let n_max = 2.0 * 1e-4_f64;
let scalar = build_scalar_plan(Arc::clone(&res_plan), &sigmas[0], n_max);
let model_with = make_precomp_for_scalar(
energies.clone(),
sigmas.clone(),
vec![0],
Some(Arc::clone(&res_plan)),
Some(Arc::clone(&scalar)),
);
let model_without =
make_precomp_for_scalar(energies, sigmas, vec![0], Some(Arc::clone(&res_plan)), None);
for bad_n in [f64::NAN, f64::INFINITY, -1e-6_f64] {
let n = [bad_n];
let t_with = model_with.evaluate(&n).unwrap();
let t_without = model_without.evaluate(&n).unwrap();
for (i, (a, b)) in t_with.iter().zip(t_without.iter()).enumerate() {
assert_eq!(
a.to_bits(),
b.to_bits(),
"n = {bad_n}: scalar guard must fall back byte-exactly; row {i}",
);
}
}
}
#[test]
fn scalar_density_within_box_direct_guard() {
let n_grid = 16_usize;
let (_energies, res_plan, _matrix) = synthetic_resolution_setup(n_grid, 2);
let sigmas = synthetic_sigmas(n_grid, 1);
let n_max = 1e-4_f64;
let plan =
nereids_physics::surrogate::ScalarChebyshevPlan::build(res_plan, &sigmas[0], n_max, 16)
.expect("build");
assert!(scalar_density_within_box(&plan, 0.0));
assert!(scalar_density_within_box(&plan, 0.5 * n_max));
assert!(scalar_density_within_box(&plan, n_max));
assert!(!scalar_density_within_box(
&plan,
n_max * (1.0 + f64::EPSILON)
));
assert!(!scalar_density_within_box(&plan, 1.01 * n_max));
assert!(!scalar_density_within_box(&plan, 1.5 * n_max));
assert!(!scalar_density_within_box(&plan, 2.0 * n_max));
assert!(!scalar_density_within_box(&plan, f64::NAN));
assert!(!scalar_density_within_box(&plan, f64::INFINITY));
assert!(!scalar_density_within_box(&plan, f64::NEG_INFINITY));
assert!(!scalar_density_within_box(&plan, -1e-9));
}
#[test]
fn density_param_indices_sorted_by_value() {
assert_eq!(density_param_indices(&[0, 0, 0]), vec![0]);
assert_eq!(density_param_indices(&[0, 1, 2, 3]), vec![0, 1, 2, 3]);
assert_eq!(density_param_indices(&[1, 0, 1]), vec![0, 1]);
assert_eq!(density_param_indices(&[3, 1, 2, 0, 2]), vec![0, 1, 2, 3]);
}
#[test]
fn normalized_identity_matches_inner() {
let xs = vec![
vec![1.0, 2.0, 3.0], vec![0.5, 0.5, 0.5], ];
let inner_ref = make_precomputed(xs.clone(), vec![0, 1]);
let inner_wrap = make_precomputed(xs, vec![0, 1]);
let energies = [4.0, 9.0, 16.0];
let model = NormalizedTransmissionModel::new(inner_wrap, &energies, 2, 3, 4, 5);
let params = [0.2, 0.4, 1.0, 0.0, 0.0, 0.0];
let y_norm = model.evaluate(¶ms).unwrap();
let y_inner = inner_ref.evaluate(¶ms).unwrap();
for (a, b) in y_norm.iter().zip(y_inner.iter()) {
assert!(
(a - b).abs() < 1e-12,
"identity normalization should match inner: {} vs {}",
a,
b
);
}
}
#[test]
fn normalized_formula_correct() {
let xs = vec![vec![1.0, 2.0, 3.0]];
let inner_ref = make_precomputed(xs.clone(), vec![0]);
let inner_wrap = make_precomputed(xs, vec![0]);
let energies = [4.0, 9.0, 16.0]; let model = NormalizedTransmissionModel::new(inner_wrap, &energies, 1, 2, 3, 4);
let anorm = 0.95;
let back_a = 0.01;
let back_b = 0.02;
let back_c = 0.005;
let density = 0.3;
let params = [density, anorm, back_a, back_b, back_c];
let y = model.evaluate(¶ms).unwrap();
let t_inner = inner_ref.evaluate(¶ms).unwrap();
for (i, (&yi, &ti)) in y.iter().zip(t_inner.iter()).enumerate() {
let sqrt_e = energies[i].sqrt();
let expected = anorm * ti + back_a + back_b / sqrt_e + back_c * sqrt_e;
assert!(
(yi - expected).abs() < 1e-12,
"E[{i}]: got {yi}, expected {expected}"
);
}
}
#[test]
fn normalized_analytical_jacobian_matches_fd() {
let xs = vec![
vec![1.0, 2.0, 3.0], vec![0.5, 0.5, 0.5], ];
let inner = make_precomputed(xs, vec![0, 1]);
let energies = [4.0, 9.0, 16.0];
let model = NormalizedTransmissionModel::new(inner, &energies, 2, 3, 4, 5);
let params = [0.2, 0.4, 0.95, 0.01, 0.02, 0.005];
let y = model.evaluate(¶ms).unwrap();
let free: Vec<usize> = (0..6).collect();
let jac = model
.analytical_jacobian(¶ms, &free, &y)
.expect("analytical_jacobian should return Some");
assert_eq!(jac.nrows, 3);
assert_eq!(jac.ncols, 6);
let h = 1e-7;
for (col, &p_idx) in free.iter().enumerate() {
let mut p_plus = params;
let mut p_minus = params;
p_plus[p_idx] += h;
p_minus[p_idx] -= h;
let y_plus = model.evaluate(&p_plus).unwrap();
let y_minus = model.evaluate(&p_minus).unwrap();
for i in 0..3 {
let fd = (y_plus[i] - y_minus[i]) / (2.0 * h);
let ana = jac.get(i, col);
let err = (fd - ana).abs();
let scale = fd.abs().max(ana.abs()).max(1e-10);
assert!(
err / scale < 1e-4,
"Jacobian mismatch (row {i}, col {col}): FD={fd:.8}, analytical={ana:.8}, \
rel_err={:.6}",
err / scale,
);
}
}
}
#[test]
fn normalized_jacobian_partial_free() {
let xs = vec![vec![1.0, 2.0, 3.0]];
let inner = make_precomputed(xs, vec![0]);
let energies = [4.0, 9.0, 16.0];
let model = NormalizedTransmissionModel::new(inner, &energies, 1, 2, 3, 4);
let params = [0.3, 0.95, 0.01, 0.0, 0.0];
let y = model.evaluate(¶ms).unwrap();
let free = vec![0usize, 1usize];
let jac = model
.analytical_jacobian(¶ms, &free, &y)
.expect("should return Some for partial free");
assert_eq!(jac.nrows, 3);
assert_eq!(jac.ncols, 2);
let h = 1e-7;
for (col, &p_idx) in free.iter().enumerate() {
let mut p_plus = params;
let mut p_minus = params;
p_plus[p_idx] += h;
p_minus[p_idx] -= h;
let y_plus = model.evaluate(&p_plus).unwrap();
let y_minus = model.evaluate(&p_minus).unwrap();
for i in 0..3 {
let fd = (y_plus[i] - y_minus[i]) / (2.0 * h);
let ana = jac.get(i, col);
let err = (fd - ana).abs();
let scale = fd.abs().max(ana.abs()).max(1e-10);
assert!(
err / scale < 1e-4,
"Jacobian mismatch (row {i}, col {col}): FD={fd:.8}, analytical={ana:.8}"
);
}
}
}
#[test]
fn normalized_fit_recovers_anorm_and_backa() {
let xs = vec![vec![1.0, 2.0, 3.0, 2.0, 1.5]];
let inner = make_precomputed(xs, vec![0]);
let energies = [4.0, 9.0, 16.0, 25.0, 36.0];
let model = NormalizedTransmissionModel::new(inner, &energies, 1, 2, 3, 4);
let true_density = 0.2;
let true_anorm = 0.95;
let true_back_a = 0.02;
let true_params = [true_density, true_anorm, true_back_a, 0.0, 0.0];
let y_obs = model.evaluate(&true_params).unwrap();
let sigma = vec![0.001; y_obs.len()];
let mut params = ParameterSet::new(vec![
FitParameter::non_negative("density", 0.1),
FitParameter {
name: "anorm".into(),
value: 1.0,
lower: 0.5,
upper: 1.5,
fixed: false,
},
FitParameter::unbounded("back_a", 0.0),
FitParameter::fixed("back_b", 0.0),
FitParameter::fixed("back_c", 0.0),
]);
let config = LmConfig {
max_iter: 200,
..LmConfig::default()
};
let result = lm::levenberg_marquardt(&model, &y_obs, &sigma, &mut params, &config).unwrap();
assert!(result.converged, "Fit should converge");
let fit_density = result.params[0];
let fit_anorm = result.params[1];
let fit_back_a = result.params[2];
assert!(
(fit_density - true_density).abs() / true_density < 0.01,
"density: fitted={fit_density}, true={true_density}"
);
assert!(
(fit_anorm - true_anorm).abs() / true_anorm < 0.01,
"anorm: fitted={fit_anorm}, true={true_anorm}"
);
assert!(
(fit_back_a - true_back_a).abs() < 0.001,
"back_a: fitted={fit_back_a}, true={true_back_a}"
);
}
#[test]
fn forward_model_predict_equals_fit_model_evaluate_precomputed() {
use crate::forward_model::ForwardModel;
let xs = vec![vec![1.0, 2.0, 3.0, 2.0, 1.5]];
let model = make_precomputed(xs, vec![0]);
let params = [0.001];
let fm_result = model.evaluate(¶ms).unwrap();
let fwd_result = model.predict(¶ms).unwrap();
assert_eq!(fm_result, fwd_result);
assert_eq!(model.n_data(), 5);
assert_eq!(model.n_params(), 1);
}
#[test]
fn forward_model_predict_equals_fit_model_evaluate_normalized() {
use crate::forward_model::ForwardModel;
let xs = vec![vec![1.0, 2.0, 3.0, 2.0, 1.5]];
let inner = make_precomputed(xs, vec![0]);
let energies = [4.0, 9.0, 16.0, 25.0, 36.0];
let model = NormalizedTransmissionModel::new(inner, &energies, 1, 2, 3, 4);
let params = [0.001, 0.95, 0.01, 0.0, 0.0];
let fm_result = model.evaluate(¶ms).unwrap();
let fwd_result = model.predict(¶ms).unwrap();
assert_eq!(fm_result, fwd_result);
assert_eq!(model.n_data(), 5);
assert_eq!(model.n_params(), 5);
}
#[test]
fn forward_model_jacobian_columns_match_precomputed() {
use crate::forward_model::ForwardModel;
let xs = vec![vec![1.0, 2.0, 3.0], vec![0.5, 1.5, 2.5]];
let model = make_precomputed(xs, vec![0, 1]);
let params = [0.001, 0.002];
let y = model.predict(¶ms).unwrap();
let free_indices = vec![0, 1];
let jac = model
.jacobian(¶ms, &free_indices, &y)
.expect("analytical jacobian should be available");
assert_eq!(jac.len(), 2); assert_eq!(jac[0].len(), 3); }
#[test]
fn precomputed_with_resolution_matches_forward_model() {
use nereids_physics::resolution::ResolutionFunction;
let data = u238_single_resonance();
let thickness = 0.0005;
let temperature = 300.0;
let energies: Vec<f64> = (0..401).map(|i| 4.0 + (i as f64) * 0.015).collect();
let inst = Arc::new(InstrumentParams {
resolution: ResolutionFunction::Gaussian(
nereids_physics::resolution::ResolutionParams::new(25.0, 0.5, 0.005, 0.0).unwrap(),
),
});
let sample = SampleParams::new(temperature, vec![(data.clone(), thickness)]).unwrap();
let t_forward = transmission::forward_model(&energies, &sample, Some(&inst)).unwrap();
let xs = transmission::broadened_cross_sections(
&energies,
std::slice::from_ref(&data),
temperature,
Some(&inst), None,
)
.unwrap();
let model = PrecomputedTransmissionModel {
cross_sections: Arc::new(xs),
density_indices: Arc::new(vec![0]),
energies: Some(Arc::new(energies.clone())),
instrument: Some(Arc::clone(&inst)),
resolution_plan: None,
sparse_cubature_plan: None,
sparse_scalar_plan: None,
work_layout: None,
};
let t_precomputed = model.evaluate(&[thickness]).unwrap();
let interior = 20..energies.len() - 20;
let mut max_err = 0.0f64;
for i in interior {
let err = (t_forward[i] - t_precomputed[i]).abs();
max_err = max_err.max(err);
}
assert!(
max_err < 0.02,
"PrecomputedTransmissionModel with resolution should match \
forward_model. Max error = {max_err}"
);
}
#[test]
fn precomputed_without_resolution_unchanged() {
let model_no_res = make_precomputed(
vec![vec![100.0, 200.0, 50.0]], vec![0],
);
let params = [0.001f64]; let t = model_no_res.evaluate(¶ms).unwrap();
let expected: Vec<f64> = [100.0, 200.0, 50.0]
.iter()
.map(|&sigma| (-params[0] * sigma).exp())
.collect();
for (i, (&ti, &ei)) in t.iter().zip(expected.iter()).enumerate() {
assert!(
(ti - ei).abs() < 1e-14,
"No-resolution mismatch at bin {i}: got {ti}, expected {ei}"
);
}
let y = model_no_res.evaluate(¶ms).unwrap();
assert!(
model_no_res
.analytical_jacobian(¶ms, &[0], &y)
.is_some(),
"Analytical Jacobian must be available when instrument is None"
);
}
#[test]
fn precomputed_jacobian_with_resolution_matches_fd() {
use nereids_physics::resolution::ResolutionFunction;
let data = u238_single_resonance();
let temperature = 300.0;
let energies: Vec<f64> = (0..201).map(|i| 4.0 + (i as f64) * 0.025).collect();
let inst = Arc::new(InstrumentParams {
resolution: ResolutionFunction::Gaussian(
nereids_physics::resolution::ResolutionParams::new(25.0, 0.5, 0.005, 0.0).unwrap(),
),
});
let xs = transmission::broadened_cross_sections(
&energies,
std::slice::from_ref(&data),
temperature,
Some(&inst),
None,
)
.unwrap();
let model = PrecomputedTransmissionModel {
cross_sections: Arc::new(xs),
density_indices: Arc::new(vec![0]),
energies: Some(Arc::new(energies.clone())),
instrument: Some(Arc::clone(&inst)),
resolution_plan: None,
sparse_cubature_plan: None,
sparse_scalar_plan: None,
work_layout: None,
};
let params = [0.0005f64];
let y = model.evaluate(¶ms).unwrap();
let jac = model
.analytical_jacobian(¶ms, &[0], &y)
.expect("analytical Jacobian must be available with resolution");
let h = 1e-7;
let y_plus = model.evaluate(&[params[0] + h]).unwrap();
let y_minus = model.evaluate(&[params[0] - h]).unwrap();
let interior = 20..energies.len() - 20;
let mut max_rel_err = 0.0f64;
for i in interior {
let fd = (y_plus[i] - y_minus[i]) / (2.0 * h);
let ana = jac.get(i, 0);
let denom = fd.abs().max(ana.abs()).max(1e-30);
max_rel_err = max_rel_err.max((ana - fd).abs() / denom);
}
assert!(
max_rel_err < 0.01,
"PrecomputedTM analytical Jacobian with resolution vs FD: \
max relative error = {max_rel_err}"
);
}
#[test]
fn precomputed_jacobian_grouped_with_resolution_matches_fd() {
use nereids_physics::resolution::ResolutionFunction;
let energies: Vec<f64> = (0..100).map(|i| 1.0 + i as f64 * 0.1).collect();
let inst = Arc::new(InstrumentParams {
resolution: ResolutionFunction::Gaussian(
nereids_physics::resolution::ResolutionParams::new(25.0, 0.5, 0.005, 0.0).unwrap(),
),
});
let xs = vec![vec![10.0; 100], vec![5.0; 100]];
let model = PrecomputedTransmissionModel {
cross_sections: Arc::new(xs),
density_indices: Arc::new(vec![0, 0]), energies: Some(Arc::new(energies.clone())),
instrument: Some(Arc::clone(&inst)),
resolution_plan: None,
sparse_cubature_plan: None,
sparse_scalar_plan: None,
work_layout: None,
};
let params = [0.001f64];
let y = model.evaluate(¶ms).unwrap();
let jac = model
.analytical_jacobian(¶ms, &[0], &y)
.expect("analytical Jacobian must be available");
let h = 1e-7;
let y_plus = model.evaluate(&[params[0] + h]).unwrap();
let y_minus = model.evaluate(&[params[0] - h]).unwrap();
let mut max_rel_err = 0.0f64;
for i in 10..energies.len() - 10 {
let fd = (y_plus[i] - y_minus[i]) / (2.0 * h);
let ana = jac.get(i, 0);
let denom = fd.abs().max(ana.abs()).max(1e-30);
max_rel_err = max_rel_err.max((ana - fd).abs() / denom);
}
assert!(
max_rel_err < 0.01,
"Grouped PrecomputedTM analytical Jacobian with resolution vs FD: \
max relative error = {max_rel_err}"
);
}
#[test]
fn transmission_fit_model_jacobian_with_resolution_matches_fd() {
use nereids_physics::resolution::ResolutionFunction;
let data = u238_single_resonance();
let energies: Vec<f64> = (0..201).map(|i| 4.0 + (i as f64) * 0.025).collect();
let inst = Arc::new(InstrumentParams {
resolution: ResolutionFunction::Gaussian(
nereids_physics::resolution::ResolutionParams::new(25.0, 0.5, 0.005, 0.0).unwrap(),
),
});
let model = TransmissionFitModel::new(
energies.clone(),
vec![data],
300.0,
Some(inst),
(vec![0], vec![1.0]),
Some(1), None,
)
.unwrap();
let params = [0.0005f64, 300.0];
let y = model.evaluate(¶ms).unwrap();
let free = vec![0usize, 1usize];
let jac = model
.analytical_jacobian(¶ms, &free, &y)
.expect("analytical Jacobian must be available with resolution");
let h_density = 1e-7;
let h_temp = 0.01;
for (col, (&fp_idx, &h)) in free.iter().zip([h_density, h_temp].iter()).enumerate() {
let mut p_plus = params;
let mut p_minus = params;
p_plus[fp_idx] += h;
p_minus[fp_idx] -= h;
let y_plus = model.evaluate(&p_plus).unwrap();
let y_minus = model.evaluate(&p_minus).unwrap();
let interior = 20..energies.len() - 20;
let mut max_rel_err = 0.0f64;
for i in interior {
let fd = (y_plus[i] - y_minus[i]) / (2.0 * h);
let ana = jac.get(i, col);
let denom = fd.abs().max(ana.abs()).max(1e-30);
max_rel_err = max_rel_err.max((ana - fd).abs() / denom);
}
let label = if col == 0 { "density" } else { "temperature" };
assert!(
max_rel_err < 0.05,
"TransmissionFitModel {label} column with resolution vs FD: \
max relative error = {max_rel_err}"
);
}
}
#[test]
fn transmission_fit_model_jacobian_available_without_resolution() {
let data = u238_single_resonance();
let energies: Vec<f64> = (0..101).map(|i| 4.0 + (i as f64) * 0.05).collect();
let model = TransmissionFitModel::new(
energies,
vec![data],
300.0,
None,
(vec![0], vec![1.0]),
Some(1),
None,
)
.unwrap();
let params = [0.0005, 300.0];
let y = model.evaluate(¶ms).unwrap();
assert!(
model.analytical_jacobian(¶ms, &[0, 1], &y).is_some(),
"TransmissionFitModel analytical Jacobian must be available \
when resolution is disabled"
);
}
#[test]
fn transmission_fit_model_temp_path_with_resolution_matches_forward_model() {
use nereids_physics::resolution::ResolutionFunction;
let data = u238_single_resonance();
let thickness = 0.0005;
let temperature = 300.0;
let energies: Vec<f64> = (0..401).map(|i| 4.0 + (i as f64) * 0.015).collect();
let inst = Arc::new(InstrumentParams {
resolution: ResolutionFunction::Gaussian(
nereids_physics::resolution::ResolutionParams::new(25.0, 0.5, 0.005, 0.0).unwrap(),
),
});
let sample = SampleParams::new(temperature, vec![(data.clone(), thickness)]).unwrap();
let t_ref = transmission::forward_model(&energies, &sample, Some(&inst)).unwrap();
let model = TransmissionFitModel::new(
energies.clone(),
vec![data],
temperature,
Some(Arc::clone(&inst)),
(vec![0], vec![1.0]),
Some(1), None,
)
.unwrap();
let t_model = model.evaluate(&[thickness, temperature]).unwrap();
let interior = 20..energies.len() - 20;
let mut max_err = 0.0f64;
for i in interior {
max_err = max_err.max((t_ref[i] - t_model[i]).abs());
}
assert!(
max_err < 0.02,
"TransmissionFitModel temperature path with resolution should match \
forward_model. Max error = {max_err}"
);
}
#[test]
fn transmission_fit_model_temp_path_no_resolution_unchanged() {
let data = u238_single_resonance();
let thickness = 0.0005;
let temperature = 300.0;
let energies: Vec<f64> = (0..201).map(|i| 4.0 + (i as f64) * 0.025).collect();
let sample = SampleParams::new(temperature, vec![(data.clone(), thickness)]).unwrap();
let t_ref = transmission::forward_model(&energies, &sample, None).unwrap();
let model = TransmissionFitModel::new(
energies.clone(),
vec![data],
temperature,
None,
(vec![0], vec![1.0]),
Some(1),
None,
)
.unwrap();
let t_model = model.evaluate(&[thickness, temperature]).unwrap();
for (i, (&r, &m)) in t_ref.iter().zip(t_model.iter()).enumerate() {
assert!(
(r - m).abs() < 1e-12,
"No-resolution mismatch at E[{i}]={}: ref={r}, model={m}",
energies[i]
);
}
}
#[test]
fn issue_608_precomputed_aux_grid_resolution_matches_forward_model() {
use nereids_physics::resolution::ResolutionFunction;
let data = u238_single_resonance();
let thickness = 0.0005;
let temperature = 300.0;
let energies: Vec<f64> = (0..401).map(|i| 4.0 + (i as f64) * 0.015).collect();
let inst = Arc::new(InstrumentParams {
resolution: ResolutionFunction::Gaussian(
nereids_physics::resolution::ResolutionParams::new(25.0, 0.5, 0.005, 0.0).unwrap(),
),
});
let sample = SampleParams::new(temperature, vec![(data.clone(), thickness)]).unwrap();
let t_ref = transmission::forward_model(&energies, &sample, Some(&inst)).unwrap();
let t_nores = transmission::forward_model(&energies, &sample, None).unwrap();
let broaden = max_abs_diff(&t_ref, &t_nores);
assert!(
broaden > 1e-3 * max_abs(&t_nores),
"resolution kernel must broaden the spectrum non-trivially (got {broaden:.3e})"
);
let working = transmission::broadened_cross_sections_on_working_grid(
&energies,
std::slice::from_ref(&data),
temperature,
Some(&inst),
None,
)
.unwrap();
assert!(
!working.layout.is_identity(),
"Gaussian resolution must build a non-identity auxiliary grid — else \
this test does not exercise the #608 fix"
);
let model_fixed = PrecomputedTransmissionModel {
cross_sections: Arc::new(working.sigma),
density_indices: Arc::new(vec![0]),
energies: Some(Arc::new(energies.clone())),
instrument: Some(Arc::clone(&inst)),
resolution_plan: None,
sparse_cubature_plan: None,
sparse_scalar_plan: None,
work_layout: Some(Arc::new(working.layout)),
};
let t_fixed = model_fixed.evaluate(&[thickness]).unwrap();
let xs_data = transmission::broadened_cross_sections(
&energies,
std::slice::from_ref(&data),
temperature,
Some(&inst),
None,
)
.unwrap();
let model_old = PrecomputedTransmissionModel {
cross_sections: Arc::new(xs_data),
density_indices: Arc::new(vec![0]),
energies: Some(Arc::new(energies.clone())),
instrument: Some(Arc::clone(&inst)),
resolution_plan: None,
sparse_cubature_plan: None,
sparse_scalar_plan: None,
work_layout: None,
};
let t_old = model_old.evaluate(&[thickness]).unwrap();
let err_fixed = max_abs_diff(&t_fixed, &t_ref);
let err_old = max_abs_diff(&t_old, &t_ref);
assert!(
err_fixed < 1e-9,
"aux-grid PrecomputedTransmissionModel must match forward_model to \
machine precision over the full grid (got {err_fixed:.3e})"
);
assert!(
err_old > 1e-4 && err_old > 1e4 * err_fixed.max(1e-15),
"old coarse-grid path should differ from forward_model far more than \
the fixed path (old={err_old:.3e}, fixed={err_fixed:.3e})"
);
}
#[test]
fn issue_608_precomputed_aux_grid_jacobian_matches_fd() {
use nereids_physics::resolution::ResolutionFunction;
let data = u238_single_resonance();
let thickness = 0.0005;
let temperature = 300.0;
let energies: Vec<f64> = (0..401).map(|i| 4.0 + (i as f64) * 0.015).collect();
let inst = Arc::new(InstrumentParams {
resolution: ResolutionFunction::Gaussian(
nereids_physics::resolution::ResolutionParams::new(25.0, 0.5, 0.005, 0.0).unwrap(),
),
});
let working = transmission::broadened_cross_sections_on_working_grid(
&energies,
std::slice::from_ref(&data),
temperature,
Some(&inst),
None,
)
.unwrap();
let model = PrecomputedTransmissionModel {
cross_sections: Arc::new(working.sigma),
density_indices: Arc::new(vec![0]),
energies: Some(Arc::new(energies.clone())),
instrument: Some(Arc::clone(&inst)),
resolution_plan: None,
sparse_cubature_plan: None,
sparse_scalar_plan: None,
work_layout: Some(Arc::new(working.layout)),
};
let params = [thickness];
let free = [0usize];
let y0 = model.evaluate(¶ms).unwrap();
let jac = model
.analytical_jacobian(¶ms, &free, &y0)
.expect("analytical jacobian must be available with resolution + aux grid");
let h = 1e-7;
let mut pp = params;
let mut pm = params;
pp[0] += h;
pm[0] -= h;
let yp = model.evaluate(&pp).unwrap();
let ym = model.evaluate(&pm).unwrap();
let mut scale = 0.0f64;
let mut max_err = 0.0f64;
for i in 0..y0.len() {
let fd = (yp[i] - ym[i]) / (2.0 * h);
let an = jac.get(i, 0);
scale = scale.max(an.abs());
max_err = max_err.max((fd - an).abs());
}
let rel = max_err / scale.max(1e-30);
assert!(
rel < 1e-6,
"analytical density Jacobian must match central FD (rel err {rel:.3e})"
);
}
#[test]
fn issue_608_transmission_fit_temp_path_aux_grid_matches_forward_model() {
use nereids_physics::resolution::ResolutionFunction;
let data = u238_single_resonance();
let thickness = 0.0005;
let temperature = 300.0;
let energies: Vec<f64> = (0..401).map(|i| 4.0 + (i as f64) * 0.015).collect();
let inst = Arc::new(InstrumentParams {
resolution: ResolutionFunction::Gaussian(
nereids_physics::resolution::ResolutionParams::new(25.0, 0.5, 0.005, 0.0).unwrap(),
),
});
let sample = SampleParams::new(temperature, vec![(data.clone(), thickness)]).unwrap();
let t_ref = transmission::forward_model(&energies, &sample, Some(&inst)).unwrap();
let t_nores = transmission::forward_model(&energies, &sample, None).unwrap();
let broaden = max_abs_diff(&t_ref, &t_nores);
assert!(
broaden > 1e-3 * max_abs(&t_nores),
"resolution kernel must broaden the spectrum non-trivially (got {broaden:.3e})"
);
let model = TransmissionFitModel::new(
energies.clone(),
vec![data],
temperature,
Some(Arc::clone(&inst)),
(vec![0], vec![1.0]),
Some(1), None,
)
.unwrap();
let t_model = model.evaluate(&[thickness, temperature]).unwrap();
let err = max_abs_diff(&t_model, &t_ref);
assert!(
err < 1e-9,
"aux-grid TransmissionFitModel temperature path must match \
forward_model over the full grid (got {err:.3e})"
);
}
#[test]
fn issue_608_transmission_fit_temp_path_jacobian_matches_fd() {
use nereids_physics::resolution::ResolutionFunction;
let data = u238_single_resonance();
let thickness = 0.0005;
let temperature = 300.0;
let energies: Vec<f64> = (0..401).map(|i| 4.0 + (i as f64) * 0.015).collect();
let inst = Arc::new(InstrumentParams {
resolution: ResolutionFunction::Gaussian(
nereids_physics::resolution::ResolutionParams::new(25.0, 0.5, 0.005, 0.0).unwrap(),
),
});
let model = TransmissionFitModel::new(
energies.clone(),
vec![data],
temperature,
Some(Arc::clone(&inst)),
(vec![0], vec![1.0]),
Some(1),
None,
)
.unwrap();
let params = [thickness, temperature];
let y0 = model.evaluate(¶ms).unwrap();
let free = [0usize, 1usize];
let jac = model
.analytical_jacobian(¶ms, &free, &y0)
.expect("analytical jacobian must be available with resolution + aux grid");
let steps = [1e-7, 1e-2];
for (col, &p_idx) in free.iter().enumerate() {
let h = steps[col];
let mut pp = params;
let mut pm = params;
pp[p_idx] += h;
pm[p_idx] -= h;
let yp = model.evaluate(&pp).unwrap();
let ym = model.evaluate(&pm).unwrap();
let mut scale = 0.0f64;
let mut max_err = 0.0f64;
for i in 0..y0.len() {
let fd = (yp[i] - ym[i]) / (2.0 * h);
let an = jac.get(i, col);
scale = scale.max(an.abs());
max_err = max_err.max((fd - an).abs());
}
let rel = max_err / scale.max(1e-30);
assert!(
rel < 1e-5,
"analytical Jacobian column {col} must match central FD (rel err {rel:.3e})"
);
}
}
#[test]
fn issue_608_energy_scale_aux_grid_true_sigma_matches_forward_model() {
use nereids_physics::resolution::ResolutionFunction;
let data = u238_single_resonance();
let density = 0.01;
let energies: Vec<f64> = (0..121).map(|i| 5.0 + (i as f64) * 0.015).collect();
let inst = Arc::new(InstrumentParams {
resolution: ResolutionFunction::Gaussian(
nereids_physics::resolution::ResolutionParams::new(25.0, 0.5, 0.005, 0.0).unwrap(),
),
});
let model = make_energy_scale_u238(energies.clone(), Some(Arc::clone(&inst)));
let t_es = model.evaluate(&[density, 0.0, 1.0]).unwrap();
let sample = SampleParams::new(300.0, vec![(data, density)]).unwrap();
let t_ref = transmission::forward_model(&energies, &sample, Some(&inst)).unwrap();
let t_nores = transmission::forward_model(&energies, &sample, None).unwrap();
let broaden = max_abs_diff(&t_ref, &t_nores);
assert!(
broaden > 1e-3 * max_abs(&t_nores),
"resolution kernel must broaden the spectrum non-trivially (got {broaden:.3e})"
);
let err = max_abs_diff(&t_es, &t_ref);
assert!(
err < 1e-9,
"EnergyScale identity-calibration evaluate must match forward_model to \
machine precision over the full grid (got {err:.3e})"
);
}
#[test]
fn issue_608_energy_scale_grouped_density_matches_forward_model() {
use nereids_endf::resonance::test_support::synthetic_single_resonance;
use nereids_physics::resolution::ResolutionFunction;
let iso0 = u238_single_resonance(); let iso1 = synthetic_single_resonance(72, 178, 176.0, 7.5); let density = 0.01_f64;
let ratios = [0.7_f64, 0.3_f64];
let energies: Vec<f64> = (0..201).map(|i| 5.0 + (i as f64) * 0.02).collect();
let inst = Arc::new(InstrumentParams {
resolution: ResolutionFunction::Gaussian(
nereids_physics::resolution::ResolutionParams::new(25.0, 0.5, 0.005, 0.0).unwrap(),
),
});
let model = EnergyScaleTransmissionModel::new(
Arc::new(vec![iso0.clone(), iso1.clone()]),
Arc::new(vec![0, 0]), Arc::new(vec![ratios[0], ratios[1]]),
300.0,
energies.clone(),
25.0,
1, 2, Some(Arc::clone(&inst)),
);
let params = [density, 0.0, 1.0]; let t_es = model.evaluate(¶ms).unwrap();
let sample = SampleParams::new(
300.0,
vec![
(iso0.clone(), density * ratios[0]),
(iso1.clone(), density * ratios[1]),
],
)
.unwrap();
let t_ref = transmission::forward_model(&energies, &sample, Some(&inst)).unwrap();
let t_nores = transmission::forward_model(&energies, &sample, None).unwrap();
assert!(
max_abs_diff(&t_ref, &t_nores) > 1e-3 * max_abs(&t_nores),
"resolution kernel must broaden the grouped spectrum non-trivially"
);
let model_swapped = EnergyScaleTransmissionModel::new(
Arc::new(vec![iso0.clone(), iso1.clone()]),
Arc::new(vec![0, 0]),
Arc::new(vec![ratios[1], ratios[0]]), 300.0,
energies.clone(),
25.0,
1,
2,
Some(Arc::clone(&inst)),
);
let t_swapped = model_swapped.evaluate(¶ms).unwrap();
assert!(
max_abs_diff(&t_es, &t_swapped) > 1e-4,
"swapping the two density ratios must change T (else the test could \
not distinguish the ratio→isotope assignment)"
);
let err = max_abs_diff(&t_es, &t_ref);
assert!(
err < 1e-9,
"grouped EnergyScale (2 isotopes → 1 density param, ratios {ratios:?}) \
must match forward_model with per-isotope effective densities to \
machine precision (got {err:.3e})"
);
let free = vec![0usize];
let jac = model
.analytical_jacobian(¶ms, &free, &t_es)
.expect("Jacobian should be available");
let h = 1e-7;
let mut pp = params;
let mut pm = params;
pp[0] += h;
pm[0] -= h;
let yp = model.evaluate(&pp).unwrap();
let ym = model.evaluate(&pm).unwrap();
for row in 0..energies.len() {
let fd = (yp[row] - ym[row]) / (2.0 * h);
let anal = jac.get(row, 0);
let abs_err = (anal - fd).abs();
let rel_err = abs_err / fd.abs().max(1e-15);
assert!(
rel_err < 1e-3 || abs_err < 1e-8,
"grouped density col bin {row}: anal={anal:.6e} fd={fd:.6e} rel={rel_err:.2e}"
);
}
}
#[test]
fn transmission_fit_model_temp_path_resolution_makes_difference() {
use nereids_physics::resolution::ResolutionFunction;
let data = u238_single_resonance();
let thickness = 0.0005;
let temperature = 300.0;
let energies: Vec<f64> = (0..401).map(|i| 4.0 + (i as f64) * 0.015).collect();
let inst = Arc::new(InstrumentParams {
resolution: ResolutionFunction::Gaussian(
nereids_physics::resolution::ResolutionParams::new(25.0, 0.5, 0.005, 0.0).unwrap(),
),
});
let model_res = TransmissionFitModel::new(
energies.clone(),
vec![data.clone()],
temperature,
Some(inst),
(vec![0], vec![1.0]),
Some(1),
None,
)
.unwrap();
let t_res = model_res.evaluate(&[thickness, temperature]).unwrap();
let model_no = TransmissionFitModel::new(
energies.clone(),
vec![data],
temperature,
None,
(vec![0], vec![1.0]),
Some(1),
None,
)
.unwrap();
let t_no = model_no.evaluate(&[thickness, temperature]).unwrap();
let interior = 20..energies.len() - 20;
let max_diff: f64 = interior
.map(|i| (t_res[i] - t_no[i]).abs())
.fold(0.0f64, f64::max);
assert!(
max_diff > 1e-4,
"Resolution should make a measurable difference in the temperature \
path, but max diff = {max_diff}"
);
}
#[test]
fn exponential_evaluate_formula_correct() {
let xs = vec![vec![1.0, 2.0, 3.0]];
let inner = make_precomputed(xs, vec![0]);
let energies = [4.0, 9.0, 25.0];
let model =
NormalizedTransmissionModel::new_with_exponential(inner, &energies, 1, 2, 3, 4, 5, 6);
let density = 0.1;
let anorm = 1.02;
let back_a = 0.01;
let back_b = 0.005;
let back_c = 0.002;
let back_d = 0.05;
let back_f = 3.0;
let params = [density, anorm, back_a, back_b, back_c, back_d, back_f];
let y = model.evaluate(¶ms).unwrap();
let xs_vals = [1.0, 2.0, 3.0];
let sqrt_e = [2.0, 3.0, 5.0];
for i in 0..3 {
let t_inner = (-density * xs_vals[i]).exp();
let expected = anorm * t_inner
+ back_a
+ back_b / sqrt_e[i]
+ back_c * sqrt_e[i]
+ back_d * (-back_f / sqrt_e[i]).exp();
assert!(
(y[i] - expected).abs() < 1e-12,
"bin {i}: got {}, expected {expected}",
y[i]
);
}
}
#[test]
fn exponential_jacobian_matches_finite_difference() {
let xs = vec![vec![1.0, 2.0, 3.0, 0.5, 1.5]];
let inner = make_precomputed(xs, vec![0]);
let energies = [0.1, 1.0, 4.0, 25.0, 100.0];
let model =
NormalizedTransmissionModel::new_with_exponential(inner, &energies, 1, 2, 3, 4, 5, 6);
let params = [0.1, 1.02, 0.01, 0.005, 0.002, 0.05, 3.0];
let y = model.evaluate(¶ms).unwrap();
let free_indices: Vec<usize> = (0..7).collect();
let jac = model
.analytical_jacobian(¶ms, &free_indices, &y)
.expect("analytical Jacobian should be available");
let h = 1e-7;
for (col, &pidx) in free_indices.iter().enumerate() {
let mut p_plus = params.to_vec();
let mut p_minus = params.to_vec();
p_plus[pidx] += h;
p_minus[pidx] -= h;
let y_plus = model.evaluate(&p_plus).unwrap();
let y_minus = model.evaluate(&p_minus).unwrap();
for row in 0..energies.len() {
let fd = (y_plus[row] - y_minus[row]) / (2.0 * h);
let anal = jac.get(row, col);
let abs_err = (anal - fd).abs();
let rel_err = abs_err / fd.abs().max(1e-15);
assert!(
rel_err < 1e-5 || abs_err < 1e-10,
"param {pidx} (col {col}), bin {row}: analytical={anal:.10e}, \
fd={fd:.10e}, rel_err={rel_err:.2e}"
);
}
}
}
#[test]
fn exponential_fit_recovers_all_params() {
let xs = vec![vec![1.0, 2.0, 3.0, 2.0, 1.5, 0.8, 1.2, 2.5]];
let inner = make_precomputed(xs, vec![0]);
let energies = [0.5, 1.0, 4.0, 9.0, 16.0, 25.0, 36.0, 64.0];
let model =
NormalizedTransmissionModel::new_with_exponential(inner, &energies, 1, 2, 3, 4, 5, 6);
let true_density = 0.15;
let true_anorm = 1.02;
let true_back_a = 0.01;
let true_back_b = 0.005;
let true_back_c = 0.002;
let true_back_d = 0.03;
let true_back_f = 2.0;
let true_params = [
true_density,
true_anorm,
true_back_a,
true_back_b,
true_back_c,
true_back_d,
true_back_f,
];
let y_obs = model.evaluate(&true_params).unwrap();
let sigma = vec![0.001; y_obs.len()];
let mut params = ParameterSet::new(vec![
FitParameter::non_negative("density", 0.1),
FitParameter {
name: "anorm".into(),
value: 1.0,
lower: 0.5,
upper: 1.5,
fixed: false,
},
FitParameter {
name: "back_a".into(),
value: 0.0,
lower: -0.5,
upper: 0.5,
fixed: false,
},
FitParameter {
name: "back_b".into(),
value: 0.0,
lower: -0.5,
upper: 0.5,
fixed: false,
},
FitParameter {
name: "back_c".into(),
value: 0.0,
lower: -0.5,
upper: 0.5,
fixed: false,
},
FitParameter {
name: "back_d".into(),
value: 0.01,
lower: 0.0,
upper: 1.0,
fixed: false,
},
FitParameter {
name: "back_f".into(),
value: 1.0,
lower: 0.0,
upper: 100.0,
fixed: false,
},
]);
let config = LmConfig {
max_iter: 500,
..LmConfig::default()
};
let result = lm::levenberg_marquardt(&model, &y_obs, &sigma, &mut params, &config).unwrap();
assert!(result.converged, "Fit should converge");
let fitted = &result.params;
let check = |name, fitted_val: f64, true_val: f64, tol: f64| {
let err = (fitted_val - true_val).abs();
let rel = err / true_val.abs().max(1e-10);
assert!(
rel < tol || err < 1e-6,
"{name}: fitted={fitted_val:.6}, true={true_val:.6}, rel_err={rel:.4}"
);
};
check("density", fitted[0], true_density, 0.10);
check("anorm", fitted[1], true_anorm, 0.10);
check("back_a", fitted[2], true_back_a, 0.10);
check("back_b", fitted[3], true_back_b, 0.10);
check("back_c", fitted[4], true_back_c, 0.10);
check("back_d", fitted[5], true_back_d, 0.10);
check("back_f", fitted[6], true_back_f, 0.10);
}
fn make_energy_scale_u238(
energies: Vec<f64>,
instrument: Option<Arc<InstrumentParams>>,
) -> EnergyScaleTransmissionModel {
EnergyScaleTransmissionModel::new(
Arc::new(vec![u238_single_resonance()]),
Arc::new(vec![0]),
Arc::new(vec![1.0]),
300.0,
energies,
25.0,
1,
2,
instrument,
)
}
#[test]
fn energy_scale_corrected_energies() {
let energies = vec![10.0, 20.0, 50.0, 100.0, 200.0];
let model = make_energy_scale_u238(energies.clone(), None);
let e_corr = model.corrected_energies(0.0, 1.0);
for (i, (&nom, &corr)) in energies.iter().zip(e_corr.iter()).enumerate() {
assert!(
(nom - corr).abs() / nom < 1e-10,
"bin {i}: nominal={nom}, corrected={corr}"
);
}
let e_corr_ls = model.corrected_energies(0.0, 1.005);
for (i, (&nom, &corr)) in energies.iter().zip(e_corr_ls.iter()).enumerate() {
assert!(
corr > nom,
"bin {i}: l_scale=1.005 should increase energy, got nom={nom}, corr={corr}"
);
}
let e_corr_t0 = model.corrected_energies(1.0, 1.0);
for (i, (&nom, &corr)) in energies.iter().zip(e_corr_t0.iter()).enumerate() {
assert!(
corr > nom,
"bin {i}: t0=1.0 should increase energy, got nom={nom}, corr={corr}"
);
}
}
#[test]
fn energy_scale_evaluate_identity() {
let energies: Vec<f64> = (0..201).map(|i| 4.0 + (i as f64) * 0.03).collect();
let density = 0.01;
let model_es = make_energy_scale_u238(energies.clone(), None);
let y_es = model_es.evaluate(&[density, 0.0, 1.0]).unwrap();
let sample = SampleParams::new(300.0, vec![(u238_single_resonance(), density)]).unwrap();
let y_ref = transmission::forward_model(&energies, &sample, None).unwrap();
for (i, (&a, &b)) in y_es.iter().zip(y_ref.iter()).enumerate() {
assert!(
(a - b).abs() < 1e-10,
"bin {i}: energy_scale={a}, forward_model={b}"
);
}
}
#[test]
fn energy_scale_jacobian_density_matches_fd() {
let energies: Vec<f64> = (0..101).map(|i| 4.0 + (i as f64) * 0.06).collect();
let model = make_energy_scale_u238(energies.clone(), None);
let params = [0.01, 0.5, 1.002]; let y = model.evaluate(¶ms).unwrap();
let free = vec![0];
let jac = model
.analytical_jacobian(¶ms, &free, &y)
.expect("Jacobian should be available");
let h = 1e-7;
for (col, &pidx) in free.iter().enumerate() {
let mut pp = params.to_vec();
let mut pm = params.to_vec();
pp[pidx] += h;
pm[pidx] -= h;
let yp = model.evaluate(&pp).unwrap();
let ym = model.evaluate(&pm).unwrap();
for row in 0..energies.len() {
let fd = (yp[row] - ym[row]) / (2.0 * h);
let anal = jac.get(row, col);
let abs_err = (anal - fd).abs();
let rel_err = abs_err / fd.abs().max(1e-15);
assert!(
rel_err < 1e-3 || abs_err < 1e-8,
"param {pidx} col {col} bin {row}: anal={anal:.6e} fd={fd:.6e} rel={rel_err:.2e}"
);
}
}
}
#[test]
fn energy_scale_fit_recovers_l_scale() {
let energies: Vec<f64> = (0..200).map(|i| 4.0 + (i as f64) * 0.03).collect();
let true_density = 0.001;
let true_ls = 1.003;
let model = make_energy_scale_u238(energies, None);
let true_params = [true_density, 0.0, true_ls];
let y_obs = model.evaluate(&true_params).unwrap();
let sigma = vec![0.001; y_obs.len()];
let mut params = ParameterSet::new(vec![
FitParameter::non_negative("density", 0.0005),
FitParameter::fixed("t0", 0.0),
FitParameter {
name: "l_scale".into(),
value: 1.0,
lower: 0.99,
upper: 1.01,
fixed: false,
},
]);
let config = LmConfig {
max_iter: 200,
..LmConfig::default()
};
let result = lm::levenberg_marquardt(&model, &y_obs, &sigma, &mut params, &config).unwrap();
assert!(result.converged, "Fit should converge");
let f = &result.params;
assert!(
(f[0] - true_density).abs() / true_density < 0.05,
"density: fitted={}, true={true_density}",
f[0]
);
assert!(
(f[2] - true_ls).abs() < 0.001,
"l_scale: fitted={}, true={true_ls}",
f[2]
);
}
#[test]
fn partial_gal_no_resolution_matches_fd2() {
let energies: Vec<f64> = (0..101).map(|i| 4.0 + (i as f64) * 0.06).collect();
let mut model = make_energy_scale_u238(energies.clone(), None)
.with_jacobian_method(EnergyScaleJacobianMethod::FiniteDifference);
let params = [0.001, 0.05, 1.002]; let free = vec![0, 1, 2];
let jac_fd2 = model
.analytical_jacobian(¶ms, &free, &model.evaluate(¶ms).unwrap())
.expect("FD2 Jacobian should be available");
model = model.with_jacobian_method(EnergyScaleJacobianMethod::PartialGal);
let jac_pg = model
.analytical_jacobian(¶ms, &free, &model.evaluate(¶ms).unwrap())
.expect("partial-GAL Jacobian should be available");
for i in 0..energies.len() {
let fd2 = jac_fd2.get(i, 0);
let pg = jac_pg.get(i, 0);
assert!(
(fd2 - pg).abs() < 1e-15,
"density bin {i}: fd2={fd2:.6e} pg={pg:.6e}"
);
}
for i in 0..energies.len() {
let fd2 = jac_fd2.get(i, 1);
let pg = jac_pg.get(i, 1);
assert!(
(fd2 - pg).abs() < 1e-15,
"t0 bin {i}: fd2={fd2:.6e} pg={pg:.6e}"
);
}
let mut num_sq = 0.0_f64;
let mut den_sq = 0.0_f64;
for i in 0..energies.len() {
let fd2 = jac_fd2.get(i, 2);
let pg = jac_pg.get(i, 2);
let diff = pg - fd2;
num_sq += diff * diff;
den_sq += fd2 * fd2;
}
let rel_l2 = (num_sq / den_sq.max(1e-30)).sqrt();
assert!(
rel_l2 < 2.5e-2,
"L_scale rank-1 vs FD2 rel L₂ = {rel_l2:.3e} (expected ≪ 1 without \
resolution — the rank-1 identity is exact up to FD truncation)"
);
}
#[test]
fn partial_gal_l_scale_only_falls_through_to_fd() {
let energies: Vec<f64> = (0..101).map(|i| 4.0 + (i as f64) * 0.06).collect();
let model = make_energy_scale_u238(energies.clone(), None)
.with_jacobian_method(EnergyScaleJacobianMethod::PartialGal);
let params = [0.001, 0.0, 1.002];
let free = vec![0, 2]; let y = model.evaluate(¶ms).unwrap();
let jac = model
.analytical_jacobian(¶ms, &free, &y)
.expect("Jacobian should be available even when t0 not free");
let h = 1e-7;
let mut pp = params.to_vec();
let mut pm = params.to_vec();
pp[2] += h;
pm[2] -= h;
let yp = model.evaluate(&pp).unwrap();
let ym = model.evaluate(&pm).unwrap();
for i in 0..energies.len() {
let fd = (yp[i] - ym[i]) / (2.0 * h);
let anal = jac.get(i, 1);
let abs_err = (anal - fd).abs();
let rel_err = abs_err / fd.abs().max(1e-15);
assert!(
rel_err < 1e-3 || abs_err < 1e-8,
"L_scale bin {i}: anal={anal:.6e} fd={fd:.6e} rel={rel_err:.2e}"
);
}
}
#[test]
fn partial_gal_l_scale_zero_falls_through_to_finite_jacobian() {
let energies: Vec<f64> = (0..101).map(|i| 4.0 + (i as f64) * 0.06).collect();
let mut model = make_energy_scale_u238(energies.clone(), None)
.with_jacobian_method(EnergyScaleJacobianMethod::FiniteDifference);
let params = [0.001, 0.05, 1e-13]; let free = vec![0, 1, 2];
let jac_fd2 = model
.analytical_jacobian(¶ms, &free, &model.evaluate(¶ms).unwrap())
.expect("FD2 Jacobian should be available at l_scale = 0");
model = model.with_jacobian_method(EnergyScaleJacobianMethod::PartialGal);
let jac_pg = model
.analytical_jacobian(¶ms, &free, &model.evaluate(¶ms).unwrap())
.expect("partial-GAL Jacobian should be available at l_scale = 0 (fallthrough to FD)");
for i in 0..jac_pg.nrows {
for c in 0..jac_pg.ncols {
let v = jac_pg.get(i, c);
assert!(
v.is_finite(),
"partial-GAL Jacobian must be finite at l_scale = 0; \
got non-finite at ({i},{c}) = {v}"
);
}
}
for c in 0..jac_pg.ncols {
for i in 0..jac_pg.nrows {
let fd2 = jac_fd2.get(i, c);
let pg = jac_pg.get(i, c);
let abs_err = (fd2 - pg).abs();
let rel_err = abs_err / fd2.abs().max(1e-15);
assert!(
rel_err < 1e-3 || abs_err < 1e-8,
"partial-GAL must match FD2 at l_scale = 0; \
col {c} bin {i}: fd2={fd2:.6e} pg={pg:.6e} rel={rel_err:.2e}"
);
}
}
}
#[test]
fn partial_gal_with_resolution_matches_fd2() {
use nereids_physics::resolution::{ResolutionFunction, ResolutionParams};
const PARTIAL_GAL_REL_L2_TOLERANCE: f64 = 1.5e-5;
let energies: Vec<f64> = (0..101).map(|i| 4.0 + (i as f64) * 0.06).collect();
let instrument = Some(Arc::new(InstrumentParams {
resolution: ResolutionFunction::Gaussian(
ResolutionParams::new(25.0, 0.5, 0.005, 0.0).unwrap(),
),
}));
let mut model = make_energy_scale_u238(energies.clone(), instrument)
.with_jacobian_method(EnergyScaleJacobianMethod::FiniteDifference);
let params = [0.001, 0.05, 1.002]; let free = vec![0, 1, 2];
let model_no_resolution = make_energy_scale_u238(energies.clone(), None)
.with_jacobian_method(EnergyScaleJacobianMethod::FiniteDifference);
let t_no_res = model_no_resolution.evaluate(¶ms).unwrap();
let t_with_res = model.evaluate(¶ms).unwrap();
let diff_inf = t_no_res
.iter()
.zip(t_with_res.iter())
.map(|(a, b)| (a - b).abs())
.fold(0.0_f64, f64::max);
let t_inf = t_no_res.iter().map(|x| x.abs()).fold(0.0_f64, f64::max);
assert!(
diff_inf > 1e-3 * t_inf,
"resolution kernel must broaden the spectrum nontrivially \
(got ||T_kernel - T_none||_∞ = {diff_inf:.3e}, ||T_none||_∞ = {t_inf:.3e}, \
ratio = {ratio:.3e}); widen the kernel or sharpen the resonance",
ratio = diff_inf / t_inf.max(1e-30),
);
let jac_fd2 = model
.analytical_jacobian(¶ms, &free, &model.evaluate(¶ms).unwrap())
.expect("FD2 Jacobian should be available with resolution kernel");
model = model.with_jacobian_method(EnergyScaleJacobianMethod::PartialGal);
let jac_pg = model
.analytical_jacobian(¶ms, &free, &model.evaluate(¶ms).unwrap())
.expect("partial-GAL Jacobian should be available with resolution kernel");
for i in 0..energies.len() {
let fd2 = jac_fd2.get(i, 0);
let pg = jac_pg.get(i, 0);
let abs_err = (fd2 - pg).abs();
let rel_err = abs_err / fd2.abs().max(1e-15);
assert!(
rel_err < 1e-3 || abs_err < 1e-8,
"density bin {i}: fd2={fd2:.6e} pg={pg:.6e} rel={rel_err:.2e}"
);
}
for i in 0..energies.len() {
let fd2 = jac_fd2.get(i, 1);
let pg = jac_pg.get(i, 1);
let abs_err = (fd2 - pg).abs();
let rel_err = abs_err / fd2.abs().max(1e-15);
assert!(
rel_err < 1e-3 || abs_err < 1e-8,
"t0 bin {i}: fd2={fd2:.6e} pg={pg:.6e} rel={rel_err:.2e}"
);
}
let mut num_sq = 0.0_f64;
let mut den_sq = 0.0_f64;
for i in 0..energies.len() {
let fd2 = jac_fd2.get(i, 2);
let pg = jac_pg.get(i, 2);
let diff = pg - fd2;
num_sq += diff * diff;
den_sq += fd2 * fd2;
}
let rel_l2 = (num_sq / den_sq.max(1e-30)).sqrt();
assert!(
rel_l2 < PARTIAL_GAL_REL_L2_TOLERANCE,
"L_scale rel L₂ = {rel_l2:.4e} exceeds tolerance {tol:.4e}; \
tighten or loosen `PARTIAL_GAL_REL_L2_TOLERANCE` (see rustdoc)",
tol = PARTIAL_GAL_REL_L2_TOLERANCE,
);
}
}