1use super::*;
6use gam_row_macros::row_atom;
7
8row_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
53pub(crate) trait LocationScaleJointPsiFamily: Clone + Send + Sync + 'static {
68 type Direction: Send + Sync + 'static;
70
71 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 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 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
285pub(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 pub(crate) standardized_residual: Array1<f64>,
433 pub(crate) inv_sigma: Array1<f64>,
435 pub(crate) kappa: Array1<f64>,
439}
440
441pub struct GaussianJointRowProgram<'a> {
450 rows: &'a GaussianJointRowScalars,
451}
452
453impl<'a> GaussianJointRowProgram<'a> {
454 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 #[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 #[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 #[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
549pub(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
615pub(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
627pub(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
652pub(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
769pub(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
822pub(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
856pub(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 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 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
939pub(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
960pub(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 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 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 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
1152pub(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 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 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 use super::*;
1511 use ndarray::array;
1512
1513 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 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 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 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 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 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 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 #[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 #[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 #[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 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}