Skip to main content

gam_models/gamlss/gaussian/
joint_psi.rs

1// Real concern-organized submodule of the gamlss family stack.
2// Cross-module items are re-exported flat through the parent (`gamlss.rs`),
3// so `use super::*;` makes the sibling-concern symbols this module references
4// resolve through the parent namespace.
5use super::*;
6use gam_row_macros::row_atom;
7
8// Stable local coordinates for the Gaussian location-scale row NLL. At the
9// expansion point `delta_mu = delta_eta = 0`, the log-b link gives
10// `sigma(eta + delta_eta) / sigma(eta) = (1-kappa) + kappa*exp(delta_eta)`.
11// The perturbed standardized residual is therefore
12// `(standardized_residual - delta_mu*inv_sigma) / scale_ratio`. Every runtime
13// constant is already certified by `gaussian_diagonal_row_kernel`, so this one
14// expression retains the production extreme-value semantics while build-time
15// differentiation emits exact observed H, contracted t3, and contracted t4.
16row_atom! {
17    fn gaussian_normalized_row [generic, order2, third, fourth](
18        delta_mu,
19        delta_eta;
20        obs_weight,
21        standardized_residual,
22        inv_sigma,
23        kappa
24    ) {
25        obs_weight * ln((1.0 - kappa) + kappa * exp(delta_eta))
26            + 0.5
27                * obs_weight
28                * (standardized_residual - delta_mu * inv_sigma)
29                * (standardized_residual - delta_mu * inv_sigma)
30                / ((1.0 - kappa) + kappa * exp(delta_eta))
31                / ((1.0 - kappa) + kappa * exp(delta_eta))
32    }
33}
34
35pub(crate) struct LocationScaleJointPsiDirection {
36    pub(crate) block_idx: usize,
37    pub(crate) local_idx: usize,
38    pub(crate) x_primary_psi: PsiDesignMap,
39    pub(crate) x_ls_psi: PsiDesignMap,
40    pub(crate) z_primary_psi: Array1<f64>,
41    pub(crate) z_ls_psi: Array1<f64>,
42}
43
44pub(crate) struct LocationScaleJointPsiSecondDrifts {
45    pub(crate) x_primary_ab_action: Option<CustomFamilyPsiSecondDesignAction>,
46    pub(crate) x_ls_ab_action: Option<CustomFamilyPsiSecondDesignAction>,
47    pub(crate) x_primary_ab: Option<Array2<f64>>,
48    pub(crate) x_ls_ab: Option<Array2<f64>>,
49    pub(crate) z_primary_ab: Array1<f64>,
50    pub(crate) z_ls_ab: Array1<f64>,
51}
52
53/// Shared interface that the Gaussian and Binomial location-scale families (and
54/// their wiggle variants) expose to the unified joint ψ workspace.
55///
56/// The four families are structurally identical at the workspace level: each
57/// owns two dense block designs (location + log-scale), produces a per-ψ
58/// direction, and assembles second-order ψ terms and a ψ-Hessian directional
59/// derivative from those parts. They differ only in (1) the concrete
60/// [`Direction`](Self::Direction) struct produced (Gaussian vs Binomial field
61/// names), (2) the family-name fragment in the dense-designs error message, and
62/// (3) whether an optional Horvitz–Thompson outer-row subsample is threaded
63/// into the per-row weight arrays (Gaussian does; Binomial ignores it and runs
64/// the full-data exact path). This single trait gives the generic
65/// [`LocationScaleJointPsiWorkspace`] one dispatch surface; each family's impl
66/// is a thin delegation to inherent methods it already owns.
67pub(crate) trait LocationScaleJointPsiFamily: Clone + Send + Sync + 'static {
68    /// Per-ψ joint direction produced by this family.
69    type Direction: Send + Sync + 'static;
70
71    /// Family-name fragment used in the workspace's dense-designs error
72    /// message so the originating family stays visible after unification.
73    const LABEL: &'static str;
74
75    fn ws_policy(&self) -> &gam_runtime::resource::ResourcePolicy;
76
77    fn ws_exact_joint_dense_block_designs<'a>(
78        &'a self,
79        specs: Option<&'a [ParameterBlockSpec]>,
80    ) -> Result<Option<(Cow<'a, Array2<f64>>, Cow<'a, Array2<f64>>)>, String>;
81
82    fn ws_psi_direction(
83        &self,
84        block_states: &[ParameterBlockState],
85        derivative_blocks: &[Vec<crate::custom_family::CustomFamilyBlockPsiDerivative>],
86        psi_index: usize,
87        design_loc: &Array2<f64>,
88        design_scale: &Array2<f64>,
89        policy: &gam_runtime::resource::ResourcePolicy,
90    ) -> Result<Option<Self::Direction>, String>;
91
92    fn ws_psi_second_order_terms_from_parts(
93        &self,
94        block_states: &[ParameterBlockState],
95        derivative_blocks: &[Vec<crate::custom_family::CustomFamilyBlockPsiDerivative>],
96        psi_a: &Self::Direction,
97        psi_b: &Self::Direction,
98        design_loc: &Array2<f64>,
99        design_scale: &Array2<f64>,
100        subsample: Option<&[crate::outer_subsample::WeightedOuterRow]>,
101    ) -> Result<ExactNewtonJointPsiSecondOrderTerms, String>;
102
103    fn ws_psi_hessian_directional_from_parts(
104        &self,
105        block_states: &[ParameterBlockState],
106        psi_dir: &Self::Direction,
107        d_beta_flat: &Array1<f64>,
108        design_loc: &Array2<f64>,
109        design_scale: &Array2<f64>,
110        subsample: Option<&[crate::outer_subsample::WeightedOuterRow]>,
111    ) -> Result<Array2<f64>, String>;
112}
113
114impl LocationScaleJointPsiFamily for GaussianLocationScaleFamily {
115    type Direction = LocationScaleJointPsiDirection;
116    const LABEL: &'static str = "GaussianLocationScaleFamily";
117
118    fn ws_policy(&self) -> &gam_runtime::resource::ResourcePolicy {
119        &self.policy
120    }
121
122    fn ws_exact_joint_dense_block_designs<'a>(
123        &'a self,
124        specs: Option<&'a [ParameterBlockSpec]>,
125    ) -> Result<Option<(Cow<'a, Array2<f64>>, Cow<'a, Array2<f64>>)>, String> {
126        self.exact_joint_dense_block_designs(specs)
127    }
128
129    fn ws_psi_direction(
130        &self,
131        block_states: &[ParameterBlockState],
132        derivative_blocks: &[Vec<crate::custom_family::CustomFamilyBlockPsiDerivative>],
133        psi_index: usize,
134        design_loc: &Array2<f64>,
135        design_scale: &Array2<f64>,
136        policy: &gam_runtime::resource::ResourcePolicy,
137    ) -> Result<Option<LocationScaleJointPsiDirection>, String> {
138        self.exact_newton_joint_psi_direction(
139            block_states,
140            derivative_blocks,
141            psi_index,
142            design_loc,
143            design_scale,
144            policy,
145        )
146    }
147
148    fn ws_psi_second_order_terms_from_parts(
149        &self,
150        block_states: &[ParameterBlockState],
151        derivative_blocks: &[Vec<crate::custom_family::CustomFamilyBlockPsiDerivative>],
152        psi_a: &LocationScaleJointPsiDirection,
153        psi_b: &LocationScaleJointPsiDirection,
154        design_loc: &Array2<f64>,
155        design_scale: &Array2<f64>,
156        subsample: Option<&[crate::outer_subsample::WeightedOuterRow]>,
157    ) -> Result<ExactNewtonJointPsiSecondOrderTerms, String> {
158        self.exact_newton_joint_psisecond_order_terms_from_parts(
159            block_states,
160            derivative_blocks,
161            psi_a,
162            psi_b,
163            design_loc,
164            design_scale,
165            subsample,
166        )
167    }
168
169    fn ws_psi_hessian_directional_from_parts(
170        &self,
171        block_states: &[ParameterBlockState],
172        psi_dir: &LocationScaleJointPsiDirection,
173        d_beta_flat: &Array1<f64>,
174        design_loc: &Array2<f64>,
175        design_scale: &Array2<f64>,
176        subsample: Option<&[crate::outer_subsample::WeightedOuterRow]>,
177    ) -> Result<Array2<f64>, String> {
178        self.exact_newton_joint_psihessian_directional_derivative_from_parts(
179            block_states,
180            psi_dir,
181            d_beta_flat,
182            design_loc,
183            design_scale,
184            subsample,
185        )
186    }
187}
188
189impl LocationScaleJointPsiFamily for GaussianLocationScaleWiggleFamily {
190    type Direction = LocationScaleJointPsiDirection;
191    const LABEL: &'static str = "GaussianLocationScaleWiggleFamily";
192
193    fn ws_policy(&self) -> &gam_runtime::resource::ResourcePolicy {
194        &self.policy
195    }
196
197    fn ws_exact_joint_dense_block_designs<'a>(
198        &'a self,
199        specs: Option<&'a [ParameterBlockSpec]>,
200    ) -> Result<Option<(Cow<'a, Array2<f64>>, Cow<'a, Array2<f64>>)>, String> {
201        self.exact_joint_dense_block_designs(specs)
202    }
203
204    fn ws_psi_direction(
205        &self,
206        block_states: &[ParameterBlockState],
207        derivative_blocks: &[Vec<crate::custom_family::CustomFamilyBlockPsiDerivative>],
208        psi_index: usize,
209        design_loc: &Array2<f64>,
210        design_scale: &Array2<f64>,
211        policy: &gam_runtime::resource::ResourcePolicy,
212    ) -> Result<Option<LocationScaleJointPsiDirection>, String> {
213        self.exact_newton_joint_psi_direction(
214            block_states,
215            derivative_blocks,
216            psi_index,
217            design_loc,
218            design_scale,
219            policy,
220        )
221    }
222
223    fn ws_psi_second_order_terms_from_parts(
224        &self,
225        block_states: &[ParameterBlockState],
226        derivative_blocks: &[Vec<crate::custom_family::CustomFamilyBlockPsiDerivative>],
227        psi_a: &LocationScaleJointPsiDirection,
228        psi_b: &LocationScaleJointPsiDirection,
229        design_loc: &Array2<f64>,
230        design_scale: &Array2<f64>,
231        _: Option<&[crate::outer_subsample::WeightedOuterRow]>,
232    ) -> Result<ExactNewtonJointPsiSecondOrderTerms, String> {
233        // Wiggle ψ path: full-data exact (= trivially unbiased). The
234        // wiggle-specific second-order from-parts function inlines 30+
235        // per-row coefficient arrays (`coeff_mm{,_a,_b,_ab}`,
236        // `coeff_ml{,_a,_b,_ab}`, `coeff_ll{,_a,_b,_ab}`, `a{,_a,_b,_ab}`,
237        // `c{,_a,_b,_ab}`, `l{,_a,_b,_ab}`, `dw_{a,b,ab}`, `s_mu*`, `s_ls*`,
238        // `s_w*`, ...) instead of packing them into a struct like the
239        // non-wiggle GLS path's `GaussianJointPsi{First,Second}Weights`.
240        // Each is row-linear in `rows.{w,m,n,kappa,...}` and the direction
241        // vectors so HT masking is theoretically clean, but threading a mask
242        // across that many call sites is brittle (any missed array silently
243        // biases the estimator). The outer score remains unbiased without
244        // touching the wiggle ψ path: HT-unbiased LL
245        // (`log_likelihood_only_with_options`) + HT-unbiased ρ-Hessian
246        // (`exact_newton_joint_hessian_workspace_with_options`) +
247        // exact-unbiased ψ (this path) = unbiased. Broadening to the wiggle
248        // ψ path is a follow-up that should refactor the inline arrays into
249        // `WiggleJointPsi{First,Second}Weights` structs mirroring
250        // `GaussianJointPsi{First,Second}Weights` so a single
251        // `apply_ht_mask_wiggle*` helper can mask everything in one place.
252        self.exact_newton_joint_psisecond_order_terms_from_parts(
253            block_states,
254            derivative_blocks,
255            psi_a,
256            psi_b,
257            design_loc,
258            design_scale,
259        )
260    }
261
262    fn ws_psi_hessian_directional_from_parts(
263        &self,
264        block_states: &[ParameterBlockState],
265        psi_dir: &LocationScaleJointPsiDirection,
266        d_beta_flat: &Array1<f64>,
267        design_loc: &Array2<f64>,
268        design_scale: &Array2<f64>,
269        _: Option<&[crate::outer_subsample::WeightedOuterRow]>,
270    ) -> Result<Array2<f64>, String> {
271        // Same rationale as `ws_psi_second_order_terms_from_parts` above:
272        // the wiggle ψ-Hessian directional-derivative function also inlines
273        // dozens of per-row arrays. Full-data is exact (= trivially
274        // unbiased), so the total outer score remains unbiased.
275        self.exact_newton_joint_psihessian_directional_derivative_from_parts(
276            block_states,
277            psi_dir,
278            d_beta_flat,
279            design_loc,
280            design_scale,
281        )
282    }
283}
284
285/// Generic joint exact-Newton ψ workspace shared by every location-scale
286/// family (Gaussian / Binomial, with or without a wiggle block) via the
287/// [`LocationScaleJointPsiFamily`] trait.
288///
289/// The workspace owns the two dense block designs as `Arc<Array2<f64>>` (the
290/// per-family `ws_exact_joint_dense_block_designs` hands back a `Cow`, which is
291/// materialized once here), the per-ψ direction cache, and an optional
292/// Horvitz–Thompson outer-row subsample. When the subsample is `Some`, every
293/// per-row weight array produced inside the second-order ψ Hessian and the
294/// ψ-Hessian directional-derivative computations is masked: each sampled row's
295/// contribution is scaled by `WeightedOuterRow.weight = 1/π_i` and non-sampled
296/// rows are zeroed. Because every downstream assembly is row-linear in those
297/// arrays, the resulting ψ score and ψ Hessian remain unbiased estimators of
298/// the full-data quantities. Families that do not thread the subsample (the
299/// Binomial families) construct with `new` and the field stays `None`.
300pub(crate) struct LocationScaleJointPsiWorkspace<F: LocationScaleJointPsiFamily> {
301    pub(crate) family: F,
302    pub(crate) block_states: Vec<ParameterBlockState>,
303    pub(crate) derivative_blocks: Vec<Vec<CustomFamilyBlockPsiDerivative>>,
304    pub(crate) design_loc: Arc<Array2<f64>>,
305    pub(crate) design_scale: Arc<Array2<f64>>,
306    pub(crate) psi_directions: ExactNewtonJointPsiDirectCache<F::Direction>,
307    pub(crate) outer_score_subsample: Option<Arc<crate::outer_subsample::OuterScoreSubsample>>,
308}
309
310impl<F: LocationScaleJointPsiFamily> LocationScaleJointPsiWorkspace<F> {
311    pub(crate) fn new(
312        family: F,
313        block_states: Vec<ParameterBlockState>,
314        specs: &[ParameterBlockSpec],
315        derivative_blocks: Vec<Vec<CustomFamilyBlockPsiDerivative>>,
316    ) -> Result<Self, String> {
317        Self::new_with_subsample(family, block_states, specs, derivative_blocks, None)
318    }
319
320    pub(crate) fn new_with_subsample(
321        family: F,
322        block_states: Vec<ParameterBlockState>,
323        specs: &[ParameterBlockSpec],
324        derivative_blocks: Vec<Vec<CustomFamilyBlockPsiDerivative>>,
325        outer_score_subsample: Option<Arc<crate::outer_subsample::OuterScoreSubsample>>,
326    ) -> Result<Self, String> {
327        let Some((design_loc, design_scale)) =
328            family.ws_exact_joint_dense_block_designs(Some(specs))?
329        else {
330            return Err(GamlssError::UnsupportedConfiguration {
331                reason: format!(
332                    "{} exact joint psi workspace requires dense block designs",
333                    F::LABEL,
334                ),
335            }
336            .into());
337        };
338        let design_loc = shared_dense_arc(design_loc.as_ref());
339        let design_scale = shared_dense_arc(design_scale.as_ref());
340        let psi_dim = derivative_blocks.iter().map(Vec::len).sum();
341        Ok(Self {
342            family,
343            block_states,
344            derivative_blocks,
345            design_loc,
346            design_scale,
347            psi_directions: ExactNewtonJointPsiDirectCache::new(psi_dim),
348            outer_score_subsample,
349        })
350    }
351
352    pub(crate) fn psi_direction(
353        &self,
354        psi_index: usize,
355    ) -> Result<Option<Arc<F::Direction>>, String> {
356        self.psi_directions.get_or_try_init(psi_index, || {
357            self.family.ws_psi_direction(
358                &self.block_states,
359                &self.derivative_blocks,
360                psi_index,
361                self.design_loc.as_ref(),
362                self.design_scale.as_ref(),
363                self.family.ws_policy(),
364            )
365        })
366    }
367
368    pub(crate) fn subsample_rows(&self) -> Option<&[crate::outer_subsample::WeightedOuterRow]> {
369        self.outer_score_subsample
370            .as_ref()
371            .map(|s| s.rows.as_ref().as_slice())
372    }
373}
374
375impl<F> ExactNewtonJointPsiWorkspace for LocationScaleJointPsiWorkspace<F>
376where
377    F: LocationScaleJointPsiFamily,
378{
379    fn second_order_terms(
380        &self,
381        psi_i: usize,
382        psi_j: usize,
383    ) -> Result<Option<ExactNewtonJointPsiSecondOrderTerms>, String> {
384        let Some(dir_i) = self.psi_direction(psi_i)? else {
385            return Ok(None);
386        };
387        let Some(dir_j) = self.psi_direction(psi_j)? else {
388            return Ok(None);
389        };
390        Ok(Some(self.family.ws_psi_second_order_terms_from_parts(
391            &self.block_states,
392            &self.derivative_blocks,
393            dir_i.as_ref(),
394            dir_j.as_ref(),
395            self.design_loc.as_ref(),
396            self.design_scale.as_ref(),
397            self.subsample_rows(),
398        )?))
399    }
400
401    fn hessian_directional_derivative(
402        &self,
403        psi_index: usize,
404        d_beta_flat: &Array1<f64>,
405    ) -> Result<Option<gam_problem::DriftDerivResult>, String> {
406        let Some(dir) = self.psi_direction(psi_index)? else {
407            return Ok(None);
408        };
409        Ok(Some(gam_problem::DriftDerivResult::Dense(
410            self.family.ws_psi_hessian_directional_from_parts(
411                &self.block_states,
412                dir.as_ref(),
413                d_beta_flat,
414                self.design_loc.as_ref(),
415                self.design_scale.as_ref(),
416                self.subsample_rows(),
417            )?,
418        )))
419    }
420}
421
422pub(crate) type GaussianLocationScaleExactNewtonJointPsiWorkspace =
423    LocationScaleJointPsiWorkspace<GaussianLocationScaleFamily>;
424
425pub(crate) type GaussianLocationScaleWiggleExactNewtonJointPsiWorkspace =
426    LocationScaleJointPsiWorkspace<GaussianLocationScaleWiggleFamily>;
427
428#[derive(Clone)]
429pub struct GaussianJointRowScalars {
430    pub(crate) obs_weight: Array1<f64>,
431    /// Stable `(y - mu) / sigma` at the expansion point.
432    pub(crate) standardized_residual: Array1<f64>,
433    /// Stable `1 / sigma` at the expansion point.
434    pub(crate) inv_sigma: Array1<f64>,
435    /// κ = (dσ/dη_ls)/σ for the active sigma link.
436    /// The cross Hessian block H_{μ,ls} carries an overall κ factor and the
437    /// scale-scale block H_{ls,ls} carries κ².
438    pub(crate) kappa: Array1<f64>,
439}
440
441/// Production [`gam_math::jet_tower::RowProgram`] for the normalized Gaussian
442/// location-scale row NLL.
443///
444/// The program borrows the exact certified row constants consumed by the live
445/// observed-Hessian and directional-weight paths. Its generic evaluator and
446/// those specialized order-2/third/fourth paths are emitted from the same
447/// [`gaussian_normalized_row`] declaration, so the parity oracle cannot retain
448/// an independent copy of the likelihood expression.
449pub struct GaussianJointRowProgram<'a> {
450    rows: &'a GaussianJointRowScalars,
451}
452
453impl<'a> GaussianJointRowProgram<'a> {
454    /// Bind the generic row program to one certified production scalar batch.
455    pub fn new(rows: &'a GaussianJointRowScalars) -> Self {
456        Self { rows }
457    }
458
459    fn require_row(&self, row: usize) -> Result<(), String> {
460        if row >= self.rows.obs_weight.len() {
461            return Err(format!(
462                "GaussianJointRowProgram row {row} out of range for {} rows",
463                self.rows.obs_weight.len()
464            ));
465        }
466        Ok(())
467    }
468
469    /// Symbolically lowered value/gradient/Hessian for one certified row.
470    ///
471    /// The concrete sparsity bits are part of the generated function's type:
472    /// both score channels and all three packed Hessian channels are live.
473    #[inline(always)]
474    pub(crate) fn row_order2(
475        &self,
476        row: usize,
477    ) -> gam_math::jet_scalar::StaticOrder2Atom<2, 3, 3, 7> {
478        gaussian_normalized_row_order2(
479            0.0,
480            0.0,
481            self.rows.obs_weight[row],
482            self.rows.standardized_residual[row],
483            self.rows.inv_sigma[row],
484            self.rows.kappa[row],
485        )
486    }
487
488    /// Symbolically lowered Hessian derivative in one predictor direction.
489    #[inline(always)]
490    pub(crate) fn row_third_contracted(&self, row: usize, direction: &[f64; 2]) -> [[f64; 2]; 2] {
491        gaussian_normalized_row_third_contracted(
492            0.0,
493            0.0,
494            self.rows.obs_weight[row],
495            self.rows.standardized_residual[row],
496            self.rows.inv_sigma[row],
497            self.rows.kappa[row],
498            direction,
499        )
500    }
501
502    /// Symbolically lowered mixed derivative of the row Hessian.
503    #[inline(always)]
504    pub(crate) fn row_fourth_contracted(
505        &self,
506        row: usize,
507        direction_u: &[f64; 2],
508        direction_v: &[f64; 2],
509    ) -> [[f64; 2]; 2] {
510        gaussian_normalized_row_fourth_contracted(
511            0.0,
512            0.0,
513            self.rows.obs_weight[row],
514            self.rows.standardized_residual[row],
515            self.rows.inv_sigma[row],
516            self.rows.kappa[row],
517            direction_u,
518            direction_v,
519        )
520    }
521}
522
523#[inline(always)]
524fn matrix_vector_2(matrix: &[[f64; 2]; 2], vector: &[f64; 2]) -> [f64; 2] {
525    [
526        matrix[0][0] * vector[0] + matrix[0][1] * vector[1],
527        matrix[1][0] * vector[0] + matrix[1][1] * vector[1],
528    ]
529}
530
531#[inline(always)]
532fn dot_2(left: &[f64; 2], right: &[f64; 2]) -> f64 {
533    left[0] * right[0] + left[1] * right[1]
534}
535
536#[inline(always)]
537fn add_vector_2(left: [f64; 2], right: [f64; 2]) -> [f64; 2] {
538    [left[0] + right[0], left[1] + right[1]]
539}
540
541#[inline(always)]
542fn add_matrix_2(left: [[f64; 2]; 2], right: [[f64; 2]; 2]) -> [[f64; 2]; 2] {
543    [
544        [left[0][0] + right[0][0], left[0][1] + right[0][1]],
545        [left[1][0] + right[1][0], left[1][1] + right[1][1]],
546    ]
547}
548
549/// One order of generated Gaussian row geometry, stored in neutral predictor
550/// coordinates. At order zero these are `(g, H)`; in a first/second tower they
551/// are the corresponding directional derivatives of `(g, H)`.
552pub(crate) struct GaussianRowChannels {
553    pub(crate) gradient_mu: Array1<f64>,
554    pub(crate) gradient_ls: Array1<f64>,
555    pub(crate) hessian_mm: Array1<f64>,
556    pub(crate) hessian_ml: Array1<f64>,
557    pub(crate) hessian_ll: Array1<f64>,
558}
559
560impl GaussianRowChannels {
561    fn zeros(n: usize) -> Self {
562        Self {
563            gradient_mu: Array1::zeros(n),
564            gradient_ls: Array1::zeros(n),
565            hessian_mm: Array1::zeros(n),
566            hessian_ml: Array1::zeros(n),
567            hessian_ll: Array1::zeros(n),
568        }
569    }
570}
571
572pub(crate) struct GaussianRowFirstTower {
573    pub(crate) base: GaussianRowChannels,
574    pub(crate) first: GaussianRowChannels,
575}
576
577pub(crate) struct GaussianRowSecondTower {
578    pub(crate) base: GaussianRowChannels,
579    pub(crate) first_a: GaussianRowChannels,
580    pub(crate) first_b: GaussianRowChannels,
581    pub(crate) second: GaussianRowChannels,
582}
583
584fn write_base_channels(
585    channels: &mut GaussianRowChannels,
586    row: usize,
587    atom: &gam_math::jet_scalar::StaticOrder2Atom<2, 3, 3, 7>,
588) -> [[f64; 2]; 2] {
589    let gradient = atom.gradient();
590    let hessian = [
591        [atom.hessian_at(0, 0), atom.hessian_at(0, 1)],
592        [atom.hessian_at(1, 0), atom.hessian_at(1, 1)],
593    ];
594    channels.gradient_mu[row] = gradient[0];
595    channels.gradient_ls[row] = gradient[1];
596    channels.hessian_mm[row] = hessian[0][0];
597    channels.hessian_ml[row] = hessian[0][1];
598    channels.hessian_ll[row] = hessian[1][1];
599    hessian
600}
601
602fn write_directional_channels(
603    channels: &mut GaussianRowChannels,
604    row: usize,
605    gradient: [f64; 2],
606    hessian: [[f64; 2]; 2],
607) {
608    channels.gradient_mu[row] = gradient[0];
609    channels.gradient_ls[row] = gradient[1];
610    channels.hessian_mm[row] = hessian[0][0];
611    channels.hessian_ml[row] = hessian[0][1];
612    channels.hessian_ll[row] = hessian[1][1];
613}
614
615/// Generated Gaussian row gradient and Hessian with no directional scratch.
616pub(crate) fn gaussian_row_channels(rows: &GaussianJointRowScalars) -> GaussianRowChannels {
617    let n = rows.obs_weight.len();
618    let program = GaussianJointRowProgram::new(rows);
619    let mut base = GaussianRowChannels::zeros(n);
620    for row in 0..n {
621        let atom = program.row_order2(row);
622        write_base_channels(&mut base, row, &atom);
623    }
624    base
625}
626
627/// Generated `(g, H)` plus its first derivative along a rowwise predictor
628/// direction. The row atom is evaluated once per row.
629pub(crate) fn gaussian_row_first_tower(
630    rows: &GaussianJointRowScalars,
631    direction_mu: &Array1<f64>,
632    direction_ls: &Array1<f64>,
633) -> GaussianRowFirstTower {
634    let n = rows.obs_weight.len();
635    let program = GaussianJointRowProgram::new(rows);
636    let mut base = GaussianRowChannels::zeros(n);
637    let mut first = GaussianRowChannels::zeros(n);
638    for row in 0..n {
639        let direction = [direction_mu[row], direction_ls[row]];
640        let atom = program.row_order2(row);
641        let hessian = write_base_channels(&mut base, row, &atom);
642        write_directional_channels(
643            &mut first,
644            row,
645            matrix_vector_2(&hessian, &direction),
646            program.row_third_contracted(row, &direction),
647        );
648    }
649    GaussianRowFirstTower { base, first }
650}
651
652/// Generated `(g, H)` tower through the mixed second derivative along two
653/// rowwise directions, including a possibly nonzero mixed predictor leg.
654pub(crate) fn gaussian_row_second_tower(
655    rows: &GaussianJointRowScalars,
656    direction_a_mu: &Array1<f64>,
657    direction_a_ls: &Array1<f64>,
658    direction_b_mu: &Array1<f64>,
659    direction_b_ls: &Array1<f64>,
660    direction_ab_mu: &Array1<f64>,
661    direction_ab_ls: &Array1<f64>,
662) -> GaussianRowSecondTower {
663    let n = rows.obs_weight.len();
664    let program = GaussianJointRowProgram::new(rows);
665    let mut base = GaussianRowChannels::zeros(n);
666    let mut first_a = GaussianRowChannels::zeros(n);
667    let mut first_b = GaussianRowChannels::zeros(n);
668    let mut second = GaussianRowChannels::zeros(n);
669    for row in 0..n {
670        let direction_a = [direction_a_mu[row], direction_a_ls[row]];
671        let direction_b = [direction_b_mu[row], direction_b_ls[row]];
672        let direction_ab = [direction_ab_mu[row], direction_ab_ls[row]];
673        let atom = program.row_order2(row);
674        let hessian = write_base_channels(&mut base, row, &atom);
675        let hessian_a = program.row_third_contracted(row, &direction_a);
676        let hessian_b = program.row_third_contracted(row, &direction_b);
677        write_directional_channels(
678            &mut first_a,
679            row,
680            matrix_vector_2(&hessian, &direction_a),
681            hessian_a,
682        );
683        write_directional_channels(
684            &mut first_b,
685            row,
686            matrix_vector_2(&hessian, &direction_b),
687            hessian_b,
688        );
689        write_directional_channels(
690            &mut second,
691            row,
692            add_vector_2(
693                matrix_vector_2(&hessian_a, &direction_b),
694                matrix_vector_2(&hessian, &direction_ab),
695            ),
696            add_matrix_2(
697                program.row_fourth_contracted(row, &direction_a, &direction_b),
698                program.row_third_contracted(row, &direction_ab),
699            ),
700        );
701    }
702    GaussianRowSecondTower {
703        base,
704        first_a,
705        first_b,
706        second,
707    }
708}
709
710impl gam_math::jet_tower::RowProgram<2> for GaussianJointRowProgram<'_> {
711    fn n_rows(&self) -> usize {
712        self.rows.obs_weight.len()
713    }
714
715    fn primaries(&self, row: usize) -> Result<[f64; 2], String> {
716        self.require_row(row)?;
717        Ok([0.0, 0.0])
718    }
719
720    fn eval<S: gam_math::jet_scalar::JetScalar<2>>(
721        &self,
722        row: usize,
723        p: &[S; 2],
724    ) -> Result<S, String> {
725        self.require_row(row)?;
726        Ok(gaussian_normalized_row(
727            &p[0],
728            &p[1],
729            self.rows.obs_weight[row],
730            self.rows.standardized_residual[row],
731            self.rows.inv_sigma[row],
732            self.rows.kappa[row],
733        ))
734    }
735}
736
737pub(crate) struct GaussianJointPsiFirstWeights {
738    pub(crate) objective_psirow: Array1<f64>,
739    pub(crate) scoremu: Array1<f64>,
740    pub(crate) score_ls: Array1<f64>,
741    pub(crate) dscoremu: Array1<f64>,
742    pub(crate) dscore_ls: Array1<f64>,
743    pub(crate) hmumu: Array1<f64>,
744    pub(crate) hmu_ls: Array1<f64>,
745    pub(crate) h_ls_ls: Array1<f64>,
746    pub(crate) dhmumu: Array1<f64>,
747    pub(crate) dhmu_ls: Array1<f64>,
748    pub(crate) dh_ls_ls: Array1<f64>,
749}
750
751pub(crate) struct GaussianJointPsiSecondWeights {
752    pub(crate) objective_psi_psirow: Array1<f64>,
753    pub(crate) d2scoremu: Array1<f64>,
754    pub(crate) d2score_ls: Array1<f64>,
755    pub(crate) d2hmumu: Array1<f64>,
756    pub(crate) d2hmu_ls: Array1<f64>,
757    pub(crate) d2h_ls_ls: Array1<f64>,
758}
759
760pub(crate) struct GaussianJointPsiMixedDriftWeights {
761    pub(crate) dhmumu_u: Array1<f64>,
762    pub(crate) dhmu_ls_u: Array1<f64>,
763    pub(crate) dh_ls_ls_u: Array1<f64>,
764    pub(crate) d2hmumu: Array1<f64>,
765    pub(crate) d2hmu_ls: Array1<f64>,
766    pub(crate) d2h_ls_ls: Array1<f64>,
767}
768
769/// Apply a Horvitz–Thompson outer-row subsample mask to every per-row array
770/// of a `GaussianJointPsiFirstWeights` in place: each sampled row's
771/// contribution is multiplied by `WeightedOuterRow.weight = 1/π_i` and all
772/// non-sampled rows are zeroed. Every downstream assembly
773/// (`gaussian_joint_psi*_fromweights`,
774/// `build_two_block_custom_family_joint_psi_operator_from_actions`) consumes
775/// these arrays row-linearly via `Xᵀ diag(W) Y` and `weighted_crossprod_psi_maps`,
776/// so the resulting first-order ψ score and Hessian remain unbiased estimators
777/// of the full-data quantities.
778pub(crate) fn apply_ht_mask_first(
779    weights: &mut GaussianJointPsiFirstWeights,
780    rows: &[crate::outer_subsample::WeightedOuterRow],
781) {
782    let n = weights.objective_psirow.len();
783    let mut obj = Array1::<f64>::zeros(n);
784    let mut smu = Array1::<f64>::zeros(n);
785    let mut sls = Array1::<f64>::zeros(n);
786    let mut dsmu = Array1::<f64>::zeros(n);
787    let mut dsls = Array1::<f64>::zeros(n);
788    let mut hmm = Array1::<f64>::zeros(n);
789    let mut hml = Array1::<f64>::zeros(n);
790    let mut hll = Array1::<f64>::zeros(n);
791    let mut dhmm = Array1::<f64>::zeros(n);
792    let mut dhml = Array1::<f64>::zeros(n);
793    let mut dhll = Array1::<f64>::zeros(n);
794    for r in rows {
795        let i = r.index;
796        let w = r.weight;
797        obj[i] = weights.objective_psirow[i] * w;
798        smu[i] = weights.scoremu[i] * w;
799        sls[i] = weights.score_ls[i] * w;
800        dsmu[i] = weights.dscoremu[i] * w;
801        dsls[i] = weights.dscore_ls[i] * w;
802        hmm[i] = weights.hmumu[i] * w;
803        hml[i] = weights.hmu_ls[i] * w;
804        hll[i] = weights.h_ls_ls[i] * w;
805        dhmm[i] = weights.dhmumu[i] * w;
806        dhml[i] = weights.dhmu_ls[i] * w;
807        dhll[i] = weights.dh_ls_ls[i] * w;
808    }
809    weights.objective_psirow = obj;
810    weights.scoremu = smu;
811    weights.score_ls = sls;
812    weights.dscoremu = dsmu;
813    weights.dscore_ls = dsls;
814    weights.hmumu = hmm;
815    weights.hmu_ls = hml;
816    weights.h_ls_ls = hll;
817    weights.dhmumu = dhmm;
818    weights.dhmu_ls = dhml;
819    weights.dh_ls_ls = dhll;
820}
821
822/// HT mask for `GaussianJointPsiSecondWeights`. Same semantics as
823/// `apply_ht_mask_first`: each per-row contribution is scaled by 1/π_i and
824/// non-sampled rows are zeroed. Consumed row-linearly by
825/// `gaussian_joint_psisecondhessian_fromweights` and the `score_psi_psi`
826/// `fast_atv(_, d2score_*)` reductions.
827pub(crate) fn apply_ht_mask_second(
828    weights: &mut GaussianJointPsiSecondWeights,
829    rows: &[crate::outer_subsample::WeightedOuterRow],
830) {
831    let n = weights.objective_psi_psirow.len();
832    let mut obj = Array1::<f64>::zeros(n);
833    let mut d2smu = Array1::<f64>::zeros(n);
834    let mut d2sls = Array1::<f64>::zeros(n);
835    let mut d2hmm = Array1::<f64>::zeros(n);
836    let mut d2hml = Array1::<f64>::zeros(n);
837    let mut d2hll = Array1::<f64>::zeros(n);
838    for r in rows {
839        let i = r.index;
840        let w = r.weight;
841        obj[i] = weights.objective_psi_psirow[i] * w;
842        d2smu[i] = weights.d2scoremu[i] * w;
843        d2sls[i] = weights.d2score_ls[i] * w;
844        d2hmm[i] = weights.d2hmumu[i] * w;
845        d2hml[i] = weights.d2hmu_ls[i] * w;
846        d2hll[i] = weights.d2h_ls_ls[i] * w;
847    }
848    weights.objective_psi_psirow = obj;
849    weights.d2scoremu = d2smu;
850    weights.d2score_ls = d2sls;
851    weights.d2hmumu = d2hmm;
852    weights.d2hmu_ls = d2hml;
853    weights.d2h_ls_ls = d2hll;
854}
855
856/// HT mask for `GaussianJointPsiMixedDriftWeights`. Same semantics as the
857/// other `apply_ht_mask_*` helpers; consumed row-linearly by
858/// `gaussian_joint_psi_mixedhessian_drift_fromweights`.
859pub(crate) fn apply_ht_mask_mixed(
860    weights: &mut GaussianJointPsiMixedDriftWeights,
861    rows: &[crate::outer_subsample::WeightedOuterRow],
862) {
863    let n = weights.dhmumu_u.len();
864    let mut dhmm_u = Array1::<f64>::zeros(n);
865    let mut dhml_u = Array1::<f64>::zeros(n);
866    let mut dhll_u = Array1::<f64>::zeros(n);
867    let mut d2hmm = Array1::<f64>::zeros(n);
868    let mut d2hml = Array1::<f64>::zeros(n);
869    let mut d2hll = Array1::<f64>::zeros(n);
870    for r in rows {
871        let i = r.index;
872        let w = r.weight;
873        dhmm_u[i] = weights.dhmumu_u[i] * w;
874        dhml_u[i] = weights.dhmu_ls_u[i] * w;
875        dhll_u[i] = weights.dh_ls_ls_u[i] * w;
876        d2hmm[i] = weights.d2hmumu[i] * w;
877        d2hml[i] = weights.d2hmu_ls[i] * w;
878        d2hll[i] = weights.d2h_ls_ls[i] * w;
879    }
880    weights.dhmumu_u = dhmm_u;
881    weights.dhmu_ls_u = dhml_u;
882    weights.dh_ls_ls_u = dhll_u;
883    weights.d2hmumu = d2hmm;
884    weights.d2hmu_ls = d2hml;
885    weights.d2h_ls_ls = d2hll;
886}
887
888pub(crate) fn gaussian_jointrow_scalars(
889    y: &Array1<f64>,
890    etamu: &Array1<f64>,
891    eta_ls: &Array1<f64>,
892    weights: &Array1<f64>,
893) -> Result<GaussianJointRowScalars, String> {
894    let nobs = y.len();
895    if etamu.len() != nobs || eta_ls.len() != nobs || weights.len() != nobs {
896        return Err(GamlssError::DimensionMismatch {
897            reason: "Gaussian joint row scalar input size mismatch".to_string(),
898        }
899        .into());
900    }
901    let mut obs_weight = Array1::<f64>::uninit(nobs);
902    let mut standardized_residual = Array1::<f64>::uninit(nobs);
903    let mut inv_sigma = Array1::<f64>::uninit(nobs);
904    let mut kappa = Array1::<f64>::uninit(nobs);
905    let ln2pi = (2.0 * std::f64::consts::PI).ln();
906    // Compute into an indexed temporary first. Parallel collection preserves
907    // row order; scanning the results afterward reports the smallest failing
908    // row deterministically and publishes no partially initialized scalar set.
909    let certified: Vec<Result<GaussianDiagonalRowKernel, String>> = (0..nobs)
910        .into_par_iter()
911        .map(|i| gaussian_diagonal_row_kernel(i, y[i], etamu[i], eta_ls[i], weights[i], ln2pi))
912        .collect();
913    for (i, row) in certified.into_iter().enumerate() {
914        let row = row?;
915        obs_weight[i].write(weights[i]);
916        standardized_residual[i].write(row.standardized_residual);
917        inv_sigma[i].write(row.inv_sigma);
918        kappa[i].write(row.kappa);
919    }
920    // SAFETY: every `MaybeUninit` slot in each of these arrays was written
921    // exactly once in the `for i in 0..nobs` loop above; no slot is read,
922    // moved, or dropped before this point.
923    let (obs_weight, standardized_residual, inv_sigma, kappa) = unsafe {
924        (
925            obs_weight.assume_init(),
926            standardized_residual.assume_init(),
927            inv_sigma.assume_init(),
928            kappa.assume_init(),
929        )
930    };
931    Ok(GaussianJointRowScalars {
932        obs_weight,
933        standardized_residual,
934        inv_sigma,
935        kappa,
936    })
937}
938
939/// Live third-order observed-Hessian contraction emitted from the same stable
940/// [`gaussian_normalized_row`] expression as the observed Hessian.
941pub(crate) fn gaussian_joint_first_directionalweights(
942    scalars: &GaussianJointRowScalars,
943    dotmu: &Array1<f64>,
944    dot_eta: &Array1<f64>,
945) -> (Array1<f64>, Array1<f64>, Array1<f64>) {
946    let nobs = scalars.obs_weight.len();
947    let program = GaussianJointRowProgram::new(scalars);
948    let mut w_u = Array1::<f64>::zeros(nobs);
949    let mut c_u = Array1::<f64>::zeros(nobs);
950    let mut d_u = Array1::<f64>::zeros(nobs);
951    for i in 0..nobs {
952        let matrix = program.row_third_contracted(i, &[dotmu[i], dot_eta[i]]);
953        w_u[i] = matrix[0][0];
954        c_u[i] = matrix[0][1];
955        d_u[i] = matrix[1][1];
956    }
957    (w_u, c_u, d_u)
958}
959
960/// Live fourth-order observed-Hessian contraction emitted from the same stable
961/// [`gaussian_normalized_row`] expression as every lower curvature channel.
962pub(crate) fn gaussian_jointsecond_directionalweights(
963    scalars: &GaussianJointRowScalars,
964    dotmu_u: &Array1<f64>,
965    dot_eta_u: &Array1<f64>,
966    dotmuv: &Array1<f64>,
967    dot_etav: &Array1<f64>,
968) -> (Array1<f64>, Array1<f64>, Array1<f64>) {
969    let nobs = scalars.obs_weight.len();
970    let program = GaussianJointRowProgram::new(scalars);
971    let mut w_uv = Array1::<f64>::zeros(nobs);
972    let mut c_uv = Array1::<f64>::zeros(nobs);
973    let mut d_uv = Array1::<f64>::zeros(nobs);
974    for i in 0..nobs {
975        let matrix = program.row_fourth_contracted(
976            i,
977            &[dotmu_u[i], dot_eta_u[i]],
978            &[dotmuv[i], dot_etav[i]],
979        );
980        w_uv[i] = matrix[0][0];
981        c_uv[i] = matrix[0][1];
982        d_uv[i] = matrix[1][1];
983    }
984    (w_uv, c_uv, d_uv)
985}
986
987pub(crate) fn gaussian_joint_psi_firstweights(
988    scalars: &GaussianJointRowScalars,
989    mu_a: &Array1<f64>,
990    eta_a: &Array1<f64>,
991) -> GaussianJointPsiFirstWeights {
992    let nobs = scalars.obs_weight.len();
993    let program = GaussianJointRowProgram::new(scalars);
994    let mut objective_psirow = Array1::<f64>::uninit(nobs);
995    let mut scoremu = Array1::<f64>::uninit(nobs);
996    let mut score_ls = Array1::<f64>::uninit(nobs);
997    let mut dscoremu = Array1::<f64>::uninit(nobs);
998    let mut dscore_ls = Array1::<f64>::uninit(nobs);
999    let mut hmumu = Array1::<f64>::uninit(nobs);
1000    let mut hmu_ls = Array1::<f64>::uninit(nobs);
1001    let mut h_ls_ls = Array1::<f64>::uninit(nobs);
1002    let mut dhmumu = Array1::<f64>::uninit(nobs);
1003    let mut dhmu_ls = Array1::<f64>::uninit(nobs);
1004    let mut dh_ls_ls = Array1::<f64>::uninit(nobs);
1005    for i in 0..nobs {
1006        let direction = [mu_a[i], eta_a[i]];
1007        let atom = program.row_order2(i);
1008        let score = atom.gradient();
1009        let hessian = [
1010            [atom.hessian_at(0, 0), atom.hessian_at(0, 1)],
1011            [atom.hessian_at(1, 0), atom.hessian_at(1, 1)],
1012        ];
1013        let score_direction = matrix_vector_2(&hessian, &direction);
1014        let hessian_direction = program.row_third_contracted(i, &direction);
1015        objective_psirow[i].write(dot_2(&score, &direction));
1016        scoremu[i].write(score[0]);
1017        score_ls[i].write(score[1]);
1018        dscoremu[i].write(score_direction[0]);
1019        dscore_ls[i].write(score_direction[1]);
1020        hmumu[i].write(hessian[0][0]);
1021        hmu_ls[i].write(hessian[0][1]);
1022        h_ls_ls[i].write(hessian[1][1]);
1023        dhmumu[i].write(hessian_direction[0][0]);
1024        dhmu_ls[i].write(hessian_direction[0][1]);
1025        dh_ls_ls[i].write(hessian_direction[1][1]);
1026    }
1027    // SAFETY: every `MaybeUninit` slot in each field array was written
1028    // exactly once inside the `for i in 0..nobs` loop above.
1029    unsafe {
1030        GaussianJointPsiFirstWeights {
1031            objective_psirow: objective_psirow.assume_init(),
1032            scoremu: scoremu.assume_init(),
1033            score_ls: score_ls.assume_init(),
1034            dscoremu: dscoremu.assume_init(),
1035            dscore_ls: dscore_ls.assume_init(),
1036            hmumu: hmumu.assume_init(),
1037            hmu_ls: hmu_ls.assume_init(),
1038            h_ls_ls: h_ls_ls.assume_init(),
1039            dhmumu: dhmumu.assume_init(),
1040            dhmu_ls: dhmu_ls.assume_init(),
1041            dh_ls_ls: dh_ls_ls.assume_init(),
1042        }
1043    }
1044}
1045
1046pub(crate) fn gaussian_joint_psisecondweights(
1047    scalars: &GaussianJointRowScalars,
1048    mu_a: &Array1<f64>,
1049    eta_a: &Array1<f64>,
1050    mu_b: &Array1<f64>,
1051    eta_b: &Array1<f64>,
1052    mu_ab: &Array1<f64>,
1053    eta_ab: &Array1<f64>,
1054) -> GaussianJointPsiSecondWeights {
1055    let nobs = scalars.obs_weight.len();
1056    let program = GaussianJointRowProgram::new(scalars);
1057    let mut objective_psi_psirow = Array1::<f64>::uninit(nobs);
1058    let mut d2scoremu = Array1::<f64>::uninit(nobs);
1059    let mut d2score_ls = Array1::<f64>::uninit(nobs);
1060    let mut d2hmumu = Array1::<f64>::uninit(nobs);
1061    let mut d2hmu_ls = Array1::<f64>::uninit(nobs);
1062    let mut d2h_ls_ls = Array1::<f64>::uninit(nobs);
1063    for i in 0..nobs {
1064        let direction_a = [mu_a[i], eta_a[i]];
1065        let direction_b = [mu_b[i], eta_b[i]];
1066        let direction_ab = [mu_ab[i], eta_ab[i]];
1067        let atom = program.row_order2(i);
1068        let score = atom.gradient();
1069        let hessian = [
1070            [atom.hessian_at(0, 0), atom.hessian_at(0, 1)],
1071            [atom.hessian_at(1, 0), atom.hessian_at(1, 1)],
1072        ];
1073        let hessian_a = program.row_third_contracted(i, &direction_a);
1074        let hessian_ab = program.row_third_contracted(i, &direction_ab);
1075        let hessian_a_b = program.row_fourth_contracted(i, &direction_a, &direction_b);
1076        let score_second = add_vector_2(
1077            matrix_vector_2(&hessian_a, &direction_b),
1078            matrix_vector_2(&hessian, &direction_ab),
1079        );
1080        let hessian_second = add_matrix_2(hessian_a_b, hessian_ab);
1081        objective_psi_psirow[i].write(
1082            dot_2(&direction_a, &matrix_vector_2(&hessian, &direction_b))
1083                + dot_2(&score, &direction_ab),
1084        );
1085        d2scoremu[i].write(score_second[0]);
1086        d2score_ls[i].write(score_second[1]);
1087        d2hmumu[i].write(hessian_second[0][0]);
1088        d2hmu_ls[i].write(hessian_second[0][1]);
1089        d2h_ls_ls[i].write(hessian_second[1][1]);
1090    }
1091    // SAFETY: every `MaybeUninit` slot in each field array was written
1092    // exactly once inside the `for i in 0..nobs` loop above.
1093    unsafe {
1094        GaussianJointPsiSecondWeights {
1095            objective_psi_psirow: objective_psi_psirow.assume_init(),
1096            d2scoremu: d2scoremu.assume_init(),
1097            d2score_ls: d2score_ls.assume_init(),
1098            d2hmumu: d2hmumu.assume_init(),
1099            d2hmu_ls: d2hmu_ls.assume_init(),
1100            d2h_ls_ls: d2h_ls_ls.assume_init(),
1101        }
1102    }
1103}
1104
1105pub(crate) fn gaussian_joint_psi_mixed_driftweights(
1106    scalars: &GaussianJointRowScalars,
1107    dot_mu: &Array1<f64>,
1108    dot_eta: &Array1<f64>,
1109    mu_a: &Array1<f64>,
1110    eta_a: &Array1<f64>,
1111    dot_mu_a: &Array1<f64>,
1112    dot_eta_a: &Array1<f64>,
1113) -> GaussianJointPsiMixedDriftWeights {
1114    let nobs = scalars.obs_weight.len();
1115    let program = GaussianJointRowProgram::new(scalars);
1116    let mut dhmumu_u = Array1::<f64>::uninit(nobs);
1117    let mut dhmu_ls_u = Array1::<f64>::uninit(nobs);
1118    let mut dh_ls_ls_u = Array1::<f64>::uninit(nobs);
1119    let mut d2hmumu = Array1::<f64>::uninit(nobs);
1120    let mut d2hmu_ls = Array1::<f64>::uninit(nobs);
1121    let mut d2h_ls_ls = Array1::<f64>::uninit(nobs);
1122    for i in 0..nobs {
1123        let drift = [dot_mu[i], dot_eta[i]];
1124        let psi = [mu_a[i], eta_a[i]];
1125        let mixed_direction = [dot_mu_a[i], dot_eta_a[i]];
1126        let hessian_drift = program.row_third_contracted(i, &drift);
1127        let hessian_mixed = add_matrix_2(
1128            program.row_fourth_contracted(i, &drift, &psi),
1129            program.row_third_contracted(i, &mixed_direction),
1130        );
1131        dhmumu_u[i].write(hessian_drift[0][0]);
1132        dhmu_ls_u[i].write(hessian_drift[0][1]);
1133        dh_ls_ls_u[i].write(hessian_drift[1][1]);
1134        d2hmumu[i].write(hessian_mixed[0][0]);
1135        d2hmu_ls[i].write(hessian_mixed[0][1]);
1136        d2h_ls_ls[i].write(hessian_mixed[1][1]);
1137    }
1138    // SAFETY: every `MaybeUninit` slot in each field array was written
1139    // exactly once inside the `for i in 0..nobs` loop above.
1140    unsafe {
1141        GaussianJointPsiMixedDriftWeights {
1142            dhmumu_u: dhmumu_u.assume_init(),
1143            dhmu_ls_u: dhmu_ls_u.assume_init(),
1144            dh_ls_ls_u: dh_ls_ls_u.assume_init(),
1145            d2hmumu: d2hmumu.assume_init(),
1146            d2hmu_ls: d2hmu_ls.assume_init(),
1147            d2h_ls_ls: d2h_ls_ls.assume_init(),
1148        }
1149    }
1150}
1151
1152/// Canonical Gaussian location-scale OBSERVED joint-Hessian row coefficients
1153/// `(mm, ml, ll)` — the SINGLE source of truth for this curvature, shared by
1154/// every representation that assembles the value Hessian (the dense
1155/// `exact_newton_joint_hessian_from_designs` and the matrix-free
1156/// `GaussianLocationScaleHessianWorkspace`). Exact second derivatives of the
1157/// row NLL (`r = y−μ`, `w = a/σ²`, `m = rw`, `n = r²w`, `κ = dlogσ/dη`):
1158///   `mm = ∂²ℓ/∂η_μ²      = w`             (observed ≡ expected — exact),
1159///   `ml = ∂²ℓ/∂η_μ∂η_ls  = 2κm`           (expectation 0 at the truth),
1160///   `ll = ∂²ℓ/∂η_ls²     = κ′(a−n) + 2κ²n` (expectation 2κ²a).
1161/// The LAML criterion `−½log|H+S|` requires the OBSERVED penalized Hessian at
1162/// β̂ (Wood–Pya–Säfken 2016): the earlier block-Fisher object (#684/#566)
1163/// zeroed `ml` and expected `ll`, which drops the cross-block Schur deficit
1164/// `H_σμ(H_μμ+S_μ)⁻¹H_μσ` and the fitted-residual shrinkage `E[n̂]≈a(1−h_μ)`
1165/// — both overstate σ-block information and bias λ̂_σ upward on the flat scale
1166/// surface (#1561: log-σ over-smoothing; same dof genus as #2133). At a
1167/// true-null/flat σ surface `n→a`, `m→0`, so observed → Fisher and null
1168/// behavior is unchanged (SPEC: defaults recover the null). Indefiniteness of
1169/// the observed joint Hessian is handled by the existing #365 modified-Newton
1170/// reflection on the inner path and the spectral PD-floor on the criterion
1171/// log-det. Routing every path through this one constructor keeps the #684
1172/// cross-block drift structurally impossible.
1173pub(crate) fn gaussian_locscale_observed_joint_row_coeffs(
1174    rows: &GaussianJointRowScalars,
1175) -> (Array1<f64>, Array1<f64>, Array1<f64>) {
1176    let n = rows.obs_weight.len();
1177    let program = GaussianJointRowProgram::new(rows);
1178    let mut mm = Array1::<f64>::zeros(n);
1179    let mut ml = Array1::<f64>::zeros(n);
1180    let mut ll = Array1::<f64>::zeros(n);
1181    for row in 0..n {
1182        let atom = program.row_order2(row);
1183        mm[row] = atom.hessian_at(0, 0);
1184        ml[row] = atom.hessian_at(0, 1);
1185        ll[row] = atom.hessian_at(1, 1);
1186    }
1187    (mm, ml, ll)
1188}
1189
1190pub(crate) fn gaussian_joint_hessian_from_designs(
1191    xmu: &DenseOrOperator<'_>,
1192    x_ls: &DenseOrOperator<'_>,
1193    hmumu_coeff: &Array1<f64>,
1194    hmu_ls_coeff: &Array1<f64>,
1195    h_ls_ls_coeff: &Array1<f64>,
1196) -> Result<Array2<f64>, String> {
1197    if xmu.nrows() != hmumu_coeff.len()
1198        || xmu.nrows() != hmu_ls_coeff.len()
1199        || xmu.nrows() != h_ls_ls_coeff.len()
1200        || x_ls.nrows() != xmu.nrows()
1201    {
1202        return Err(GamlssError::DimensionMismatch { reason: format!(
1203            "gaussian_joint_hessian_from_designs dimension mismatch: xmu {}x{}, x_ls {}x{}, coeffs {}/{}/{}",
1204            xmu.nrows(),
1205            xmu.ncols(),
1206            x_ls.nrows(),
1207            x_ls.ncols(),
1208            hmumu_coeff.len(),
1209            hmu_ls_coeff.len(),
1210            h_ls_ls_coeff.len()
1211        ) }.into());
1212    }
1213
1214    let n = xmu.nrows();
1215    let pmu = xmu.ncols();
1216    let p_ls = x_ls.ncols();
1217    let total = pmu + p_ls;
1218    let mut out = Array2::<f64>::zeros((total, total));
1219    for rows in exact_design_row_chunks(n, pmu.max(p_ls)) {
1220        let xmu_chunk = xmu.row_chunk(rows.clone())?;
1221        let xls_chunk = x_ls.row_chunk(rows.clone())?;
1222        let hmumu = hmumu_coeff.slice(s![rows.clone()]);
1223        let hmu_ls = hmu_ls_coeff.slice(s![rows.clone()]);
1224        let h_ls_ls = h_ls_ls_coeff.slice(s![rows.clone()]);
1225        let chunk_hessian =
1226            fast_joint_hessian_2x2(&xmu_chunk, &xls_chunk, &hmumu, &hmu_ls, &h_ls_ls);
1227        out += &chunk_hessian;
1228    }
1229    Ok(out)
1230}
1231
1232pub(crate) fn gaussian_joint_psihessian_fromweights(
1233    xmu: &Array2<f64>,
1234    x_ls: &Array2<f64>,
1235    xmu_psi: CustomFamilyPsiLinearMapRef<'_>,
1236    x_ls_psi: CustomFamilyPsiLinearMapRef<'_>,
1237    weights: &GaussianJointPsiFirstWeights,
1238) -> Result<Array2<f64>, String> {
1239    // For the symmetric blocks (hmumu, h_ls_ls), the pair
1240    //   X_psi^T D X  and  X^T D X_psi
1241    // are transposes of each other, so compute one and add its transpose.
1242    let a_mu = weighted_crossprod_psi_maps(
1243        xmu_psi,
1244        weights.hmumu.view(),
1245        CustomFamilyPsiLinearMapRef::Dense(xmu),
1246    )?;
1247    let hmumu = &a_mu + &a_mu.t() + &xt_diag_x_dense(xmu, &weights.dhmumu)?;
1248    let hmu_ls = weighted_crossprod_psi_maps(
1249        xmu_psi,
1250        weights.hmu_ls.view(),
1251        CustomFamilyPsiLinearMapRef::Dense(x_ls),
1252    )? + &weighted_crossprod_psi_maps(
1253        CustomFamilyPsiLinearMapRef::Dense(xmu),
1254        weights.hmu_ls.view(),
1255        x_ls_psi,
1256    )? + &xt_diag_y_dense(xmu, &weights.dhmu_ls, x_ls)?;
1257    let a_ls = weighted_crossprod_psi_maps(
1258        x_ls_psi,
1259        weights.h_ls_ls.view(),
1260        CustomFamilyPsiLinearMapRef::Dense(x_ls),
1261    )?;
1262    let h_ls_ls = &a_ls + &a_ls.t() + &xt_diag_x_dense(x_ls, &weights.dh_ls_ls)?;
1263    Ok(gaussian_pack_joint_symmetrichessian(
1264        &hmumu, &hmu_ls, &h_ls_ls,
1265    ))
1266}
1267
1268pub(crate) fn build_two_block_custom_family_joint_psi_operator_from_actions(
1269    left_action: Option<CustomFamilyPsiDesignAction>,
1270    right_action: Option<CustomFamilyPsiDesignAction>,
1271    left_range: std::ops::Range<usize>,
1272    right_range: std::ops::Range<usize>,
1273    left_design: &Array2<f64>,
1274    right_design: &Array2<f64>,
1275    left_weights: &Array1<f64>,
1276    cross_weights: &Array1<f64>,
1277    right_weights: &Array1<f64>,
1278    left_drift_weights: &Array1<f64>,
1279    cross_drift_weights: &Array1<f64>,
1280    right_drift_weights: &Array1<f64>,
1281) -> Result<Option<std::sync::Arc<dyn gam_problem::HyperOperator>>, String> {
1282    if left_action.is_none() && right_action.is_none() {
1283        return Ok(None);
1284    }
1285
1286    let total = left_design.ncols() + right_design.ncols();
1287    let channels = vec![
1288        CustomFamilyJointDesignChannel::new(left_range, shared_dense_arc(left_design), left_action),
1289        CustomFamilyJointDesignChannel::new(
1290            right_range,
1291            shared_dense_arc(right_design),
1292            right_action,
1293        ),
1294    ];
1295    let pair_contributions = vec![
1296        CustomFamilyJointDesignPairContribution::new(
1297            0,
1298            0,
1299            left_weights.clone(),
1300            left_drift_weights.clone(),
1301        ),
1302        CustomFamilyJointDesignPairContribution::new(
1303            0,
1304            1,
1305            cross_weights.clone(),
1306            cross_drift_weights.clone(),
1307        ),
1308        CustomFamilyJointDesignPairContribution::new(
1309            1,
1310            0,
1311            cross_weights.clone(),
1312            cross_drift_weights.clone(),
1313        ),
1314        CustomFamilyJointDesignPairContribution::new(
1315            1,
1316            1,
1317            right_weights.clone(),
1318            right_drift_weights.clone(),
1319        ),
1320    ];
1321
1322    Ok(Some(std::sync::Arc::new(
1323        CustomFamilyJointPsiOperator::new(total, channels, pair_contributions),
1324    )))
1325}
1326
1327pub(crate) fn gaussian_joint_psisecondhessian_fromweights(
1328    xmu: &Array2<f64>,
1329    x_ls: &Array2<f64>,
1330    xmu_i: CustomFamilyPsiLinearMapRef<'_>,
1331    x_ls_i: CustomFamilyPsiLinearMapRef<'_>,
1332    xmu_j: CustomFamilyPsiLinearMapRef<'_>,
1333    x_ls_j: CustomFamilyPsiLinearMapRef<'_>,
1334    xmu_ab: CustomFamilyPsiLinearMapRef<'_>,
1335    x_ls_ab: CustomFamilyPsiLinearMapRef<'_>,
1336    weights_i: &GaussianJointPsiFirstWeights,
1337    weights_j: &GaussianJointPsiFirstWeights,
1338    secondweights: &GaussianJointPsiSecondWeights,
1339) -> Result<Array2<f64>, String> {
1340    // Exploit transpose symmetry: X_a^T D X_b and X_b^T D X_a are transposes.
1341    // For each such pair in the symmetric blocks (hmumu, h_ls_ls), compute one
1342    // and add its transpose, halving the number of O(np²) products.
1343    let a_ab_mu = weighted_crossprod_psi_maps(
1344        xmu_ab,
1345        weights_i.hmumu.view(),
1346        CustomFamilyPsiLinearMapRef::Dense(xmu),
1347    )?;
1348    let a_ij_mu = weighted_crossprod_psi_maps(xmu_i, weights_i.hmumu.view(), xmu_j)?;
1349    let a_iwj_mu = weighted_crossprod_psi_maps(
1350        xmu_i,
1351        weights_j.dhmumu.view(),
1352        CustomFamilyPsiLinearMapRef::Dense(xmu),
1353    )?;
1354    let a_jwi_mu = weighted_crossprod_psi_maps(
1355        xmu_j,
1356        weights_i.dhmumu.view(),
1357        CustomFamilyPsiLinearMapRef::Dense(xmu),
1358    )?;
1359    let hmumu = &a_ab_mu
1360        + &a_ab_mu.t()
1361        + &a_ij_mu
1362        + a_ij_mu.t()
1363        + &a_iwj_mu
1364        + a_iwj_mu.t()
1365        + &a_jwi_mu
1366        + a_jwi_mu.t()
1367        + &xt_diag_x_dense(xmu, &secondweights.d2hmumu)?;
1368    let hmu_ls = weighted_crossprod_psi_maps(
1369        xmu_ab,
1370        weights_i.hmu_ls.view(),
1371        CustomFamilyPsiLinearMapRef::Dense(x_ls),
1372    )? + &weighted_crossprod_psi_maps(xmu_i, weights_i.hmu_ls.view(), x_ls_j)?
1373        + &weighted_crossprod_psi_maps(xmu_j, weights_i.hmu_ls.view(), x_ls_i)?
1374        + &weighted_crossprod_psi_maps(
1375            xmu_i,
1376            weights_j.dhmu_ls.view(),
1377            CustomFamilyPsiLinearMapRef::Dense(x_ls),
1378        )?
1379        + &weighted_crossprod_psi_maps(
1380            xmu_j,
1381            weights_i.dhmu_ls.view(),
1382            CustomFamilyPsiLinearMapRef::Dense(x_ls),
1383        )?
1384        + &weighted_crossprod_psi_maps(
1385            CustomFamilyPsiLinearMapRef::Dense(xmu),
1386            weights_i.dhmu_ls.view(),
1387            x_ls_j,
1388        )?
1389        + &weighted_crossprod_psi_maps(
1390            CustomFamilyPsiLinearMapRef::Dense(xmu),
1391            weights_j.dhmu_ls.view(),
1392            x_ls_i,
1393        )?
1394        + &xt_diag_y_dense(xmu, &secondweights.d2hmu_ls, x_ls)?
1395        + &weighted_crossprod_psi_maps(
1396            CustomFamilyPsiLinearMapRef::Dense(xmu),
1397            weights_i.hmu_ls.view(),
1398            x_ls_ab,
1399        )?;
1400    let a_ab_ls = weighted_crossprod_psi_maps(
1401        x_ls_ab,
1402        weights_i.h_ls_ls.view(),
1403        CustomFamilyPsiLinearMapRef::Dense(x_ls),
1404    )?;
1405    let a_ij_ls = weighted_crossprod_psi_maps(x_ls_i, weights_i.h_ls_ls.view(), x_ls_j)?;
1406    let a_iwj_ls = weighted_crossprod_psi_maps(
1407        x_ls_i,
1408        weights_j.dh_ls_ls.view(),
1409        CustomFamilyPsiLinearMapRef::Dense(x_ls),
1410    )?;
1411    let a_jwi_ls = weighted_crossprod_psi_maps(
1412        x_ls_j,
1413        weights_i.dh_ls_ls.view(),
1414        CustomFamilyPsiLinearMapRef::Dense(x_ls),
1415    )?;
1416    let h_ls_ls = &a_ab_ls
1417        + &a_ab_ls.t()
1418        + &a_ij_ls
1419        + a_ij_ls.t()
1420        + &a_iwj_ls
1421        + a_iwj_ls.t()
1422        + &a_jwi_ls
1423        + a_jwi_ls.t()
1424        + &xt_diag_x_dense(x_ls, &secondweights.d2h_ls_ls)?;
1425    Ok(gaussian_pack_joint_symmetrichessian(
1426        &hmumu, &hmu_ls, &h_ls_ls,
1427    ))
1428}
1429
1430pub(crate) fn gaussian_joint_psi_mixedhessian_drift_fromweights(
1431    xmu: &Array2<f64>,
1432    x_ls: &Array2<f64>,
1433    xmu_psi: CustomFamilyPsiLinearMapRef<'_>,
1434    x_ls_psi: CustomFamilyPsiLinearMapRef<'_>,
1435    mixedweights: &GaussianJointPsiMixedDriftWeights,
1436) -> Result<Array2<f64>, String> {
1437    let a_mu = weighted_crossprod_psi_maps(
1438        xmu_psi,
1439        mixedweights.dhmumu_u.view(),
1440        CustomFamilyPsiLinearMapRef::Dense(xmu),
1441    )?;
1442    let hmumu = &a_mu + &a_mu.t() + &xt_diag_x_dense(xmu, &mixedweights.d2hmumu)?;
1443    let hmu_ls = weighted_crossprod_psi_maps(
1444        xmu_psi,
1445        mixedweights.dhmu_ls_u.view(),
1446        CustomFamilyPsiLinearMapRef::Dense(x_ls),
1447    )? + &weighted_crossprod_psi_maps(
1448        CustomFamilyPsiLinearMapRef::Dense(xmu),
1449        mixedweights.dhmu_ls_u.view(),
1450        x_ls_psi,
1451    )? + &xt_diag_y_dense(xmu, &mixedweights.d2hmu_ls, x_ls)?;
1452    let a_ls = weighted_crossprod_psi_maps(
1453        x_ls_psi,
1454        mixedweights.dh_ls_ls_u.view(),
1455        CustomFamilyPsiLinearMapRef::Dense(x_ls),
1456    )?;
1457    let h_ls_ls = &a_ls + &a_ls.t() + &xt_diag_x_dense(x_ls, &mixedweights.d2h_ls_ls)?;
1458    Ok(gaussian_pack_joint_symmetrichessian(
1459        &hmumu, &hmu_ls, &h_ls_ls,
1460    ))
1461}
1462
1463#[inline]
1464pub(crate) fn exp_sigma_derivs_up_to_fourth_array(
1465    eta: ArrayView1<'_, f64>,
1466) -> (
1467    Array1<f64>,
1468    Array1<f64>,
1469    Array1<f64>,
1470    Array1<f64>,
1471    Array1<f64>,
1472) {
1473    use rayon::iter::{IntoParallelIterator, ParallelIterator};
1474    let n = eta.len();
1475    let tuples: Vec<(f64, f64, f64, f64, f64)> = (0..n)
1476        .into_par_iter()
1477        .map(|i| exp_sigma_derivs_up_to_fourth_scalar(eta[i]))
1478        .collect();
1479    let mut sigma = Array1::<f64>::zeros(n);
1480    let mut d1 = Array1::<f64>::zeros(n);
1481    let mut d2 = Array1::<f64>::zeros(n);
1482    let mut d3 = Array1::<f64>::zeros(n);
1483    let mut d4 = Array1::<f64>::zeros(n);
1484    for (i, (s_i, d1_i, d2_i, d3_i, d4_i)) in tuples.into_iter().enumerate() {
1485        sigma[i] = s_i;
1486        d1[i] = d1_i;
1487        d2[i] = d2_i;
1488        d3[i] = d3_i;
1489        d4[i] = d4_i;
1490    }
1491    (sigma, d1, d2, d3, d4)
1492}
1493
1494#[cfg(test)]
1495mod observed_single_source_oracle_tests {
1496    //! #932 doctrine oracle for the generated Gaussian location-scale row
1497    //! program and its OBSERVED joint-Hessian tower.
1498    //!
1499    //! Every non-wiggle live score/Hessian/ψ channel is a chain-rule projection
1500    //! of the symbolic order2/third/fourth lowerings emitted from
1501    //! `gaussian_normalized_row`. The generic nested-jet evaluator and the
1502    //! likelihood-only finite differences below are independent witnesses.
1503    //!
1504    //! MECHANICAL SOURCE (no hand math reused):
1505    //!  * The per-row negative log-likelihood is `ρ(μ,η)=−ℓ(μ,η)`, evaluated by
1506    //!    the production row kernel `gaussian_diagonal_row_kernel`.
1507    //!  * Its OBSERVED 2×2 Hessian in `(μ,η_ls)` and every directional chain are
1508    //!    taken by central finite differences of that likelihood alone.
1509
1510    use super::*;
1511    use ndarray::array;
1512
1513    /// Row negative log-likelihood from the production kernel (likelihood only,
1514    /// no curvature coefficients involved).
1515    fn row_nll(y: f64, mu: f64, eta_ls: f64, a: f64) -> f64 {
1516        let ln2pi = (2.0 * std::f64::consts::PI).ln();
1517        -gaussian_diagonal_row_kernel(0, y, mu, eta_ls, a, ln2pi)
1518            .expect("representable Gaussian oracle row")
1519            .log_likelihood
1520    }
1521
1522    fn gradient_fd(y: f64, mu: f64, eta_ls: f64, a: f64, h: f64) -> [f64; 2] {
1523        [
1524            (row_nll(y, mu + h, eta_ls, a) - row_nll(y, mu - h, eta_ls, a)) / (2.0 * h),
1525            (row_nll(y, mu, eta_ls + h, a) - row_nll(y, mu, eta_ls - h, a)) / (2.0 * h),
1526        ]
1527    }
1528
1529    /// Observed 2×2 Hessian of the row NLL in `(μ, η_ls)` by central FD.
1530    /// Returns `(H_μμ, H_{μ,ls}, H_{ls,ls})`.
1531    fn observed_hessian_fd(y: f64, mu: f64, eta_ls: f64, a: f64, h: f64) -> (f64, f64, f64) {
1532        let hmm = (row_nll(y, mu + h, eta_ls, a) - 2.0 * row_nll(y, mu, eta_ls, a)
1533            + row_nll(y, mu - h, eta_ls, a))
1534            / (h * h);
1535        let hll = (row_nll(y, mu, eta_ls + h, a) - 2.0 * row_nll(y, mu, eta_ls, a)
1536            + row_nll(y, mu, eta_ls - h, a))
1537            / (h * h);
1538        let hml = (row_nll(y, mu + h, eta_ls + h, a)
1539            - row_nll(y, mu + h, eta_ls - h, a)
1540            - row_nll(y, mu - h, eta_ls + h, a)
1541            + row_nll(y, mu - h, eta_ls - h, a))
1542            / (4.0 * h * h);
1543        (hmm, hml, hll)
1544    }
1545
1546    fn hessian_fd(y: f64, mu: f64, eta_ls: f64, a: f64, h: f64) -> [[f64; 2]; 2] {
1547        let (mm, ml, ll) = observed_hessian_fd(y, mu, eta_ls, a, h);
1548        [[mm, ml], [ml, ll]]
1549    }
1550
1551    fn path_point(
1552        base: [f64; 2],
1553        direction_a: [f64; 2],
1554        direction_b: [f64; 2],
1555        direction_ab: [f64; 2],
1556        s: f64,
1557        t: f64,
1558    ) -> [f64; 2] {
1559        [
1560            base[0] + s * direction_a[0] + t * direction_b[0] + s * t * direction_ab[0],
1561            base[1] + s * direction_a[1] + t * direction_b[1] + s * t * direction_ab[1],
1562        ]
1563    }
1564
1565    fn mixed_value_fd(
1566        y: f64,
1567        base: [f64; 2],
1568        a: f64,
1569        direction_a: [f64; 2],
1570        direction_b: [f64; 2],
1571        direction_ab: [f64; 2],
1572        h: f64,
1573    ) -> f64 {
1574        let pp = path_point(base, direction_a, direction_b, direction_ab, h, h);
1575        let pm = path_point(base, direction_a, direction_b, direction_ab, h, -h);
1576        let mp = path_point(base, direction_a, direction_b, direction_ab, -h, h);
1577        let mm = path_point(base, direction_a, direction_b, direction_ab, -h, -h);
1578        (row_nll(y, pp[0], pp[1], a) - row_nll(y, pm[0], pm[1], a) - row_nll(y, mp[0], mp[1], a)
1579            + row_nll(y, mm[0], mm[1], a))
1580            / (4.0 * h * h)
1581    }
1582
1583    fn mixed_gradient_fd(
1584        y: f64,
1585        base: [f64; 2],
1586        a: f64,
1587        direction_a: [f64; 2],
1588        direction_b: [f64; 2],
1589        direction_ab: [f64; 2],
1590        outer_h: f64,
1591        inner_h: f64,
1592    ) -> [f64; 2] {
1593        let pp = path_point(
1594            base,
1595            direction_a,
1596            direction_b,
1597            direction_ab,
1598            outer_h,
1599            outer_h,
1600        );
1601        let pm = path_point(
1602            base,
1603            direction_a,
1604            direction_b,
1605            direction_ab,
1606            outer_h,
1607            -outer_h,
1608        );
1609        let mp = path_point(
1610            base,
1611            direction_a,
1612            direction_b,
1613            direction_ab,
1614            -outer_h,
1615            outer_h,
1616        );
1617        let mm = path_point(
1618            base,
1619            direction_a,
1620            direction_b,
1621            direction_ab,
1622            -outer_h,
1623            -outer_h,
1624        );
1625        let gpp = gradient_fd(y, pp[0], pp[1], a, inner_h);
1626        let gpm = gradient_fd(y, pm[0], pm[1], a, inner_h);
1627        let gmp = gradient_fd(y, mp[0], mp[1], a, inner_h);
1628        let gmm = gradient_fd(y, mm[0], mm[1], a, inner_h);
1629        std::array::from_fn(|axis| {
1630            (gpp[axis] - gpm[axis] - gmp[axis] + gmm[axis]) / (4.0 * outer_h * outer_h)
1631        })
1632    }
1633
1634    fn mixed_hessian_fd(
1635        y: f64,
1636        base: [f64; 2],
1637        a: f64,
1638        direction_a: [f64; 2],
1639        direction_b: [f64; 2],
1640        direction_ab: [f64; 2],
1641        outer_h: f64,
1642        inner_h: f64,
1643    ) -> [[f64; 2]; 2] {
1644        let pp = path_point(
1645            base,
1646            direction_a,
1647            direction_b,
1648            direction_ab,
1649            outer_h,
1650            outer_h,
1651        );
1652        let pm = path_point(
1653            base,
1654            direction_a,
1655            direction_b,
1656            direction_ab,
1657            outer_h,
1658            -outer_h,
1659        );
1660        let mp = path_point(
1661            base,
1662            direction_a,
1663            direction_b,
1664            direction_ab,
1665            -outer_h,
1666            outer_h,
1667        );
1668        let mm = path_point(
1669            base,
1670            direction_a,
1671            direction_b,
1672            direction_ab,
1673            -outer_h,
1674            -outer_h,
1675        );
1676        let hpp = hessian_fd(y, pp[0], pp[1], a, inner_h);
1677        let hpm = hessian_fd(y, pm[0], pm[1], a, inner_h);
1678        let hmp = hessian_fd(y, mp[0], mp[1], a, inner_h);
1679        let hmm = hessian_fd(y, mm[0], mm[1], a, inner_h);
1680        std::array::from_fn(|row| {
1681            std::array::from_fn(|column| {
1682                (hpp[row][column] - hpm[row][column] - hmp[row][column] + hmm[row][column])
1683                    / (4.0 * outer_h * outer_h)
1684            })
1685        })
1686    }
1687
1688    fn assert_close(actual: f64, expected: f64, tolerance: f64, label: &str) {
1689        let band = tolerance * actual.abs().max(expected.abs()).max(1.0);
1690        assert!(
1691            (actual - expected).abs() <= band,
1692            "{label}: actual={actual:+.15e} expected={expected:+.15e} band={band:.3e}"
1693        );
1694    }
1695
1696    fn assert_matrix_close(
1697        actual: &[[f64; 2]; 2],
1698        expected: &[[f64; 2]; 2],
1699        tolerance: f64,
1700        label: &str,
1701    ) {
1702        for row in 0..2 {
1703            for column in 0..2 {
1704                assert_close(
1705                    actual[row][column],
1706                    expected[row][column],
1707                    tolerance,
1708                    &format!("{label}[{row},{column}]"),
1709                );
1710            }
1711        }
1712    }
1713
1714    /// Production observed joint-Hessian coefficients for a single row.
1715    fn production_observed_row(y: f64, mu: f64, eta_ls: f64, a: f64) -> (f64, f64, f64) {
1716        let rows = gaussian_jointrow_scalars(&array![y], &array![mu], &array![eta_ls], &array![a])
1717            .expect("row scalars");
1718        let (mm, ml, ll) = gaussian_locscale_observed_joint_row_coeffs(&rows);
1719        (mm[0], ml[0], ll[0])
1720    }
1721
1722    #[test]
1723    fn generated_gaussian_psi_chain_matches_generic_nested_jet_all_channels_932() {
1724        use gam_math::jet_tower::{
1725            program_fourth_contracted, program_row_kernel, program_third_contracted,
1726        };
1727
1728        let rows =
1729            gaussian_jointrow_scalars(&array![0.55], &array![0.3], &array![-0.4], &array![1.7])
1730                .expect("row scalars");
1731        let program = GaussianJointRowProgram::new(&rows);
1732        let direction_a = [0.5, -0.7];
1733        let direction_b = [-0.3, 0.8];
1734        let direction_ab = [0.2, -0.15];
1735        let drift = [0.6, 0.25];
1736        let psi = [-0.4, 0.75];
1737        let drift_psi = [0.12, -0.18];
1738
1739        let atom = program.row_order2(0);
1740        let (jet_value, jet_score, jet_hessian) =
1741            program_row_kernel(&program, 0).expect("generic order2");
1742        assert_close(atom.value(), jet_value, 1e-12, "row value");
1743        for axis in 0..2 {
1744            assert_close(
1745                atom.gradient()[axis],
1746                jet_score[axis],
1747                1e-12,
1748                &format!("row score[{axis}]"),
1749            );
1750        }
1751        let generated_hessian = [
1752            [atom.hessian_at(0, 0), atom.hessian_at(0, 1)],
1753            [atom.hessian_at(1, 0), atom.hessian_at(1, 1)],
1754        ];
1755        assert_matrix_close(&generated_hessian, &jet_hessian, 1e-12, "row Hessian");
1756
1757        let jet_third_a =
1758            program_third_contracted(&program, 0, &direction_a).expect("generic third a");
1759        let jet_third_ab =
1760            program_third_contracted(&program, 0, &direction_ab).expect("generic third ab");
1761        let jet_fourth_ab = program_fourth_contracted(&program, 0, &direction_a, &direction_b)
1762            .expect("generic fourth ab");
1763        assert_matrix_close(
1764            &program.row_third_contracted(0, &direction_a),
1765            &jet_third_a,
1766            1e-9,
1767            "generated dH",
1768        );
1769        assert_matrix_close(
1770            &program.row_fourth_contracted(0, &direction_a, &direction_b),
1771            &jet_fourth_ab,
1772            1e-9,
1773            "generated d2H",
1774        );
1775
1776        let first = gaussian_joint_psi_firstweights(
1777            &rows,
1778            &array![direction_a[0]],
1779            &array![direction_a[1]],
1780        );
1781        let expected_score_a = matrix_vector_2(&jet_hessian, &direction_a);
1782        assert_close(
1783            first.objective_psirow[0],
1784            dot_2(&jet_score, &direction_a),
1785            1e-9,
1786            "first psi objective",
1787        );
1788        assert_close(first.scoremu[0], jet_score[0], 1e-12, "first score mu");
1789        assert_close(first.score_ls[0], jet_score[1], 1e-12, "first score ls");
1790        assert_close(
1791            first.dscoremu[0],
1792            expected_score_a[0],
1793            1e-9,
1794            "first dscore mu",
1795        );
1796        assert_close(
1797            first.dscore_ls[0],
1798            expected_score_a[1],
1799            1e-9,
1800            "first dscore ls",
1801        );
1802        assert_close(first.hmumu[0], jet_hessian[0][0], 1e-12, "first H mm");
1803        assert_close(first.hmu_ls[0], jet_hessian[0][1], 1e-12, "first H ml");
1804        assert_close(first.h_ls_ls[0], jet_hessian[1][1], 1e-12, "first H ll");
1805        assert_close(first.dhmumu[0], jet_third_a[0][0], 1e-9, "first dH mm");
1806        assert_close(first.dhmu_ls[0], jet_third_a[0][1], 1e-9, "first dH ml");
1807        assert_close(first.dh_ls_ls[0], jet_third_a[1][1], 1e-9, "first dH ll");
1808
1809        let second = gaussian_joint_psisecondweights(
1810            &rows,
1811            &array![direction_a[0]],
1812            &array![direction_a[1]],
1813            &array![direction_b[0]],
1814            &array![direction_b[1]],
1815            &array![direction_ab[0]],
1816            &array![direction_ab[1]],
1817        );
1818        let expected_second_score = add_vector_2(
1819            matrix_vector_2(&jet_third_a, &direction_b),
1820            matrix_vector_2(&jet_hessian, &direction_ab),
1821        );
1822        let expected_second_hessian = add_matrix_2(jet_fourth_ab, jet_third_ab);
1823        assert_close(
1824            second.objective_psi_psirow[0],
1825            dot_2(&direction_a, &matrix_vector_2(&jet_hessian, &direction_b))
1826                + dot_2(&jet_score, &direction_ab),
1827            1e-9,
1828            "second psi objective",
1829        );
1830        assert_close(
1831            second.d2scoremu[0],
1832            expected_second_score[0],
1833            1e-9,
1834            "second psi score mu",
1835        );
1836        assert_close(
1837            second.d2score_ls[0],
1838            expected_second_score[1],
1839            1e-9,
1840            "second psi score ls",
1841        );
1842        assert_close(
1843            second.d2hmumu[0],
1844            expected_second_hessian[0][0],
1845            1e-9,
1846            "second psi H mm",
1847        );
1848        assert_close(
1849            second.d2hmu_ls[0],
1850            expected_second_hessian[0][1],
1851            1e-9,
1852            "second psi H ml",
1853        );
1854        assert_close(
1855            second.d2h_ls_ls[0],
1856            expected_second_hessian[1][1],
1857            1e-9,
1858            "second psi H ll",
1859        );
1860
1861        let mixed = gaussian_joint_psi_mixed_driftweights(
1862            &rows,
1863            &array![drift[0]],
1864            &array![drift[1]],
1865            &array![psi[0]],
1866            &array![psi[1]],
1867            &array![drift_psi[0]],
1868            &array![drift_psi[1]],
1869        );
1870        let jet_third_drift =
1871            program_third_contracted(&program, 0, &drift).expect("generic third drift");
1872        let jet_mixed_hessian = add_matrix_2(
1873            program_fourth_contracted(&program, 0, &drift, &psi).expect("generic fourth mixed"),
1874            program_third_contracted(&program, 0, &drift_psi)
1875                .expect("generic third mixed direction"),
1876        );
1877        assert_close(
1878            mixed.dhmumu_u[0],
1879            jet_third_drift[0][0],
1880            1e-9,
1881            "mixed dH mm",
1882        );
1883        assert_close(
1884            mixed.dhmu_ls_u[0],
1885            jet_third_drift[0][1],
1886            1e-9,
1887            "mixed dH ml",
1888        );
1889        assert_close(
1890            mixed.dh_ls_ls_u[0],
1891            jet_third_drift[1][1],
1892            1e-9,
1893            "mixed dH ll",
1894        );
1895        assert_close(
1896            mixed.d2hmumu[0],
1897            jet_mixed_hessian[0][0],
1898            1e-9,
1899            "mixed d2H mm",
1900        );
1901        assert_close(
1902            mixed.d2hmu_ls[0],
1903            jet_mixed_hessian[0][1],
1904            1e-9,
1905            "mixed d2H ml",
1906        );
1907        assert_close(
1908            mixed.d2h_ls_ls[0],
1909            jet_mixed_hessian[1][1],
1910            1e-9,
1911            "mixed d2H ll",
1912        );
1913    }
1914
1915    #[test]
1916    fn generated_gaussian_psi_chain_matches_likelihood_finite_differences_932() {
1917        let y = 0.55;
1918        let base = [0.3, -0.4];
1919        let weight = 1.7;
1920        let direction_a = [0.5, -0.7];
1921        let direction_b = [-0.3, 0.8];
1922        let direction_ab = [0.2, -0.15];
1923        let drift = [0.6, 0.25];
1924        let psi = [-0.4, 0.75];
1925        let drift_psi = [0.12, -0.18];
1926        let rows = gaussian_jointrow_scalars(
1927            &array![y],
1928            &array![base[0]],
1929            &array![base[1]],
1930            &array![weight],
1931        )
1932        .expect("row scalars");
1933        let program = GaussianJointRowProgram::new(&rows);
1934        let atom = program.row_order2(0);
1935
1936        let normalized_value =
1937            0.5 * weight * rows.standardized_residual[0] * rows.standardized_residual[0];
1938        assert_close(
1939            atom.value(),
1940            normalized_value,
1941            1e-12,
1942            "normalized row value",
1943        );
1944        let score_fd = gradient_fd(y, base[0], base[1], weight, 2e-5);
1945        let hessian_fd0 = hessian_fd(y, base[0], base[1], weight, 1e-3);
1946        for axis in 0..2 {
1947            assert_close(
1948                atom.gradient()[axis],
1949                score_fd[axis],
1950                2e-8,
1951                &format!("score FD {axis}"),
1952            );
1953        }
1954        let atom_hessian = [
1955            [atom.hessian_at(0, 0), atom.hessian_at(0, 1)],
1956            [atom.hessian_at(1, 0), atom.hessian_at(1, 1)],
1957        ];
1958        assert_matrix_close(&atom_hessian, &hessian_fd0, 3e-6, "Hessian FD");
1959
1960        let first = gaussian_joint_psi_firstweights(
1961            &rows,
1962            &array![direction_a[0]],
1963            &array![direction_a[1]],
1964        );
1965        let first_step = 1e-2;
1966        let plus = [
1967            base[0] + first_step * direction_a[0],
1968            base[1] + first_step * direction_a[1],
1969        ];
1970        let minus = [
1971            base[0] - first_step * direction_a[0],
1972            base[1] - first_step * direction_a[1],
1973        ];
1974        let objective_first_fd = (row_nll(y, plus[0], plus[1], weight)
1975            - row_nll(y, minus[0], minus[1], weight))
1976            / (2.0 * first_step);
1977        let gradient_plus = gradient_fd(y, plus[0], plus[1], weight, 2e-5);
1978        let gradient_minus = gradient_fd(y, minus[0], minus[1], weight, 2e-5);
1979        let score_first_fd: [f64; 2] = std::array::from_fn(|axis| {
1980            (gradient_plus[axis] - gradient_minus[axis]) / (2.0 * first_step)
1981        });
1982        let hessian_plus = hessian_fd(y, plus[0], plus[1], weight, 1e-3);
1983        let hessian_minus = hessian_fd(y, minus[0], minus[1], weight, 1e-3);
1984        let hessian_first_fd: [[f64; 2]; 2] = std::array::from_fn(|row| {
1985            std::array::from_fn(|column| {
1986                (hessian_plus[row][column] - hessian_minus[row][column]) / (2.0 * first_step)
1987            })
1988        });
1989        assert_close(
1990            first.objective_psirow[0],
1991            objective_first_fd,
1992            2e-4,
1993            "first objective FD",
1994        );
1995        assert_close(
1996            first.dscoremu[0],
1997            score_first_fd[0],
1998            2e-4,
1999            "first score mu FD",
2000        );
2001        assert_close(
2002            first.dscore_ls[0],
2003            score_first_fd[1],
2004            2e-4,
2005            "first score ls FD",
2006        );
2007        assert_close(
2008            first.dhmumu[0],
2009            hessian_first_fd[0][0],
2010            2e-3,
2011            "first H mm FD",
2012        );
2013        assert_close(
2014            first.dhmu_ls[0],
2015            hessian_first_fd[0][1],
2016            2e-3,
2017            "first H ml FD",
2018        );
2019        assert_close(
2020            first.dh_ls_ls[0],
2021            hessian_first_fd[1][1],
2022            2e-3,
2023            "first H ll FD",
2024        );
2025
2026        let second = gaussian_joint_psisecondweights(
2027            &rows,
2028            &array![direction_a[0]],
2029            &array![direction_a[1]],
2030            &array![direction_b[0]],
2031            &array![direction_b[1]],
2032            &array![direction_ab[0]],
2033            &array![direction_ab[1]],
2034        );
2035        let objective_second_fd = mixed_value_fd(
2036            y,
2037            base,
2038            weight,
2039            direction_a,
2040            direction_b,
2041            direction_ab,
2042            5e-3,
2043        );
2044        let score_second_fd = mixed_gradient_fd(
2045            y,
2046            base,
2047            weight,
2048            direction_a,
2049            direction_b,
2050            direction_ab,
2051            1e-2,
2052            2e-5,
2053        );
2054        let hessian_second_fd = mixed_hessian_fd(
2055            y,
2056            base,
2057            weight,
2058            direction_a,
2059            direction_b,
2060            direction_ab,
2061            2e-2,
2062            1e-3,
2063        );
2064        assert_close(
2065            second.objective_psi_psirow[0],
2066            objective_second_fd,
2067            5e-4,
2068            "second objective FD",
2069        );
2070        assert_close(
2071            second.d2scoremu[0],
2072            score_second_fd[0],
2073            2e-3,
2074            "second score mu FD",
2075        );
2076        assert_close(
2077            second.d2score_ls[0],
2078            score_second_fd[1],
2079            2e-3,
2080            "second score ls FD",
2081        );
2082        assert_close(
2083            second.d2hmumu[0],
2084            hessian_second_fd[0][0],
2085            3e-2,
2086            "second H mm FD",
2087        );
2088        assert_close(
2089            second.d2hmu_ls[0],
2090            hessian_second_fd[0][1],
2091            3e-2,
2092            "second H ml FD",
2093        );
2094        assert_close(
2095            second.d2h_ls_ls[0],
2096            hessian_second_fd[1][1],
2097            3e-2,
2098            "second H ll FD",
2099        );
2100
2101        let mixed = gaussian_joint_psi_mixed_driftweights(
2102            &rows,
2103            &array![drift[0]],
2104            &array![drift[1]],
2105            &array![psi[0]],
2106            &array![psi[1]],
2107            &array![drift_psi[0]],
2108            &array![drift_psi[1]],
2109        );
2110        let drift_plus = [
2111            base[0] + first_step * drift[0],
2112            base[1] + first_step * drift[1],
2113        ];
2114        let drift_minus = [
2115            base[0] - first_step * drift[0],
2116            base[1] - first_step * drift[1],
2117        ];
2118        let drift_hessian_plus = hessian_fd(y, drift_plus[0], drift_plus[1], weight, 1e-3);
2119        let drift_hessian_minus = hessian_fd(y, drift_minus[0], drift_minus[1], weight, 1e-3);
2120        let drift_hessian_fd: [[f64; 2]; 2] = std::array::from_fn(|row| {
2121            std::array::from_fn(|column| {
2122                (drift_hessian_plus[row][column] - drift_hessian_minus[row][column])
2123                    / (2.0 * first_step)
2124            })
2125        });
2126        let mixed_hessian = mixed_hessian_fd(y, base, weight, drift, psi, drift_psi, 2e-2, 1e-3);
2127        assert_close(
2128            mixed.dhmumu_u[0],
2129            drift_hessian_fd[0][0],
2130            2e-3,
2131            "mixed dH mm FD",
2132        );
2133        assert_close(
2134            mixed.dhmu_ls_u[0],
2135            drift_hessian_fd[0][1],
2136            2e-3,
2137            "mixed dH ml FD",
2138        );
2139        assert_close(
2140            mixed.dh_ls_ls_u[0],
2141            drift_hessian_fd[1][1],
2142            2e-3,
2143            "mixed dH ll FD",
2144        );
2145        assert_close(
2146            mixed.d2hmumu[0],
2147            mixed_hessian[0][0],
2148            3e-2,
2149            "mixed d2H mm FD",
2150        );
2151        assert_close(
2152            mixed.d2hmu_ls[0],
2153            mixed_hessian[0][1],
2154            3e-2,
2155            "mixed d2H ml FD",
2156        );
2157        assert_close(
2158            mixed.d2h_ls_ls[0],
2159            mixed_hessian[1][1],
2160            3e-2,
2161            "mixed d2H ll FD",
2162        );
2163    }
2164
2165    #[test]
2166    fn observed_joint_row_coeffs_match_likelihood_fd_single_source() {
2167        // Residual-dependent cases: y ≠ μ so the cross and (ls,ls) observed
2168        // weights are material (not the Fisher limit m→0, n→a).
2169        let cases = [
2170            (0.3_f64, -0.4_f64, 1.0_f64, 0.55_f64),
2171            (-1.2, 0.7, 2.5, -0.9),
2172            (0.0, 1.5, 0.4, 0.35),
2173            (2.4, -1.1, 0.8, 1.7),
2174            (-0.6, 0.2, 3.3, -1.1),
2175        ];
2176        let h = 1e-4;
2177        for &(mu, eta_ls, a, y) in &cases {
2178            let (mm_hand, ml_hand, ll_hand) = production_observed_row(y, mu, eta_ls, a);
2179            let (mm_fd, ml_fd, ll_fd) = observed_hessian_fd(y, mu, eta_ls, a, h);
2180            assert!(
2181                (mm_hand - mm_fd).abs() <= 1e-5 * mm_hand.abs().max(1.0),
2182                "H_μμ observed μ={mu} η={eta_ls} y={y}: hand={mm_hand} fd={mm_fd}"
2183            );
2184            assert!(
2185                (ml_hand - ml_fd).abs() <= 1e-5 * ml_hand.abs().max(1.0),
2186                "H_μls observed μ={mu} η={eta_ls} y={y}: hand={ml_hand} fd={ml_fd}"
2187            );
2188            assert!(
2189                (ll_hand - ll_fd).abs() <= 1e-5 * ll_hand.abs().max(1.0),
2190                "H_lsls observed μ={mu} η={eta_ls} y={y}: hand={ll_hand} fd={ll_fd}"
2191            );
2192        }
2193    }
2194
2195    #[test]
2196    fn first_directional_weights_match_observed_finite_difference() {
2197        // Pin (w_u, c_u, d_u) to a central FD of the observed coefficients
2198        // along the β-direction (μ += t·ξμ, η_ls += t·ξls).
2199        let cases = [
2200            (0.3_f64, -0.4_f64, 1.0_f64, 0.5_f64, -0.7_f64, 0.55_f64),
2201            (-1.2, 0.7, 2.5, 1.1, 0.3, -0.9),
2202            (0.0, 1.5, 0.4, -0.2, 0.9, 0.35),
2203            (2.4, -1.1, 0.8, 0.6, -0.4, 1.7),
2204        ];
2205        let t = 1e-6;
2206        for &(mu, eta_ls, a, xi_mu, xi_ls, y) in &cases {
2207            let rows =
2208                gaussian_jointrow_scalars(&array![y], &array![mu], &array![eta_ls], &array![a])
2209                    .expect("row scalars");
2210            let (w_u, c_u, d_u) =
2211                gaussian_joint_first_directionalweights(&rows, &array![xi_mu], &array![xi_ls]);
2212            let coeffs_at =
2213                |m: f64, e: f64| -> (f64, f64, f64) { production_observed_row(y, m, e, a) };
2214            let (mmp, mlp, llp) = coeffs_at(mu + t * xi_mu, eta_ls + t * xi_ls);
2215            let (mmm, mlm, llm) = coeffs_at(mu - t * xi_mu, eta_ls - t * xi_ls);
2216            let fd_w = (mmp - mmm) / (2.0 * t);
2217            let fd_c = (mlp - mlm) / (2.0 * t);
2218            let fd_d = (llp - llm) / (2.0 * t);
2219            assert!(
2220                (w_u[0] - fd_w).abs() <= 1e-5 * fd_w.abs().max(1.0),
2221                "dH_μμ drift μ={mu} η={eta_ls}: hand={} fd={fd_w}",
2222                w_u[0]
2223            );
2224            assert!(
2225                (c_u[0] - fd_c).abs() <= 1e-5 * fd_c.abs().max(1.0),
2226                "dH_μls drift μ={mu} η={eta_ls}: hand={} fd={fd_c}",
2227                c_u[0]
2228            );
2229            assert!(
2230                (d_u[0] - fd_d).abs() <= 1e-5 * fd_d.abs().max(1.0),
2231                "dH_lsls drift μ={mu} η={eta_ls}: hand={} fd={fd_d}",
2232                d_u[0]
2233            );
2234        }
2235    }
2236
2237    #[test]
2238    fn second_directional_weights_match_first_directional_finite_difference() {
2239        // Pin (w_uv, c_uv, d_uv) to a central FD of the FIRST-directional
2240        // observed drifts along the v-direction.
2241        let cases = [
2242            (
2243                0.3_f64, -0.4_f64, 1.0_f64, 0.5_f64, -0.7_f64, 0.8_f64, 0.2_f64, 0.55_f64,
2244            ),
2245            (-1.2, 0.7, 2.5, 1.1, 0.3, -0.6, 0.9, -0.9),
2246            (0.0, 1.5, 0.4, -0.2, 0.9, 0.4, -0.5, 0.35),
2247        ];
2248        let t = 1e-5;
2249        for &(mu, eta_ls, a, xi_mu_u, xi_ls_u, xi_mu_v, xi_ls_v, y) in &cases {
2250            let rows =
2251                gaussian_jointrow_scalars(&array![y], &array![mu], &array![eta_ls], &array![a])
2252                    .expect("row scalars");
2253            let (w_uv, c_uv, d_uv) = gaussian_jointsecond_directionalweights(
2254                &rows,
2255                &array![xi_mu_u],
2256                &array![xi_ls_u],
2257                &array![xi_mu_v],
2258                &array![xi_ls_v],
2259            );
2260            let first_at = |m: f64, e: f64| -> (f64, f64, f64) {
2261                let r = gaussian_jointrow_scalars(&array![y], &array![m], &array![e], &array![a])
2262                    .expect("row scalars");
2263                let (w_u, c_u, d_u) =
2264                    gaussian_joint_first_directionalweights(&r, &array![xi_mu_u], &array![xi_ls_u]);
2265                (w_u[0], c_u[0], d_u[0])
2266            };
2267            let (wp, cp, dp) = first_at(mu + t * xi_mu_v, eta_ls + t * xi_ls_v);
2268            let (wm, cm, dm) = first_at(mu - t * xi_mu_v, eta_ls - t * xi_ls_v);
2269            let fd_w = (wp - wm) / (2.0 * t);
2270            let fd_c = (cp - cm) / (2.0 * t);
2271            let fd_d = (dp - dm) / (2.0 * t);
2272            assert!(
2273                (w_uv[0] - fd_w).abs() <= 1e-4 * fd_w.abs().max(1.0),
2274                "d²H_μμ drift μ={mu} η={eta_ls}: hand={} fd={fd_w}",
2275                w_uv[0]
2276            );
2277            assert!(
2278                (c_uv[0] - fd_c).abs() <= 1e-4 * fd_c.abs().max(1.0),
2279                "d²H_μls drift μ={mu} η={eta_ls}: hand={} fd={fd_c}",
2280                c_uv[0]
2281            );
2282            assert!(
2283                (d_uv[0] - fd_d).abs() <= 1e-4 * fd_d.abs().max(1.0),
2284                "d²H_lsls drift μ={mu} η={eta_ls}: hand={} fd={fd_d}",
2285                d_uv[0]
2286            );
2287        }
2288    }
2289
2290    /// #932 jet oracle for the LIVE gaulss third/fourth directional Hessian
2291    /// drifts. The FD tests above are the independent numerical witness; these
2292    /// two pin the SAME live hand closed forms
2293    /// (`gaussian_joint_first_directionalweights` /
2294    /// `gaussian_jointsecond_directionalweights`) against the universal gam-math
2295    /// Taylor jet — the mechanical single source — at ≤1e-9, closing the audit
2296    /// gap that the gam-math gaulss oracle covered only value/∇/observed-H.
2297    mod jet_third_fourth_oracle {
2298        use super::*;
2299        use gam_math::jet_tower::{program_fourth_contracted, program_third_contracted};
2300
2301        fn close(hand: f64, jet: f64, label: &str) {
2302            let band = 1e-9 + 1e-9 * hand.abs().max(jet.abs());
2303            assert!(
2304                (hand - jet).abs() <= band,
2305                "{label}: hand {hand:+.15e} vs jet {jet:+.15e} (band {band:.3e})"
2306            );
2307        }
2308
2309        /// `gaussian_joint_first_directionalweights` (the LIVE third-order ∂_dir
2310        /// of the observed Hessian) equals the jet's contracted third at ≤1e-9.
2311        #[test]
2312        fn first_directional_weights_match_jet_third() {
2313            let cases = [
2314                (0.3_f64, -0.4_f64, 1.0_f64, 0.5_f64, -0.7_f64, 0.55_f64),
2315                (-1.2, 0.7, 2.5, 1.1, 0.3, -0.9),
2316                (0.0, 1.5, 0.4, -0.2, 0.9, 0.35),
2317                (2.4, -1.1, 0.8, 0.6, -0.4, 1.7),
2318                (-0.6, 0.2, 3.3, 0.8, -1.0, -1.1),
2319            ];
2320            for &(mu, eta_ls, a, xi_mu, xi_ls, y) in &cases {
2321                let rows =
2322                    gaussian_jointrow_scalars(&array![y], &array![mu], &array![eta_ls], &array![a])
2323                        .expect("row scalars");
2324                let (w_u, c_u, d_u) =
2325                    gaussian_joint_first_directionalweights(&rows, &array![xi_mu], &array![xi_ls]);
2326                let prog = crate::gamlss::GaussianJointRowProgram::new(&rows);
2327                let jt = program_third_contracted(&prog, 0, &[xi_mu, xi_ls]).expect("jet third");
2328                close(w_u[0], jt[0][0], &format!("dH_μμ μ={mu} η={eta_ls}"));
2329                close(c_u[0], jt[0][1], &format!("dH_μls μ={mu} η={eta_ls}"));
2330                close(c_u[0], jt[1][0], &format!("dH_lsμ μ={mu} η={eta_ls}"));
2331                close(d_u[0], jt[1][1], &format!("dH_lsls μ={mu} η={eta_ls}"));
2332            }
2333        }
2334
2335        /// `gaussian_jointsecond_directionalweights` (the LIVE fourth-order
2336        /// ∂_u∂_v of the observed Hessian) equals the jet's contracted fourth at
2337        /// ≤1e-9.
2338        #[test]
2339        fn second_directional_weights_match_jet_fourth() {
2340            let cases = [
2341                (
2342                    0.3_f64, -0.4_f64, 1.0_f64, 0.5_f64, -0.7_f64, 0.8_f64, 0.2_f64, 0.55_f64,
2343                ),
2344                (-1.2, 0.7, 2.5, 1.1, 0.3, -0.6, 0.9, -0.9),
2345                (0.0, 1.5, 0.4, -0.2, 0.9, 0.4, -0.5, 0.35),
2346                (2.4, -1.1, 0.8, 0.6, -0.4, -0.3, 0.7, 1.7),
2347            ];
2348            for &(mu, eta_ls, a, xi_mu_u, xi_ls_u, xi_mu_v, xi_ls_v, y) in &cases {
2349                let rows =
2350                    gaussian_jointrow_scalars(&array![y], &array![mu], &array![eta_ls], &array![a])
2351                        .expect("row scalars");
2352                let (w_uv, c_uv, d_uv) = gaussian_jointsecond_directionalweights(
2353                    &rows,
2354                    &array![xi_mu_u],
2355                    &array![xi_ls_u],
2356                    &array![xi_mu_v],
2357                    &array![xi_ls_v],
2358                );
2359                let prog = crate::gamlss::GaussianJointRowProgram::new(&rows);
2360                let jt =
2361                    program_fourth_contracted(&prog, 0, &[xi_mu_u, xi_ls_u], &[xi_mu_v, xi_ls_v])
2362                        .expect("jet fourth");
2363                close(w_uv[0], jt[0][0], &format!("d²H_μμ μ={mu} η={eta_ls}"));
2364                close(c_uv[0], jt[0][1], &format!("d²H_μls μ={mu} η={eta_ls}"));
2365                close(c_uv[0], jt[1][0], &format!("d²H_lsμ μ={mu} η={eta_ls}"));
2366                close(d_uv[0], jt[1][1], &format!("d²H_lsls μ={mu} η={eta_ls}"));
2367            }
2368        }
2369
2370        /// #932 release speed gate for the Gaussian location-scale joint row.
2371        /// The production structure-compiled order-2/third/fourth lowerings
2372        /// ([`GaussianJointRowProgram::row_order2`] /
2373        /// `row_third_contracted` / `row_fourth_contracted`, all emitted from
2374        /// the one [`gaussian_normalized_row`] declaration) are timed against
2375        /// the generic gam-math forward-mode jet tower ([`program_row_kernel`]
2376        /// / [`program_third_contracted`] / [`program_fourth_contracted`]) —
2377        /// the naive automatic-differentiation baseline the retained
2378        /// specialization must beat; #932 keeps no separate `cfg(test)` hand
2379        /// restatement for this family. Emits one harness-parsed
2380        /// `hand_over_production` token (generic-tower time over production
2381        /// time) per derivative channel; the MSI release harness fails closed
2382        /// whenever any measured cell is `<= 1`.
2383        ///
2384        /// The batch carries distinct certified per-row scalars and distinct
2385        /// per-row directions, so the optimizer cannot hoist the pure row
2386        /// calls out of the sweeps, and the finite checksum over every
2387        /// returned channel keeps each sweep live without
2388        /// `std::hint::black_box`.
2389        #[test]
2390        fn release_measure_gaussian_joint_specialized_vs_generic_tower_932() {
2391            use gam_math::jet_tower::program_row_kernel;
2392            use std::time::Instant;
2393
2394            const ROWS: usize = 512;
2395            let y = Array1::from_shape_fn(ROWS, |i| {
2396                let f = i as f64;
2397                1.4 * (f * 0.17 + 0.3).sin() - 0.5 * (f * 0.09).cos()
2398            });
2399            let mu = Array1::from_shape_fn(ROWS, |i| {
2400                let f = i as f64;
2401                0.9 * (f * 0.13 + 0.7).cos() + 0.3 * (f * 0.05).sin()
2402            });
2403            let eta_ls = Array1::from_shape_fn(ROWS, |i| {
2404                let f = i as f64;
2405                0.8 * (f * 0.11 + 0.2).sin() - 0.25 * (f * 0.07).cos()
2406            });
2407            let weight = Array1::from_shape_fn(ROWS, |i| {
2408                let f = i as f64;
2409                0.6 + 0.4 * (f * 0.19 + 1.0).sin().abs()
2410            });
2411            let rows = gaussian_jointrow_scalars(&y, &mu, &eta_ls, &weight)
2412                .expect("certified release-measure Gaussian joint rows");
2413            let program = crate::gamlss::GaussianJointRowProgram::new(&rows);
2414            let dir_u: Vec<[f64; 2]> = (0..ROWS)
2415                .map(|i| {
2416                    let f = i as f64;
2417                    [
2418                        0.7 * (f * 0.23 + 0.4).cos() - 0.2 * (f * 0.03).sin(),
2419                        -0.6 * (f * 0.29 + 0.1).sin() + 0.25 * (f * 0.15).cos(),
2420                    ]
2421                })
2422                .collect();
2423            let dir_v: Vec<[f64; 2]> = (0..ROWS)
2424                .map(|i| {
2425                    let f = i as f64;
2426                    [
2427                        -0.5 * (f * 0.21 + 0.9).sin() + 0.3 * (f * 0.06).cos(),
2428                        0.8 * (f * 0.27 + 0.5).cos() - 0.15 * (f * 0.04).sin(),
2429                    ]
2430                })
2431                .collect();
2432
2433            // Warm both paths and pin equal V/G/H plus equal contracted
2434            // third/fourth, so each timed pair measures equal work.
2435            for row in 0..ROWS {
2436                let atom = program.row_order2(row);
2437                let (tower_value, tower_gradient, tower_hessian) =
2438                    program_row_kernel(&program, row).expect("tower warm kernel");
2439                close(atom.value(), tower_value, "release-measure value parity");
2440                let production_gradient = atom.gradient();
2441                for a in 0..2 {
2442                    close(
2443                        production_gradient[a],
2444                        tower_gradient[a],
2445                        "release-measure gradient parity",
2446                    );
2447                    for b in 0..2 {
2448                        close(
2449                            atom.hessian_at(a, b),
2450                            tower_hessian[a][b],
2451                            "release-measure hessian parity",
2452                        );
2453                    }
2454                }
2455                let production_third = program.row_third_contracted(row, &dir_u[row]);
2456                let tower_third = program_third_contracted(&program, row, &dir_u[row])
2457                    .expect("tower warm third");
2458                let production_fourth =
2459                    program.row_fourth_contracted(row, &dir_u[row], &dir_v[row]);
2460                let tower_fourth =
2461                    program_fourth_contracted(&program, row, &dir_u[row], &dir_v[row])
2462                        .expect("tower warm fourth");
2463                for a in 0..2 {
2464                    for b in 0..2 {
2465                        close(
2466                            production_third[a][b],
2467                            tower_third[a][b],
2468                            "release-measure third parity",
2469                        );
2470                        close(
2471                            production_fourth[a][b],
2472                            tower_fourth[a][b],
2473                            "release-measure fourth parity",
2474                        );
2475                    }
2476                }
2477            }
2478
2479            let best_secs = |sweep: &mut dyn FnMut() -> f64| -> f64 {
2480                let mut best = f64::INFINITY;
2481                for _ in 0..5 {
2482                    let started = Instant::now();
2483                    let checksum = sweep();
2484                    assert!(
2485                        checksum.is_finite(),
2486                        "Gaussian joint release-measure checksum must stay finite"
2487                    );
2488                    best = best.min(started.elapsed().as_secs_f64());
2489                }
2490                best
2491            };
2492
2493            let mut production_order2_sweep = || {
2494                let mut checksum = 0.0_f64;
2495                for row in 0..ROWS {
2496                    let atom = program.row_order2(row);
2497                    checksum += atom.value() + atom.gradient()[0] + atom.hessian_at(0, 0);
2498                }
2499                checksum
2500            };
2501            let production_order2_secs = best_secs(&mut production_order2_sweep);
2502            let mut tower_order2_sweep = || {
2503                let mut checksum = 0.0_f64;
2504                for row in 0..ROWS {
2505                    let (value, gradient, hessian) =
2506                        program_row_kernel(&program, row).expect("tower kernel");
2507                    checksum += value + gradient[0] + hessian[0][0];
2508                }
2509                checksum
2510            };
2511            let tower_order2_secs = best_secs(&mut tower_order2_sweep);
2512
2513            let mut production_third_sweep = || {
2514                let mut checksum = 0.0_f64;
2515                for row in 0..ROWS {
2516                    let third = program.row_third_contracted(row, &dir_u[row]);
2517                    checksum += third[0][0] + third[0][1] + third[1][1];
2518                }
2519                checksum
2520            };
2521            let production_third_secs = best_secs(&mut production_third_sweep);
2522            let mut tower_third_sweep = || {
2523                let mut checksum = 0.0_f64;
2524                for row in 0..ROWS {
2525                    let third = program_third_contracted(&program, row, &dir_u[row])
2526                        .expect("tower third");
2527                    checksum += third[0][0] + third[0][1] + third[1][1];
2528                }
2529                checksum
2530            };
2531            let tower_third_secs = best_secs(&mut tower_third_sweep);
2532
2533            let mut production_fourth_sweep = || {
2534                let mut checksum = 0.0_f64;
2535                for row in 0..ROWS {
2536                    let fourth = program.row_fourth_contracted(row, &dir_u[row], &dir_v[row]);
2537                    checksum += fourth[0][0] + fourth[0][1] + fourth[1][1];
2538                }
2539                checksum
2540            };
2541            let production_fourth_secs = best_secs(&mut production_fourth_sweep);
2542            let mut tower_fourth_sweep = || {
2543                let mut checksum = 0.0_f64;
2544                for row in 0..ROWS {
2545                    let fourth =
2546                        program_fourth_contracted(&program, row, &dir_u[row], &dir_v[row])
2547                            .expect("tower fourth");
2548                    checksum += fourth[0][0] + fourth[0][1] + fourth[1][1];
2549                }
2550                checksum
2551            };
2552            let tower_fourth_secs = best_secs(&mut tower_fourth_sweep);
2553
2554            for (channel, production_secs, tower_secs) in [
2555                ("order2", production_order2_secs, tower_order2_secs),
2556                ("third", production_third_secs, tower_third_secs),
2557                ("fourth", production_fourth_secs, tower_fourth_secs),
2558            ] {
2559                let production_ns = production_secs * 1e9 / ROWS as f64;
2560                let tower_ns = tower_secs * 1e9 / ROWS as f64;
2561                eprintln!(
2562                    "GAUSSIAN-JOINT-RELEASE-932 channel={channel} rows={ROWS} \
2563                     production_ns={production_ns:.3} generic_tower_ns={tower_ns:.3} \
2564                     hand_over_production={:.6}",
2565                    tower_ns / production_ns,
2566                );
2567            }
2568        }
2569    }
2570}