Skip to main content

gam_solve/reml/
mod.rs

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;
23// #1521 carve: promoted `pub(crate)` -> `pub` so the extracted
24// `gam-custom-family` crate reaches the Jeffreys-subspace items it consumes.
25pub 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/// Cached state from the most recent successful PIRLS solve, populated by
101/// `updatewarm_start_from` and consumed by the IFT-based warm-start
102/// predictor (`RemlState::predict_warm_start_beta_ift_with_outcome`).
103/// See the field doc on `RemlState::ift_warm_start_cache` for the math.
104#[derive(Clone)]
105pub(crate) struct IftWarmStartCache {
106    /// β at the converged solve, in ORIGINAL basis. Mirror of
107    /// `warm_start_beta` stashed alongside the H factor for atomic
108    /// consistency under concurrent reads (the predictor needs both
109    /// β and H from the SAME solve; reading them from two locks risks
110    /// a torn pair if a fresh solve lands between reads).
111    pub beta_original: ndarray::Array1<f64>,
112    /// ρ at which the solve occurred. Mirror of `warm_start_rho`,
113    /// stashed for the same atomic-consistency reason.
114    pub rho: ndarray::Array1<f64>,
115    /// Penalized Hessian H_pen at the converged β, in TRANSFORMED basis.
116    /// The IFT predictor factors this on demand; basis transforms run in
117    /// transformed basis for numerical stability.
118    pub penalized_hessian_transformed: gam_linalg::matrix::SymmetricMatrix,
119    /// Reparameterization matrix qs converting between transformed
120    /// (column) basis and original basis: `β_orig = qs · β_tfd`,
121    /// `H_orig = qs · H_tfd · qs^T`.
122    pub qs: ndarray::Array2<f64>,
123    /// True when the PIRLS result was already in original basis
124    /// (`OriginalSparseNative`) — in which case `qs` is the identity
125    /// and the IFT predictor can skip the basis-conversion ops.
126    pub frame_was_original: bool,
127    /// Per-penalty precomputation `S_k · β_cur[cp.col_range]`,
128    /// indexed in lockstep with `RemlObjectiveState::canonical_penalties`.
129    /// Each entry is the local-block mat-vec the IFT predictor would
130    /// otherwise recompute on every predict call. With H_pen factor
131    /// caching (commit ec18559d) the per-call cost dropped from
132    /// `O(p³)` Cholesky to `O(p²) ≈ k · O(block²)` rhs construction;
133    /// at large-scale CTN (p ≈ several thousand) that's several ms
134    /// per predict call still being paid. By stashing `S_k · β_cur`
135    /// at cache-write time the predictor's per-call work drops to
136    /// just the `Δρ_k · e^{ρ_k} · sb_block` accumulation, which is
137    /// `O(p)` rather than `O(p²)`.
138    ///
139    /// `None` when the cache predates this commit's writer hook (e.g.,
140    /// transient state during invalidation); the predictor falls back
141    /// to recomputing the mat-vec when this is `None` or the length
142    /// mismatches `canonical_penalties.len()`.
143    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    /// True when the τ-τ exact-Hessian path cannot be assembled at all and the
184    /// eval must fall back to value-and-gradient mode (forcing
185    /// `HessianValue::Unavailable`).
186    ///
187    /// This is the *only* remaining capability gate: the previous
188    /// implementation also forced gradient-only when the design used implicit
189    /// multi-dim Duchon storage or when the dense τ-cache plan would exceed
190    /// the budget.  Both of those are now *cost* gates, not capability gates
191    /// — the unified evaluator's `prefer_outer_hessian_operator(n, p, k)`
192    /// selects the matrix-free `HessianValue::Operator` representation in
193    /// exactly the regimes where the dense cache would be unaffordable, and
194    /// the planner routes operator returns through `run_operator_trust_region`
195    /// (or basis-probes them when `dim ≤ OUTER_HVP_MATERIALIZE_MAX_DIM`).
196    /// Forcing gradient-only would have prevented the operator representation
197    /// from ever being requested, defeating that routing; hence the
198    /// `implicit_multidim_duchon` and cost-bytes clauses are deliberately
199    /// gone.
200    ///
201    /// Firth-pair-terms-unavailability remains a capability gate: when the
202    /// Firth-aware derivative provider cannot produce the τ-τ pair
203    /// corrections at all, no representation choice can substitute.  At every
204    /// production call site this flag is hardcoded `false` (the
205    /// `hphi_tau_tau_partial_apply` + `d_beta_hphi_tau_partial_apply`
206    /// primitives now cover the gap), so this method effectively returns
207    /// `false` in production.  We retain the field and method signature
208    /// unchanged so future Firth corner cases have a single, surfaced place
209    /// to land.
210    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    /// Shorthand for the canonical Binomial-Logit `GlmLikelihoodSpec` used by
355    /// the REML test fixtures.
356    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    /// Shorthand for the canonical Gaussian-Identity `GlmLikelihoodSpec` used
364    /// by the REML test fixtures.
365    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        // Multi-dim Duchon implicit storage used to force gradient-only,
475        // because the τ-cache materialization plan was infeasible.  The
476        // unified evaluator now elects the matrix-free
477        // `HessianValue::Operator` representation in this regime via
478        // `prefer_outer_hessian_operator`, so the planner can route to the
479        // operator trust-region (or basis-probe to dense for small K) — the
480        // capability is preserved and gradient-only must NOT engage.
481        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        // The dense τ-cache plan exceeds the budget, but cost is no longer a
517        // capability gate: the eval-side selects the matrix-free operator
518        // representation in exactly this regime, and the planner routes
519        // accordingly.  `prefer_gradient_only` must NOT force `Unavailable`
520        // here.
521        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    /// Common shape for the design-motion + penalty-motion REML test fixtures
554    /// (Gaussian-identity and binomial-logit at present): both carry the
555    /// same `(y, w, X, S0, cfg, ρ)` plus a perturbation pair, and need the
556    /// same three helpers (`state`, `state_perturbed`, `fd_directional_gradient`).
557    /// Per-fixture `new()` constructors fill the fields with family-specific
558    /// data; the helpers below are shared via default impls so every fixture
559    /// pays the boilerplate once.
560    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        /// Central FD approximation to the directional cost derivative at ρ.
589        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        // The residual correction implemented by the unified evaluator is the
674        // fixed-dispersion LAML identity. Gaussian-identity uses profiled-scale
675        // REML and deliberately does not enter that gate, so use a genuine
676        // fixed-dispersion design-moving model here.
677        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    /// The adaptive IFT step-cap controller belongs to the `RemlState` that
789    /// produced it, not to the ADDRESS that state happened to occupy.
790    ///
791    /// It used to live in a process-global `HashMap` keyed by
792    /// `self as *const _ as usize`. Every constructing caller makes the state a
793    /// stack local, so two fits reached from the same call site land on the
794    /// same frame slot — and with no `Drop` to evict the entry, the second fit
795    /// read the first fit's quality history and shrunken step cap. That is a
796    /// silent cross-fit dependence: the same data fitted twice in one process
797    /// takes a different IFT step schedule the second time.
798    ///
799    /// The probe builds a state, drives the controller off its defaults, drops
800    /// it, and builds a second state in the same slot. The second state must
801    /// see the documented defaults — the caller-supplied cap, and no flat
802    /// fallback.
803    #[test]
804    fn adaptive_ift_controller_does_not_survive_into_the_next_state_in_the_same_slot() {
805        // `default_cap` is deliberately a value the controller can never
806        // produce on its own, so reading it back is proof of a fresh slot
807        // rather than a coincidence.
808        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                // A prediction quality far above the flat-fallback band: the
837                // controller shrinks the cap and arms the flat fallback. This
838                // is precisely the state that used to be inherited.
839                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    /// #2379: the Gaussian profiled-diagonal seed helper honors its (validated)
873    /// ρ-box by CLAMPING into it — never silently swapping or escaping it. The
874    /// helper now takes an `OrderedRhoBounds`, so an inverted box is impossible
875    /// to hand it (refused upstream at construction); this pins that a valid box
876    /// with the upper bound below the natural profiled optimum clamps the emitted
877    /// seed to that upper bound rather than producing an out-of-box value.
878    #[test]
879    fn gaussian_profiled_diagonal_seed_clamps_into_its_validated_box() {
880        // A smooth-ish Gaussian-identity design whose profiled REML optimum for
881        // the summed penalty sits at a moderate, finite ρ.
882        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        // A wide, ordered box: the emitted seed is finite and inside it. Its
917        // value is the natural profiled ρ (nothing binds).
918        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        // A box whose UPPER bound is pinned strictly below the natural optimum:
934        // the profiled ρ must be clamped DOWN to that upper bound, proving the box
935        // is honored (a silent swap would instead have solved a different box).
936        // Derive the cap from the measured optimum so the assertion is robust to
937        // the exact fixture geometry; `cap_lo < cap_hi` and both are finite.
938        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    /// Regression (issue #893): for a fixed-dispersion family a uniform prior
1074    /// weight `w = c` is *exact* `c`-fold row replication. The two encodings must
1075    /// therefore present a byte-identical LAML smoothing-selection surface — both
1076    /// the cost `V(ρ)` and its gradient `∇V(ρ)` — because every term (penalised
1077    /// deviance `D_p`, the working cross-product `XᵀWX`, the log-determinants)
1078    /// is a sum of per-observation contributions that is identical whether a row
1079    /// carries weight `c` or is stacked `c` times. This locks the *surface*
1080    /// invariant that #893 ultimately reduces to: when the surfaces coincide,
1081    /// the only remaining requirement for `λ̂(w=c) = λ̂(c×)` is that the outer
1082    /// optimiser resolve the shared optimum (handled by the tightened outer
1083    /// tolerance in `workflow.rs`). A regression that reintroduced a
1084    /// row-count-vs-weight-sum asymmetry into the inner solve or the cost would
1085    /// break this directly, independent of the optimiser tolerance.
1086    #[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            // Deterministic non-negative integer counts near exp(eta).
1106            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        // Replicated design (c literal copies of each row).
1116        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            // Costs and gradients must coincide to optimiser precision; the only
1142            // admissible difference is f64 summation order over n vs c·n rows.
1143            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    /// Regression (issue #893): the geometric-mean log-weight ρ-anchor
1161    /// ([`RemlState::rho_weight_anchor`]) is a *profiled*-dispersion construct
1162    /// (issue #877). For a fixed-dispersion family the optimum does not slide by
1163    /// `log c` under a weight rescale in a way the prior should track, and a
1164    /// nonzero anchor would evaluate the regularising ρ-prior at *different*
1165    /// coordinates for the `w=c` vs `c×` encodings — breaking the very
1166    /// equivalence #893 requires. The anchor must therefore be exactly `0` for a
1167    /// fixed-dispersion family and the geometric mean for Gaussian-identity.
1168    #[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        // All weights = c: geometric-mean log-weight = ln(c) ≠ 0.
1184        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    /// #2545 — the barrier gradient the CERTIFICATE subtracts is evaluated at
1206    /// the weight-anchored coordinate `ρ̃ = ρ − rho_weight_anchor`, exactly like
1207    /// the one the CRITERION adds.
1208    ///
1209    /// This gate exists because of a specific near-miss. The subtrahend's closed
1210    /// form was validated as `w·a·tanh(a·ρ) = 1.332439e-7` against a measurement
1211    /// of `1.332439e-7` — a seven-figure match, on an UNWEIGHTED fixture, where
1212    /// the anchor is exactly `0` and the raw-ρ and anchored forms are the same
1213    /// function. A formula validated only where one of its inputs is zero has not
1214    /// been validated in that input. Had the certificate subtracted the raw-ρ
1215    /// form, every WEIGHTED fit would have had a subtrahend that is not the term
1216    /// the criterion added, reintroducing exactly the weight dependence
1217    /// `rho_weight_anchor` exists to remove (#877: λ̂ is invariant to a global
1218    /// rescale `w → c·w`) — trading a rail that cannot certify for a certificate
1219    /// that moves when you rescale the weights.
1220    ///
1221    /// So the fixture is deliberately the weighted one (all weights `c = 4`,
1222    /// Gaussian-identity, anchor `ln 4 ≈ 1.3863`) and the probe ρ's are chosen
1223    /// where `tanh` is NOT saturated: at ρ=5 the anchored value is `5.970e-8`
1224    /// against the raw form's `7.769e-8`, a 30% gap. At ρ=30 the two agree to
1225    /// 3e-5 relative, which is why a saturated probe could never have caught this.
1226    #[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            // The discriminator: at the unsaturated probes the raw-rho form is a
1273            // different number, so a subtrahend that dropped the anchor could not
1274            // pass the assertion above.
1275            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        // And it is the SAME emission the criterion adds, not a parallel one:
1286        // `build_prior` reads `soft_rho_guard_prior_atom(rho).gradient()`.
1287        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        // The invariance itself, not just the formula. Under `w → c·w` the
1297        // pure-REML optimum slides `ρ̂ → ρ̂ + log c` and the anchor slides by the
1298        // same `log c`, so ρ̃ — and therefore everything the certificate
1299        // subtracts — is UNCHANGED at the corresponding coordinate. This is the
1300        // property #877 bought and the one a raw-ρ subtrahend would have spent:
1301        // it is asserted here as a bitwise identity between two states built on
1302        // weights that differ by a factor of `c`.
1303        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        //
1312        // The bar is roundoff, not bit-identity, and the distinction is the
1313        // arithmetic rather than a hedge: the rescaled arm evaluates
1314        // `(rho + ln c) - anchor` where the unrescaled one evaluates `rho - 0`,
1315        // and `(x + h) - h != x` in f64. Measured spread at these probes is
1316        // 4e-17 relative (5 ULP at rho=-3, exact at the other two), so 1e-12 is
1317        // five orders inside the measurement while still refusing the 30% gap a
1318        // dropped anchor would open at rho=5.
1319        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    /// #1575 multi-slot outer-eval cache correctness oracle.
1503    ///
1504    /// A memoization is only safe if a hit returns *exactly* what the miss path
1505    /// stored. This pins three properties of the bounded LRU:
1506    ///   1. round-trip fidelity — a hit is `f64::to_bits`-identical in cost AND
1507    ///      every gradient component to the value stored on the miss path;
1508    ///   2. no aliasing — distinct rho-keys never return each other's eval;
1509    ///   3. honest eviction — once an evicted key is requested again it MISSES
1510    ///      (so the caller recomputes) rather than returning a stale neighbour.
1511    #[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        // Helper: a deterministic OuterEval whose bits encode `seed`, so any
1516        // cross-key contamination is detectable bit-for-bit.
1517        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        // (1) Round-trip fidelity: store at rho_a, then a forced hit must equal
1535        // the stored eval bit-for-bit (the "hit == miss" guarantee).
1536        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        // (2) No aliasing: a second, distinct rho must return its OWN eval, and
1554        // the first key must still return the first eval untouched.
1555        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        // (3) Honest eviction: overflow the LRU with fresh keys. The
1578        // least-recently-used entry must be evicted and then MISS (forcing a
1579        // recompute), while a still-resident key returns its exact stored bits.
1580        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        // Cache is exactly full; key[0] is the least-recently-used.
1592        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        // One more distinct key evicts the LRU (key[0]).
1602        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        // Build a minimal logit RemlState, populate the cross-call PIRLS LRU
1643        // by evaluating the outer objective at one rho, then verify that
1644        // reset_outer_seed_state wipes that LRU (alongside the eval bundle
1645        // and warm-start signals). This pins down the cross-attempt
1646        // cleanup contract that a budget-bump retry relies on.
1647        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        // Trigger a full outer eval so execute_pirls_if_needed inserts at
1663        // least one entry into the cross-call PIRLS LRU.
1664        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        // #1448 regression: the NB outer θ↔λ alternation loop
1701        // (solver/estimate/optimizer.rs) re-runs the ρ search after each θ
1702        // refresh by (a) re-freezing the λ-search θ at θ_final into
1703        // `frozen_negbin_theta`, then (b) calling `reset_outer_seed_state()` to
1704        // drop the caches keyed to the old θ. Step (b) MUST NOT clear the freeze
1705        // set in step (a): the capture in `solve_for_unified_rho` only writes the
1706        // frozen slot when it is 0, so if the reset zeroed it the next round would
1707        // re-derive θ from the seed η and the loop would never reach the (ρ, θ)
1708        // joint fixed point — silently regressing #1448 back to a single
1709        // freeze→refresh pass.
1710        //
1711        // This pins the load-bearing distinction between `reset_outer_seed_state`
1712        // (alternation-round reset, freeze SURVIVES) and the surface-refresh reset
1713        // (new design, freeze re-zeroed). The end-to-end convergence on a real NB
1714        // fit is exercised by the public-API path; here we lock the invariant the
1715        // loop depends on, next to `reset_outer_seed_state_clears_pirls_cache`.
1716        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        // Simulate the alternation loop's re-freeze step: pin θ_final.
1733        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        // The alternation loop's per-round reset.
1744        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        // Use one directional hyperparameter τ with a penalty perturbation:
1857        // S(τ) = S + τ S_τ.
1858        // Keep X_τ = 0 so this identity test remains valid in both non-Firth
1859        // and Firth-logit modes.
1860        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        // Tight inner tolerance: the envelope theorem requires an exact inner
1868        // P-IRLS optimum; 1e-10 leaves enough residual gradient to cause ~12%
1869        // V_tau mismatch on this small (n=8) logistic problem.
1870        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        // B from implicit solve:
1880        //   H B = X_τ^T g - X^T W(X_τ β̂) - S_τ β̂.
1881        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        // H_τ from exact total derivative:
1890        //   H_τ = X_τ^T W X + X^T W X_τ + X^T W_τ X + S_τ,
1891        // with W_τ provided by the family directional curvature callback.
1892        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        // Fit-block stationarity cancellation:
1911        //   -ℓ_β^T B + β̂^T S B = 0.
1912        // Here S is the effective penalty in the inner Hessian surface:
1913        //   S = H - X^T W X.
1914        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        // Finite differences in τ against re-fit objective and mode.
1919        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        // Rank-deficient X: the 4th column is 2x the 2nd column.
2002        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        // Rank-deficient Firth logit needs more inner iterations to converge
2030        // tightly enough for the envelope-theorem derivative tests.
2031        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        // Rank-deficient design: col4 = 2*col2.
2166        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        // Exact reduced-space Firth gradient:
2187        //   gradPhi = 0.5 Xᵀ (w' ⊙ h), with h = diag(X_r K_r X_rᵀ).
2188        let gradphi = 0.5 * x.t().dot(&(&op.w1 * &op.h_diag));
2189
2190        // Check (I - QQᵀ) gradPhi ≈ 0.
2191        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    /// Regression for gam#1821: the analytic ρ-gradient of the Firth-corrected
2261    /// LAML cost must equal the central finite difference of that SAME cost,
2262    /// evaluated END-TO-END through the inner P-IRLS solve (`compute_cost` /
2263    /// `compute_gradient`), for a genuinely Firth-active (near-separable) fit.
2264    ///
2265    /// This exercises the branch the earlier operator-level Firth FD tests never
2266    /// touched: the LM line-search that produces β̂. The dense LAML gradient uses
2267    /// the envelope identity, which holds ONLY when β̂ satisfies the *Firth*-KKT
2268    /// stationarity `∇(−ℓ+½βᵀSβ) = ∇Φ`. When `GamWorkingModel::update_candidate`
2269    /// built line-search candidates with Firth disabled, the candidate/accepted
2270    /// `WorkingState` dropped the `−2·½log|XᵀWX|` Jeffreys term, the objective
2271    /// the line search compared (candidate vs `current_penalized`) was
2272    /// inconsistent, and — because the accepted state IS the candidate and
2273    /// convergence is certified on `accepted_state.gradient` — the inner solve
2274    /// settled at the ordinary penalized MLE (`∇(−ℓ+½βᵀSβ)=0`) instead. At that
2275    /// wrong mode the envelope breaks and the analytic ρ-gradient disagrees with
2276    /// the FD of the cost by `O((∇Φ)ᵀ∂β̂/∂ρ)` — a large (percent-level) desync
2277    /// that appears ONLY under `firth_bias_reduction`. A tight FD-vs-analytic
2278    /// bound therefore fails iff the inner solve regresses off the Firth mode.
2279    #[test]
2280    pub(crate) fn firth_logit_rho_gradient_matches_finite_difference_through_inner_solve() {
2281        // Near-separable n=3 logit: Firth is materially active (the ordinary
2282        // penalized MLE and the Firth-penalized mode differ enough that any
2283        // envelope break shows up far above the 1e-4 tolerance).
2284        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        // Full-rank identity penalty on both coefficients (single ρ).
2288        let s0 = array![[1.0, 0.0], [0.0, 1.0]];
2289        // Tight inner tolerance so β̂(ρ ± δ) and β̂(ρ) are all at the converged
2290        // Firth-KKT mode; otherwise the FD would capture β̂'s residual motion.
2291        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        // Mixed block must be nontrivial for at least one supplied direction.
2462        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        // FD via physical perturbation of design/penalty matrices (matching
2516        // the V_tau FD pattern).  For column j we perturb X and S₀ along
2517        // direction j, build fresh states, and evaluate the τ-gradient for
2518        // every direction i at those perturbed states.
2519        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        // The hyper direction declares a second-order penalty derivative
2556        // against base penalty index 1, but the configured ρ vector has
2557        // dimension 1 (so only index 0 is valid).  The pair-callback
2558        // builder in `build_tau_penalty_derivative_data` is responsible for
2559        // validating both first- and second-order penalty indices against
2560        // `rho.len()`; this test pins that contract.
2561        //
2562        // We deliberately keep `firth_bias_reduction = true` here so the
2563        // call site exercises the full Firth/Tierney–Kadane outer pipeline:
2564        // PIRLS + ext-coord construction + pair-callback assembly.  With
2565        // analytic c/d propagation now wired in
2566        // `tk_direct_gradient_from_cd_and_design`, there is no longer any
2567        // FD-fallback rejection on this path, so the out-of-bounds error
2568        // fired by the pair-callback builder is the first failure the
2569        // joint evaluator surfaces — and that is exactly what we want this
2570        // test to assert.
2571        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        // FD via physical perturbation of design/penalty matrices (matching
2656        // the V_tau FD pattern).  For column j we perturb X and S₀ along
2657        // direction j, build fresh states, and evaluate the τ-gradient for
2658        // every direction i at those perturbed states.
2659        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    // ── Profiled Gaussian REML coverage for design-moving τ-directions ──
2694    //
2695    // The existing directional-hyper tests all use BinomialLogit, which has
2696    // DispersionHandling::Fixed.  These tests validate the profiled Gaussian
2697    // path (DispersionHandling::ProfiledGaussian) with design-moving
2698    // τ-directions, where the profiled scale φ̂ = D_p/(n−M) depends on ρ
2699    // and the envelope-theorem rescaling by (n−M)/D_p must be correct.
2700
2701    /// Shared test fixture for profiled Gaussian REML tests.
2702    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        /// Design-moving τ-direction (non-zero X_τ, zero S_τ).
2710        pub(crate) x_tau_design: Array2<f64>,
2711        /// Penalty-only τ-direction (zero X_τ, non-zero S_τ).
2712        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        // Validate the ττ Hessian block under profiled Gaussian REML with
2830        // both a penalty-only and a design-moving direction.
2831        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        // Finite-difference Hessian: perturb each direction, re-evaluate
2853        // gradient of all directions at perturbed states.
2854        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    // ── Non-Gaussian + design-motion: IFT Hessian-drift coverage ────────
2883    //
2884    // For non-Gaussian links (logit, probit, cloglog, ...), H = X'W(η)X + S
2885    // depends on β̂ through η = Xβ̂.  When ψ moves the design, the total
2886    // Hessian drift dH/dψ includes an IFT contribution from dβ̂/dψ:
2887    //
2888    //   dH/dψ = [explicit at fixed β] + X' diag(c ⊙ X(-v_i)) X
2889    //
2890    // where v_i = H⁻¹ g_i.  The standard GLM path handles this via
2891    // `hessian_derivative_correction(v_i)`.  This test validates that the
2892    // gradient is correct for logit + design-moving ψ, which would fail if
2893    // the IFT correction were missing.
2894
2895    #[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        // Design-moving direction with non-zero X_τ.
2917        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        // Hessian-level validation for logit + design-motion.
2957        // The IFT correction enters the trace term through
2958        // hessian_derivative_correction(v_i), so the Hessian is the most
2959        // sensitive test of whether the correction is applied correctly.
2960        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        // Two directions: one penalty-only, one design-moving.
2979        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    // ── Larger non-Gaussian + design-motion fixture (n=30, p=5) ────────
3039    //
3040    // Validates the IFT correction (hessian_derivative_correction) at a
3041    // scale large enough that the correction is numerically non-trivial:
3042    // with n=30 and p=5, the logistic Hessian W(η) is far from identity
3043    // and the IFT term dβ̂/dψ contributes meaningfully.
3044
3045    /// Shared test fixture for binomial-logit REML with design-moving
3046    /// ψ-coordinates, n=30, p=5.
3047    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        /// Design-moving τ-direction: non-zero X_τ, zero S_τ.
3055        pub(crate) x_tau_design: Array2<f64>,
3056        /// Penalty-only τ-direction: zero X_τ, non-zero S_τ.
3057        pub(crate) s_tau_penalty: Array2<f64>,
3058    }
3059
3060    impl BinomialLogitDesignMotionFixture {
3061        pub(crate) fn new() -> Self {
3062            // Binary response with roughly balanced classes.
3063            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            // Design matrix: intercept + 4 covariate columns with varied magnitudes.
3068            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            // Penalty matrix: zero on intercept, SPD on remaining 4 columns.
3101            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            // Design-moving direction: perturb covariate columns, leave
3110            // intercept untouched.
3111            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            // Penalty-only direction: non-zero S_τ, symmetric, zero on intercept.
3144            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    // ── n=30, p=5 binomial-logit design-motion gradient tests ────────
3186
3187    #[test]
3188    pub(crate) fn binomial_logit_n30_design_moving_gradient_matches_fd() {
3189        // Pure design-motion: X_τ ≠ 0, S_τ = 0.
3190        // The IFT correction is essential here: because the family is
3191        // binomial-logit, the working weights W(η) depend on β̂, so
3192        // when X moves with ψ, the implicit derivative dβ̂/dψ enters
3193        // the total Hessian drift.  Without hessian_derivative_correction
3194        // the analytic gradient would disagree with FD.
3195        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        // Penalty-only direction: X_τ = 0, S_τ ≠ 0.
3222        // Serves as a baseline: the IFT correction should still be
3223        // present (since H depends on β̂ through W(η)), but the
3224        // explicit X_τ contribution is zero.
3225        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        // Joint direction: both X_τ ≠ 0 and S_τ ≠ 0 simultaneously.
3252        // This is the hardest case: the analytic gradient must correctly
3253        // combine the explicit penalty drift, the explicit design drift,
3254        // and the IFT Hessian-drift correction.
3255        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        // Hessian-level validation with two τ-directions: one
3281        // penalty-only and one design-moving.  The ττ Hessian block is
3282        // the most sensitive test of the IFT correction because errors
3283        // in the correction accumulate quadratically in the trace term.
3284        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        // Validate at a non-trivial smoothing parameter ρ = log(λ) = 1.5,
3336        // so the penalty term λS is scaled up and the balance between
3337        // likelihood and penalty is different from ρ=0.
3338        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        // FD at the shifted ρ: perturb X, re-solve inner, evaluate cost.
3356        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        // Regression lock for the `PenaltySubspaceTrace` pseudo-logdet
3373        // kernel installed by the rank-deficient LAML fix (see
3374        // `PenaltySubspaceTrace` and `intrinsic_hessian_pseudo_logdet_parts`;
3375        // since #901 the cost is the intrinsic `½ log|H_pen|₊` and the kernel
3376        // is the spectral `H_pen⁺`, exact for every drift direction).
3377        //
3378        // The sibling `binomial_logit_n30_design_moving_hessian_matches_fd`
3379        // passes pre- AND post-fix because its FD reference differentiates
3380        // the *analytic gradient* — any self-consistent (if wrong) gradient
3381        // kernel gives a self-consistent Hessian under re-differentiation,
3382        // so that test cannot distinguish full-space from projected traces.
3383        // It passed under the buggy kernel because the same leakage entered
3384        // both sides of the ratio and cancelled.
3385        //
3386        // Here we FD-differentiate `compute_cost` TWICE and compare against
3387        // the analytic Hessian.  Central second differences expose every
3388        // disagreement between `½ log|U_Sᵀ H U_S|_+` (used by the cost) and
3389        // `½ tr(G_ε(H) · Ḣ)` / `−½ tr(G_ε Ḣ_i G_ε Ḣ_j)` (the full-space
3390        // traces that the gradient and Hessian used before the projection
3391        // fix).  Under the buggy kernel the IFT correction
3392        // `D_β H[v] = X' diag(c ⊙ X v) X` leaks onto `null(S)` — X's
3393        // all-ones intercept column sits there — and that leakage enters
3394        // the analytic Hessian but not the cost's projected logdet.
3395        //
3396        // Direction mix chosen to maximise the null-space leakage pathway:
3397        //   τ_0 = penalty-only (X_τ = 0, S_τ ≠ 0)  → v_0 = H⁻¹(−S_τ β̂) is
3398        //         concentrated in range(S_+), but `D_β H[v_0]` has rows and
3399        //         columns on the intercept because `X[:,0] = 1_n`.
3400        //   τ_1 = design-moving (X_τ ≠ 0 on non-intercept columns, S_τ = 0)
3401        //         → `v_1` also picks up the intercept via `X'WX_τβ̂`, and
3402        //         the base drift `X_τᵀWX + XᵀWX_τ` straddles range(S_+) /
3403        //         null(S).
3404        // Both pure directions AND the mixed partial load the Schur correction,
3405        // so any of the three entries can catch a regression.
3406        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        // Analytic Hessian block.
3420        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        // Cost-level FD reference.  Central second differences give O(h²)
3429        // accuracy; the step is sized so the physical perturbation on X / S
3430        // stays near `1e-5` (same scale as the gradient tests).
3431        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        // Evaluate `compute_cost` at `(a · τ_0, b · τ_1)` multipliers.
3451        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
3513/// Mechanical surface every `DerivativeMatrixStorage` variant must expose so
3514/// the `HyperDesignDerivative` / `HyperPenaltyDerivative` wrappers can dispatch
3515/// with a single per-call `storage_dispatch!`. Each backend owns its variant's
3516/// substantive math; the wrappers contain only one-line routing.
3517///
3518/// `design_*` variants treat the backend as an X-style operator (rows index
3519/// data, columns index coefficients); `penalty_*` variants treat the backend
3520/// as a square `p×p` penalty in the global coefficient frame. The Embedded
3521/// case is the only variant whose two views genuinely differ (local rows vs
3522/// total_dim square), which is why the two role-specific methods both live in
3523/// one trait rather than two parallel traits.
3524trait 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    /// Default materialises through `design_transformed` then `.dot(u)`;
3550    /// implicit/latent-coordinate backends override with a direct-operator
3551    /// path that skips the dense materialisation.
3552    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    /// Default materialises through `design_transformed` then `.t().dot(v)`;
3561    /// implicit/latent-coordinate backends override with a direct path.
3562    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
3582/// Fans `expr` over the four `DerivativeMatrixStorage` variants in one place
3583/// so every wrapper method is a single dispatch line — the compiler enforces
3584/// exhaustiveness here, so adding a new variant produces one hard error at
3585/// this site rather than a silent miss in any of the (currently 16) ladders.
3586macro_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/// Which derivative level the implicit operator should compute.
3611#[derive(Clone, Copy, Debug)]
3612pub enum ImplicitDerivLevel {
3613    /// ∂X/∂ψ_d
3614    First(usize),
3615    /// ∂²X/∂ψ_d²
3616    SecondDiag(usize),
3617    /// ∂²X/∂ψ_d∂ψ_e
3618    SecondCross(usize, usize),
3619}
3620
3621/// Lazy implicit operator storage: delegates matvecs to the
3622/// `ImplicitDesignPsiDerivative` and materializes dense form only on demand.
3623#[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    /// Cached dense materialization (lazy, populated on first call to ops that need the full matrix).
3630    ///
3631    /// Rayon-safe: `materialize_local` calls `materialize_first` / `_second_diag`
3632    /// / `_second_cross` on the implicit basis-derivative operator, which for
3633    /// streaming bases dispatches `(0..nc).into_par_iter().for_each(...)`. A plain
3634    /// `std::sync::OnceLock` here would deadlock if `materialize_dense` were first
3635    /// called concurrently from inside another rayon par_iter — racing workers
3636    /// would park on the OnceLock's OS condvar, leaving the leader's nested
3637    /// par_iter without workers. `RayonSafeOnce` runs init lock-free.
3638    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        // Zero penalty derivative: `amp · 0 = 0`, so `amp` scales nothing and
4040        // `target` is left unchanged. Validate it is finite so a bad scale
4041        // surfaces here rather than silently no-op'ing on a NaN/inf amplitude.
4042        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    /// If this derivative uses implicit storage at the first-derivative level,
4580    /// return the shared implicit operator and the axis index.
4581    ///
4582    /// Returns `None` for dense/embedded storage or for second-derivative levels.
4583    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    // Canonical penalty representation: every tau direction is decomposed into
4675    // base-penalty derivatives. There is no separate "assembled total" path.
4676    pub(crate) penalty_first_components: Vec<PenaltyDerivativeComponent>,
4677    // Optional pairwise second hyper-derivatives against all tau directions.
4678    // If provided, each vector must have length psi_dim and hold an optional
4679    // X_{tau_i,tau_j} entry in original coordinates.
4680    pub(crate) x_tau_tau_original: Option<Vec<Option<HyperDesignDerivative>>>,
4681    // Pairwise second derivatives are stored in the same canonical base-penalty
4682    // decomposition as the first derivatives.
4683    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    /// Whether this coordinate is penalty-like (B_i = ∂H/∂τ_i is PSD).
4694    /// True for τ (penalty scaling) coordinates; false for ψ (design-moving,
4695    /// anisotropic length-scale) coordinates. Controls EFS eligibility.
4696    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    /// Mark this coordinate as non-penalty-like (design-moving).
4775    /// EFS will skip it; use Newton/BFGS for these coordinates.
4776    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    /// Whether this coordinate's design derivative uses implicit storage at the
4819    /// first-derivative level.
4820    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    /// Extract the implicit design derivative operator and axis, if available.
4830    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    /// Design nonzeros, or `None` when the design has no sparse representation
4924    /// to count them from. Not the same statement as "the design has none":
4925    /// `design_not_sparse` fires whenever `as_sparse()` returns `None`, which
4926    /// includes a lazy operator wrapping a genuinely sparse matrix — and this
4927    /// field used to render that case as `nnz_x=0`, a measurement-shaped zero
4928    /// for a design measured at 583,674 nonzeros one gate later.
4929    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    // Exact Firth/Jeffreys objects on the identifiable subspace.
4947    //
4948    // Let X in R^{n×p} potentially be rank-deficient with rank r.
4949    // With optional fixed observation weights a_i >= 0 we define A = diag(a),
4950    // choose an orthonormal coefficient-space basis Q for the identifiable
4951    // subspace of A^{1/2} X, and set:
4952    //   X_r := A^{1/2} X Q          (A = I when no fixed observation weights),
4953    //   W   := diag(w), with w_i = mu_i (1 - mu_i), 0 < w_i <= 1/4 for finite logit eta,
4954    //   I_r := X_rᵀ W X_r,
4955    //   S_r := X_rᵀ X_r.
4956    //
4957    // Firth term is represented as:
4958    //   Phi(beta) = 0.5 log |I_r(beta)| - 0.5 log |S_r|,
4959    // which is exactly
4960    //   0.5 log |Uᵀ W U|
4961    // for the canonical orthonormalized identifiable design
4962    //   U = X_r S_r^{-1/2}.
4963    // This removes the raw-basis term from explicit reduced designs while
4964    // keeping the same identifiable-subspace hat matrix and beta derivatives,
4965    // because S_r is fixed with respect to beta.
4966    //
4967    // Mapping back to the full p-space uses:
4968    //   I_+^dagger = Q I_r^{-1} Qᵀ.
4969    //
4970    // We store reduced-space factors so all derivatives can be evaluated exactly
4971    // without materializing dense n×n matrices M = X K Xᵀ or P = M⊙M.
4972    pub(crate) x_dense: Array2<f64>,
4973    pub(crate) x_dense_t: Array2<f64>,
4974    // Orthonormal coefficient-space basis for the identifiable subspace,
4975    // built from the retained eigenspace of (A^{1/2} X)ᵀ(A^{1/2} X).
4976    pub(crate) q_basis: Array2<f64>,
4977    // Reduced identifiable design. With fixed observation weights a_i this is
4978    // diag(sqrt(a_i)) X Q; otherwise it is X Q.
4979    pub(crate) x_reduced: Array2<f64>,
4980    // Optional fixed case-weight square roots used when the Jeffreys/Firth
4981    // operator is formed from Xᵀ diag(case_weight ⊙ w(η)) X rather than
4982    // Xᵀ diag(w(η)) X. The exact directional tau derivatives must project and
4983    // row-scale with the same weights so the reduced Fisher, hat diagonals,
4984    // and tau kernels all live on one consistent identifiable subspace.
4985    pub(crate) observation_weight_sqrt: Option<Array1<f64>>,
4986    // I_r^{-1}
4987    pub(crate) k_reduced: Array2<f64>,
4988    // diag(S_r^{-1}) with S_r = X_rᵀ X_r. In the current canonical reduced
4989    // basis this completely characterizes the metric inverse, because Q
4990    // diagonalizes the design Gram. It is used to remove the reduced-coordinate
4991    // basis term from Phi_tau when the design moves.
4992    pub(crate) x_metric_reduced_inv_diag: Array1<f64>,
4993    // 0.5 (log|I_r| - log|S_r|) at the current eta.
4994    pub(crate) half_log_det: f64,
4995    // h = diag(M), M = X_r K_r X_r'
4996    pub(crate) h_diag: Array1<f64>,
4997    // Logistic Fisher-weight eta-derivatives: w', w'', w''', w'''' as n-vectors.
4998    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    // B = diag(w') X used in D Hphi and D^2 Hphi contractions.
5004    pub(crate) b_base: Array2<f64>,
5005    // Cached invariant contraction P*B where P = (X_r K_r X_r') ⊙ (X_r K_r X_r').
5006    // This avoids recomputing the same O(n r^2 p) block in every directional call.
5007    pub(crate) p_b_base: Array2<f64>,
5008}
5009
5010/// β-independent (design-only) factor of the Firth/Jeffreys operator.
5011///
5012/// Everything stored here depends ONLY on the fixed design `X` and the fixed
5013/// prior/observation weights `a_i` — NOT on the current linear predictor `η`
5014/// (i.e. NOT on β). For a single inner PIRLS solve the design and prior weights
5015/// are constant while `η` changes every Newton iteration, so this factor can be
5016/// built once per solve and reused, hoisting the O(n·p²) Gram, the O(p³)
5017/// identifiable-subspace eigendecomposition, and the two n×p design clones out
5018/// of the per-iteration hot path (#1575).
5019///
5020/// The β-dependent remainder (Fisher weights `w(η)`, reduced Fisher
5021/// `I_r = X_rᵀ W X_r`, its inverse `K_r`, the hat diagonal `h`, and the
5022/// half-log-determinant) is rebuilt per iteration from this factor via
5023/// [`FirthDenseOperator::build_from_design_factor`]. The design-only work stays
5024/// hoisted while the full per-state operator supplies both PIRLS diagnostics
5025/// and the exact Jeffreys coefficient curvature.
5026#[derive(Clone)]
5027pub(crate) struct FirthDesignFactor {
5028    // Raw design and its transpose (the operator stores owned copies).
5029    pub(crate) x_dense: Array2<f64>,
5030    pub(crate) x_dense_t: Array2<f64>,
5031    // Orthonormal identifiable-subspace basis Q of (A^{1/2} X)ᵀ(A^{1/2} X).
5032    pub(crate) q_basis: Array2<f64>,
5033    // Reduced identifiable design X_r = A^{1/2} X Q.
5034    pub(crate) x_reduced: Array2<f64>,
5035    // Fixed case-weight square roots (sqrt(a_i)), if any.
5036    pub(crate) observation_weight_sqrt: Option<Array1<f64>>,
5037    // Retained positive spectrum of the design Gram = S_r diagonal.
5038    pub(crate) metric_spectrum: Array1<f64>,
5039    // diag(S_r^{-1}); precomputed reciprocal of `metric_spectrum`.
5040    pub(crate) x_metric_reduced_inv_diag: Array1<f64>,
5041    // rank r = ncols(q_basis); n = nrows(x_dense).
5042    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    // B_u = diag(w'' ⊙ δη_u) X is represented by the row-scaling vector only.
5053    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    // Reduced design drift X_{tau,r} = X_tau Q used in exact design-moving
5063    // Hadamard-Gram contractions.
5064    pub(crate) x_tau_reduced: Array2<f64>,
5065    pub(super) dot_i_partial: Array2<f64>,
5066    // Reduced Fisher inverse drift:
5067    //   dot(K_r) = -K_r dot(I_r) K_r
5068    // where dot(I_r) includes explicit X_tau and weight drift at beta-fixed.
5069    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/// Pair-level (τ_i × τ_j) exact Firth bundle at fixed β.
5080///
5081/// Mirrors `FirthTauExactKernel` but for the 2nd-order cross
5082/// derivatives:
5083///   Phi_{τ_i τ_j}|β  (scalar, `phi_tau_tau_partial`)
5084///   (gphi)_{τ_i τ_j}|β (p-vector, `gphi_tau_tau`)
5085///
5086/// Carries an optional `tau_tau_kernel` so pair callbacks can chain
5087/// into Primitive A (`hphi_tau_tau_partial_apply`) for the operator-
5088/// valued Hessian 2nd drift without recomputing shared reduced Grams.
5089///
5090#[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/// Prepared state for `∂²H_φ/∂τ_i ∂τ_j |_β` (Primitive A).
5098///
5099/// Carries both τ-direction reduced designs, their η̇ vectors, and the
5100/// reduced-coordinate drifts (İ, K̇, ḣ) for i and j so the apply step can
5101/// form M̈_{ij}, K̈_{ij}, ḧ_{ij}, Γ̈_{ij}, and B̈_{ij} matrix-free.  Fields
5102/// are filled in by 13b; kept with a neutral internal shape so downstream
5103/// pair callbacks can hold the kernel across the pair dispatch.
5104///
5105/// Wired into the pair-callback's `b_operator` via
5106/// `FirthAugmentedPairHyperOperator`, and produced by both
5107/// `hphi_tau_tau_partial_prepare_from_partials` and
5108/// `exact_tau_tau_kernel` (the scalar/p-vector companion).
5109#[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/// Prepared state for `D_β((H_φ)_τ|_β)[v]` (Primitive B).
5126///
5127/// Carries the τ-kernel pieces (x_tau_reduced, İ, K̇, ḣ), the
5128/// β-direction quantities (δη_v, A_v, dh_v, b-chain), and the mixed
5129/// β-τ pieces (D_β(K̇_τ)[v], D_β(ḣ_τ)[v], δη_{τ,v}) so the apply
5130/// step collapses to the 9-term β-τ expansion without recomputing
5131/// shared reduced Grams.
5132#[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/// Holds the state for the outer REML optimization and supplies cost and
5149/// gradient evaluations to the `opt` optimizer.
5150///
5151/// The `cache` field uses `RefCell` to enable interior mutability. This is a crucial
5152/// performance optimization. The `cost_andgrad` closure required by the BFGS
5153/// optimizer takes an immutable reference `&self`. However, we want to cache the
5154/// results of the expensive P-IRLS computation to avoid re-calculating the fit
5155/// for the same `rho` vector, which can happen during the line search.
5156/// `RefCell` allows us to mutate the cache through a `&self` reference,
5157/// making this optimization possible while adhering to the optimizer's API.
5158#[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    /// The routing verdict this bundle was built under, carried WITH the
5164    /// quantities it was decided from (#2465 instance 4). The bundle used to
5165    /// hold the bare `RemlGeometry` label, so every consumer that reported
5166    /// `backend {:?}` reported a two-valued enum and nothing that could
5167    /// falsify it: `select_reml_geometry` measures a penalized-Hessian
5168    /// density against `SPARSE_HESSIAN_MAX_DENSITY` on one of its six routes
5169    /// and never measures it on the other five, and the label is identical
5170    /// across all six. Storing the decision rather than its outcome makes a
5171    /// bundle unrepresentable without the basis for its own label.
5172    pub(crate) geometry: SparseRemlDecision,
5173    /// The exact H_total matrix used for LAML cost computation.
5174    /// For Firth: effective Hessian minus hphi (plus any barrier curvature).
5175    /// For non-Firth: the effective Hessian itself (plus any barrier curvature).
5176    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    /// Cached FirthDenseOperator built from the original (non-reparameterized)
5180    /// design matrix, for use by the sparse evaluation path.
5181    pub(crate) firth_dense_operator_original: Option<Arc<FirthDenseOperator>>,
5182    /// The ONE original-frame penalty pseudo-logdet factorization for this
5183    /// evaluation point (#931 atom discipline). `log|Σ λ_k S_k|₊`'s VALUE,
5184    /// ρ-derivatives, τ/ψ components, and ρ×τ cross blocks are all
5185    /// contractions of this single eigendecomposition; the ρ-side criterion
5186    /// assembly (`dense_penalty_logdet_derivs`, the sparse det2 path) and the
5187    /// original-basis hyper-coordinate builders share it through
5188    /// [`EvalShared::penalty_pseudologdet_original`]. Building a second
5189    /// factorization of the same Sλ for the same evaluation point is the
5190    /// objective↔gradient desync surface (#748/#752/#901) this cell removes:
5191    /// the ridge and positive-eigenspace threshold are decided exactly once.
5192    /// (The transformed-frame pair-callback path builds its own object — it
5193    /// factorizes the canonical-TRANSFORMED, possibly constraint-projected
5194    /// penalties, a genuinely different matrix, not a duplicate of this one.)
5195    pub(crate) penalty_pseudologdet: std::sync::OnceLock<Arc<penalty_logdet::PenaltyPseudologdet>>,
5196    /// #2644: the root-scale `log|H|` for this evaluation point, decided once.
5197    ///
5198    /// `None` inside the cell means the upgrade was examined and declined (the
5199    /// assembled spectrum already resolves the criterion, or the root could not
5200    /// be shown to reproduce this `H`); an empty cell means it has not been
5201    /// examined yet. Caching it here is what keeps the extra `O(n·p²)` Gram to
5202    /// once per ρ rather than once per value/gradient/Hessian call at that ρ.
5203    pub(crate) root_scale_hessian_logdet: std::sync::OnceLock<Option<f64>>,
5204    /// The penalty components the criterion APPLIES, `S̃_k = Π S_k Π`, in the
5205    /// ORIGINAL coefficient frame (#2454).
5206    ///
5207    /// `RemlState::canonical_penalties` are the penalties as the term layer
5208    /// built them. The penalty the model is actually fitted under is the
5209    /// reparameterization's split-projected `S̃(λ) = Π(Σ_k λ_k S_k)Π`: that is
5210    /// what the inner solve minimizes, what `H` contains, and what the reported
5211    /// penalty energy `‖Eβ̂‖²` measures. `log|S|₊`, its rank, its ρ- and
5212    /// ψ-derivatives, and the `M_p` that sets the REML denominator must all be
5213    /// taken on THAT penalty, or the two halves of the LAML ratio saturate at
5214    /// different rates and the criterion acquires an asymptotic ρ-slope of
5215    /// `½(rank(S̃) − rank(S))` — unbounded below, no interior optimum.
5216    ///
5217    /// `Π` is λ-invariant (it comes from the balanced penalty, not the weighted
5218    /// one), so this set is a function of the inner solve's reparameterization
5219    /// alone and is computed once per bundle. It is the same `Arc` as the input
5220    /// when the split declares nothing null or the projection falls below the
5221    /// roots' own representation noise, which is the ordinary case.
5222    pub(crate) applied_canonical_penalties:
5223        std::sync::OnceLock<Arc<Vec<gam_terms::construction::CanonicalPenalty>>>,
5224    /// Per-evaluation-point cache of the canonical penalty score vectors
5225    /// `S_k β̂` evaluated at this bundle's inner mode `β̂ =
5226    /// pirls_result.beta_transformed` (unscaled by λ_k). These depend ONLY
5227    /// on the inner solution carried by this bundle and the `RemlState`'s
5228    /// fixed `canonical_penalties` — never on which ρ-coordinate or eval
5229    /// mode the assembly is running — so they are computed exactly once per
5230    /// inner solution and shared by every assemble call that reuses the
5231    /// bundle (cost + gradient evaluations at the same ρ, EFS, synthetic-ext
5232    /// value probes). Exact hoist, not an approximation: every consumer sees
5233    /// literally the same vectors it previously recomputed. Initialized via
5234    /// plain ndarray matvecs (no rayon inside the `OnceLock` closure — the
5235    /// `get_or_init`+`into_par_iter` deadlock trap does not apply).
5236    pub(crate) penalty_scores_at_mode: std::sync::OnceLock<Arc<Vec<Array1<f64>>>>,
5237    /// Per-evaluation-point cache of the #784 block-local Laplace-to-sampling
5238    /// correction `TkCorrectionTerms { value, gradient }`. The correction is a
5239    /// deterministic function of ONLY this bundle's converged inner state
5240    /// (`pirls_result`, `h_total`), the `RemlState`'s fixed
5241    /// `canonical_penalties`, and the bundle's ρ — never of the eval `mode`:
5242    /// the diagnostic eigendecomposition, the fixed-seed importance sampler,
5243    /// and the (b)–(d) gradient channels all read mode-invariant fields, and
5244    /// the term carries no Hessian, so the value+gradient are identical for the
5245    /// value-only, value+gradient, and value+gradient+Hessian assemble calls
5246    /// that share this bundle at a single ρ. The expensive path (eigendecomp +
5247    /// O(draws·n·m) sampler) previously reran on every one of those 2–3 calls
5248    /// per outer iteration; hoisting it onto the bundle computes it exactly
5249    /// once per inner solution (exact hoist, identical values — #784, #1082).
5250    /// Keyed only on the external-coordinate count `n_ext`: with no ψ
5251    /// coordinates (`n_ext == 0`) the correction engages; with ψ present the
5252    /// seam declines (returns the cheap zero), and n_ext is fixed for a fit, so
5253    /// a single cell suffices.
5254    /// The third slot carries the #2623 ρ-block audit record for the SAME
5255    /// computation, so a later assemble call at this ρ that reads the cache can
5256    /// re-publish it. Without that, the audit window — which is cleared at the
5257    /// start of every assemble call — would report the second and third calls at
5258    /// an engaged ρ as DECLINED, and an FD row asserting engagement would fail
5259    /// on a fit where the splice ran. `None` when the splice declined or when
5260    /// the audit was disarmed (the production case, which allocates nothing).
5261    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
5268/// The penalty components the criterion APPLIES, `S̃_k = Π S_k Π`, for an inner
5269/// solve that ran under `reparam` (#2454).
5270///
5271/// A free function rather than a method, because the sparse-exact lane has to
5272/// build this BEFORE its `EvalShared` exists — and building it twice from two
5273/// places is exactly the split this whole issue is about. `EvalShared` caches
5274/// the result of this call; the sparse builder seeds that cache with the same
5275/// object it used itself.
5276///
5277/// The identity `Arc` comes back whenever the split declares nothing null, the
5278/// penalty list is empty, or every projection falls below the roots' own
5279/// representation noise — so the ordinary fit allocates nothing and keeps each
5280/// penalty's `op` handle, block chart and cached spectrum.
5281pub(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    // `project_canonical` returns the penalty itself when the projection is
5302    // below the root's own noise, so an all-unchanged result IS the identity
5303    // and is handed back as the original `Arc` rather than as a copy.
5304    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    /// Lazily build — once per evaluation point — the original-frame
5325    /// [`PenaltyPseudologdet`](penalty_logdet::PenaltyPseudologdet) of
5326    /// `Σ λ_k S_k` and hand every caller the SAME factorization.
5327    ///
5328    /// This is the #931 port of the penalty-logdet term: value, ρ-first /
5329    /// ρ-second derivatives, τ-gradient components, τ×τ and ρ×τ Hessian
5330    /// blocks are all projections of one eigendecomposition, so no pair of
5331    /// consumers can disagree about the ridge or the positive-eigenspace
5332    /// threshold. The ridge is read from this bundle's `ridge_passport` —
5333    /// the single place that convention is decided.
5334    ///
5335    /// `lambdas` must be the λ = exp(ρ) vector of this bundle's evaluation
5336    /// point and `p` the original-basis coefficient dimension; on a cache
5337    /// hit both are checked against the stored object where representable.
5338    /// The penalty components the criterion APPLIES, `S̃_k = Π S_k Π`, in the
5339    /// ORIGINAL frame — see [`EvalShared::applied_canonical_penalties`] the
5340    /// field for why every criterion consumer must read these rather than the
5341    /// term layer's raw blocks (#2454).
5342    ///
5343    /// Computed once per bundle and shared. Returns the SAME `Arc` as the input
5344    /// whenever the split declares nothing null or the projection falls below
5345    /// the roots' own representation noise, so the ordinary path allocates
5346    /// nothing and keeps every penalty's `op` handle and block chart.
5347    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        // `log|S|₊` is one half of the LAML ratio `½(log|H| − log|S|₊)`; the
5385        // other half carries the split-projected penalty, so this one must too.
5386        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            // A concurrent caller initialized the cell first; both objects
5399            // were built from identical inputs — return the canonical winner
5400            // so every consumer holds literally the same factorization.
5401            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    /// One-line rendering of the quantities this routing verdict was decided
5421    /// from, for emission beside the `backend=`/`geometry=` label itself
5422    /// (#2465). `nnz_x`/`nnz_h_est`/`density_h_est` read `na` on the routes
5423    /// that decide before the corresponding structure is measured — that
5424    /// absence is itself the basis, and it is what tells
5425    /// `penalized_hessian_too_dense` (a density genuinely measured above the
5426    /// threshold) apart from `design_not_sparse` and its four siblings, which
5427    /// never measure one.
5428    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
5447/// LRU cache keyed by sanitized ρ vectors that holds compacted PIRLS results
5448/// for warm-starting outer line searches and revisited evaluations.
5449///
5450/// Eviction is byte-budgeted rather than entry-count-budgeted: each entry
5451/// records its own estimated footprint (the surviving n-length vectors plus
5452/// the two p×p Hessians plus per-entry overhead) and the cache evicts in
5453/// LRU order until the running total fits under the budget. An entry that
5454/// individually exceeds the budget is rejected silently rather than poisoning
5455/// the cache.
5456pub(crate) struct PirlsLruCache {
5457    // Stored tuple: (compacted result, last-touched clock, estimated bytes).
5458    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        // Refuse entries that on their own already exceed the entire budget;
5488        // caching one would force eviction of every other entry without
5489        // leaving room for the new one anyway.
5490        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    /// Build a cache key from the transformed-E matrix and ridge passport.
5564    /// `E` is hashed by exact f64 bits (column-major), so the key is bit-exact
5565    /// and avoids float-Hash issues; the ridge passport is hashed via its
5566    /// `Hash` impl. Two calls at the same `(E, ridge)` yield equal keys.
5567    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
5592/// Estimate the in-cache footprint of a (compacted) PIRLS result.
5593///
5594/// Mirrors what `compact_for_reml_cache` keeps:
5595/// * six surviving n-length f64 arrays (final_eta, solveweights,
5596///   solveworking_response, solvemu, solve_c_array, solve_d_array);
5597/// * the p-length coefficient vector;
5598/// * the two p×p Hessians (dense or CSC sparse);
5599/// * the `ReparamResult` payload — the dominant scaling term beyond n, since
5600///   it carries `s_transformed`, `qs`, and `e_transformed` as p×p / rank×p
5601///   matrices.
5602/// A small constant overhead absorbs scalar fields, enum discriminants, and
5603/// the HashMap entry. This errs on the conservative side: overestimation
5604/// causes earlier eviction, never under-counting that would let the cache
5605/// silently exceed the byte budget.
5606pub(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            // CSC sparse: f64 values + usize row indices + usize column pointers.
5632            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
5639/// Capacity (number of distinct rho-points) of the outer-eval reuse LRU.
5640///
5641/// Sized to comfortably span a binomial seed grid's local revisit window
5642/// (baseline + isotropic shifts + per-axis refinements) plus a few
5643/// line-search trial points without unbounded growth. Each slot holds one
5644/// `OuterEval` (a scalar cost, a length-k gradient, an optional inner-beta
5645/// hint, and a usually-`Unavailable` Hessian), so the footprint is tiny.
5646pub(crate) const OUTER_EVAL_LRU_CAPACITY: usize = 8;
5647
5648/// Bounded least-recently-used cache of converged outer REML evaluations,
5649/// keyed by sanitized rho-bits plus the active inner-solve caps.
5650///
5651/// CORRECTNESS: the key starts with `Vec<u64>` of `f64::to_bits` (with ±0
5652/// canonicalized) and appends the screening and outer P-IRLS caps. Every
5653/// other input to `OuterEval` — design matrix, prior weights, offset, penalty
5654/// structure, link/SAS/mixture state, Firth/Jeffreys configuration, and the
5655/// rho-prior — is immutable for the lifetime of the state that owns the cache.
5656/// Inner fidelity is not immutable: search-time caps deliberately change it.
5657/// Including those caps prevents a full-fidelity terminal request from
5658/// replaying a coarse search entry at the same rho. Distinct rho/fidelity
5659/// states never alias because lookups compare the full key vector.
5660pub(crate) struct OuterEvalLru {
5661    capacity: usize,
5662    /// Front = least-recently-used, back = most-recently-used.
5663    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    /// Returns a clone of the eval stored under `key`, if present, promoting it
5675    /// to most-recently-used. A miss returns `None` so the caller recomputes —
5676    /// never a stale value from a different key.
5677    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    /// Inserts (or refreshes) the eval for `key` as most-recently-used,
5686    /// evicting the least-recently-used entry once capacity is exceeded.
5687    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
5706/// Centralized cache/memoization owner for REML evaluations.
5707///
5708/// This keeps cache-key identity, bundle reuse, and invalidation policy out of
5709/// the math kernels so objective/derivative routines can stay algebra-focused.
5710pub(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    /// Most-recently-*stored* outer eval (single slot). Retained verbatim so
5715    /// `previous_outer_gradient_norm` keeps its exact "immediately previous
5716    /// distinct eval" semantics, independent of the multi-slot reuse cache.
5717    pub(crate) current_outer_eval: RwLock<Option<(Vec<u64>, OuterEval)>>,
5718    /// Bounded multi-slot LRU of converged outer evaluations keyed by the
5719    /// sanitized rho-bits and the active inner-solve caps (#1575/#2309).
5720    ///
5721    /// For a frozen `RemlState` (fixed design, prior weights, offset, penalty
5722    /// structure, link state, Firth/Jeffreys configuration, and rho-prior — all
5723    /// of which are immutable for the lifetime of the state that owns this
5724    /// manager and therefore the lifetime of the cache), the remaining
5725    /// result-determining state is `(rho, screening_cap, outer_cap)`. The cap
5726    /// suffix is essential because a search-time partial inner mode and the
5727    /// terminal uncapped mode can share bit-identical rho.
5728    /// The binomial REML fit performs ~20-32 seed-grid pre-solves plus
5729    /// line-search revisits; with only the single `current_outer_eval` slot,
5730    /// any revisit to an earlier rho re-ran a full n-sized P-IRLS. This LRU
5731    /// returns the stored cost/gradient for those revisited rho-points.
5732    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    /// Memoizing wrapper for `PenaltySubspace` construction.
5749    ///
5750    /// The penalty-subspace eigendecomposition is shape-invariant: any two
5751    /// outer evaluations at the same `(E_transformed, ridge_passport)` produce
5752    /// bit-identical subspaces. The single-slot cache amortizes consecutive
5753    /// fixed-S queries (rank, logdet, trace) within a single outer iter.
5754    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        // The LRU is the authoritative multi-slot store; it always contains the
5800        // most-recently-stored eval too (kept in sync by `store_outer_eval`), so
5801        // a single LRU probe subsumes the old single-slot fast path while also
5802        // serving revisited (non-immediate) rho-points. `get` is a tiny linear
5803        // scan (capacity is `OUTER_EVAL_LRU_CAPACITY`) that promotes the hit to
5804        // most-recently-used; hence the write lock.
5805        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            // Keep the single-slot mirror for `previous_outer_gradient_norm`,
5814            // whose "immediately previous distinct eval" contract reads it
5815            // directly and must stay byte-for-byte unchanged.
5816            *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
5850/// Reusable scratch/runtime memory that should not be part of mathematical
5851/// state invariants.
5852pub(crate) struct RemlArena {
5853    pub(crate) cost_eval_count: RwLock<u64>,
5854    /// Number of *actual* full-n inner P-IRLS solves performed (#1575).
5855    ///
5856    /// Distinct from `cost_eval_count`, which counts every outer cost/gradient
5857    /// REQUEST including single-slot cache hits and prior short-circuits. This
5858    /// counts only the cache-missing `prepare_eval_bundlewithkey` calls — i.e.
5859    /// the genuinely expensive `O(n·p²)` inner solves the #1575 slowdown is
5860    /// about ("~150 outer cost evals each running a full n-sized P-IRLS"). A
5861    /// healthy warm-started fit performs roughly 2 inner solves per outer
5862    /// cost-eval (one value, one gradient/Hessian), so a large ratio between
5863    /// the two signals broken warm-starting or duplicate solving. This is pure
5864    /// observability: it never feeds back into the optimization and changes no
5865    /// fitted value.
5866    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    /// Canonicalized block-local penalties with pre-computed roots.
5886    /// This is the single canonical penalty representation — no full-width
5887    /// `rank × p` roots are stored separately.
5888    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    /// Explicit prior on log smoothing parameters used by the REML/LAML objective.
5900    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    /// Two-point ρ-trajectory used for second-order warm-start
5906    /// extrapolation: when the outer optimizer asks for a fit at a new
5907    /// ρ, we have `β(ρ_k)` (in `warm_start_beta`) and `β(ρ_{k-1})` (in
5908    /// `prev_warm_start_beta`). The implicit β(ρ) trajectory is locally
5909    /// linear under the FOC ∇F(β,ρ)=0, so a tangent-line prediction
5910    /// `β_predict(ρ_new) = β_k + α · (β_k − β_{k-1})` where α is the
5911    /// projection of `(ρ_new − ρ_k)` onto `(ρ_k − ρ_{k-1})` gives a
5912    /// better seed than the flat `β_k` alone — replacing PIRLS warm-
5913    /// start "use last β as-is" with a real tangent-prediction step.
5914    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    /// The #784 block-local correction's ADMISSION, frozen for the fit: `0` is
5918    /// "not yet admitted anywhere", `m + 1` is "admitted, with a block of `m`
5919    /// curvature-heavy directions" (#2748).
5920    ///
5921    /// The correction used to be admitted or declined by predicates evaluated
5922    /// at ρ — a skewness threshold `τ(n_eff)`, a quadrature-resolution test,
5923    /// and a block-dimension cap. A criterion term switched on and off by a
5924    /// function of ρ is not a function of ρ: it has a jump of the FULL `|Δ_b|`
5925    /// across every switching surface, and the outer search descends onto one,
5926    /// because the ON region's minimum is on its boundary. Measured on
5927    /// `haberman_5yr` (#2748): `max|γ| = 0.125` against `τ = 0.125` at the
5928    /// refused checkpoint, `V` jumping `1.744282e2 → 1.744593e2` — exactly
5929    /// `Δ_b = 3.1144e-2` — between two adjacent line-search trial points while
5930    /// `|g| = 2.045e-2`. No sufficient-decrease test can pass across a jump
5931    /// larger than `c₁·α·|gᵀd|` for any `α`, so `StepSizeTooSmall after 50
5932    /// attempt(s)` is the only reachable outcome.
5933    ///
5934    /// Freezing the admission is what makes the spliced objective a function.
5935    /// It also keeps the spliced GRADIENT exact: the four channels differentiate
5936    /// `Δ_b` at a FIXED block, and any admission that depends on ρ contributes
5937    /// a `∂Δ_b/∂(admission)·∂(admission)/∂ρ` term the channels do not carry —
5938    /// the same objective↔gradient desync this site already declines the splice
5939    /// over for ψ coordinates and for the Beta family.
5940    ///
5941    /// Latched on FIRST admission rather than at the first evaluation, so a fit
5942    /// whose correction never engages is bit-identical to the pre-#2748 fit and
5943    /// a fit that engaged consistently keeps the same block it always had; only
5944    /// the fits that were toggling change.
5945    pub(crate) block_correction_admission: AtomicUsize,
5946    /// Adaptive IFT step-cap controller, the hypergradient budget controller,
5947    /// and the two mode-response caches.
5948    ///
5949    /// These four lived in process-global `HashMap`s keyed by
5950    /// `self as *const _ as usize` — the state's own ADDRESS. An address is not
5951    /// an identity: it is reused the moment the state is dropped, so a fit
5952    /// whose `RemlState` landed on a recycled slot inherited the previous
5953    /// fit's quality history, budget history, and cached mode-response
5954    /// columns. `RemlState` is a stack local in every constructing caller
5955    /// (`estimate/evaluation.rs`, `estimate/optimizer.rs`, …), so consecutive
5956    /// fits on one thread reuse the same frame address near-deterministically,
5957    /// and there is no `Drop` to evict the dead entry. Two fits running
5958    /// concurrently in one process shared the tables outright.
5959    ///
5960    /// Holding them as fields makes the lifetime exactly right by
5961    /// construction: the state dies with its owner, concurrent fits cannot
5962    /// see each other, and no key can alias. They sit beside
5963    /// `ift_warm_start_cache` / `ift_cached_factor`, which were already
5964    /// per-state interior-mutability fields — these were the odd ones out.
5965    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    /// Outer-aware inner-PIRLS iteration cap for the main descent loop.
5975    ///
5976    /// Distinct from `screening_max_inner_iterations`, which is used during
5977    /// seed selection and toggles a side-effect bundle (cache writes,
5978    /// warm-start updates, KKT enforcement all suppressed). This atomic is
5979    /// purely a cap — when nonzero, the inner Newton loop is capped at
5980    /// `min(this, full_max_iterations)`, but cache writes and warm-start
5981    /// updates remain enabled. Driven by the outer optimizer to coarsen
5982    /// inner solves at early outer iterations when ρ is far from converged,
5983    /// and lifted back to full at the final accepted iter (otherwise the
5984    /// returned β would be biased by the loose cap).
5985    ///
5986    /// Both atomics are honored together as `min(screening_cap, outer_cap)`
5987    /// when both are nonzero. Default 0 (no cap from this source).
5988    pub(crate) outer_inner_cap: Arc<AtomicUsize>,
5989
5990    /// Inner-PIRLS feedback signal driven by `execute_pirls_if_needed` after
5991    /// each NON-screening solve. Stores the iteration count at which the
5992    /// inner Newton stopped, plus a flag indicating whether it converged
5993    /// (vs. hit the iteration cap). The outer first-/second-order bridges
5994    /// read these atomics to drive an adaptive `inner_cap_schedule`: the
5995    /// next outer iter's inner cap becomes `last_iters + small_margin`
5996    /// when the previous solve converged, or a geometric backoff when it
5997    /// hit the cap. This replaces the older hardcoded iter-tier schedule
5998    /// (3/5/10/20) with a cap that follows the inner solver's actual
5999    /// convergence behavior — Eisenstat-Walker style for the inner
6000    /// quadratic loop. Default 0 / false (no signal yet — first outer
6001    /// iter falls back to a coarse iter-count tier).
6002    pub(crate) last_inner_iters: Arc<AtomicUsize>,
6003    pub(crate) last_inner_converged: Arc<AtomicBool>,
6004
6005    /// Cached state from the most recent successful PIRLS solve, used by
6006    /// the IFT-based warm-start predictor.
6007    ///
6008    /// The implicit-function theorem applied to the FOC ∇_β F(β,ρ)=0
6009    /// gives `dβ/dρ_k = -H_pen^{-1} · (e^{ρ_k} · S_k · β)`. A first-order
6010    /// Taylor predictor reads
6011    /// `β_predict(ρ_new) = β_cur − Σ_k Δρ_k · H_pen^{-1} · (e^{ρ_cur_k} · S_k · β_cur)`.
6012    /// This is a strict superset of the tangent-line predictor's
6013    /// requirements: works after a single successful solve (tangent-line
6014    /// needs two prior fits), and gives the EXACT first-order Jacobian
6015    /// of the implicit β(ρ) trajectory rather than a secant proxy along one
6016    /// ρ-direction.
6017    ///
6018    /// Populated in `updatewarm_start_from` when PIRLS converges; cleared
6019    /// on failure, on `reset_surface`, and on link-state changes.
6020    pub(crate) ift_warm_start_cache: RwLock<Option<IftWarmStartCache>>,
6021
6022    /// Persisted Levenberg-Marquardt damping coefficient from the most
6023    /// recent successful PIRLS solve, bit-packed into an `AtomicU64`
6024    /// (`f64::to_bits` low 64 bits). Read at the start of
6025    /// `execute_pirls_if_needed` and written into the
6026    /// `PirlsConfig::initial_lm_lambda` hint so the inner Newton seeds
6027    /// `λ_LM` near the damping the previous solve discovered, instead
6028    /// of cold-starting at `1e-6` and burning 4-6 halving steps to
6029    /// recover. `0` (the default) signals "no hint"; the inner solver
6030    /// clamps any positive hint into `[1e-6, 1e-3]` so a stale value
6031    /// cannot destabilize the next solve. Reset on `reset_surface` and
6032    /// on failed solves.
6033    pub(crate) last_pirls_lm_lambda: Arc<AtomicU64>,
6034
6035    /// Negative-Binomial overdispersion `theta` frozen for the smoothing-
6036    /// parameter (λ) search (#1082), bit-packed `f64` (`f64::to_bits`). `0`
6037    /// (the default) signals "not yet frozen". On the first non-screening
6038    /// λ-search inner solve of an estimated-θ NB fit, the maximum-likelihood θ
6039    /// at the solve's converged η is computed once and stored here; every subsequent
6040    /// λ-search evaluation pins the inner solve to this value via
6041    /// `GlmLikelihoodSpec::with_negbin_theta_frozen_for_search`, so the REML
6042    /// criterion `F(ρ) = REML(ρ, θ_frozen)` is a stationary function of ρ and
6043    /// the outer optimizer converges instead of chasing the per-eval θ drift
6044    /// that the estimated path injects. The single final reported fit still
6045    /// ML-refreshes θ at the converged η. Reset on `reset_surface`.
6046    pub(crate) frozen_negbin_theta: Arc<AtomicU64>,
6047
6048    /// Tweedie exponential-dispersion `phi` frozen for the smoothing-parameter
6049    /// (λ) search (#1477), bit-packed `f64` (`f64::to_bits`). `0` (the default)
6050    /// signals "not yet frozen". On the first non-screening λ-search inner solve
6051    /// of an estimated-φ Tweedie fit, the converged-η Pearson `phî` is captured
6052    /// once and stored here; every subsequent λ-search evaluation pins the inner
6053    /// solve to this value via
6054    /// `GlmLikelihoodSpec::with_tweedie_phi_frozen_for_search`, so the REML
6055    /// criterion `F(ρ) = REML(ρ, φ_frozen)` is a stationary function of ρ. The
6056    /// Tweedie LAML omits the `phi`-dependent saddlepoint normalizer, so a `phi`
6057    /// drifting with each warm-start η lets the criterion reward dispersion
6058    /// inflation and rail a double-penalty null-space `λ` to the box bound (the
6059    /// #1477 boundary blow-up). The single final reported fit still
6060    /// Pearson-refreshes `phi` at the converged η. Reset on `reset_surface`.
6061    pub(crate) frozen_tweedie_phi: Arc<AtomicU64>,
6062
6063    /// Gamma shape `k = 1/φ` frozen for the smoothing-parameter (λ) search
6064    /// (#1074), bit-packed `f64` (`f64::to_bits`). `0` (the default) signals
6065    /// "not yet frozen". On the first non-screening λ-search inner solve of an
6066    /// estimated-shape Gamma fit, the seed's converged-η MLE `k̂` is captured
6067    /// once and stored here; every subsequent λ-search evaluation pins the inner
6068    /// solve to this value via
6069    /// `GlmLikelihoodSpec::with_gamma_shape_frozen_for_search`, so the REML
6070    /// criterion `F(ρ) = REML(ρ, k_frozen)` is a stationary function of ρ. With
6071    /// `k` estimated the inner solver re-derives it from each warm-start η, and
6072    /// because the Gamma working weight is `W = prior·k` and the
6073    /// omitting-constants log-likelihood is `−k·½D`, a `k` swinging with η makes
6074    /// BOTH the curvature `H = k·XᵀX + λS` and the data-fit `k·½D` jump with ρ —
6075    /// the criterion grows deterministic spikes that floor the projected
6076    /// gradient and rail `λ` to the over-smoothed corner (the #1074 te/Gamma
6077    /// tensor under-recovery). The single final reported fit still ML-refreshes
6078    /// `k` at the converged η. Reset on `reset_surface`.
6079    pub(crate) frozen_gamma_shape: Arc<AtomicU64>,
6080
6081    /// Beta-regression precision `phi` frozen for the smoothing-parameter (λ)
6082    /// search (#2369), bit-packed `f64` (`f64::to_bits`). `0` (the default)
6083    /// signals "not yet frozen". On the first non-screening λ-search inner solve
6084    /// of an estimated-φ Beta fit, the seed's converged-η Pearson `phî` is
6085    /// captured once and stored here; every subsequent λ-search evaluation pins
6086    /// the inner solve to this value via
6087    /// `GlmLikelihoodSpec::with_beta_phi_frozen_for_search`, so the REML
6088    /// criterion `F(ρ) = REML(ρ, φ_frozen)` is a stationary function of ρ. With
6089    /// `phi` estimated the inner solver re-derives it from each warm-start η, and
6090    /// because the Beta precision does not factor out of the digamma mean score
6091    /// `∂ℓ/∂β = φ·Σ xᵢ(y*ᵢ − μ*ᵢ)`, a `phi` swinging with η makes both β̂(ρ) and
6092    /// the REML data-fit / log-det terms jump with ρ while the analytic outer
6093    /// gradient holds `phi` fixed — the projected gradient floors above tolerance
6094    /// and the optimizer refuses ("NOT STATIONARY"), the family-unusable #2369
6095    /// signature, identical to the sibling Gamma/Tweedie/NB drift. The single
6096    /// final reported fit still Pearson-refreshes `phi` at the converged η. Reset
6097    /// on `reset_surface`.
6098    pub(crate) frozen_beta_phi: Arc<AtomicU64>,
6099
6100    /// Last observed IFT-prediction residual (`‖β_converged − β_predicted‖
6101    /// / ‖β_converged‖`) from the most recent non-screening solve where
6102    /// the predictor was actually consumed. Bit-packed `f64` (low 64
6103    /// bits via `f64::to_bits`).
6104    ///
6105    /// "No signal yet" is encoded as a NaN bit-pattern
6106    /// (`IFT_RESIDUAL_NO_SIGNAL_BITS`). The original `0` sentinel
6107    /// collided with `f64::to_bits(0.0) == 0` — a true residual of
6108    /// exactly 0 (degenerate but mathematically possible if every
6109    /// β_predicted_i matched β_converged_i to bit-equality) would
6110    /// have been indistinguishable from "predictor never reported".
6111    /// NaN's self-inequality makes the sentinel unambiguous: any
6112    /// stored finite non-negative value is genuine signal.
6113    ///
6114    /// Read by `predict_warm_start_beta_ift_with_outcome` to drive the adaptive
6115    /// |Δρ| cap (`adaptive_ift_max_drho`): a small residual loosens
6116    /// the cap, a large one tightens it. Replaces the previous
6117    /// hardcoded `IFT_WARM_START_MAX_DRHO = 2.0` constant with a
6118    /// data-driven policy, so the predictor adapts to the empirical
6119    /// faithfulness of the linearization at this surface's scale.
6120    /// Reset on `reset_surface` and on failed solves.
6121    pub(crate) last_ift_prediction_residual: Arc<AtomicU64>,
6122
6123    /// Last observed gain ratio of the accepted LM step
6124    /// (`actual_reduction / predicted_reduction`) from the most recent
6125    /// non-screening PIRLS solve. Bit-packed `f64` with the same NaN
6126    /// sentinel discipline as `last_ift_prediction_residual`: NaN bits
6127    /// (`IFT_RESIDUAL_NO_SIGNAL_BITS`) encode "no signal yet" so a
6128    /// recorded ratio of exactly 0 (degenerate but possible) doesn't
6129    /// collide with the no-signal token.
6130    ///
6131    /// Used by `first_order_inner_cap_schedule` as a third quality
6132    /// signal alongside `last_iters` and `last_converged`. A small
6133    /// `accept_rho` (model overstating predicted reduction) is a hint
6134    /// the next iter's inner Newton may need extra margin even when
6135    /// the previous solve converged in few iters. Reset on
6136    /// `reset_surface` and on failed solves.
6137    pub(crate) last_pirls_accept_rho: Arc<AtomicU64>,
6138
6139    /// Cached Cholesky factorization of `IftWarmStartCache::penalized_hessian_transformed`.
6140    /// Lazily computed on the first IFT predict call after a fresh
6141    /// `updatewarm_start_from`, then reused by every subsequent
6142    /// predict call until the IFT cache is invalidated. At large-scale
6143    /// scale where p can reach several thousand, the dense Cholesky
6144    /// is O(p³)/3 — multiple seconds per refactor — so caching saves
6145    /// real wall time across the typical 5-10 IFT predict calls per
6146    /// outer fit. Reset jointly with `ift_warm_start_cache` (on
6147    /// reset_surface, on link-state changes, on failed PIRLS solves,
6148    /// and whenever a new H_pen replaces the cached one).
6149    pub(crate) ift_cached_factor: RwLock<Option<Arc<dyn gam_linalg::matrix::FactorizedSystem>>>,
6150
6151    /// When set, the penalties have Kronecker (tensor-product) structure and
6152    /// the REML evaluator can use O(∏q_j) logdet instead of O(p³) eigendecomposition.
6153    /// Populated via `set_kronecker_penalty_system` after construction.
6154    pub(crate) kronecker_penalty_system: Option<gam_terms::smooth::KroneckerPenaltySystem>,
6155    /// Full Kronecker factored basis (marginal designs + penalties + dims).
6156    /// Used by P-IRLS for factored reparameterization.
6157    pub(crate) kronecker_factored: Option<gam_terms::basis::KroneckerFactoredBasis>,
6158
6159    /// Precomputed `(XᵀWX, XᵀW(y − offset))` for the Gaussian + Identity
6160    /// outer REML loop, populated once before the outer optimizer when the
6161    /// family / link / constraint preconditions hold and the design supports
6162    /// the Identity short-circuit at `pirls.rs:6237`. When present, each
6163    /// inner `solve_penalized_least_squares_implicit` reads these matrices
6164    /// instead of restreaming the O(N·p²) GEMM and O(N·p) matvec per outer
6165    /// iteration — the penalty `λ·S` is still added per-λ.
6166    ///
6167    /// Invalidated jointly with the design in `reset_surface`.
6168    pub(crate) gaussian_fixed_cache: RwLock<Option<Arc<crate::pirls::GaussianFixedCache>>>,
6169    /// Once-built row placeholders for fixed-design Gaussian value-only
6170    /// evaluations. These rows depend only on the immutable
6171    /// `(offset, y, weights, link)` surface, so every distinct rho probe can
6172    /// share them while its exact deviance/score comes from the Gaussian
6173    /// sufficient statistics. Full derivative/final-fit evaluations never
6174    /// consume this slot and continue to materialize their exact `X beta`
6175    /// rows (#2435).
6176    pub(crate) gaussian_cost_only_frozen_rows:
6177        RwLock<Option<Arc<crate::pirls::GaussianFrozenRows>>>,
6178    /// Conditioned-frame exact ψ-derivatives `(∂XᵀWX/∂ψ, ∂XᵀW(y−offset)/∂ψ)`
6179    /// for the SINGLE design-moving spatial hyperparameter (#1033b), assembled
6180    /// n-free from the certified Chebyshev ψ-Gram tensor and installed beside
6181    /// `gaussian_fixed_cache` at the same in-window trial. When present the
6182    /// Gaussian-identity ψ-gradient HyperCoord (`a_j`, `g_j`, dense `B_j`) is
6183    /// formed from these k×k objects instead of realizing and contracting the
6184    /// n×k ∂X/∂ψ slab — retiring the second per-trial n-pass. Lives in the
6185    /// SAME conditioned column frame as `gaussian_fixed_cache.xtwx_orig`, so
6186    /// the hyper-coord builder transforms it by the per-eval Qs/free-basis the
6187    /// same way it transforms the streamed Gram. Invalidated with the design.
6188    pub(crate) gaussian_psi_gram_deriv:
6189        RwLock<Option<Arc<(ndarray::Array2<f64>, ndarray::Array1<f64>)>>>,
6190    /// Conditioned-frame exact ψ-derivative pair `(∂XᵀWX/∂ψ, ∂XᵀW(y−offset)/∂ψ)`
6191    /// for the SINGLE design-moving spatial hyperparameter in the GLM (frozen-W)
6192    /// lane (#1033 / #1111), assembled n-free from
6193    /// [`crate::glm_sufficient_lane::FrozenWeightGramTensor::gradient_pair_if_sound`]
6194    /// and installed beside `glm_first_step_gram` at the same in-window
6195    /// drift-OK trial. When present, the GLM ψ-gradient HyperCoord serves its
6196    /// envelope `a_j` and score `g_j` from these k×k objects instead of
6197    /// realizing and contracting the n×k ∂X/∂ψ slab — the second per-trial
6198    /// n-pass. Unlike the Gaussian lane the Hessian curvature `B_j` is NOT
6199    /// served from the tensor: for a GLM the per-trial `B_j` term
6200    /// `X_τᵀWX + XᵀWX_τ` is irreducibly n-dependent (the moving working weight
6201    /// `W` does not factor out of a frozen-W k×k object), so `B_j` keeps the
6202    /// exact streamed slab (#1033). Lives in the SAME conditioned column frame
6203    /// as `glm_first_step_gram` / `gaussian_fixed_cache.xtwx_orig`, so the
6204    /// hyper-coord builder transforms it by the per-eval Qs/free-basis the same
6205    /// way. NOT family-gated (the GLM lane's own slot). Invalidated with the
6206    /// design.
6207    pub(crate) glm_psi_gram_deriv:
6208        RwLock<Option<Arc<(ndarray::Array2<f64>, ndarray::Array1<f64>)>>>,
6209    /// Frozen-weight first-Fisher-step data-fit Gram `XᵀWX` for the GLM
6210    /// design-moving ψ-sweep (#1111 / #1033 mechanism (c)), in the conditioned
6211    /// (original / `x_fit`) column frame — the SAME frame as
6212    /// `gaussian_fixed_cache.xtwx_orig`.
6213    ///
6214    /// Assembled n-free per in-window ψ-trial from the certified frozen-weight
6215    /// Chebyshev tensor ([`crate::glm_sufficient_lane::FrozenWeightGramTensor`])
6216    /// and installed only when the trial's converged working weight has not
6217    /// drifted past tolerance from the frozen snapshot. When present, the GLM
6218    /// inner P-IRLS serves its FIRST Fisher-scoring iteration's `XᵀWX` from this
6219    /// cache instead of restreaming the O(N·p²) weighted cross-product — the
6220    /// dominant per-trial n-term in a large-n Poisson/Binomial κ-sweep. The
6221    /// penalty `Sλ` is still added per-λ on top, and every subsequent inner
6222    /// iteration restreams the true (moving) `W`, so the converged β̂ is
6223    /// unchanged; only the first-iteration Gram build is elided. Unlike
6224    /// `gaussian_fixed_cache` this is NOT family-gated — it is the GLM lane's
6225    /// own slot, consumed once per inner solve. Invalidated with the design in
6226    /// `reset_surface`.
6227    pub(crate) glm_first_step_gram: RwLock<Option<Arc<ndarray::Array2<f64>>>>,
6228    /// Previous successful non-Gaussian fixed-design data-fit Gram `XᵀWX` in
6229    /// the conditioned original frame, keyed to `warm_start_beta`.
6230    ///
6231    /// When the next outer trial uses a flat warm start, its first PIRLS
6232    /// curvature build evaluates at the same `η = Xβ` as the previous converged
6233    /// solve, so the Hessian weights and `XᵀWX` are identical. Reusing this
6234    /// original-frame Gram skips one dense `O(n·p²)` pass per warm-started
6235    /// trial while still letting later PIRLS iterations restream the moving
6236    /// weights. IFT/tangent-predicted starts do not consume this cache.
6237    ///
6238    /// The cached dense `p×p` Gram is coupled to a [`MemoryGovernor`] ledger
6239    /// reservation ([`gam_runtime::resource::Governed`]) so a long-lived cache
6240    /// entry never holds dense bytes off-ledger; a reservation refusal under
6241    /// joint pressure skips caching (the streamed path remains available).
6242    pub(crate) flat_glm_first_step_gram:
6243        RwLock<Option<gam_runtime::resource::Governed<Arc<ndarray::Array2<f64>>>>>,
6244    /// Stable disk-cache key for the current realized REML surface. Computed
6245    /// lazily because it hashes the row-chunked design and data vectors.
6246    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    /// Ensures the process attempts at most one disk restore per surface.
6251    pub(crate) persistent_warm_start_loaded: AtomicBool,
6252    /// Scoped counter disabling all disk persistence from canonical anchors and
6253    /// cost-only posterior/probe evaluations. In-memory warm starts still
6254    /// update; loads, writes, and eviction sweeps are suppressed together so a
6255    /// probe cannot inherit cache history or mutate it.
6256    pub(crate) persistent_warm_start_store_suppression: AtomicUsize,
6257    /// Caller-configured cross-process warm-start capability.
6258    ///
6259    /// Empty by default: in-memory warm starts remain on, while no ambient
6260    /// filesystem root is discovered. The estimate constructor attaches the
6261    /// explicit capability at most once after nuisance anchoring; clones of the
6262    /// capability share its lazy/opened store.
6263    pub(crate) persistent_warm_start_store:
6264        std::sync::OnceLock<gam_runtime::warm_start::ConfiguredWarmStartStore>,
6265    /// #1033: memoized fit-invariant O(n) response/weight scalars.
6266    ///
6267    /// `gaussian_weight_log_sum_half` (`½·Σ log wᵢ`),
6268    /// `gaussian_dp_floor_scale` (the weighted null deviance `D₀`), the
6269    /// positive-weight observation count, and `rho_weight_anchor`
6270    /// (`mean(log wᵢ)`) are pure functions of the borrowed `(y, weights)` —
6271    /// fields `reset_surface` NEVER reassigns — so they are constant for the
6272    /// whole life of the `RemlState`. Memoizing them once per fit keeps every
6273    /// subsequent outer trial in coefficient space. Plain scalar closures, no
6274    /// rayon inside `get_or_init` (no deadlock trap).
6275    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}