Skip to main content

nereids_fitting/
lm.rs

1//! Shared fitting types: `FlatMatrix` and the `FitModel` trait.
2
3use crate::error::FittingError;
4
5/// Row-major flat matrix.
6#[derive(Debug, Clone)]
7pub struct FlatMatrix {
8    /// Row-major storage: `data[i * ncols + j]`.
9    pub data: Vec<f64>,
10    /// Number of rows.
11    pub nrows: usize,
12    /// Number of columns.
13    pub ncols: usize,
14}
15
16impl FlatMatrix {
17    /// A zero-filled matrix.
18    pub fn zeros(nrows: usize, ncols: usize) -> Self {
19        let len = nrows
20            .checked_mul(ncols)
21            .expect("FlatMatrix dimensions overflow usize");
22        Self {
23            data: vec![0.0; len],
24            nrows,
25            ncols,
26        }
27    }
28
29    /// Element at `(row, col)`.
30    #[inline(always)]
31    pub fn get(&self, row: usize, col: usize) -> f64 {
32        debug_assert!(row < self.nrows && col < self.ncols);
33        self.data[row * self.ncols + col]
34    }
35
36    /// Mutable element at `(row, col)`.
37    #[inline(always)]
38    pub fn get_mut(&mut self, row: usize, col: usize) -> &mut f64 {
39        debug_assert!(row < self.nrows && col < self.ncols);
40        &mut self.data[row * self.ncols + col]
41    }
42}
43
44/// A model that can be fitted to data.
45pub trait FitModel {
46    /// The model's predictions at `params`, one per data point.
47    fn evaluate(&self, params: &[f64]) -> Result<Vec<f64>, FittingError>;
48
49    /// The Jacobian at `params` if available analytically, else `None`.
50    fn analytical_jacobian(
51        &self,
52        _params: &[f64],
53        _free_param_indices: &[usize],
54        _y_current: &[f64],
55    ) -> Option<FlatMatrix> {
56        None
57    }
58}
59
60impl<M: FitModel + ?Sized> FitModel for &M {
61    fn evaluate(&self, params: &[f64]) -> Result<Vec<f64>, FittingError> {
62        (**self).evaluate(params)
63    }
64
65    fn analytical_jacobian(
66        &self,
67        params: &[f64],
68        free_param_indices: &[usize],
69        y_current: &[f64],
70    ) -> Option<FlatMatrix> {
71        (**self).analytical_jacobian(params, free_param_indices, y_current)
72    }
73}
74
75impl<M: FitModel + ?Sized> FitModel for Box<M> {
76    fn evaluate(&self, params: &[f64]) -> Result<Vec<f64>, FittingError> {
77        (**self).evaluate(params)
78    }
79
80    fn analytical_jacobian(
81        &self,
82        params: &[f64],
83        free_param_indices: &[usize],
84        y_current: &[f64],
85    ) -> Option<FlatMatrix> {
86        (**self).analytical_jacobian(params, free_param_indices, y_current)
87    }
88}