1use self::inner_strategy::GeometryBackendKind;
2use super::*;
3use crate::pirls::PIRLS_CACHE_BYTE_BUDGET;
4use crate::pirls::assemble_and_factor_sparse_penalized_system;
5use gam_linalg::sparse_exact::SparseExactFactor;
6use gam_problem::OuterEval;
7use gam_problem::SasLinkState;
8use gam_terms::basis::LocalDesignJacobianProvider;
9use ndarray::{Array1, Array2, s};
10use std::collections::{HashMap, VecDeque};
11use std::ops::Range;
12use std::sync::atomic::{AtomicBool, AtomicU64, AtomicUsize};
13use std::sync::{Arc, RwLock};
14
15pub mod assembly;
16pub mod atoms;
17pub mod boundary_laml;
18pub(crate) mod continuation;
19pub(crate) mod eval;
20mod firth;
21pub(super) mod hyper;
22mod inner_strategy;
23pub mod jeffreys_subspace;
26pub(crate) mod laml_logdet;
27pub mod outer_eval;
28pub mod penalty_logdet;
29pub mod per_atom_efs;
30pub mod reml_outer_engine;
31pub mod reparameterized_inner;
32mod rho_key;
33mod sparse_exact_penalty;
34mod trace;
35
36pub(crate) use sparse_exact_penalty::sparse_penalty_block_count_from_canonical;
37
38pub(crate) const EXACT_TAU_TAU_HESSIAN_DENSE_CACHE_BUDGET_BYTES: usize = 512 * 1024 * 1024;
39pub(crate) const FIRTH_MAX_OBSERVATIONS: usize = 20_000;
40pub(crate) const FIRTH_MAX_COEFFICIENTS: usize = 256;
41pub(crate) const FIRTH_MAX_LINEAR_WORK: usize = 2_000_000;
42pub(crate) const FIRTH_MAX_QUADRATIC_WORK: usize = 100_000_000;
43pub(crate) const PERSISTENT_LATENT_VALUES_CACHE_CAPACITY: usize = 8;
44
45#[derive(Debug)]
46pub(crate) struct PersistentLatentValuesCache {
47 pub(crate) entries: HashMap<String, Array2<f64>>,
48 pub(crate) lru: VecDeque<String>,
49 pub(crate) capacity: usize,
50}
51
52impl Default for PersistentLatentValuesCache {
53 fn default() -> Self {
54 Self {
55 entries: HashMap::new(),
56 lru: VecDeque::new(),
57 capacity: PERSISTENT_LATENT_VALUES_CACHE_CAPACITY,
58 }
59 }
60}
61
62impl PersistentLatentValuesCache {
63 pub(crate) fn lookup(
64 &mut self,
65 key: &str,
66 n_obs: usize,
67 latent_dim: usize,
68 ) -> Option<Array2<f64>> {
69 let values = self.entries.get(key)?;
70 if values.dim() != (n_obs, latent_dim) {
71 return None;
72 }
73 let values = values.clone();
74 self.touch(key.to_string());
75 Some(values)
76 }
77
78 pub(crate) fn insert(&mut self, key: String, values: Array2<f64>) {
79 if values.iter().any(|value| !value.is_finite()) {
80 return;
81 }
82 self.entries.insert(key.clone(), values);
83 self.touch(key);
84 while self.entries.len() > self.capacity {
85 let Some(evicted) = self.lru.pop_front() else {
86 break;
87 };
88 self.entries.remove(&evicted);
89 }
90 }
91
92 pub(crate) fn touch(&mut self, key: String) {
93 if let Some(index) = self.lru.iter().position(|queued| queued == &key) {
94 self.lru.remove(index);
95 }
96 self.lru.push_back(key);
97 }
98}
99
100#[derive(Clone)]
105pub(crate) struct IftWarmStartCache {
106 pub beta_original: ndarray::Array1<f64>,
112 pub rho: ndarray::Array1<f64>,
115 pub penalized_hessian_transformed: gam_linalg::matrix::SymmetricMatrix,
119 pub qs: ndarray::Array2<f64>,
123 pub frame_was_original: bool,
127 pub lambda_s_beta_blocks: Option<Vec<ndarray::Array1<f64>>>,
144}
145
146#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
147pub(crate) struct TauTauPlanEstimate {
148 pub(crate) dense_x_bytes: usize,
149 pub(crate) first_order_tau_bytes: usize,
150 pub(crate) second_order_tau_bytes: usize,
151 pub(crate) penalty_first_bytes: usize,
152 pub(crate) penalty_pair_bytes: usize,
153 pub(crate) rho_tau_penalty_bytes: usize,
154 pub(crate) vector_cache_bytes: usize,
155 pub(crate) weighted_scratch_bytes: usize,
156}
157
158impl TauTauPlanEstimate {
159 pub(crate) fn total_bytes(self) -> usize {
160 self.dense_x_bytes
161 .saturating_add(self.first_order_tau_bytes)
162 .saturating_add(self.second_order_tau_bytes)
163 .saturating_add(self.penalty_first_bytes)
164 .saturating_add(self.penalty_pair_bytes)
165 .saturating_add(self.rho_tau_penalty_bytes)
166 .saturating_add(self.vector_cache_bytes)
167 .saturating_add(self.weighted_scratch_bytes)
168 }
169}
170
171#[derive(Clone, Copy, Debug, PartialEq, Eq)]
172pub(crate) struct TauTauHessianPolicy {
173 pub(crate) any_has_implicit: bool,
174 pub(crate) implicit_multidim_duchon: bool,
175 pub(crate) estimated_dense_tau_cache_bytes: usize,
176 pub(crate) gradient_plan: TauTauPlanEstimate,
177 pub(crate) hessian_plan: TauTauPlanEstimate,
178 pub(crate) budget_bytes: usize,
179 pub(crate) firth_pair_terms_unavailable: bool,
180}
181
182impl TauTauHessianPolicy {
183 pub(crate) fn prefer_gradient_only(self) -> bool {
211 self.firth_pair_terms_unavailable
212 }
213}
214
215pub(crate) fn exact_tau_tau_hessian_policy_with_firth(
216 n_obs: usize,
217 p_coeff: usize,
218 hyper_dirs: &[DirectionalHyperParam],
219 firth_pair_terms_unavailable: bool,
220) -> TauTauHessianPolicy {
221 let f64_bytes = std::mem::size_of::<f64>();
222 let dense_matrix_bytes =
223 |rows: usize, cols: usize| -> usize { rows.saturating_mul(cols).saturating_mul(f64_bytes) };
224 let dense_design_bytes = dense_matrix_bytes(n_obs, p_coeff);
225 let dense_penalty_bytes = dense_matrix_bytes(p_coeff, p_coeff);
226 let psi_dim = hyper_dirs.len();
227 let implicit_n_axes = hyper_dirs
228 .iter()
229 .find_map(DirectionalHyperParam::implicit_axis_count_hint)
230 .unwrap_or(0);
231 let gradient_uses_implicit_design = hyper_dirs
232 .iter()
233 .any(DirectionalHyperParam::has_implicit_operator)
234 && gam_terms::basis::should_use_implicit_operators_with_policy(
235 n_obs,
236 p_coeff,
237 implicit_n_axes,
238 &gam_runtime::resource::ResourcePolicy::default_library(),
239 );
240 let dense_first_order_count = hyper_dirs
241 .iter()
242 .filter(|dir| !dir.has_implicit_operator())
243 .count();
244 let first_penalty_component_count = hyper_dirs
245 .iter()
246 .map(DirectionalHyperParam::penalty_first_component_count)
247 .sum::<usize>();
248
249 let mut dense_second_order_count = 0usize;
250 let mut penalty_pair_count = 0usize;
251 for i in 0..psi_dim {
252 for j in i..psi_dim {
253 if hyper_dirs[i]
254 .x_tau_tau_entry_at(j)
255 .or_else(|| hyper_dirs[j].x_tau_tau_entry_at(i))
256 .is_some_and(|entry| !entry.uses_implicit_storage())
257 {
258 dense_second_order_count += if i == j { 1 } else { 2 };
259 }
260 if hyper_dirs[i].has_penaltysecond_pair_at(j)
261 || hyper_dirs[j].has_penaltysecond_pair_at(i)
262 {
263 penalty_pair_count += if i == j { 1 } else { 2 };
264 }
265 }
266 }
267
268 let gradient_dense_first_order_count = if gradient_uses_implicit_design {
269 dense_first_order_count
270 } else {
271 psi_dim
272 };
273 let gradient_needs_dense_x =
274 firth_pair_terms_unavailable || gradient_dense_first_order_count > 0;
275 let gradient_plan = TauTauPlanEstimate {
276 dense_x_bytes: if gradient_needs_dense_x {
277 dense_design_bytes
278 } else {
279 0
280 },
281 first_order_tau_bytes: if gradient_dense_first_order_count > 0 {
282 dense_design_bytes
283 } else {
284 0
285 },
286 second_order_tau_bytes: 0,
287 penalty_first_bytes: psi_dim.saturating_mul(dense_penalty_bytes),
288 penalty_pair_bytes: 0,
289 rho_tau_penalty_bytes: 0,
290 vector_cache_bytes: n_obs.saturating_mul(f64_bytes),
291 weighted_scratch_bytes: dense_penalty_bytes,
292 };
293 let hessian_plan = TauTauPlanEstimate {
294 dense_x_bytes: if psi_dim > 0 { dense_design_bytes } else { 0 },
295 first_order_tau_bytes: dense_first_order_count.saturating_mul(dense_design_bytes),
296 second_order_tau_bytes: dense_second_order_count.saturating_mul(dense_design_bytes),
297 penalty_first_bytes: psi_dim.saturating_mul(dense_penalty_bytes),
298 penalty_pair_bytes: penalty_pair_count.saturating_mul(dense_penalty_bytes),
299 rho_tau_penalty_bytes: first_penalty_component_count
300 .saturating_mul(2)
301 .saturating_mul(dense_penalty_bytes),
302 vector_cache_bytes: psi_dim.saturating_mul(n_obs).saturating_mul(f64_bytes),
303 weighted_scratch_bytes: dense_penalty_bytes,
304 };
305 let any_has_implicit = hyper_dirs
306 .iter()
307 .any(DirectionalHyperParam::has_implicit_operator);
308 let implicit_multidim_duchon = hyper_dirs
309 .iter()
310 .any(DirectionalHyperParam::has_implicit_multidim_duchon);
311 let estimated_dense_tau_cache_bytes = hessian_plan
312 .first_order_tau_bytes
313 .saturating_add(hessian_plan.second_order_tau_bytes);
314 TauTauHessianPolicy {
315 any_has_implicit,
316 implicit_multidim_duchon,
317 estimated_dense_tau_cache_bytes,
318 gradient_plan,
319 hessian_plan,
320 budget_bytes: EXACT_TAU_TAU_HESSIAN_DENSE_CACHE_BUDGET_BYTES,
321 firth_pair_terms_unavailable: firth_pair_terms_unavailable && !hyper_dirs.is_empty(),
322 }
323}
324
325pub(crate) fn firth_problem_scale_allows(n_obs: usize, p_coeff: usize) -> bool {
326 let linear_work = n_obs.saturating_mul(p_coeff);
327 let quadratic_work = linear_work.saturating_mul(p_coeff);
328 n_obs <= FIRTH_MAX_OBSERVATIONS
329 && p_coeff <= FIRTH_MAX_COEFFICIENTS
330 && linear_work <= FIRTH_MAX_LINEAR_WORK
331 && quadratic_work <= FIRTH_MAX_QUADRATIC_WORK
332}
333
334#[cfg(test)]
335mod tests {
336 use super::atoms::CriterionAtom;
337 use super::{
338 DirectionalHyperParam, EvalCacheManager, EvalShared, HyperDesignDerivative,
339 HyperPenaltyDerivative, ImplicitDerivLevel, RemlConfig, RemlState,
340 };
341 use crate::estimate::EstimationError;
342 use crate::pirls::PirlsCoordinateFrame;
343 use faer::Side;
344 use gam_linalg::faer_ndarray::FaerCholesky;
345 use gam_linalg::matrix::symmetrize_in_place;
346 use gam_problem::{
347 GlmLikelihoodSpec, InverseLink, LikelihoodSpec, ResponseFamily, StandardLink,
348 };
349 use gam_problem::{HessianValue, OuterEval};
350 use gam_terms::basis::{ImplicitDesignPsiDerivative, RadialScalarKind};
351 use ndarray::{Array1, Array2, array, s};
352 use std::sync::Arc;
353
354 pub(crate) fn binomial_logit_glm_spec() -> GlmLikelihoodSpec {
357 GlmLikelihoodSpec::canonical(LikelihoodSpec::new(
358 ResponseFamily::Binomial,
359 InverseLink::Standard(StandardLink::Logit),
360 ))
361 }
362
363 pub(crate) fn gaussian_identity_glm_spec() -> GlmLikelihoodSpec {
366 GlmLikelihoodSpec::canonical(LikelihoodSpec::new(
367 ResponseFamily::Gaussian,
368 InverseLink::Standard(StandardLink::Identity),
369 ))
370 }
371
372 impl DirectionalHyperParam {
373 pub(super) fn new(
374 x_tau_original: Array2<f64>,
375 penalty_first_components: Vec<(usize, Array2<f64>)>,
376 x_tau_tau_original: Option<Vec<Option<Array2<f64>>>>,
377 penaltysecond_components: Option<Vec<Option<Vec<(usize, Array2<f64>)>>>>,
378 ) -> Result<Self, EstimationError> {
379 let x_tau_tau_original = x_tau_tau_original.map(|rows| {
380 rows.into_iter()
381 .map(|entry| entry.map(HyperDesignDerivative::from))
382 .collect::<Vec<_>>()
383 });
384 let penalty_first_components = penalty_first_components
385 .into_iter()
386 .map(|(idx, matrix)| (idx, HyperPenaltyDerivative::from(matrix)))
387 .collect();
388 let penaltysecond_components = penaltysecond_components.map(|rows| {
389 rows.into_iter()
390 .map(|row| {
391 row.map(|components| {
392 components
393 .into_iter()
394 .map(|(idx, matrix)| (idx, HyperPenaltyDerivative::from(matrix)))
395 .collect::<Vec<_>>()
396 })
397 })
398 .collect::<Vec<_>>()
399 });
400 Self::new_compact(
401 HyperDesignDerivative::from(x_tau_original),
402 penalty_first_components,
403 x_tau_tau_original,
404 penaltysecond_components,
405 )
406 }
407
408 pub(super) fn single_penalty(
409 penalty_index: usize,
410 x_tau_original: Array2<f64>,
411 s_tau_original: Array2<f64>,
412 x_tau_tau_original: Option<Vec<Option<Array2<f64>>>>,
413 s_tau_tau_original: Option<Vec<Option<Array2<f64>>>>,
414 ) -> Result<Self, EstimationError> {
415 let penaltysecond_components = s_tau_tau_original.map(|rows| {
416 rows.into_iter()
417 .map(|mat| mat.map(|mat| vec![(penalty_index, mat)]))
418 .collect::<Vec<_>>()
419 });
420 Self::new(
421 x_tau_original,
422 vec![(penalty_index, s_tau_original)],
423 x_tau_tau_original,
424 penaltysecond_components,
425 )
426 }
427 }
428
429 #[test]
430 pub(crate) fn firth_problem_scale_gate_blocks_large_quadratic_work() {
431 assert!(super::firth_problem_scale_allows(2_000, 200));
432 assert!(!super::firth_problem_scale_allows(4_800, 241));
433 assert!(!super::firth_problem_scale_allows(4_800, 433));
434 }
435
436 #[test]
437 pub(crate) fn tau_tau_hessian_policy_prefers_gradient_only_for_implicit_tau() {
438 let operator = ImplicitDesignPsiDerivative::new(
439 array![1.0, 2.0, 3.0, 4.0],
440 array![0.5, -1.0, 1.5, 2.0],
441 array![0.1, 0.2, 0.3, 0.4],
442 array![[1.0, 0.2], [0.5, 0.1], [1.5, 0.3], [2.0, 0.4]],
443 None,
444 None,
445 2,
446 2,
447 1,
448 2,
449 );
450 let dir = DirectionalHyperParam::new_compact(
451 HyperDesignDerivative::from_implicit(
452 Arc::new(operator),
453 ImplicitDerivLevel::First(0),
454 1..4,
455 5,
456 ),
457 Vec::new(),
458 None,
459 None,
460 )
461 .expect("implicit directional hyperparam");
462 let policy = super::exact_tau_tau_hessian_policy_with_firth(10, 5, &[dir], false);
463 assert!(policy.any_has_implicit);
464 assert_eq!(
465 policy.gradient_plan.dense_x_bytes,
466 10 * 5 * std::mem::size_of::<f64>()
467 );
468 assert!(!policy.prefer_gradient_only());
469 }
470
471 #[test]
472 pub(crate) fn tau_tau_hessian_policy_does_not_force_gradient_only_for_implicit_multidim_duchon()
473 {
474 let operator = ImplicitDesignPsiDerivative::new_streaming(
482 Arc::new(array![[0.0, 0.0], [1.0, 0.2]]),
483 Arc::new(array![[0.0, 0.0], [1.0, 1.0]]),
484 vec![0.0, 0.0],
485 RadialScalarKind::PureDuchon {
486 block_order: 1,
487 p_order: 0,
488 s_order: 0,
489 dim: 2,
490 },
491 None,
492 None,
493 0,
494 );
495 let dir = DirectionalHyperParam::new_compact(
496 HyperDesignDerivative::from_implicit(
497 Arc::new(operator),
498 ImplicitDerivLevel::First(0),
499 0..2,
500 2,
501 ),
502 Vec::new(),
503 None,
504 None,
505 )
506 .expect("implicit duchon directional hyperparam");
507 let policy = super::exact_tau_tau_hessian_policy_with_firth(10, 5, &[dir], false);
508 assert!(policy.any_has_implicit);
509 assert!(policy.implicit_multidim_duchon);
510 assert!(!policy.prefer_gradient_only());
511 }
512
513 #[test]
514 pub(crate) fn tau_tau_hessian_policy_does_not_force_gradient_only_when_cache_budget_is_exceeded()
515 {
516 let dirs = (0..16)
522 .map(|_| {
523 DirectionalHyperParam::new_compact(
524 HyperDesignDerivative::from(Array2::<f64>::zeros((2, 2))),
525 Vec::new(),
526 None,
527 None,
528 )
529 .expect("dense directional hyperparam")
530 })
531 .collect::<Vec<_>>();
532 let policy = super::exact_tau_tau_hessian_policy_with_firth(320_000, 71, &dirs, false);
533 assert!(!policy.any_has_implicit);
534 assert!(policy.hessian_plan.total_bytes() > policy.budget_bytes);
535 assert!(policy.hessian_plan.total_bytes() > policy.gradient_plan.total_bytes());
536 assert!(!policy.prefer_gradient_only());
537 }
538
539 #[test]
540 pub(crate) fn tau_tau_hessian_policy_prefers_gradient_only_for_firth_pair_gap() {
541 let dir = DirectionalHyperParam::new_compact(
542 HyperDesignDerivative::from(Array2::<f64>::zeros((2, 2))),
543 Vec::new(),
544 None,
545 None,
546 )
547 .expect("dense directional hyperparam");
548 let policy = super::exact_tau_tau_hessian_policy_with_firth(10, 5, &[dir], true);
549 assert!(policy.firth_pair_terms_unavailable);
550 assert!(policy.prefer_gradient_only());
551 }
552
553 trait LogitDesignMotionFixture {
561 fn y(&self) -> &Array1<f64>;
562 fn w(&self) -> &Array1<f64>;
563 fn x(&self) -> &Array2<f64>;
564 fn s0(&self) -> &Array2<f64>;
565 fn cfg(&self) -> &RemlConfig;
566 fn rho(&self) -> &Array1<f64>;
567
568 fn state(&self) -> RemlState<'_> {
569 build_logit_state(self.y(), self.w(), self.x(), self.s0(), self.cfg())
570 }
571
572 fn state_perturbed(
573 &self,
574 x_tau: &Array2<f64>,
575 s_tau: &Array2<f64>,
576 eps: f64,
577 ) -> (RemlState<'_>, RemlState<'_>) {
578 let x_plus = self.x() + &x_tau.mapv(|v| eps * v);
579 let x_minus = self.x() - &x_tau.mapv(|v| eps * v);
580 let s_plus = self.s0() + &s_tau.mapv(|v| eps * v);
581 let s_minus = self.s0() - &s_tau.mapv(|v| eps * v);
582 (
583 build_logit_state(self.y(), self.w(), &x_plus, &s_plus, self.cfg()),
584 build_logit_state(self.y(), self.w(), &x_minus, &s_minus, self.cfg()),
585 )
586 }
587
588 fn fd_directional_gradient(&self, x_tau: &Array2<f64>, s_tau: &Array2<f64>) -> f64 {
590 let h = 2e-5;
591 let (state_plus, state_minus) = self.state_perturbed(x_tau, s_tau, h);
592 let v_plus = state_plus.compute_cost(self.rho()).expect("cost+");
593 let v_minus = state_minus.compute_cost(self.rho()).expect("cost-");
594 (v_plus - v_minus) / (2.0 * h)
595 }
596 }
597
598 pub(crate) fn build_logit_state<'a>(
599 y: &'a Array1<f64>,
600 w: &'a Array1<f64>,
601 x: &Array2<f64>,
602 s: &Array2<f64>,
603 cfg: &'a RemlConfig,
604 ) -> RemlState<'a> {
605 use crate::estimate::PenaltySpec;
606 let p = x.ncols();
607 let offset = Array1::<f64>::zeros(y.len());
608 let spec = PenaltySpec::Dense(s.clone());
609 let canonical =
610 gam_terms::construction::canonicalize_penalty_specs(&[spec], &[1], p, "test")
611 .map(|(canonical, _)| canonical)
612 .expect("canonicalize");
613 RemlState::newwith_offset(
614 y.view(),
615 x.clone(),
616 w.view(),
617 offset.view(),
618 canonical,
619 p,
620 cfg,
621 Some(vec![1]),
622 None,
623 None,
624 )
625 .expect("state")
626 }
627
628 fn bundle_with_inner_kkt_residual(bundle: &EvalShared, residual: Array1<f64>) -> EvalShared {
629 let mut pirls_result = bundle.pirls_result.as_ref().clone();
630 pirls_result.lastgradient_norm = residual.dot(&residual).sqrt();
631 pirls_result.penalized_gradient_transformed = residual;
632 let mut cloned = bundle.clone();
633 cloned.pirls_result = Arc::new(pirls_result);
634 cloned
635 }
636
637 fn evaluate_synthetic_psi_value_without_inner_kkt(
638 state: &RemlState<'_>,
639 rho: &Array1<f64>,
640 bundle: &EvalShared,
641 ) -> f64 {
642 let mode = super::reml_outer_engine::EvalMode::ValueOnly;
643 let mut assembly = state
644 .build_auto_assembly(rho, bundle, mode, true, false)
645 .expect("uncorrected synthetic-psi assembly");
646 assert!(
647 matches!(
648 &assembly.dispersion,
649 super::reml_outer_engine::DispersionHandling::Fixed { .. }
650 ),
651 "the #2305 bridge fixture must exercise the fixed-dispersion LAML identity"
652 );
653 let p_dim = assembly.beta.len();
654 assembly.ext_coords = vec![super::reml_outer_engine::HyperCoord {
655 a: 0.0,
656 g: Array1::zeros(p_dim),
657 drift: super::reml_outer_engine::HyperCoordDrift::none(),
658 ld_s: 0.0,
659 b_depends_on_beta: false,
660 is_penalty_like: false,
661 firth_g: None,
662 tk_eta_fixed: None,
663 tk_x_fixed: None,
664 }];
665 state
666 .assemble_and_evaluate(rho, bundle, mode, assembly)
667 .expect("uncorrected synthetic-psi value")
668 .cost
669 }
670
671 #[test]
672 fn psi_value_bridge_corrects_nonstationary_inner_mode_2305() {
673 let y = array![0.0, 1.0, 0.0, 1.0, 0.0, 1.0];
678 let w = Array1::<f64>::ones(y.len());
679 let x = array![
680 [1.0, -1.0, 0.2],
681 [1.0, -0.6, -0.4],
682 [1.0, -0.2, 0.7],
683 [1.0, 0.3, -0.5],
684 [1.0, 0.8, 0.1],
685 [1.0, 1.2, 0.6],
686 ];
687 let s = array![[0.0, 0.0, 0.0], [0.0, 1.0, 0.15], [0.0, 0.15, 0.8]];
688 let cfg = RemlConfig::external(binomial_logit_glm_spec(), 1e-12, false);
689 let state = build_logit_state(&y, &w, &x, &s, &cfg);
690 let rho = array![0.2];
691 let base = state.obtain_eval_bundle(&rho).expect("base inner mode");
692 let residual = array![0.18, -0.11, 0.07];
693 let capped = bundle_with_inner_kkt_residual(&base, residual);
694
695 let corrected = state
696 .evaluate_unified_value_only_with_synthetic_ext_count(&rho, &capped, 1, false)
697 .expect("psi value with inner-KKT correction");
698 let exact_kkt_assumption =
699 evaluate_synthetic_psi_value_without_inner_kkt(&state, &rho, &capped);
700 let residual_energy = corrected
701 .ift_residual_energy
702 .expect("design-moving bridge must attach the nonstationary KKT residual");
703
704 assert!(
705 residual_energy.abs() > 1e-10,
706 "the synthetic nonstationary residual must produce a material fixed-dispersion \
707 IFT correction, got residual_energy={residual_energy:.12e}"
708 );
709 assert_eq!(
710 corrected.cost,
711 exact_kkt_assumption - residual_energy,
712 "the psi value bridge must apply exactly the correction reported by the unified \
713 evaluator: corrected={:.12e}, exact-kkt={exact_kkt_assumption:.12e}, \
714 residual-energy={residual_energy:.12e}",
715 corrected.cost,
716 );
717 }
718
719 #[test]
720 fn psi_value_bridge_correction_vanishes_at_stationarity_2305() {
721 let y = array![0.0, 1.0, 0.0, 1.0, 0.0, 1.0];
722 let w = Array1::<f64>::ones(y.len());
723 let x = array![
724 [1.0, -1.0, 0.2],
725 [1.0, -0.6, -0.4],
726 [1.0, -0.2, 0.7],
727 [1.0, 0.3, -0.5],
728 [1.0, 0.8, 0.1],
729 [1.0, 1.2, 0.6],
730 ];
731 let s = array![[0.0, 0.0, 0.0], [0.0, 1.0, 0.15], [0.0, 0.15, 0.8]];
732 let cfg = RemlConfig::external(binomial_logit_glm_spec(), 1e-12, false);
733 let state = build_logit_state(&y, &w, &x, &s, &cfg);
734 let rho = array![0.2];
735 let base = state.obtain_eval_bundle(&rho).expect("base inner mode");
736 let stationary = bundle_with_inner_kkt_residual(&base, Array1::zeros(x.ncols()));
737
738 let corrected = state
739 .evaluate_unified_value_only_with_synthetic_ext_count(&rho, &stationary, 1, false)
740 .expect("stationary psi value with correction enabled");
741 let exact_kkt_assumption =
742 evaluate_synthetic_psi_value_without_inner_kkt(&state, &rho, &stationary);
743
744 assert_eq!(
745 corrected.ift_residual_energy,
746 Some(0.0),
747 "the design-moving bridge must attach the residual, whose correction is exactly \
748 zero at stationarity"
749 );
750 assert_eq!(
751 corrected.cost, exact_kkt_assumption,
752 "the generic inner-KKT correction must be exactly zero at a stationary inner mode"
753 );
754 }
755
756 #[test]
757 fn repeated_penalty_ranges_keep_analytic_outer_hessian() {
758 let y = array![0.2, -0.1, 0.3, 0.0];
759 let w = Array1::<f64>::ones(y.len());
760 let x = array![[1.0, -0.7], [1.0, -0.2], [1.0, 0.3], [1.0, 0.9]];
761 let offset = Array1::<f64>::zeros(y.len());
762 let cfg = RemlConfig::external(gaussian_identity_glm_spec(), 1e-10, false);
763 let p = x.ncols();
764 let canonical = vec![
765 gam_terms::construction::CanonicalPenalty::from_dense_root(array![[0.0, 1.0]], p),
766 gam_terms::construction::CanonicalPenalty::from_dense_root(array![[1.0, 0.0]], p),
767 ];
768 let state = RemlState::newwith_offset(
769 y.view(),
770 x,
771 w.view(),
772 offset.view(),
773 canonical,
774 p,
775 &cfg,
776 Some(vec![1, 1]),
777 None,
778 None,
779 )
780 .expect("state");
781
782 assert!(
783 state.analytic_outer_hessian_enabled(),
784 "double-penalty-style repeated coefficient ranges must still route to exact Hessian"
785 );
786 }
787
788 #[test]
804 fn adaptive_ift_controller_does_not_survive_into_the_next_state_in_the_same_slot() {
805 const DEFAULT_CAP: f64 = 7.5;
809
810 fn probe(record: bool) -> (usize, f64, bool) {
811 let y = array![0.2, -0.1, 0.3, 0.0];
812 let w = Array1::<f64>::ones(y.len());
813 let x = array![[1.0, -0.7], [1.0, -0.2], [1.0, 0.3], [1.0, 0.9]];
814 let offset = Array1::<f64>::zeros(y.len());
815 let cfg = RemlConfig::external(gaussian_identity_glm_spec(), 1e-10, false);
816 let p = x.ncols();
817 let canonical = vec![gam_terms::construction::CanonicalPenalty::from_dense_root(
818 array![[0.0, 1.0]],
819 p,
820 )];
821 let state = RemlState::newwith_offset(
822 y.view(),
823 x,
824 w.view(),
825 offset.view(),
826 canonical,
827 p,
828 &cfg,
829 Some(vec![1]),
830 None,
831 None,
832 )
833 .expect("state");
834 let address = &state as *const RemlState<'_> as usize;
835 if record {
836 let shrunk = state
840 .record_ift_prediction_quality(1.0e3, 0.25)
841 .expect("a finite quality and a positive cap must update the controller");
842 assert!(
843 shrunk < 0.25,
844 "a quality above the flat-fallback band must shrink the step cap; got {shrunk}"
845 );
846 return (address, shrunk, true);
847 }
848 let cap = state.ift_quality_step_cap(DEFAULT_CAP);
849 let flat = state.take_ift_quality_flat_override();
850 (address, cap, flat)
851 }
852
853 let (first_address, _, _) = probe(true);
854 let (second_address, cap, flat) = probe(false);
855
856 assert_eq!(
857 first_address, second_address,
858 "the probe only exercises the defect when the second state reuses the first's slot; \
859 both calls are to the same fn at the same depth, so the frame offsets must coincide"
860 );
861 assert_eq!(
862 cap, DEFAULT_CAP,
863 "a freshly built state must return the caller's default step cap, not the cap the \
864 previous state at this address shrank to"
865 );
866 assert!(
867 !flat,
868 "a freshly built state must not inherit the previous state's armed flat fallback"
869 );
870 }
871
872 #[test]
879 fn gaussian_profiled_diagonal_seed_clamps_into_its_validated_box() {
880 let n = 40usize;
883 let y = Array1::from_iter((0..n).map(|i| {
884 let t = (i as f64 + 0.5) / n as f64;
885 (std::f64::consts::TAU * t).sin() + 0.05 * (i as f64 % 3.0 - 1.0)
886 }));
887 let w = Array1::<f64>::ones(n);
888 let mut x = Array2::<f64>::zeros((n, 3));
889 for i in 0..n {
890 let t = (i as f64 + 0.5) / n as f64;
891 x[[i, 0]] = 1.0;
892 x[[i, 1]] = t;
893 x[[i, 2]] = t * t;
894 }
895 let offset = Array1::<f64>::zeros(n);
896 let cfg = RemlConfig::external(gaussian_identity_glm_spec(), 1e-10, false);
897 let p = x.ncols();
898 let canonical = vec![gam_terms::construction::CanonicalPenalty::from_dense_root(
899 array![[0.0, 1.0, 0.0], [0.0, 0.0, 1.0]],
900 p,
901 )];
902 let state = RemlState::newwith_offset(
903 y.view(),
904 x,
905 w.view(),
906 offset.view(),
907 canonical,
908 p,
909 &cfg,
910 Some(vec![1]),
911 None,
912 None,
913 )
914 .expect("state");
915
916 let wide = gam_problem::OrderedRhoBounds::new(-12.0, 12.0)
919 .expect("ordered rho bounds: the fixture passes lo < hi");
920 let seed_wide = state
921 .analytic_gaussian_profiled_diagonal_rho(wide)
922 .expect("no error")
923 .expect("gaussian-identity profiled diagonal returns a seed");
924 let natural = seed_wide[0];
925 for &r in seed_wide.iter() {
926 assert!(r.is_finite(), "seed coordinate is finite");
927 assert!(
928 (-12.0..=12.0).contains(&r),
929 "seed {r} stays inside the wide box"
930 );
931 }
932
933 let cap_hi = natural - 2.0;
939 let cap_lo = natural - 10.0;
940 let capped = gam_problem::OrderedRhoBounds::new(cap_lo, cap_hi)
941 .expect("ordered rho bounds: the fixture passes lo < hi");
942 let seed_capped = state
943 .analytic_gaussian_profiled_diagonal_rho(capped)
944 .expect("no error")
945 .expect("seed present");
946 assert!(
947 seed_capped.iter().all(|&r| (r - cap_hi).abs() < 1e-9),
948 "capped seed {seed_capped:?} clamps to the binding upper bound {cap_hi}"
949 );
950 }
951
952 #[test]
953 fn canonical_logit_firth_declines_exact_tk_hessian_when_row_pair_work_is_large() {
954 let n = 2_000usize;
955 let p = 28usize;
956 let y = Array1::from_iter((0..n).map(|i| if i % 3 == 0 { 1.0 } else { 0.0 }));
957 let w = Array1::<f64>::ones(n);
958 let mut x = Array2::<f64>::zeros((n, p));
959 for i in 0..n {
960 let t = (i as f64 + 0.5) / n as f64;
961 x[[i, 0]] = 1.0;
962 for j in 1..p {
963 x[[i, j]] = ((j as f64) * std::f64::consts::TAU * t).sin()
964 + 0.25 * (((j + 1) as f64) * std::f64::consts::TAU * t).cos();
965 }
966 }
967 let mut s = Array2::<f64>::zeros((p, p));
968 for j in 1..p {
969 s[[j, j]] = 1.0;
970 }
971 let cfg = RemlConfig::external(binomial_logit_glm_spec(), 1e-10, true);
972 let state = build_logit_state(&y, &w, &x, &s, &cfg);
973
974 assert!(
975 !RemlState::firth_tk_exact_hessian_scale_allows(n, p),
976 "fixture must sit beyond the O(n²·p) exact-Hessian budget"
977 );
978 assert!(
979 !state.analytic_outer_hessian_enabled(),
980 "large canonical-logit Firth fits should keep exact value/gradient but route outer curvature to BFGS"
981 );
982 }
983
984 #[test]
985 fn canonical_logit_firth_keeps_exact_tk_hessian_for_small_separation_guards() {
986 let n = 40usize;
987 let p = 6usize;
988 let y = Array1::from_iter((0..n).map(|i| if i >= n / 2 { 1.0 } else { 0.0 }));
989 let w = Array1::<f64>::ones(n);
990 let mut x = Array2::<f64>::zeros((n, p));
991 for i in 0..n {
992 let t = (i as f64) / (n - 1) as f64;
993 x[[i, 0]] = 1.0;
994 for j in 1..p {
995 x[[i, j]] = t.powi(j as i32);
996 }
997 }
998 let mut s = Array2::<f64>::zeros((p, p));
999 for j in 1..p {
1000 s[[j, j]] = 1.0;
1001 }
1002 let cfg = RemlConfig::external(binomial_logit_glm_spec(), 1e-10, true);
1003 let state = build_logit_state(&y, &w, &x, &s, &cfg);
1004
1005 assert!(RemlState::firth_tk_exact_hessian_scale_allows(n, p));
1006 assert!(
1007 state.analytic_outer_hessian_enabled(),
1008 "small Firth rescue fits should keep exact TK Hessian curvature"
1009 );
1010 }
1011
1012 #[test]
1013 fn nonlogit_firth_keeps_tk_value_and_gradient() {
1014 let y = array![0.0, 1.0, 0.0, 1.0, 1.0, 0.0, 1.0, 0.0];
1015 let w = Array1::<f64>::ones(y.len());
1016 let x = array![
1017 [1.0, -1.0, 0.3],
1018 [1.0, -0.7, -0.2],
1019 [1.0, -0.3, 0.4],
1020 [1.0, 0.0, -0.5],
1021 [1.0, 0.2, 0.6],
1022 [1.0, 0.6, -0.4],
1023 [1.0, 0.9, 0.2],
1024 [1.0, 1.3, -0.1],
1025 ];
1026 let s = array![[0.0, 0.0, 0.0], [0.0, 1.2, 0.1], [0.0, 0.1, 0.7]];
1027 let rho = array![0.15];
1028
1029 for link in [StandardLink::Probit, StandardLink::CLogLog] {
1030 let likelihood = GlmLikelihoodSpec::canonical(LikelihoodSpec::new(
1031 ResponseFamily::Binomial,
1032 InverseLink::Standard(link),
1033 ));
1034 let cfg = RemlConfig::external(likelihood, 1e-9, true).with_max_iterations(500);
1035 let state = build_logit_state(&y, &w, &x, &s, &cfg);
1036 assert!(
1037 !state.analytic_outer_hessian_enabled(),
1038 "{link:?} should use BFGS curvature until exact f_obs is available"
1039 );
1040
1041 let bundle = state
1042 .obtain_eval_bundle(&rho)
1043 .expect("non-logit Firth bundle");
1044 let atom = state
1045 .tierney_kadane_terms(
1046 &rho,
1047 &bundle,
1048 super::reml_outer_engine::EvalMode::ValueAndGradient,
1049 &[],
1050 )
1051 .expect("non-logit TK correction");
1052 let value = CriterionAtom::value(&atom);
1053 let gradient = atom.gradient().expect("TK gradient");
1054 assert!(
1055 value.is_finite() && value.abs() > 1e-12,
1056 "{link:?} must receive a material finite TK correction, got {value}"
1057 );
1058 assert_eq!(gradient.len(), rho.len());
1059 assert!(
1060 gradient.iter().all(|entry| entry.is_finite()),
1061 "{link:?} TK gradient must be finite: {gradient:?}"
1062 );
1063 }
1064 }
1065
1066 pub(crate) fn poisson_log_glm_spec() -> GlmLikelihoodSpec {
1067 GlmLikelihoodSpec::canonical(LikelihoodSpec::new(
1068 ResponseFamily::Poisson,
1069 InverseLink::Standard(StandardLink::Log),
1070 ))
1071 }
1072
1073 #[test]
1087 pub(crate) fn fixed_dispersion_laml_surface_is_replication_invariant() {
1088 let n = 200usize;
1089 let p = 8usize;
1090 let c = 3usize;
1091 let mut x = Array2::<f64>::zeros((n, p));
1092 let mut y = Array1::<f64>::zeros(n);
1093 for i in 0..n {
1094 let t = (i as f64) / ((n - 1) as f64);
1095 let tau = std::f64::consts::TAU;
1096 x[[i, 0]] = 1.0;
1097 x[[i, 1]] = t;
1098 x[[i, 2]] = (tau * t).sin();
1099 x[[i, 3]] = (tau * t).cos();
1100 x[[i, 4]] = (2.0 * tau * t).sin();
1101 x[[i, 5]] = (2.0 * tau * t).cos();
1102 x[[i, 6]] = (3.0 * tau * t).sin();
1103 x[[i, 7]] = (3.0 * tau * t).cos();
1104 let eta = 0.3 + 0.9 * (1.4 * (t - 0.5)).sin();
1105 y[i] = (eta.exp() + 0.5 * ((i as f64) * 2.399_963).sin())
1107 .round()
1108 .max(0.0);
1109 }
1110 let mut s = Array2::<f64>::zeros((p, p));
1111 for j in 1..p {
1112 s[[j, j]] = 1.0;
1113 }
1114
1115 let mut x_rep = Array2::<f64>::zeros((n * c, p));
1117 let mut y_rep = Array1::<f64>::zeros(n * c);
1118 for r in 0..c {
1119 for i in 0..n {
1120 let row = r * n + i;
1121 for j in 0..p {
1122 x_rep[[row, j]] = x[[i, j]];
1123 }
1124 y_rep[row] = y[i];
1125 }
1126 }
1127
1128 let w_weighted = Array1::<f64>::from_elem(n, c as f64);
1129 let w_rep = Array1::<f64>::ones(n * c);
1130
1131 let cfg = RemlConfig::external(poisson_log_glm_spec(), 1e-10, false);
1132 let st_w = build_logit_state(&y, &w_weighted, &x, &s, &cfg);
1133 let st_r = build_logit_state(&y_rep, &w_rep, &x_rep, &s, &cfg);
1134
1135 for &rho in &[-2.0_f64, -1.0, 0.0, 1.0, 2.0, 3.0, 4.0, 5.0] {
1136 let r = Array1::from_elem(1, rho);
1137 let cw = st_w.compute_cost(&r).expect("weighted cost");
1138 let cr = st_r.compute_cost(&r).expect("replicated cost");
1139 let gw = st_w.compute_gradient(&r).expect("weighted grad");
1140 let gr = st_r.compute_gradient(&r).expect("replicated grad");
1141 assert!(
1144 (cw - cr).abs() <= 1e-9 * (1.0 + cw.abs()),
1145 "LAML cost differs between w=c and c× replication at rho={rho}: \
1146 cost_w={cw:.12e} cost_r={cr:.12e} diff={:.3e}",
1147 cw - cr
1148 );
1149 assert!(
1150 (gw[0] - gr[0]).abs() <= 1e-9 * (1.0 + gw[0].abs()),
1151 "LAML gradient differs between w=c and c× replication at rho={rho}: \
1152 g_w={:.12e} g_r={:.12e} diff={:.3e}",
1153 gw[0],
1154 gr[0],
1155 gw[0] - gr[0]
1156 );
1157 }
1158 }
1159
1160 #[test]
1169 pub(crate) fn rho_weight_anchor_is_zero_for_fixed_dispersion() {
1170 let n = 50usize;
1171 let p = 3usize;
1172 let mut x = Array2::<f64>::zeros((n, p));
1173 let mut y = Array1::<f64>::zeros(n);
1174 for i in 0..n {
1175 let t = (i as f64) / ((n - 1) as f64);
1176 x[[i, 0]] = 1.0;
1177 x[[i, 1]] = t;
1178 x[[i, 2]] = t * t;
1179 y[i] = (1.0 + (3.0 * t).sin()).round().max(0.0);
1180 }
1181 let mut s = Array2::<f64>::zeros((p, p));
1182 s[[2, 2]] = 1.0;
1183 let c = 4.0_f64;
1185 let w = Array1::<f64>::from_elem(n, c);
1186
1187 let cfg_pois = RemlConfig::external(poisson_log_glm_spec(), 1e-10, false);
1188 let st_pois = build_logit_state(&y, &w, &x, &s, &cfg_pois);
1189 assert_eq!(
1190 st_pois.rho_weight_anchor(),
1191 0.0,
1192 "fixed-dispersion (Poisson) anchor must be 0, not the geometric-mean log-weight"
1193 );
1194
1195 let cfg_gauss = RemlConfig::external(gaussian_identity_glm_spec(), 1e-10, false);
1196 let st_gauss = build_logit_state(&y, &w, &x, &s, &cfg_gauss);
1197 assert!(
1198 (st_gauss.rho_weight_anchor() - c.ln()).abs() <= 1e-12,
1199 "Gaussian-identity (profiled) anchor must be the geometric-mean log-weight ln(c)={:.6}, got {:.6}",
1200 c.ln(),
1201 st_gauss.rho_weight_anchor()
1202 );
1203 }
1204
1205 #[test]
1227 pub(crate) fn soft_rho_guard_gradient_is_evaluated_at_the_weight_anchor() {
1228 let n = 50usize;
1229 let p = 3usize;
1230 let mut x = Array2::<f64>::zeros((n, p));
1231 let mut y = Array1::<f64>::zeros(n);
1232 for i in 0..n {
1233 let t = (i as f64) / ((n - 1) as f64);
1234 x[[i, 0]] = 1.0;
1235 x[[i, 1]] = t;
1236 x[[i, 2]] = t * t;
1237 y[i] = (1.0 + (3.0 * t).sin()).round().max(0.0);
1238 }
1239 let mut s = Array2::<f64>::zeros((p, p));
1240 s[[2, 2]] = 1.0;
1241 let c = 4.0_f64;
1242 let w = Array1::<f64>::from_elem(n, c);
1243 let cfg = RemlConfig::external(gaussian_identity_glm_spec(), 1e-10, false);
1244 let state = build_logit_state(&y, &w, &x, &s, &cfg);
1245 let anchor = state.rho_weight_anchor();
1246 assert!(
1247 (anchor - c.ln()).abs() <= 1e-12 && anchor.abs() > 1.0,
1248 "this gate is only meaningful where the anchor is materially nonzero; \
1249 got {anchor:.6e}"
1250 );
1251
1252 let weight = crate::estimate::RHO_SOFT_PRIOR_WEIGHT;
1253 let sharpness = crate::estimate::RHO_SOFT_PRIOR_SHARPNESS;
1254 let bound = crate::estimate::RHO_BOUND;
1255 let a = sharpness / bound;
1256 let rho = array![5.0, -3.0, 30.0];
1257 let published = state.soft_rho_guard_gradient(&rho);
1258 assert_eq!(
1259 published.len(),
1260 rho.len(),
1261 "the published barrier gradient must be one entry per rho coordinate"
1262 );
1263 for (k, &value) in rho.iter().enumerate() {
1264 let anchored = weight * a * (a * (value - anchor)).tanh();
1265 let raw = weight * a * (a * value).tanh();
1266 assert!(
1267 (published[k] - anchored).abs() <= 1e-18,
1268 "coordinate {k} (rho={value}): the published barrier gradient must be \
1269 the ANCHORED closed form {anchored:.12e}, got {:.12e}",
1270 published[k]
1271 );
1272 if value.abs() < 10.0 {
1276 assert!(
1277 (raw - anchored).abs() > 1e-9 * anchored.abs().max(1e-30),
1278 "coordinate {k} (rho={value}) must DISCRIMINATE the anchored form \
1279 from the raw-rho one, else this gate cannot catch a dropped \
1280 anchor: anchored={anchored:.12e} raw={raw:.12e}"
1281 );
1282 }
1283 }
1284
1285 let atom = state.soft_rho_guard_prior_atom(&rho);
1288 assert_eq!(
1289 published,
1290 *atom.gradient(),
1291 "the published barrier gradient must be the criterion's own atom \
1292 emission bit for bit, so the certificate's subtraction cannot drift \
1293 from the criterion's addition"
1294 );
1295
1296 let unit_weights = Array1::<f64>::ones(n);
1304 let unit_state = build_logit_state(&y, &unit_weights, &x, &s, &cfg);
1305 assert_eq!(
1306 unit_state.rho_weight_anchor(),
1307 0.0,
1308 "the unweighted arm must anchor at 0, else this pair does not isolate \
1309 the rescale"
1310 );
1311 let slid = &rho + c.ln();
1320 let rescaled = state.soft_rho_guard_gradient(&slid);
1321 let unrescaled = unit_state.soft_rho_guard_gradient(&rho);
1322 for k in 0..rho.len() {
1323 let scale = unrescaled[k].abs().max(1.0e-30);
1324 assert!(
1325 (rescaled[k] - unrescaled[k]).abs() <= 1.0e-12 * scale,
1326 "coordinate {k}: the barrier gradient the certificate subtracts must \
1327 be the SAME at corresponding coordinates under a global weight \
1328 rescale w -> c*w (#877/#2545) — a subtrahend evaluated at raw rho \
1329 would differ here and make the certificate weight-dependent. \
1330 rescaled={:.17e} unrescaled={:.17e} rel={:.3e}",
1331 rescaled[k],
1332 unrescaled[k],
1333 (rescaled[k] - unrescaled[k]).abs() / scale
1334 );
1335 }
1336 }
1337
1338 pub(crate) fn beta_original_from_bundle(bundle: &EvalShared) -> Array1<f64> {
1339 let pr = bundle.pirls_result.as_ref();
1340 match pr.coordinate_frame {
1341 PirlsCoordinateFrame::OriginalSparseNative => pr.beta_transformed.as_ref().clone(),
1342 PirlsCoordinateFrame::TransformedQs => {
1343 pr.reparam_result.qs.dot(pr.beta_transformed.as_ref())
1344 }
1345 }
1346 }
1347
1348 pub(crate) fn compute_joint_hypercostgradienthessian(
1349 state: &RemlState<'_>,
1350 theta: &Array1<f64>,
1351 rho_dim: usize,
1352 hyper_dirs: &[DirectionalHyperParam],
1353 ) -> Result<(f64, Array1<f64>, Array2<f64>), EstimationError> {
1354 let (cost, gradient, hessian) = state.compute_joint_hyper_eval_with_order(
1355 theta,
1356 rho_dim,
1357 hyper_dirs,
1358 crate::rho_optimizer::OuterEvalOrder::ValueGradientHessian,
1359 )?;
1360 Ok((
1361 cost,
1362 gradient,
1363 hessian
1364 .materialize_dense()
1365 .map_err(|error| EstimationError::RemlOptimizationFailed(error.to_string()))?
1366 .ok_or_else(|| {
1367 EstimationError::RemlOptimizationFailed(
1368 "joint hyper Hessian requested but unavailable".to_string(),
1369 )
1370 })?,
1371 ))
1372 }
1373
1374 pub(crate) fn h_original_from_bundle(bundle: &EvalShared) -> Array2<f64> {
1375 let pr = bundle.pirls_result.as_ref();
1376 match pr.coordinate_frame {
1377 PirlsCoordinateFrame::OriginalSparseNative => bundle.h_total.as_ref().clone(),
1378 PirlsCoordinateFrame::TransformedQs => {
1379 let qs = &pr.reparam_result.qs;
1380 let tmp = gam_linalg::faer_ndarray::fast_ab(qs, bundle.h_total.as_ref());
1381 gam_linalg::faer_ndarray::fast_abt(&tmp, qs)
1382 }
1383 }
1384 }
1385
1386 pub(crate) fn single_directional_tau_gradient(
1387 state: &RemlState<'_>,
1388 rho: &Array1<f64>,
1389 hyper: DirectionalHyperParam,
1390 ) -> Result<f64, EstimationError> {
1391 let mut theta = Array1::<f64>::zeros(rho.len() + 1);
1392 theta.slice_mut(s![..rho.len()]).assign(rho);
1393 let (_, gradient, _) = state.compute_joint_hyper_eval_with_order(
1394 &theta,
1395 rho.len(),
1396 &[hyper],
1397 crate::rho_optimizer::OuterEvalOrder::ValueAndGradient,
1398 )?;
1399 Ok(gradient[rho.len()])
1400 }
1401
1402 pub(crate) fn fd_directional_tau_cost_gradient(
1403 y: &Array1<f64>,
1404 w: &Array1<f64>,
1405 x: &Array2<f64>,
1406 s0: &Array2<f64>,
1407 cfg: &RemlConfig,
1408 rho: &Array1<f64>,
1409 x_tau: &Array2<f64>,
1410 s_tau: &Array2<f64>,
1411 ) -> f64 {
1412 let h = 2e-5;
1413 let x_plus = x + &x_tau.mapv(|v| h * v);
1414 let x_minus = x - &x_tau.mapv(|v| h * v);
1415 let s_plus = s0 + &s_tau.mapv(|v| h * v);
1416 let s_minus = s0 - &s_tau.mapv(|v| h * v);
1417 let state_plus = build_logit_state(y, w, &x_plus, &s_plus, cfg);
1418 let state_minus = build_logit_state(y, w, &x_minus, &s_minus, cfg);
1419 let v_plus = state_plus.compute_cost(rho).expect("cost+");
1420 let v_minus = state_minus.compute_cost(rho).expect("cost-");
1421 (v_plus - v_minus) / (2.0 * h)
1422 }
1423
1424 pub(crate) fn directional_tau_hessian_fd_reference(
1425 y: &Array1<f64>,
1426 w: &Array1<f64>,
1427 x: &Array2<f64>,
1428 s0: &Array2<f64>,
1429 cfg: &RemlConfig,
1430 rho: &Array1<f64>,
1431 hyper_dirs: &[DirectionalHyperParam],
1432 x_tau_mats: &[Array2<f64>],
1433 s_tau_mats: &[Array2<f64>],
1434 ) -> Array2<f64> {
1435 assert_eq!(hyper_dirs.len(), x_tau_mats.len());
1436 assert_eq!(hyper_dirs.len(), s_tau_mats.len());
1437
1438 const TARGET_PHYSICAL_STEP: f64 = 1e-5;
1439
1440 let n_dirs = hyper_dirs.len();
1441 let mut h_ttfd = Array2::<f64>::zeros((n_dirs, n_dirs));
1442 for j in 0..n_dirs {
1443 let direction_scale = x_tau_mats[j]
1444 .iter()
1445 .chain(s_tau_mats[j].iter())
1446 .fold(0.0_f64, |acc, value| acc.max(value.abs()));
1447 let h = if direction_scale > 0.0 {
1448 TARGET_PHYSICAL_STEP / direction_scale
1449 } else {
1450 TARGET_PHYSICAL_STEP
1451 };
1452
1453 let x_plus = x + &x_tau_mats[j].mapv(|v| h * v);
1454 let x_minus = x - &x_tau_mats[j].mapv(|v| h * v);
1455 let s_plus = s0 + &s_tau_mats[j].mapv(|v| h * v);
1456 let s_minus = s0 - &s_tau_mats[j].mapv(|v| h * v);
1457
1458 let state_plus = build_logit_state(y, w, &x_plus, &s_plus, cfg);
1459 let state_minus = build_logit_state(y, w, &x_minus, &s_minus, cfg);
1460 for i in 0..n_dirs {
1461 let g_plus =
1462 single_directional_tau_gradient(&state_plus, rho, hyper_dirs[i].clone())
1463 .expect("g+ for FD");
1464 let g_minus =
1465 single_directional_tau_gradient(&state_minus, rho, hyper_dirs[i].clone())
1466 .expect("g- for FD");
1467 h_ttfd[[i, j]] = (g_plus - g_minus) / (2.0 * h);
1468 }
1469 }
1470 symmetrize_in_place(&mut h_ttfd);
1471 h_ttfd
1472 }
1473
1474 #[test]
1475 pub(crate) fn eval_cache_manager_stores_first_order_outer_eval() {
1476 let cache = EvalCacheManager::new();
1477 let rho = array![0.25, -0.0];
1478 let rho_key = super::rho_key::sanitized_rhokey(&rho);
1479 let eval = OuterEval {
1480 cost: 3.5,
1481 gradient: array![1.0, -2.0],
1482 hessian: HessianValue::Unavailable,
1483 inner_beta_hint: None,
1484 };
1485
1486 cache.store_outer_eval(&rho_key, &eval);
1487
1488 let cached = cache
1489 .cached_outer_eval(&rho_key)
1490 .expect("first-order outer eval should be cached");
1491 assert_eq!(cached.cost, eval.cost);
1492 assert_eq!(cached.gradient, eval.gradient);
1493 assert!(matches!(cached.hessian, HessianValue::Unavailable));
1494
1495 cache.invalidate_eval_bundle();
1496 assert!(
1497 cache.cached_outer_eval(&rho_key).is_none(),
1498 "invalidating the bundle should clear the outer-eval cache too"
1499 );
1500 }
1501
1502 #[test]
1512 pub(crate) fn outer_eval_lru_hit_is_bit_identical_and_evicts_honestly_1575() {
1513 use super::OUTER_EVAL_LRU_CAPACITY;
1514
1515 let make_eval = |seed: f64| OuterEval {
1518 cost: (seed * std::f64::consts::PI).sin() / 3.0 - seed,
1519 gradient: array![seed, -seed * 2.0, seed.recip()],
1520 hessian: HessianValue::Unavailable,
1521 inner_beta_hint: Some(array![seed + 0.5, seed - 0.5]),
1522 };
1523 let bits_eq = |a: &OuterEval, b: &OuterEval| -> bool {
1524 a.cost.to_bits() == b.cost.to_bits()
1525 && a.gradient.len() == b.gradient.len()
1526 && a.gradient
1527 .iter()
1528 .zip(b.gradient.iter())
1529 .all(|(x, y)| x.to_bits() == y.to_bits())
1530 };
1531
1532 let cache = EvalCacheManager::new();
1533
1534 let rho_a = array![0.25, -1.5];
1537 let key_a = super::rho_key::sanitized_rhokey(&rho_a);
1538 let eval_a = make_eval(0.25);
1539 cache.store_outer_eval(&key_a, &eval_a);
1540 let hit_a = cache
1541 .cached_outer_eval(&key_a)
1542 .expect("stored rho_a must hit");
1543 assert!(
1544 bits_eq(&hit_a, &eval_a),
1545 "cache hit must be bit-identical (cost+gradient) to the stored miss-path eval"
1546 );
1547 assert_eq!(
1548 hit_a.inner_beta_hint.as_ref().map(|b| b.to_vec()),
1549 eval_a.inner_beta_hint.as_ref().map(|b| b.to_vec()),
1550 "inner_beta_hint must round-trip unchanged"
1551 );
1552
1553 let rho_b = array![0.25, -1.4999999999999998];
1556 let key_b = super::rho_key::sanitized_rhokey(&rho_b);
1557 assert_ne!(key_a, key_b, "the two rho-keys must differ");
1558 let eval_b = make_eval(7.0);
1559 cache.store_outer_eval(&key_b, &eval_b);
1560 assert!(
1561 bits_eq(
1562 &cache.cached_outer_eval(&key_b).expect("rho_b must hit"),
1563 &eval_b
1564 ),
1565 "rho_b must return its own eval, not rho_a's"
1566 );
1567 assert!(
1568 bits_eq(
1569 &cache
1570 .cached_outer_eval(&key_a)
1571 .expect("rho_a must still hit"),
1572 &eval_a
1573 ),
1574 "rho_a must be unaffected by the rho_b insert"
1575 );
1576
1577 let cache = EvalCacheManager::new();
1581 let mut keys = Vec::new();
1582 let mut evals = Vec::new();
1583 for i in 0..OUTER_EVAL_LRU_CAPACITY {
1584 let rho = array![i as f64, -(i as f64)];
1585 let key = super::rho_key::sanitized_rhokey(&rho);
1586 let eval = make_eval(i as f64 + 0.123);
1587 cache.store_outer_eval(&key, &eval);
1588 keys.push(key);
1589 evals.push(eval);
1590 }
1591 assert_eq!(
1593 cache
1594 .outer_eval_lru
1595 .read()
1596 .expect("outer-eval LRU lock is poisoned: a writer panicked while holding it")
1597 .entries
1598 .len(),
1599 OUTER_EVAL_LRU_CAPACITY
1600 );
1601 let rho_overflow = array![999.0, -999.0];
1603 let key_overflow = super::rho_key::sanitized_rhokey(&rho_overflow);
1604 let eval_overflow = make_eval(42.0);
1605 cache.store_outer_eval(&key_overflow, &eval_overflow);
1606 assert_eq!(
1607 cache
1608 .outer_eval_lru
1609 .read()
1610 .expect("outer-eval LRU lock is poisoned: a writer panicked while holding it")
1611 .entries
1612 .len(),
1613 OUTER_EVAL_LRU_CAPACITY,
1614 "capacity must stay bounded"
1615 );
1616 assert!(
1617 cache.cached_outer_eval(&keys[0]).is_none(),
1618 "the least-recently-used key must be evicted and now MISS (recompute), not return stale"
1619 );
1620 assert!(
1621 bits_eq(
1622 &cache
1623 .cached_outer_eval(&keys[1])
1624 .expect("a still-resident key must hit"),
1625 &evals[1]
1626 ),
1627 "a still-resident key must return its exact stored bits"
1628 );
1629 assert!(
1630 bits_eq(
1631 &cache
1632 .cached_outer_eval(&key_overflow)
1633 .expect("the freshest key must hit"),
1634 &eval_overflow
1635 ),
1636 "the freshest key must hit with its own eval"
1637 );
1638 }
1639
1640 #[test]
1641 pub(crate) fn reset_outer_seed_state_clears_pirls_cache() {
1642 let y = array![0.0, 1.0, 1.0, 0.0, 0.0, 1.0];
1648 let w = Array1::<f64>::ones(y.len());
1649 let x = array![
1650 [1.0, -1.0, 0.2],
1651 [1.0, -0.5, -0.4],
1652 [1.0, 0.0, 0.7],
1653 [1.0, 0.4, -0.3],
1654 [1.0, 0.9, 0.1],
1655 [1.0, 1.3, -0.6],
1656 ];
1657 let s0 = array![[0.0, 0.0, 0.0], [0.0, 1.1, 0.15], [0.0, 0.15, 0.8],];
1658 let rho = array![0.0];
1659 let cfg = RemlConfig::external(binomial_logit_glm_spec(), 1e-10, false);
1660 let state = build_logit_state(&y, &w, &x, &s0, &cfg);
1661
1662 state
1665 .compute_outer_eval_with_order(
1666 &rho,
1667 crate::rho_optimizer::OuterEvalOrder::ValueAndGradient,
1668 )
1669 .expect("outer eval should succeed");
1670
1671 let populated_len = state
1672 .cache_manager
1673 .pirls_cache
1674 .read()
1675 .expect("PIRLS cache lock is poisoned: a writer panicked while holding it")
1676 .map
1677 .len();
1678 assert!(
1679 populated_len > 0,
1680 "evaluating the outer objective should populate the PIRLS LRU, got {populated_len}"
1681 );
1682
1683 state.reset_outer_seed_state();
1684
1685 let cleared_len = state
1686 .cache_manager
1687 .pirls_cache
1688 .read()
1689 .expect("PIRLS cache lock is poisoned: a writer panicked while holding it")
1690 .map
1691 .len();
1692 assert_eq!(
1693 cleared_len, 0,
1694 "reset_outer_seed_state must clear the cross-call PIRLS LRU; got {cleared_len} entries"
1695 );
1696 }
1697
1698 #[test]
1699 pub(crate) fn reset_outer_seed_state_preserves_frozen_negbin_theta_1448() {
1700 use std::sync::atomic::Ordering;
1717
1718 let y = array![0.0, 1.0, 1.0, 0.0, 0.0, 1.0];
1719 let w = Array1::<f64>::ones(y.len());
1720 let x = array![
1721 [1.0, -1.0, 0.2],
1722 [1.0, -0.5, -0.4],
1723 [1.0, 0.0, 0.7],
1724 [1.0, 0.4, -0.3],
1725 [1.0, 0.9, 0.1],
1726 [1.0, 1.3, -0.6],
1727 ];
1728 let s0 = array![[0.0, 0.0, 0.0], [0.0, 1.1, 0.15], [0.0, 0.15, 0.8],];
1729 let cfg = RemlConfig::external(binomial_logit_glm_spec(), 1e-10, false);
1730 let state = build_logit_state(&y, &w, &x, &s0, &cfg);
1731
1732 let theta_final_bits = 2.5_f64.to_bits();
1734 state
1735 .frozen_negbin_theta
1736 .store(theta_final_bits, Ordering::Relaxed);
1737 assert_eq!(
1738 state.frozen_negbin_theta.load(Ordering::Relaxed),
1739 theta_final_bits,
1740 "precondition: the re-freeze stores θ_final into the frozen slot"
1741 );
1742
1743 state.reset_outer_seed_state();
1745
1746 assert_eq!(
1747 state.frozen_negbin_theta.load(Ordering::Relaxed),
1748 theta_final_bits,
1749 "reset_outer_seed_state (alternation-round reset) must PRESERVE the \
1750 re-frozen NB θ; clearing it would defeat the #1448 θ↔λ alternation \
1751 (the next ρ search would re-derive θ from the seed and never reach \
1752 the joint fixed point)"
1753 );
1754 }
1755
1756 #[test]
1757 pub(crate) fn implicit_hyper_design_derivative_respects_full_model_embedding() {
1758 let operator = ImplicitDesignPsiDerivative::new(
1759 array![1.0, 2.0, 3.0, 4.0],
1760 array![0.5, -1.0, 1.5, 2.0],
1761 array![0.1, 0.2, 0.3, 0.4],
1762 array![[1.0, 0.2], [0.5, 0.1], [1.5, 0.3], [2.0, 0.4]],
1763 None,
1764 None,
1765 2,
1766 2,
1767 1,
1768 2,
1769 );
1770 let local = operator
1771 .materialize_first(0)
1772 .expect("materialized first derivative");
1773 assert_eq!(
1774 local.ncols(),
1775 3,
1776 "operator-local derivative should stay smooth-local"
1777 );
1778
1779 let implicit = HyperDesignDerivative::from_implicit(
1780 Arc::new(operator),
1781 ImplicitDerivLevel::First(0),
1782 1..4,
1783 5,
1784 );
1785 let embedded = HyperDesignDerivative::from_embedded(local.clone(), 1..4, 5);
1786
1787 assert_eq!(implicit.nrows(), embedded.nrows());
1788 assert_eq!(implicit.ncols(), 5);
1789 assert_eq!(implicit.materialize(), embedded.materialize());
1790
1791 let u = array![7.0, 1.5, -2.0, 0.25, -3.0];
1792 let v = array![0.75, -1.25];
1793 assert_eq!(
1794 implicit.forward_mul_original(&u).expect("implicit forward"),
1795 embedded.forward_mul_original(&u).expect("embedded forward")
1796 );
1797 assert_eq!(
1798 implicit
1799 .transpose_mul_original(&v)
1800 .expect("implicit transpose"),
1801 embedded
1802 .transpose_mul_original(&v)
1803 .expect("embedded transpose")
1804 );
1805
1806 let qs = array![
1807 [1.0, 0.0, 0.0],
1808 [0.0, 1.0, 0.0],
1809 [0.0, 0.5, 0.5],
1810 [0.0, 0.0, 1.0],
1811 [0.0, 0.0, 0.0],
1812 ];
1813 assert_eq!(
1814 implicit
1815 .transformed(&qs, None)
1816 .expect("implicit transformed"),
1817 embedded
1818 .transformed(&qs, None)
1819 .expect("embedded transformed")
1820 );
1821 let u_transformed = array![1.0, -0.5, 2.0];
1822 assert_eq!(
1823 implicit
1824 .transformed_forward_mul(&qs, None, &u_transformed)
1825 .expect("implicit transformed forward"),
1826 embedded
1827 .transformed_forward_mul(&qs, None, &u_transformed)
1828 .expect("embedded transformed forward")
1829 );
1830 assert_eq!(
1831 implicit
1832 .transformed_transpose_mul(&qs, None, &v)
1833 .expect("implicit transformed transpose"),
1834 embedded
1835 .transformed_transpose_mul(&qs, None, &v)
1836 .expect("embedded transformed transpose")
1837 );
1838 }
1839
1840 #[test]
1841 pub(crate) fn directional_hyper_identities_match_finite_differences_logit() {
1842 let y = array![0.0, 1.0, 0.0, 1.0, 0.0, 1.0, 1.0, 0.0];
1843 let w = Array1::<f64>::ones(y.len());
1844 let x = array![
1845 [1.0, -1.2, 0.3],
1846 [1.0, -0.8, -0.4],
1847 [1.0, -0.3, 0.7],
1848 [1.0, 0.1, -0.9],
1849 [1.0, 0.5, 0.2],
1850 [1.0, 0.9, -0.1],
1851 [1.0, 1.3, 0.8],
1852 [1.0, 1.7, -0.6],
1853 ];
1854 let s0 = array![[0.0, 0.0, 0.0], [0.0, 1.2, 0.2], [0.0, 0.2, 0.9],];
1855
1856 let x_tau = Array2::<f64>::zeros(x.raw_dim());
1861 let s_tau = array![[0.0, 0.0, 0.0], [0.0, 0.25, 0.04], [0.0, 0.04, 0.15],];
1862 let hyper =
1863 DirectionalHyperParam::single_penalty(0, x_tau.clone(), s_tau.clone(), None, None)
1864 .expect("single-penalty hyper direction");
1865 let rho = array![0.0];
1866
1867 let cfg = RemlConfig::external(binomial_logit_glm_spec(), 1e-14, false);
1871 let state = build_logit_state(&y, &w, &x, &s0, &cfg);
1872 let bundle = state.obtain_eval_bundle(&rho).expect("bundle");
1873 let pr = bundle.pirls_result.as_ref();
1874
1875 let beta = beta_original_from_bundle(&bundle);
1876 let h_orig = h_original_from_bundle(&bundle);
1877 let u = &pr.solveweights * &(&pr.solveworking_response - &pr.final_eta);
1878
1879 let x_tau_beta = gam_linalg::faer_ndarray::fast_av(&x_tau, &beta);
1882 let weighted_x_tau_beta = &pr.finalweights * &x_tau_beta;
1883 let rhs = gam_linalg::faer_ndarray::fast_atv(&x_tau, &u)
1884 - gam_linalg::faer_ndarray::fast_atv(&x, &weighted_x_tau_beta)
1885 - s_tau.dot(&beta);
1886 let chol = h_orig.cholesky(Side::Lower).expect("chol(H)");
1887 let b_analytic = chol.solvevec(&rhs);
1888
1889 let eta_dot = &x_tau_beta + &gam_linalg::faer_ndarray::fast_av(&x, &b_analytic);
1893 let w_direction = crate::pirls::directionalworking_curvature_from_c_array(
1894 &pr.solve_c_array.to_owned(),
1895 &eta_dot,
1896 );
1897 let wx = RemlState::row_scale(&x, &pr.finalweights.to_owned());
1898 let wx_tau = RemlState::row_scale(&x_tau, &pr.finalweights.to_owned());
1899 let mut xwtau_x = x.clone();
1900 match w_direction {
1901 crate::pirls::DirectionalWorkingCurvature::Diagonal(diag) => {
1902 xwtau_x = RemlState::row_scale(&xwtau_x, &diag);
1903 }
1904 }
1905 let mut h_tau_analytic = gam_linalg::faer_ndarray::fast_atb(&x_tau, &wx);
1906 h_tau_analytic += &gam_linalg::faer_ndarray::fast_atb(&x, &wx_tau);
1907 h_tau_analytic += &gam_linalg::faer_ndarray::fast_atb(&x, &xwtau_x);
1908 h_tau_analytic += &s_tau;
1909
1910 let ell_beta = gam_linalg::faer_ndarray::fast_atv(&x, &u);
1915 let s_eff = &h_orig - &gam_linalg::faer_ndarray::fast_atb(&x, &wx);
1916 let cancellation = -ell_beta.dot(&b_analytic) + beta.dot(&s_eff.dot(&b_analytic));
1917
1918 let h = 2e-5;
1920 let x_plus = &x + &(x_tau.mapv(|v| h * v));
1921 let x_minus = &x - &(x_tau.mapv(|v| h * v));
1922 let s_plus = &s0 + &(s_tau.mapv(|v| h * v));
1923 let s_minus = &s0 - &(s_tau.mapv(|v| h * v));
1924
1925 let state_plus = build_logit_state(&y, &w, &x_plus, &s_plus, &cfg);
1926 let state_minus = build_logit_state(&y, &w, &x_minus, &s_minus, &cfg);
1927 let bundle_plus = state_plus.obtain_eval_bundle(&rho).expect("bundle+");
1928 let bundle_minus = state_minus.obtain_eval_bundle(&rho).expect("bundle-");
1929 let beta_plus = beta_original_from_bundle(&bundle_plus);
1930 let beta_minus = beta_original_from_bundle(&bundle_minus);
1931 let bfd = (&beta_plus - &beta_minus).mapv(|v| v / (2.0 * h));
1932
1933 let h_plus = h_original_from_bundle(&bundle_plus);
1934 let h_minus = h_original_from_bundle(&bundle_minus);
1935 let h_taufd = (&h_plus - &h_minus).mapv(|v| v / (2.0 * h));
1936
1937 let v_plus = state_plus.compute_cost(&rho).expect("cost+");
1938 let v_minus = state_minus.compute_cost(&rho).expect("cost-");
1939 let v_taufd = (v_plus - v_minus) / (2.0 * h);
1940
1941 let v_tau_analytic = single_directional_tau_gradient(&state, &rho, hyper.clone())
1942 .expect("analytic directional gradient");
1943
1944 let b_num = (&b_analytic - &bfd).mapv(|v| v * v).sum().sqrt();
1945 let b_den = bfd.mapv(|v| v * v).sum().sqrt().max(1e-12);
1946 let b_rel = b_num / b_den;
1947 for i in 0..b_analytic.len() {
1948 assert_eq!(
1949 b_analytic[i].signum(),
1950 bfd[i].signum(),
1951 "B sign mismatch at i={i}: analytic={} fd={}",
1952 b_analytic[i],
1953 bfd[i]
1954 );
1955 }
1956 assert!(
1957 b_rel < 2e-2,
1958 "B implicit solve mismatch vs FD: rel={b_rel:.3e}, num={b_num:.3e}, den={b_den:.3e}"
1959 );
1960
1961 let dh_num = (&h_tau_analytic - &h_taufd).mapv(|v| v * v).sum().sqrt();
1962 let dh_den = h_taufd.mapv(|v| v * v).sum().sqrt().max(1e-12);
1963 let dh_rel = dh_num / dh_den;
1964 for i in 0..h_tau_analytic.nrows() {
1965 for j in 0..h_tau_analytic.ncols() {
1966 assert_eq!(
1967 h_tau_analytic[[i, j]].signum(),
1968 h_taufd[[i, j]].signum(),
1969 "H_tau sign mismatch at ({i},{j}): analytic={} fd={}",
1970 h_tau_analytic[[i, j]],
1971 h_taufd[[i, j]]
1972 );
1973 }
1974 }
1975 assert!(
1976 dh_rel < 3e-2,
1977 "H_tau mismatch vs FD: rel={dh_rel:.3e}, num={dh_num:.3e}, den={dh_den:.3e}"
1978 );
1979
1980 let v_abs = (v_tau_analytic - v_taufd).abs();
1981 let v_rel = v_abs / v_taufd.abs().max(1e-10);
1982 assert_eq!(
1983 v_tau_analytic.signum(),
1984 v_taufd.signum(),
1985 "V_tau sign mismatch: analytic={v_tau_analytic:.6e}, fd={v_taufd:.6e}"
1986 );
1987 assert!(
1988 v_rel < 2e-2,
1989 "V_tau mismatch vs FD: rel={v_rel:.3e}, abs={v_abs:.3e}, analytic={v_tau_analytic:.6e}, fd={v_taufd:.6e}"
1990 );
1991
1992 assert!(
1993 cancellation.abs() < 1e-10,
1994 "stationarity cancellation failed: | -ell_beta^T B + beta^T S B | = {:.3e}",
1995 cancellation.abs()
1996 );
1997 }
1998
1999 #[test]
2000 pub(crate) fn firth_exacthessian_includes_analytic_tk_second_derivatives() {
2001 let y = array![0.0, 1.0, 0.0, 1.0, 1.0, 0.0, 1.0, 0.0, 1.0, 0.0];
2003 let w = Array1::<f64>::ones(y.len());
2004 let x = array![
2005 [1.0, -1.2, 0.4, -2.4],
2006 [1.0, -0.9, -0.1, -1.8],
2007 [1.0, -0.6, 0.3, -1.2],
2008 [1.0, -0.2, -0.4, -0.4],
2009 [1.0, 0.1, 0.5, 0.2],
2010 [1.0, 0.4, -0.6, 0.8],
2011 [1.0, 0.8, 0.2, 1.6],
2012 [1.0, 1.1, -0.3, 2.2],
2013 [1.0, 1.4, 0.7, 2.8],
2014 [1.0, 1.7, -0.2, 3.4],
2015 ];
2016 let s0 = array![
2017 [0.0, 0.0, 0.0, 0.0],
2018 [0.0, 1.5, 0.2, 0.0],
2019 [0.0, 0.2, 1.0, 0.0],
2020 [0.0, 0.0, 0.0, 0.5],
2021 ];
2022 let s1 = array![
2023 [0.0, 0.0, 0.0, 0.0],
2024 [0.0, 0.8, -0.1, 0.0],
2025 [0.0, -0.1, 0.6, 0.0],
2026 [0.0, 0.0, 0.0, 0.3],
2027 ];
2028 let offset = Array1::<f64>::zeros(y.len());
2029 let cfg =
2032 RemlConfig::external(binomial_logit_glm_spec(), 1e-9, true).with_max_iterations(500);
2033 let p = x.ncols();
2034 use crate::estimate::PenaltySpec;
2035 let specs = vec![PenaltySpec::Dense(s0), PenaltySpec::Dense(s1)];
2036 let canonical =
2037 gam_terms::construction::canonicalize_penalty_specs(&specs, &[1, 1], p, "test")
2038 .map(|(canonical, _)| canonical)
2039 .expect("canonicalize");
2040 let state = RemlState::newwith_offset(
2041 y.view(),
2042 x.clone(),
2043 w.view(),
2044 offset.view(),
2045 canonical,
2046 p,
2047 &cfg,
2048 Some(vec![1, 1]),
2049 None,
2050 None,
2051 )
2052 .expect("state");
2053 let rho = array![0.1, -0.2];
2054 assert!(
2055 state.analytic_outer_hessian_enabled(),
2056 "Firth logit should no longer disable analytic outer Hessian planning"
2057 );
2058 let outer = state
2059 .compute_outer_eval_with_order(
2060 &rho,
2061 crate::rho_optimizer::OuterEvalOrder::ValueGradientHessian,
2062 )
2063 .expect("outer Hessian eval should succeed");
2064 assert!(
2065 outer.hessian.is_analytic(),
2066 "outer planner should request and return an analytic Hessian"
2067 );
2068 let bundle = state.obtain_eval_bundle(&rho).expect("exact firth bundle");
2069 let h_dense = state
2070 .compute_lamlhessian_exact_from_bundle(&rho, &bundle)
2071 .expect("Firth exact Hessian should include analytic TK second derivatives");
2072 assert_eq!(h_dense.raw_dim(), ndarray::Ix2(2, 2));
2073 assert!(
2074 h_dense.iter().all(|value| value.is_finite()),
2075 "Hessian should be finite: {h_dense:?}"
2076 );
2077 }
2078
2079 #[test]
2080 pub(crate) fn firth_outer_hessian_matches_gradient_finite_difference_with_tk_terms() {
2081 let y = array![0.0, 1.0, 0.0, 1.0, 1.0, 0.0, 1.0, 0.0];
2082 let w = Array1::<f64>::ones(y.len());
2083 let x = array![
2084 [1.0, -1.0, 0.3],
2085 [1.0, -0.7, -0.2],
2086 [1.0, -0.3, 0.4],
2087 [1.0, 0.0, -0.5],
2088 [1.0, 0.2, 0.6],
2089 [1.0, 0.6, -0.4],
2090 [1.0, 0.9, 0.2],
2091 [1.0, 1.3, -0.1],
2092 ];
2093 let s0 = array![[0.0, 0.0, 0.0], [0.0, 1.2, 0.1], [0.0, 0.1, 0.7],];
2094 let s1 = array![[0.0, 0.0, 0.0], [0.0, 0.4, -0.05], [0.0, -0.05, 0.9],];
2095 let cfg =
2096 RemlConfig::external(binomial_logit_glm_spec(), 1e-9, true).with_max_iterations(500);
2097 let p_dim = x.ncols();
2098 use crate::estimate::PenaltySpec;
2099 let specs = vec![PenaltySpec::Dense(s0), PenaltySpec::Dense(s1)];
2100 let canonical =
2101 gam_terms::construction::canonicalize_penalty_specs(&specs, &[1, 1], p_dim, "test")
2102 .map(|(canonical, _)| canonical)
2103 .expect("canonicalize");
2104 let offset = Array1::<f64>::zeros(y.len());
2105 let state = RemlState::newwith_offset(
2106 y.view(),
2107 x.clone(),
2108 w.view(),
2109 offset.view(),
2110 canonical,
2111 p_dim,
2112 &cfg,
2113 Some(vec![1, 1]),
2114 None,
2115 None,
2116 )
2117 .expect("state");
2118 let rho = array![0.15, -0.25];
2119 let eval = state
2120 .compute_outer_eval_with_order(
2121 &rho,
2122 crate::rho_optimizer::OuterEvalOrder::ValueGradientHessian,
2123 )
2124 .expect("analytic Hessian eval");
2125 let h = match eval.hessian {
2126 HessianValue::Dense(hessian) => hessian,
2127 HessianValue::Operator(_) | HessianValue::Unavailable => {
2128 panic!("expected dense analytic Hessian")
2129 }
2130 };
2131 let delta = 2.0e-5;
2132 for col in 0..rho.len() {
2133 let mut rp = rho.clone();
2134 let mut rm = rho.clone();
2135 rp[col] += delta;
2136 rm[col] -= delta;
2137 let gp = state
2138 .compute_outer_eval_with_order(
2139 &rp,
2140 crate::rho_optimizer::OuterEvalOrder::ValueAndGradient,
2141 )
2142 .expect("plus grad")
2143 .gradient;
2144 let gm = state
2145 .compute_outer_eval_with_order(
2146 &rm,
2147 crate::rho_optimizer::OuterEvalOrder::ValueAndGradient,
2148 )
2149 .expect("minus grad")
2150 .gradient;
2151 for row in 0..rho.len() {
2152 let fd = (gp[row] - gm[row]) / (2.0 * delta);
2153 let an = h[[row, col]];
2154 let rel = (fd - an).abs() / fd.abs().max(an.abs()).max(1e-6);
2155 assert!(
2156 rel < 2.0e-3,
2157 "Hessian mismatch ({row},{col}): analytic={an:.9e}, fd={fd:.9e}, rel={rel:.3e}"
2158 );
2159 }
2160 }
2161 }
2162
2163 #[test]
2164 pub(crate) fn firthgradient_lives_in_design_column_space_under_rank_deficiency() {
2165 let x = array![
2167 [1.0, -1.2, 0.4, -2.4],
2168 [1.0, -0.9, -0.1, -1.8],
2169 [1.0, -0.6, 0.3, -1.2],
2170 [1.0, -0.2, -0.4, -0.4],
2171 [1.0, 0.1, 0.5, 0.2],
2172 [1.0, 0.4, -0.6, 0.8],
2173 [1.0, 0.8, 0.2, 1.6],
2174 [1.0, 1.1, -0.3, 2.2],
2175 ];
2176 let beta = array![0.1, -0.2, 0.3, 0.05];
2177 let eta = x.dot(&beta);
2178 let op = super::RemlState::build_firth_dense_operator_for_link(
2179 &gam_problem::InverseLink::Standard(gam_problem::StandardLink::Logit),
2180 &x,
2181 &eta,
2182 ndarray::Array1::ones(x.nrows()).view(),
2183 )
2184 .expect("firth operator");
2185
2186 let gradphi = 0.5 * x.t().dot(&(&op.w1 * &op.h_diag));
2189
2190 let q = &op.q_basis;
2192 let proj = q.dot(&q.t().dot(&gradphi));
2193 let resid = &gradphi - &proj;
2194 let rel =
2195 resid.mapv(|v| v * v).sum().sqrt() / gradphi.mapv(|v| v * v).sum().sqrt().max(1e-12);
2196 assert!(
2197 rel < 1e-10,
2198 "Firth gradient should lie in Col(Xᵀ): rel residual={rel:.3e}"
2199 );
2200 }
2201
2202 #[test]
2203 pub(crate) fn firth_logit_directional_hypergradient_accepts_penalty_only_with_full_tk_gradient()
2204 {
2205 let y = array![0.0, 1.0, 0.0, 1.0, 0.0, 1.0];
2206 let w = Array1::<f64>::ones(y.len());
2207 let x = array![
2208 [1.0, -1.1, 0.2],
2209 [1.0, -0.6, -0.3],
2210 [1.0, -0.1, 0.5],
2211 [1.0, 0.3, -0.7],
2212 [1.0, 0.8, 0.1],
2213 [1.0, 1.2, -0.4],
2214 ];
2215 let s0 = array![[0.0, 0.0, 0.0], [0.0, 1.0, 0.1], [0.0, 0.1, 0.8],];
2216 let hyper = DirectionalHyperParam::single_penalty(
2217 0,
2218 Array2::<f64>::zeros((x.nrows(), x.ncols())),
2219 array![[0.0, 0.0, 0.0], [0.0, 0.2, 0.03], [0.0, 0.03, 0.12],],
2220 None,
2221 None,
2222 )
2223 .expect("single-penalty hyper direction");
2224 let rho = array![0.0];
2225 let cfg = RemlConfig::external(binomial_logit_glm_spec(), 1e-8, true);
2226 let state = build_logit_state(&y, &w, &x, &s0, &cfg);
2227 let gradient = single_directional_tau_gradient(&state, &rho, hyper)
2228 .expect("Firth penalty-only directional gradient should use analytic TK propagation");
2229 assert!(gradient.is_finite(), "gradient={gradient}");
2230 let fd = fd_directional_tau_cost_gradient(
2231 &y,
2232 &w,
2233 &x,
2234 &s0,
2235 &cfg,
2236 &rho,
2237 &Array2::<f64>::zeros((x.nrows(), x.ncols())),
2238 &array![[0.0, 0.0, 0.0], [0.0, 0.2, 0.03], [0.0, 0.03, 0.12],],
2239 );
2240 let rel = (gradient - fd).abs() / gradient.abs().max(fd.abs()).max(1.0e-10);
2241 assert!(
2242 rel < 1.0e-3,
2243 "Firth penalty-only directional gradient mismatch: analytic={gradient:.12e}, fd={fd:.12e}, rel={rel:.3e}"
2244 );
2245
2246 let efs_hyper = DirectionalHyperParam::single_penalty(
2247 0,
2248 Array2::<f64>::zeros((x.nrows(), x.ncols())),
2249 array![[0.0, 0.0, 0.0], [0.0, 0.2, 0.03], [0.0, 0.03, 0.12],],
2250 None,
2251 None,
2252 )
2253 .expect("single-penalty EFS hyper direction");
2254 let efs = state
2255 .compute_efs_steps_with_psi_ext(&rho, &[efs_hyper])
2256 .expect("Firth penalty-only EFS should use analytic TK propagation");
2257 assert!(efs.cost.is_finite(), "efs cost={}", efs.cost);
2258 }
2259
2260 #[test]
2280 pub(crate) fn firth_logit_rho_gradient_matches_finite_difference_through_inner_solve() {
2281 let x = array![[1.0, -6.0], [1.0, 0.2], [1.0, 5.8]];
2285 let y = array![0.0, 0.0, 1.0];
2286 let w = Array1::<f64>::ones(y.len());
2287 let s0 = array![[1.0, 0.0], [0.0, 1.0]];
2289 let cfg = RemlConfig::external(binomial_logit_glm_spec(), 1e-12, true);
2292 let state = build_logit_state(&y, &w, &x, &s0, &cfg);
2293 let delta = 1e-4_f64;
2294 for &rho in &[-0.6_f64, -0.3, 0.0, 0.3, 0.6] {
2295 let r = array![rho];
2296 let analytic = state
2297 .compute_gradient(&r)
2298 .expect("Firth LAML ρ-gradient should evaluate")[0];
2299 let cost_plus = state
2300 .compute_cost(&array![rho + delta])
2301 .expect("Firth LAML cost(ρ+δ) should evaluate");
2302 let cost_minus = state
2303 .compute_cost(&array![rho - delta])
2304 .expect("Firth LAML cost(ρ−δ) should evaluate");
2305 let fd = (cost_plus - cost_minus) / (2.0 * delta);
2306 let rel = (fd - analytic).abs() / fd.abs().max(1e-3);
2307 assert!(
2308 analytic.is_finite() && fd.is_finite(),
2309 "non-finite Firth ρ-gradient at rho={rho:+.3}: fd={fd:+.6e}, analytic={analytic:+.6e}"
2310 );
2311 assert!(
2312 rel < 1e-4,
2313 "Firth ρ-gradient FD desync at rho={rho:+.3}: fd={fd:+.6e}, analytic={analytic:+.6e}, rel={rel:.3e} (>= 1e-4). \
2314 The inner P-IRLS likely converged off the Firth-KKT mode (gam#1821)."
2315 );
2316 }
2317 }
2318
2319 #[test]
2320 pub(crate) fn firth_logit_directional_hypergradient_accepts_design_moving_with_full_tk_gradient()
2321 {
2322 let y = array![0.0, 1.0, 0.0, 1.0, 0.0, 1.0];
2323 let w = Array1::<f64>::ones(y.len());
2324 let x = array![
2325 [1.0, -1.1, 0.2],
2326 [1.0, -0.6, -0.3],
2327 [1.0, -0.1, 0.5],
2328 [1.0, 0.3, -0.7],
2329 [1.0, 0.8, 0.1],
2330 [1.0, 1.2, -0.4],
2331 ];
2332 let s0 = array![[0.0, 0.0, 0.0], [0.0, 1.0, 0.1], [0.0, 0.1, 0.8],];
2333 let hyper = DirectionalHyperParam::single_penalty(
2334 0,
2335 Array2::from_elem((x.nrows(), x.ncols()), 1e-3),
2336 Array2::<f64>::zeros((x.ncols(), x.ncols())),
2337 None,
2338 None,
2339 )
2340 .expect("single-penalty hyper direction");
2341 let rho = array![0.0];
2342 let cfg = RemlConfig::external(binomial_logit_glm_spec(), 1e-8, true);
2343 let state = build_logit_state(&y, &w, &x, &s0, &cfg);
2344 let gradient = single_directional_tau_gradient(&state, &rho, hyper)
2345 .expect("Firth design-moving directional gradient should use analytic TK propagation");
2346 assert!(gradient.is_finite(), "gradient={gradient}");
2347 let x_tau = Array2::from_elem((x.nrows(), x.ncols()), 1e-3);
2348 let s_tau = Array2::<f64>::zeros((x.ncols(), x.ncols()));
2349 let fd = fd_directional_tau_cost_gradient(&y, &w, &x, &s0, &cfg, &rho, &x_tau, &s_tau);
2350 let rel = (gradient - fd).abs() / gradient.abs().max(fd.abs()).max(1.0e-10);
2351 assert!(
2352 rel < 2.0e-2,
2353 "Firth design-moving directional gradient mismatch: analytic={gradient:.12e}, fd={fd:.12e}, rel={rel:.3e}"
2354 );
2355 }
2356
2357 #[test]
2358 pub(crate) fn firth_logit_hybrid_efs_accepts_full_tk_psi_gradient() {
2359 let y = array![0.0, 1.0, 0.0, 1.0, 0.0, 1.0];
2360 let w = Array1::<f64>::ones(y.len());
2361 let x = array![
2362 [1.0, -1.1, 0.2],
2363 [1.0, -0.6, -0.3],
2364 [1.0, -0.1, 0.5],
2365 [1.0, 0.3, -0.7],
2366 [1.0, 0.8, 0.1],
2367 [1.0, 1.2, -0.4],
2368 ];
2369 let s0 = array![[0.0, 0.0, 0.0], [0.0, 1.0, 0.1], [0.0, 0.1, 0.8],];
2370 let hyper_dirs = vec![
2371 DirectionalHyperParam::single_penalty(
2372 0,
2373 Array2::from_shape_fn((x.nrows(), x.ncols()), |(i, j)| {
2374 1e-3 * ((i + 1) as f64) * ((j + 2) as f64)
2375 }),
2376 Array2::<f64>::zeros((x.ncols(), x.ncols())),
2377 None,
2378 None,
2379 )
2380 .expect("design-moving hyper direction"),
2381 ];
2382 let rho = array![0.0];
2383 let cfg = RemlConfig::external(binomial_logit_glm_spec(), 1e-8, true);
2384 let state = build_logit_state(&y, &w, &x, &s0, &cfg);
2385
2386 let full = state
2387 .evaluate_unified_with_psi_ext(
2388 &rho,
2389 None,
2390 crate::estimate::reml::reml_outer_engine::EvalMode::ValueAndGradient,
2391 &hyper_dirs,
2392 )
2393 .expect("full Firth psi gradient should use analytic TK propagation");
2394 assert!(full.cost.is_finite(), "full cost={}", full.cost);
2395 let full_grad = full.gradient.expect("gradient should be present");
2396 assert!(
2397 full_grad.iter().all(|value| value.is_finite()),
2398 "full gradient={full_grad:?}"
2399 );
2400
2401 let efs = state
2402 .compute_efs_steps_with_psi_ext(&rho, &hyper_dirs)
2403 .expect("hybrid EFS should use analytic TK propagation");
2404 assert!(efs.cost.is_finite(), "efs cost={}", efs.cost);
2405 }
2406
2407 #[test]
2408 pub(crate) fn joint_hyperhessianwires_mixed_blocks() {
2409 let y = array![0.0, 1.0, 0.0, 1.0, 0.0, 1.0, 1.0, 0.0];
2410 let w = Array1::<f64>::ones(y.len());
2411 let x = array![
2412 [1.0, -1.2, 0.3],
2413 [1.0, -0.8, -0.4],
2414 [1.0, -0.3, 0.7],
2415 [1.0, 0.1, -0.9],
2416 [1.0, 0.5, 0.2],
2417 [1.0, 0.9, -0.1],
2418 [1.0, 1.3, 0.8],
2419 [1.0, 1.7, -0.6],
2420 ];
2421 let s0 = array![[0.0, 0.0, 0.0], [0.0, 1.2, 0.2], [0.0, 0.2, 0.9],];
2422 let cfg =
2423 RemlConfig::external(binomial_logit_glm_spec(), 1e-10, false).with_max_iterations(500);
2424 let state = build_logit_state(&y, &w, &x, &s0, &cfg);
2425 let rho = array![0.0];
2426 let theta = array![0.0, 0.0, 0.0];
2427 let hyper_dirs = vec![
2428 DirectionalHyperParam::single_penalty(
2429 0,
2430 Array2::<f64>::zeros((x.nrows(), x.ncols())),
2431 array![[0.0, 0.0, 0.0], [0.0, 0.2, 0.01], [0.0, 0.01, 0.15],],
2432 None,
2433 None,
2434 )
2435 .expect("single-penalty hyper direction"),
2436 DirectionalHyperParam::single_penalty(
2437 0,
2438 Array2::from_elem((x.nrows(), x.ncols()), 2e-4),
2439 Array2::<f64>::zeros((x.ncols(), x.ncols())),
2440 None,
2441 None,
2442 )
2443 .expect("single-penalty hyper direction"),
2444 ];
2445
2446 let (_, _, h) =
2447 compute_joint_hypercostgradienthessian(&state, &theta, rho.len(), &hyper_dirs)
2448 .expect("joint hyper cost+gradient+hessian");
2449 assert_eq!(h.nrows(), theta.len());
2450 assert_eq!(h.ncols(), theta.len());
2451 assert!(h.iter().all(|v| v.is_finite()));
2452 for i in 0..h.nrows() {
2453 for j in 0..i {
2454 let diff = (h[[i, j]] - h[[j, i]]).abs();
2455 assert!(
2456 diff < 1e-6,
2457 "joint hessian asymmetry at ({i},{j}): {diff:.3e}"
2458 );
2459 }
2460 }
2461 let mixed_0 = h[[0, 1]];
2463 let mixed_1 = h[[0, 2]];
2464 assert!(
2465 mixed_0.is_finite() && mixed_1.is_finite(),
2466 "mixed blocks must be finite"
2467 );
2468 }
2469
2470 #[test]
2471 pub(crate) fn joint_tau_tau_linear_dirs_matchfd_reference_away_fromzero_psi() {
2472 let y = array![0.0, 1.0, 0.0, 1.0, 0.0, 1.0, 1.0, 0.0];
2473 let w = Array1::<f64>::ones(y.len());
2474 let x = array![
2475 [1.0, -1.2, 0.3],
2476 [1.0, -0.8, -0.4],
2477 [1.0, -0.3, 0.7],
2478 [1.0, 0.1, -0.9],
2479 [1.0, 0.5, 0.2],
2480 [1.0, 0.9, -0.1],
2481 [1.0, 1.3, 0.8],
2482 [1.0, 1.7, -0.6],
2483 ];
2484 let s0 = array![[0.0, 0.0, 0.0], [0.0, 1.2, 0.2], [0.0, 0.2, 0.9],];
2485 let cfg =
2486 RemlConfig::external(binomial_logit_glm_spec(), 1e-10, false).with_max_iterations(500);
2487 let state = build_logit_state(&y, &w, &x, &s0, &cfg);
2488 let rho = array![0.0];
2489 let psi = array![0.7, -0.4];
2490 let theta = array![rho[0], psi[0], psi[1]];
2491 let hyper_dirs = vec![
2492 DirectionalHyperParam::single_penalty(
2493 0,
2494 Array2::<f64>::zeros((x.nrows(), x.ncols())),
2495 array![[0.0, 0.0, 0.0], [0.0, 0.2, 0.01], [0.0, 0.01, 0.15],],
2496 None,
2497 None,
2498 )
2499 .expect("linear tau direction"),
2500 DirectionalHyperParam::single_penalty(
2501 0,
2502 Array2::from_elem((x.nrows(), x.ncols()), 2e-4),
2503 Array2::<f64>::zeros((x.ncols(), x.ncols())),
2504 None,
2505 None,
2506 )
2507 .expect("linear tau direction"),
2508 ];
2509
2510 let (_, _, h_full) =
2511 compute_joint_hypercostgradienthessian(&state, &theta, rho.len(), &hyper_dirs)
2512 .expect("joint hyper cost+gradient+hessian");
2513 let h_tt_analytic = h_full.slice(s![rho.len().., rho.len()..]).to_owned();
2514
2515 let x_tau_mats: Vec<Array2<f64>> = vec![
2520 Array2::<f64>::zeros((x.nrows(), x.ncols())),
2521 Array2::from_elem((x.nrows(), x.ncols()), 2e-4),
2522 ];
2523 let s_tau_mats: Vec<Array2<f64>> = vec![
2524 array![[0.0, 0.0, 0.0], [0.0, 0.2, 0.01], [0.0, 0.01, 0.15]],
2525 Array2::<f64>::zeros((x.ncols(), x.ncols())),
2526 ];
2527
2528 let h_ttfd = directional_tau_hessian_fd_reference(
2529 &y,
2530 &w,
2531 &x,
2532 &s0,
2533 &cfg,
2534 &rho,
2535 &hyper_dirs,
2536 &x_tau_mats,
2537 &s_tau_mats,
2538 );
2539
2540 let num = (&h_tt_analytic - &h_ttfd)
2541 .iter()
2542 .map(|v| v * v)
2543 .sum::<f64>()
2544 .sqrt();
2545 let den = h_ttfd.iter().map(|v| v * v).sum::<f64>().sqrt().max(1e-10);
2546 let rel = num / den;
2547 assert!(
2548 rel < 1e-4,
2549 "linear-dir joint tau-tau block deviates from FD reference away from zero psi: rel={rel:.3e}, analytic={h_tt_analytic:?}, fd={h_ttfd:?}"
2550 );
2551 }
2552
2553 #[test]
2554 pub(crate) fn joint_hypervalidation_rejects_out_of_boundssecond_order_penalty_index() {
2555 let y = array![0.0, 1.0, 0.0, 1.0];
2572 let w = Array1::<f64>::ones(y.len());
2573 let x = array![
2574 [1.0, -0.5, 0.2],
2575 [1.0, -0.1, -0.3],
2576 [1.0, 0.4, 0.6],
2577 [1.0, 0.9, -0.2],
2578 ];
2579 let s0 = array![[0.0, 0.0, 0.0], [0.0, 1.0, 0.1], [0.0, 0.1, 0.8],];
2580 let cfg = RemlConfig::external(binomial_logit_glm_spec(), 1e-10, true);
2581 let state = build_logit_state(&y, &w, &x, &s0, &cfg);
2582 let theta = array![0.0, 0.0];
2583 let hyper_dirs = vec![
2584 DirectionalHyperParam::new(
2585 Array2::<f64>::zeros((x.nrows(), x.ncols())),
2586 vec![(0, Array2::<f64>::zeros((x.ncols(), x.ncols())))],
2587 None,
2588 Some(vec![Some(vec![(1, Array2::<f64>::eye(x.ncols()))])]),
2589 )
2590 .expect("hyper direction with invalid second-order penalty index"),
2591 ];
2592
2593 let msg = match compute_joint_hypercostgradienthessian(&state, &theta, 1, &hyper_dirs) {
2594 Ok(_) => panic!("invalid second-order penalty index should be rejected"),
2595 Err(err) => err.to_string(),
2596 };
2597 assert!(
2598 msg.contains("out of bounds") || msg.contains("penalty_index"),
2599 "unexpected validation error: {msg}"
2600 );
2601 }
2602
2603 #[test]
2604 pub(crate) fn joint_tau_tau_analytic_matchesfd_reference() {
2605 let y = array![0.0, 1.0, 0.0, 1.0, 0.0, 1.0, 1.0, 0.0];
2606 let w = Array1::<f64>::ones(y.len());
2607 let x = array![
2608 [1.0, -1.2, 0.3],
2609 [1.0, -0.8, -0.4],
2610 [1.0, -0.3, 0.7],
2611 [1.0, 0.1, -0.9],
2612 [1.0, 0.5, 0.2],
2613 [1.0, 0.9, -0.1],
2614 [1.0, 1.3, 0.8],
2615 [1.0, 1.7, -0.6],
2616 ];
2617 let s0 = array![[0.0, 0.0, 0.0], [0.0, 1.2, 0.2], [0.0, 0.2, 0.9],];
2618 let cfg =
2619 RemlConfig::external(binomial_logit_glm_spec(), 1e-10, false).with_max_iterations(500);
2620 let state = build_logit_state(&y, &w, &x, &s0, &cfg);
2621 let rho = array![0.0];
2622 let psi = array![0.0, 0.0];
2623 let hyper_dirs = vec![
2624 DirectionalHyperParam::single_penalty(
2625 0,
2626 Array2::<f64>::zeros((x.nrows(), x.ncols())),
2627 array![[0.0, 0.0, 0.0], [0.0, 0.2, 0.01], [0.0, 0.01, 0.15],],
2628 None,
2629 None,
2630 )
2631 .expect("single-penalty hyper direction"),
2632 DirectionalHyperParam::single_penalty(
2633 0,
2634 Array2::from_elem((x.nrows(), x.ncols()), 2e-4),
2635 Array2::<f64>::zeros((x.ncols(), x.ncols())),
2636 None,
2637 None,
2638 )
2639 .expect("single-penalty hyper direction"),
2640 ];
2641
2642 let theta = {
2643 let mut t = Array1::<f64>::zeros(rho.len() + psi.len());
2644 t.slice_mut(s![..rho.len()]).assign(&rho);
2645 t.slice_mut(s![rho.len()..]).assign(&psi);
2646 t
2647 };
2648 let (_, _, h_full) =
2649 compute_joint_hypercostgradienthessian(&state, &theta, rho.len(), &hyper_dirs)
2650 .expect("joint hyper cost+gradient+hessian");
2651 let h_tt_analytic = h_full.slice(s![rho.len().., rho.len()..]).to_owned();
2652 assert_eq!(h_tt_analytic.nrows(), hyper_dirs.len());
2653 assert_eq!(h_tt_analytic.ncols(), hyper_dirs.len());
2654
2655 let x_tau_mats: Vec<Array2<f64>> = vec![
2660 Array2::<f64>::zeros((x.nrows(), x.ncols())),
2661 Array2::from_elem((x.nrows(), x.ncols()), 2e-4),
2662 ];
2663 let s_tau_mats: Vec<Array2<f64>> = vec![
2664 array![[0.0, 0.0, 0.0], [0.0, 0.2, 0.01], [0.0, 0.01, 0.15]],
2665 Array2::<f64>::zeros((x.ncols(), x.ncols())),
2666 ];
2667
2668 let h_ttfd = directional_tau_hessian_fd_reference(
2669 &y,
2670 &w,
2671 &x,
2672 &s0,
2673 &cfg,
2674 &rho,
2675 &hyper_dirs,
2676 &x_tau_mats,
2677 &s_tau_mats,
2678 );
2679
2680 let num = (&h_tt_analytic - &h_ttfd)
2681 .iter()
2682 .map(|v| v * v)
2683 .sum::<f64>()
2684 .sqrt();
2685 let den = h_ttfd.iter().map(|v| v * v).sum::<f64>().sqrt().max(1e-10);
2686 let rel = num / den;
2687 assert!(
2688 rel < 1e-4,
2689 "analytic tau-tau block deviates from FD reference: rel={rel:.3e}, analytic={h_tt_analytic:?}, fd={h_ttfd:?}"
2690 );
2691 }
2692
2693 pub(crate) struct GaussianRemlFixture {
2703 pub(crate) y: Array1<f64>,
2704 pub(crate) w: Array1<f64>,
2705 pub(crate) x: Array2<f64>,
2706 pub(crate) s0: Array2<f64>,
2707 pub(crate) cfg: RemlConfig,
2708 pub(crate) rho: Array1<f64>,
2709 pub(crate) x_tau_design: Array2<f64>,
2711 pub(crate) s_tau_penalty: Array2<f64>,
2713 }
2714
2715 impl GaussianRemlFixture {
2716 pub(crate) fn new() -> Self {
2717 let y = array![0.5, 1.2, -0.3, 0.8, 1.1, -0.6, 0.9, 0.1, -0.2, 0.7];
2718 let x = array![
2719 [1.0, -1.2, 0.3],
2720 [1.0, -0.8, -0.4],
2721 [1.0, -0.3, 0.7],
2722 [1.0, 0.1, -0.9],
2723 [1.0, 0.5, 0.2],
2724 [1.0, 0.9, -0.1],
2725 [1.0, 1.3, 0.8],
2726 [1.0, 1.7, -0.6],
2727 [1.0, -0.5, 0.5],
2728 [1.0, 0.3, -0.3],
2729 ];
2730 Self {
2731 w: Array1::<f64>::ones(y.len()),
2732 y,
2733 x: x.clone(),
2734 s0: array![[0.0, 0.0, 0.0], [0.0, 1.2, 0.2], [0.0, 0.2, 0.9]],
2735 cfg: RemlConfig::external(gaussian_identity_glm_spec(), 1e-14, false),
2736 rho: array![0.0],
2737 x_tau_design: array![
2738 [0.0, 1e-3, -2e-3],
2739 [0.0, -3e-3, 1e-3],
2740 [0.0, 2e-3, 0.5e-3],
2741 [0.0, -1e-3, 3e-3],
2742 [0.0, 0.5e-3, -1e-3],
2743 [0.0, 1.5e-3, 2e-3],
2744 [0.0, -2e-3, -0.5e-3],
2745 [0.0, 3e-3, 1e-3],
2746 [0.0, -0.5e-3, 2e-3],
2747 [0.0, 1e-3, -1.5e-3],
2748 ],
2749 s_tau_penalty: array![[0.0, 0.0, 0.0], [0.0, 0.25, 0.04], [0.0, 0.04, 0.15]],
2750 }
2751 }
2752 }
2753
2754 impl LogitDesignMotionFixture for GaussianRemlFixture {
2755 fn y(&self) -> &Array1<f64> {
2756 &self.y
2757 }
2758 fn w(&self) -> &Array1<f64> {
2759 &self.w
2760 }
2761 fn x(&self) -> &Array2<f64> {
2762 &self.x
2763 }
2764 fn s0(&self) -> &Array2<f64> {
2765 &self.s0
2766 }
2767 fn cfg(&self) -> &RemlConfig {
2768 &self.cfg
2769 }
2770 fn rho(&self) -> &Array1<f64> {
2771 &self.rho
2772 }
2773 }
2774
2775 #[test]
2776 pub(crate) fn profiled_gaussian_design_moving_gradient_matches_fd() {
2777 let f = GaussianRemlFixture::new();
2778 let state = f.state();
2779 let s_tau = Array2::<f64>::zeros((3, 3));
2780 let hyper = DirectionalHyperParam::single_penalty(
2781 0,
2782 f.x_tau_design.clone(),
2783 s_tau.clone(),
2784 None,
2785 None,
2786 )
2787 .expect("design-moving hyper direction");
2788
2789 let v_tau_analytic = single_directional_tau_gradient(&state, &f.rho, hyper)
2790 .expect("analytic directional gradient");
2791 let v_taufd = f.fd_directional_gradient(&f.x_tau_design, &s_tau);
2792
2793 let v_rel = (v_tau_analytic - v_taufd).abs() / v_taufd.abs().max(1e-10);
2794 assert!(
2795 v_rel < 1e-3,
2796 "Gaussian REML design-moving V_tau mismatch: rel={v_rel:.3e}, \
2797 analytic={v_tau_analytic:.6e}, fd={v_taufd:.6e}"
2798 );
2799 }
2800
2801 #[test]
2802 pub(crate) fn profiled_gaussian_penalty_only_gradient_matches_fd() {
2803 let f = GaussianRemlFixture::new();
2804 let state = f.state();
2805 let x_tau = Array2::<f64>::zeros(f.x.raw_dim());
2806 let hyper = DirectionalHyperParam::single_penalty(
2807 0,
2808 x_tau.clone(),
2809 f.s_tau_penalty.clone(),
2810 None,
2811 None,
2812 )
2813 .expect("penalty-only hyper direction");
2814
2815 let v_tau_analytic = single_directional_tau_gradient(&state, &f.rho, hyper)
2816 .expect("analytic directional gradient");
2817 let v_taufd = f.fd_directional_gradient(&x_tau, &f.s_tau_penalty);
2818
2819 let v_rel = (v_tau_analytic - v_taufd).abs() / v_taufd.abs().max(1e-10);
2820 assert!(
2821 v_rel < 1e-3,
2822 "Gaussian REML penalty-only V_tau mismatch: rel={v_rel:.3e}, \
2823 analytic={v_tau_analytic:.6e}, fd={v_taufd:.6e}"
2824 );
2825 }
2826
2827 #[test]
2828 pub(crate) fn profiled_gaussian_joint_hessian_matches_fd() {
2829 let f = GaussianRemlFixture::new();
2832 let x_tau_0 = Array2::<f64>::zeros(f.x.raw_dim());
2833 let s_tau_0 = f.s_tau_penalty.clone();
2834 let x_tau_1 = f.x_tau_design.clone();
2835 let s_tau_1 = Array2::<f64>::zeros((3, 3));
2836
2837 let hyper_dirs = vec![
2838 DirectionalHyperParam::single_penalty(0, x_tau_0.clone(), s_tau_0.clone(), None, None)
2839 .expect("penalty-only direction"),
2840 DirectionalHyperParam::single_penalty(0, x_tau_1.clone(), s_tau_1.clone(), None, None)
2841 .expect("design-moving direction"),
2842 ];
2843
2844 let state = f.state();
2845 let mut theta = Array1::<f64>::zeros(f.rho.len() + hyper_dirs.len());
2846 theta.slice_mut(s![..f.rho.len()]).assign(&f.rho);
2847 let (_, _, h_full) =
2848 compute_joint_hypercostgradienthessian(&state, &theta, f.rho.len(), &hyper_dirs)
2849 .expect("joint cost+gradient+hessian");
2850 let h_tt_analytic = h_full.slice(s![f.rho.len().., f.rho.len()..]).to_owned();
2851
2852 let x_tau_mats = vec![x_tau_0.clone(), x_tau_1.clone()];
2855 let s_tau_mats = vec![s_tau_0.clone(), s_tau_1.clone()];
2856 let h_ttfd = directional_tau_hessian_fd_reference(
2857 &f.y,
2858 &f.w,
2859 &f.x,
2860 &f.s0,
2861 &f.cfg,
2862 &f.rho,
2863 &hyper_dirs,
2864 &x_tau_mats,
2865 &s_tau_mats,
2866 );
2867
2868 let num = (&h_tt_analytic - &h_ttfd)
2869 .iter()
2870 .map(|v| v * v)
2871 .sum::<f64>()
2872 .sqrt();
2873 let den = h_ttfd.iter().map(|v| v * v).sum::<f64>().sqrt().max(1e-10);
2874 let rel = num / den;
2875 assert!(
2876 rel < 1e-4,
2877 "Gaussian REML tau-tau Hessian mismatch: rel={rel:.3e}, \
2878 analytic={h_tt_analytic:?}, fd={h_ttfd:?}"
2879 );
2880 }
2881
2882 #[test]
2896 pub(crate) fn logit_design_moving_gradient_matches_fd() {
2897 let y = array![0.0, 1.0, 0.0, 1.0, 0.0, 1.0, 1.0, 0.0, 1.0, 0.0];
2898 let w = Array1::<f64>::ones(y.len());
2899 let x = array![
2900 [1.0, -1.2, 0.3],
2901 [1.0, -0.8, -0.4],
2902 [1.0, -0.3, 0.7],
2903 [1.0, 0.1, -0.9],
2904 [1.0, 0.5, 0.2],
2905 [1.0, 0.9, -0.1],
2906 [1.0, 1.3, 0.8],
2907 [1.0, 1.7, -0.6],
2908 [1.0, -0.5, 0.5],
2909 [1.0, 0.3, -0.3],
2910 ];
2911 let s0 = array![[0.0, 0.0, 0.0], [0.0, 1.2, 0.2], [0.0, 0.2, 0.9]];
2912 let cfg = RemlConfig::external(binomial_logit_glm_spec(), 1e-14, false);
2913 let state = build_logit_state(&y, &w, &x, &s0, &cfg);
2914 let rho = array![0.0];
2915
2916 let x_tau = array![
2918 [0.0, 1e-3, -2e-3],
2919 [0.0, -3e-3, 1e-3],
2920 [0.0, 2e-3, 0.5e-3],
2921 [0.0, -1e-3, 3e-3],
2922 [0.0, 0.5e-3, -1e-3],
2923 [0.0, 1.5e-3, 2e-3],
2924 [0.0, -2e-3, -0.5e-3],
2925 [0.0, 3e-3, 1e-3],
2926 [0.0, -0.5e-3, 2e-3],
2927 [0.0, 1e-3, -1.5e-3],
2928 ];
2929 let s_tau = Array2::<f64>::zeros((3, 3));
2930 let hyper =
2931 DirectionalHyperParam::single_penalty(0, x_tau.clone(), s_tau.clone(), None, None)
2932 .expect("design-moving hyper direction");
2933
2934 let v_tau_analytic = single_directional_tau_gradient(&state, &rho, hyper)
2935 .expect("analytic directional gradient");
2936
2937 let h = 2e-5;
2938 let x_plus = &x + &x_tau.mapv(|v| h * v);
2939 let x_minus = &x - &x_tau.mapv(|v| h * v);
2940 let state_plus = build_logit_state(&y, &w, &x_plus, &s0, &cfg);
2941 let state_minus = build_logit_state(&y, &w, &x_minus, &s0, &cfg);
2942 let v_plus = state_plus.compute_cost(&rho).expect("cost+");
2943 let v_minus = state_minus.compute_cost(&rho).expect("cost-");
2944 let v_taufd = (v_plus - v_minus) / (2.0 * h);
2945
2946 let v_rel = (v_tau_analytic - v_taufd).abs() / v_taufd.abs().max(1e-10);
2947 assert!(
2948 v_rel < 1e-3,
2949 "Logit REML design-moving V_tau mismatch: rel={v_rel:.3e}, \
2950 analytic={v_tau_analytic:.6e}, fd={v_taufd:.6e}"
2951 );
2952 }
2953
2954 #[test]
2955 pub(crate) fn logit_design_moving_hessian_matches_fd() {
2956 let y = array![0.0, 1.0, 0.0, 1.0, 0.0, 1.0, 1.0, 0.0, 1.0, 0.0];
2961 let w = Array1::<f64>::ones(y.len());
2962 let x = array![
2963 [1.0, -1.2, 0.3],
2964 [1.0, -0.8, -0.4],
2965 [1.0, -0.3, 0.7],
2966 [1.0, 0.1, -0.9],
2967 [1.0, 0.5, 0.2],
2968 [1.0, 0.9, -0.1],
2969 [1.0, 1.3, 0.8],
2970 [1.0, 1.7, -0.6],
2971 [1.0, -0.5, 0.5],
2972 [1.0, 0.3, -0.3],
2973 ];
2974 let s0 = array![[0.0, 0.0, 0.0], [0.0, 1.2, 0.2], [0.0, 0.2, 0.9]];
2975 let cfg = RemlConfig::external(binomial_logit_glm_spec(), 1e-14, false);
2976 let rho = array![0.0];
2977
2978 let x_tau_0 = Array2::<f64>::zeros(x.raw_dim());
2980 let s_tau_0 = array![[0.0, 0.0, 0.0], [0.0, 0.25, 0.04], [0.0, 0.04, 0.15]];
2981 let x_tau_1 = array![
2982 [0.0, 1e-3, -2e-3],
2983 [0.0, -3e-3, 1e-3],
2984 [0.0, 2e-3, 0.5e-3],
2985 [0.0, -1e-3, 3e-3],
2986 [0.0, 0.5e-3, -1e-3],
2987 [0.0, 1.5e-3, 2e-3],
2988 [0.0, -2e-3, -0.5e-3],
2989 [0.0, 3e-3, 1e-3],
2990 [0.0, -0.5e-3, 2e-3],
2991 [0.0, 1e-3, -1.5e-3],
2992 ];
2993 let s_tau_1 = Array2::<f64>::zeros((3, 3));
2994
2995 let hyper_dirs = vec![
2996 DirectionalHyperParam::single_penalty(0, x_tau_0.clone(), s_tau_0.clone(), None, None)
2997 .expect("penalty-only direction"),
2998 DirectionalHyperParam::single_penalty(0, x_tau_1.clone(), s_tau_1.clone(), None, None)
2999 .expect("design-moving direction"),
3000 ];
3001
3002 let state = build_logit_state(&y, &w, &x, &s0, &cfg);
3003 let mut theta = Array1::<f64>::zeros(rho.len() + hyper_dirs.len());
3004 theta.slice_mut(s![..rho.len()]).assign(&rho);
3005 let (_, _, h_full) =
3006 compute_joint_hypercostgradienthessian(&state, &theta, rho.len(), &hyper_dirs)
3007 .expect("joint cost+gradient+hessian");
3008 let h_tt_analytic = h_full.slice(s![rho.len().., rho.len()..]).to_owned();
3009
3010 let x_tau_mats = vec![x_tau_0.clone(), x_tau_1.clone()];
3011 let s_tau_mats = vec![s_tau_0.clone(), s_tau_1.clone()];
3012 let h_ttfd = directional_tau_hessian_fd_reference(
3013 &y,
3014 &w,
3015 &x,
3016 &s0,
3017 &cfg,
3018 &rho,
3019 &hyper_dirs,
3020 &x_tau_mats,
3021 &s_tau_mats,
3022 );
3023
3024 let num = (&h_tt_analytic - &h_ttfd)
3025 .iter()
3026 .map(|v| v * v)
3027 .sum::<f64>()
3028 .sqrt();
3029 let den = h_ttfd.iter().map(|v| v * v).sum::<f64>().sqrt().max(1e-10);
3030 let rel = num / den;
3031 assert!(
3032 rel < 1e-4,
3033 "Logit REML design-moving tau-tau Hessian mismatch: rel={rel:.3e}, \
3034 analytic={h_tt_analytic:?}, fd={h_ttfd:?}"
3035 );
3036 }
3037
3038 pub(crate) struct BinomialLogitDesignMotionFixture {
3048 pub(crate) y: Array1<f64>,
3049 pub(crate) w: Array1<f64>,
3050 pub(crate) x: Array2<f64>,
3051 pub(crate) s0: Array2<f64>,
3052 pub(crate) cfg: RemlConfig,
3053 pub(crate) rho: Array1<f64>,
3054 pub(crate) x_tau_design: Array2<f64>,
3056 pub(crate) s_tau_penalty: Array2<f64>,
3058 }
3059
3060 impl BinomialLogitDesignMotionFixture {
3061 pub(crate) fn new() -> Self {
3062 let y = array![
3064 1.0, 0.0, 1.0, 1.0, 0.0, 0.0, 1.0, 0.0, 1.0, 0.0, 1.0, 1.0, 0.0, 0.0, 1.0, 0.0,
3065 1.0, 1.0, 0.0, 1.0, 0.0, 0.0, 1.0, 1.0, 0.0, 1.0, 0.0, 1.0, 0.0, 1.0
3066 ];
3067 let x = array![
3069 [1.0, -1.50, 0.42, 0.88, -0.31],
3070 [1.0, -1.12, -0.65, 0.14, 1.23],
3071 [1.0, -0.80, 1.10, -0.53, 0.07],
3072 [1.0, -0.55, -0.22, 1.40, -0.90],
3073 [1.0, -0.30, 0.73, -1.05, 0.44],
3074 [1.0, -0.05, -1.33, 0.60, 0.81],
3075 [1.0, 0.18, 0.55, -0.27, -1.15],
3076 [1.0, 0.42, -0.90, 1.12, 0.33],
3077 [1.0, 0.70, 1.28, -0.78, -0.56],
3078 [1.0, 0.95, -0.18, 0.45, 1.40],
3079 [1.0, 1.20, 0.66, -1.30, -0.02],
3080 [1.0, 1.45, -1.05, 0.22, 0.68],
3081 [1.0, -1.35, 0.90, 0.55, -0.43],
3082 [1.0, -0.98, -0.40, -0.88, 1.05],
3083 [1.0, -0.62, 1.42, 0.30, -0.70],
3084 [1.0, -0.28, -0.77, -1.18, 0.52],
3085 [1.0, 0.05, 0.15, 0.95, -1.35],
3086 [1.0, 0.33, -1.20, -0.40, 0.18],
3087 [1.0, 0.60, 0.82, 1.25, -0.85],
3088 [1.0, 0.88, -0.50, -0.65, 1.10],
3089 [1.0, 1.15, 1.05, 0.10, -0.22],
3090 [1.0, -1.22, -0.95, 0.72, 0.90],
3091 [1.0, -0.75, 0.38, -1.42, 0.15],
3092 [1.0, -0.42, -1.15, 0.50, -1.08],
3093 [1.0, -0.10, 0.60, -0.15, 0.75],
3094 [1.0, 0.25, -0.28, 1.05, -0.48],
3095 [1.0, 0.52, 1.35, -0.92, 0.30],
3096 [1.0, 0.80, -0.70, 0.38, 1.20],
3097 [1.0, 1.08, 0.48, -0.60, -0.95],
3098 [1.0, 1.35, -0.55, 0.85, 0.42]
3099 ];
3100 let s0 = array![
3102 [0.0, 0.0, 0.0, 0.0, 0.0],
3103 [0.0, 1.40, 0.15, 0.05, -0.10],
3104 [0.0, 0.15, 1.10, -0.20, 0.08],
3105 [0.0, 0.05, -0.20, 0.95, 0.12],
3106 [0.0, -0.10, 0.08, 0.12, 1.25]
3107 ];
3108 let cfg = RemlConfig::external(binomial_logit_glm_spec(), 1e-14, false);
3109 let x_tau_design = array![
3112 [0.0, 1.2e-3, -0.8e-3, 0.5e-3, -1.5e-3],
3113 [0.0, -2.0e-3, 1.4e-3, -0.3e-3, 0.9e-3],
3114 [0.0, 0.6e-3, -1.1e-3, 1.8e-3, -0.4e-3],
3115 [0.0, -1.3e-3, 0.7e-3, -1.0e-3, 2.1e-3],
3116 [0.0, 0.9e-3, -0.5e-3, 0.2e-3, -0.8e-3],
3117 [0.0, -0.4e-3, 1.8e-3, -1.5e-3, 0.3e-3],
3118 [0.0, 1.5e-3, -1.3e-3, 0.8e-3, -1.1e-3],
3119 [0.0, -0.7e-3, 0.4e-3, -2.0e-3, 1.6e-3],
3120 [0.0, 2.2e-3, -0.9e-3, 1.3e-3, -0.6e-3],
3121 [0.0, -1.0e-3, 1.6e-3, -0.7e-3, 0.5e-3],
3122 [0.0, 0.3e-3, -2.1e-3, 1.1e-3, -1.8e-3],
3123 [0.0, -1.8e-3, 0.2e-3, -0.4e-3, 1.3e-3],
3124 [0.0, 1.1e-3, -1.5e-3, 2.0e-3, -0.2e-3],
3125 [0.0, -0.5e-3, 0.9e-3, -1.2e-3, 0.7e-3],
3126 [0.0, 1.7e-3, -0.3e-3, 0.6e-3, -2.0e-3],
3127 [0.0, -1.4e-3, 1.1e-3, -0.9e-3, 0.4e-3],
3128 [0.0, 0.8e-3, -1.7e-3, 1.5e-3, -0.1e-3],
3129 [0.0, -0.2e-3, 0.6e-3, -1.8e-3, 1.0e-3],
3130 [0.0, 1.4e-3, -0.4e-3, 0.3e-3, -1.3e-3],
3131 [0.0, -0.9e-3, 2.0e-3, -0.5e-3, 0.8e-3],
3132 [0.0, 0.5e-3, -1.0e-3, 1.6e-3, -0.7e-3],
3133 [0.0, -2.1e-3, 0.3e-3, -0.8e-3, 1.5e-3],
3134 [0.0, 0.7e-3, -1.8e-3, 0.9e-3, -0.3e-3],
3135 [0.0, -0.6e-3, 1.3e-3, -2.2e-3, 1.1e-3],
3136 [0.0, 1.9e-3, -0.7e-3, 0.4e-3, -0.9e-3],
3137 [0.0, -1.1e-3, 0.5e-3, -1.4e-3, 2.2e-3],
3138 [0.0, 0.4e-3, -1.6e-3, 1.2e-3, -0.5e-3],
3139 [0.0, -1.6e-3, 0.8e-3, -0.1e-3, 0.6e-3],
3140 [0.0, 1.3e-3, -2.2e-3, 0.7e-3, -1.4e-3],
3141 [0.0, -0.3e-3, 1.0e-3, -1.6e-3, 1.8e-3]
3142 ];
3143 let s_tau_penalty = array![
3145 [0.0, 0.0, 0.0, 0.0, 0.0],
3146 [0.0, 0.30, 0.05, -0.02, 0.04],
3147 [0.0, 0.05, 0.22, 0.03, -0.01],
3148 [0.0, -0.02, 0.03, 0.18, 0.06],
3149 [0.0, 0.04, -0.01, 0.06, 0.26]
3150 ];
3151 Self {
3152 w: Array1::<f64>::ones(y.len()),
3153 y,
3154 x,
3155 s0,
3156 cfg,
3157 rho: array![0.0],
3158 x_tau_design,
3159 s_tau_penalty,
3160 }
3161 }
3162 }
3163
3164 impl LogitDesignMotionFixture for BinomialLogitDesignMotionFixture {
3165 fn y(&self) -> &Array1<f64> {
3166 &self.y
3167 }
3168 fn w(&self) -> &Array1<f64> {
3169 &self.w
3170 }
3171 fn x(&self) -> &Array2<f64> {
3172 &self.x
3173 }
3174 fn s0(&self) -> &Array2<f64> {
3175 &self.s0
3176 }
3177 fn cfg(&self) -> &RemlConfig {
3178 &self.cfg
3179 }
3180 fn rho(&self) -> &Array1<f64> {
3181 &self.rho
3182 }
3183 }
3184
3185 #[test]
3188 pub(crate) fn binomial_logit_n30_design_moving_gradient_matches_fd() {
3189 let f = BinomialLogitDesignMotionFixture::new();
3196 let state = f.state();
3197 let s_tau = Array2::<f64>::zeros((5, 5));
3198 let hyper = DirectionalHyperParam::single_penalty(
3199 0,
3200 f.x_tau_design.clone(),
3201 s_tau.clone(),
3202 None,
3203 None,
3204 )
3205 .expect("design-moving hyper direction");
3206
3207 let v_tau_analytic = single_directional_tau_gradient(&state, &f.rho, hyper)
3208 .expect("analytic directional gradient");
3209 let v_tau_fd = f.fd_directional_gradient(&f.x_tau_design, &s_tau);
3210
3211 let v_rel = (v_tau_analytic - v_tau_fd).abs() / v_tau_fd.abs().max(1e-10);
3212 assert!(
3213 v_rel < 1e-3,
3214 "Binomial-logit n=30 design-moving gradient mismatch: rel={v_rel:.3e}, \
3215 analytic={v_tau_analytic:.6e}, fd={v_tau_fd:.6e}"
3216 );
3217 }
3218
3219 #[test]
3220 pub(crate) fn binomial_logit_n30_penalty_only_gradient_matches_fd() {
3221 let f = BinomialLogitDesignMotionFixture::new();
3226 let state = f.state();
3227 let x_tau = Array2::<f64>::zeros(f.x.raw_dim());
3228 let hyper = DirectionalHyperParam::single_penalty(
3229 0,
3230 x_tau.clone(),
3231 f.s_tau_penalty.clone(),
3232 None,
3233 None,
3234 )
3235 .expect("penalty-only hyper direction");
3236
3237 let v_tau_analytic = single_directional_tau_gradient(&state, &f.rho, hyper)
3238 .expect("analytic directional gradient");
3239 let v_tau_fd = f.fd_directional_gradient(&x_tau, &f.s_tau_penalty);
3240
3241 let v_rel = (v_tau_analytic - v_tau_fd).abs() / v_tau_fd.abs().max(1e-10);
3242 assert!(
3243 v_rel < 1e-3,
3244 "Binomial-logit n=30 penalty-only gradient mismatch: rel={v_rel:.3e}, \
3245 analytic={v_tau_analytic:.6e}, fd={v_tau_fd:.6e}"
3246 );
3247 }
3248
3249 #[test]
3250 pub(crate) fn binomial_logit_n30_joint_design_penalty_gradient_matches_fd() {
3251 let f = BinomialLogitDesignMotionFixture::new();
3256 let state = f.state();
3257 let hyper = DirectionalHyperParam::single_penalty(
3258 0,
3259 f.x_tau_design.clone(),
3260 f.s_tau_penalty.clone(),
3261 None,
3262 None,
3263 )
3264 .expect("joint design+penalty hyper direction");
3265
3266 let v_tau_analytic = single_directional_tau_gradient(&state, &f.rho, hyper)
3267 .expect("analytic directional gradient");
3268 let v_tau_fd = f.fd_directional_gradient(&f.x_tau_design, &f.s_tau_penalty);
3269
3270 let v_rel = (v_tau_analytic - v_tau_fd).abs() / v_tau_fd.abs().max(1e-10);
3271 assert!(
3272 v_rel < 1e-3,
3273 "Binomial-logit n=30 joint design+penalty gradient mismatch: rel={v_rel:.3e}, \
3274 analytic={v_tau_analytic:.6e}, fd={v_tau_fd:.6e}"
3275 );
3276 }
3277
3278 #[test]
3279 pub(crate) fn binomial_logit_n30_design_moving_hessian_matches_fd() {
3280 let f = BinomialLogitDesignMotionFixture::new();
3285 let x_tau_0 = Array2::<f64>::zeros(f.x.raw_dim());
3286 let s_tau_0 = f.s_tau_penalty.clone();
3287 let x_tau_1 = f.x_tau_design.clone();
3288 let s_tau_1 = Array2::<f64>::zeros((5, 5));
3289
3290 let hyper_dirs = vec![
3291 DirectionalHyperParam::single_penalty(0, x_tau_0.clone(), s_tau_0.clone(), None, None)
3292 .expect("penalty-only direction"),
3293 DirectionalHyperParam::single_penalty(0, x_tau_1.clone(), s_tau_1.clone(), None, None)
3294 .expect("design-moving direction"),
3295 ];
3296
3297 let state = f.state();
3298 let mut theta = Array1::<f64>::zeros(f.rho.len() + hyper_dirs.len());
3299 theta.slice_mut(s![..f.rho.len()]).assign(&f.rho);
3300 let (_, _, h_full) =
3301 compute_joint_hypercostgradienthessian(&state, &theta, f.rho.len(), &hyper_dirs)
3302 .expect("joint cost+gradient+hessian");
3303 let h_tt_analytic = h_full.slice(s![f.rho.len().., f.rho.len()..]).to_owned();
3304
3305 let x_tau_mats = vec![x_tau_0.clone(), x_tau_1.clone()];
3306 let s_tau_mats = vec![s_tau_0.clone(), s_tau_1.clone()];
3307 let h_tt_fd = directional_tau_hessian_fd_reference(
3308 &f.y,
3309 &f.w,
3310 &f.x,
3311 &f.s0,
3312 &f.cfg,
3313 &f.rho,
3314 &hyper_dirs,
3315 &x_tau_mats,
3316 &s_tau_mats,
3317 );
3318
3319 let num = (&h_tt_analytic - &h_tt_fd)
3320 .iter()
3321 .map(|v| v * v)
3322 .sum::<f64>()
3323 .sqrt();
3324 let den = h_tt_fd.iter().map(|v| v * v).sum::<f64>().sqrt().max(1e-10);
3325 let rel = num / den;
3326 assert!(
3327 rel < 1e-4,
3328 "Binomial-logit n=30 tau-tau Hessian mismatch: rel={rel:.3e}, \
3329 analytic={h_tt_analytic:?}, fd={h_tt_fd:?}"
3330 );
3331 }
3332
3333 #[test]
3334 pub(crate) fn binomial_logit_n30_nonzero_rho_design_moving_gradient_matches_fd() {
3335 let f = BinomialLogitDesignMotionFixture::new();
3339 let rho = array![1.5];
3340 let s_tau = Array2::<f64>::zeros((5, 5));
3341
3342 let state = f.state();
3343 let hyper = DirectionalHyperParam::single_penalty(
3344 0,
3345 f.x_tau_design.clone(),
3346 s_tau.clone(),
3347 None,
3348 None,
3349 )
3350 .expect("design-moving hyper direction");
3351
3352 let v_tau_analytic = single_directional_tau_gradient(&state, &rho, hyper)
3353 .expect("analytic directional gradient");
3354
3355 let h = 2e-5;
3357 let (state_plus, state_minus) = f.state_perturbed(&f.x_tau_design, &s_tau, h);
3358 let v_plus = state_plus.compute_cost(&rho).expect("cost+");
3359 let v_minus = state_minus.compute_cost(&rho).expect("cost-");
3360 let v_tau_fd = (v_plus - v_minus) / (2.0 * h);
3361
3362 let v_rel = (v_tau_analytic - v_tau_fd).abs() / v_tau_fd.abs().max(1e-10);
3363 assert!(
3364 v_rel < 1e-3,
3365 "Binomial-logit n=30 rho=1.5 design-moving gradient mismatch: rel={v_rel:.3e}, \
3366 analytic={v_tau_analytic:.6e}, fd={v_tau_fd:.6e}"
3367 );
3368 }
3369
3370 #[test]
3371 pub(crate) fn binomial_logit_n30_rank_deficient_hessian_matches_cost_fd() {
3372 let f = BinomialLogitDesignMotionFixture::new();
3407 let x_tau_0 = Array2::<f64>::zeros(f.x.raw_dim());
3408 let s_tau_0 = f.s_tau_penalty.clone();
3409 let x_tau_1 = f.x_tau_design.clone();
3410 let s_tau_1 = Array2::<f64>::zeros((5, 5));
3411
3412 let hyper_dirs = vec![
3413 DirectionalHyperParam::single_penalty(0, x_tau_0.clone(), s_tau_0.clone(), None, None)
3414 .expect("penalty-only direction"),
3415 DirectionalHyperParam::single_penalty(0, x_tau_1.clone(), s_tau_1.clone(), None, None)
3416 .expect("design-moving direction"),
3417 ];
3418
3419 let state = f.state();
3421 let mut theta = Array1::<f64>::zeros(f.rho.len() + hyper_dirs.len());
3422 theta.slice_mut(s![..f.rho.len()]).assign(&f.rho);
3423 let (_, _, h_full) =
3424 compute_joint_hypercostgradienthessian(&state, &theta, f.rho.len(), &hyper_dirs)
3425 .expect("joint cost+gradient+hessian");
3426 let h_tt_analytic = h_full.slice(s![f.rho.len().., f.rho.len()..]).to_owned();
3427
3428 const TARGET_PHYSICAL_STEP: f64 = 1e-5;
3432 let x_tau_mats = [&x_tau_0, &x_tau_1];
3433 let s_tau_mats = [&s_tau_0, &s_tau_1];
3434 let steps: [f64; 2] = {
3435 let mut steps = [0.0; 2];
3436 for (j, step) in steps.iter_mut().enumerate() {
3437 let scale = x_tau_mats[j]
3438 .iter()
3439 .chain(s_tau_mats[j].iter())
3440 .fold(0.0_f64, |acc, value| acc.max(value.abs()));
3441 *step = if scale > 0.0 {
3442 TARGET_PHYSICAL_STEP / scale
3443 } else {
3444 TARGET_PHYSICAL_STEP
3445 };
3446 }
3447 steps
3448 };
3449
3450 let eval_cost = |a: f64, b: f64| -> f64 {
3452 let x_eval = &f.x
3453 + &x_tau_mats[0].mapv(|v| a * steps[0] * v)
3454 + &x_tau_mats[1].mapv(|v| b * steps[1] * v);
3455 let s_eval = &f.s0
3456 + &s_tau_mats[0].mapv(|v| a * steps[0] * v)
3457 + &s_tau_mats[1].mapv(|v| b * steps[1] * v);
3458 let st = build_logit_state(&f.y, &f.w, &x_eval, &s_eval, &f.cfg);
3459 st.compute_cost(&f.rho).expect("cost eval")
3460 };
3461
3462 let v_00 = eval_cost(0.0, 0.0);
3463 let v_p0 = eval_cost(1.0, 0.0);
3464 let v_m0 = eval_cost(-1.0, 0.0);
3465 let v_0p = eval_cost(0.0, 1.0);
3466 let v_0m = eval_cost(0.0, -1.0);
3467 let v_pp = eval_cost(1.0, 1.0);
3468 let v_pm = eval_cost(1.0, -1.0);
3469 let v_mp = eval_cost(-1.0, 1.0);
3470 let v_mm = eval_cost(-1.0, -1.0);
3471
3472 let h00_fd = (v_p0 - 2.0 * v_00 + v_m0) / (steps[0] * steps[0]);
3473 let h11_fd = (v_0p - 2.0 * v_00 + v_0m) / (steps[1] * steps[1]);
3474 let h01_fd = (v_pp - v_pm - v_mp + v_mm) / (4.0 * steps[0] * steps[1]);
3475
3476 let h_tt_fd = array![[h00_fd, h01_fd], [h01_fd, h11_fd]];
3477
3478 let num = (&h_tt_analytic - &h_tt_fd)
3479 .iter()
3480 .map(|v| v * v)
3481 .sum::<f64>()
3482 .sqrt();
3483 let den = h_tt_fd.iter().map(|v| v * v).sum::<f64>().sqrt().max(1e-10);
3484 let rel = num / den;
3485
3486 assert!(
3487 rel < 3e-3,
3488 "Binomial-logit n=30 rank-deficient Hessian vs cost-FD mismatch: rel={rel:.3e}, \
3489 analytic={h_tt_analytic:?}, fd={h_tt_fd:?}"
3490 );
3491 }
3492}
3493
3494#[derive(Clone, Copy, Debug)]
3495pub(crate) enum RemlGeometry {
3496 DenseSpectral,
3497 SparseExactSpd,
3498}
3499
3500trait PenalizedGeometry {
3501 fn backend_kind(&self) -> GeometryBackendKind;
3502}
3503
3504#[derive(Clone)]
3505pub(crate) enum DerivativeMatrixStorage {
3506 Dense(Array2<f64>),
3507 Zero(ZeroDerivativeMatrix),
3508 Embedded(EmbeddedDerivativeMatrix),
3509 Implicit(ImplicitDerivativeOp),
3510 LatentCoord(LatentCoordDerivativeOp),
3511}
3512
3513trait DerivativeStorageBackend {
3525 fn resident_byte_count(&self) -> usize;
3526 fn design_nrows(&self) -> usize;
3527 fn design_ncols(&self) -> usize;
3528 fn penalty_dim(&self) -> usize;
3529 fn uses_implicit_storage(&self) -> bool;
3530 fn any_nonzero(&self) -> bool;
3531 fn materialize(&self) -> Array2<f64>;
3532 fn implicit_first_axis_info(
3533 &self,
3534 ) -> Option<(
3535 std::sync::Arc<gam_terms::basis::ImplicitDesignPsiDerivative>,
3536 usize,
3537 )>;
3538 fn implicit_axis_count_hint(&self) -> Option<usize>;
3539 fn design_forward_mul_original(&self, u: &Array1<f64>) -> Result<Array1<f64>, EstimationError>;
3540 fn design_transpose_mul_original(
3541 &self,
3542 v: &Array1<f64>,
3543 ) -> Result<Array1<f64>, EstimationError>;
3544 fn design_transformed(
3545 &self,
3546 qs: &Array2<f64>,
3547 free_basis_opt: Option<&Array2<f64>>,
3548 ) -> Result<Array2<f64>, EstimationError>;
3549 fn design_transformed_forward_mul(
3553 &self,
3554 qs: &Array2<f64>,
3555 free_basis_opt: Option<&Array2<f64>>,
3556 u: &Array1<f64>,
3557 ) -> Result<Array1<f64>, EstimationError> {
3558 Ok(self.design_transformed(qs, free_basis_opt)?.dot(u))
3559 }
3560 fn design_transformed_transpose_mul(
3563 &self,
3564 qs: &Array2<f64>,
3565 free_basis_opt: Option<&Array2<f64>>,
3566 v: &Array1<f64>,
3567 ) -> Result<Array1<f64>, EstimationError> {
3568 Ok(self.design_transformed(qs, free_basis_opt)?.t().dot(v))
3569 }
3570 fn penalty_transformed(
3571 &self,
3572 qs: &Array2<f64>,
3573 free_basis_opt: Option<&Array2<f64>>,
3574 ) -> Result<Array2<f64>, EstimationError>;
3575 fn penalty_scaled_add_to(
3576 &self,
3577 target: &mut Array2<f64>,
3578 amp: f64,
3579 ) -> Result<(), EstimationError>;
3580}
3581
3582macro_rules! storage_dispatch {
3587 ($scrutinee:expr, $backend:ident => $body:expr) => {
3588 match $scrutinee {
3589 DerivativeMatrixStorage::Dense($backend) => $body,
3590 DerivativeMatrixStorage::Zero($backend) => $body,
3591 DerivativeMatrixStorage::Embedded($backend) => $body,
3592 DerivativeMatrixStorage::Implicit($backend) => $body,
3593 DerivativeMatrixStorage::LatentCoord($backend) => $body,
3594 }
3595 };
3596}
3597
3598#[derive(Clone)]
3599pub(crate) struct ZeroDerivativeMatrix {
3600 rows: usize,
3601 cols: usize,
3602}
3603
3604impl ZeroDerivativeMatrix {
3605 pub(crate) fn new(rows: usize, cols: usize) -> Self {
3606 Self { rows, cols }
3607 }
3608}
3609
3610#[derive(Clone, Copy, Debug)]
3612pub enum ImplicitDerivLevel {
3613 First(usize),
3615 SecondDiag(usize),
3617 SecondCross(usize, usize),
3619}
3620
3621#[derive(Clone)]
3624pub(crate) struct ImplicitDerivativeOp {
3625 pub(crate) operator: std::sync::Arc<gam_terms::basis::ImplicitDesignPsiDerivative>,
3626 pub(crate) level: ImplicitDerivLevel,
3627 pub(crate) global_range: Range<usize>,
3628 pub(crate) total_dim: usize,
3629 pub(crate) cached_dense: std::sync::Arc<gam_runtime::resource::RayonSafeOnce<Array2<f64>>>,
3639}
3640
3641#[derive(Clone)]
3642pub(crate) struct LatentCoordDerivativeOp {
3643 pub(crate) operator: std::sync::Arc<gam_terms::basis::LatentCoordDesignDerivative>,
3644 pub(crate) flat_axis: usize,
3645 pub(crate) global_range: Range<usize>,
3646 pub(crate) total_dim: usize,
3647 pub(crate) cached_dense: std::sync::Arc<gam_runtime::resource::RayonSafeOnce<Array2<f64>>>,
3648}
3649
3650impl LatentCoordDerivativeOp {
3651 pub(crate) fn materialize_local(&self) -> Array2<f64> {
3652 self.operator.materialize_axis(self.flat_axis).expect(
3653 "radial scalar evaluation failed during latent-coordinate derivative materialization",
3654 )
3655 }
3656
3657 pub(crate) fn materialize_dense(&self) -> &Array2<f64> {
3658 self.cached_dense.get_or_compute(|| {
3659 let local = self.materialize_local();
3660 let mut out = Array2::<f64>::zeros((local.nrows(), self.total_dim));
3661 out.slice_mut(s![.., self.global_range.clone()])
3662 .assign(&local);
3663 out
3664 })
3665 }
3666
3667 pub(crate) fn nrows(&self) -> usize {
3668 self.operator.n_data()
3669 }
3670
3671 pub(crate) fn ncols(&self) -> usize {
3672 self.total_dim
3673 }
3674
3675 pub(crate) fn transpose_mul(&self, v: &Array1<f64>) -> Array1<f64> {
3676 let local = self
3677 .operator
3678 .transpose_mul_axis(self.flat_axis, &v.view())
3679 .expect(
3680 "radial scalar evaluation failed during latent-coordinate derivative transpose_mul",
3681 );
3682 let mut out = Array1::<f64>::zeros(self.total_dim);
3683 out.slice_mut(s![self.global_range.clone()]).assign(&local);
3684 out
3685 }
3686
3687 pub(crate) fn forward_mul(&self, u: &Array1<f64>) -> Array1<f64> {
3688 let u_local = u.slice(s![self.global_range.clone()]).to_owned();
3689 self.operator
3690 .forward_mul_axis(self.flat_axis, &u_local.view())
3691 .expect(
3692 "radial scalar evaluation failed during latent-coordinate derivative forward_mul",
3693 )
3694 }
3695}
3696
3697impl ImplicitDerivativeOp {
3698 pub(crate) fn materialize_local(&self) -> Array2<f64> {
3699 match self.level {
3700 ImplicitDerivLevel::First(axis) => self.operator.materialize_first(axis).expect(
3701 "radial scalar evaluation failed during implicit derivative materialization",
3702 ),
3703 ImplicitDerivLevel::SecondDiag(axis) => {
3704 self.operator.materialize_second_diag(axis).expect(
3705 "radial scalar evaluation failed during implicit derivative materialization",
3706 )
3707 }
3708 ImplicitDerivLevel::SecondCross(d, e) => {
3709 self.operator.materialize_second_cross(d, e).expect(
3710 "radial scalar evaluation failed during implicit derivative materialization",
3711 )
3712 }
3713 }
3714 }
3715
3716 pub(crate) fn materialize_dense(&self) -> &Array2<f64> {
3717 self.cached_dense.get_or_compute(|| {
3718 let local = self.materialize_local();
3719 let mut out = Array2::<f64>::zeros((local.nrows(), self.total_dim));
3720 out.slice_mut(s![.., self.global_range.clone()])
3721 .assign(&local);
3722 out
3723 })
3724 }
3725
3726 pub(crate) fn nrows(&self) -> usize {
3727 self.operator.n_data()
3728 }
3729
3730 pub(crate) fn ncols(&self) -> usize {
3731 self.total_dim
3732 }
3733
3734 pub(crate) fn transpose_mul(&self, v: &Array1<f64>) -> Array1<f64> {
3735 let local = match self.level {
3736 ImplicitDerivLevel::First(axis) => self
3737 .operator
3738 .transpose_mul(axis, &v.view())
3739 .expect("radial scalar evaluation failed during implicit derivative transpose_mul"),
3740 ImplicitDerivLevel::SecondDiag(axis) => self
3741 .operator
3742 .transpose_mul_second_diag(axis, &v.view())
3743 .expect("radial scalar evaluation failed during implicit derivative transpose_mul"),
3744 ImplicitDerivLevel::SecondCross(d, e) => self
3745 .operator
3746 .transpose_mul_second_cross(d, e, &v.view())
3747 .expect("radial scalar evaluation failed during implicit derivative transpose_mul"),
3748 };
3749 let mut out = Array1::<f64>::zeros(self.total_dim);
3750 out.slice_mut(s![self.global_range.clone()]).assign(&local);
3751 out
3752 }
3753
3754 pub(crate) fn forward_mul(&self, u: &Array1<f64>) -> Array1<f64> {
3755 let u_local = u.slice(s![self.global_range.clone()]).to_owned();
3756 match self.level {
3757 ImplicitDerivLevel::First(axis) => self
3758 .operator
3759 .forward_mul(axis, &u_local.view())
3760 .expect("radial scalar evaluation failed during implicit derivative forward_mul"),
3761 ImplicitDerivLevel::SecondDiag(axis) => self
3762 .operator
3763 .forward_mul_second_diag(axis, &u_local.view())
3764 .expect("radial scalar evaluation failed during implicit derivative forward_mul"),
3765 ImplicitDerivLevel::SecondCross(d, e) => self
3766 .operator
3767 .forward_mul_second_cross(d, e, &u_local.view())
3768 .expect("radial scalar evaluation failed during implicit derivative forward_mul"),
3769 }
3770 }
3771}
3772
3773#[derive(Clone)]
3774pub(crate) struct EmbeddedDerivativeMatrix {
3775 pub(crate) local: Array2<f64>,
3776 pub(crate) global_range: Range<usize>,
3777 pub(crate) total_dim: usize,
3778}
3779
3780impl EmbeddedDerivativeMatrix {
3781 pub(crate) fn new(local: Array2<f64>, global_range: Range<usize>, total_dim: usize) -> Self {
3782 Self {
3783 local,
3784 global_range,
3785 total_dim,
3786 }
3787 }
3788}
3789
3790impl DerivativeStorageBackend for Array2<f64> {
3791 fn resident_byte_count(&self) -> usize {
3792 self.len().saturating_mul(std::mem::size_of::<f64>())
3793 }
3794 fn design_nrows(&self) -> usize {
3795 Array2::nrows(self)
3796 }
3797 fn design_ncols(&self) -> usize {
3798 Array2::ncols(self)
3799 }
3800 fn penalty_dim(&self) -> usize {
3801 Array2::nrows(self)
3802 }
3803 fn uses_implicit_storage(&self) -> bool {
3804 false
3805 }
3806 fn any_nonzero(&self) -> bool {
3807 self.iter().any(|v| *v != 0.0)
3808 }
3809 fn materialize(&self) -> Array2<f64> {
3810 self.clone()
3811 }
3812 fn implicit_first_axis_info(
3813 &self,
3814 ) -> Option<(
3815 std::sync::Arc<gam_terms::basis::ImplicitDesignPsiDerivative>,
3816 usize,
3817 )> {
3818 None
3819 }
3820 fn implicit_axis_count_hint(&self) -> Option<usize> {
3821 None
3822 }
3823
3824 fn design_forward_mul_original(&self, u: &Array1<f64>) -> Result<Array1<f64>, EstimationError> {
3825 if Array2::ncols(self) != u.len() {
3826 crate::bail_invalid_estim!(
3827 "dense hyper design derivative forward_mul_original width mismatch: matrix={}x{}, vector={}",
3828 Array2::nrows(self),
3829 Array2::ncols(self),
3830 u.len()
3831 );
3832 }
3833 Ok(self.dot(u))
3834 }
3835
3836 fn design_transpose_mul_original(
3837 &self,
3838 v: &Array1<f64>,
3839 ) -> Result<Array1<f64>, EstimationError> {
3840 if Array2::nrows(self) != v.len() {
3841 crate::bail_invalid_estim!(
3842 "dense hyper design derivative transpose_mul_original height mismatch: matrix={}x{}, vector={}",
3843 Array2::nrows(self),
3844 Array2::ncols(self),
3845 v.len()
3846 );
3847 }
3848 Ok(self.t().dot(v))
3849 }
3850
3851 fn design_transformed(
3852 &self,
3853 qs: &Array2<f64>,
3854 free_basis_opt: Option<&Array2<f64>>,
3855 ) -> Result<Array2<f64>, EstimationError> {
3856 Ok(gam_linalg::matrix::DenseRightProductView::new(self)
3857 .with_factor(qs)
3858 .with_optional_factor(free_basis_opt)
3859 .materialize())
3860 }
3861
3862 fn penalty_transformed(
3863 &self,
3864 qs: &Array2<f64>,
3865 free_basis_opt: Option<&Array2<f64>>,
3866 ) -> Result<Array2<f64>, EstimationError> {
3867 let mut transformed = qs.t().dot(self).dot(qs);
3868 if let Some(z) = free_basis_opt {
3869 transformed = z.t().dot(&transformed).dot(z);
3870 }
3871 Ok(transformed)
3872 }
3873
3874 fn penalty_scaled_add_to(
3875 &self,
3876 target: &mut Array2<f64>,
3877 amp: f64,
3878 ) -> Result<(), EstimationError> {
3879 if target.raw_dim() != self.raw_dim() {
3880 crate::bail_invalid_estim!(
3881 "dense hyper penalty derivative shape mismatch: target={}x{}, matrix={}x{}",
3882 target.nrows(),
3883 target.ncols(),
3884 Array2::nrows(self),
3885 Array2::ncols(self)
3886 );
3887 }
3888 target.scaled_add(amp, self);
3889 Ok(())
3890 }
3891}
3892
3893impl DerivativeStorageBackend for ZeroDerivativeMatrix {
3894 fn resident_byte_count(&self) -> usize {
3895 0
3896 }
3897 fn design_nrows(&self) -> usize {
3898 self.rows
3899 }
3900 fn design_ncols(&self) -> usize {
3901 self.cols
3902 }
3903 fn penalty_dim(&self) -> usize {
3904 self.cols
3905 }
3906 fn uses_implicit_storage(&self) -> bool {
3907 false
3908 }
3909 fn any_nonzero(&self) -> bool {
3910 false
3911 }
3912 fn materialize(&self) -> Array2<f64> {
3913 Array2::<f64>::zeros((self.rows, self.cols))
3914 }
3915 fn implicit_first_axis_info(
3916 &self,
3917 ) -> Option<(
3918 std::sync::Arc<gam_terms::basis::ImplicitDesignPsiDerivative>,
3919 usize,
3920 )> {
3921 None
3922 }
3923 fn implicit_axis_count_hint(&self) -> Option<usize> {
3924 None
3925 }
3926
3927 fn design_forward_mul_original(&self, u: &Array1<f64>) -> Result<Array1<f64>, EstimationError> {
3928 if self.cols != u.len() {
3929 crate::bail_invalid_estim!(
3930 "zero hyper design derivative forward_mul_original width mismatch: matrix={}x{}, vector={}",
3931 self.rows,
3932 self.cols,
3933 u.len()
3934 );
3935 }
3936 Ok(Array1::<f64>::zeros(self.rows))
3937 }
3938
3939 fn design_transpose_mul_original(
3940 &self,
3941 v: &Array1<f64>,
3942 ) -> Result<Array1<f64>, EstimationError> {
3943 if self.rows != v.len() {
3944 crate::bail_invalid_estim!(
3945 "zero hyper design derivative transpose_mul_original height mismatch: matrix={}x{}, vector={}",
3946 self.rows,
3947 self.cols,
3948 v.len()
3949 );
3950 }
3951 Ok(Array1::<f64>::zeros(self.cols))
3952 }
3953
3954 fn design_transformed(
3955 &self,
3956 qs: &Array2<f64>,
3957 free_basis_opt: Option<&Array2<f64>>,
3958 ) -> Result<Array2<f64>, EstimationError> {
3959 if self.cols != qs.nrows() {
3960 crate::bail_invalid_estim!(
3961 "zero design derivative width mismatch: total_cols={}, qs rows={}",
3962 self.cols,
3963 qs.nrows()
3964 );
3965 }
3966 let cols = free_basis_opt.map_or(qs.ncols(), |z| z.ncols());
3967 Ok(Array2::<f64>::zeros((self.rows, cols)))
3968 }
3969
3970 fn design_transformed_forward_mul(
3971 &self,
3972 qs: &Array2<f64>,
3973 free_basis_opt: Option<&Array2<f64>>,
3974 u: &Array1<f64>,
3975 ) -> Result<Array1<f64>, EstimationError> {
3976 if self.cols != qs.nrows() {
3977 crate::bail_invalid_estim!(
3978 "zero design derivative width mismatch: total_cols={}, qs rows={}",
3979 self.cols,
3980 qs.nrows()
3981 );
3982 }
3983 let cols = free_basis_opt.map_or(qs.ncols(), |z| z.ncols());
3984 if u.len() != cols {
3985 crate::bail_invalid_estim!(
3986 "zero design derivative transformed forward width mismatch: expected {}, vector={}",
3987 cols,
3988 u.len()
3989 );
3990 }
3991 Ok(Array1::<f64>::zeros(self.rows))
3992 }
3993
3994 fn design_transformed_transpose_mul(
3995 &self,
3996 qs: &Array2<f64>,
3997 free_basis_opt: Option<&Array2<f64>>,
3998 v: &Array1<f64>,
3999 ) -> Result<Array1<f64>, EstimationError> {
4000 if self.rows != v.len() {
4001 crate::bail_invalid_estim!(
4002 "zero design derivative transpose height mismatch: matrix rows={}, vector={}",
4003 self.rows,
4004 v.len()
4005 );
4006 }
4007 if self.cols != qs.nrows() {
4008 crate::bail_invalid_estim!(
4009 "zero design derivative width mismatch: total_cols={}, qs rows={}",
4010 self.cols,
4011 qs.nrows()
4012 );
4013 }
4014 let cols = free_basis_opt.map_or(qs.ncols(), |z| z.ncols());
4015 Ok(Array1::<f64>::zeros(cols))
4016 }
4017
4018 fn penalty_transformed(
4019 &self,
4020 qs: &Array2<f64>,
4021 free_basis_opt: Option<&Array2<f64>>,
4022 ) -> Result<Array2<f64>, EstimationError> {
4023 if self.cols != qs.nrows() {
4024 crate::bail_invalid_estim!(
4025 "zero penalty derivative width mismatch: total_dim={}, qs rows={}",
4026 self.cols,
4027 qs.nrows()
4028 );
4029 }
4030 let cols = free_basis_opt.map_or(qs.ncols(), |z| z.ncols());
4031 Ok(Array2::<f64>::zeros((cols, cols)))
4032 }
4033
4034 fn penalty_scaled_add_to(
4035 &self,
4036 target: &mut Array2<f64>,
4037 amp: f64,
4038 ) -> Result<(), EstimationError> {
4039 if !amp.is_finite() {
4043 crate::bail_invalid_estim!(
4044 "zero hyper penalty derivative received non-finite amp={amp}"
4045 );
4046 }
4047 if target.nrows() != self.cols || target.ncols() != self.cols {
4048 crate::bail_invalid_estim!(
4049 "zero hyper penalty derivative shape mismatch: target={}x{}, expected {}x{}",
4050 target.nrows(),
4051 target.ncols(),
4052 self.cols,
4053 self.cols
4054 );
4055 }
4056 Ok(())
4057 }
4058}
4059
4060impl DerivativeStorageBackend for EmbeddedDerivativeMatrix {
4061 fn resident_byte_count(&self) -> usize {
4062 self.local.len().saturating_mul(std::mem::size_of::<f64>())
4063 }
4064 fn design_nrows(&self) -> usize {
4065 self.local.nrows()
4066 }
4067 fn design_ncols(&self) -> usize {
4068 self.total_dim
4069 }
4070 fn penalty_dim(&self) -> usize {
4071 self.total_dim
4072 }
4073 fn uses_implicit_storage(&self) -> bool {
4074 false
4075 }
4076 fn any_nonzero(&self) -> bool {
4077 self.local.iter().any(|v| *v != 0.0)
4078 }
4079 fn materialize(&self) -> Array2<f64> {
4080 let mut dense = Array2::<f64>::zeros((self.local.nrows(), self.total_dim));
4081 dense
4082 .slice_mut(s![.., self.global_range.clone()])
4083 .assign(&self.local);
4084 dense
4085 }
4086 fn implicit_first_axis_info(
4087 &self,
4088 ) -> Option<(
4089 std::sync::Arc<gam_terms::basis::ImplicitDesignPsiDerivative>,
4090 usize,
4091 )> {
4092 None
4093 }
4094 fn implicit_axis_count_hint(&self) -> Option<usize> {
4095 None
4096 }
4097
4098 fn design_forward_mul_original(&self, u: &Array1<f64>) -> Result<Array1<f64>, EstimationError> {
4099 if self.total_dim != u.len() {
4100 crate::bail_invalid_estim!(
4101 "embedded hyper design derivative forward_mul_original width mismatch: total_dim={}, vector={}",
4102 self.total_dim,
4103 u.len()
4104 );
4105 }
4106 let u_local = u.slice(s![self.global_range.clone()]).to_owned();
4107 Ok(self.local.dot(&u_local))
4108 }
4109
4110 fn design_transpose_mul_original(
4111 &self,
4112 v: &Array1<f64>,
4113 ) -> Result<Array1<f64>, EstimationError> {
4114 if self.local.nrows() != v.len() {
4115 crate::bail_invalid_estim!(
4116 "embedded hyper design derivative transpose_mul_original height mismatch: local_rows={}, vector={}",
4117 self.local.nrows(),
4118 v.len()
4119 );
4120 }
4121 let mut out = Array1::<f64>::zeros(self.total_dim);
4122 let pulled = self.local.t().dot(v);
4123 out.slice_mut(s![self.global_range.clone()]).assign(&pulled);
4124 Ok(out)
4125 }
4126
4127 fn design_transformed(
4128 &self,
4129 qs: &Array2<f64>,
4130 free_basis_opt: Option<&Array2<f64>>,
4131 ) -> Result<Array2<f64>, EstimationError> {
4132 if self.total_dim != qs.nrows() {
4133 crate::bail_invalid_estim!(
4134 "embedded design derivative width mismatch: total_cols={}, qs rows={}",
4135 self.total_dim,
4136 qs.nrows()
4137 );
4138 }
4139 let qs_local = qs.slice(s![self.global_range.clone(), ..]);
4140 let mut transformed = self.local.dot(&qs_local);
4141 if let Some(z) = free_basis_opt {
4142 transformed = transformed.dot(z);
4143 }
4144 Ok(transformed)
4145 }
4146
4147 fn penalty_transformed(
4148 &self,
4149 qs: &Array2<f64>,
4150 free_basis_opt: Option<&Array2<f64>>,
4151 ) -> Result<Array2<f64>, EstimationError> {
4152 if self.total_dim != qs.nrows() {
4153 crate::bail_invalid_estim!(
4154 "embedded penalty derivative width mismatch: total_dim={}, qs rows={}",
4155 self.total_dim,
4156 qs.nrows()
4157 );
4158 }
4159 let qs_local = qs.slice(s![self.global_range.clone(), ..]);
4160 let mut transformed = qs_local.t().dot(&self.local).dot(&qs_local);
4161 if let Some(z) = free_basis_opt {
4162 transformed = z.t().dot(&transformed).dot(z);
4163 }
4164 Ok(transformed)
4165 }
4166
4167 fn penalty_scaled_add_to(
4168 &self,
4169 target: &mut Array2<f64>,
4170 amp: f64,
4171 ) -> Result<(), EstimationError> {
4172 if target.nrows() != self.total_dim || target.ncols() != self.total_dim {
4173 crate::bail_invalid_estim!(
4174 "embedded hyper penalty derivative shape mismatch: target={}x{}, expected {}x{}",
4175 target.nrows(),
4176 target.ncols(),
4177 self.total_dim,
4178 self.total_dim
4179 );
4180 }
4181 target
4182 .slice_mut(s![self.global_range.clone(), self.global_range.clone()])
4183 .scaled_add(amp, &self.local);
4184 Ok(())
4185 }
4186}
4187
4188impl DerivativeStorageBackend for ImplicitDerivativeOp {
4189 fn resident_byte_count(&self) -> usize {
4190 0
4191 }
4192 fn design_nrows(&self) -> usize {
4193 self.nrows()
4194 }
4195 fn design_ncols(&self) -> usize {
4196 self.ncols()
4197 }
4198 fn penalty_dim(&self) -> usize {
4199 self.nrows()
4200 }
4201 fn uses_implicit_storage(&self) -> bool {
4202 true
4203 }
4204 fn any_nonzero(&self) -> bool {
4205 true
4206 }
4207 fn materialize(&self) -> Array2<f64> {
4208 self.materialize_dense().clone()
4209 }
4210 fn implicit_first_axis_info(
4211 &self,
4212 ) -> Option<(
4213 std::sync::Arc<gam_terms::basis::ImplicitDesignPsiDerivative>,
4214 usize,
4215 )> {
4216 match self.level {
4217 ImplicitDerivLevel::First(axis) => Some((self.operator.clone(), axis)),
4218 _ => None,
4219 }
4220 }
4221 fn implicit_axis_count_hint(&self) -> Option<usize> {
4222 Some(self.operator.n_axes())
4223 }
4224
4225 fn design_forward_mul_original(&self, u: &Array1<f64>) -> Result<Array1<f64>, EstimationError> {
4226 if self.ncols() != u.len() {
4227 crate::bail_invalid_estim!(
4228 "implicit hyper design derivative forward_mul_original width mismatch: operator_cols={}, vector={}",
4229 self.ncols(),
4230 u.len()
4231 );
4232 }
4233 Ok(self.forward_mul(u))
4234 }
4235
4236 fn design_transpose_mul_original(
4237 &self,
4238 v: &Array1<f64>,
4239 ) -> Result<Array1<f64>, EstimationError> {
4240 if self.nrows() != v.len() {
4241 crate::bail_invalid_estim!(
4242 "implicit hyper design derivative transpose_mul_original height mismatch: operator_rows={}, vector={}",
4243 self.nrows(),
4244 v.len()
4245 );
4246 }
4247 Ok(self.transpose_mul(v))
4248 }
4249
4250 fn design_transformed(
4251 &self,
4252 qs: &Array2<f64>,
4253 free_basis_opt: Option<&Array2<f64>>,
4254 ) -> Result<Array2<f64>, EstimationError> {
4255 let dense = self.materialize_dense();
4256 Ok(gam_linalg::matrix::DenseRightProductView::new(dense)
4257 .with_factor(qs)
4258 .with_optional_factor(free_basis_opt)
4259 .materialize())
4260 }
4261
4262 fn design_transformed_forward_mul(
4263 &self,
4264 qs: &Array2<f64>,
4265 free_basis_opt: Option<&Array2<f64>>,
4266 u: &Array1<f64>,
4267 ) -> Result<Array1<f64>, EstimationError> {
4268 let mut right = if let Some(z) = free_basis_opt {
4269 z.dot(u)
4270 } else {
4271 u.clone()
4272 };
4273 right = qs.dot(&right);
4274 Ok(self.forward_mul(&right))
4275 }
4276
4277 fn design_transformed_transpose_mul(
4278 &self,
4279 qs: &Array2<f64>,
4280 free_basis_opt: Option<&Array2<f64>>,
4281 v: &Array1<f64>,
4282 ) -> Result<Array1<f64>, EstimationError> {
4283 let mut pulled = qs.t().dot(&self.transpose_mul(v));
4284 if let Some(z) = free_basis_opt {
4285 pulled = z.t().dot(&pulled);
4286 }
4287 Ok(pulled)
4288 }
4289
4290 fn penalty_transformed(
4291 &self,
4292 qs: &Array2<f64>,
4293 free_basis_opt: Option<&Array2<f64>>,
4294 ) -> Result<Array2<f64>, EstimationError> {
4295 let dense = self.materialize_dense();
4296 let mut transformed = qs.t().dot(dense).dot(qs);
4297 if let Some(z) = free_basis_opt {
4298 transformed = z.t().dot(&transformed).dot(z);
4299 }
4300 Ok(transformed)
4301 }
4302
4303 fn penalty_scaled_add_to(
4304 &self,
4305 target: &mut Array2<f64>,
4306 amp: f64,
4307 ) -> Result<(), EstimationError> {
4308 let dense = self.materialize_dense();
4309 if target.raw_dim() != dense.raw_dim() {
4310 crate::bail_invalid_estim!(
4311 "implicit hyper penalty derivative shape mismatch: target={}x{}, matrix={}x{}",
4312 target.nrows(),
4313 target.ncols(),
4314 dense.nrows(),
4315 dense.ncols()
4316 );
4317 }
4318 target.scaled_add(amp, dense);
4319 Ok(())
4320 }
4321}
4322
4323impl DerivativeStorageBackend for LatentCoordDerivativeOp {
4324 fn resident_byte_count(&self) -> usize {
4325 0
4326 }
4327 fn design_nrows(&self) -> usize {
4328 self.nrows()
4329 }
4330 fn design_ncols(&self) -> usize {
4331 self.ncols()
4332 }
4333 fn penalty_dim(&self) -> usize {
4334 self.nrows()
4335 }
4336 fn uses_implicit_storage(&self) -> bool {
4337 true
4338 }
4339 fn any_nonzero(&self) -> bool {
4340 true
4341 }
4342 fn materialize(&self) -> Array2<f64> {
4343 self.materialize_dense().clone()
4344 }
4345 fn implicit_first_axis_info(
4346 &self,
4347 ) -> Option<(
4348 std::sync::Arc<gam_terms::basis::ImplicitDesignPsiDerivative>,
4349 usize,
4350 )> {
4351 None
4352 }
4353 fn implicit_axis_count_hint(&self) -> Option<usize> {
4354 Some(self.operator.n_axes())
4355 }
4356
4357 fn design_forward_mul_original(&self, u: &Array1<f64>) -> Result<Array1<f64>, EstimationError> {
4358 if self.ncols() != u.len() {
4359 crate::bail_invalid_estim!(
4360 "latent-coordinate hyper design derivative forward_mul_original width mismatch: operator_cols={}, vector={}",
4361 self.ncols(),
4362 u.len()
4363 );
4364 }
4365 Ok(self.forward_mul(u))
4366 }
4367
4368 fn design_transpose_mul_original(
4369 &self,
4370 v: &Array1<f64>,
4371 ) -> Result<Array1<f64>, EstimationError> {
4372 if self.nrows() != v.len() {
4373 crate::bail_invalid_estim!(
4374 "latent-coordinate hyper design derivative transpose_mul_original height mismatch: operator_rows={}, vector={}",
4375 self.nrows(),
4376 v.len()
4377 );
4378 }
4379 Ok(self.transpose_mul(v))
4380 }
4381
4382 fn design_transformed(
4383 &self,
4384 qs: &Array2<f64>,
4385 free_basis_opt: Option<&Array2<f64>>,
4386 ) -> Result<Array2<f64>, EstimationError> {
4387 let dense = self.materialize_dense();
4388 Ok(gam_linalg::matrix::DenseRightProductView::new(dense)
4389 .with_factor(qs)
4390 .with_optional_factor(free_basis_opt)
4391 .materialize())
4392 }
4393
4394 fn design_transformed_forward_mul(
4395 &self,
4396 qs: &Array2<f64>,
4397 free_basis_opt: Option<&Array2<f64>>,
4398 u: &Array1<f64>,
4399 ) -> Result<Array1<f64>, EstimationError> {
4400 let mut right = if let Some(z) = free_basis_opt {
4401 z.dot(u)
4402 } else {
4403 u.clone()
4404 };
4405 right = qs.dot(&right);
4406 Ok(self.forward_mul(&right))
4407 }
4408
4409 fn design_transformed_transpose_mul(
4410 &self,
4411 qs: &Array2<f64>,
4412 free_basis_opt: Option<&Array2<f64>>,
4413 v: &Array1<f64>,
4414 ) -> Result<Array1<f64>, EstimationError> {
4415 let mut pulled = qs.t().dot(&self.transpose_mul(v));
4416 if let Some(z) = free_basis_opt {
4417 pulled = z.t().dot(&pulled);
4418 }
4419 Ok(pulled)
4420 }
4421
4422 fn penalty_transformed(
4423 &self,
4424 qs: &Array2<f64>,
4425 free_basis_opt: Option<&Array2<f64>>,
4426 ) -> Result<Array2<f64>, EstimationError> {
4427 let dense = self.materialize_dense();
4428 let mut transformed = qs.t().dot(dense).dot(qs);
4429 if let Some(z) = free_basis_opt {
4430 transformed = z.t().dot(&transformed).dot(z);
4431 }
4432 Ok(transformed)
4433 }
4434
4435 fn penalty_scaled_add_to(
4436 &self,
4437 target: &mut Array2<f64>,
4438 amp: f64,
4439 ) -> Result<(), EstimationError> {
4440 let dense = self.materialize_dense();
4441 if target.raw_dim() != dense.raw_dim() {
4442 crate::bail_invalid_estim!(
4443 "latent-coordinate hyper penalty derivative shape mismatch: target={}x{}, matrix={}x{}",
4444 target.nrows(),
4445 target.ncols(),
4446 dense.nrows(),
4447 dense.ncols()
4448 );
4449 }
4450 target.scaled_add(amp, dense);
4451 Ok(())
4452 }
4453}
4454
4455#[derive(Clone)]
4456pub struct HyperDesignDerivative {
4457 pub(crate) storage: DerivativeMatrixStorage,
4458}
4459
4460impl HyperDesignDerivative {
4461 pub fn zero(nrows: usize, ncols: usize) -> Self {
4462 Self {
4463 storage: DerivativeMatrixStorage::Zero(ZeroDerivativeMatrix::new(nrows, ncols)),
4464 }
4465 }
4466
4467 pub fn from_embedded(
4468 local: Array2<f64>,
4469 global_range: Range<usize>,
4470 total_cols: usize,
4471 ) -> Self {
4472 Self {
4473 storage: DerivativeMatrixStorage::Embedded(EmbeddedDerivativeMatrix::new(
4474 local,
4475 global_range,
4476 total_cols,
4477 )),
4478 }
4479 }
4480
4481 pub fn from_implicit(
4482 operator: std::sync::Arc<gam_terms::basis::ImplicitDesignPsiDerivative>,
4483 level: ImplicitDerivLevel,
4484 global_range: Range<usize>,
4485 total_cols: usize,
4486 ) -> Self {
4487 Self {
4488 storage: DerivativeMatrixStorage::Implicit(ImplicitDerivativeOp {
4489 operator,
4490 level,
4491 global_range,
4492 total_dim: total_cols,
4493 cached_dense: std::sync::Arc::new(gam_runtime::resource::RayonSafeOnce::new()),
4494 }),
4495 }
4496 }
4497
4498 pub fn from_latent_coord(
4499 operator: std::sync::Arc<gam_terms::basis::LatentCoordDesignDerivative>,
4500 flat_axis: usize,
4501 global_range: Range<usize>,
4502 total_cols: usize,
4503 ) -> Self {
4504 Self {
4505 storage: DerivativeMatrixStorage::LatentCoord(LatentCoordDerivativeOp {
4506 operator,
4507 flat_axis,
4508 global_range,
4509 total_dim: total_cols,
4510 cached_dense: std::sync::Arc::new(gam_runtime::resource::RayonSafeOnce::new()),
4511 }),
4512 }
4513 }
4514
4515 pub(crate) fn resident_byte_count(&self) -> usize {
4516 storage_dispatch!(&self.storage, b => b.resident_byte_count())
4517 }
4518
4519 pub(crate) fn nrows(&self) -> usize {
4520 storage_dispatch!(&self.storage, b => b.design_nrows())
4521 }
4522
4523 pub(crate) fn ncols(&self) -> usize {
4524 storage_dispatch!(&self.storage, b => b.design_ncols())
4525 }
4526
4527 pub(crate) fn uses_implicit_storage(&self) -> bool {
4528 storage_dispatch!(&self.storage, b => b.uses_implicit_storage())
4529 }
4530
4531 pub(crate) fn materialize(&self) -> Array2<f64> {
4532 storage_dispatch!(&self.storage, b => b.materialize())
4533 }
4534
4535 pub(crate) fn any_nonzero(&self) -> bool {
4536 storage_dispatch!(&self.storage, b => b.any_nonzero())
4537 }
4538
4539 pub(crate) fn forward_mul_original(
4540 &self,
4541 u: &Array1<f64>,
4542 ) -> Result<Array1<f64>, EstimationError> {
4543 storage_dispatch!(&self.storage, b => b.design_forward_mul_original(u))
4544 }
4545
4546 pub(crate) fn transpose_mul_original(
4547 &self,
4548 v: &Array1<f64>,
4549 ) -> Result<Array1<f64>, EstimationError> {
4550 storage_dispatch!(&self.storage, b => b.design_transpose_mul_original(v))
4551 }
4552
4553 pub(crate) fn transformed(
4554 &self,
4555 qs: &Array2<f64>,
4556 free_basis_opt: Option<&Array2<f64>>,
4557 ) -> Result<Array2<f64>, EstimationError> {
4558 storage_dispatch!(&self.storage, b => b.design_transformed(qs, free_basis_opt))
4559 }
4560
4561 pub(crate) fn transformed_forward_mul(
4562 &self,
4563 qs: &Array2<f64>,
4564 free_basis_opt: Option<&Array2<f64>>,
4565 u: &Array1<f64>,
4566 ) -> Result<Array1<f64>, EstimationError> {
4567 storage_dispatch!(&self.storage, b => b.design_transformed_forward_mul(qs, free_basis_opt, u))
4568 }
4569
4570 pub(crate) fn transformed_transpose_mul(
4571 &self,
4572 qs: &Array2<f64>,
4573 free_basis_opt: Option<&Array2<f64>>,
4574 v: &Array1<f64>,
4575 ) -> Result<Array1<f64>, EstimationError> {
4576 storage_dispatch!(&self.storage, b => b.design_transformed_transpose_mul(qs, free_basis_opt, v))
4577 }
4578
4579 pub(crate) fn implicit_first_axis_info(
4584 &self,
4585 ) -> Option<(
4586 std::sync::Arc<gam_terms::basis::ImplicitDesignPsiDerivative>,
4587 usize,
4588 )> {
4589 storage_dispatch!(&self.storage, b => b.implicit_first_axis_info())
4590 }
4591
4592 pub(crate) fn implicit_axis_count_hint(&self) -> Option<usize> {
4593 storage_dispatch!(&self.storage, b => b.implicit_axis_count_hint())
4594 }
4595}
4596
4597impl From<Array2<f64>> for HyperDesignDerivative {
4598 fn from(value: Array2<f64>) -> Self {
4599 Self {
4600 storage: DerivativeMatrixStorage::Dense(value),
4601 }
4602 }
4603}
4604
4605#[derive(Clone)]
4606pub struct HyperPenaltyDerivative {
4607 pub(crate) storage: DerivativeMatrixStorage,
4608}
4609
4610impl HyperPenaltyDerivative {
4611 pub fn from_embedded(local: Array2<f64>, global_range: Range<usize>, total_dim: usize) -> Self {
4612 Self {
4613 storage: DerivativeMatrixStorage::Embedded(EmbeddedDerivativeMatrix::new(
4614 local,
4615 global_range,
4616 total_dim,
4617 )),
4618 }
4619 }
4620
4621 pub(crate) fn resident_byte_count(&self) -> usize {
4622 storage_dispatch!(&self.storage, b => b.resident_byte_count())
4623 }
4624
4625 pub(crate) fn nrows(&self) -> usize {
4626 storage_dispatch!(&self.storage, b => b.penalty_dim())
4627 }
4628
4629 pub(crate) fn ncols(&self) -> usize {
4630 self.nrows()
4631 }
4632
4633 pub(crate) fn scaled_materialize(&self, amp: f64) -> Array2<f64> {
4634 let mut out = Array2::<f64>::zeros((self.nrows(), self.ncols()));
4635 self.scaled_add_to(&mut out, amp)
4636 .expect("scaled materialize uses matching target shape");
4637 out
4638 }
4639
4640 pub(crate) fn transformed(
4641 &self,
4642 qs: &Array2<f64>,
4643 free_basis_opt: Option<&Array2<f64>>,
4644 ) -> Result<Array2<f64>, EstimationError> {
4645 storage_dispatch!(&self.storage, b => b.penalty_transformed(qs, free_basis_opt))
4646 }
4647
4648 pub(crate) fn scaled_add_to(
4649 &self,
4650 target: &mut Array2<f64>,
4651 amp: f64,
4652 ) -> Result<(), EstimationError> {
4653 storage_dispatch!(&self.storage, b => b.penalty_scaled_add_to(target, amp))
4654 }
4655}
4656
4657impl From<Array2<f64>> for HyperPenaltyDerivative {
4658 fn from(value: Array2<f64>) -> Self {
4659 Self {
4660 storage: DerivativeMatrixStorage::Dense(value),
4661 }
4662 }
4663}
4664
4665#[derive(Clone)]
4666pub struct PenaltyDerivativeComponent {
4667 pub penalty_index: usize,
4668 pub matrix: HyperPenaltyDerivative,
4669}
4670
4671#[derive(Clone)]
4672pub struct DirectionalHyperParam {
4673 pub(crate) x_tau_original: HyperDesignDerivative,
4674 pub(crate) penalty_first_components: Vec<PenaltyDerivativeComponent>,
4677 pub(crate) x_tau_tau_original: Option<Vec<Option<HyperDesignDerivative>>>,
4681 pub(crate) penaltysecond_components: Option<Vec<Option<Vec<PenaltyDerivativeComponent>>>>,
4684 pub(crate) penaltysecond_component_provider: Option<
4685 std::sync::Arc<
4686 dyn Fn(usize) -> Result<Option<Vec<PenaltyDerivativeComponent>>, EstimationError>
4687 + Send
4688 + Sync
4689 + 'static,
4690 >,
4691 >,
4692 pub(crate) penaltysecond_partner_indices: Option<std::sync::Arc<[usize]>>,
4693 pub(crate) is_penalty_like: bool,
4697}
4698
4699impl DirectionalHyperParam {
4700 pub(crate) fn resident_byte_count(&self) -> usize {
4701 let mut bytes = self.x_tau_original.resident_byte_count();
4702 for component in &self.penalty_first_components {
4703 bytes = bytes.saturating_add(component.matrix.resident_byte_count());
4704 }
4705 if let Some(entries) = self.x_tau_tau_original.as_ref() {
4706 for entry in entries.iter().flatten() {
4707 bytes = bytes.saturating_add(entry.resident_byte_count());
4708 }
4709 }
4710 if let Some(rows) = self.penaltysecond_components.as_ref() {
4711 for components in rows.iter().flatten() {
4712 for component in components {
4713 bytes = bytes.saturating_add(component.matrix.resident_byte_count());
4714 }
4715 }
4716 }
4717 bytes
4718 }
4719
4720 pub(crate) fn canonicalize_penalty_components(
4721 components: Vec<(usize, HyperPenaltyDerivative)>,
4722 ) -> Result<Vec<PenaltyDerivativeComponent>, EstimationError> {
4723 let mut out: Vec<PenaltyDerivativeComponent> = Vec::with_capacity(components.len());
4724 for (penalty_index, matrix) in components {
4725 if out.iter().any(|c| c.penalty_index == penalty_index) {
4726 crate::bail_invalid_estim!(
4727 "duplicate penalty derivative component for penalty {}",
4728 penalty_index
4729 );
4730 }
4731 out.push(PenaltyDerivativeComponent {
4732 penalty_index,
4733 matrix,
4734 });
4735 }
4736 Ok(out)
4737 }
4738
4739 pub fn new_compact(
4740 x_tau_original: HyperDesignDerivative,
4741 penalty_first_components: Vec<(usize, HyperPenaltyDerivative)>,
4742 x_tau_tau_original: Option<Vec<Option<HyperDesignDerivative>>>,
4743 penaltysecond_components: Option<Vec<Option<Vec<(usize, HyperPenaltyDerivative)>>>>,
4744 ) -> Result<Self, EstimationError> {
4745 let is_penalty_like = !x_tau_original.any_nonzero();
4746 let penalty_first_components =
4747 Self::canonicalize_penalty_components(penalty_first_components)?;
4748 let penaltysecond_components = match penaltysecond_components {
4749 Some(rows) => {
4750 let mut out = Vec::with_capacity(rows.len());
4751 for row in rows {
4752 out.push(match row {
4753 Some(components) => {
4754 Some(Self::canonicalize_penalty_components(components)?)
4755 }
4756 None => None,
4757 });
4758 }
4759 Some(out)
4760 }
4761 None => None,
4762 };
4763 Ok(Self {
4764 x_tau_original,
4765 penalty_first_components,
4766 x_tau_tau_original,
4767 penaltysecond_components,
4768 penaltysecond_component_provider: None,
4769 penaltysecond_partner_indices: None,
4770 is_penalty_like,
4771 })
4772 }
4773
4774 pub fn not_penalty_like(mut self) -> Self {
4777 self.is_penalty_like = false;
4778 self
4779 }
4780
4781 pub fn with_penaltysecond_component_provider(
4782 mut self,
4783 provider: std::sync::Arc<
4784 dyn Fn(usize) -> Result<Option<Vec<PenaltyDerivativeComponent>>, EstimationError>
4785 + Send
4786 + Sync
4787 + 'static,
4788 >,
4789 ) -> Self {
4790 self.penaltysecond_component_provider = Some(provider);
4791 self
4792 }
4793
4794 pub fn with_penaltysecond_partner_indices(mut self, partners: Vec<usize>) -> Self {
4795 self.penaltysecond_partner_indices = Some(std::sync::Arc::from(partners));
4796 self
4797 }
4798
4799 pub(crate) fn x_tau_dense(&self) -> Array2<f64> {
4800 self.x_tau_original.materialize()
4801 }
4802
4803 pub(crate) fn transformed_x_tau(
4804 &self,
4805 qs: &Array2<f64>,
4806 free_basis_opt: Option<&Array2<f64>>,
4807 ) -> Result<Array2<f64>, EstimationError> {
4808 self.x_tau_original.transformed(qs, free_basis_opt)
4809 }
4810
4811 pub(crate) fn x_tau_tau_entry_at(&self, j: usize) -> Option<HyperDesignDerivative> {
4812 self.x_tau_tau_original
4813 .as_ref()
4814 .and_then(|rows| rows.get(j))
4815 .and_then(|entry| entry.clone())
4816 }
4817
4818 pub(crate) fn has_implicit_operator(&self) -> bool {
4821 self.x_tau_original.uses_implicit_storage()
4822 }
4823
4824 pub(crate) fn has_implicit_multidim_duchon(&self) -> bool {
4825 self.implicit_first_axis_info()
4826 .is_some_and(|(op, _)| op.n_axes() > 1 && op.is_duchon_family())
4827 }
4828
4829 pub(crate) fn implicit_first_axis_info(
4831 &self,
4832 ) -> Option<(
4833 std::sync::Arc<gam_terms::basis::ImplicitDesignPsiDerivative>,
4834 usize,
4835 )> {
4836 self.x_tau_original.implicit_first_axis_info()
4837 }
4838
4839 pub(crate) fn implicit_axis_count_hint(&self) -> Option<usize> {
4840 self.x_tau_original.implicit_axis_count_hint()
4841 }
4842
4843 pub(crate) fn penalty_first_components(&self) -> &[PenaltyDerivativeComponent] {
4844 &self.penalty_first_components
4845 }
4846
4847 pub(crate) fn penalty_total_at(
4848 &self,
4849 rho: &Array1<f64>,
4850 p: usize,
4851 ) -> Result<Array2<f64>, EstimationError> {
4852 let mut out = Array2::<f64>::zeros((p, p));
4853 for component in &self.penalty_first_components {
4854 if component.matrix.nrows() != p || component.matrix.ncols() != p {
4855 crate::bail_invalid_estim!(
4856 "S_tau shape mismatch for penalty {}: expected {}x{}, got {}x{}",
4857 component.penalty_index,
4858 p,
4859 p,
4860 component.matrix.nrows(),
4861 component.matrix.ncols()
4862 );
4863 }
4864 if component.penalty_index >= rho.len() {
4865 crate::bail_invalid_estim!(
4866 "penalty_index {} out of bounds for rho dimension {}",
4867 component.penalty_index,
4868 rho.len()
4869 );
4870 }
4871 let lambda = gam_problem::checked_exp_log_strength(rho[component.penalty_index])
4872 .map_err(|error| EstimationError::InvalidInput(error.to_string()))?;
4873 component.matrix.scaled_add_to(&mut out, lambda)?;
4874 }
4875 Ok(out)
4876 }
4877
4878 pub(crate) fn penaltysecond_components_for(
4879 &self,
4880 j: usize,
4881 ) -> Result<Option<Vec<PenaltyDerivativeComponent>>, EstimationError> {
4882 if let Some(components) = self
4883 .penaltysecond_components
4884 .as_ref()
4885 .and_then(|rows| rows.get(j))
4886 .and_then(|row| row.clone())
4887 {
4888 return Ok(Some(components));
4889 }
4890 if let Some(provider) = self.penaltysecond_component_provider.as_ref() {
4891 return provider(j);
4892 }
4893 Ok(None)
4894 }
4895
4896 pub(crate) fn penaltysecond_componentrows(
4897 &self,
4898 ) -> Option<&[Option<Vec<PenaltyDerivativeComponent>>]> {
4899 self.penaltysecond_components.as_deref()
4900 }
4901
4902 pub(crate) fn penalty_first_component_count(&self) -> usize {
4903 self.penalty_first_components.len()
4904 }
4905
4906 pub(crate) fn has_penaltysecond_pair_at(&self, j: usize) -> bool {
4907 self.penaltysecond_components
4908 .as_ref()
4909 .and_then(|rows| rows.get(j))
4910 .is_some_and(Option::is_some)
4911 || self
4912 .penaltysecond_partner_indices
4913 .as_ref()
4914 .is_some_and(|partners| partners.contains(&j))
4915 }
4916}
4917
4918#[derive(Clone, Debug)]
4919pub(crate) struct SparseRemlDecision {
4920 pub(crate) geometry: RemlGeometry,
4921 pub(crate) reason: &'static str,
4922 pub(crate) p: usize,
4923 pub(crate) nnz_x: Option<usize>,
4930 pub(crate) nnz_h_upper_est: Option<usize>,
4931 pub(crate) density_h_upper_est: Option<f64>,
4932}
4933
4934#[derive(Clone)]
4935pub(crate) struct SparseExactEvalData {
4936 pub(crate) factor: Arc<SparseExactFactor>,
4937 pub(crate) takahashi: Option<Arc<gam_linalg::sparse_exact::TakahashiInverse>>,
4938 pub(crate) logdet_h: f64,
4939 pub(crate) logdet_s_pos: f64,
4940 pub(crate) penalty_rank: usize,
4941 pub(crate) det1_values: Arc<Array1<f64>>,
4942}
4943
4944#[derive(Clone)]
4945pub struct FirthDenseOperator {
4946 pub(crate) x_dense: Array2<f64>,
4973 pub(crate) x_dense_t: Array2<f64>,
4974 pub(crate) q_basis: Array2<f64>,
4977 pub(crate) x_reduced: Array2<f64>,
4980 pub(crate) observation_weight_sqrt: Option<Array1<f64>>,
4986 pub(crate) k_reduced: Array2<f64>,
4988 pub(crate) x_metric_reduced_inv_diag: Array1<f64>,
4993 pub(crate) half_log_det: f64,
4995 pub(crate) h_diag: Array1<f64>,
4997 pub(crate) w: Array1<f64>,
4999 pub(crate) w1: Array1<f64>,
5000 pub(crate) w2: Array1<f64>,
5001 pub(crate) w3: Array1<f64>,
5002 pub(crate) w4: Array1<f64>,
5003 pub(crate) b_base: Array2<f64>,
5005 pub(crate) p_b_base: Array2<f64>,
5008}
5009
5010#[derive(Clone)]
5027pub(crate) struct FirthDesignFactor {
5028 pub(crate) x_dense: Array2<f64>,
5030 pub(crate) x_dense_t: Array2<f64>,
5031 pub(crate) q_basis: Array2<f64>,
5033 pub(crate) x_reduced: Array2<f64>,
5035 pub(crate) observation_weight_sqrt: Option<Array1<f64>>,
5037 pub(crate) metric_spectrum: Array1<f64>,
5039 pub(crate) x_metric_reduced_inv_diag: Array1<f64>,
5041 pub(crate) r: usize,
5043 pub(crate) n: usize,
5044}
5045
5046#[derive(Clone)]
5047pub(crate) struct FirthDirection {
5048 pub(crate) deta: Array1<f64>,
5049 pub(crate) g_u_reduced: Array2<f64>,
5050 pub(crate) a_u_reduced: Array2<f64>,
5051 pub(crate) dh: Array1<f64>,
5052 pub(crate) b_uvec: Array1<f64>,
5054}
5055
5056#[derive(Clone)]
5057pub(crate) struct FirthTauPartialKernel {
5058 pub(super) deta_partial: Array1<f64>,
5059 pub(crate) dotw1: Array1<f64>,
5060 pub(crate) dotw2: Array1<f64>,
5061 pub(crate) dot_h_partial: Array1<f64>,
5062 pub(crate) x_tau_reduced: Array2<f64>,
5065 pub(super) dot_i_partial: Array2<f64>,
5066 pub(crate) dot_k_reduced: Array2<f64>,
5070}
5071
5072#[derive(Clone)]
5073pub(crate) struct FirthTauExactKernel {
5074 pub(crate) gphi_tau: Array1<f64>,
5075 pub(crate) phi_tau_partial: f64,
5076 pub(crate) tau_kernel: Option<FirthTauPartialKernel>,
5077}
5078
5079#[derive(Clone)]
5091pub(crate) struct FirthTauTauExactKernel {
5092 pub(super) phi_tau_tau_partial: f64,
5093 pub(super) gphi_tau_tau: Array1<f64>,
5094 pub(super) tau_tau_kernel: Option<FirthTauTauPartialKernel>,
5095}
5096
5097#[derive(Clone, Default)]
5110pub(crate) struct FirthTauTauPartialKernel {
5111 pub(super) x_tau_i_reduced: Array2<f64>,
5112 pub(super) x_tau_j_reduced: Array2<f64>,
5113 pub(super) deta_i_partial: Array1<f64>,
5114 pub(super) deta_j_partial: Array1<f64>,
5115 pub(super) dot_h_i_partial: Array1<f64>,
5116 pub(super) dot_h_j_partial: Array1<f64>,
5117 pub(super) dot_k_i_reduced: Array2<f64>,
5118 pub(super) dot_k_j_reduced: Array2<f64>,
5119 pub(super) dot_i_i_partial: Array2<f64>,
5120 pub(super) dot_i_j_partial: Array2<f64>,
5121 pub(super) x_tau_tau_reduced: Option<Array2<f64>>,
5122 pub(super) deta_ij_partial: Option<Array1<f64>>,
5123}
5124
5125#[derive(Clone, Default)]
5133pub(crate) struct FirthTauBetaPartialKernel {
5134 pub(super) x_tau_reduced: Array2<f64>,
5135 pub(super) deta_partial: Array1<f64>,
5136 pub(super) dot_h_partial: Array1<f64>,
5137 pub(super) dot_i_partial: Array2<f64>,
5138 pub(super) dot_k_reduced: Array2<f64>,
5139 pub(super) deta_v: Array1<f64>,
5140 pub(super) deta_tau_v: Array1<f64>,
5141 pub(super) a_v_reduced: Array2<f64>,
5142 pub(super) dh_v: Array1<f64>,
5143 pub(super) b_vvec: Array1<f64>,
5144 pub(super) d_beta_dot_k: Array2<f64>,
5145 pub(super) d_beta_dot_h: Array1<f64>,
5146}
5147
5148#[derive(Clone)]
5159pub(crate) struct EvalShared {
5160 pub(crate) key: Option<Vec<u64>>,
5161 pub(crate) pirls_result: Arc<PirlsResult>,
5162 pub(crate) ridge_passport: RidgePassport,
5163 pub(crate) geometry: SparseRemlDecision,
5173 pub(crate) h_total: Arc<Array2<f64>>,
5177 pub(crate) sparse_exact: Option<Arc<SparseExactEvalData>>,
5178 pub(crate) firth_dense_operator: Option<Arc<FirthDenseOperator>>,
5179 pub(crate) firth_dense_operator_original: Option<Arc<FirthDenseOperator>>,
5182 pub(crate) penalty_pseudologdet: std::sync::OnceLock<Arc<penalty_logdet::PenaltyPseudologdet>>,
5196 pub(crate) root_scale_hessian_logdet: std::sync::OnceLock<Option<f64>>,
5204 pub(crate) applied_canonical_penalties:
5223 std::sync::OnceLock<Arc<Vec<gam_terms::construction::CanonicalPenalty>>>,
5224 pub(crate) penalty_scores_at_mode: std::sync::OnceLock<Arc<Vec<Array1<f64>>>>,
5237 pub(crate) block_local_correction: std::sync::OnceLock<(
5262 usize,
5263 Arc<outer_eval::TkCorrectionTerms>,
5264 Option<crate::estimate::outer_eval_capture::QuadratureMarginalAudit>,
5265 )>,
5266}
5267
5268pub(crate) fn applied_canonical_penalties_for(
5282 reparam: &gam_terms::construction::ReparamResult,
5283 canonical_penalties: &Arc<Vec<gam_terms::construction::CanonicalPenalty>>,
5284) -> Result<Arc<Vec<gam_terms::construction::CanonicalPenalty>>, EstimationError> {
5285 let split = reparam.null_split();
5286 if split.declared_null_dim() == 0 || canonical_penalties.is_empty() {
5287 return Ok(Arc::clone(canonical_penalties));
5288 }
5289 let projected = canonical_penalties
5290 .iter()
5291 .map(|penalty| {
5292 split.project_canonical(penalty, gam_terms::construction::PenaltyFrame::Original)
5293 })
5294 .collect::<Result<Vec<_>, _>>()
5295 .map_err(|error| {
5296 EstimationError::LayoutError(format!(
5297 "projecting the canonical penalties onto the reparameterization's penalized \
5298 subspace failed: {error}"
5299 ))
5300 })?;
5301 if projected
5305 .iter()
5306 .zip(canonical_penalties.iter())
5307 .all(|(a, b)| a.root == b.root && a.col_range == b.col_range)
5308 {
5309 Ok(Arc::clone(canonical_penalties))
5310 } else {
5311 Ok(Arc::new(projected))
5312 }
5313}
5314
5315impl EvalShared {
5316 pub(crate) fn matches(&self, key: &Option<Vec<u64>>) -> bool {
5317 match (&self.key, key) {
5318 (None, None) => true,
5319 (Some(a), Some(b)) => a == b,
5320 _ => false,
5321 }
5322 }
5323
5324 pub(crate) fn applied_canonical_penalties(
5348 &self,
5349 canonical_penalties: &Arc<Vec<gam_terms::construction::CanonicalPenalty>>,
5350 ) -> Result<Arc<Vec<gam_terms::construction::CanonicalPenalty>>, EstimationError> {
5351 if let Some(applied) = self.applied_canonical_penalties.get() {
5352 return Ok(Arc::clone(applied));
5353 }
5354 let applied = applied_canonical_penalties_for(
5355 &self.pirls_result.reparam_result,
5356 canonical_penalties,
5357 )?;
5358 match self.applied_canonical_penalties.set(Arc::clone(&applied)) {
5359 Ok(()) => Ok(applied),
5360 Err(_) => Ok(Arc::clone(
5361 self.applied_canonical_penalties
5362 .get()
5363 .expect("OnceLock set raced, so it is initialized"),
5364 )),
5365 }
5366 }
5367
5368 pub(crate) fn penalty_pseudologdet_original(
5369 &self,
5370 canonical_penalties: &Arc<Vec<gam_terms::construction::CanonicalPenalty>>,
5371 lambdas: &[f64],
5372 p: usize,
5373 ) -> Result<Arc<penalty_logdet::PenaltyPseudologdet>, EstimationError> {
5374 if let Some(pld) = self.penalty_pseudologdet.get() {
5375 if pld.dim() != p {
5376 return Err(EstimationError::LayoutError(format!(
5377 "shared penalty pseudo-logdet frame mismatch: cached p={}, requested p={}",
5378 pld.dim(),
5379 p
5380 )));
5381 }
5382 return Ok(Arc::clone(pld));
5383 }
5384 let applied = self.applied_canonical_penalties(canonical_penalties)?;
5387 let pld = Arc::new(
5388 penalty_logdet::PenaltyPseudologdet::from_penalties(
5389 &applied,
5390 lambdas,
5391 self.ridge_passport.penalty_logdet_ridge(),
5392 p,
5393 )
5394 .map_err(EstimationError::InvalidInput)?,
5395 );
5396 match self.penalty_pseudologdet.set(Arc::clone(&pld)) {
5397 Ok(()) => Ok(pld),
5398 Err(_) => Ok(Arc::clone(
5402 self.penalty_pseudologdet
5403 .get()
5404 .expect("OnceLock set raced, so it is initialized"),
5405 )),
5406 }
5407 }
5408}
5409
5410impl PenalizedGeometry for EvalShared {
5411 fn backend_kind(&self) -> GeometryBackendKind {
5412 match self.geometry.geometry {
5413 RemlGeometry::DenseSpectral => GeometryBackendKind::DenseSpectral,
5414 RemlGeometry::SparseExactSpd => GeometryBackendKind::SparseExactSpd,
5415 }
5416 }
5417}
5418
5419impl SparseRemlDecision {
5420 pub(crate) fn basis(&self) -> String {
5429 format!(
5430 "reason={} p={} nnz_x={} nnz_h_est={} density_h_est={} threshold={:.4}",
5431 self.reason,
5432 self.p,
5433 self.nnz_x
5434 .map(|value| value.to_string())
5435 .unwrap_or_else(|| "na".to_string()),
5436 self.nnz_h_upper_est
5437 .map(|value| value.to_string())
5438 .unwrap_or_else(|| "na".to_string()),
5439 self.density_h_upper_est
5440 .map(|value| format!("{value:.4}"))
5441 .unwrap_or_else(|| "na".to_string()),
5442 RemlState::SPARSE_HESSIAN_MAX_DENSITY,
5443 )
5444 }
5445}
5446
5447pub(crate) struct PirlsLruCache {
5457 pub(crate) map: HashMap<Vec<u64>, (Arc<PirlsResult>, u64, usize)>,
5459 pub(crate) byte_budget: usize,
5460 pub(crate) current_bytes: usize,
5461 pub(crate) clock: u64,
5462}
5463
5464impl PirlsLruCache {
5465 pub(crate) fn new(byte_budget: usize) -> Self {
5466 Self {
5467 map: HashMap::new(),
5468 byte_budget: byte_budget.max(1),
5469 current_bytes: 0,
5470 clock: 0,
5471 }
5472 }
5473
5474 pub(crate) fn get(&mut self, key: &Vec<u64>) -> Option<Arc<PirlsResult>> {
5475 if let Some(entry) = self.map.get_mut(key) {
5476 self.clock += 1;
5477 entry.1 = self.clock;
5478 Some(entry.0.clone())
5479 } else {
5480 None
5481 }
5482 }
5483
5484 pub(crate) fn insert(&mut self, key: Vec<u64>, value: Arc<PirlsResult>) {
5485 self.clock += 1;
5486 let bytes = pirls_result_cache_bytes(&value);
5487 if bytes > self.byte_budget {
5491 if let Some((_, _, prev_bytes)) = self.map.remove(&key) {
5492 self.current_bytes = self.current_bytes.saturating_sub(prev_bytes);
5493 }
5494 return;
5495 }
5496 if let Some((_, _, prev_bytes)) = self.map.remove(&key) {
5497 self.current_bytes = self.current_bytes.saturating_sub(prev_bytes);
5498 }
5499 while self.current_bytes + bytes > self.byte_budget {
5500 let evict_key = self
5501 .map
5502 .iter()
5503 .min_by_key(|(_, (_, ts, _))| *ts)
5504 .map(|(k, _)| k.clone());
5505 match evict_key {
5506 Some(k) => {
5507 if let Some((_, _, evict_bytes)) = self.map.remove(&k) {
5508 self.current_bytes = self.current_bytes.saturating_sub(evict_bytes);
5509 }
5510 }
5511 None => break,
5512 }
5513 }
5514 self.current_bytes += bytes;
5515 self.map.insert(key, (value, self.clock, bytes));
5516 }
5517
5518 pub(crate) fn clear(&mut self) {
5519 self.map.clear();
5520 self.current_bytes = 0;
5521 }
5522}
5523
5524#[derive(Clone, Copy, PartialEq, Eq)]
5525pub(crate) struct PenaltySubspaceCacheKey {
5526 pub(crate) penalty_matrix_fingerprint: u64,
5527 pub(crate) ridge_passport_signature: u64,
5528}
5529
5530pub(crate) struct PenaltySubspaceCache {
5531 pub(crate) entry: Option<(PenaltySubspaceCacheKey, Arc<outer_eval::PenaltySubspace>)>,
5532}
5533
5534impl PenaltySubspaceCache {
5535 pub(crate) fn new() -> Self {
5536 Self { entry: None }
5537 }
5538
5539 pub(crate) fn get(
5540 &self,
5541 key: &PenaltySubspaceCacheKey,
5542 ) -> Option<Arc<outer_eval::PenaltySubspace>> {
5543 self.entry
5544 .as_ref()
5545 .filter(|(cached_key, _)| cached_key == key)
5546 .map(|(_, value)| value.clone())
5547 }
5548
5549 pub(crate) fn insert(
5550 &mut self,
5551 key: PenaltySubspaceCacheKey,
5552 value: Arc<outer_eval::PenaltySubspace>,
5553 ) {
5554 self.entry = Some((key, value));
5555 }
5556
5557 pub(crate) fn clear(&mut self) {
5558 self.entry = None;
5559 }
5560}
5561
5562impl PenaltySubspaceCacheKey {
5563 pub(crate) fn from_inputs(
5568 e_transformed: &ndarray::Array2<f64>,
5569 ridge_passport: &gam_problem::RidgePassport,
5570 ) -> Self {
5571 use std::collections::hash_map::DefaultHasher;
5572 use std::hash::{Hash, Hasher};
5573 let mut hasher = DefaultHasher::new();
5574 e_transformed.nrows().hash(&mut hasher);
5575 e_transformed.ncols().hash(&mut hasher);
5576 for value in e_transformed.iter() {
5577 value.to_bits().hash(&mut hasher);
5578 }
5579 let penalty_matrix_fingerprint = hasher.finish();
5580 let mut ridge_hasher = DefaultHasher::new();
5581 ridge_passport.delta().to_bits().hash(&mut ridge_hasher);
5582 ridge_passport.matrix_form().hash(&mut ridge_hasher);
5583 ridge_passport.policy().hash(&mut ridge_hasher);
5584 let ridge_passport_signature = ridge_hasher.finish();
5585 Self {
5586 penalty_matrix_fingerprint,
5587 ridge_passport_signature,
5588 }
5589 }
5590}
5591
5592pub(crate) fn pirls_result_cache_bytes(result: &PirlsResult) -> usize {
5607 use std::mem::size_of;
5608 let n_array_elems = result.final_eta.len()
5609 + result.solveweights.len()
5610 + result.solveworking_response.len()
5611 + result.solvemu.len()
5612 + result.solve_c_array.len()
5613 + result.solve_d_array.len();
5614 let p = result.beta_transformed.0.len();
5615 let pen_h = symmetric_matrix_cache_bytes(&result.penalized_hessian_transformed);
5616 let stab_h = symmetric_matrix_cache_bytes(&result.stabilizedhessian_transformed);
5617 let reparam = (result.reparam_result.s_transformed.len()
5618 + result.reparam_result.qs.len()
5619 + result.reparam_result.e_transformed.len()
5620 + result.reparam_result.det1.len())
5621 * size_of::<f64>();
5622 n_array_elems * size_of::<f64>() + p * size_of::<f64>() + pen_h + stab_h + reparam + 1024
5623}
5624
5625pub(crate) fn symmetric_matrix_cache_bytes(m: &gam_linalg::matrix::SymmetricMatrix) -> usize {
5626 use gam_linalg::matrix::SymmetricMatrix;
5627 use std::mem::size_of;
5628 match m {
5629 SymmetricMatrix::Dense(a) => a.len() * size_of::<f64>(),
5630 SymmetricMatrix::Sparse(s) => {
5631 let (symbolic, values) = s.parts();
5633 values.len() * (size_of::<f64>() + size_of::<usize>())
5634 + std::mem::size_of_val(symbolic.col_ptr())
5635 }
5636 }
5637}
5638
5639pub(crate) const OUTER_EVAL_LRU_CAPACITY: usize = 8;
5647
5648pub(crate) struct OuterEvalLru {
5661 capacity: usize,
5662 entries: std::collections::VecDeque<(Vec<u64>, OuterEval)>,
5664}
5665
5666impl OuterEvalLru {
5667 pub(crate) fn new(capacity: usize) -> Self {
5668 Self {
5669 capacity: capacity.max(1),
5670 entries: std::collections::VecDeque::new(),
5671 }
5672 }
5673
5674 pub(crate) fn get(&mut self, key: &[u64]) -> Option<OuterEval> {
5678 let pos = self.entries.iter().position(|(k, _)| k.as_slice() == key)?;
5679 let entry = self.entries.remove(pos)?;
5680 let eval = entry.1.clone();
5681 self.entries.push_back(entry);
5682 Some(eval)
5683 }
5684
5685 pub(crate) fn insert(&mut self, key: Vec<u64>, eval: OuterEval) {
5688 if let Some(pos) = self
5689 .entries
5690 .iter()
5691 .position(|(k, _)| k.as_slice() == key.as_slice())
5692 {
5693 self.entries.remove(pos);
5694 }
5695 self.entries.push_back((key, eval));
5696 while self.entries.len() > self.capacity {
5697 self.entries.pop_front();
5698 }
5699 }
5700
5701 pub(crate) fn clear(&mut self) {
5702 self.entries.clear();
5703 }
5704}
5705
5706pub(crate) struct EvalCacheManager {
5711 pub(crate) pirls_cache: RwLock<PirlsLruCache>,
5712 pub(crate) penalty_subspace_cache: RwLock<PenaltySubspaceCache>,
5713 pub(crate) current_eval_bundle: RwLock<Option<EvalShared>>,
5714 pub(crate) current_outer_eval: RwLock<Option<(Vec<u64>, OuterEval)>>,
5718 pub(crate) outer_eval_lru: RwLock<OuterEvalLru>,
5733 pub(crate) pirls_cache_enabled: AtomicBool,
5734}
5735
5736impl EvalCacheManager {
5737 pub(crate) fn new() -> Self {
5738 Self {
5739 pirls_cache: RwLock::new(PirlsLruCache::new(PIRLS_CACHE_BYTE_BUDGET)),
5740 penalty_subspace_cache: RwLock::new(PenaltySubspaceCache::new()),
5741 current_eval_bundle: RwLock::new(None),
5742 current_outer_eval: RwLock::new(None),
5743 outer_eval_lru: RwLock::new(OuterEvalLru::new(OUTER_EVAL_LRU_CAPACITY)),
5744 pirls_cache_enabled: AtomicBool::new(true),
5745 }
5746 }
5747
5748 pub(super) fn cached_penalty_subspace<F>(
5755 &self,
5756 e_transformed: &ndarray::Array2<f64>,
5757 ridge_passport: &gam_problem::RidgePassport,
5758 build: F,
5759 ) -> Result<Arc<outer_eval::PenaltySubspace>, EstimationError>
5760 where
5761 F: FnOnce() -> Result<outer_eval::PenaltySubspace, EstimationError>,
5762 {
5763 let key = PenaltySubspaceCacheKey::from_inputs(e_transformed, ridge_passport);
5764 if let Some(hit) = self
5765 .penalty_subspace_cache
5766 .read()
5767 .expect("penalty-subspace cache lock is poisoned: a writer panicked while holding it")
5768 .get(&key)
5769 {
5770 return Ok(hit);
5771 }
5772 let value = Arc::new(build()?);
5773 self.penalty_subspace_cache
5774 .write()
5775 .expect("penalty-subspace cache lock is poisoned: a writer panicked while holding it")
5776 .insert(key, value.clone());
5777 Ok(value)
5778 }
5779
5780 pub(crate) fn cached_eval_bundle(&self, key: &Option<Vec<u64>>) -> Option<EvalShared> {
5781 let guard = self
5782 .current_eval_bundle
5783 .read()
5784 .expect("current eval bundle lock is poisoned: a writer panicked while holding it");
5785 let bundle: &EvalShared = guard.as_ref()?;
5786 bundle.matches(key).then(|| bundle.clone())
5787 }
5788
5789 pub(crate) fn store_eval_bundle(&self, bundle: EvalShared) {
5790 *self
5791 .current_eval_bundle
5792 .write()
5793 .expect("current eval bundle lock is poisoned: a writer panicked while holding it") =
5794 Some(bundle);
5795 }
5796
5797 pub(crate) fn cached_outer_eval(&self, key: &Option<Vec<u64>>) -> Option<OuterEval> {
5798 let key = key.as_ref()?;
5799 self.outer_eval_lru
5806 .write()
5807 .expect("outer-eval LRU lock is poisoned: a writer panicked while holding it")
5808 .get(key)
5809 }
5810
5811 pub(crate) fn store_outer_eval(&self, key: &Option<Vec<u64>>, eval: &OuterEval) {
5812 if let Some(key) = key.clone() {
5813 *self.current_outer_eval.write().expect(
5817 "current outer eval lock is poisoned: a writer panicked while holding it",
5818 ) = Some((key.clone(), eval.clone()));
5819 self.outer_eval_lru
5820 .write()
5821 .expect("outer-eval LRU lock is poisoned: a writer panicked while holding it")
5822 .insert(key, eval.clone());
5823 }
5824 }
5825
5826 pub(crate) fn invalidate_eval_bundle(&self) {
5827 self.current_eval_bundle
5828 .write()
5829 .expect("current eval bundle lock is poisoned: a writer panicked while holding it")
5830 .take();
5831 self.current_outer_eval
5832 .write()
5833 .expect("current outer eval lock is poisoned: a writer panicked while holding it")
5834 .take();
5835 self.outer_eval_lru
5836 .write()
5837 .expect("outer-eval LRU lock is poisoned: a writer panicked while holding it")
5838 .clear();
5839 }
5840
5841 pub(crate) fn clear_eval_and_factor_caches(&self) {
5842 self.invalidate_eval_bundle();
5843 self.penalty_subspace_cache
5844 .write()
5845 .expect("penalty-subspace cache lock is poisoned: a writer panicked while holding it")
5846 .clear();
5847 }
5848}
5849
5850pub(crate) struct RemlArena {
5853 pub(crate) cost_eval_count: RwLock<u64>,
5854 pub(crate) inner_pirls_solve_count: AtomicU64,
5867 pub(crate) lastgradient_used_stochastic_fallback: AtomicBool,
5868}
5869
5870impl RemlArena {
5871 pub(crate) fn new() -> Self {
5872 Self {
5873 cost_eval_count: RwLock::new(0),
5874 inner_pirls_solve_count: AtomicU64::new(0),
5875 lastgradient_used_stochastic_fallback: AtomicBool::new(false),
5876 }
5877 }
5878}
5879
5880pub(crate) struct RemlState<'a> {
5881 pub(crate) y: ArrayView1<'a, f64>,
5882 pub(crate) x: DesignMatrix,
5883 pub(crate) weights: ArrayView1<'a, f64>,
5884 pub(crate) offset: Array1<f64>,
5885 pub(crate) canonical_penalties: Arc<Vec<gam_terms::construction::CanonicalPenalty>>,
5889 pub(crate) balanced_penalty_root: Array2<f64>,
5890 pub(crate) reparam_invariant: ReparamInvariant,
5891 pub(crate) sparse_penalty_block_count: Option<usize>,
5892 pub(crate) p: usize,
5893 pub(crate) config: Arc<RemlConfig>,
5894 pub(crate) runtime_mixture_link_state: Option<gam_problem::MixtureLinkState>,
5895 pub(crate) runtime_sas_link_state: Option<SasLinkState>,
5896 pub(crate) nullspace_dims: Vec<usize>,
5897 pub(crate) coefficient_lower_bounds: Option<Array1<f64>>,
5898 pub(crate) linear_constraints: Option<crate::pirls::LinearInequalityConstraints>,
5899 pub(crate) rho_prior: gam_problem::RhoPrior,
5901
5902 pub(crate) cache_manager: EvalCacheManager,
5903 pub(crate) arena: RemlArena,
5904 pub(crate) warm_start_beta: RwLock<Option<Coefficients>>,
5905 pub(crate) warm_start_rho: RwLock<Option<Array1<f64>>>,
5915 pub(crate) prev_warm_start_beta: RwLock<Option<Coefficients>>,
5916 pub(crate) prev_warm_start_rho: RwLock<Option<Array1<f64>>>,
5917 pub(crate) block_correction_admission: AtomicUsize,
5946 pub(crate) ift_quality_runtime: std::sync::Mutex<outer_eval::IftQualityRuntimeState>,
5966 pub(crate) hypergradient_runtime:
5967 std::sync::Mutex<Option<outer_eval::HyperGradientRuntimeState>>,
5968 pub(crate) ift_mode_response_slot:
5969 std::sync::Mutex<Option<outer_eval::IftModeResponseRuntimeCache>>,
5970 pub(crate) ift_joint_mode_response_slot:
5971 std::sync::Mutex<Option<outer_eval::IftJointModeResponseRuntimeCache>>,
5972 pub(crate) warm_start_enabled: AtomicBool,
5973 pub(crate) screening_max_inner_iterations: Arc<AtomicUsize>,
5974 pub(crate) outer_inner_cap: Arc<AtomicUsize>,
5989
5990 pub(crate) last_inner_iters: Arc<AtomicUsize>,
6003 pub(crate) last_inner_converged: Arc<AtomicBool>,
6004
6005 pub(crate) ift_warm_start_cache: RwLock<Option<IftWarmStartCache>>,
6021
6022 pub(crate) last_pirls_lm_lambda: Arc<AtomicU64>,
6034
6035 pub(crate) frozen_negbin_theta: Arc<AtomicU64>,
6047
6048 pub(crate) frozen_tweedie_phi: Arc<AtomicU64>,
6062
6063 pub(crate) frozen_gamma_shape: Arc<AtomicU64>,
6080
6081 pub(crate) frozen_beta_phi: Arc<AtomicU64>,
6099
6100 pub(crate) last_ift_prediction_residual: Arc<AtomicU64>,
6122
6123 pub(crate) last_pirls_accept_rho: Arc<AtomicU64>,
6138
6139 pub(crate) ift_cached_factor: RwLock<Option<Arc<dyn gam_linalg::matrix::FactorizedSystem>>>,
6150
6151 pub(crate) kronecker_penalty_system: Option<gam_terms::smooth::KroneckerPenaltySystem>,
6155 pub(crate) kronecker_factored: Option<gam_terms::basis::KroneckerFactoredBasis>,
6158
6159 pub(crate) gaussian_fixed_cache: RwLock<Option<Arc<crate::pirls::GaussianFixedCache>>>,
6169 pub(crate) gaussian_cost_only_frozen_rows:
6177 RwLock<Option<Arc<crate::pirls::GaussianFrozenRows>>>,
6178 pub(crate) gaussian_psi_gram_deriv:
6189 RwLock<Option<Arc<(ndarray::Array2<f64>, ndarray::Array1<f64>)>>>,
6190 pub(crate) glm_psi_gram_deriv:
6208 RwLock<Option<Arc<(ndarray::Array2<f64>, ndarray::Array1<f64>)>>>,
6209 pub(crate) glm_first_step_gram: RwLock<Option<Arc<ndarray::Array2<f64>>>>,
6228 pub(crate) flat_glm_first_step_gram:
6243 RwLock<Option<gam_runtime::resource::Governed<Arc<ndarray::Array2<f64>>>>>,
6244 pub(crate) persistent_warm_start_key: RwLock<Option<String>>,
6247 pub(crate) persistent_latent_values_fingerprint: Option<u64>,
6248 pub(crate) persistent_latent_values_cache: RwLock<PersistentLatentValuesCache>,
6249 pub(crate) analytic_penalty_registry_fingerprint: u64,
6250 pub(crate) persistent_warm_start_loaded: AtomicBool,
6252 pub(crate) persistent_warm_start_store_suppression: AtomicUsize,
6257 pub(crate) persistent_warm_start_store:
6264 std::sync::OnceLock<gam_runtime::warm_start::ConfiguredWarmStartStore>,
6265 pub(crate) gaussian_weight_log_sum_half_cache: std::sync::OnceLock<f64>,
6276 pub(crate) gaussian_dp_floor_scale_cache: std::sync::OnceLock<f64>,
6277 pub(crate) positive_weight_observation_count_cache: std::sync::OnceLock<usize>,
6278 pub(crate) rho_weight_anchor_cache: std::sync::OnceLock<f64>,
6279}