nereids_fitting/
exact_count_model.rs

1//! Exact fixed-instrument model for a joint open/sample count likelihood.
2
3use nereids_core::constants::PIVOT_FLOOR;
4use nereids_physics::counts_response::DetectorBinResponseMatrix;
5
6use crate::error::FittingError;
7use crate::lm::{FitModel, FlatMatrix};
8
9/// Maps a true-energy transmission model into measured detector-bin ratios.
10///
11/// For a fixed response matrix `R` and fixed incident fluence weights `F`
12/// (`F_j = w_j ε(E_j) Φ(E_j)` — detector efficiency folded in, pipeline-map
13/// R5·7), the model supplied to the joint-Poisson objective is evaluated as
14///
15/// ```text
16/// T_eff,i = sum_j F_j T_j R_ij / sum_j F_j R_ij.
17/// ```
18///
19/// The open and sample arms are therefore broadened separately. This is not
20/// the invalid post-hoc shortcut `R[T]`.
21///
22/// # Dead-bin precondition
23///
24/// A detector bin whose open-arm expectation `sum_j F_j R_ij` is at or below
25/// the pivot floor is a **dead bin**: the ratio is undefined there, so
26/// `evaluate` fills it with `T_eff = 1.0` and `analytical_jacobian` zeroes
27/// its row. That filler is harmless only when the caller guarantees dead
28/// bins carry **no observed counts** — a joint-Poisson objective then ignores
29/// the bin entirely (its profiled rate is zero). Callers fitting observed
30/// data MUST enforce that invariant before constructing this model, as
31/// `fit_counts_joint_poisson` does with its occupied-bin pre-check: observed
32/// counts in a dead bin would otherwise be silently compared against an
33/// arbitrary `T = 1` instead of raising an error.
34pub struct ExactTwoArmRatioModel {
35    inner: Box<dyn FitModel>,
36    response: DetectorBinResponseMatrix,
37    scaled_fluence_weights: Vec<f64>,
38    open_expectation: Vec<f64>,
39    n_true_energies: usize,
40    n_detector_bins: usize,
41}
42
43impl ExactTwoArmRatioModel {
44    /// Build the reusable fixed-response model.
45    pub fn new(
46        inner: Box<dyn FitModel>,
47        response: DetectorBinResponseMatrix,
48        incident_fluence_weights: &[f64],
49    ) -> Result<Self, FittingError> {
50        if incident_fluence_weights.len() != response.n_true_energies() {
51            return Err(FittingError::LengthMismatch {
52                expected: response.n_true_energies(),
53                actual: incident_fluence_weights.len(),
54                field: "incident_fluence_weights",
55            });
56        }
57        for (index, &fluence) in incident_fluence_weights.iter().enumerate() {
58            if !fluence.is_finite() || fluence < 0.0 {
59                return Err(FittingError::InvalidConfig(format!(
60                    "incident_fluence_weights[{index}] must be finite and >= 0, got {fluence}"
61                )));
62            }
63        }
64        let fluence_scale = incident_fluence_weights
65            .iter()
66            .copied()
67            .fold(0.0_f64, f64::max);
68        if fluence_scale <= 0.0 {
69            return Err(FittingError::InvalidConfig(
70                "incident_fluence_weights must contain at least one positive value".into(),
71            ));
72        }
73
74        let n_true_energies = response.n_true_energies();
75        let n_detector_bins = response.n_detector_bins();
76        let scaled_fluence_weights: Vec<f64> = incident_fluence_weights
77            .iter()
78            .map(|fluence| fluence / fluence_scale)
79            .collect();
80        let mut open_expectation = vec![0.0; n_detector_bins];
81        let mut open_compensation = vec![0.0; n_detector_bins];
82        for (true_index, &fluence) in scaled_fluence_weights.iter().enumerate() {
83            for (detector_bin, probability) in response.row_entries(true_index) {
84                let value = fluence * probability;
85                compensated_add(
86                    &mut open_expectation[detector_bin],
87                    &mut open_compensation[detector_bin],
88                    value,
89                );
90            }
91        }
92        for detector_bin in 0..n_detector_bins {
93            open_expectation[detector_bin] += open_compensation[detector_bin];
94        }
95
96        Ok(Self {
97            inner,
98            response,
99            scaled_fluence_weights,
100            open_expectation,
101            n_true_energies,
102            n_detector_bins,
103        })
104    }
105
106    /// Expected open-arm source shape in detector-bin order.
107    pub fn open_expectation(&self) -> &[f64] {
108        &self.open_expectation
109    }
110
111    fn map_true_energy_values(&self, values: &[f64]) -> Result<Vec<f64>, FittingError> {
112        if values.len() != self.n_true_energies {
113            return Err(FittingError::LengthMismatch {
114                expected: self.n_true_energies,
115                actual: values.len(),
116                field: "true_energy_model",
117            });
118        }
119        let mut sample = vec![0.0; self.n_detector_bins];
120        let mut compensation = vec![0.0; self.n_detector_bins];
121        for (true_index, &value) in values.iter().enumerate() {
122            if !value.is_finite() {
123                return Err(FittingError::EvaluationFailed(format!(
124                    "true-energy model output[{true_index}] is not finite: {value}"
125                )));
126            }
127            let fluence = self.scaled_fluence_weights[true_index];
128            for (detector_bin, probability) in self.response.row_entries(true_index) {
129                compensated_add(
130                    &mut sample[detector_bin],
131                    &mut compensation[detector_bin],
132                    value * fluence * probability,
133                );
134            }
135        }
136
137        Ok(sample
138            .into_iter()
139            .zip(compensation)
140            .zip(&self.open_expectation)
141            .map(|((sum, correction), &open)| {
142                if open > PIVOT_FLOOR {
143                    (sum + correction) / open
144                } else {
145                    1.0
146                }
147            })
148            .collect())
149    }
150}
151
152impl FitModel for ExactTwoArmRatioModel {
153    fn evaluate(&self, params: &[f64]) -> Result<Vec<f64>, FittingError> {
154        let true_transmission = self.inner.evaluate(params)?;
155        self.map_true_energy_values(&true_transmission)
156    }
157
158    fn analytical_jacobian(
159        &self,
160        params: &[f64],
161        free_param_indices: &[usize],
162        _y_current: &[f64],
163    ) -> Option<FlatMatrix> {
164        let true_transmission = self.inner.evaluate(params).ok()?;
165        if true_transmission.len() != self.n_true_energies {
166            return None;
167        }
168        let inner_jacobian =
169            self.inner
170                .analytical_jacobian(params, free_param_indices, &true_transmission)?;
171        if inner_jacobian.nrows != self.n_true_energies
172            || inner_jacobian.ncols != free_param_indices.len()
173        {
174            return None;
175        }
176
177        let mut output = FlatMatrix::zeros(self.n_detector_bins, free_param_indices.len());
178        for true_index in 0..self.n_true_energies {
179            let fluence = self.scaled_fluence_weights[true_index];
180            for (detector_bin, probability) in self.response.row_entries(true_index) {
181                let open = self.open_expectation[detector_bin];
182                if open <= PIVOT_FLOOR {
183                    continue;
184                }
185                let weight = fluence * probability / open;
186                for column in 0..free_param_indices.len() {
187                    *output.get_mut(detector_bin, column) +=
188                        weight * inner_jacobian.get(true_index, column);
189                }
190            }
191        }
192        Some(output)
193    }
194}
195
196#[inline]
197fn compensated_add(sum: &mut f64, compensation: &mut f64, value: f64) {
198    let next = *sum + value;
199    if sum.abs() >= value.abs() {
200        *compensation += (*sum - next) + value;
201    } else {
202        *compensation += (value - next) + *sum;
203    }
204    *sum = next;
205}
206
207#[cfg(test)]
208mod tests {
209    use std::sync::Arc;
210
211    use nereids_physics::resolution::{ResolutionFunction, TOF_FACTOR, TabulatedResolution};
212
213    use super::*;
214
215    struct LinearTrueEnergyModel;
216
217    impl FitModel for LinearTrueEnergyModel {
218        fn evaluate(&self, params: &[f64]) -> Result<Vec<f64>, FittingError> {
219            Ok(vec![params[0], 1.0 - 0.5 * params[0]])
220        }
221
222        fn analytical_jacobian(
223            &self,
224            _params: &[f64],
225            free_param_indices: &[usize],
226            _y_current: &[f64],
227        ) -> Option<FlatMatrix> {
228            if free_param_indices != [0] {
229                return None;
230            }
231            Some(FlatMatrix {
232                data: vec![1.0, -0.5],
233                nrows: 2,
234                ncols: 1,
235            })
236        }
237    }
238
239    fn triangle_response() -> ResolutionFunction {
240        ResolutionFunction::Tabulated(Arc::new(
241            TabulatedResolution::from_kernels(
242                vec![25.0],
243                vec![(vec![-1.0, 0.0, 1.0], vec![0.0, 1.0, 0.0])],
244                25.0,
245            )
246            .expect("valid triangle response"),
247        ))
248    }
249
250    #[test]
251    fn maps_values_and_analytical_jacobian_through_separate_arms() {
252        let arrival_0 = TOF_FACTOR * 25.0 / 25.0_f64.sqrt();
253        let arrival_1 = arrival_0 + 1.0;
254        let energy_1 = (TOF_FACTOR * 25.0 / arrival_1).powi(2);
255        let response = DetectorBinResponseMatrix::new(
256            &[25.0, energy_1],
257            &[arrival_0 - 1.0, arrival_0, arrival_0 + 1.0, arrival_0 + 2.0],
258            0.0,
259            &triangle_response(),
260        )
261        .expect("valid exact response");
262        let model =
263            ExactTwoArmRatioModel::new(Box::new(LinearTrueEnergyModel), response, &[100.0, 200.0])
264                .expect("valid exact ratio model");
265
266        let values = model.evaluate(&[0.2]).expect("model evaluation");
267        let expected = [0.2, 2.0 / 3.0, 0.9];
268        for (got, want) in values.iter().zip(expected) {
269            assert!((got - want).abs() < 2.0e-12, "{got} != {want}");
270        }
271
272        let jacobian = model
273            .analytical_jacobian(&[0.2], &[0], &values)
274            .expect("analytical Jacobian");
275        let expected_jacobian = [1.0, 0.0, -0.5];
276        for (row, want) in expected_jacobian.into_iter().enumerate() {
277            assert!((jacobian.get(row, 0) - want).abs() < 2.0e-12);
278        }
279    }
280}
nereids_fitting/exact_count_model.rs

nereids_fitting/
exact_count_model.rs
