gam_terms/analytic_penalties/
normalized_gram.rs

1//! Closed-form derivatives of normalized cross-Gram energies.
2//!
3//! Both decoder coherence and subspace overlap are a scalar monomial of three
4//! polynomial statistics. Keeping that composition explicit provides their
5//! gradient, Hessian diagonal, Hessian action and contracted third derivative
6//! without coordinate probes, a third-order tensor, or automatic differentiation.
7
8use ndarray::{Array1, Array2, ArrayView1, ArrayView2, s};
9
10#[derive(Clone, Copy, Debug)]
11pub enum GramNormalization {
12    /// ||X Y'||² / (||X||² ||Y||²).
13    DecoderNorm,
14    /// ||X Y'||² / (||X X'|| ||Y Y'||).
15    SelfGramNorm,
16}
17
18pub struct NormalizedCrossGram {
19    left: Array2<f64>,
20    right: Array2<f64>,
21    cross: Array2<f64>,
22    left_gram: Array2<f64>,
23    right_gram: Array2<f64>,
24    normalization: GramNormalization,
25    gradients: [Array1<f64>; 3],
26    diagonals: [Array1<f64>; 3],
27    scalar_first: [f64; 3],
28    scalar_second: [[f64; 3]; 3],
29    scalar_third: [[[f64; 3]; 3]; 3],
30}
31
32fn flatten_pair(left: Array2<f64>, right: Array2<f64>) -> Array1<f64> {
33    Array1::from_iter(left.iter().chain(right.iter()).copied())
34}
35
36impl NormalizedCrossGram {
37    /// Positive normalizers define this smooth stratum. A zero decoder has no
38    /// normalized direction, and is reported explicitly to the owning prior.
39    pub fn new(
40        left: ArrayView2<'_, f64>,
41        right: ArrayView2<'_, f64>,
42        normalization: GramNormalization,
43    ) -> Option<Self> {
44        assert_eq!(left.ncols(), right.ncols());
45        let cross = left.dot(&right.t());
46        let left_gram = left.dot(&left.t());
47        let right_gram = right.dot(&right.t());
48        let (left_stat, right_stat, exponent) = match normalization {
49            GramNormalization::DecoderNorm => (
50                left.iter().map(|x| x * x).sum::<f64>(),
51                right.iter().map(|x| x * x).sum::<f64>(),
52                -1.0,
53            ),
54            GramNormalization::SelfGramNorm => (
55                left_gram.iter().map(|x| x * x).sum::<f64>(),
56                right_gram.iter().map(|x| x * x).sum::<f64>(),
57                -0.5,
58            ),
59        };
60        if !(left_stat > 0.0 && right_stat > 0.0) {
61            return None;
62        }
63        let energy = cross.iter().map(|x| x * x).sum::<f64>();
64        // Derivatives of phi(E,u,v)=E*u^p*v^p. Multi-index orders above
65        // one in E vanish. Falling factorials give all remaining derivatives.
66        let scalar_derivative = |indices: &[usize]| {
67            let mut orders = [0_usize; 3];
68            for &index in indices {
69                orders[index] += 1;
70            }
71            if orders[0] > 1 {
72                return 0.0;
73            }
74            let power_derivative = |value: f64, order: usize| {
75                let coefficient = (0..order).map(|i| exponent - i as f64).product::<f64>();
76                coefficient * value.powf(exponent - order as f64)
77            };
78            (if orders[0] == 0 { energy } else { 1.0 })
79                * power_derivative(left_stat, orders[1])
80                * power_derivative(right_stat, orders[2])
81        };
82        let scalar_first = std::array::from_fn(|i| scalar_derivative(&[i]));
83        let scalar_second =
84            std::array::from_fn(|i| std::array::from_fn(|j| scalar_derivative(&[i, j])));
85        let scalar_third = std::array::from_fn(|i| {
86            std::array::from_fn(|j| std::array::from_fn(|k| scalar_derivative(&[i, j, k])))
87        });
88        let energy_gradient = flatten_pair(cross.dot(&right) * 2.0, cross.t().dot(&left) * 2.0);
89        let left_gradient = match normalization {
90            GramNormalization::DecoderNorm => left.to_owned() * 2.0,
91            GramNormalization::SelfGramNorm => left_gram.dot(&left) * 4.0,
92        };
93        let right_gradient = match normalization {
94            GramNormalization::DecoderNorm => right.to_owned() * 2.0,
95            GramNormalization::SelfGramNorm => right_gram.dot(&right) * 4.0,
96        };
97        let left_columns =
98            Array1::from_shape_fn(left.ncols(), |col| left.column(col).dot(&left.column(col)));
99        let right_columns = Array1::from_shape_fn(right.ncols(), |col| {
100            right.column(col).dot(&right.column(col))
101        });
102        let energy_diagonal = flatten_pair(
103            Array2::from_shape_fn(left.dim(), |(_, col)| 2.0 * right_columns[col]),
104            Array2::from_shape_fn(right.dim(), |(_, col)| 2.0 * left_columns[col]),
105        );
106        let left_diagonal = Array2::from_shape_fn(left.dim(), |(row, col)| match normalization {
107            GramNormalization::DecoderNorm => 2.0,
108            GramNormalization::SelfGramNorm => {
109                4.0 * (left_columns[col]
110                    + left[[row, col]] * left[[row, col]]
111                    + left_gram[[row, row]])
112            }
113        });
114        let right_diagonal = Array2::from_shape_fn(right.dim(), |(row, col)| match normalization {
115            GramNormalization::DecoderNorm => 2.0,
116            GramNormalization::SelfGramNorm => {
117                4.0 * (right_columns[col]
118                    + right[[row, col]] * right[[row, col]]
119                    + right_gram[[row, row]])
120            }
121        });
122        Some(Self {
123            left: left.to_owned(),
124            right: right.to_owned(),
125            cross,
126            left_gram,
127            right_gram,
128            normalization,
129            gradients: [
130                energy_gradient,
131                flatten_pair(left_gradient, Array2::zeros(right.dim())),
132                flatten_pair(Array2::zeros(left.dim()), right_gradient),
133            ],
134            diagonals: [
135                energy_diagonal,
136                flatten_pair(left_diagonal, Array2::zeros(right.dim())),
137                flatten_pair(Array2::zeros(left.dim()), right_diagonal),
138            ],
139            scalar_first,
140            scalar_second,
141            scalar_third,
142        })
143    }
144
145    pub fn diagonal(&self) -> Array1<f64> {
146        let mut out = Array1::zeros(self.gradients[0].len());
147        for i in 0..3 {
148            out.scaled_add(self.scalar_first[i], &self.diagonals[i]);
149            for j in 0..3 {
150                out.scaled_add(
151                    self.scalar_second[i][j],
152                    &(&self.gradients[i] * &self.gradients[j]),
153                );
154            }
155        }
156        out
157    }
158
159    /// The diagonal of the frozen-normalizer cross-Gram Gauss–Newton matrix.
160    /// This is the positive curvature used by decoder incoherence assembly.
161    pub fn gauss_newton_diagonal(&self) -> Array1<f64> {
162        &self.diagonals[0] * self.scalar_first[0]
163    }
164
165    pub fn gauss_newton_action(&self, direction: ArrayView1<'_, f64>) -> Array1<f64> {
166        let (l, r) = self.split(direction);
167        let dc = l.dot(&self.right.t()) + self.left.dot(&r.t());
168        flatten_pair(dc.dot(&self.right), dc.t().dot(&self.left)) * (2.0 * self.scalar_first[0])
169    }
170
171    /// Gradient of left' B right for that installed majorizer. Its live
172    /// normalizer is differentiated too; the directional vectors stay fixed.
173    pub fn gauss_newton_bilinear_gradient(
174        &self,
175        left: ArrayView1<'_, f64>,
176        right: ArrayView1<'_, f64>,
177    ) -> Array1<f64> {
178        let (lx, ly) = self.split(left);
179        let (rx, ry) = self.split(right);
180        let cl = lx.dot(&self.right.t()) + self.left.dot(&ly.t());
181        let cr = rx.dot(&self.right.t()) + self.left.dot(&ry.t());
182        let inner = cl.iter().zip(cr.iter()).map(|(&a, &b)| a * b).sum::<f64>();
183        let mut out = flatten_pair(cr.dot(&ly) + cl.dot(&ry), cr.t().dot(&lx) + cl.t().dot(&rx))
184            * (2.0 * self.scalar_first[0]);
185        for j in 0..3 {
186            out.scaled_add(2.0 * inner * self.scalar_second[0][j], &self.gradients[j]);
187        }
188        out
189    }
190
191    fn split<'a>(
192        &self,
193        direction: ArrayView1<'a, f64>,
194    ) -> (ArrayView2<'a, f64>, ArrayView2<'a, f64>) {
195        assert_eq!(direction.len(), self.left.len() + self.right.len());
196        let left = direction
197            .slice_move(s![..self.left.len()])
198            .into_shape_with_order(self.left.dim())
199            .expect("validated left direction length matches the captured decoder layout");
200        let right = direction
201            .slice_move(s![self.left.len()..])
202            .into_shape_with_order(self.right.dim())
203            .expect("validated right direction length matches the captured decoder layout");
204        (left, right)
205    }
206
207    fn statistic_hessians(&self, direction: ArrayView1<'_, f64>) -> [Array1<f64>; 3] {
208        let (l, r) = self.split(direction);
209        let dc = l.dot(&self.right.t()) + self.left.dot(&r.t());
210        let energy = flatten_pair(
211            (dc.dot(&self.right) + self.cross.dot(&r)) * 2.0,
212            (dc.t().dot(&self.left) + self.cross.t().dot(&l)) * 2.0,
213        );
214        let normalizer_action = |x: &Array2<f64>, gram: &Array2<f64>, v: ArrayView2<'_, f64>| {
215            match self.normalization {
216                GramNormalization::DecoderNorm => v.to_owned() * 2.0,
217                GramNormalization::SelfGramNorm => {
218                    let one_leg = v.dot(&x.t());
219                    ((&one_leg + &one_leg.t()).dot(x) + gram.dot(&v)) * 4.0
220                }
221            }
222        };
223        [
224            energy,
225            flatten_pair(
226                normalizer_action(&self.left, &self.left_gram, l),
227                Array2::zeros(self.right.dim()),
228            ),
229            flatten_pair(
230                Array2::zeros(self.left.dim()),
231                normalizer_action(&self.right, &self.right_gram, r),
232            ),
233        ]
234    }
235
236    pub fn hessian_action(&self, direction: ArrayView1<'_, f64>) -> Array1<f64> {
237        let actions = self.statistic_hessians(direction);
238        let mut out = Array1::zeros(direction.len());
239        for i in 0..3 {
240            out.scaled_add(self.scalar_first[i], &actions[i]);
241            for j in 0..3 {
242                out.scaled_add(
243                    self.scalar_second[i][j] * self.gradients[j].dot(&direction),
244                    &self.gradients[i],
245                );
246            }
247        }
248        out
249    }
250
251    /// Gradient of left' H right, evaluated by the analytic third-order chain
252    /// rule of the three polynomial statistics. Both direction vectors are fixed.
253    pub fn third_bilinear(
254        &self,
255        left: ArrayView1<'_, f64>,
256        right: ArrayView1<'_, f64>,
257    ) -> Array1<f64> {
258        let (lx, ly) = self.split(left);
259        let (rx, ry) = self.split(right);
260        let cl = lx.dot(&self.right.t()) + self.left.dot(&ly.t());
261        let cr = rx.dot(&self.right.t()) + self.left.dot(&ry.t());
262        let clr = lx.dot(&ry.t()) + rx.dot(&ly.t());
263        let energy_third = flatten_pair(
264            (clr.dot(&self.right) + cl.dot(&ry) + cr.dot(&ly)) * 2.0,
265            (clr.t().dot(&self.left) + cl.t().dot(&rx) + cr.t().dot(&lx)) * 2.0,
266        );
267        let norm_third =
268            |x: &Array2<f64>, l: ArrayView2<'_, f64>, r: ArrayView2<'_, f64>| match self
269                .normalization
270            {
271                GramNormalization::DecoderNorm => Array2::zeros(x.dim()),
272                GramNormalization::SelfGramNorm => {
273                    let dl = l.dot(&x.t());
274                    let dr = r.dot(&x.t());
275                    let mixed = l.dot(&r.t());
276                    ((&mixed + &mixed.t()).dot(x)
277                        + (&dl + &dl.t()).dot(&r)
278                        + (&dr + &dr.t()).dot(&l))
279                        * 4.0
280                }
281            };
282        let thirds = [
283            energy_third,
284            flatten_pair(
285                norm_third(&self.left, lx, rx),
286                Array2::zeros(self.right.dim()),
287            ),
288            flatten_pair(
289                Array2::zeros(self.left.dim()),
290                norm_third(&self.right, ly, ry),
291            ),
292        ];
293        let h_left = self.statistic_hessians(left);
294        let h_right = self.statistic_hessians(right);
295        let dl: [f64; 3] = std::array::from_fn(|i| self.gradients[i].dot(&left));
296        let dr: [f64; 3] = std::array::from_fn(|i| self.gradients[i].dot(&right));
297        let mixed: [f64; 3] = std::array::from_fn(|i| left.dot(&h_right[i]));
298        let mut out = Array1::zeros(left.len());
299        for i in 0..3 {
300            out.scaled_add(self.scalar_first[i], &thirds[i]);
301            for j in 0..3 {
302                let coefficient = self.scalar_second[i][j];
303                out.scaled_add(coefficient * mixed[i], &self.gradients[j]);
304                out.scaled_add(coefficient * dr[j], &h_left[i]);
305                out.scaled_add(coefficient * dl[j], &h_right[i]);
306                for k in 0..3 {
307                    out.scaled_add(
308                        self.scalar_third[i][j][k] * dl[i] * dr[j],
309                        &self.gradients[k],
310                    );
311                }
312            }
313        }
314        out
315    }
316}
gam_terms/analytic_penalties/normalized_gram.rs

gam_terms/analytic_penalties/
normalized_gram.rs
