Skip to main content

gam_models/transformation_normal/
alo_replay.rs

1use super::chart::{CtnRowBases, CtnRowFloors, ctn_component_sensitivity, ctn_row_geometry};
2use crate::inference::model::TransformationNormalParameterization;
3use ndarray::{Array1, Array2, ArrayView1};
4
5/// Complete local state for one saved transformation-normal likelihood row.
6pub struct TransformationNormalAloRowInput<'a> {
7    pub response_value_basis: &'a [f64],
8    pub response_derivative_basis: &'a [f64],
9    pub response_lower_basis: &'a [f64],
10    pub response_upper_basis: &'a [f64],
11    pub alpha: &'a [f64],
12    pub additive_offset: f64,
13    pub response_floor_offset: f64,
14    pub response_lower_floor_offset: f64,
15    pub response_upper_floor_offset: f64,
16    pub prior_weight: f64,
17}
18
19/// Exact negative-log-likelihood derivatives in the affine local coordinates
20/// `alpha_k(x) = covariate_row(x) beta_k`.
21#[derive(Clone, Debug, PartialEq)]
22pub struct TransformationNormalAloRowGeometry {
23    pub negative_log_likelihood: f64,
24    pub nll_score: Array1<f64>,
25    pub observed_hessian: Array2<f64>,
26}
27
28fn validate_row(input: &TransformationNormalAloRowInput<'_>) -> Result<usize, String> {
29    let dimension = input.alpha.len();
30    if dimension == 0
31        || input.response_value_basis.len() != dimension
32        || input.response_derivative_basis.len() != dimension
33        || input.response_lower_basis.len() != dimension
34        || input.response_upper_basis.len() != dimension
35    {
36        return Err(format!(
37            "transformation-normal ALO row dimension mismatch: alpha={dimension}, value={}, derivative={}, lower={}, upper={}",
38            input.response_value_basis.len(),
39            input.response_derivative_basis.len(),
40            input.response_lower_basis.len(),
41            input.response_upper_basis.len(),
42        ));
43    }
44    if !input.prior_weight.is_finite() || input.prior_weight < 0.0 {
45        return Err(format!(
46            "transformation-normal ALO prior weight must be finite and non-negative, got {}",
47            input.prior_weight
48        ));
49    }
50    if input
51        .response_value_basis
52        .iter()
53        .chain(input.response_derivative_basis)
54        .chain(input.response_lower_basis)
55        .chain(input.response_upper_basis)
56        .chain(input.alpha)
57        .copied()
58        .chain([
59            input.additive_offset,
60            input.response_floor_offset,
61            input.response_lower_floor_offset,
62            input.response_upper_floor_offset,
63        ])
64        .any(|value| !value.is_finite())
65    {
66        return Err("transformation-normal ALO row state must be finite".to_string());
67    }
68    Ok(dimension)
69}
70
71/// Replay one row of the fitted finite-support SCOP likelihood.
72///
73/// This is the row factorization of the same score and negative Hessian used by
74/// `TransformationNormalFamily`: every component is affine in direct-alpha
75/// coordinates, the monotonicity derivative floor is exact, and both
76/// transformed support endpoints contribute through the normalized Gaussian
77/// mass. Feasibility of the shape coordinates is owned by the fitted model's
78/// Khatri-Rao cone before this row replay is called.
79pub fn transformation_normal_alo_row_geometry(
80    input: TransformationNormalAloRowInput<'_>,
81) -> Result<TransformationNormalAloRowGeometry, String> {
82    let dimension = validate_row(&input)?;
83    if input.prior_weight == 0.0 {
84        return Ok(TransformationNormalAloRowGeometry {
85            negative_log_likelihood: 0.0,
86            nll_score: Array1::zeros(dimension),
87            observed_hessian: Array2::zeros((dimension, dimension)),
88        });
89    }
90
91    // One chart, one evaluator (gam#2680): the saved-model ALO replay reads the
92    // same coefficients as the fit and must read them the same way.
93    let chart = TransformationNormalParameterization::DirectAlpha;
94    let geometry = ctn_row_geometry(
95        chart,
96        ArrayView1::from(input.alpha),
97        CtnRowBases {
98            value: ArrayView1::from(input.response_value_basis),
99            derivative: ArrayView1::from(input.response_derivative_basis),
100            lower: ArrayView1::from(input.response_lower_basis),
101            upper: ArrayView1::from(input.response_upper_basis),
102        },
103        CtnRowFloors {
104            additive_offset: input.additive_offset,
105            value_floor: input.response_floor_offset,
106            lower_floor: input.response_lower_floor_offset,
107            upper_floor: input.response_upper_floor_offset,
108        },
109    );
110    let (h, h_prime, lower, upper) = (
111        geometry.h,
112        geometry.h_prime,
113        geometry.lower,
114        geometry.upper,
115    );
116    if !(h.is_finite() && h_prime.is_finite() && lower.is_finite() && upper.is_finite()) {
117        return Err(format!(
118            "transformation-normal ALO row transform is non-finite: h={h}, h_prime={h_prime}, lower={lower}, upper={upper}"
119        ));
120    }
121    if h_prime <= 0.0 {
122        return Err(format!(
123            "transformation-normal ALO row derivative must be positive, got {h_prime}"
124        ));
125    }
126    // gam#2600: the fitted likelihood is the untruncated MLT density, so the ALO
127    // replay uses the same one — `f(y) = φ(h)·h'`, with no renormalization by
128    // the mass between the saved support endpoints.
129    let weight = input.prior_weight;
130    let negative_log_likelihood =
131        weight * (0.5 * h * h + 0.5 * (2.0 * std::f64::consts::PI).ln() - h_prime.ln());
132
133    let mut dh = vec![0.0; dimension];
134    let mut dh_prime = vec![0.0; dimension];
135    let mut dlower = vec![0.0; dimension];
136    let mut dupper = vec![0.0; dimension];
137    for component in 0..dimension {
138        dh[component] =
139            ctn_component_sensitivity(chart, ArrayView1::from(input.response_value_basis), component);
140        dh_prime[component] = ctn_component_sensitivity(
141            chart,
142            ArrayView1::from(input.response_derivative_basis),
143            component,
144        );
145        dlower[component] =
146            ctn_component_sensitivity(chart, ArrayView1::from(input.response_lower_basis), component);
147        dupper[component] =
148            ctn_component_sensitivity(chart, ArrayView1::from(input.response_upper_basis), component);
149    }
150
151    let inverse_h_prime = 1.0 / h_prime;
152    let inverse_h_prime_squared = inverse_h_prime * inverse_h_prime;
153    let mut nll_score = Array1::<f64>::zeros(dimension);
154    let mut observed_hessian = Array2::<f64>::zeros((dimension, dimension));
155    for left in 0..dimension {
156        nll_score[left] = weight * (h * dh[left] - dh_prime[left] * inverse_h_prime);
157        for right in 0..dimension {
158            observed_hessian[[left, right]] = weight
159                * (dh[left] * dh[right]
160                    + dh_prime[left] * dh_prime[right] * inverse_h_prime_squared);
161        }
162    }
163    if !negative_log_likelihood.is_finite()
164        || nll_score.iter().any(|value| !value.is_finite())
165        || observed_hessian.iter().any(|value| !value.is_finite())
166    {
167        return Err("transformation-normal ALO row geometry is non-finite".to_string());
168    }
169    Ok(TransformationNormalAloRowGeometry {
170        negative_log_likelihood,
171        nll_score,
172        observed_hessian,
173    })
174}
175
176#[cfg(test)]
177mod tests {
178    use super::*;
179    use crate::transformation_normal::TRANSFORMATION_MONOTONICITY_EPS;
180
181    /// gam#2600: the replayed row density is the untruncated MLT density
182    /// `log φ(h) + log h'`. The endpoint bases and floors the row input carries
183    /// still define the certified support of the saved model, but they are no
184    /// longer a term of the likelihood — so this independent reconstruction
185    /// does not mention them, and the fixture below feeds deliberately
186    /// asymmetric ones (`[1.0, 0.1]` / `[1.0, 0.9]`, floors `−0.04` / `0.06`)
187    /// so that any endpoint contribution surviving in the production path shows
188    /// up here as a mismatch rather than cancelling.
189    fn scalar_nll(alpha: [f64; 2]) -> f64 {
190        let value = [1.0, 0.4];
191        let derivative = [0.0, 0.7];
192        let offset = -0.15;
193        let floor = 0.02;
194        let weight = 1.3;
195        let h = value[0] * alpha[0] + value[1] * alpha[1] + offset + floor;
196        let h_prime = TRANSFORMATION_MONOTONICITY_EPS
197            + derivative[0] * alpha[0]
198            + derivative[1] * alpha[1];
199        weight * (0.5 * h * h + 0.5 * (2.0 * std::f64::consts::PI).ln() - h_prime.ln())
200    }
201
202    #[test]
203    fn saved_transformation_row_geometry_matches_independent_scalar_finite_difference() {
204        let alpha: [f64; 2] = [0.25, 0.8];
205        let geometry = transformation_normal_alo_row_geometry(TransformationNormalAloRowInput {
206            response_value_basis: &[1.0, 0.4],
207            response_derivative_basis: &[0.0, 0.7],
208            response_lower_basis: &[1.0, 0.1],
209            response_upper_basis: &[1.0, 0.9],
210            alpha: &alpha,
211            additive_offset: -0.15,
212            response_floor_offset: 0.02,
213            response_lower_floor_offset: -0.04,
214            response_upper_floor_offset: 0.06,
215            prior_weight: 1.3,
216        })
217        .expect("saved transformation-normal row must replay");
218        let step = 2.0e-5;
219        let base = scalar_nll(alpha);
220        assert!((geometry.negative_log_likelihood - base).abs() <= 2.0e-13);
221        for axis in 0..2 {
222            let mut plus = alpha;
223            let mut minus = alpha;
224            plus[axis] += step;
225            minus[axis] -= step;
226            let gradient_fd = (scalar_nll(plus) - scalar_nll(minus)) / (2.0 * step);
227            assert!(
228                (geometry.nll_score[axis] - gradient_fd).abs() <= 2.0e-8,
229                "score[{axis}] analytic={} fd={gradient_fd}",
230                geometry.nll_score[axis]
231            );
232            for other in 0..2 {
233                let mut pp = alpha;
234                let mut pm = alpha;
235                let mut mp = alpha;
236                let mut mm = alpha;
237                pp[axis] += step;
238                pp[other] += step;
239                pm[axis] += step;
240                pm[other] -= step;
241                mp[axis] -= step;
242                mp[other] += step;
243                mm[axis] -= step;
244                mm[other] -= step;
245                let hessian_fd = (scalar_nll(pp) - scalar_nll(pm) - scalar_nll(mp)
246                    + scalar_nll(mm))
247                    / (4.0 * step * step);
248                assert!(
249                    (geometry.observed_hessian[[axis, other]] - hessian_fd).abs() <= 3.0e-6,
250                    "hessian[{axis},{other}] analytic={} fd={hessian_fd}",
251                    geometry.observed_hessian[[axis, other]]
252                );
253            }
254        }
255        assert!(
256            (geometry.observed_hessian[[0, 0]] - geometry.nll_score[0].powi(2)).abs() > 1.0e-3,
257            "observed curvature must remain distinct from score covariance"
258        );
259    }
260}