gam_problem/
coefficient_prior_mean.rs

1//! Neutral programmatic prior-mean type for a coefficient penalty block.
2//!
3//! Lives in `gam-problem` so the penalty contract can carry a centering vector
4//! without depending on `solver`'s `EstimationError`. Evaluation failures are
5//! reported through the neutral [`PriorMeanError`]; callers map this into their
6//! own error flow (e.g. `EstimationError::InvalidInput`).
7
8use std::sync::Arc;
9
10use ndarray::Array1;
11
12/// Neutral error for prior-mean evaluation failures.
13///
14/// Carries the human-readable message; callers in the solver crate map this
15/// into `EstimationError::InvalidInput` to preserve end-to-end behavior.
16#[derive(Debug, Clone)]
17pub struct PriorMeanError(pub String);
18
19impl std::fmt::Display for PriorMeanError {
20    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
21        f.write_str(&self.0)
22    }
23}
24
25impl std::error::Error for PriorMeanError {}
26
27/// Programmatic prior mean for a coefficient penalty block.
28///
29/// The mean is evaluated once during penalty canonicalization and then enters
30/// the solver as the centering vector in `(beta - mean)' S (beta - mean)`.
31#[derive(Clone, Default)]
32pub enum CoefficientPriorMean {
33    #[default]
34    Zero,
35    Scalar(f64),
36    Constant(Array1<f64>),
37    Functional {
38        metadata: Array1<f64>,
39        evaluator: Arc<dyn Fn(&Array1<f64>) -> Array1<f64> + Send + Sync>,
40    },
41    /// Covariate-functional mean `mu(a) = amplitude * K(a)` for a coefficient block.
42    ///
43    /// Formula-level coefficient groups pass their row/covariate metadata as
44    /// `covariates`; the user-supplied kernel returns the block-sized basis
45    /// vector `K(a)` and the scalar amplitude supplies `alpha`.
46    KernelBasis {
47        covariates: Array1<f64>,
48        amplitude: f64,
49        kernel: Arc<dyn Fn(&Array1<f64>) -> Array1<f64> + Send + Sync>,
50    },
51}
52
53impl std::fmt::Debug for CoefficientPriorMean {
54    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
55        match self {
56            Self::Zero => f.write_str("Zero"),
57            Self::Scalar(value) => f.debug_tuple("Scalar").field(value).finish(),
58            Self::Constant(values) => f
59                .debug_tuple("Constant")
60                .field(&format_args!("len={}", values.len()))
61                .finish(),
62            Self::Functional { metadata, .. } => f
63                .debug_struct("Functional")
64                .field("metadata_len", &metadata.len())
65                .finish_non_exhaustive(),
66            Self::KernelBasis {
67                covariates,
68                amplitude,
69                ..
70            } => f
71                .debug_struct("KernelBasis")
72                .field("covariate_len", &covariates.len())
73                .field("amplitude", amplitude)
74                .finish_non_exhaustive(),
75        }
76    }
77}
78
79impl CoefficientPriorMean {
80    pub const fn scalar(value: f64) -> Self {
81        Self::Scalar(value)
82    }
83
84    pub fn constant(values: Array1<f64>) -> Self {
85        Self::Constant(values)
86    }
87
88    pub fn evaluate(&self, block_dim: usize, context: &str) -> Result<Array1<f64>, PriorMeanError> {
89        let values = match self {
90            Self::Zero => Array1::zeros(block_dim),
91            Self::Scalar(value) => {
92                if !value.is_finite() {
93                    return Err(PriorMeanError(format!(
94                        "{context}: coefficient prior mean scalar must be finite, got {value}"
95                    )));
96                }
97                Array1::from_elem(block_dim, *value)
98            }
99            Self::Constant(values) => values.clone(),
100            Self::Functional {
101                metadata,
102                evaluator,
103            } => evaluator(metadata),
104            Self::KernelBasis {
105                covariates,
106                amplitude,
107                kernel,
108            } => {
109                if !amplitude.is_finite() {
110                    return Err(PriorMeanError(format!(
111                        "{context}: coefficient prior mean amplitude must be finite, got {amplitude}"
112                    )));
113                }
114                let mut values = kernel(covariates);
115                values *= *amplitude;
116                values
117            }
118        };
119        if values.len() != block_dim {
120            return Err(PriorMeanError(format!(
121                "{context}: coefficient prior mean length must be {block_dim}, got {}",
122                values.len()
123            )));
124        }
125        if values.iter().any(|&value| !value.is_finite()) {
126            return Err(PriorMeanError(format!(
127                "{context}: coefficient prior mean contains non-finite values"
128            )));
129        }
130        Ok(values)
131    }
132}
133
134#[cfg(test)]
135mod tests {
136    use super::*;
137    use ndarray::array;
138
139    #[test]
140    fn zero_variant_returns_zeros_of_requested_len() {
141        let m = CoefficientPriorMean::Zero;
142        let v = m.evaluate(4, "ctx").unwrap();
143        assert_eq!(v.len(), 4);
144        assert!(v.iter().all(|&x| x == 0.0));
145    }
146
147    #[test]
148    fn scalar_fills_vector_with_constant() {
149        let m = CoefficientPriorMean::scalar(3.0);
150        let v = m.evaluate(3, "ctx").unwrap();
151        assert_eq!(v.len(), 3);
152        assert!(v.iter().all(|&x| x == 3.0));
153    }
154
155    #[test]
156    fn scalar_nan_returns_error() {
157        let m = CoefficientPriorMean::scalar(f64::NAN);
158        assert!(m.evaluate(2, "ctx").is_err());
159    }
160
161    #[test]
162    fn scalar_infinite_returns_error() {
163        let m = CoefficientPriorMean::scalar(f64::INFINITY);
164        assert!(m.evaluate(2, "ctx").is_err());
165    }
166
167    #[test]
168    fn constant_variant_clones_vector() {
169        let arr = array![1.0_f64, 2.0, 3.0];
170        let m = CoefficientPriorMean::constant(arr.clone());
171        let v = m.evaluate(3, "ctx").unwrap();
172        assert_eq!(v, arr);
173    }
174
175    #[test]
176    fn constant_dimension_mismatch_returns_error() {
177        let arr = array![1.0_f64, 2.0];
178        let m = CoefficientPriorMean::constant(arr);
179        assert!(m.evaluate(5, "ctx").is_err());
180    }
181
182    #[test]
183    fn default_is_zero_variant() {
184        let m = CoefficientPriorMean::default();
185        let v = m.evaluate(5, "ctx").unwrap();
186        assert!(v.iter().all(|&x| x == 0.0));
187    }
188
189    #[test]
190    fn error_message_includes_context() {
191        let m = CoefficientPriorMean::scalar(f64::NAN);
192        let err = m.evaluate(1, "myctx").unwrap_err();
193        let msg = err.to_string();
194        assert!(msg.contains("myctx"), "error should mention context: {msg}");
195    }
196}
gam_problem/coefficient_prior_mean.rs

gam_problem/
coefficient_prior_mean.rs
