gam_models/transformation_normal/
alo_replay.rs1use super::chart::{CtnRowBases, CtnRowFloors, ctn_component_sensitivity, ctn_row_geometry};
2use crate::inference::model::TransformationNormalParameterization;
3use ndarray::{Array1, Array2, ArrayView1};
4
5pub 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#[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
71pub 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 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 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 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}