Skip to main content

gam_models/fit_orchestration/drivers/
design_construction.rs

1// #1521: `build_term_collection_design` and its term-design subgraph were
2// relocated DOWN into `gam_terms::smooth` (see `gam_terms::smooth::term_design`).
3// The joint-build variants below STAY here: they return a `gam_solve`
4// `EstimationError` / call `freeze_term_collection_from_design`
5// (`spatial_optimization.rs`), so they belong to the gam-models orchestration
6// tier. They reach the relocated `build_term_collection_design_inner` /
7// `build_term_collection_design` via the module's `use gam_terms::smooth::*`.
8pub fn build_term_collection_designs_joint(
9    data: ArrayView2<'_, f64>,
10    specs: &[TermCollectionSpec],
11) -> Result<Vec<TermCollectionDesign>, BasisError> {
12    for spec in specs {
13        validate_term_collection_finite_inputs(data, spec)?;
14    }
15    let smooth_blocks = specs
16        .iter()
17        .map(|spec| spec.smooth_terms.clone())
18        .collect::<Vec<_>>();
19    let planned_blocks = plan_joint_spatial_centers_for_term_blocks(data, &smooth_blocks)?;
20    let mut out = Vec::with_capacity(specs.len());
21    for (spec, planned_terms) in specs.iter().zip(planned_blocks.into_iter()) {
22        let mut planned_spec = spec.clone();
23        planned_spec.smooth_terms = planned_terms;
24        out.push(build_term_collection_design_inner(data, &planned_spec)?);
25    }
26    Ok(out)
27}
28
29pub fn build_term_collection_designs_and_freeze_joint(
30    data: ArrayView2<'_, f64>,
31    specs: &[TermCollectionSpec],
32) -> Result<(Vec<TermCollectionDesign>, Vec<TermCollectionSpec>), EstimationError> {
33    let designs = build_term_collection_designs_joint(data, specs)?;
34    let mut resolved_specs = Vec::with_capacity(specs.len());
35    for (spec, design) in specs.iter().zip(designs.iter()) {
36        resolved_specs.push(freeze_term_collection_from_design(spec, design)?);
37    }
38    Ok((designs, resolved_specs))
39}
40
41pub fn fit_term_collection_forspec(
42    data: ArrayView2<'_, f64>,
43    y: ArrayView1<'_, f64>,
44    weights: ArrayView1<'_, f64>,
45    offset: ArrayView1<'_, f64>,
46    spec: &TermCollectionSpec,
47    family: LikelihoodSpec,
48    options: &FitOptions,
49) -> Result<FittedTermCollection, EstimationError> {
50    fit_term_collection_forspecwith_heuristic_lambdas(
51        data, y, weights, offset, spec, None, family, options,
52    )
53}
54
55pub fn fit_term_collection_with_coefficient_groups(
56    data: ArrayView2<'_, f64>,
57    y: ArrayView1<'_, f64>,
58    weights: ArrayView1<'_, f64>,
59    offset: ArrayView1<'_, f64>,
60    spec: &TermCollectionSpec,
61    groups: &[CoefficientGroupSpec],
62    family: LikelihoodSpec,
63    options: &FitOptions,
64) -> Result<FittedTermCollection, EstimationError> {
65    if groups.is_empty() {
66        return fit_term_collection_forspec(data, y, weights, offset, spec, family, options);
67    }
68    let design = build_term_collection_design_with_policy(data, spec, &options.resource_policy)?;
69    let base_fit_opts = adaptive_fit_options_base(options, &design);
70    let realized = design
71        .realize_coefficient_groups(groups, &base_fit_opts.rho_prior)
72        .map_err(EstimationError::BasisError)?;
73    let mut grouped_options = base_fit_opts.clone();
74    grouped_options.rho_prior = realized.rho_prior;
75    let fitted = FittedTermCollection {
76        fit: gam_solve::estimate::fit_gam_with_penalty_specs(
77            design.design.clone(),
78            y,
79            weights,
80            offset,
81            realized.penalty_specs,
82            realized.nullspace_dims,
83            family.clone(),
84            &grouped_options,
85        )?,
86        design,
87        adaptive_diagnostics: None,
88    };
89    enforce_term_constraint_feasibility(&fitted.design, &fitted.fit)?;
90    Ok(fitted)
91}
92
93pub fn fit_term_collection_with_penalty_block_gamma_prior_callback<F>(
94    data: ArrayView2<'_, f64>,
95    y: ArrayView1<'_, f64>,
96    weights: ArrayView1<'_, f64>,
97    offset: ArrayView1<'_, f64>,
98    spec: &TermCollectionSpec,
99    callback: F,
100    family: LikelihoodSpec,
101    options: &FitOptions,
102) -> Result<FittedTermCollection, EstimationError>
103where
104    F: FnMut(&PenaltyBlockGammaPriorMetadata<'_>) -> Option<(f64, f64)>,
105{
106    let design = build_term_collection_design_with_policy(data, spec, &options.resource_policy)?;
107    let mut fit_opts = adaptive_fit_options_base(options, &design);
108    fit_opts.rho_prior = realize_penalty_block_gamma_priors(&design, callback)
109        .map_err(EstimationError::BasisError)?;
110    let fitted = FittedTermCollection {
111        fit: fit_gamwith_heuristic_lambdas(
112            design.design.clone(),
113            y,
114            weights,
115            offset,
116            &design.penalties,
117            None,
118            family.clone(),
119            &fit_opts,
120        )?,
121        design,
122        adaptive_diagnostics: None,
123    };
124    enforce_term_constraint_feasibility(&fitted.design, &fitted.fit)?;
125    Ok(fitted)
126}
127
128pub fn fit_term_collection_with_penalty_block_gamma_priors(
129    data: ArrayView2<'_, f64>,
130    y: ArrayView1<'_, f64>,
131    weights: ArrayView1<'_, f64>,
132    offset: ArrayView1<'_, f64>,
133    spec: &TermCollectionSpec,
134    priors: &[(String, f64, f64)],
135    family: LikelihoodSpec,
136    options: &FitOptions,
137) -> Result<FittedTermCollection, EstimationError> {
138    let design = build_term_collection_design_with_policy(data, spec, &options.resource_policy)?;
139    let mut fit_opts = adaptive_fit_options_base(options, &design);
140    fit_opts.rho_prior = realize_keyed_penalty_block_gamma_priors(&design, priors)
141        .map_err(EstimationError::BasisError)?;
142    let fitted = FittedTermCollection {
143        fit: fit_gamwith_heuristic_lambdas(
144            design.design.clone(),
145            y,
146            weights,
147            offset,
148            &design.penalties,
149            None,
150            family.clone(),
151            &fit_opts,
152        )?,
153        design,
154        adaptive_diagnostics: None,
155    };
156    enforce_term_constraint_feasibility(&fitted.design, &fitted.fit)?;
157    Ok(fitted)
158}
159
160pub fn fit_term_collection_with_coefficient_groups_and_penalty_block_gamma_priors(
161    data: ArrayView2<'_, f64>,
162    y: ArrayView1<'_, f64>,
163    weights: ArrayView1<'_, f64>,
164    offset: ArrayView1<'_, f64>,
165    spec: &TermCollectionSpec,
166    groups: &[CoefficientGroupSpec],
167    priors: &[(String, f64, f64)],
168    family: LikelihoodSpec,
169    options: &FitOptions,
170) -> Result<FittedTermCollection, EstimationError> {
171    if groups.is_empty() {
172        return fit_term_collection_with_penalty_block_gamma_priors(
173            data, y, weights, offset, spec, priors, family, options,
174        );
175    }
176    if priors.is_empty() {
177        return fit_term_collection_with_coefficient_groups(
178            data, y, weights, offset, spec, groups, family, options,
179        );
180    }
181
182    // The base design already emits one term-named function-space ridge per
183    // recoverable linear effect, so keyed priors and coefficient groups address
184    // the same authoritative λ coordinates as every other fit path.
185    let design = build_term_collection_design_with_policy(data, spec, &options.resource_policy)?;
186    let base_fit_opts = adaptive_fit_options_base(options, &design);
187    let base_rho_prior = realize_keyed_penalty_block_gamma_priors(&design, priors)
188        .map_err(EstimationError::BasisError)?;
189    let realized = design
190        .realize_coefficient_groups(groups, &base_rho_prior)
191        .map_err(EstimationError::BasisError)?;
192    let mut grouped_options = base_fit_opts.clone();
193    grouped_options.rho_prior = realized.rho_prior;
194    let fitted = FittedTermCollection {
195        fit: gam_solve::estimate::fit_gam_with_penalty_specs(
196            design.design.clone(),
197            y,
198            weights,
199            offset,
200            realized.penalty_specs,
201            realized.nullspace_dims,
202            family.clone(),
203            &grouped_options,
204        )?,
205        design,
206        adaptive_diagnostics: None,
207    };
208    enforce_term_constraint_feasibility(&fitted.design, &fitted.fit)?;
209    Ok(fitted)
210}
211
212fn fit_term_collection_forspecwith_heuristic_lambdas(
213    data: ArrayView2<'_, f64>,
214    y: ArrayView1<'_, f64>,
215    weights: ArrayView1<'_, f64>,
216    offset: ArrayView1<'_, f64>,
217    spec: &TermCollectionSpec,
218    heuristic_lambdas: Option<&[f64]>,
219    family: LikelihoodSpec,
220    options: &FitOptions,
221) -> Result<FittedTermCollection, EstimationError> {
222    let adaptive_opts = options.adaptive_regularization.clone().unwrap_or_default();
223    let resolved_spec;
224    let design_spec = if adaptive_opts.enabled {
225        resolved_spec = ensure_matern_adaptive_center_resolution(spec, data.nrows());
226        &resolved_spec
227    } else {
228        spec
229    };
230    let base_design =
231        build_term_collection_design_with_policy(data, design_spec, &options.resource_policy)?;
232    fit_term_collection_on_realized_design(
233        y,
234        weights,
235        offset,
236        design_spec,
237        &base_design,
238        heuristic_lambdas,
239        family,
240        options,
241    )
242}
243
244fn ensure_matern_adaptive_center_resolution(
245    spec: &TermCollectionSpec,
246    n_rows: usize,
247) -> TermCollectionSpec {
248    let mut out = spec.clone();
249    for term in &mut out.smooth_terms {
250        let gam_terms::smooth::SmoothBasisSpec::Matern {
251            feature_cols,
252            spec: matern,
253            ..
254        } = &mut term.basis
255        else {
256            continue;
257        };
258        if let gam_terms::basis::CenterStrategy::FarthestPoint { num_centers } =
259            &mut matern.center_strategy
260        {
261            // Exact spatial-adaptive regularization estimates three operator
262            // weights from the fitted Matérn field and its first/second
263            // collocation derivatives.  That is a richer hyperproblem than the
264            // ordinary quadratic Matérn fit: with fewer centers than the
265            // coordinate dimension's linear scale, the radial span cannot carry
266            // even low-order directional structure, so REML can only explain the
267            // signal by pushing the adaptive operator weights into the
268            // over-smoothed mean basin.  Treat user-supplied FarthestPoint counts
269            // as a lower bound for this exact-adaptive path and ensure a modest
270            // O(d) collocation resolution.  Existing larger bases are left
271            // untouched, and the cap at n_rows preserves the reduced-rank
272            // contract.
273            let min_centers = (4 * feature_cols.len()).min(n_rows).max(*num_centers);
274            *num_centers = min_centers;
275        }
276    }
277    out
278}
279
280fn has_bounded_linear_terms(spec: &TermCollectionSpec) -> bool {
281    spec.linear_terms.iter().any(|term| {
282        matches!(
283            term.coefficient_geometry,
284            LinearCoefficientGeometry::Bounded { .. }
285        )
286    })
287}
288
289fn fit_term_collection_on_realized_design(
290    y: ArrayView1<'_, f64>,
291    weights: ArrayView1<'_, f64>,
292    offset: ArrayView1<'_, f64>,
293    spec: &TermCollectionSpec,
294    design: &TermCollectionDesign,
295    heuristic_lambdas: Option<&[f64]>,
296    family: LikelihoodSpec,
297    options: &FitOptions,
298) -> Result<FittedTermCollection, EstimationError> {
299    if has_bounded_linear_terms(spec) {
300        return fit_bounded_term_collection_with_design(
301            y,
302            weights,
303            offset,
304            spec,
305            design,
306            heuristic_lambdas,
307            family,
308            options,
309        );
310    }
311    let mut base_fit_opts = adaptive_fit_options_base(options, design);
312    // Lift the symmetric log-λ cap off the smoothing coordinates of
313    // well-determined Gaussian-identity B-spline / thin-plate / tensor smooths so
314    // REML can drive λ to the value the data wants — including λ → ∞ when a
315    // term's signal lives in its penalty null space (#1271 single-penalty tp/ps,
316    // #1266 double-penalty selection). Length-safe: only fires when the inner ρ
317    // aligns 1:1 with the penalty blocks (see `relax_smoothing_rho_prior`).
318    base_fit_opts.rho_prior = relax_smoothing_rho_prior(options, design);
319    let fitted = FittedTermCollection {
320        fit: fit_gamwith_heuristic_lambdas(
321            design.design.clone(),
322            y,
323            weights,
324            offset,
325            &design.penalties,
326            heuristic_lambdas,
327            family.clone(),
328            &base_fit_opts,
329        )?,
330        design: design.clone(),
331        adaptive_diagnostics: None,
332    };
333    enforce_term_constraint_feasibility(&fitted.design, &fitted.fit)?;
334
335    let adaptive_opts = options.adaptive_regularization.clone().unwrap_or_default();
336    if !adaptive_opts.enabled {
337        return Ok(fitted);
338    }
339    let runtime_caches = extract_spatial_operator_runtime_caches(spec, &fitted.design)?;
340    if runtime_caches.is_empty() {
341        return Ok(fitted);
342    }
343    // Spatial-adaptive overlay always runs when the operator caches are
344    // non-empty. Catastrophic-overfit protection lives in the operator-log-λ
345    // box bound (Fix B at the BFGS bounds construction), which caps maximum
346    // unpenalization regardless of n. Production fits at n≈300K must run the
347    // overlay; the previous n-gate (n < max(4·p_total, 200)) silently skipped
348    // it for any small-n test, contradicting that contract.
349    fit_term_collectionwith_exact_spatial_adaptive_regularization(
350        fitted,
351        y,
352        weights,
353        offset,
354        family,
355        options,
356        &runtime_caches,
357    )
358}
359
360#[derive(Clone)]
361struct SpatialOperatorRuntimeCache {
362    termname: String,
363    feature_cols: Vec<usize>,
364    coeff_global_range: Range<usize>,
365    mass_penalty_global_idx: usize,
366    tension_penalty_global_idx: usize,
367    stiffness_penalty_global_idx: usize,
368    d0: Array2<f64>,
369    d1: Array2<f64>,
370    d2: Array2<f64>,
371    collocation_points: Array2<f64>,
372    dimension: usize,
373}
374
375#[derive(Clone)]
376struct SpatialAdaptiveWeights {
377    inv_magweight: Array1<f64>,
378    invgradweight: Array1<f64>,
379    inv_lapweight: Array1<f64>,
380}
381
382#[derive(Clone)]
383struct CharbonnierScalarBlockState {
384    signal: Array1<f64>,
385    radius: Array1<f64>,
386    epsilon: f64,
387}
388
389impl CharbonnierScalarBlockState {
390    fn from_signal(signal: Array1<f64>, epsilon: f64) -> Self {
391        let eps = epsilon.max(1e-12);
392        let radius = signal.mapv(|t| (t * t + eps * eps).sqrt());
393        Self {
394            signal,
395            radius,
396            epsilon: eps,
397        }
398    }
399
400    fn absolute_signal(&self) -> Array1<f64> {
401        self.signal.mapv(f64::abs)
402    }
403
404    fn penalty_value(&self) -> f64 {
405        self.radius.iter().map(|r| r - self.epsilon).sum::<f64>()
406    }
407
408    fn betagradient_coeff(&self) -> Array1<f64> {
409        Array1::from_iter(
410            self.signal
411                .iter()
412                .zip(self.radius.iter())
413                .map(|(t, r)| t / r),
414        )
415    }
416
417    fn betahessian_diag(&self) -> Array1<f64> {
418        let eps2 = self.epsilon * self.epsilon;
419        self.radius.mapv(|r| eps2 / r.powi(3))
420    }
421
422    fn log_epsilon_gradient_terms(&self) -> Array1<f64> {
423        let epsilon = self.epsilon;
424        let eps2 = epsilon * epsilon;
425        self.radius.mapv(|r| eps2 / r - epsilon)
426    }
427
428    fn log_epsilon_betagradient_coeff(&self) -> Array1<f64> {
429        let eps2 = self.epsilon * self.epsilon;
430        Array1::from_iter(
431            self.signal
432                .iter()
433                .zip(self.radius.iter())
434                .map(|(t, r)| -eps2 * t / r.powi(3)),
435        )
436    }
437
438    fn log_epsilon_hessian_terms(&self) -> Array1<f64> {
439        let epsilon = self.epsilon;
440        let eps2 = epsilon * epsilon;
441        let eps4 = eps2 * eps2;
442        self.radius
443            .mapv(|r| 2.0 * eps2 / r - eps4 / r.powi(3) - epsilon)
444    }
445
446    fn surrogateweights_posterior_snr(
447        &self,
448        variance: &Array1<f64>,
449        weight_floor: f64,
450        weight_ceiling: f64,
451    ) -> (Array1<f64>, Array1<f64>) {
452        // Posterior-SNR (credible-magnitude) reweighting of the scalar MM
453        // majorizer.
454        //
455        // The magnitude-only surrogate weight uses the *point-estimate* radius
456        //
457        //   r_k^mag = sqrt( t_k^2 + eps^2 ),   t_k = (D0 beta_hat)_k,
458        //   w_k     = 1 / r_k^mag.
459        //
460        // The weight multiplies the local quadratic surrogate penalty
461        // w_k (D0 beta)^2, so a *small* w_k leaves the response un-penalized
462        // (treated as a genuine feature) and a *large* w_k pulls it toward zero
463        // (enforces flatness). The failure of the point-estimate radius is that
464        // a response t_k which is large only because it is poorly determined
465        // gets a tiny weight and is left un-penalized — the weight chases noise
466        // in low-information regions.
467        //
468        // Resolution via the posterior second moment under the working-Laplace
469        // posterior beta ~ N(beta_hat, Sigma_beta), Sigma_beta = H^{-1}: the
470        // variance of the response is
471        //
472        //   Var( (D0 beta)_k ) = (D0 Sigma_beta D0^T)_kk >= 0,
473        //
474        // and the *credible* (noise-floor-corrected) squared magnitude is
475        //
476        //   t_k^credible^2 = max( t_k^2 - Var(...)_k , 0 ),
477        //   r_k^snr        = sqrt( t_k^credible^2 + eps^2 ),
478        //   w_k            = 1 / r_k^snr.
479        //
480        // The principled fix evaluates the MM weight at the *credible* (noise-
481        // floor-corrected) squared magnitude rather than the raw point estimate.
482        // Under the working-Laplace posterior `beta ~ N(beta_hat, Sigma_beta)`,
483        // `Sigma_beta = H^{-1}`, the response `t_k = (D0 beta)_k` has posterior
484        // mean `t_hat_k` and variance `V_k = (D0 Sigma_beta D0^T)_kk >= 0`. The
485        // expected squared response is `E[t_k^2] = t_hat_k^2 + V_k`, so the part
486        // of `t_hat_k^2` that exceeds the noise floor `V_k` is the credibly real
487        // squared magnitude
488        //
489        //   t_k^credible^2 = max( t_hat_k^2 - V_k , 0 ),
490        //   r_k^snr        = sqrt( t_k^credible^2 + eps^2 ),   w_k = 1 / r_k^snr.
491        //
492        // This is the correct realization of the intent. Where the point
493        // estimate is a *credible* edge (t_hat^2 >> V) the credible magnitude is
494        // ~|t_hat| and the weight is essentially `1/|t_hat|` (left un-penalized,
495        // edge preserved). Where the large point-estimate magnitude is *noise*
496        // (t_hat^2 <~ V) the credible magnitude collapses to 0 and the weight
497        // rises to `1/eps` (extra smoothing, noise suppressed). The weight is
498        // monotone non-decreasing in `V`, and is bounded above by `1/eps` — the
499        // *same* ceiling the magnitude-only weight `1/sqrt(t^2 + eps^2)` already
500        // attains at `t = 0` (and clamped by `weight_ceiling`), so it is not an
501        // unbounded blow-up: it only moves the noise-dominated rows to the flat-
502        // response weight they would have had with a credible estimate of zero
503        // curvature. The earlier delta-method form `f + ½ f'' V` was non-monotone
504        // (`f''` flips sign at `2t^2 = eps^2`) and unbounded in `V`, which left
505        // noisy rows under-penalized and was the source of the SNR regression.
506        // With `V == 0` everywhere this degrades exactly to `surrogateweights`
507        // (`1/sqrt(t^2 + eps^2)`), so any covariance-unavailable path is
508        // unchanged.
509        let eps2 = self.epsilon * self.epsilon;
510        let weight = Array1::from_iter(self.signal.iter().zip(variance.iter()).map(|(&t, &v)| {
511            let credible2 = (t * t - v.max(0.0)).max(0.0);
512            let r = (credible2 + eps2).sqrt();
513            (1.0 / r).clamp(weight_floor, weight_ceiling)
514        }));
515        let invweight = weight.mapv(|u| 1.0 / u);
516        (weight, invweight)
517    }
518
519    fn directionalhessian_diag(&self, direction_signal: &Array1<f64>) -> Array1<f64> {
520        // Scalar-image directional third derivative:
521        //
522        // If t(beta) = A beta and
523        //   H(beta) = A^T diag( eps^2 / (t_k(beta)^2 + eps^2)^(3/2) ) A,
524        // then for q = A u,
525        //
526        //   D(H)[u]
527        //   = A^T diag( -3 eps^2 t_k q_k / (t_k^2 + eps^2)^(5/2) ) A.
528        //
529        // This is one of the exact P_{beta,beta,beta}[u] terms needed by the
530        // Laplace hypergradient
531        //
532        //   d/dtheta log det H = tr(H^{-1} Hdot_theta),
533        //   Hdot_theta = J_{beta,beta,theta} + D_beta(H)[beta_theta].
534        let eps2 = self.epsilon * self.epsilon;
535        Array1::from_iter(
536            self.signal
537                .iter()
538                .zip(direction_signal.iter())
539                .zip(self.radius.iter())
540                .map(|((t, q), r)| -3.0 * eps2 * t * q / r.powi(5)),
541        )
542    }
543
544    /// Exact scalar-image fourth derivative contracted along two coefficient
545    /// directions: with `t(β)=Aβ`, `H(β)=Aᵀ diag(ψ''(t_k)) A`,
546    /// `ψ''(t)=ε²/r³`, the second directional derivative of `H` along
547    /// `(u, v)` (signals `q1=A u`, `q2=A v`) is
548    /// `Aᵀ diag( ψ''''(t_k) q1_k q2_k ) A`, with
549    /// `ψ''''(t) = -3 ε² / r⁵ + 15 ε² t² / r⁷`.
550    fn second_directionalhessian_diag(
551        &self,
552        direction1_signal: &Array1<f64>,
553        direction2_signal: &Array1<f64>,
554    ) -> Array1<f64> {
555        let eps2 = self.epsilon * self.epsilon;
556        Array1::from_iter(
557            self.signal
558                .iter()
559                .zip(direction1_signal.iter())
560                .zip(direction2_signal.iter())
561                .zip(self.radius.iter())
562                .map(|(((t, q1), q2), r)| {
563                    let r2 = r * r;
564                    let psi4 = -3.0 * eps2 / r.powi(5) + 15.0 * eps2 * t * t / (r.powi(5) * r2);
565                    psi4 * q1 * q2
566                }),
567        )
568    }
569
570    fn log_epsilon_betahessian_diag(&self) -> Array1<f64> {
571        let eps2 = self.epsilon * self.epsilon;
572        let eps4 = eps2 * eps2;
573        Array1::from_iter(
574            self.signal
575                .iter()
576                .zip(self.radius.iter())
577                .map(|(_, r)| 2.0 * eps2 / r.powi(3) - 3.0 * eps4 / r.powi(5)),
578        )
579    }
580
581    fn log_epsilon_beta_mixed_second_coeff(&self) -> Array1<f64> {
582        let eps2 = self.epsilon * self.epsilon;
583        Array1::from_iter(
584            self.signal
585                .iter()
586                .zip(self.radius.iter())
587                .map(|(t, r)| eps2 * t * (eps2 - 2.0 * t * t) / r.powi(5)),
588        )
589    }
590
591    fn log_epsilon_betahessian_second_diag(&self) -> Array1<f64> {
592        let eps2 = self.epsilon * self.epsilon;
593        let eps4 = eps2 * eps2;
594        let eps6 = eps4 * eps2;
595        Array1::from_iter(
596            self.radius.iter().map(|r| {
597                4.0 * eps2 / r.powi(3) - 18.0 * eps4 / r.powi(5) + 15.0 * eps6 / r.powi(7)
598            }),
599        )
600    }
601
602    fn log_epsilon_betahessian_directional_diag(
603        &self,
604        direction_signal: &Array1<f64>,
605    ) -> Array1<f64> {
606        let eps2 = self.epsilon * self.epsilon;
607        let eps4 = eps2 * eps2;
608        Array1::from_iter(
609            self.signal
610                .iter()
611                .zip(direction_signal.iter())
612                .zip(self.radius.iter())
613                .map(|((t, q), r)| (-6.0 * eps2 * t / r.powi(5) + 15.0 * eps4 * t / r.powi(7)) * q),
614        )
615    }
616}
617
618#[derive(Clone)]
619struct CharbonnierGroupedBlockState {
620    norm: Array1<f64>,
621    radius: Array1<f64>,
622    signal_blocks: Array2<f64>,
623    epsilon: f64,
624}
625
626impl CharbonnierGroupedBlockState {
627    fn from_signal_blocks(signal_blocks: Array2<f64>, epsilon: f64) -> Self {
628        let eps = epsilon.max(1e-12);
629        let norm = Array1::from_iter(
630            signal_blocks
631                .rows()
632                .into_iter()
633                .map(|row| row.iter().map(|v| v * v).sum::<f64>().sqrt()),
634        );
635        let radius = norm.mapv(|g| (g * g + eps * eps).sqrt());
636        Self {
637            norm,
638            radius,
639            signal_blocks,
640            epsilon: eps,
641        }
642    }
643
644    fn penalty_value(&self) -> f64 {
645        self.radius.iter().map(|r| r - self.epsilon).sum::<f64>()
646    }
647
648    fn norm_signal(&self) -> Array1<f64> {
649        self.norm.clone()
650    }
651
652    fn betagradient_blocks(&self) -> Array2<f64> {
653        let mut out = self.signal_blocks.clone();
654        for (k, mut row) in out.rows_mut().into_iter().enumerate() {
655            let scale = 1.0 / self.radius[k];
656            row.mapv_inplace(|v| v * scale);
657        }
658        out
659    }
660
661    fn betahessian_blocks(&self) -> Vec<Array2<f64>> {
662        let mut out = Vec::with_capacity(self.signal_blocks.nrows());
663        for (k, row) in self.signal_blocks.rows().into_iter().enumerate() {
664            let dim = row.len();
665            let mut block = Array2::<f64>::eye(dim);
666            block.mapv_inplace(|v| v / self.radius[k]);
667            for i in 0..dim {
668                for j in 0..dim {
669                    block[[i, j]] -= row[i] * row[j] / self.radius[k].powi(3);
670                }
671            }
672            out.push(block);
673        }
674        out
675    }
676
677    fn log_epsilon_gradient_terms(&self) -> Array1<f64> {
678        let epsilon = self.epsilon;
679        let eps2 = epsilon * epsilon;
680        self.radius.mapv(|r| eps2 / r - epsilon)
681    }
682
683    fn log_epsilon_betagradient_blocks(&self) -> Array2<f64> {
684        let mut out = self.signal_blocks.clone();
685        let eps2 = self.epsilon * self.epsilon;
686        for (k, mut row) in out.rows_mut().into_iter().enumerate() {
687            let scale = -eps2 / self.radius[k].powi(3);
688            row.mapv_inplace(|v| v * scale);
689        }
690        out
691    }
692
693    fn log_epsilon_hessian_terms(&self) -> Array1<f64> {
694        let epsilon = self.epsilon;
695        let eps2 = epsilon * epsilon;
696        let eps4 = eps2 * eps2;
697        self.radius
698            .mapv(|r| 2.0 * eps2 / r - eps4 / r.powi(3) - epsilon)
699    }
700
701    fn surrogateweights_posterior_snr(
702        &self,
703        variance: &Array1<f64>,
704        weight_floor: f64,
705        weight_ceiling: f64,
706    ) -> (Array1<f64>, Array1<f64>) {
707        // Grouped posterior-SNR (credible-magnitude) reweighting.
708        //
709        // The magnitude-only grouped surrogate weight uses the point-estimate
710        // block norm
711        //
712        //   g_k     = ||v_k||_2,   v_k = G_k beta_hat,
713        //   r_k^mag = sqrt( g_k^2 + eps^2 ),
714        //   w_k     = 1 / r_k^mag.
715        //
716        // The posterior covariance of the *block* response v_k = G_k beta under
717        // beta ~ N(beta_hat, Sigma_beta), Sigma_beta = H^{-1}, has total trace
718        //
719        //   Cov(v_k)     = G_k Sigma_beta G_k^T   (a block_dim x block_dim block),
720        //   variance[k]  = tr(Cov(v_k)) = sum_axis ( G_k[axis] Sigma_beta G_k[axis]^T ),
721        //
722        // i.e. the variance aggregated over the axis-block in the same way
723        // `norm` aggregates ||v_k||^2. As for the scalar block, we deflate the
724        // squared block norm by this noise floor to obtain the credible squared
725        // magnitude and shrink poorly-determined responses toward zero:
726        //
727        //   g_k^credible^2 = max( g_k^2 - tr(Cov(v_k)) , 0 ),
728        //   r_k^snr        = sqrt( g_k^credible^2 + eps^2 ),   w_k = 1 / r_k^snr.
729        //
730        // A block whose norm is credibly large (g_k^2 >> tr Cov) keeps a small
731        // weight (real feature, left un-penalized); a block whose norm is
732        // dominated by posterior variance has its credible norm collapse to 0,
733        // raising the weight to `1/eps` (noise suppressed). The weight is
734        // monotone non-decreasing in `tr Cov` and bounded above by `1/eps` — the
735        // same ceiling the magnitude-only weight already attains at `g = 0`
736        // (and clamped by `weight_ceiling`), so it is not an unbounded blow-up.
737        //
738        // This evaluates the grouped MM weight `f(v) = (||v||^2 + eps^2)^{-1/2}`
739        // at the credible block norm rather than at the raw point estimate. The
740        // expected squared block norm under `v_k ~ N(v_hat_k, C_k)` is
741        // `E[||v_k||^2] = ||v_hat_k||^2 + tr(C_k)`, so the credibly-real squared
742        // norm is `max(g_k^2 - tr(C_k), 0)`, identical in form to the scalar
743        // path (`block_dim == 1` recovers it exactly). The earlier delta-method
744        // correction `½ Σ ∂²f · C` was non-monotone (its sign flips with the
745        // Hessian of `f`) and unbounded in `tr C`, which under-penalized noisy
746        // blocks and was the source of the SNR regression. With `tr C == 0` it
747        // recovers `1/sqrt(g^2 + eps^2)`.
748        let eps2 = self.epsilon * self.epsilon;
749        let weight = Array1::from_iter(self.norm.iter().zip(variance.iter()).map(|(&g, &v)| {
750            let credible2 = (g * g - v.max(0.0)).max(0.0);
751            let r = (credible2 + eps2).sqrt();
752            (1.0 / r).clamp(weight_floor, weight_ceiling)
753        }));
754        let invweight = weight.mapv(|u| 1.0 / u);
755        (weight, invweight)
756    }
757
758    fn directionalhessian_blocks(&self, direction_blocks: &Array2<f64>) -> Vec<Array2<f64>> {
759        // Exact grouped directional third derivative for the slope penalty.
760        //
761        // For each collocation block k:
762        //   v_k = G_k beta,
763        //   q_k = G_k u,
764        //   r_k = sqrt(||v_k||^2 + eps^2),
765        //
766        // the exact Hessian block for psi(g; eps) = sqrt(g^2 + eps^2) - eps is
767        //   B_k,
768        //   B_k = (1 / r_k) I - v_k v_k^T / r_k^3.
769        //
770        // Differentiating B_k along u gives
771        //   M_k(u)
772        //   = -(v_k^T q_k / r_k^3) I
773        //     - (q_k v_k^T + v_k q_k^T) / r_k^3
774        //     + 3 (v_k^T q_k) v_k v_k^T / r_k^5.
775        //
776        // This expression must be symmetric because it is the directional
777        // derivative of the symmetric matrix
778        //
779        //   B_k = (1 / r_k) I - v_k v_k^T / r_k^3.
780        //
781        // The full directional penalty Hessian map is then
782        //   D(H_g)[u] = lambda_g * sum_k G_k^T M_k(u) G_k.
783        let mut out = Vec::with_capacity(self.signal_blocks.nrows());
784        for (k, (v, q)) in self
785            .signal_blocks
786            .rows()
787            .into_iter()
788            .zip(direction_blocks.rows().into_iter())
789            .enumerate()
790        {
791            let dim = v.len();
792            let dot = v.iter().zip(q.iter()).map(|(a, b)| a * b).sum::<f64>();
793            let r3 = self.radius[k].powi(3);
794            let r5 = self.radius[k].powi(5);
795            let mut block = Array2::<f64>::eye(dim);
796            block.mapv_inplace(|x| -dot * x / r3);
797            for i in 0..dim {
798                for j in 0..dim {
799                    block[[i, j]] -= (q[i] * v[j] + v[i] * q[j]) / r3;
800                    block[[i, j]] += 3.0 * dot * v[i] * v[j] / r5;
801                }
802            }
803            out.push(block);
804        }
805        out
806    }
807
808    /// Exact grouped second directional derivative of the slope/curvature block
809    /// Hessian `B_k = (1/r_k) I − v_k v_kᵀ / r_k³` along two coefficient
810    /// directions, with per-block signal images `a_k = G_k u`, `b_k = G_k w`.
811    ///
812    /// `B_k`'s first directional derivative along `a` is
813    ///   `M_k(a) = −(v·a/r³) I − (a vᵀ + v aᵀ)/r³ + 3 (v·a) v vᵀ/r⁵`
814    /// (see `directionalhessian_blocks`). Differentiating `M_k(a)` once more
815    /// along `b` (i.e. `v ← v + t b`) gives the symmetric block
816    ///   `N_k(a,b) = (−a·b/r³ + 3 (v·a)(v·b)/r⁵) I`
817    ///            `  − (a bᵀ + b aᵀ)/r³`
818    ///            `  + 3 (v·b)(a vᵀ + v aᵀ)/r⁵`
819    ///            `  + 3 (a·b) v vᵀ/r⁵`
820    ///            `  + 3 (v·a)(b vᵀ + v bᵀ)/r⁵`
821    ///            `  − 15 (v·a)(v·b) v vᵀ/r⁷`,
822    /// so `D²_β H_g[u,w] = λ_g Σ_k G_kᵀ N_k(a_k,b_k) G_k`. `N_k` is symmetric in
823    /// `a ↔ b`, matching `D²H[u,w] = D²H[w,u]`.
824    fn second_directionalhessian_blocks(
825        &self,
826        direction1_blocks: &Array2<f64>,
827        direction2_blocks: &Array2<f64>,
828    ) -> Vec<Array2<f64>> {
829        let mut out = Vec::with_capacity(self.signal_blocks.nrows());
830        for ((k, v), (a, b)) in self.signal_blocks.rows().into_iter().enumerate().zip(
831            direction1_blocks
832                .rows()
833                .into_iter()
834                .zip(direction2_blocks.rows().into_iter()),
835        ) {
836            let dim = v.len();
837            let dot = |x: ndarray::ArrayView1<'_, f64>, y: ndarray::ArrayView1<'_, f64>| {
838                x.iter().zip(y.iter()).map(|(p, q)| p * q).sum::<f64>()
839            };
840            let sa = dot(v, a);
841            let sb = dot(v, b);
842            let ab = dot(a, b);
843            let r = self.radius[k];
844            let r3 = r.powi(3);
845            let r5 = r.powi(5);
846            let r7 = r5 * r * r;
847            let diag = -ab / r3 + 3.0 * sa * sb / r5;
848            let mut block = Array2::<f64>::eye(dim);
849            block.mapv_inplace(|x| diag * x);
850            for i in 0..dim {
851                for j in 0..dim {
852                    block[[i, j]] -= (a[i] * b[j] + b[i] * a[j]) / r3;
853                    block[[i, j]] += 3.0 * sb * (a[i] * v[j] + v[i] * a[j]) / r5;
854                    block[[i, j]] += 3.0 * ab * v[i] * v[j] / r5;
855                    block[[i, j]] += 3.0 * sa * (b[i] * v[j] + v[i] * b[j]) / r5;
856                    block[[i, j]] -= 15.0 * sa * sb * v[i] * v[j] / r7;
857                }
858            }
859            out.push(block);
860        }
861        out
862    }
863
864    fn log_epsilon_betahessian_blocks(&self) -> Vec<Array2<f64>> {
865        let mut out = Vec::with_capacity(self.signal_blocks.nrows());
866        for (k, row) in self.signal_blocks.rows().into_iter().enumerate() {
867            let dim = row.len();
868            let r3 = self.radius[k].powi(3);
869            let r5 = self.radius[k].powi(5);
870            let mut block = Array2::<f64>::eye(dim);
871            let eps2 = self.epsilon * self.epsilon;
872            block.mapv_inplace(|v| -eps2 * v / r3);
873            for i in 0..dim {
874                for j in 0..dim {
875                    block[[i, j]] += 3.0 * eps2 * row[i] * row[j] / r5;
876                }
877            }
878            out.push(block);
879        }
880        out
881    }
882
883    fn log_epsilon_beta_mixed_second_blocks(&self) -> Array2<f64> {
884        let mut out = self.signal_blocks.clone();
885        let eps2 = self.epsilon * self.epsilon;
886        for (k, mut row) in out.rows_mut().into_iter().enumerate() {
887            let norm2 = self.norm[k] * self.norm[k];
888            let scale = eps2 * (eps2 - 2.0 * norm2) / self.radius[k].powi(5);
889            row.mapv_inplace(|v| v * scale);
890        }
891        out
892    }
893
894    fn log_epsilon_betahessian_second_blocks(&self) -> Vec<Array2<f64>> {
895        let mut out = Vec::with_capacity(self.signal_blocks.nrows());
896        let eps2 = self.epsilon * self.epsilon;
897        for (k, row) in self.signal_blocks.rows().into_iter().enumerate() {
898            let dim = row.len();
899            let norm2 = self.norm[k] * self.norm[k];
900            let r5 = self.radius[k].powi(5);
901            let r7 = self.radius[k].powi(7);
902            let mut block = Array2::<f64>::eye(dim);
903            block.mapv_inplace(|v| eps2 * (eps2 - 2.0 * norm2) * v / r5);
904            for i in 0..dim {
905                for j in 0..dim {
906                    block[[i, j]] += 3.0 * eps2 * (2.0 * norm2 - 3.0 * eps2) * row[i] * row[j] / r7;
907                }
908            }
909            out.push(block);
910        }
911        out
912    }
913
914    fn log_epsilon_betahessian_directional_blocks(
915        &self,
916        direction_blocks: &Array2<f64>,
917    ) -> Vec<Array2<f64>> {
918        let mut out = Vec::with_capacity(self.signal_blocks.nrows());
919        let eps2 = self.epsilon * self.epsilon;
920        for (k, (v, q)) in self
921            .signal_blocks
922            .rows()
923            .into_iter()
924            .zip(direction_blocks.rows().into_iter())
925            .enumerate()
926        {
927            let dim = v.len();
928            let dot = v.iter().zip(q.iter()).map(|(a, b)| a * b).sum::<f64>();
929            let r5 = self.radius[k].powi(5);
930            let r7 = self.radius[k].powi(7);
931            let mut block = Array2::<f64>::eye(dim);
932            block.mapv_inplace(|x| 3.0 * eps2 * dot * x / r5);
933            for i in 0..dim {
934                for j in 0..dim {
935                    block[[i, j]] += 3.0 * eps2 * (q[i] * v[j] + v[i] * q[j]) / r5;
936                    block[[i, j]] -= 15.0 * eps2 * dot * v[i] * v[j] / r7;
937                }
938            }
939            out.push(block);
940        }
941        out
942    }
943}
944
945fn scalar_operatorgradient(operator: &Array2<f64>, coeff: &Array1<f64>) -> Array1<f64> {
946    operator.t().dot(coeff)
947}
948
949fn scalar_operatorhessian(operator: &Array2<f64>, diag: &Array1<f64>) -> Array2<f64> {
950    let mut weighted = operator.clone();
951    for (k, &w) in diag.iter().enumerate() {
952        weighted.row_mut(k).mapv_inplace(|v| v * w);
953    }
954    let gram = operator.t().dot(&weighted);
955    (&gram + &gram.t().to_owned()) * 0.5
956}
957
958fn grouped_operatorgradient(
959    d1: &Array2<f64>,
960    dimension: usize,
961    blocks: &Array2<f64>,
962) -> Result<Array1<f64>, EstimationError> {
963    if blocks.ncols() != dimension {
964        crate::bail_invalid_estim!(
965            "grouped gradient block dimension mismatch: got {}, expected {dimension}",
966            blocks.ncols()
967        );
968    }
969    if d1.nrows() != blocks.nrows() * dimension {
970        crate::bail_invalid_estim!(
971            "grouped gradient row mismatch: D1 has {} rows, blocks imply {}",
972            d1.nrows(),
973            blocks.nrows() * dimension
974        );
975    }
976    let mut out = Array1::<f64>::zeros(d1.ncols());
977    for k in 0..blocks.nrows() {
978        let gk = d1
979            .slice(s![k * dimension..(k + 1) * dimension, ..])
980            .to_owned();
981        out += &gk.t().dot(&blocks.row(k));
982    }
983    Ok(out)
984}
985
986fn grouped_operatorhessian(
987    d1: &Array2<f64>,
988    dimension: usize,
989    blocks: &[Array2<f64>],
990) -> Result<Array2<f64>, EstimationError> {
991    if d1.nrows() != blocks.len() * dimension {
992        crate::bail_invalid_estim!(
993            "grouped Hessian row mismatch: D1 has {} rows, blocks imply {}",
994            d1.nrows(),
995            blocks.len() * dimension
996        );
997    }
998    let p = d1.ncols();
999    let mut out = Array2::<f64>::zeros((p, p));
1000    for (k, block) in blocks.iter().enumerate() {
1001        if block.nrows() != dimension || block.ncols() != dimension {
1002            crate::bail_invalid_estim!(
1003                "grouped Hessian block {k} has shape {}x{}, expected {}x{}",
1004                block.nrows(),
1005                block.ncols(),
1006                dimension,
1007                dimension
1008            );
1009        }
1010        let gk = d1
1011            .slice(s![k * dimension..(k + 1) * dimension, ..])
1012            .to_owned();
1013        out += &gk.t().dot(&block.dot(&gk));
1014    }
1015    Ok((&out + &out.t().to_owned()) * 0.5)
1016}
1017
1018#[derive(Clone)]
1019struct SpatialPenaltyExactState {
1020    magnitude: CharbonnierScalarBlockState,
1021    gradient: CharbonnierGroupedBlockState,
1022    curvature: CharbonnierGroupedBlockState,
1023}
1024
1025fn collocationgradient_blocks(
1026    gradrows: &Array1<f64>,
1027    dimension: usize,
1028) -> Result<Array2<f64>, EstimationError> {
1029    if dimension == 0 || !gradrows.len().is_multiple_of(dimension) {
1030        crate::bail_invalid_estim!(
1031            "invalid collocation gradient layout: rows={}, dimension={dimension}",
1032            gradrows.len()
1033        );
1034    }
1035    let p = gradrows.len() / dimension;
1036    let mut out = Array2::<f64>::zeros((p, dimension));
1037    for k in 0..p {
1038        for axis in 0..dimension {
1039            out[[k, axis]] = gradrows[k * dimension + axis];
1040        }
1041    }
1042    Ok(out)
1043}
1044
1045fn collocationhessian_blocks(
1046    hessianrows: &Array1<f64>,
1047    dimension: usize,
1048) -> Result<Array2<f64>, EstimationError> {
1049    let block_dim = dimension.checked_mul(dimension).ok_or_else(|| {
1050        EstimationError::InvalidInput("invalid collocation Hessian dimension overflow".to_string())
1051    })?;
1052    if block_dim == 0 || !hessianrows.len().is_multiple_of(block_dim) {
1053        crate::bail_invalid_estim!(
1054            "invalid collocation Hessian layout: rows={}, dimension={dimension}",
1055            hessianrows.len()
1056        );
1057    }
1058    let p = hessianrows.len() / block_dim;
1059    let mut out = Array2::<f64>::zeros((p, block_dim));
1060    for k in 0..p {
1061        for idx in 0..block_dim {
1062            out[[k, idx]] = hessianrows[k * block_dim + idx];
1063        }
1064    }
1065    Ok(out)
1066}
1067
1068impl SpatialPenaltyExactState {
1069    fn from_beta_local(
1070        beta_local: ArrayView1<'_, f64>,
1071        cache: &SpatialOperatorRuntimeCache,
1072        epsilons: [f64; 3],
1073    ) -> Result<Self, EstimationError> {
1074        // Exact collocation-state extraction for the three Charbonnier penalty blocks.
1075        //
1076        // For one spatial smooth term with coefficient vector beta_local, the exact
1077        // operator-decomposition penalty is built from three collocation images:
1078        //
1079        //   magnitude:  f = D0 beta_local
1080        //   slope:      v_k = G_k beta_local
1081        //   curvature:  H_k = D2_k beta_local
1082        //
1083        // where the gradient operator is stored in row-stacked form:
1084        //
1085        //   D1 beta_local in R^(P * d),
1086        //   row layout = (point 0, axis 0..d-1), (point 1, axis 0..d-1), ...
1087        //   D2 beta_local in R^(P * d * d),
1088        //   row layout = (point, Hessian axis_a, Hessian axis_b).
1089        //
1090        // so we first reshape that stacked vector into the grouped block array
1091        //
1092        //   [v_0^T
1093        //    ...
1094        //    v_(P-1)^T]  in R^(P x d).
1095        //
1096        // The three exact Charbonnier block states then carry:
1097        //   - the raw operator signals,
1098        //   - their radii sqrt(signal^2 + eps^2) or sqrt(||v_k||^2 + eps^2),
1099        //   - and all exact derivatives derived from those radii.
1100        //
1101        // This is the canonical translation from coefficient-space beta to the
1102        // penalty-side mathematical objects used throughout the implementation.
1103        let gradientrows = cache.d1.dot(&beta_local);
1104        let hessianrows = cache.d2.dot(&beta_local);
1105        Ok(Self {
1106            magnitude: CharbonnierScalarBlockState::from_signal(
1107                cache.d0.dot(&beta_local),
1108                epsilons[0],
1109            ),
1110            gradient: CharbonnierGroupedBlockState::from_signal_blocks(
1111                collocationgradient_blocks(&gradientrows, cache.dimension)?,
1112                epsilons[1],
1113            ),
1114            curvature: CharbonnierGroupedBlockState::from_signal_blocks(
1115                collocationhessian_blocks(&hessianrows, cache.dimension)?,
1116                epsilons[2],
1117            ),
1118        })
1119    }
1120
1121    fn absolute_collocation_magnitudes(&self) -> (Array1<f64>, Array1<f64>, Array1<f64>) {
1122        (
1123            self.magnitude.absolute_signal(),
1124            self.gradient.norm_signal(),
1125            self.curvature.norm_signal(),
1126        )
1127    }
1128}
1129
1130fn robust_epsilon_from_samples(values: &[f64], min_epsilon_cfg: f64) -> f64 {
1131    if values.is_empty() {
1132        return min_epsilon_cfg.max(1e-12);
1133    }
1134    let mut clean = values
1135        .iter()
1136        .copied()
1137        .filter(|v| v.is_finite() && *v >= 0.0)
1138        .collect::<Vec<_>>();
1139    if clean.is_empty() {
1140        return min_epsilon_cfg.max(1e-12);
1141    }
1142    clean.sort_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal));
1143
1144    let n = clean.len();
1145    let median = quantile_from_sorted(&clean, 0.5);
1146    let q75 = quantile_from_sorted(&clean, 0.75);
1147    let q95 = quantile_from_sorted(&clean, 0.95);
1148
1149    let mut abs_dev = clean
1150        .iter()
1151        .map(|v| (v - median).abs())
1152        .filter(|v| v.is_finite())
1153        .collect::<Vec<_>>();
1154    abs_dev.sort_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal));
1155    let mad = 1.4826 * quantile_from_sorted(&abs_dev, 0.5);
1156
1157    // Charbonnier/MM requires eps bounded away from zero:
1158    //   u(t0) = 1 / (2*sqrt(t0^2 + eps^2)) ~ 1/(2*eps) near t0=0.
1159    // Use robust pilot scale:
1160    //   s = max(median(z), 1.4826*MAD(z), Q75(z)).
1161    // If s is tiny (<= delta), fallback to:
1162    //   s <- max(Q95(z), RMS(z)).
1163    // If still tiny, fallback to absolute floor s_min.
1164    // Then eps = kappa * s.
1165    // Primary robust scale: s = max(median, 1.4826*MAD, Q75).
1166    let mut scale = median.max(mad).max(q75);
1167
1168    // Safety threshold delta and absolute floor s_min.
1169    let delta = (f64::EPSILON.sqrt() * q95.max(1.0))
1170        .max(min_epsilon_cfg)
1171        .max(1e-12);
1172    let s_min = min_epsilon_cfg.max(1e-12);
1173
1174    // If robust scale is tiny, use high-quantile / RMS fallback.
1175    if scale <= delta {
1176        let rms = (clean.iter().map(|v| v * v).sum::<f64>() / n as f64).sqrt();
1177        scale = q95.max(rms);
1178    }
1179    if scale <= delta {
1180        scale = s_min;
1181    }
1182
1183    // Start near the observed operator scale so the optimizer begins in a
1184    // neutral regime where both quadratic and linear behavior are reachable.
1185    let kappa = 1.0_f64;
1186    (kappa * scale).max(s_min)
1187}
1188
1189fn extract_spatial_operator_runtime_caches(
1190    spec: &TermCollectionSpec,
1191    design: &TermCollectionDesign,
1192) -> Result<Vec<SpatialOperatorRuntimeCache>, EstimationError> {
1193    let smooth_start = design
1194        .design
1195        .ncols()
1196        .saturating_sub(design.smooth.total_smooth_cols());
1197    let mut out = Vec::<SpatialOperatorRuntimeCache>::new();
1198    for (term_idx, (termspec, term_fit)) in spec
1199        .smooth_terms
1200        .iter()
1201        .zip(design.smooth.terms.iter())
1202        .enumerate()
1203    {
1204        let Some(global_base_idx) = smooth_term_penalty_index(spec, design, term_idx) else {
1205            continue;
1206        };
1207        let mut active_local_idx = 0usize;
1208        let mut mass_local_idx = None;
1209        let mut tension_local_idx = None;
1210        let mut stiffness_local_idx = None;
1211        let mut mass_norm = None;
1212        let mut tension_norm = None;
1213        let mut stiffness_norm = None;
1214        for info in &term_fit.penaltyinfo_local {
1215            if !info.active {
1216                continue;
1217            }
1218            match info.source {
1219                PenaltySource::OperatorMass => {
1220                    mass_local_idx = Some(active_local_idx);
1221                    mass_norm = Some(info.normalization_scale);
1222                }
1223                PenaltySource::OperatorTension => {
1224                    tension_local_idx = Some(active_local_idx);
1225                    tension_norm = Some(info.normalization_scale);
1226                }
1227                PenaltySource::OperatorStiffness => {
1228                    stiffness_local_idx = Some(active_local_idx);
1229                    stiffness_norm = Some(info.normalization_scale);
1230                }
1231                _ => {}
1232            }
1233            active_local_idx += 1;
1234        }
1235        // The Charbonnier adaptive overlay rebuilds the {mass, tension,
1236        // stiffness} D-operator triplet from explicit collocation derivatives
1237        // and reweights all three channels in tandem; the stiffness slot in
1238        // particular is the D2 second-derivative operator. A term that does
1239        // NOT ship an explicit Stiffness penalty (pure Duchon's RKHS-Primary-
1240        // curvature layout — `DuchonOperatorPenaltySpec::default()`) has no
1241        // matching shipped penalty for the Charbonnier D2 surrogate to reweight,
1242        // so applying the overlay would smuggle a fresh D2 collocation
1243        // operator into a basis whose curvature is the RKHS Primary Gram (a
1244        // different mathematical object). Without an explicit Stiffness
1245        // channel the term must be skipped — the runtime cache for the
1246        // adaptive overlay simply doesn't apply.
1247        let (
1248            Some(mass_local),
1249            Some(tension_local),
1250            Some(stiffness_local),
1251            Some(mass_scale),
1252            Some(tension_scale),
1253            Some(stiffness_scale),
1254        ) = (
1255            mass_local_idx,
1256            tension_local_idx,
1257            stiffness_local_idx,
1258            mass_norm,
1259            tension_norm,
1260            stiffness_norm,
1261        )
1262        else {
1263            continue;
1264        };
1265        let mass_global_idx = global_base_idx + mass_local;
1266        let tension_global_idx = global_base_idx + tension_local;
1267        let stiffness_global_idx = global_base_idx + stiffness_local;
1268
1269        let (feature_cols, mut d0, mut d1, mut d2, collocation_points, dim, center_mass_rows) =
1270            match (&termspec.basis, &term_fit.metadata) {
1271                (
1272                    SmoothBasisSpec::Matern { feature_cols, .. },
1273                    BasisMetadata::Matern {
1274                        centers,
1275                        length_scale,
1276                        nu,
1277                        include_intercept,
1278                        identifiability_transform,
1279                        aniso_log_scales,
1280                        input_scales,
1281                        ..
1282                    },
1283                ) => {
1284                    // Match the σ_geom-compensated effective length scale the
1285                    // design (and shipped penalties) use against the standardized
1286                    // centers; the raw metadata length_scale lives in original
1287                    // coordinates and would put this overlay on a different kernel
1288                    // range than the penalties it scales (#706).
1289                    let collocation_length_scale = match input_scales.as_deref() {
1290                        Some(scales) => {
1291                            compensate_length_scale_for_standardization(*length_scale, scales)
1292                        }
1293                        None => *length_scale,
1294                    };
1295                    let ops = build_matern_collocation_operator_matrices(
1296                        centers.view(),
1297                        None,
1298                        collocation_length_scale,
1299                        *nu,
1300                        *include_intercept,
1301                        identifiability_transform.as_ref().map(|z| z.view()),
1302                        aniso_log_scales.as_deref(),
1303                    )?;
1304                    (
1305                        feature_cols.clone(),
1306                        ops.d0,
1307                        ops.d1,
1308                        ops.d2,
1309                        ops.collocation_points,
1310                        centers.ncols(),
1311                        false,
1312                    )
1313                }
1314                (
1315                    SmoothBasisSpec::Duchon { feature_cols, .. },
1316                    BasisMetadata::Duchon {
1317                        centers,
1318                        length_scale,
1319                        power,
1320                        nullspace_order,
1321                        identifiability_transform,
1322                        input_scales,
1323                        aniso_log_scales,
1324                        operator_collocation_points: Some(collocation_points),
1325                        ..
1326                    },
1327                ) => {
1328                    let collocation_length_scale = match (length_scale, input_scales.as_deref()) {
1329                        (Some(ls), Some(scales)) => {
1330                            Some(compensate_length_scale_for_standardization(*ls, scales))
1331                        }
1332                        (Some(ls), None) => Some(*ls),
1333                        (None, _) => None,
1334                    };
1335                    let ops =
1336                        gam_terms::basis::build_duchon_collocation_operator_matriceswithworkspace(
1337                            centers.view(),
1338                            collocation_points.view(),
1339                            None,
1340                            collocation_length_scale,
1341                            *power,
1342                            *nullspace_order,
1343                            aniso_log_scales.as_deref(),
1344                            identifiability_transform.as_ref().map(|z| z.view()),
1345                            2,
1346                            None,
1347                            &mut BasisWorkspace::default(),
1348                        )?;
1349                    (
1350                        feature_cols.clone(),
1351                        ops.d0,
1352                        ops.d1,
1353                        ops.d2,
1354                        ops.collocation_points,
1355                        centers.ncols(),
1356                        true,
1357                    )
1358                }
1359                _ => continue,
1360            };
1361        if center_mass_rows && d0.nrows() > 0 && d0.ncols() > 0 {
1362            let means = d0.sum_axis(Axis(0)).mapv(|v| v / d0.nrows() as f64);
1363            for mut row in d0.rows_mut() {
1364                row -= &means;
1365            }
1366        }
1367
1368        // Runtime operator caches must live on the same normalized penalty scale as the
1369        // shipped design penalties. The basis builders normalize S0=D0'D0, S1=D1'D1, and
1370        // S2=D2'D2 before exposing them as smoothing blocks, recording the corresponding
1371        // Frobenius norms in penaltyinfo_local.normalization_scale. If the exact adaptive
1372        // path uses raw collocation operators here, then its Charbonnier penalties live on a
1373        // different geometry from the ordinary Matérn/Duchon penalties:
1374        //
1375        //   raw quadratic limit:        beta' (D'D) beta
1376        //   shipped design penalty:     beta' (D'D / c) beta
1377        //
1378        // The correct operator-level normalization is therefore
1379        //
1380        //   D_norm = D / sqrt(c),
1381        //
1382        // so that D_norm' D_norm = (D'D)/c matches the design penalty exactly. Without this,
1383        // adaptive lambdas compensate for hidden operator-scale mismatches and are no longer
1384        // comparable to the baseline smoothing parameters.
1385        let mass_scale = mass_scale.max(1e-12).sqrt();
1386        let tension_scale = tension_scale.max(1e-12).sqrt();
1387        let stiffness_scale = stiffness_scale.max(1e-12).sqrt();
1388        d0.mapv_inplace(|v| v / mass_scale);
1389        d1.mapv_inplace(|v| v / tension_scale);
1390        d2.mapv_inplace(|v| v / stiffness_scale);
1391
1392        let coeff_global_range =
1393            (smooth_start + term_fit.coeff_range.start)..(smooth_start + term_fit.coeff_range.end);
1394        if d0.ncols() != coeff_global_range.len()
1395            || d1.ncols() != coeff_global_range.len()
1396            || d2.ncols() != coeff_global_range.len()
1397        {
1398            crate::bail_invalid_estim!(
1399                "spatial operator dimension mismatch for term '{}': D0 cols={}, D1 cols={}, D2 cols={}, coeffs={}",
1400                term_fit.name,
1401                d0.ncols(),
1402                d1.ncols(),
1403                d2.ncols(),
1404                coeff_global_range.len()
1405            );
1406        }
1407        out.push(SpatialOperatorRuntimeCache {
1408            termname: term_fit.name.clone(),
1409            feature_cols,
1410            coeff_global_range,
1411            mass_penalty_global_idx: mass_global_idx,
1412            tension_penalty_global_idx: tension_global_idx,
1413            stiffness_penalty_global_idx: stiffness_global_idx,
1414            d0,
1415            d1,
1416            d2,
1417            collocation_points,
1418            dimension: dim,
1419        });
1420    }
1421    Ok(out)
1422}
1423
1424/// Posterior variance of a scalar collocation operator response under the
1425/// working-Laplace posterior `beta ~ N(beta_hat, Sigma_local)`.
1426///
1427/// For operator row `D_k` (one row of `D0`) acting on the term-local coefficient
1428/// block, `Var((D beta)_k) = D_k Sigma_local D_k^T = (D Sigma_local D^T)_kk`.
1429/// We compute it without forming `D Sigma D^T` densely: for each row we evaluate
1430/// `s_k = Sigma_local D_k^T` (one matrix-vector product) and then `D_k . s_k`.
1431/// `Sigma_local` is the sub-block of the global conditional covariance
1432/// `Sigma_beta = H^{-1}` indexed by the term's `coeff_global_range`, i.e. the
1433/// covariance proxy is the already-materialized inner working-Laplace inverse;
1434/// no second factorization is formed.
1435fn scalar_operator_response_variance(
1436    operator: &Array2<f64>,
1437    cov_local: &Array2<f64>,
1438) -> Array1<f64> {
1439    Array1::from_iter(operator.rows().into_iter().map(|row| {
1440        let s = cov_local.dot(&row);
1441        row.dot(&s).max(0.0)
1442    }))
1443}
1444
1445/// Posterior second-moment variance aggregated over each grouped collocation
1446/// block (gradient/curvature). The grouped operator is stored row-stacked with
1447/// `block_dim` rows per collocation point (`d` axes for the gradient, `d*d` for
1448/// the Hessian). For block `k`,
1449///
1450///   v_k = G_k beta,   Cov(v_k) = G_k Sigma_local G_k^T   (block_dim x block_dim),
1451///   variance_k = tr(Cov(v_k)) = sum_axis ( G_k[axis] Sigma_local G_k[axis]^T ),
1452///
1453/// which matches how `CharbonnierGroupedBlockState::norm` aggregates
1454/// `||v_k||^2 = sum_axis (G_k[axis] beta)^2` across the axis-block.
1455fn grouped_operator_response_variance(
1456    operator: &Array2<f64>,
1457    block_dim: usize,
1458    cov_local: &Array2<f64>,
1459) -> Result<Array1<f64>, EstimationError> {
1460    if block_dim == 0 || !operator.nrows().is_multiple_of(block_dim) {
1461        crate::bail_invalid_estim!(
1462            "grouped variance row layout invalid: rows={}, block_dim={block_dim}",
1463            operator.nrows()
1464        );
1465    }
1466    let p = operator.nrows() / block_dim;
1467    let mut out = Array1::<f64>::zeros(p);
1468    for k in 0..p {
1469        let mut acc = 0.0;
1470        for axis in 0..block_dim {
1471            let row = operator.row(k * block_dim + axis);
1472            let s = cov_local.dot(&row);
1473            acc += row.dot(&s);
1474        }
1475        out[k] = acc.max(0.0);
1476    }
1477    Ok(out)
1478}
1479
1480fn compute_spatial_adaptiveweights_for_beta(
1481    beta: &Array1<f64>,
1482    caches: &[SpatialOperatorRuntimeCache],
1483    epsilon_0: f64,
1484    epsilon_g: f64,
1485    epsilon_c: f64,
1486    weight_floor: f64,
1487    weight_ceiling: f64,
1488    beta_covariance: Option<&Array2<f64>>,
1489) -> Result<Vec<SpatialAdaptiveWeights>, EstimationError> {
1490    // Charbonnier / pseudo-Huber MM derivation (per collocation scalar t):
1491    //   psi(t; eps) = sqrt(t^2 + eps^2) - eps
1492    // and for reference t0 the tangent majorizer in t^2 gives:
1493    //   psi(t) <= 0.5 * w(t0) * t^2 + const(t0),
1494    //   w(t0) = 1 / sqrt(t0^2 + eps^2).
1495    //
1496    // We apply this to:
1497    //   t = f_k = |f(z_k)|             (magnitude),
1498    //   t = g_k = ||nabla f(z_k)||_2   (gradient magnitude),
1499    //   t = c_k = ||D²f(z_k)||_F       (full Hessian curvature),
1500    // both computed from beta^(t-1).
1501    //
1502    // These w values define the quadratic surrogate penalties:
1503    //   K0 = D0_con^T W_0 D0_con,  W_0 = diag(w_0)
1504    //   K1 = D1_con^T W_g D1_con,  W_g = diag(w_g) \otimes I_d  (k,axis order)
1505    //   K2 = D2_con^T W_c D2_con,  W_c = diag(w_c) \otimes I_(d*d).
1506    //
1507    // We clamp w directly, then derive inv_w=1/w for diagnostics and row scaling.
1508    //
1509    // Posterior-SNR reweighting (magic by default): when the inner working-Laplace
1510    // conditional covariance `Sigma_beta = H^{-1}` is available we replace the
1511    // squared point-estimate radius `t_k^2 + eps^2` by the credible (noise-floor-
1512    // corrected) second moment `max(t_k^2 - Var((D beta)_k), 0) + eps^2`, with
1513    // `Var = (D Sigma_beta D^T)_kk`. This stops the weight from leaving derivatives
1514    // un-penalized just because they are large but poorly determined: such
1515    // responses are shrunk toward zero (large weight, strong smoothing), while
1516    // credibly large derivatives (real edges) keep their small weight. `Sigma_beta`
1517    // here is the already-formed inner Hessian inverse from the final exact-family
1518    // solve — no second factorization is built; we only reuse the materialized
1519    // covariance. When the covariance is unavailable (`None`) the variance is zero
1520    // and this degrades *exactly* to the old magnitude-only radius.
1521    caches
1522        .iter()
1523        .map(|cache| {
1524            let beta_local = beta.slice(s![cache.coeff_global_range.clone()]);
1525            let exact = SpatialPenaltyExactState::from_beta_local(
1526                beta_local,
1527                cache,
1528                [epsilon_0, epsilon_g, epsilon_c],
1529            )?;
1530            let cov_local = beta_covariance.map(|cov| {
1531                cov.slice(s![
1532                    cache.coeff_global_range.clone(),
1533                    cache.coeff_global_range.clone()
1534                ])
1535                .to_owned()
1536            });
1537            let dim = cache.dimension;
1538            let (var_0, var_g, var_c) = match cov_local.as_ref() {
1539                Some(cov) => (
1540                    scalar_operator_response_variance(&cache.d0, cov),
1541                    grouped_operator_response_variance(&cache.d1, dim, cov)?,
1542                    grouped_operator_response_variance(&cache.d2, dim * dim, cov)?,
1543                ),
1544                None => (
1545                    Array1::<f64>::zeros(exact.magnitude.signal.len()),
1546                    Array1::<f64>::zeros(exact.gradient.norm.len()),
1547                    Array1::<f64>::zeros(exact.curvature.norm.len()),
1548                ),
1549            };
1550            let (_, inv_0) = exact.magnitude.surrogateweights_posterior_snr(
1551                &var_0,
1552                weight_floor,
1553                weight_ceiling,
1554            );
1555            let (_, inv_g) =
1556                exact
1557                    .gradient
1558                    .surrogateweights_posterior_snr(&var_g, weight_floor, weight_ceiling);
1559            let (_, inv_c) = exact.curvature.surrogateweights_posterior_snr(
1560                &var_c,
1561                weight_floor,
1562                weight_ceiling,
1563            );
1564            Ok(SpatialAdaptiveWeights {
1565                inv_magweight: inv_0,
1566                invgradweight: inv_g,
1567                inv_lapweight: inv_c,
1568            })
1569        })
1570        .collect()
1571}
1572
1573fn compute_initial_epsilons(
1574    beta: &Array1<f64>,
1575    caches: &[SpatialOperatorRuntimeCache],
1576    min_epsilon: f64,
1577) -> Result<(f64, f64, f64), EstimationError> {
1578    let mut fvals = Vec::<f64>::new();
1579    let mut gvals = Vec::<f64>::new();
1580    let mut cvals = Vec::<f64>::new();
1581    for cache in caches {
1582        let beta_local = beta.slice(s![cache.coeff_global_range.clone()]);
1583        let exact = SpatialPenaltyExactState::from_beta_local(
1584            beta_local,
1585            cache,
1586            [min_epsilon, min_epsilon, min_epsilon],
1587        )?;
1588        let (f, g, c) = exact.absolute_collocation_magnitudes();
1589        fvals.extend(f.iter().copied());
1590        gvals.extend(g.iter().copied());
1591        cvals.extend(c.iter().copied());
1592    }
1593    // Robust epsilon initialization from pilot magnitudes:
1594    //   s = max(median(z), 1.4826*MAD(z), Q75(z)),
1595    //   if s is tiny then fallback to max(Q95(z), RMS(z)),
1596    //   if still tiny then use absolute floor min_epsilon.
1597    // Epsilon is then kappa * s.
1598    let eps_0 = robust_epsilon_from_samples(&fvals, min_epsilon);
1599    let eps_g = robust_epsilon_from_samples(&gvals, min_epsilon);
1600    let eps_c = robust_epsilon_from_samples(&cvals, min_epsilon);
1601    Ok((eps_0, eps_g, eps_c))
1602}
1603
1604fn exact_spatial_adaptive_penalty_index_set(
1605    caches: &[SpatialOperatorRuntimeCache],
1606) -> BTreeSet<usize> {
1607    let mut out = BTreeSet::new();
1608    for cache in caches {
1609        out.insert(cache.mass_penalty_global_idx);
1610        out.insert(cache.tension_penalty_global_idx);
1611        out.insert(cache.stiffness_penalty_global_idx);
1612    }
1613    out
1614}
1615
1616fn build_spatial_adaptive_hyperspecs(cache_count: usize) -> Vec<SpatialAdaptiveHyperSpec> {
1617    let mut out = Vec::with_capacity(cache_count * 3 + 3);
1618    for cache_index in 0..cache_count {
1619        out.push(SpatialAdaptiveHyperSpec {
1620            cache_index,
1621            kind: SpatialAdaptiveHyperKind::LogLambdaMagnitude,
1622        });
1623        out.push(SpatialAdaptiveHyperSpec {
1624            cache_index,
1625            kind: SpatialAdaptiveHyperKind::LogLambdaGradient,
1626        });
1627        out.push(SpatialAdaptiveHyperSpec {
1628            cache_index,
1629            kind: SpatialAdaptiveHyperKind::LogLambdaCurvature,
1630        });
1631    }
1632    out.push(SpatialAdaptiveHyperSpec {
1633        cache_index: 0,
1634        kind: SpatialAdaptiveHyperKind::LogEpsilonMagnitude,
1635    });
1636    out.push(SpatialAdaptiveHyperSpec {
1637        cache_index: 0,
1638        kind: SpatialAdaptiveHyperKind::LogEpsilonGradient,
1639    });
1640    out.push(SpatialAdaptiveHyperSpec {
1641        cache_index: 0,
1642        kind: SpatialAdaptiveHyperKind::LogEpsilonCurvature,
1643    });
1644    out
1645}
1646
1647fn penalty_matrixwith_local_block(
1648    total_dim: usize,
1649    coeff_range: Range<usize>,
1650    local: &Array2<f64>,
1651) -> Array2<f64> {
1652    let mut out = Array2::<f64>::zeros((total_dim, total_dim));
1653    out.slice_mut(s![coeff_range.clone(), coeff_range])
1654        .assign(local);
1655    out
1656}
1657
1658fn fit_term_collectionwith_exact_spatial_adaptive_regularization(
1659    baseline: FittedTermCollection,
1660    y: ArrayView1<'_, f64>,
1661    weights: ArrayView1<'_, f64>,
1662    offset: ArrayView1<'_, f64>,
1663    family: LikelihoodSpec,
1664    options: &FitOptions,
1665    runtime_caches: &[SpatialOperatorRuntimeCache],
1666) -> Result<FittedTermCollection, EstimationError> {
1667    // Exact adaptive-regularization hyperfit.
1668    //
1669    // This replaces the old MM-plus-approximate hyperfit with the
1670    // exact pseudo-Laplace objective agreed in the math notes:
1671    //
1672    //   L_tilde(theta)
1673    //   = J(beta_hat(theta); theta) + 0.5 log det H(beta_hat(theta), theta),
1674    //
1675    // where:
1676    //   - beta_hat(theta) is the exact inner mode of the true nonquadratic
1677    //     Charbonnier-penalized objective,
1678    //   - theta contains:
1679    //       * retained quadratic log-lambdas for non-adaptive penalties,
1680    //       * one log-lambda per adaptive operator block,
1681    //       * three global log-epsilons shared by every adaptive spatial term,
1682    //   - H is the exact beta-Hessian of the true objective at the mode.
1683    //
1684    // Implementation structure:
1685    //   1. keep ordinary quadratic penalties that are unrelated to adaptive
1686    //      spatial terms in the standard outer-rho path;
1687    //   2. move the adaptive Charbonnier penalties into a one-block exact-Newton
1688    //      custom family so the inner solve uses the real model rather than an
1689    //      MM surrogate;
1690    //   3. expose exact psi-gradients for adaptive log-lambda / log-epsilon
1691    //      coordinates through the custom-family pseudo-Laplace hook;
1692    //   4. refit once at the optimized hyperparameters with all penalties frozen
1693    //      inside the exact family, so covariance and final diagnostics are
1694    //      computed on the same exact surface.
1695    let adaptive_opts = options.adaptive_regularization.clone().unwrap_or_default();
1696    let adaptive_penalty_indices = exact_spatial_adaptive_penalty_index_set(runtime_caches);
1697    let p_total = baseline.design.design.ncols();
1698    struct RetainedPenaltySetup {
1699        global_idx: usize,
1700        global_penalty: Array2<f64>,
1701        nullspace_dim: usize,
1702        log_lambda: f64,
1703        col_range: Range<usize>,
1704        hessian_piece: Array2<f64>,
1705    }
1706    use rayon::iter::{IndexedParallelIterator, IntoParallelRefIterator, ParallelIterator};
1707    let retained_setups = baseline
1708        .design
1709        .penalties
1710        .par_iter()
1711        .enumerate()
1712        .map(|(idx, bp)| {
1713            if adaptive_penalty_indices.contains(&idx) {
1714                return None;
1715            }
1716            let lambda = baseline.fit.lambdas[idx];
1717            Some(RetainedPenaltySetup {
1718                global_idx: idx,
1719                global_penalty: bp.to_global(p_total),
1720                nullspace_dim: baseline
1721                    .design
1722                    .nullspace_dims
1723                    .get(idx)
1724                    .copied()
1725                    .unwrap_or(0),
1726                log_lambda: lambda.max(1e-12).ln(),
1727                col_range: bp.col_range.clone(),
1728                hessian_piece: bp.local.mapv(|v| lambda * v),
1729            })
1730        })
1731        .collect::<Vec<_>>();
1732    let retained_count = retained_setups
1733        .iter()
1734        .filter(|setup| setup.is_some())
1735        .count();
1736    let mut retained_penalties = Vec::<Array2<f64>>::with_capacity(retained_count);
1737    let mut retained_nullspace_dims = Vec::<usize>::with_capacity(retained_count);
1738    let mut retained_log_lambdas = Vec::<f64>::with_capacity(retained_count);
1739    let mut retained_global_indices = Vec::<usize>::with_capacity(retained_count);
1740    let mut fixed_quadratichessian = Array2::<f64>::zeros((p_total, p_total));
1741    for setup in retained_setups.into_iter().flatten() {
1742        retained_penalties.push(setup.global_penalty);
1743        retained_nullspace_dims.push(setup.nullspace_dim);
1744        retained_log_lambdas.push(setup.log_lambda);
1745        retained_global_indices.push(setup.global_idx);
1746        fixed_quadratichessian
1747            .slice_mut(s![setup.col_range.clone(), setup.col_range])
1748            .scaled_add(1.0, &setup.hessian_piece);
1749    }
1750
1751    let (eps_0_init, eps_g_init, eps_c_init) = compute_initial_epsilons(
1752        &baseline.fit.beta,
1753        runtime_caches,
1754        adaptive_opts.min_epsilon,
1755    )?;
1756    let mut initial_theta =
1757        Array1::<f64>::zeros(retained_penalties.len() + runtime_caches.len() * 3 + 3);
1758    for (idx, value) in retained_log_lambdas.iter().enumerate() {
1759        initial_theta[idx] = *value;
1760    }
1761    let adaptive_log_lambda_components = runtime_caches
1762        .par_iter()
1763        .map(|cache| {
1764            [
1765                baseline.fit.lambdas[cache.mass_penalty_global_idx]
1766                    .max(1e-12)
1767                    .ln(),
1768                baseline.fit.lambdas[cache.tension_penalty_global_idx]
1769                    .max(1e-12)
1770                    .ln(),
1771                baseline.fit.lambdas[cache.stiffness_penalty_global_idx]
1772                    .max(1e-12)
1773                    .ln(),
1774            ]
1775        })
1776        .collect::<Vec<_>>();
1777    let mut at = retained_penalties.len();
1778    for logs in &adaptive_log_lambda_components {
1779        initial_theta[at] = logs[0];
1780        initial_theta[at + 1] = logs[1];
1781        initial_theta[at + 2] = logs[2];
1782        at += 3;
1783    }
1784    initial_theta[at] = eps_0_init.max(adaptive_opts.min_epsilon).ln();
1785    initial_theta[at + 1] = eps_g_init.max(adaptive_opts.min_epsilon).ln();
1786    initial_theta[at + 2] = eps_c_init.max(adaptive_opts.min_epsilon).ln();
1787
1788    let hyperspecs = build_spatial_adaptive_hyperspecs(runtime_caches.len());
1789    let zero_psi_op: std::sync::Arc<dyn gam_custom_family::CustomFamilyPsiDerivativeOperator> =
1790        std::sync::Arc::new(gam_custom_family::ZeroPsiDerivativeOperator::new(
1791            baseline.design.design.nrows(),
1792            baseline.design.design.ncols(),
1793        ));
1794    let derivative_blocks = vec![
1795        hyperspecs
1796            .par_iter()
1797            .map(|_| CustomFamilyBlockPsiDerivative {
1798                penalty_index: None,
1799                x_psi: Array2::<f64>::zeros((0, 0)),
1800                s_psi: Array2::<f64>::zeros((0, 0)),
1801                s_psi_components: None,
1802                s_psi_penalty_components: None,
1803                x_psi_psi: None,
1804                s_psi_psi: None,
1805                s_psi_psi_components: None,
1806                s_psi_psi_penalty_components: None,
1807                implicit_operator: Some(std::sync::Arc::clone(&zero_psi_op)),
1808                implicit_axis: 0,
1809                implicit_group_id: None,
1810            })
1811            .collect::<Vec<_>>(),
1812    ];
1813
1814    let mixture_link_state = options
1815        .mixture_link
1816        .clone()
1817        .as_ref()
1818        .map(state_fromspec)
1819        .transpose()
1820        .map_err(EstimationError::InvalidInput)?;
1821    let sas_link_state = options
1822        .sas_link
1823        .map(|spec| {
1824            if family.is_binomial_beta_logistic() {
1825                state_from_beta_logisticspec(spec)
1826            } else {
1827                state_from_sasspec(spec)
1828            }
1829        })
1830        .transpose()
1831        .map_err(EstimationError::InvalidInput)?;
1832    let latent_cloglog_state = options.latent_cloglog;
1833    let shared_y = Arc::new(y.to_owned());
1834    let sharedweights = Arc::new(weights.to_owned());
1835    let shared_design = baseline
1836        .design
1837        .design
1838        .try_to_dense_arc("spatial adaptive exact hyperfit design")
1839        .map_err(EstimationError::InvalidInput)?;
1840    let shared_offset = Arc::new(offset.to_owned());
1841    let shared_runtime_caches = Arc::new(runtime_caches.to_vec());
1842    let shared_hyperspecs = Arc::new(hyperspecs.clone());
1843    let zero_quadratic = Arc::new(Array2::<f64>::zeros((
1844        baseline.design.design.ncols(),
1845        baseline.design.design.ncols(),
1846    )));
1847    let base_family = SpatialAdaptiveExactFamily {
1848        family: family.clone(),
1849        latent_cloglog_state,
1850        mixture_link_state: mixture_link_state.clone(),
1851        sas_link_state,
1852        y: shared_y.clone(),
1853        weights: sharedweights.clone(),
1854        design: shared_design.clone(),
1855        offset: shared_offset.clone(),
1856        linear_constraints: baseline.design.linear_constraints.clone(),
1857        runtime_caches: shared_runtime_caches.clone(),
1858        adaptive_params: Vec::new(),
1859        fixed_quadratichessian: zero_quadratic.clone(),
1860        hyperspecs: shared_hyperspecs.clone(),
1861        exact_eval_cache: Arc::new(Mutex::new(None)),
1862    };
1863
1864    let rho_dim = retained_penalties.len();
1865    let operator_slots_end = rho_dim + runtime_caches.len() * 3;
1866    // Every slot's box is `initial_theta[idx] ± WINDOW` clamped into a
1867    // per-slot [floor, cap]. Retained-λ previously used a scale-blind
1868    // ±30 absolute interval, which on small-n / weakly-identified Duchon
1869    // fits let those lambdas wander to the exp(-30) floor and produce
1870    // near-interpolant solutions. Anchoring on baseline log-λ inherits the
1871    // baseline REML's scale calibration so the overlay can only refine
1872    // within an exp(±6) ≈ 400× band of the well-posed baseline regime,
1873    // matching the discipline already applied to operator and epsilon
1874    // slots.
1875    const UNIFIED_LOG_WINDOW: f64 = 6.0;
1876    const RETAINED_LAMBDA_LOG_LOWER_FLOOR: f64 = -30.0;
1877    const RETAINED_LAMBDA_LOG_UPPER_CAP: f64 = 30.0;
1878    const OPERATOR_LAMBDA_LOG_LOWER_FLOOR: f64 = -10.0;
1879    const OPERATOR_LAMBDA_LOG_UPPER_CAP: f64 = 30.0;
1880    let epsilon_floor_log = adaptive_opts.min_epsilon.max(1e-12).ln();
1881    let anchored_bound = |idx: usize, sign: f64| -> f64 {
1882        let raw = initial_theta[idx] + sign * UNIFIED_LOG_WINDOW;
1883        if idx < rho_dim {
1884            raw.clamp(
1885                RETAINED_LAMBDA_LOG_LOWER_FLOOR,
1886                RETAINED_LAMBDA_LOG_UPPER_CAP,
1887            )
1888        } else if idx < operator_slots_end {
1889            raw.clamp(
1890                OPERATOR_LAMBDA_LOG_LOWER_FLOOR,
1891                OPERATOR_LAMBDA_LOG_UPPER_CAP,
1892            )
1893        } else {
1894            raw.max(epsilon_floor_log)
1895        }
1896    };
1897    let eps_lower =
1898        Array1::from_iter((0..initial_theta.len()).map(|idx| anchored_bound(idx, -1.0)));
1899    let eps_upper = Array1::from_iter((0..initial_theta.len()).map(|idx| anchored_bound(idx, 1.0)));
1900    let blockspec = ParameterBlockSpec {
1901        name: "eta".to_string(),
1902        design: baseline.design.design.clone(),
1903        offset: offset.to_owned(),
1904        penalties: retained_penalties
1905            .iter()
1906            .cloned()
1907            .map(PenaltyMatrix::Dense)
1908            .collect(),
1909        nullspace_dims: retained_nullspace_dims.clone(),
1910        initial_log_lambdas: Array1::from_vec(retained_log_lambdas.clone()),
1911        initial_beta: Some(baseline.fit.beta.clone()),
1912        gauge_priority: 100,
1913        jacobian_callback: None,
1914        stacked_design: None,
1915        stacked_offset: None,
1916    };
1917    let screening_cap = Arc::new(AtomicUsize::new(0));
1918    let outer_opts = BlockwiseFitOptions {
1919        inner_max_cycles: options.max_iter,
1920        inner_tol: options.tol,
1921        outer_max_iter: options.max_iter,
1922        outer_tol: options.tol,
1923        compute_covariance: false,
1924        screening_max_inner_iterations: Some(Arc::clone(&screening_cap)),
1925        ..BlockwiseFitOptions::default()
1926    };
1927
1928    use gam_solve::rho_optimizer::OuterProblem;
1929    use gam_problem::{DeclaredHessianForm, Derivative, HessianValue, OuterEval};
1930
1931    struct SpatialAdaptiveOuterState {
1932        warm_cache: Option<CustomFamilyWarmStart>,
1933        last_eval: Option<(
1934            Array1<f64>,
1935            f64,
1936            Array1<f64>,
1937            HessianValue,
1938            CustomFamilyWarmStart,
1939        )>,
1940    }
1941
1942    let n_theta = initial_theta.len();
1943
1944    // Clamp theta to the asymmetric epsilon bounds that run_outer's symmetric
1945    // rho_bound cannot express directly.
1946    let theta_bounds = Some((eps_lower.clone(), eps_upper.clone()));
1947    let clamp_theta = {
1948        let lo = eps_lower;
1949        let hi = eps_upper;
1950        move |theta: &Array1<f64>| -> Array1<f64> {
1951            let mut clamped = theta.clone();
1952            for i in 0..clamped.len() {
1953                clamped[i] = clamped[i].clamp(lo[i], hi[i]);
1954            }
1955            clamped
1956        }
1957    };
1958
1959    let decode_theta = |theta: &Array1<f64>| -> (Array1<f64>, Vec<SpatialAdaptiveTermHyperParams>) {
1960        let rho = theta.slice(s![..rho_dim]).to_owned();
1961        let adaptive_lambda_start = rho_dim;
1962        let adaptive_lambda_end = adaptive_lambda_start + runtime_caches.len() * 3;
1963        let eps = [
1964            theta[adaptive_lambda_end].exp(),
1965            theta[adaptive_lambda_end + 1].exp(),
1966            theta[adaptive_lambda_end + 2].exp(),
1967        ];
1968        let adaptive_params = runtime_caches
1969            .iter()
1970            .enumerate()
1971            .map(|(cache_idx, _)| SpatialAdaptiveTermHyperParams {
1972                lambda: [
1973                    theta[adaptive_lambda_start + cache_idx * 3].exp(),
1974                    theta[adaptive_lambda_start + cache_idx * 3 + 1].exp(),
1975                    theta[adaptive_lambda_start + cache_idx * 3 + 2].exp(),
1976                ],
1977                epsilon: eps,
1978            })
1979            .collect::<Vec<_>>();
1980        (rho, adaptive_params)
1981    };
1982    let analytic_outer_hessian_available =
1983        gam_custom_family::joint_exact_analytic_outer_hessian_available()
1984            && base_family
1985                .exact_outer_derivative_order(std::slice::from_ref(&blockspec), &outer_opts)
1986                .has_hessian()
1987            && gam_custom_family::exact_newton_outer_geometry_supports_second_order_solver(
1988                &base_family,
1989            );
1990    // Keep the exact outer Hessian whenever the adaptive family can provide it.
1991    // The Charbonnier pseudo-Laplace surface mixes ordinary log-lambda
1992    // coordinates with adaptive λ/ε coordinates; exact curvature is the best
1993    // route when available. If a family cannot provide exact curvature, this
1994    // builder declares only the true first-order capability.
1995    let problem = OuterProblem::new(n_theta)
1996        .with_gradient(Derivative::Analytic)
1997        .with_hessian(if analytic_outer_hessian_available {
1998            DeclaredHessianForm::Either
1999        } else {
2000            DeclaredHessianForm::Unavailable
2001        })
2002        .with_fallback_policy(gam_solve::rho_optimizer::FallbackPolicy::Disabled)
2003        .with_psi_dim(n_theta.saturating_sub(rho_dim))
2004        .with_tolerance(options.tol)
2005        .with_max_iter(options.max_iter)
2006        .with_seed_config(gam_problem::SeedConfig::default())
2007        .with_screening_cap(Arc::clone(&screening_cap))
2008        .with_initial_rho(initial_theta.clone());
2009    let problem = if let Some((lo, hi)) = theta_bounds {
2010        problem.with_bounds(lo, hi)
2011    } else {
2012        problem
2013    };
2014
2015    let eval_outer = |st: &mut SpatialAdaptiveOuterState,
2016                      theta: &Array1<f64>,
2017                      order: gam_solve::rho_optimizer::OuterEvalOrder|
2018     -> Result<OuterEval, EstimationError> {
2019        let theta = clamp_theta(theta);
2020
2021        if let Some((cached_theta, cached_cost, cached_grad, cached_hess, cached_warm)) =
2022            &st.last_eval
2023            && cached_theta.len() == theta.len()
2024            && cached_theta
2025                .iter()
2026                .zip(theta.iter())
2027                .all(|(&a, &b)| (a - b).abs() <= 1e-12)
2028            && (!matches!(
2029                order,
2030                gam_solve::rho_optimizer::OuterEvalOrder::ValueGradientHessian
2031            ) || analytic_outer_hessian_available)
2032        {
2033            st.warm_cache = Some(cached_warm.clone());
2034            return Ok(OuterEval {
2035                cost: *cached_cost,
2036                gradient: cached_grad.clone(),
2037                hessian: if matches!(
2038                    order,
2039                    gam_solve::rho_optimizer::OuterEvalOrder::ValueGradientHessian
2040                ) && analytic_outer_hessian_available
2041                {
2042                    cached_hess.clone()
2043                } else {
2044                    HessianValue::Unavailable
2045                },
2046                inner_beta_hint: None,
2047            });
2048        }
2049
2050        let (rho, adaptive_params) = decode_theta(&theta);
2051        let family_eval = base_family.with_adaptive_params(adaptive_params, zero_quadratic.clone());
2052        let need_hessian = matches!(
2053            order,
2054            gam_solve::rho_optimizer::OuterEvalOrder::ValueGradientHessian
2055        ) && analytic_outer_hessian_available;
2056        let result = evaluate_custom_family_joint_hyper(
2057            &family_eval,
2058            std::slice::from_ref(&blockspec),
2059            &outer_opts,
2060            &rho,
2061            &derivative_blocks,
2062            st.warm_cache.as_ref(),
2063            if need_hessian {
2064                gam_solve::estimate::reml::reml_outer_engine::EvalMode::ValueGradientHessian
2065            } else {
2066                gam_solve::estimate::reml::reml_outer_engine::EvalMode::ValueAndGradient
2067            },
2068        )
2069        .map_err(|e| {
2070            EstimationError::RemlOptimizationFailed(format!("spatial adaptive eval failed: {e}"))
2071        })?;
2072        if !result.inner_converged {
2073            st.warm_cache = Some(result.warm_start.clone());
2074            return Err(EstimationError::RemlOptimizationFailed(
2075                "exact spatial adaptive inner solve did not converge".to_string(),
2076            ));
2077        }
2078        if !result.objective.is_finite() || result.gradient.iter().any(|v| !v.is_finite()) {
2079            return Err(EstimationError::RemlOptimizationFailed(
2080                "exact spatial adaptive objective returned non-finite values".to_string(),
2081            ));
2082        }
2083        let hessian_result = if need_hessian {
2084            if !result.outer_hessian.is_analytic() {
2085                return Err(EstimationError::RemlOptimizationFailed(
2086                    "exact spatial adaptive objective did not return an exact outer Hessian"
2087                        .to_string(),
2088                ));
2089            }
2090            match result.outer_hessian.dim() {
2091                Some(dim) if dim == theta.len() => {}
2092                Some(dim) => {
2093                    return Err(EstimationError::RemlOptimizationFailed(format!(
2094                        "exact spatial adaptive outer Hessian dimension mismatch: got {dim}, expected {}",
2095                        theta.len(),
2096                    )));
2097                }
2098                None => {
2099                    return Err(EstimationError::RemlOptimizationFailed(
2100                        "exact spatial adaptive objective did not report an outer Hessian dimension"
2101                            .to_string(),
2102                    ));
2103                }
2104            }
2105            st.last_eval = Some((
2106                theta.clone(),
2107                result.objective,
2108                result.gradient.clone(),
2109                result.outer_hessian.clone(),
2110                result.warm_start.clone(),
2111            ));
2112            result.outer_hessian
2113        } else {
2114            HessianValue::Unavailable
2115        };
2116        st.warm_cache = Some(result.warm_start);
2117        Ok(OuterEval {
2118            cost: result.objective,
2119            gradient: result.gradient,
2120            hessian: hessian_result,
2121            inner_beta_hint: None,
2122        })
2123    };
2124
2125    let mut obj = problem.build_objective_with_screening_proxy(
2126        SpatialAdaptiveOuterState {
2127            warm_cache: None,
2128            last_eval: None,
2129        },
2130        |st: &mut SpatialAdaptiveOuterState, theta: &Array1<f64>| {
2131            let theta = clamp_theta(theta);
2132            let (rho, adaptive_params) = decode_theta(&theta);
2133            let family_eval =
2134                base_family.with_adaptive_params(adaptive_params, zero_quadratic.clone());
2135            let result = evaluate_custom_family_joint_hyper(
2136                &family_eval,
2137                std::slice::from_ref(&blockspec),
2138                &outer_opts,
2139                &rho,
2140                &derivative_blocks,
2141                st.warm_cache.as_ref(),
2142                gam_solve::estimate::reml::reml_outer_engine::EvalMode::ValueOnly,
2143            )
2144            .map_err(|e| {
2145                EstimationError::RemlOptimizationFailed(format!(
2146                    "spatial adaptive cost eval failed: {e}"
2147                ))
2148            })?;
2149            if !result.inner_converged {
2150                st.warm_cache = Some(result.warm_start);
2151                return Err(EstimationError::RemlOptimizationFailed(
2152                    "exact spatial adaptive cost inner solve did not converge".to_string(),
2153                ));
2154            }
2155            st.warm_cache = Some(result.warm_start);
2156            Ok(result.objective)
2157        },
2158        |st: &mut SpatialAdaptiveOuterState, theta: &Array1<f64>| {
2159            eval_outer(
2160                st,
2161                theta,
2162                if analytic_outer_hessian_available {
2163                    gam_solve::rho_optimizer::OuterEvalOrder::ValueGradientHessian
2164                } else {
2165                    gam_solve::rho_optimizer::OuterEvalOrder::ValueAndGradient
2166                },
2167            )
2168        },
2169        |st: &mut SpatialAdaptiveOuterState,
2170         theta: &Array1<f64>,
2171         order: gam_solve::rho_optimizer::OuterEvalOrder| {
2172            eval_outer(st, theta, order)
2173        },
2174        Some(|st: &mut SpatialAdaptiveOuterState| {
2175            st.warm_cache = None;
2176            st.last_eval = None;
2177        }),
2178        Some(|st: &mut SpatialAdaptiveOuterState, theta: &Array1<f64>| {
2179            let theta = clamp_theta(theta);
2180            let (rho, adaptive_params) = decode_theta(&theta);
2181            let family_eval =
2182                base_family.with_adaptive_params(adaptive_params, zero_quadratic.clone());
2183            let result = evaluate_custom_family_joint_hyper_efs(
2184                &family_eval,
2185                std::slice::from_ref(&blockspec),
2186                &outer_opts,
2187                &rho,
2188                &derivative_blocks,
2189                st.warm_cache.as_ref(),
2190            )
2191            .map_err(|e| {
2192                EstimationError::RemlOptimizationFailed(format!(
2193                    "spatial adaptive EFS eval failed: {e}"
2194                ))
2195            })?;
2196            if !result.inner_converged {
2197                st.warm_cache = Some(result.warm_start);
2198                return Err(EstimationError::RemlOptimizationFailed(
2199                    "exact spatial adaptive EFS inner solve did not converge".to_string(),
2200                ));
2201            }
2202            st.warm_cache = Some(result.warm_start);
2203            Ok(result.efs_eval)
2204        }),
2205        // Seed-screening ranking proxy (#969). The regular cost closure
2206        // above hard-errors on a non-converged inner solve — correct for
2207        // line-search costs, but under the screening cap
2208        // (`screening_max_inner_iterations`, wired into `outer_opts`) the
2209        // inner solve is truncated BY DESIGN, so screening through that
2210        // closure rejects every seed and re-creates the all-seeds-rejected
2211        // front-door failure genus. Screening only RANKS candidates: the
2212        // penalized objective of the capped solve is a meaningful ranking
2213        // signal even unconverged (the same contract as the custom-family
2214        // labeled proxy), so accept it and let the cascade pick the best
2215        // seed; the selected seed is then fit with the full budget.
2216        |st: &mut SpatialAdaptiveOuterState, theta: &Array1<f64>| {
2217            let theta = clamp_theta(theta);
2218            let (rho, adaptive_params) = decode_theta(&theta);
2219            let family_eval =
2220                base_family.with_adaptive_params(adaptive_params, zero_quadratic.clone());
2221            let result = evaluate_custom_family_joint_hyper(
2222                &family_eval,
2223                std::slice::from_ref(&blockspec),
2224                &outer_opts,
2225                &rho,
2226                &derivative_blocks,
2227                st.warm_cache.as_ref(),
2228                gam_solve::estimate::reml::reml_outer_engine::EvalMode::ValueOnly,
2229            )
2230            .map_err(|e| {
2231                EstimationError::RemlOptimizationFailed(format!(
2232                    "spatial adaptive screening eval failed: {e}"
2233                ))
2234            })?;
2235            st.warm_cache = Some(result.warm_start);
2236            Ok(result.objective)
2237        },
2238    );
2239
2240    let outer_result = problem
2241        .run(&mut obj, "exact spatial adaptive regularization")
2242        .map_err(|e| {
2243            EstimationError::InvalidInput(format!(
2244                "exact spatial adaptive outer optimization failed: {e}"
2245            ))
2246        })?;
2247    if !outer_result.converged {
2248        // The strict absolute-floor gradient criterion (`‖g‖_proj ≤ options.tol`)
2249        // is too tight near the box-constrained boundary of the adaptive
2250        // Charbonnier pseudo-Laplace objective: as the optimizer pushes ε → ∞
2251        // (overlay-disabled corner), λ → λ_min, the Hessian's nearly-null
2252        // direction lets Cauchy/Newton accept ~e-3-magnitude probe steps that
2253        // give cost changes well below 6-digit precision, and the projected
2254        // gradient floors at numerical-noise-scale (≈ 5e-6 for n≈500, cost≈
2255        // 3e2 fits in double precision) rather than at 0. Accept the iterate
2256        // when the mgcv-style relative-to-cost criterion ‖g‖_proj ≤ τ·(1+|f|)
2257        // is satisfied — that is the textbook REML convergence rule and is
2258        // exactly what `opt::GradientTolerance::relative_to_cost(τ)` would
2259        // have enforced if this OuterProblem path had wired it through. The
2260        // strict absolute floor is retained as the primary check; the
2261        // rel-to-cost form only kicks in once the absolute one has timed out
2262        // at `max_iter`, so unconverged divergent runs (which have large |g|)
2263        // still surface as errors.
2264        let rel_to_cost_threshold = options.tol * (1.0_f64 + outer_result.final_value.abs());
2265        // Rel-to-cost acceptance requires an actual gradient measurement;
2266        // `None` (cache-hit short-circuit, gradient-free path) cannot satisfy
2267        // the mgcv-style criterion regardless of magnitude.
2268        if let Some(final_grad) = outer_result
2269            .final_grad_norm
2270            .filter(|v| v.is_finite() && *v <= rel_to_cost_threshold)
2271        {
2272            log::info!(
2273                "[spatial-adaptive] outer optimization hit max_iter={} but \
2274                 projected gradient norm {:.3e} ≤ τ·(1+|f|) = {:.3e} \
2275                 (τ={:.3e}, |f|={:.3e}); accepting iterate under the mgcv-style \
2276                 relative-to-cost REML convergence criterion.",
2277                outer_result.iterations,
2278                final_grad,
2279                rel_to_cost_threshold,
2280                options.tol,
2281                outer_result.final_value.abs(),
2282            );
2283        } else {
2284            crate::bail_invalid_estim!(
2285                "exact spatial adaptive outer optimization did not converge after {} iterations (final_objective={:.6e}, final_grad_norm={})",
2286                outer_result.iterations,
2287                outer_result.final_value,
2288                outer_result.final_grad_norm_report(),
2289            );
2290        }
2291    }
2292    let outer_iterations = outer_result.iterations;
2293    // `None` = no gradient measurement (cache-hit / gradient-free); the
2294    // authoritative convergence signal is `outer_converged`.
2295    let outer_grad_norm: Option<f64> = outer_result.final_grad_norm;
2296    let theta_star = outer_result.rho;
2297    let rho_star = theta_star.slice(s![..rho_dim]).to_owned();
2298    let adaptive_lambda_start = rho_dim;
2299    let adaptive_lambda_end = adaptive_lambda_start + runtime_caches.len() * 3;
2300    let eps_star = [
2301        theta_star[adaptive_lambda_end].exp(),
2302        theta_star[adaptive_lambda_end + 1].exp(),
2303        theta_star[adaptive_lambda_end + 2].exp(),
2304    ];
2305    let adaptive_params = runtime_caches
2306        .iter()
2307        .enumerate()
2308        .map(|(cache_idx, _)| SpatialAdaptiveTermHyperParams {
2309            lambda: [
2310                theta_star[adaptive_lambda_start + cache_idx * 3].exp(),
2311                theta_star[adaptive_lambda_start + cache_idx * 3 + 1].exp(),
2312                theta_star[adaptive_lambda_start + cache_idx * 3 + 2].exp(),
2313            ],
2314            epsilon: eps_star,
2315        })
2316        .collect::<Vec<_>>();
2317    let mut fixed_total = Array2::<f64>::zeros((
2318        baseline.design.design.ncols(),
2319        baseline.design.design.ncols(),
2320    ));
2321    for (idx, penalty) in retained_penalties.iter().enumerate() {
2322        fixed_total.scaled_add(rho_star[idx].exp(), penalty);
2323    }
2324    let final_family =
2325        base_family.with_adaptive_params(adaptive_params.clone(), Arc::new(fixed_total.clone()));
2326    let final_blockspec = ParameterBlockSpec {
2327        name: "eta".to_string(),
2328        design: baseline.design.design.clone(),
2329        offset: offset.to_owned(),
2330        penalties: vec![],
2331        nullspace_dims: vec![],
2332        initial_log_lambdas: Array1::zeros(0),
2333        initial_beta: Some(baseline.fit.beta.clone()),
2334        gauge_priority: 100,
2335        jacobian_callback: None,
2336        stacked_design: None,
2337        stacked_offset: None,
2338    };
2339    let final_fit = fit_custom_family(
2340        &final_family,
2341        &[final_blockspec],
2342        &BlockwiseFitOptions {
2343            inner_max_cycles: options.max_iter,
2344            inner_tol: options.tol,
2345            outer_max_iter: 1,
2346            outer_tol: options.tol,
2347            compute_covariance: true,
2348            ..BlockwiseFitOptions::default()
2349        },
2350    )
2351    .map_err(EstimationError::CustomFamily)?;
2352    let beta = final_fit.block_states[0].beta.clone();
2353    let final_eval = final_family
2354        .exact_evaluation(&beta)
2355        .map_err(EstimationError::InvalidInput)?;
2356    let penalized_hessian = final_eval
2357        .totalobjectivehessian(&final_family.design)
2358        .map_err(EstimationError::InvalidInput)?;
2359    let beta_covariance = final_fit.covariance_conditional.clone();
2360    let beta_standard_errors = beta_covariance
2361        .as_ref()
2362        .map(|cov| Array1::from_iter((0..cov.nrows()).map(|i| cov[[i, i]].max(0.0).sqrt())));
2363
2364    let mut full_lambdas = baseline.fit.lambdas.clone();
2365    for (idx, &global_idx) in retained_global_indices.iter().enumerate() {
2366        full_lambdas[global_idx] = rho_star[idx].exp();
2367    }
2368    for (cache_idx, cache) in runtime_caches.iter().enumerate() {
2369        full_lambdas[cache.mass_penalty_global_idx] = adaptive_params[cache_idx].lambda[0];
2370        full_lambdas[cache.tension_penalty_global_idx] = adaptive_params[cache_idx].lambda[1];
2371        full_lambdas[cache.stiffness_penalty_global_idx] = adaptive_params[cache_idx].lambda[2];
2372    }
2373
2374    let deviance = if family.is_gaussian_identity() {
2375        y.iter()
2376            .zip(final_eval.obs.mu.iter())
2377            .zip(weights.iter())
2378            .map(|((&yy, &mu), &w)| w.max(0.0) * (yy - mu) * (yy - mu))
2379            .sum()
2380    } else {
2381        -2.0 * final_eval.obs.log_likelihood
2382    };
2383    let mut local_penalty_blocks =
2384        Vec::<PenaltySpec>::with_capacity(baseline.design.penalties.len());
2385    for (global_idx, bp) in baseline.design.penalties.iter().enumerate() {
2386        if adaptive_penalty_indices.contains(&global_idx) {
2387            let cache = runtime_caches
2388                .iter()
2389                .find(|cache| {
2390                    cache.mass_penalty_global_idx == global_idx
2391                        || cache.tension_penalty_global_idx == global_idx
2392                        || cache.stiffness_penalty_global_idx == global_idx
2393                })
2394                .ok_or_else(|| {
2395                    EstimationError::InvalidInput(format!(
2396                        "missing runtime cache for adaptive penalty index {global_idx}"
2397                    ))
2398                })?;
2399            let cache_idx = runtime_caches
2400                .iter()
2401                .position(|c| {
2402                    c.mass_penalty_global_idx == global_idx
2403                        || c.tension_penalty_global_idx == global_idx
2404                        || c.stiffness_penalty_global_idx == global_idx
2405                })
2406                .ok_or_else(|| {
2407                    EstimationError::InvalidInput(format!(
2408                        "missing adaptive cache position for penalty index {global_idx}"
2409                    ))
2410                })?;
2411            let state = &final_eval.adaptive_states[cache_idx];
2412            let local = if cache.mass_penalty_global_idx == global_idx {
2413                scalar_operatorhessian(&cache.d0, &state.magnitude.betahessian_diag())
2414                    .mapv(|v| adaptive_params[cache_idx].lambda[0] * v)
2415            } else if cache.tension_penalty_global_idx == global_idx {
2416                grouped_operatorhessian(
2417                    &cache.d1,
2418                    cache.dimension,
2419                    &state.gradient.betahessian_blocks(),
2420                )?
2421                .mapv(|v| adaptive_params[cache_idx].lambda[1] * v)
2422            } else {
2423                grouped_operatorhessian(
2424                    &cache.d2,
2425                    cache.dimension * cache.dimension,
2426                    &state.curvature.betahessian_blocks(),
2427                )?
2428                .mapv(|v| adaptive_params[cache_idx].lambda[2] * v)
2429            };
2430            // Wrap the pre-scaled global penalty matrix as PenaltySpec::Dense.
2431            local_penalty_blocks.push(PenaltySpec::Dense(penalty_matrixwith_local_block(
2432                baseline.design.design.ncols(),
2433                cache.coeff_global_range.clone(),
2434                &local,
2435            )));
2436        } else {
2437            local_penalty_blocks.push(PenaltySpec::Dense(
2438                bp.to_global(p_total).mapv(|v| v * full_lambdas[global_idx]),
2439            ));
2440        }
2441    }
2442    let (edf_by_block, penalty_block_trace, edf_total) = if let Some(cov) = beta_covariance.as_ref()
2443    {
2444        exact_bounded_edf(
2445            &local_penalty_blocks,
2446            &Array1::from_elem(local_penalty_blocks.len(), 1.0),
2447            cov,
2448        )?
2449    } else {
2450        (
2451            vec![0.0; local_penalty_blocks.len()],
2452            vec![0.0; local_penalty_blocks.len()],
2453            0.0,
2454        )
2455    };
2456    let stable_penalty_term =
2457        2.0 * final_eval.adaptive_penalty_value + beta.dot(&fixed_total.dot(&beta));
2458    let standard_deviation = if family.is_gaussian_identity() {
2459        let denom = (y.len() as f64 - edf_total).max(1.0);
2460        (deviance / denom).sqrt()
2461    } else {
2462        1.0
2463    };
2464    let maps = compute_spatial_adaptiveweights_for_beta(
2465        &beta,
2466        runtime_caches,
2467        eps_star[0],
2468        eps_star[1],
2469        eps_star[2],
2470        adaptive_opts.weight_floor,
2471        adaptive_opts.weight_ceiling,
2472        // Working-Laplace conditional covariance Sigma_beta = H^{-1} from the
2473        // final exact-family solve, reused here as the posterior-SNR variance
2474        // source (no second factorization is formed).
2475        beta_covariance.as_ref(),
2476    )?
2477    .into_iter()
2478    .zip(runtime_caches.iter())
2479    .map(|(w, cache)| AdaptiveSpatialMap {
2480        termname: cache.termname.clone(),
2481        feature_cols: cache.feature_cols.clone(),
2482        collocation_points: cache.collocation_points.clone(),
2483        inv_magweight: w.inv_magweight,
2484        invgradweight: w.invgradweight,
2485        inv_lapweight: w.inv_lapweight,
2486    })
2487    .collect::<Vec<_>>();
2488    let fitted_link = if family.is_latent_cloglog() {
2489        FittedLinkState::LatentCLogLog {
2490            state: latent_cloglog_state
2491                .expect("BinomialLatentCLogLog requires an explicit latent-cloglog state"),
2492        }
2493    } else if family.is_binomial_mixture() {
2494        mixture_link_state
2495            .clone()
2496            .map(|state| FittedLinkState::Mixture {
2497                state,
2498                covariance: None,
2499            })
2500            .unwrap_or(FittedLinkState::Standard(None))
2501    } else if family.is_binomial_sas() {
2502        sas_link_state
2503            .map(|state| FittedLinkState::Sas {
2504                state,
2505                covariance: None,
2506            })
2507            .unwrap_or(FittedLinkState::Standard(None))
2508    } else if family.is_binomial_beta_logistic() {
2509        sas_link_state
2510            .map(|state| FittedLinkState::BetaLogistic {
2511                state,
2512                covariance: None,
2513            })
2514            .unwrap_or(FittedLinkState::Standard(None))
2515    } else {
2516        FittedLinkState::Standard(None)
2517    };
2518    let max_abs_eta = final_eval
2519        .obs
2520        .eta
2521        .iter()
2522        .fold(0.0_f64, |acc, &v| acc.max(v.abs()));
2523    let fitted = FittedTermCollection {
2524        fit: {
2525            let log_lambdas = full_lambdas.mapv(|v| v.max(1e-300).ln());
2526            let inf = FitInference {
2527                edf_by_block,
2528                penalty_block_trace,
2529                edf_total,
2530                smoothing_correction: None,
2531                // Boundary adapter: wrap the raw `Array2<f64>` Hessian as
2532                // `UnscaledPrecision` for the newtype storage.
2533                penalized_hessian: penalized_hessian.clone().into(),
2534                working_weights: final_eval.obs.fisherweight.clone(),
2535                working_response: {
2536                    let mut out = final_eval.obs.eta.clone();
2537                    for i in 0..out.len() {
2538                        let wi = final_eval.obs.fisherweight[i].max(1e-12);
2539                        out[i] += final_eval.obs.score[i] / wi;
2540                    }
2541                    out
2542                },
2543                reparam_qs: None,
2544                dispersion: gam_solve::estimate::Dispersion::Known(1.0),
2545                beta_covariance: beta_covariance
2546                    .clone()
2547                    .map(gam_problem::dispersion_cov::PhiScaledCovariance::from),
2548                beta_standard_errors,
2549                beta_covariance_corrected: None,
2550                beta_standard_errors_corrected: None,
2551                beta_covariance_frequentist: None,
2552                coefficient_influence: None,
2553                weighted_gram: None,
2554                bias_correction_beta: None,
2555                bias_correction_jacobian: None,
2556            };
2557            let geometry = Some(gam_solve::estimate::FitGeometry {
2558                penalized_hessian: penalized_hessian.into(),
2559                working_weights: inf.working_weights.clone(),
2560                working_response: inf.working_response.clone(),
2561            });
2562            let covariance_conditional = beta_covariance;
2563            // `final_fit` is a sealed `UnifiedFitResult`: it can only exist
2564            // because `try_from_parts` already certified inner+outer
2565            // convergence, so its outer status is convergence by construction.
2566            let pirls_status_val = gam_solve::pirls::PirlsStatus::Converged;
2567            UnifiedFitResult::try_from_parts(UnifiedFitResultParts {
2568                blocks: vec![gam_solve::estimate::FittedBlock {
2569                    beta: beta.clone(),
2570                    role: gam_problem::BlockRole::Mean,
2571                    edf: edf_total,
2572                    lambdas: full_lambdas.clone(),
2573                }],
2574                log_lambdas,
2575                lambdas: full_lambdas,
2576                likelihood_scale: family.default_scale_metadata(),
2577                likelihood_family: Some(family),
2578                log_likelihood_normalization:
2579                    gam_spec::LogLikelihoodNormalization::UserProvided,
2580                log_likelihood: final_eval.obs.log_likelihood,
2581                deviance,
2582                reml_score: final_fit.penalized_objective,
2583                stable_penalty_term,
2584                penalized_objective: final_fit.penalized_objective,
2585                used_device: false,
2586                outer_iterations,
2587                // Sealed result ⇒ outer convergence was certified at assembly.
2588                outer_converged: true,
2589                outer_gradient_norm: outer_grad_norm,
2590                standard_deviation,
2591                covariance_conditional,
2592                covariance_corrected: None,
2593                inference: Some(inf),
2594                fitted_link,
2595                geometry,
2596                block_states: Vec::new(),
2597                pirls_status: pirls_status_val,
2598                max_abs_eta,
2599                constraint_kkt: None,
2600                artifacts: gam_solve::estimate::FitArtifacts {
2601                    pirls: None,
2602                    ..Default::default()
2603                },
2604                inner_cycles: 0,
2605            })?
2606        },
2607        design: baseline.design,
2608        adaptive_diagnostics: Some(AdaptiveRegularizationDiagnostics {
2609            epsilon_0: eps_star[0],
2610            epsilon_g: eps_star[1],
2611            epsilon_c: eps_star[2],
2612            epsilon_outer_iterations: outer_iterations,
2613            mm_iterations: 0,
2614            converged: true,
2615            maps,
2616        }),
2617    };
2618    enforce_term_constraint_feasibility(&fitted.design, &fitted.fit)?;
2619    Ok(fitted)
2620}
2621
2622/// Relax the per-coordinate ρ-prior for terms running in Marra–Wood
2623/// double-penalty selection mode (#1266).
2624///
2625/// The default ρ-prior is a `Normal { mean: 0, sd: 3 }` cap on each log-λ — a
2626/// stabiliser that keeps ordinary smoothing parameters from drifting to
2627/// degenerate extremes (gam#893/#1196). For a smooth carrying a
2628/// `DoublePenaltyNullspace` block (`double_penalty = True`, the default `s(...)`
2629/// — analogous to mgcv `select = TRUE`) that cap is actively wrong: the whole
2630/// purpose of the second penalty is to let REML drive an *unsupported* term to
2631/// `EDF → 0`, which needs both the wiggliness and null-space log-λ to grow
2632/// large. The `ρ²/(2·9)` cap pulls them back toward 0, so REML settles at a
2633/// point that leaves the term under-shrunk — the smooth's EDF comes out ABOVE
2634/// the single-penalty (`double_penalty = False`) EDF instead of at or below it,
2635/// the exact contract violation in #1266. mgcv's `select = TRUE` applies no
2636/// such cap to the selection coordinates, and the lower-level term-collection
2637/// fits already converge correctly under a flat prior.
2638///
2639/// We therefore rewrite the prior to `Independent`, holding the base prior on
2640/// every ordinary coordinate but switching the coordinates of any
2641/// double-penalty term to `Flat`. Single-penalty terms are byte-for-byte
2642/// unchanged, and an already-`Flat`/already-`Independent` base prior, or a
2643/// design with no double-penalty block, is returned untouched.
2644///
2645/// The relaxed per-coordinate prior is FAMILY-AGNOSTIC: the cap-lifting of the
2646/// bending coordinate and the determinacy-gated null-space treatment apply
2647/// identically for Gaussian and non-Gaussian families. The response family / link
2648/// only matters for length-safety (it can append auxiliary trailing ρ
2649/// coordinates via dispersion / SAS / mixture / moving-κ machinery), which is
2650/// gated separately by `length_safe`; once that gate passes the inner ρ aligns
2651/// 1:1 with `penaltyinfo` regardless of family, so the same relaxation is valid
2652/// for a Tweedie / Gamma-log `ps` smooth as for a Gaussian one (#1426/#1477).
2653fn relax_smoothing_rho_prior(
2654    options: &FitOptions,
2655    design: &TermCollectionDesign,
2656) -> gam_spec::RhoPrior {
2657    use gam_terms::basis::BasisMetadata;
2658    let base = &options.rho_prior;
2659    // Only a single scalar prior that actually caps log-λ needs relaxing;
2660    // `Flat` already imposes no cap and `Independent` is assumed caller-built.
2661    if matches!(
2662        base,
2663        gam_spec::RhoPrior::Flat | gam_spec::RhoPrior::Independent(_)
2664    ) {
2665        return base.clone();
2666    }
2667    // LENGTH SAFETY (load-bearing). The per-coordinate `Independent` prior is
2668    // validated against the FULL outer ρ vector and a length disagreement
2669    // saturates the prior to `+∞`, breaking the fit. The ρ vector this prior is
2670    // attached to (the inner REML fit at a *fixed* realized design) aligns 1:1
2671    // with the penalty blocks in `design.penaltyinfo` ONLY when the fit
2672    // introduces no auxiliary trailing ρ coordinates. Such coordinates come from
2673    //   * non-Gaussian dispersion / non-identity link machinery,
2674    //   * SAS ε/δ and mixture-link parameters,
2675    //   * spatial κ length-scale optimisation that actually moves κ.
2676    // Gate to the link-aux-free case. Spatial κ optimisation (Matérn / Duchon /
2677    // sphere / curvature / measure-jet) genuinely appends a moving log-κ
2678    // coordinate AND needs the cap to stabilise it, so bail if any such term is
2679    // present. Thin-plate is the exception: its length-scale is a pure radial
2680    // SCALE that REML cannot identify (the κ optimiser converges to a no-op,
2681    // leaving `n_params = penalty-block count`), so it adds no trailing
2682    // coordinate and is safe to relax alongside the B-spline family. The response
2683    // family / link itself does NOT break length-safety (a non-Gaussian GAM with
2684    // no link-aux and no moving κ still has exactly `penaltyinfo.len()` inner ρ
2685    // coordinates), so the relaxed prior below is family-agnostic.
2686    let has_link_aux = options.sas_link.is_some()
2687        || options.optimize_sas
2688        || options.mixture_link.is_some()
2689        || options.optimize_mixture;
2690    let has_moving_kappa = design.smooth.terms.iter().any(|t| {
2691        // A PURE (scale-free) Duchon / polyharmonic smooth carries NO free length
2692        // scale: its radial scale is REML-unidentifiable, so — exactly like
2693        // thin-plate — the isotropic κ prescan skips it
2694        // (`prescan_isotropic_spatial_range_seed`: "Pure Duchon / TPS without a
2695        // length scale are skipped"), it is never assigned a `length_scale`, and it
2696        // appends NO moving log-κ ρ coordinate. The inner ρ vector then aligns 1:1
2697        // with `penaltyinfo` just as it does for `tp`, so relaxing its symmetric
2698        // cap is length-safe. Only a HYBRID Duchon-Matérn term
2699        // (`length_scale = Some`) or an ANISOTROPIC Duchon (`aniso_log_scales =
2700        // Some`) puts a genuine moving κ into the inner ρ vector and needs the cap
2701        // as a stabiliser. Treat pure Duchon as κ-free; every other spatial family
2702        // keeps the blanket exclusion.
2703        if let BasisMetadata::Duchon {
2704            length_scale,
2705            aniso_log_scales,
2706            ..
2707        } = &t.metadata
2708        {
2709            return length_scale.is_some() || aniso_log_scales.is_some();
2710        }
2711        matches!(
2712            t.metadata,
2713            BasisMetadata::Matern { .. }
2714                | BasisMetadata::Sphere { .. }
2715                | BasisMetadata::SphereHarmonics { .. }
2716                | BasisMetadata::ConstantCurvature { .. }
2717                | BasisMetadata::MeasureJet { .. }
2718        )
2719    });
2720    // LENGTH SAFETY decides only whether the inner ρ aligns 1:1 with the penalty
2721    // blocks (so an `Independent` prior is valid): it is broken by SAS/mixture
2722    // link-shape coordinates and by a moving spatial κ, NOT by the response
2723    // family or link per se. A Gamma/log (or any other non-Gaussian) GAM with no
2724    // link-aux and no moving κ has exactly `penaltyinfo.len()` ρ coordinates, so
2725    // the `DoublePenaltyNullspace` selection prior below is length-safe there too.
2726    let length_safe = !has_link_aux && !has_moving_kappa;
2727    if !length_safe {
2728        return base.clone();
2729    }
2730    let coords = &design.penaltyinfo;
2731    if coords.is_empty() {
2732        return base.clone();
2733    }
2734    // WELL-IDENTIFICATION GATE (#1089). The ρ-prior is two things at once: a
2735    // #1266/#1271-harmful symmetric cap on each smoothing log-λ, AND a
2736    // #1089-load-bearing stabiliser that makes the outer REML loop terminate on
2737    // an *under-determined* design (gam#893/#1196/#1089: the n=30 five-`ps` wine
2738    // fit has p ≈ 51 > n, so without the cap's curvature the outer criterion is
2739    // flat/degenerate in ρ-space and the loop never certifies a stationary
2740    // point). Only lift the cap when the data comfortably over-determines the
2741    // model (`n ≥ 2·p`), so the unregularised REML problem is well-posed on its
2742    // own; otherwise keep the base prior. The #1266/#1271 cases (n ≈ 800,
2743    // p ≈ 20–40) clear this by ≥20×; the #1089 wine fit (n < p) keeps its cap.
2744    let n_obs = design.design.nrows();
2745    let p_total = design.design.ncols();
2746    // REGIME of the relaxed prior on the relaxable smooth coordinates.
2747    //
2748    // * WELL-DETERMINED (`n ≥ 2·p`): the unregularised REML problem is well
2749    //   posed on its own, so the relaxable coordinates are freed to `Flat`,
2750    //   which the runtime resolves to the firth one-sided barrier — byte-flat
2751    //   on the identified side (pure REML, exactly mgcv) and only a convex wall
2752    //   against the `λ → 0` degeneracy. This is the #1266/#1271 behaviour.
2753    //
2754    // * UNDER-DETERMINED (`n < 2·p`): the design does NOT over-determine the
2755    //   model (the n≈26 five-`ps` wine fit has p > n), so the firth barrier's
2756    //   zero curvature on the identified side leaves the outer REML criterion
2757    //   flat/degenerate in ρ-space and the loop hits `max_iter` at whatever
2758    //   (under-smoothed) λ it last held — EDF rails up to ≈n, the smooths
2759    //   interpolate the training rows, and held-out prediction explodes
2760    //   (#1392: held-out R² as low as −2.5e6 on `wine_gamair`). The previous
2761    //   stabiliser kept the FULL base prior here — a symmetric
2762    //   `Normal{mean:0, sd:3}` cap. Its `ρ²/(2·9)` curvature does terminate the
2763    //   loop, but it is centred at λ=1 with a tight `sd=3`: at the REML optimum
2764    //   `ρ* ≈ 8–15` (heavy smoothing, which an over-parameterised fit needs and
2765    //   which mgcv's pure REML reaches), the cap's `ρ*/9` gradient drags λ back
2766    //   down by `O(1)` in ρ, pinning the fit in the under-smoothed regime.
2767    //
2768    //   The fix keeps a stabiliser with strictly positive curvature (so the
2769    //   loop still certifies a stationary point — the #1089 requirement) but
2770    //   WIDENS it to `sd = RELAX_UNDERDETERMINED_RHO_SD` so its gradient drag at
2771    //   the heavily-smoothed optimum is negligible (`ρ*/sd² = O(1/100)`) and
2772    //   pure REML — not the prior — chooses λ. The wide symmetric Gaussian is
2773    //   weakly informative: ±2σ spans the whole feasible ρ range (`|ρ| ≤ 30`),
2774    //   so it adds termination curvature without biasing which λ REML lands on,
2775    //   restoring the mgcv-like heavy smoothing on the over-parameterised fit.
2776    let underdetermined = n_obs < 2 * p_total;
2777    // Relaxable terms: penalized smooths whose smoothing log-λ the symmetric cap
2778    // wrongly bounds when the term's signal lives in its penalty null space — a
2779    // straight line under a bending penalty drives λ → ∞ but the cap pulls it
2780    // back, leaving spurious wiggle. mgcv caps neither. This is exactly the
2781    // B-spline family (`ps`/`cr`/`cs`/`bs`, BSpline1D), thin-plate (`tp`), and
2782    // tensor-B-spline (`te`/`ti`) smooths — single- AND double-penalty (#1266 is
2783    // the double-penalty case, #1271 the single-penalty `tp`/`ps`). EVERY penalty
2784    // coordinate such a term owns (bending wiggliness AND any null-space
2785    // shrinkage) is freed to `Flat`, which the runtime resolves to the
2786    // firth-default one-sided barrier: no high-λ cap, but still a convex wall
2787    // against the `λ → 0` under-smoothing degeneracy.
2788    let relaxable_terms: std::collections::HashSet<&str> = design
2789        .smooth
2790        .terms
2791        .iter()
2792        .filter(|t| {
2793            (matches!(
2794                t.metadata,
2795                BasisMetadata::BSpline1D { .. }
2796                    | BasisMetadata::ThinPlate { .. }
2797                    | BasisMetadata::TensorBSpline { .. }
2798            )
2799            // A PURE (scale-free) Duchon / polyharmonic smooth IS a thin-plate
2800            // spline (unidentifiable radial scale, no moving κ coordinate — see the
2801            // `has_moving_kappa` note), so its smoothing log-λ earns the SAME cap
2802            // relaxation as `tp`. A straight-line truth under a Duchon bending
2803            // penalty drives λ → ∞ (the collapse shelf mgcv `bs="ds"` rails to,
2804            // edf → null); the symmetric `Normal{0,3}` cap otherwise pins it in the
2805            // under-smoothed interior (#1867 null-recovery over-smoothing: the
2806            // summed-diagonal shelf seed b26e1cfe9 could never win because
2807            // `compute_cost` charged it the cap's ρ²/2·9 penalty). Hybrid
2808            // Duchon-Matérn (`length_scale = Some`) / anisotropic Duchon keep the
2809            // cap — their κ is a real moving coordinate that needs the stabiliser.
2810            || matches!(
2811                t.metadata,
2812                BasisMetadata::Duchon {
2813                    length_scale: None,
2814                    aniso_log_scales: None,
2815                    ..
2816                }
2817            ))
2818            // SHAPE-CONSTRAINED terms must KEEP the cap (#1380). A monotone /
2819            // convex / concave smooth carries linear-inequality constraints; at
2820            // the active boundary (e.g. a convex fit pinned at 2nd-diff = 0) the
2821            // active set collapses the penalized subspace onto the bending
2822            // penalty's own null space ({1, x}), where the smoothing log-λ is
2823            // UNIDENTIFIED. Lifting the cap to `Flat` there lets REML rail λ to
2824            // `RHO_BOUND` (zero curvature → the smooth collapses to a flat/linear
2825            // fit, R² ≈ 0 on data the constraint is correct for). The constraint
2826            // already regularizes the term, and the symmetric cap is the
2827            // #1089-style stabiliser that pins the unidentified λ — so a
2828            // shape-constrained term needs the cap KEPT, exactly the
2829            // under-determined case this gate protects. (Unconstrained #1266/#1271
2830            // selection terms still relax.)
2831            && matches!(t.shape, gam_terms::smooth::ShapeConstraint::None)
2832        })
2833        .map(|t| t.name.as_str())
2834        .collect();
2835    let any_relaxed = coords.iter().any(|info| {
2836        info.termname
2837            .as_deref()
2838            .is_some_and(|name| relaxable_terms.contains(name))
2839    });
2840    if !any_relaxed {
2841        return base.clone();
2842    }
2843    // Relaxed prior for a relaxable smooth coordinate, chosen by regime (see the
2844    // block above): the firth one-sided barrier (`Flat`) when the fit is
2845    // well-determined, a wide-but-curved symmetric Gaussian when it is
2846    // under-determined and the loop still needs termination curvature.
2847    let relaxed_prior = if underdetermined {
2848        gam_spec::RhoPrior::Normal {
2849            mean: 0.0,
2850            sd: RELAX_UNDERDETERMINED_RHO_SD,
2851        }
2852    } else {
2853        gam_spec::RhoPrior::Flat
2854    };
2855    // DOUBLE-PENALTY NULL-SPACE SELECTION (#1392, mgcv `select=TRUE`). A
2856    // double-penalty smooth carries a second `DoublePenaltyNullspace` ridge on
2857    // the term's penalty null space ({1, x} for a 1-D bend) whose only job is
2858    // selection: drive its λ UP (toward the prior's finite well-penalized mode
2859    // λ* = θ², not to ∞) to shrink the null-space (linear) component OUT when
2860    // the data does not support it, exactly as mgcv's `select=TRUE` adds a
2861    // null-space penalty. On an over-parameterized `p > n` fit
2862    // (`wine_gamair`: 5 `ps` smooths on ~26 rows) the symmetric relaxed prior
2863    // above leaves this ridge's outer score flat on the select-out side, so REML
2864    // stalls it at λ ≈ 0.11 — the null space is kept, the EDF rails up, and
2865    // held-out prediction collapses (#1392). The RANGE-space (`Primary`) bending
2866    // coordinate's smoothing selection must NOT be touched, so this select-out
2867    // bias is gated to `DoublePenaltyNullspace` coordinates only and is applied
2868    // ONLY in the under-determined regime — in the well-determined regime the
2869    // relaxable coordinates stay byte-flat (`Flat`) so a clean `n > p` fit is
2870    // unchanged (no regression on ordinary smooth recovery).
2871    //
2872    // The strong select-out PC prior is applied to the `DoublePenaltyNullspace`
2873    // coordinate ONLY in the UNDER-DETERMINED regime, where the outer score is
2874    // genuinely flat on the select-out side and REML needs the active push. In the
2875    // WELL-DETERMINED regime the null space gets the wide
2876    // `nullspace_degeneracy_prior` instead (see below) — an active select-out mode
2877    // there would over-shrink a genuinely-supported collinear null space (#1476).
2878    // The RANGE-space (`Primary`) bending coordinate is untouched (stays `Flat`
2879    // when well-determined), so ordinary single-smooth recovery is unchanged.
2880    //
2881    let nullspace_select_prior = gam_spec::RhoPrior::PenalizedComplexity {
2882        upper: NULLSPACE_SELECT_PC_UPPER,
2883        tail_prob: NULLSPACE_SELECT_PC_TAIL_PROB,
2884    };
2885    // WELL-DETERMINED NULL-SPACE DEGENERACY BREAKER (#1476). When the fit is
2886    // well-determined (`n ≥ 2·p`) the strong `nullspace_select_prior` above is the
2887    // WRONG tool for the Gaussian null-space coordinate: its finite well-penalized
2888    // mode at `λ* = θ² ≈ 8483` is an aggressive select-OUT pull that drags a
2889    // GENUINELY-SUPPORTED null space (a real linear/constant component) toward
2890    // collapse — the #1476 over-shrink. But leaving the coordinate fully `Flat`
2891    // (the previous well-determined behaviour) is the OTHER failure: under
2892    // concurvity (`s(x1)+s(x2)`, corr ≈ 0.9) the two smooths' null-space (linear)
2893    // directions are near-collinear, so the joint REML objective is essentially
2894    // FLAT along the "transfer the shared linear signal between the two smooths"
2895    // ridge; with zero curvature on that coordinate REML cannot certify an
2896    // interior stationary point and one smooth's `λ_nullspace` rails to the ρ
2897    // bound (≈1e13), annihilating its genuine linear signal to `EDF ≈ 0` while the
2898    // other absorbs it. The principled fix is NEITHER a select-out mode NOR a
2899    // flat coordinate: it is a WIDE, weakly-informative symmetric Gaussian that
2900    // contributes strictly-positive termination curvature `1/sd²` (breaking the
2901    // concurvity flat-ridge degeneracy so REML lands an interior allocation) while
2902    // its gradient `ρ/sd²` at any plausible optimum is negligible — so REML, not
2903    // the prior, chooses how the shared linear signal is split. This adds no
2904    // directional select-out bias, so it does NOT over-shrink a supported null
2905    // space (#1476); a genuinely-UNSUPPORTED null space is still selected out
2906    // because REML's own score drives its `λ` up and the weak symmetric pull
2907    // barely opposes it (#1266 irrelevant-covariate shrinkage, #1371 single-smooth
2908    // recovery preserved). The strong PC select-out remains in the
2909    // UNDER-DETERMINED regime, where the score IS flat on the select-out side and
2910    // REML needs the active push (#1392 wine `p > n`).
2911    let nullspace_degeneracy_prior = gam_spec::RhoPrior::Normal {
2912        mean: 0.0,
2913        sd: NULLSPACE_WELLDET_DEGENERACY_RHO_SD,
2914    };
2915    let per_coord = coords
2916        .iter()
2917        .map(|info| {
2918            let relax = info
2919                .termname
2920                .as_deref()
2921                .is_some_and(|name| relaxable_terms.contains(name));
2922            if !relax {
2923                return base.clone();
2924            }
2925            let is_nullspace =
2926                matches!(info.penalty.source, PenaltySource::DoublePenaltyNullspace);
2927            // The relaxed per-coordinate prior is FAMILY-AGNOSTIC: the choice
2928            // depends only on the coordinate's role (bending vs null-space
2929            // selection) and on whether the data over-determines the model, NOT
2930            // on the response family or link. (Length-safety — the only thing the
2931            // family/link can break via auxiliary ρ coordinates — is already
2932            // gated above by `length_safe`; reaching this point means the inner ρ
2933            // aligns 1:1 with `penaltyinfo` for Gaussian and non-Gaussian alike.)
2934            //
2935            // The previous code split here on `gaussian_identity` and pinned the
2936            // non-Gaussian null-space coordinate to the AGGRESSIVE PC select-out
2937            // prior in BOTH determinacy regimes. That select-out prior has a
2938            // finite well-penalized mode at λ* ≈ θ² ≈ 8483, which carves a SECOND,
2939            // deep basin into the 2-D (bending, null-space) outer REML surface at
2940            // large λ_null. On a well-determined non-Gaussian double-penalty `ps`
2941            // smooth the outer ARC then has two competing basins — the genuine
2942            // bending optimum and the prior-induced high-λ_null shelf — and the
2943            // expensive non-Gaussian multi-start lands the wrong one: the fit
2944            // ships a right-boundary blow-up (Tweedie `s(x)` pred ≈ 1.4–2.0× truth
2945            // at x=1 on data whose null space is unsupported) and, on the hard
2946            // seeds, a falsely-"converged" EDF-inflated under-smooth (#1477; the
2947            // same genus as the #1426 Gamma/log overfit). The Gaussian path does
2948            // NOT do this — #1476 deliberately switched its well-determined
2949            // null-space coordinate to the wide, weakly-informative degeneracy
2950            // prior precisely because the active select-out over-shrinks /
2951            // destabilises a well-determined fit. Non-Gaussian needs the identical
2952            // treatment, so the determinacy gate now applies to BOTH families:
2953            //
2954            //   * BENDING (range-space) coordinate → `relaxed_prior` (firth
2955            //     one-sided barrier when well-determined = pure REML = mgcv; wide
2956            //     #1089 `Normal` when under-determined).
2957            //   * NULL-SPACE selection coordinate → the AGGRESSIVE PC select-out
2958            //     ONLY when under-determined (`p > n`, #1392 wine: the outer score
2959            //     is flat on the select-out side and REML needs the active push);
2960            //     otherwise the gentle, wide degeneracy prior (#1476), which adds
2961            //     termination curvature without biasing which λ_null REML lands on
2962            //     — so a genuinely-unsupported null space is still selected out by
2963            //     REML's own score (the sin-data linear trend → λ_null large) and a
2964            //     genuinely-supported one is not over-shrunk.
2965            if is_nullspace {
2966                if underdetermined {
2967                    nullspace_select_prior.clone()
2968                } else {
2969                    nullspace_degeneracy_prior.clone()
2970                }
2971            } else {
2972                relaxed_prior.clone()
2973            }
2974        })
2975        .collect::<Vec<_>>();
2976    gam_spec::RhoPrior::Independent(per_coord)
2977}
2978
2979/// Standard deviation of the wide, weakly-informative symmetric `Normal` prior
2980/// placed on a relaxable smooth's log-λ coordinates when the fit is
2981/// under-determined (`n < 2·p`); see [`relax_smoothing_rho_prior`].
2982///
2983/// Chosen so that ±2σ spans the entire feasible ρ range (the outer optimiser
2984/// bounds `|ρ| ≤ 30`): the prior contributes strictly-positive termination
2985/// curvature `1/sd²` to the outer Hessian (the #1089 requirement that the REML
2986/// loop certify a stationary point on a `p > n` design) while its gradient drag
2987/// at the heavily-smoothed REML optimum is negligible, so pure REML — matching
2988/// mgcv — selects λ. Reducing it toward the old `sd = 3` re-introduces the
2989/// #1392 under-smoothing drag; widening it further weakens termination
2990/// curvature without further benefit.
2991const RELAX_UNDERDETERMINED_RHO_SD: f64 = 15.0;
2992
2993/// Distance-scale bound `upper` (`P(d > upper) = tail_prob` on the marginal-SD
2994/// scale `d = exp(-ρ/2)`) of the penalized-complexity prior placed on a
2995/// relaxable smooth's `DoublePenaltyNullspace` selection coordinate when the fit
2996/// is under-determined (`n < 2·p`); see [`relax_smoothing_rho_prior`].
2997///
2998/// The null-space ridge exists only to SELECT the linear/constant null-space
2999/// component out (mgcv `select=TRUE`): we want its `λ` driven UP (`d → 0`)
3000/// unless the data clearly buys the null-space wiggle. The PC prior is the
3001/// convex bowl `C(ρ) = ρ/2 + θ e^{-ρ/2}` with the steep exponential wall on the
3002/// `λ → 0` (null space kept, `d > upper`) side and a FINITE interior mode at
3003/// `ρ* = 2 ln θ` (`λ* = θ²`). A small `upper` puts that wall close in, so the
3004/// coordinate's λ is selected up toward the well-penalized mode; the data can
3005/// still keep the null space when it genuinely earns it (the over-smoothing side
3006/// of the bowl, gradient `→ +1/2` only in the far tail, pulls ρ back DOWN toward
3007/// λ* — there is no λ → ∞ runaway). `0.05` places the wall at a marginal-SD
3008/// scale two decades below unit, biasing toward select-out on the
3009/// over-parameterized `p > n` wine fit while staying weakly informative.
3010const NULLSPACE_SELECT_PC_UPPER: f64 = 0.05;
3011
3012/// Tail probability `α` (`P(d > upper) = α`) calibrating the rate
3013/// `θ = −ln(α)/upper` of the [`NULLSPACE_SELECT_PC_UPPER`] penalized-complexity
3014/// select-out prior. A small `α` makes the wall against the kept-null-space
3015/// (`λ → 0`) side steep; combined with the small `upper` it yields a strong
3016/// θ ≈ 92 so REML moves the under-determined null-space ridge off its stalled
3017/// λ ≈ 0.11 toward select-out. The PC bowl has a FINITE mode at `λ* = θ² ≈ 8483`
3018/// (`ρ* = 2 ln θ ≈ 9.05`), NOT a hard `λ → ∞` cap: beyond the mode the gradient
3019/// turns positive (approaching `+1/2` only as `ρ → +∞`) and, the objective being
3020/// minimized, pulls ρ back DOWN toward λ*. See [`relax_smoothing_rho_prior`].
3021const NULLSPACE_SELECT_PC_TAIL_PROB: f64 = 0.01;
3022
3023fn adaptive_fit_options_base(options: &FitOptions, design: &TermCollectionDesign) -> FitOptions {
3024    FitOptions {
3025        resource_policy: options.resource_policy.clone(),
3026        latent_cloglog: options.latent_cloglog,
3027        mixture_link: options.mixture_link.clone(),
3028        optimize_mixture: options.optimize_mixture,
3029        sas_link: options.sas_link,
3030        optimize_sas: options.optimize_sas,
3031        compute_inference: options.compute_inference,
3032        skip_rho_posterior_inference: options.skip_rho_posterior_inference,
3033        max_iter: options.max_iter,
3034        tol: options.tol,
3035        nullspace_dims: design.nullspace_dims.clone(),
3036        linear_constraints: design.linear_constraints.clone(),
3037        firth_bias_reduction: options.firth_bias_reduction,
3038        adaptive_regularization: None,
3039        penalty_shrinkage_floor: options.penalty_shrinkage_floor,
3040        // Propagate user-supplied rho_prior so the baseline/refit and the
3041        // joint optimizer minimize the same REML objective.
3042        rho_prior: options.rho_prior.clone(),
3043        kronecker_penalty_system: design.kronecker_penalty_system(),
3044        kronecker_factored: design
3045            .smooth
3046            .terms
3047            .iter()
3048            .find_map(|t| t.kronecker_factored.clone()),
3049        persist_warm_start_disk: options.persist_warm_start_disk,
3050    }
3051}
3052
3053fn superseded_fit_options(options: &FitOptions) -> FitOptions {
3054    let mut fit_options = options.clone();
3055    fit_options.skip_rho_posterior_inference = true;
3056    fit_options
3057}
3058
3059#[derive(Clone)]
3060struct BoundedLinearTermMeta {
3061    col_idx: usize,
3062    min: f64,
3063    max: f64,
3064    prior: BoundedCoefficientPriorSpec,
3065}
3066
3067/// β-dependent effective Jacobian for the bounded-linear fit block.
3068///
3069/// Each bounded coefficient enters the linear predictor non-linearly, as
3070/// `β = min + width·σ(θ)`, and is supplied to the solver through the family
3071/// adapter's offset rather than the linear design. To keep that contribution
3072/// out of the *linear* design the fit places a deliberately **zeroed**
3073/// placeholder column for every bounded term in the block design
3074/// (see `fit_bounded_term_collection_with_design`). The pre-fit
3075/// identifiability audit, however, assesses block rank by reading each block's
3076/// effective Jacobian — and a zeroed column reads as a structural rank
3077/// deficiency, so without this callback the audit refuses *every* bounded
3078/// model before fitting begins.
3079///
3080/// This callback reports the model's true Jacobian column for each bounded
3081/// term, `∂η_i/∂θ = (dβ/dθ)·x_i`, so the audit inspects the same geometry the
3082/// solver actually fits. Because `dβ/dθ = width·σ(θ)(1−σ(θ))` is strictly
3083/// positive for finite θ and `width > 0`, a bounded column is rank-deficient
3084/// in the audit exactly when its underlying covariate is genuinely collinear
3085/// with the rest of the design — never merely because the placeholder was
3086/// zeroed. The callback is consumed only by the identifiability audit /
3087/// canonicalisation; the inner PIRLS solve drives η through the
3088/// [`BoundedLinearFamily`] adapter, so reporting the non-zeroed Jacobian here
3089/// does not double-count the bounded contribution.
3090struct BoundedEffectiveJacobian {
3091    design: Array2<f64>,
3092    bounded_terms: Vec<BoundedLinearTermMeta>,
3093}
3094
3095impl BlockEffectiveJacobian for BoundedEffectiveJacobian {
3096    fn effective_jacobian_rows(
3097        &self,
3098        state: &FamilyLinearizationState<'_>,
3099        rows: std::ops::Range<usize>,
3100    ) -> Result<Array2<f64>, String> {
3101        let p = self.design.ncols();
3102        let n = self.design.nrows();
3103        let rows = rows.start.min(n)..rows.end.min(n);
3104        if !state.beta.is_empty() {
3105            if state.beta.len() != p {
3106                return Err(format!(
3107                    "BoundedEffectiveJacobian::effective_jacobian_at: beta length {} != design \
3108                     ncols {p}",
3109                    state.beta.len(),
3110                ));
3111            }
3112            if state.beta.iter().any(|v| v.is_nan()) {
3113                return Err(
3114                    "BoundedEffectiveJacobian::effective_jacobian_at: beta contains NaN"
3115                        .to_string(),
3116                );
3117            }
3118        }
3119        let mut jac = self
3120            .design
3121            .slice(ndarray::s![rows.start..rows.end, ..])
3122            .to_owned();
3123        for term in &self.bounded_terms {
3124            let theta = if state.beta.is_empty() {
3125                0.0
3126            } else {
3127                state.beta[term.col_idx]
3128            };
3129            let (_, _, db_dtheta, _, _) = bounded_latent_derivatives(theta, term.min, term.max);
3130            jac.column_mut(term.col_idx).mapv_inplace(|v| v * db_dtheta);
3131        }
3132        Ok(jac)
3133    }
3134}
3135
3136#[derive(Clone)]
3137struct BoundedLinearFamily {
3138    family: LikelihoodSpec,
3139    latent_cloglog_state: Option<LatentCLogLogState>,
3140    mixture_link_state: Option<MixtureLinkState>,
3141    sas_link_state: Option<SasLinkState>,
3142    y: Array1<f64>,
3143    weights: Array1<f64>,
3144    design: Array2<f64>,
3145    designzeroed: Array2<f64>,
3146    offset: Array1<f64>,
3147    bounded_terms: Vec<BoundedLinearTermMeta>,
3148}
3149
3150#[derive(Clone)]
3151struct StandardFamilyObservationState {
3152    eta: Array1<f64>,
3153    mu: Array1<f64>,
3154    score: Array1<f64>,
3155    fisherweight: Array1<f64>,
3156    neghessian_eta: Array1<f64>,
3157    neghessian_eta_derivative: Array1<f64>,
3158    log_likelihood: f64,
3159}
3160
3161fn bounded_logit(z: f64) -> f64 {
3162    let zc = z.clamp(1e-12, 1.0 - 1e-12);
3163    (zc / (1.0 - zc)).ln()
3164}
3165
3166fn stable_sigmoid(theta: f64) -> f64 {
3167    if theta >= 0.0 {
3168        let exp_neg = (-theta).exp();
3169        1.0 / (1.0 + exp_neg)
3170    } else {
3171        let exp_pos = theta.exp();
3172        exp_pos / (1.0 + exp_pos)
3173    }
3174}
3175
3176fn bounded_latent_to_user(theta: f64, min: f64, max: f64) -> (f64, f64, f64) {
3177    let z = stable_sigmoid(theta);
3178    let width = max - min;
3179    let beta = min + width * z;
3180    let db_dtheta = width * z * (1.0 - z);
3181    (beta, z, db_dtheta)
3182}
3183
3184/// Invert the bounded interval transform: given a user-scale coefficient
3185/// `beta` in the open interval `(min, max)`, return the latent coordinate
3186/// `theta` with `bounded_latent_to_user(theta, min, max).0 == beta`.
3187///
3188/// This is the exact inverse of the logistic interval map used by the bounded
3189/// custom family: `z = (beta - min)/(max - min)` (the normalized position in
3190/// the interval) and `theta = logit(z)`. The normalized position is clamped
3191/// strictly inside `(0, 1)` (mirroring `bounded_logit`) so a coefficient that
3192/// sits numerically at a boundary maps to a large-but-finite latent value
3193/// rather than `±∞`.
3194fn bounded_user_to_latent(beta: f64, min: f64, max: f64) -> f64 {
3195    let width = max - min;
3196    if width <= 0.0 || !width.is_finite() {
3197        return 0.0;
3198    }
3199    let z = (beta - min) / width;
3200    bounded_logit(z)
3201}
3202
3203/// One bounded coefficient column for posterior sampling: its position in the
3204/// (internal, conditioned) coefficient vector and the interval bounds expressed
3205/// on that same internal scale.
3206#[derive(Debug, Clone, Copy)]
3207pub struct BoundedSampleColumn {
3208    /// Column index into the internal (conditioned) coefficient vector.
3209    pub col_idx: usize,
3210    /// Lower interval bound on the internal scale.
3211    pub min: f64,
3212    /// Upper interval bound on the internal scale.
3213    pub max: f64,
3214}
3215
3216/// Exact posterior draws for a model with `bounded()` coefficients.
3217///
3218/// The bounded custom family fits each bounded coefficient as a smooth interval
3219/// transform `beta = min + (max - min)·sigmoid(theta)` of an unconstrained
3220/// latent `theta`. The Laplace approximation is *Gaussian on the latent scale*
3221/// — that is precisely the scale on which the fit treats the coefficient as an
3222/// unconstrained, locally-quadratic parameter. Sampling a Gaussian directly on
3223/// the user (bounded) scale is wrong twice over: it can place mass outside
3224/// `[min, max]`, and it discards the boundary-induced skew that the nonlinear
3225/// map produces. This routine instead draws `theta ~ N(theta_mode, H_latent^{-1})`
3226/// and pushes every draw through the *exact* interval map, so user-scale draws
3227/// always lie strictly inside the interval and carry the correct skew.
3228///
3229/// Coordinate bookkeeping. The caller supplies the user-scale mode `beta_user`
3230/// and the user-scale penalized Hessian `user_hessian` (both in *internal /
3231/// conditioned* coordinates — i.e. before `backtransform_*` to the original
3232/// data scale) together with the internal-scale bounds for each bounded column.
3233/// The user-scale Hessian relates to the latent-scale Hessian by the diagonal
3234/// delta-method Jacobian `J = diag(db/dtheta)`:
3235///   `H_user = J^{-1} H_latent J^{-1}`  ⇒  `H_latent = J H_user J`,
3236/// which is exactly the inverse of `transform_bounded_latent_precision_to_user_internal`.
3237/// Non-bounded columns have `J_ii = 1`, so they are sampled as the ordinary
3238/// Gaussian Laplace draw and returned unchanged.
3239///
3240/// Dispersion. `user_hessian` is the UNSCALED penalized Hessian `H_user`
3241/// (unit implicit dispersion). For a free-dispersion family the latent
3242/// posterior covariance is `φ̂·H_latent⁻¹`, so the caller passes
3243/// `sqrt_cov_scale = √φ̂` (the coefficient-covariance scale `√σ̂²` for a
3244/// profiled Gaussian, `1` for fixed-scale families like Binomial) and every
3245/// latent perturbation is multiplied by it. This makes the draw covariance
3246/// `sqrt_cov_scale² · H_latent⁻¹`, matching the fit's reported
3247/// `Vb = cov_scale·H_user⁻¹` exactly (gam#1514) — without it a Gaussian
3248/// bounded slope's draws were ~`1/σ̂` too wide.
3249///
3250/// Returns the draws as a `(n_draws, p)` matrix on the *internal* user scale
3251/// (still conditioned); the caller back-transforms to the original data scale
3252/// with the same conditioning it used for the point estimate.
3253pub fn sample_bounded_latent_posterior_internal(
3254    beta_user: &Array1<f64>,
3255    user_hessian: &Array2<f64>,
3256    bounded_columns: &[BoundedSampleColumn],
3257    n_draws: usize,
3258    sqrt_cov_scale: f64,
3259    base_seed: u64,
3260) -> Result<Array2<f64>, EstimationError> {
3261    let p = beta_user.len();
3262    if user_hessian.nrows() != p || user_hessian.ncols() != p {
3263        crate::bail_invalid_estim!(
3264            "bounded posterior sampling dimension mismatch: mode has {p} entries, user Hessian is {}x{}",
3265            user_hessian.nrows(),
3266            user_hessian.ncols()
3267        );
3268    }
3269
3270    // Latent mode and delta-method Jacobian, column by column.
3271    let mut theta_mode = beta_user.clone();
3272    let mut jac_diag = Array1::<f64>::ones(p);
3273    for bc in bounded_columns {
3274        if bc.col_idx >= p {
3275            crate::bail_invalid_estim!(
3276                "bounded posterior sampling: bounded column index {} out of range for {p} coefficients",
3277                bc.col_idx
3278            );
3279        }
3280        let theta_i = bounded_user_to_latent(beta_user[bc.col_idx], bc.min, bc.max);
3281        let (_, _, db_dtheta) = bounded_latent_to_user(theta_i, bc.min, bc.max);
3282        theta_mode[bc.col_idx] = theta_i;
3283        // Guard against a degenerate (numerically vanishing) Jacobian at a
3284        // coefficient pinned hard against a boundary: floor the slope so the
3285        // latent precision stays finite and the draw simply collapses onto the
3286        // boundary, which is the correct limiting posterior.
3287        jac_diag[bc.col_idx] = db_dtheta.max(1e-12);
3288    }
3289
3290    // H_latent = J H_user J  (J diagonal). This is the exact inverse of the
3291    // user-scale precision transform applied at fit time.
3292    let mut h_latent = user_hessian.clone();
3293    for i in 0..p {
3294        let ji = jac_diag[i];
3295        if ji != 1.0 {
3296            h_latent.row_mut(i).mapv_inplace(|v| v * ji);
3297            h_latent.column_mut(i).mapv_inplace(|v| v * ji);
3298        }
3299    }
3300
3301    // Draw theta ~ N(theta_mode, H_latent^{-1}) via the Cholesky of H_latent:
3302    // L Lᵀ = H_latent, solve Lᵀ δ = ε so Var(δ) = H_latent^{-1}.
3303    use gam_linalg::faer_ndarray::FaerCholesky as _;
3304    use rand::SeedableRng as _;
3305    let chol = h_latent.cholesky(faer::Side::Lower).map_err(|err| {
3306        EstimationError::InvalidInput(format!(
3307            "bounded posterior sampling: Cholesky of the latent penalized Hessian failed: {err:?}"
3308        ))
3309    })?;
3310    let l = chol.lower_triangular();
3311
3312    let mut draws = Array2::<f64>::zeros((n_draws, p));
3313    let mut eps = Array1::<f64>::zeros(p);
3314    let mut delta = Array1::<f64>::zeros(p);
3315    let mut rng = rand::rngs::StdRng::seed_from_u64(base_seed);
3316    for k in 0..n_draws {
3317        for e in eps.iter_mut() {
3318            *e = standard_normal_draw(&mut rng);
3319        }
3320        solve_lower_transpose_into(&l, &eps, &mut delta);
3321        for i in 0..p {
3322            // δ has covariance `H_latent⁻¹`; scaling by √cov_scale lifts it to
3323            // the dispersion-correct posterior covariance `cov_scale·H_latent⁻¹`.
3324            draws[(k, i)] = theta_mode[i] + sqrt_cov_scale * delta[i];
3325        }
3326        // Push bounded columns through the exact interval map so every draw is
3327        // strictly inside (min, max); leave unconstrained columns untouched.
3328        for bc in bounded_columns {
3329            let (beta_draw, _, _) = bounded_latent_to_user(draws[(k, bc.col_idx)], bc.min, bc.max);
3330            draws[(k, bc.col_idx)] = beta_draw;
3331        }
3332    }
3333
3334    Ok(draws)
3335}
3336
3337/// Box-Muller standard-normal draw (kept local so the bounded sampler does not
3338/// depend on the HMC module's RNG plumbing).
3339#[inline]
3340fn standard_normal_draw<R: rand::Rng + ?Sized>(rng: &mut R) -> f64 {
3341    use rand::RngExt as _;
3342    let u1 = rng.random::<f64>().max(1e-16);
3343    let u2 = rng.random::<f64>();
3344    (-2.0 * u1.ln()).sqrt() * (2.0 * std::f64::consts::PI * u2).cos()
3345}
3346
3347/// Solve `Lᵀ x = b` for a lower-triangular `L` (back substitution), writing the
3348/// result into `out`. Used to turn a standard-normal `b` into a draw with
3349/// covariance `(L Lᵀ)^{-1}`.
3350fn solve_lower_transpose_into(l: &Array2<f64>, b: &Array1<f64>, out: &mut Array1<f64>) {
3351    let p = l.nrows();
3352    for i in (0..p).rev() {
3353        let mut acc = b[i];
3354        for j in (i + 1)..p {
3355            acc -= l[(j, i)] * out[j];
3356        }
3357        let diag = l[(i, i)];
3358        out[i] = if diag.abs() > 0.0 { acc / diag } else { 0.0 };
3359    }
3360}
3361
3362fn bounded_latent_derivatives(theta: f64, min: f64, max: f64) -> (f64, f64, f64, f64, f64) {
3363    let z = stable_sigmoid(theta);
3364    let width = max - min;
3365    let s = z * (1.0 - z);
3366    let beta = min + width * z;
3367    let db_dtheta = width * s;
3368    let d2b_dtheta2 = width * s * (1.0 - 2.0 * z);
3369    let d3b_dtheta3 = width * s * (1.0 - 6.0 * z + 6.0 * z * z);
3370    (beta, z, db_dtheta, d2b_dtheta2, d3b_dtheta3)
3371}
3372
3373fn bounded_prior_terms(theta: f64, prior: &BoundedCoefficientPriorSpec) -> (f64, f64, f64, f64) {
3374    let (a, b) = match prior {
3375        // `None` means constrained MLE with no extra prior term on the bounded coefficient.
3376        BoundedCoefficientPriorSpec::None => return (0.0, 0.0, 0.0, 0.0),
3377        // Uniform on the normalized user-scale coefficient z in (0, 1). In latent space this is
3378        // exactly the Jacobian term for the logistic transform, up to an additive width constant.
3379        BoundedCoefficientPriorSpec::Uniform => (1.0, 1.0),
3380        BoundedCoefficientPriorSpec::Beta { a, b } => (*a, *b),
3381    };
3382    let z = stable_sigmoid(theta).clamp(1e-12, 1.0 - 1e-12);
3383    let logp = a * z.ln() + b * (1.0 - z).ln();
3384    let grad = a - (a + b) * z;
3385    let neghess = (a + b) * z * (1.0 - z);
3386    let neghess_derivative = (a + b) * z * (1.0 - z) * (1.0 - 2.0 * z);
3387    (logp, grad, neghess, neghess_derivative)
3388}
3389
3390/// Assemble the per-observation IRLS state for a standard exponential-family
3391/// GLM from the inverse-link jet (`d1 = dμ/dη`, `d2`, `d3`) and the
3392/// log-likelihood derivatives in `μ` (`lmu = ∂l/∂μ`, `lmumu`, `lmumumu`) plus
3393/// the mean-variance `var = V(μ)`. Mirrors the Binomial general arm EXACTLY
3394/// (#1615/#1616): expected Fisher weight `(dμ/dη)²/V(μ)`, observed
3395/// negative-Hessian in `η`, and its `η`-derivative. Canonical-link self-check:
3396/// for a canonical link the observed neg-Hessian equals the Fisher weight
3397/// (e.g. Poisson-log → both `= w·μ`), which this construction reproduces.
3398#[inline]
3399fn glm_eta_observation_state(
3400    w: f64,
3401    lmu: f64,
3402    lmumu: f64,
3403    lmumumu: f64,
3404    var: f64,
3405    d1: f64,
3406    d2: f64,
3407    d3: f64,
3408    mu_deriv_eps: f64,
3409) -> (f64, f64, f64, f64) {
3410    let score = w * lmu * d1;
3411    let fisherweight = (w * d1 * d1 / var).max(mu_deriv_eps);
3412    let neghessian = -w * (lmumu * d1 * d1 + lmu * d2);
3413    let neghessian_deriv = -w * (lmumumu * d1 * d1 * d1 + 3.0 * lmumu * d1 * d2 + lmu * d3);
3414    (score, fisherweight, neghessian, neghessian_deriv)
3415}
3416
3417fn evaluate_standard_familyobservations(
3418    family: LikelihoodSpec,
3419    latent_cloglog_state: Option<&LatentCLogLogState>,
3420    mixture_link_state: Option<&MixtureLinkState>,
3421    sas_link_state: Option<&SasLinkState>,
3422    y: &Array1<f64>,
3423    weights: &Array1<f64>,
3424    eta: &Array1<f64>,
3425) -> Result<StandardFamilyObservationState, EstimationError> {
3426    const PROB_EPS: f64 = 1e-10;
3427    const MU_DERIV_EPS: f64 = 1e-12;
3428    let n = y.len();
3429    if weights.len() != n || eta.len() != n {
3430        crate::bail_invalid_estim!("bounded family observation size mismatch");
3431    }
3432
3433    let mut mu = Array1::<f64>::zeros(n);
3434    let mut score = Array1::<f64>::zeros(n);
3435    let mut fisherweight = Array1::<f64>::zeros(n);
3436    let mut neghessian_eta = Array1::<f64>::zeros(n);
3437    let mut neghessian_eta_derivative = Array1::<f64>::zeros(n);
3438    let mut log_likelihood = 0.0;
3439
3440    for i in 0..n {
3441        let w = weights[i].max(0.0);
3442        let yi = y[i];
3443        let eta_i = eta[i];
3444        match (&family.response, &family.link) {
3445            (ResponseFamily::Gaussian, _) => {
3446                let resid = yi - eta_i;
3447                mu[i] = eta_i;
3448                score[i] = w * resid;
3449                fisherweight[i] = w.max(MU_DERIV_EPS);
3450                neghessian_eta[i] = w;
3451                neghessian_eta_derivative[i] = 0.0;
3452                log_likelihood += -0.5 * w * resid * resid;
3453            }
3454            (ResponseFamily::Binomial, InverseLink::Standard(StandardLink::Logit)) => {
3455                let jet = logit_inverse_link_jet5(eta_i);
3456                mu[i] = jet.mu;
3457                score[i] = w * (yi - jet.mu);
3458                fisherweight[i] = jet.d1.max(MU_DERIV_EPS);
3459                neghessian_eta[i] = jet.d1;
3460                neghessian_eta_derivative[i] = jet.d2;
3461                let logmu = -gam_linalg::utils::stable_softplus(-eta_i);
3462                let log_one_minusmu = -gam_linalg::utils::stable_softplus(eta_i);
3463                log_likelihood += w * (yi * logmu + (1.0 - yi) * log_one_minusmu);
3464            }
3465            (ResponseFamily::Binomial, _) => {
3466                let inverse_link = if let Some(state) = latent_cloglog_state {
3467                    Some(InverseLink::LatentCLogLog(*state))
3468                } else if let Some(state) = mixture_link_state {
3469                    Some(InverseLink::Mixture(state.clone()))
3470                } else {
3471                    sas_link_state.map(|state| {
3472                        if family.is_binomial_beta_logistic() {
3473                            InverseLink::BetaLogistic(*state)
3474                        } else {
3475                            InverseLink::Sas(*state)
3476                        }
3477                    })
3478                };
3479                let strategy_spec = LikelihoodSpec {
3480                    response: family.response.clone(),
3481                    link: inverse_link.clone().unwrap_or_else(|| family.link.clone()),
3482                };
3483                let jet = strategy_for_spec(&strategy_spec).inverse_link_jet(eta_i)?;
3484                let mu_i_raw = jet.mu;
3485                let dmu_deta_raw = jet.d1;
3486                let mu_i: f64 = mu_i_raw.clamp(PROB_EPS, 1.0 - PROB_EPS);
3487                let dmu_deta = dmu_deta_raw.max(MU_DERIV_EPS);
3488                let d2mu_deta2 = jet.d2;
3489                let d3mu_deta3 = jet.d3;
3490                let var = (mu_i * (1.0 - mu_i)).max(PROB_EPS);
3491                let lmu = (yi - mu_i) / var;
3492                let lmumu = -(yi / (mu_i * mu_i)) - ((1.0 - yi) / ((1.0 - mu_i) * (1.0 - mu_i)));
3493                let lmumumu = 2.0 * yi / (mu_i * mu_i * mu_i)
3494                    - 2.0 * (1.0 - yi) / ((1.0 - mu_i) * (1.0 - mu_i) * (1.0 - mu_i));
3495                mu[i] = mu_i;
3496                score[i] = w * lmu * dmu_deta;
3497                fisherweight[i] = (w * dmu_deta * dmu_deta / var).max(MU_DERIV_EPS);
3498                neghessian_eta[i] = -w * (lmumu * dmu_deta * dmu_deta + lmu * d2mu_deta2);
3499                neghessian_eta_derivative[i] = -w
3500                    * (lmumumu * dmu_deta * dmu_deta * dmu_deta
3501                        + 3.0 * lmumu * dmu_deta * d2mu_deta2
3502                        + lmu * d3mu_deta3);
3503                log_likelihood += w * (yi * mu_i.ln() + (1.0 - yi) * (1.0 - mu_i).ln());
3504            }
3505            (ResponseFamily::Poisson, _) => {
3506                // l(μ) = y·ln μ − μ (drop ln y! const) ⇒ lmu = y/μ − 1,
3507                // lmumu = −y/μ², lmumumu = 2y/μ³, V(μ) = μ. #1615/#1616.
3508                let strategy_spec = LikelihoodSpec {
3509                    response: family.response.clone(),
3510                    link: family.link.clone(),
3511                };
3512                let jet = strategy_for_spec(&strategy_spec).inverse_link_jet(eta_i)?;
3513                let mu_i = jet.mu.max(PROB_EPS);
3514                let d1 = jet.d1.max(MU_DERIV_EPS);
3515                let var = mu_i;
3516                let lmu = yi / mu_i - 1.0;
3517                let lmumu = -yi / (mu_i * mu_i);
3518                let lmumumu = 2.0 * yi / (mu_i * mu_i * mu_i);
3519                let (s, f, nh, nhd) = glm_eta_observation_state(
3520                    w, lmu, lmumu, lmumumu, var, d1, jet.d2, jet.d3, MU_DERIV_EPS,
3521                );
3522                mu[i] = mu_i;
3523                score[i] = s;
3524                fisherweight[i] = f;
3525                neghessian_eta[i] = nh;
3526                neghessian_eta_derivative[i] = nhd;
3527                log_likelihood += w * (yi * mu_i.ln() - mu_i);
3528            }
3529            (ResponseFamily::Tweedie { p }, _) => {
3530                // Fixed-p Tweedie quasi-likelihood: ∂Q/∂μ = (y−μ)/V, V(μ) = μ^p.
3531                // lmu = (y−μ)μ^{−p};  lmumu = −μ^{−p} − p(y−μ)μ^{−p−1};
3532                // lmumumu = 2p·μ^{−p−1} + p(p+1)(y−μ)μ^{−p−2};
3533                // Q = y·μ^{1−p}/(1−p) − μ^{2−p}/(2−p). #1615/#1616.
3534                let p = *p;
3535                let strategy_spec = LikelihoodSpec {
3536                    response: family.response.clone(),
3537                    link: family.link.clone(),
3538                };
3539                let jet = strategy_for_spec(&strategy_spec).inverse_link_jet(eta_i)?;
3540                let mu_i = jet.mu.max(PROB_EPS);
3541                let d1 = jet.d1.max(MU_DERIV_EPS);
3542                let var = mu_i.powf(p);
3543                let resid = yi - mu_i;
3544                let lmu = resid / var;
3545                let lmumu = -mu_i.powf(-p) - p * resid * mu_i.powf(-p - 1.0);
3546                let lmumumu =
3547                    2.0 * p * mu_i.powf(-p - 1.0) + p * (p + 1.0) * resid * mu_i.powf(-p - 2.0);
3548                let (s, f, nh, nhd) = glm_eta_observation_state(
3549                    w, lmu, lmumu, lmumumu, var, d1, jet.d2, jet.d3, MU_DERIV_EPS,
3550                );
3551                mu[i] = mu_i;
3552                score[i] = s;
3553                fisherweight[i] = f;
3554                neghessian_eta[i] = nh;
3555                neghessian_eta_derivative[i] = nhd;
3556                // Quasi-log-likelihood (p ≠ 1, 2 in the supported compound range).
3557                log_likelihood += w
3558                    * (yi * mu_i.powf(1.0 - p) / (1.0 - p) - mu_i.powf(2.0 - p) / (2.0 - p));
3559            }
3560            (ResponseFamily::NegativeBinomial { theta, .. }, _) => {
3561                // l(μ) = y·ln μ − (y+θ)·ln(μ+θ) (drop μ-independent terms) ⇒
3562                // lmu = y/μ − (y+θ)/(μ+θ), lmumu = −y/μ² + (y+θ)/(μ+θ)²,
3563                // lmumumu = 2y/μ³ − 2(y+θ)/(μ+θ)³, V(μ) = μ + μ²/θ. #1615/#1616.
3564                let theta = (*theta).max(PROB_EPS);
3565                let strategy_spec = LikelihoodSpec {
3566                    response: family.response.clone(),
3567                    link: family.link.clone(),
3568                };
3569                let jet = strategy_for_spec(&strategy_spec).inverse_link_jet(eta_i)?;
3570                let mu_i = jet.mu.max(PROB_EPS);
3571                let d1 = jet.d1.max(MU_DERIV_EPS);
3572                let mu_plus = mu_i + theta;
3573                let var = mu_i + mu_i * mu_i / theta;
3574                let lmu = yi / mu_i - (yi + theta) / mu_plus;
3575                let lmumu = -yi / (mu_i * mu_i) + (yi + theta) / (mu_plus * mu_plus);
3576                let lmumumu =
3577                    2.0 * yi / (mu_i * mu_i * mu_i) - 2.0 * (yi + theta) / (mu_plus * mu_plus * mu_plus);
3578                let (s, f, nh, nhd) = glm_eta_observation_state(
3579                    w, lmu, lmumu, lmumumu, var, d1, jet.d2, jet.d3, MU_DERIV_EPS,
3580                );
3581                mu[i] = mu_i;
3582                score[i] = s;
3583                fisherweight[i] = f;
3584                neghessian_eta[i] = nh;
3585                neghessian_eta_derivative[i] = nhd;
3586                log_likelihood += w * (yi * mu_i.ln() - (yi + theta) * mu_plus.ln());
3587            }
3588            (ResponseFamily::Beta { .. }, _) => {
3589                crate::bail_invalid_estim!(
3590                    "bounded linear terms are not supported for BetaLogit fits"
3591                );
3592            }
3593            (ResponseFamily::Gamma, _) => {
3594                // Unit-dispersion Gamma kernel l(μ) = −y/μ − ln μ ⇒
3595                // lmu = y/μ² − 1/μ, lmumu = −2y/μ³ + 1/μ²,
3596                // lmumumu = 6y/μ⁴ − 2/μ³, V(μ) = μ². #1615/#1616.
3597                let strategy_spec = LikelihoodSpec {
3598                    response: family.response.clone(),
3599                    link: family.link.clone(),
3600                };
3601                let jet = strategy_for_spec(&strategy_spec).inverse_link_jet(eta_i)?;
3602                let mu_i = jet.mu.max(PROB_EPS);
3603                let d1 = jet.d1.max(MU_DERIV_EPS);
3604                let var = mu_i * mu_i;
3605                let lmu = yi / (mu_i * mu_i) - 1.0 / mu_i;
3606                let lmumu = -2.0 * yi / (mu_i * mu_i * mu_i) + 1.0 / (mu_i * mu_i);
3607                let lmumumu =
3608                    6.0 * yi / (mu_i * mu_i * mu_i * mu_i) - 2.0 / (mu_i * mu_i * mu_i);
3609                let (s, f, nh, nhd) = glm_eta_observation_state(
3610                    w, lmu, lmumu, lmumumu, var, d1, jet.d2, jet.d3, MU_DERIV_EPS,
3611                );
3612                mu[i] = mu_i;
3613                score[i] = s;
3614                fisherweight[i] = f;
3615                neghessian_eta[i] = nh;
3616                neghessian_eta_derivative[i] = nhd;
3617                log_likelihood += w * (-(yi / mu_i) - mu_i.ln());
3618            }
3619            (ResponseFamily::RoystonParmar, _) => {
3620                crate::bail_invalid_estim!(
3621                    "bounded linear terms are not supported for survival model fits"
3622                );
3623            }
3624        }
3625    }
3626
3627    Ok(StandardFamilyObservationState {
3628        eta: eta.clone(),
3629        mu,
3630        score,
3631        fisherweight,
3632        neghessian_eta,
3633        neghessian_eta_derivative,
3634        log_likelihood,
3635    })
3636}
3637
3638#[derive(Clone, Copy, Debug, PartialEq, Eq)]
3639enum SpatialAdaptiveHyperKind {
3640    LogLambdaMagnitude,
3641    LogLambdaGradient,
3642    LogLambdaCurvature,
3643    LogEpsilonMagnitude,
3644    LogEpsilonGradient,
3645    LogEpsilonCurvature,
3646}
3647
3648impl SpatialAdaptiveHyperKind {
3649    fn component_index(self) -> usize {
3650        match self {
3651            SpatialAdaptiveHyperKind::LogLambdaMagnitude
3652            | SpatialAdaptiveHyperKind::LogEpsilonMagnitude => 0,
3653            SpatialAdaptiveHyperKind::LogLambdaGradient
3654            | SpatialAdaptiveHyperKind::LogEpsilonGradient => 1,
3655            SpatialAdaptiveHyperKind::LogLambdaCurvature
3656            | SpatialAdaptiveHyperKind::LogEpsilonCurvature => 2,
3657        }
3658    }
3659
3660    fn is_log_lambda(self) -> bool {
3661        matches!(
3662            self,
3663            SpatialAdaptiveHyperKind::LogLambdaMagnitude
3664                | SpatialAdaptiveHyperKind::LogLambdaGradient
3665                | SpatialAdaptiveHyperKind::LogLambdaCurvature
3666        )
3667    }
3668
3669    fn is_log_epsilon(self) -> bool {
3670        matches!(
3671            self,
3672            SpatialAdaptiveHyperKind::LogEpsilonMagnitude
3673                | SpatialAdaptiveHyperKind::LogEpsilonGradient
3674                | SpatialAdaptiveHyperKind::LogEpsilonCurvature
3675        )
3676    }
3677}
3678
3679#[derive(Clone, Copy, Debug)]
3680struct SpatialAdaptiveHyperSpec {
3681    cache_index: usize,
3682    kind: SpatialAdaptiveHyperKind,
3683}
3684
3685#[derive(Clone, Copy, Debug, PartialEq, Eq)]
3686enum SpatialAdaptiveExplicitSecondOrderKind {
3687    StructuralZero,
3688    LocalAlphaAlpha,
3689    LocalAlphaEta,
3690    SharedEtaEta,
3691}
3692
3693/// Penalty family selected within one adaptive smooth cache. The component index
3694/// (0/1/2) used throughout the runtime caches maps onto these three operators:
3695/// the scalar magnitude operator `d0`, the grouped gradient operator `d1`, and
3696/// the grouped curvature operator `d2`.
3697#[derive(Clone, Copy, Debug, PartialEq, Eq)]
3698enum AdaptiveComponent {
3699    Magnitude,
3700    Gradient,
3701    Curvature,
3702}
3703
3704impl AdaptiveComponent {
3705    fn from_index(index: usize) -> Result<Self, String> {
3706        match index {
3707            0 => Ok(AdaptiveComponent::Magnitude),
3708            1 => Ok(AdaptiveComponent::Gradient),
3709            2 => Ok(AdaptiveComponent::Curvature),
3710            other => Err(SmoothError::invalid_index(format!(
3711                "invalid adaptive component index {}",
3712                other
3713            ))
3714            .into()),
3715        }
3716    }
3717}
3718
3719/// Which hyper-derivative of the adaptive penalty's local pieces to assemble.
3720/// Each variant selects one accessor triple (objective scalar, beta-mixed
3721/// gradient, beta hessian) on the per-component exact state; the operator
3722/// embedding around those accessors is identical across variants.
3723#[derive(Clone, Copy, Debug, PartialEq, Eq)]
3724enum HyperDerivativeKind {
3725    /// First derivative in `log lambda` (rho): the bare penalty pieces.
3726    Rho,
3727    /// First derivative in `log epsilon`.
3728    LogEpsilonFirst,
3729    /// Second derivative in `log epsilon`.
3730    LogEpsilonSecond,
3731}
3732
3733/// Which directional-drift hyper-derivative of the adaptive penalty Hessian to
3734/// assemble: the bare rho drift, or the shared-`log epsilon` drift. Both share
3735/// the per-component direction projection, operator embedding, and global
3736/// embedding; only the directional state accessor differs.
3737#[derive(Clone, Copy, Debug, PartialEq, Eq)]
3738enum HyperDriftKind {
3739    Rho,
3740    LogEpsilon,
3741}
3742
3743impl SpatialAdaptiveHyperSpec {
3744    fn component_index(self) -> usize {
3745        self.kind.component_index()
3746    }
3747
3748    fn explicit_second_order_kind(self, other: Self) -> SpatialAdaptiveExplicitSecondOrderKind {
3749        if self.component_index() != other.component_index() {
3750            return SpatialAdaptiveExplicitSecondOrderKind::StructuralZero;
3751        }
3752        match (
3753            self.kind.is_log_lambda(),
3754            other.kind.is_log_lambda(),
3755            self.kind.is_log_epsilon(),
3756            other.kind.is_log_epsilon(),
3757        ) {
3758            (true, true, false, false) if self.cache_index == other.cache_index => {
3759                SpatialAdaptiveExplicitSecondOrderKind::LocalAlphaAlpha
3760            }
3761            (true, false, false, true) | (false, true, true, false) => {
3762                SpatialAdaptiveExplicitSecondOrderKind::LocalAlphaEta
3763            }
3764            (false, false, true, true) => SpatialAdaptiveExplicitSecondOrderKind::SharedEtaEta,
3765            _ => SpatialAdaptiveExplicitSecondOrderKind::StructuralZero,
3766        }
3767    }
3768}
3769
3770#[derive(Clone, Debug)]
3771struct SpatialAdaptiveTermHyperParams {
3772    lambda: [f64; 3],
3773    epsilon: [f64; 3],
3774}
3775
3776#[derive(Clone)]
3777struct SpatialAdaptiveExactEvaluation {
3778    obs: StandardFamilyObservationState,
3779    adaptive_states: Vec<SpatialPenaltyExactState>,
3780    adaptive_penalty_value: f64,
3781    adaptive_penaltygradient: Array1<f64>,
3782    adaptive_penaltyhessian: Array2<f64>,
3783    fixed_quadraticvalue: f64,
3784    fixed_quadraticgradient: Array1<f64>,
3785    fixed_quadratichessian: Array2<f64>,
3786}
3787
3788#[derive(Clone)]
3789struct CachedSpatialAdaptiveExactEvaluation {
3790    beta: Array1<f64>,
3791    eval: Arc<SpatialAdaptiveExactEvaluation>,
3792}
3793
3794impl SpatialAdaptiveExactEvaluation {
3795    fn total_penalty_value(&self) -> f64 {
3796        self.adaptive_penalty_value + self.fixed_quadraticvalue
3797    }
3798
3799    fn total_penaltygradient(&self) -> Array1<f64> {
3800        &self.adaptive_penaltygradient + &self.fixed_quadraticgradient
3801    }
3802
3803    fn total_penaltyhessian(&self) -> Array2<f64> {
3804        &self.adaptive_penaltyhessian + &self.fixed_quadratichessian
3805    }
3806
3807    fn totalobjectivehessian(&self, design: &Array2<f64>) -> Result<Array2<f64>, String> {
3808        let mut out = xt_diag_x_dense(design.view(), self.obs.neghessian_eta.view())?;
3809        out += &self.total_penaltyhessian();
3810        Ok(out)
3811    }
3812}
3813
3814#[derive(Clone)]
3815struct SpatialAdaptiveExactFamily {
3816    family: LikelihoodSpec,
3817    latent_cloglog_state: Option<LatentCLogLogState>,
3818    mixture_link_state: Option<MixtureLinkState>,
3819    sas_link_state: Option<SasLinkState>,
3820    y: Arc<Array1<f64>>,
3821    weights: Arc<Array1<f64>>,
3822    design: Arc<Array2<f64>>,
3823    offset: Arc<Array1<f64>>,
3824    linear_constraints: Option<LinearInequalityConstraints>,
3825    runtime_caches: Arc<Vec<SpatialOperatorRuntimeCache>>,
3826    adaptive_params: Vec<SpatialAdaptiveTermHyperParams>,
3827    fixed_quadratichessian: Arc<Array2<f64>>,
3828    hyperspecs: Arc<Vec<SpatialAdaptiveHyperSpec>>,
3829    exact_eval_cache: Arc<Mutex<Option<CachedSpatialAdaptiveExactEvaluation>>>,
3830}
3831
3832impl SpatialAdaptiveExactFamily {
3833    fn with_adaptive_params(
3834        &self,
3835        adaptive_params: Vec<SpatialAdaptiveTermHyperParams>,
3836        fixed_quadratichessian: Arc<Array2<f64>>,
3837    ) -> Self {
3838        Self {
3839            family: self.family.clone(),
3840            latent_cloglog_state: self.latent_cloglog_state,
3841            mixture_link_state: self.mixture_link_state.clone(),
3842            sas_link_state: self.sas_link_state,
3843            y: self.y.clone(),
3844            weights: self.weights.clone(),
3845            design: self.design.clone(),
3846            offset: self.offset.clone(),
3847            linear_constraints: self.linear_constraints.clone(),
3848            runtime_caches: self.runtime_caches.clone(),
3849            adaptive_params,
3850            fixed_quadratichessian,
3851            hyperspecs: self.hyperspecs.clone(),
3852            exact_eval_cache: Arc::new(Mutex::new(None)),
3853        }
3854    }
3855
3856    fn total_eta(&self, beta: &Array1<f64>) -> Array1<f64> {
3857        gam_linalg::faer_ndarray::fast_av(self.design.as_ref(), beta) + self.offset.as_ref()
3858    }
3859
3860    fn fixed_quadratic_terms(&self, beta: &Array1<f64>) -> (f64, Array1<f64>) {
3861        let grad = self.fixed_quadratichessian.dot(beta);
3862        let value = 0.5 * beta.dot(&grad);
3863        (value, grad)
3864    }
3865
3866    fn adaptive_penalty_value_only(&self, beta: &Array1<f64>) -> Result<f64, String> {
3867        let mut penalty_value = 0.0;
3868        for (cache_idx, cache) in self.runtime_caches.iter().enumerate() {
3869            let params = self.adaptive_params.get(cache_idx).ok_or_else(|| {
3870                format!(
3871                    "missing adaptive parameter block for cache {}",
3872                    cache.termname
3873                )
3874            })?;
3875            let beta_local = beta.slice(s![cache.coeff_global_range.clone()]);
3876            let state =
3877                SpatialPenaltyExactState::from_beta_local(beta_local, cache, params.epsilon)
3878                    .map_err(|e| e.to_string())?;
3879            penalty_value += params.lambda[0] * state.magnitude.penalty_value();
3880            penalty_value += params.lambda[1] * state.gradient.penalty_value();
3881            penalty_value += params.lambda[2] * state.curvature.penalty_value();
3882        }
3883        Ok(penalty_value)
3884    }
3885
3886    fn zero_hyper_parts(&self) -> (Array1<f64>, Array2<f64>) {
3887        let total_dim = self.design.ncols();
3888        (
3889            Array1::<f64>::zeros(total_dim),
3890            Array2::<f64>::zeros((total_dim, total_dim)),
3891        )
3892    }
3893
3894    fn embed_local_hyper_parts(
3895        &self,
3896        coeff_range: &Range<usize>,
3897        local_grad: &Array1<f64>,
3898        local_hess: &Array2<f64>,
3899    ) -> (Array1<f64>, Array2<f64>) {
3900        let (mut beta_mixed, mut betahessian) = self.zero_hyper_parts();
3901        beta_mixed
3902            .slice_mut(s![coeff_range.clone()])
3903            .assign(local_grad);
3904        betahessian
3905            .slice_mut(s![coeff_range.clone(), coeff_range.clone()])
3906            .assign(local_hess);
3907        (beta_mixed, betahessian)
3908    }
3909
3910    fn embed_local_hyper_hessian(
3911        &self,
3912        coeff_range: &Range<usize>,
3913        local_hess: &Array2<f64>,
3914    ) -> Array2<f64> {
3915        let total_dim = self.design.ncols();
3916        let mut out = Array2::<f64>::zeros((total_dim, total_dim));
3917        out.slice_mut(s![coeff_range.clone(), coeff_range.clone()])
3918            .assign(local_hess);
3919        out
3920    }
3921
3922    /// Unified per-block hyper-derivative assembly. Owns the shared cache /
3923    /// hyperparameter / exact-state lookup, the component -> operator selection
3924    /// (scalar magnitude `d0`, grouped gradient `d1`, grouped curvature `d2`),
3925    /// and the global embedding via [`Self::embed_local_hyper_parts`]. The only
3926    /// piece that varies with `derivative` is the per-component accessor triple
3927    /// (objective scalar, beta-mixed gradient, beta hessian) read off the exact
3928    /// state. Returns `(objective, beta_mixed, betahessian)`, each already
3929    /// scaled by the component's penalty weight `lambda`.
3930    fn adaptive_block_eval(
3931        &self,
3932        eval: &SpatialAdaptiveExactEvaluation,
3933        cache_idx: usize,
3934        component: AdaptiveComponent,
3935        derivative: HyperDerivativeKind,
3936    ) -> Result<(f64, Array1<f64>, Array2<f64>), String> {
3937        let cache = self
3938            .runtime_caches
3939            .get(cache_idx)
3940            .ok_or_else(|| format!("adaptive cache index {} out of bounds", cache_idx))?;
3941        let params = self
3942            .adaptive_params
3943            .get(cache_idx)
3944            .ok_or_else(|| format!("adaptive hyperparameter block {} out of bounds", cache_idx))?;
3945        let state = eval
3946            .adaptive_states
3947            .get(cache_idx)
3948            .ok_or_else(|| format!("adaptive exact state index {} out of bounds", cache_idx))?;
3949
3950        let (objective_local, beta_mixed_local, betahessian_local) = match component {
3951            AdaptiveComponent::Magnitude => {
3952                let lambda = params.lambda[0];
3953                let mag = &state.magnitude;
3954                let (objective, gradient_coeff, hessian_diag) = match derivative {
3955                    HyperDerivativeKind::Rho => (
3956                        mag.penalty_value(),
3957                        mag.betagradient_coeff(),
3958                        mag.betahessian_diag(),
3959                    ),
3960                    HyperDerivativeKind::LogEpsilonFirst => (
3961                        mag.log_epsilon_gradient_terms().sum(),
3962                        mag.log_epsilon_betagradient_coeff(),
3963                        mag.log_epsilon_betahessian_diag(),
3964                    ),
3965                    HyperDerivativeKind::LogEpsilonSecond => (
3966                        mag.log_epsilon_hessian_terms().sum(),
3967                        mag.log_epsilon_beta_mixed_second_coeff(),
3968                        mag.log_epsilon_betahessian_second_diag(),
3969                    ),
3970                };
3971                (
3972                    lambda * objective,
3973                    lambda * scalar_operatorgradient(&cache.d0, &gradient_coeff),
3974                    lambda * scalar_operatorhessian(&cache.d0, &hessian_diag),
3975                )
3976            }
3977            AdaptiveComponent::Gradient => {
3978                let lambda = params.lambda[1];
3979                let grad = &state.gradient;
3980                let (objective, gradient_blocks, hessian_blocks) = match derivative {
3981                    HyperDerivativeKind::Rho => (
3982                        grad.penalty_value(),
3983                        grad.betagradient_blocks(),
3984                        grad.betahessian_blocks(),
3985                    ),
3986                    HyperDerivativeKind::LogEpsilonFirst => (
3987                        grad.log_epsilon_gradient_terms().sum(),
3988                        grad.log_epsilon_betagradient_blocks(),
3989                        grad.log_epsilon_betahessian_blocks(),
3990                    ),
3991                    HyperDerivativeKind::LogEpsilonSecond => (
3992                        grad.log_epsilon_hessian_terms().sum(),
3993                        grad.log_epsilon_beta_mixed_second_blocks(),
3994                        grad.log_epsilon_betahessian_second_blocks(),
3995                    ),
3996                };
3997                (
3998                    lambda * objective,
3999                    lambda
4000                        * grouped_operatorgradient(&cache.d1, cache.dimension, &gradient_blocks)
4001                            .map_err(|e| e.to_string())?,
4002                    lambda
4003                        * grouped_operatorhessian(&cache.d1, cache.dimension, &hessian_blocks)
4004                            .map_err(|e| e.to_string())?,
4005                )
4006            }
4007            AdaptiveComponent::Curvature => {
4008                let lambda = params.lambda[2];
4009                let group = cache.dimension * cache.dimension;
4010                let curv = &state.curvature;
4011                let (objective, gradient_blocks, hessian_blocks) = match derivative {
4012                    HyperDerivativeKind::Rho => (
4013                        curv.penalty_value(),
4014                        curv.betagradient_blocks(),
4015                        curv.betahessian_blocks(),
4016                    ),
4017                    HyperDerivativeKind::LogEpsilonFirst => (
4018                        curv.log_epsilon_gradient_terms().sum(),
4019                        curv.log_epsilon_betagradient_blocks(),
4020                        curv.log_epsilon_betahessian_blocks(),
4021                    ),
4022                    HyperDerivativeKind::LogEpsilonSecond => (
4023                        curv.log_epsilon_hessian_terms().sum(),
4024                        curv.log_epsilon_beta_mixed_second_blocks(),
4025                        curv.log_epsilon_betahessian_second_blocks(),
4026                    ),
4027                };
4028                (
4029                    lambda * objective,
4030                    lambda
4031                        * grouped_operatorgradient(&cache.d2, group, &gradient_blocks)
4032                            .map_err(|e| e.to_string())?,
4033                    lambda
4034                        * grouped_operatorhessian(&cache.d2, group, &hessian_blocks)
4035                            .map_err(|e| e.to_string())?,
4036                )
4037            }
4038        };
4039
4040        let (beta_mixed, betahessian) = self.embed_local_hyper_parts(
4041            &cache.coeff_global_range,
4042            &beta_mixed_local,
4043            &betahessian_local,
4044        );
4045        Ok((objective_local, beta_mixed, betahessian))
4046    }
4047
4048    fn adaptive_shared_log_epsilon_parts(
4049        &self,
4050        eval: &SpatialAdaptiveExactEvaluation,
4051        component: usize,
4052    ) -> Result<(f64, Array1<f64>, Array2<f64>), String> {
4053        // Exact shared-log-epsilon first-order pieces:
4054        //
4055        //   J_{eta_p}         = sum_m lambda_{m,p} U_{m,p,eta},
4056        //   J_{beta,eta_p}    = sum_m lambda_{m,p} U_{m,p,beta eta},
4057        //   J_{beta,beta,eta} = sum_m lambda_{m,p} U_{m,p,beta beta eta}.
4058        self.adaptive_shared_block_eval(eval, component, HyperDerivativeKind::LogEpsilonFirst)
4059    }
4060
4061    fn adaptive_shared_log_epsilon_second_parts(
4062        &self,
4063        eval: &SpatialAdaptiveExactEvaluation,
4064        component: usize,
4065    ) -> Result<(f64, Array1<f64>, Array2<f64>), String> {
4066        // Exact shared-log-epsilon second-order pieces:
4067        //
4068        //   J_{eta_p,eta_p}            = sum_m lambda_{m,p} U_{m,p,eta eta},
4069        //   J_{beta,eta_p,eta_p}       = sum_m lambda_{m,p} U_{m,p,beta eta eta},
4070        //   J_{beta,beta,eta_p,eta_p}  = sum_m lambda_{m,p} U_{m,p,beta beta eta eta}.
4071        self.adaptive_shared_block_eval(eval, component, HyperDerivativeKind::LogEpsilonSecond)
4072    }
4073
4074    /// Sum a per-block hyper-derivative across every adaptive term for one shared
4075    /// `log epsilon` coordinate (selected by `component`). The three log-epsilon
4076    /// coordinates are shared globally by penalty type, so each contributes the
4077    /// matching component's block from every cache.
4078    fn adaptive_shared_block_eval(
4079        &self,
4080        eval: &SpatialAdaptiveExactEvaluation,
4081        component: usize,
4082        derivative: HyperDerivativeKind,
4083    ) -> Result<(f64, Array1<f64>, Array2<f64>), String> {
4084        let component = AdaptiveComponent::from_index(component)?;
4085        let (mut score, mut hessian) = self.zero_hyper_parts();
4086        let mut objective = 0.0;
4087        for cache_idx in 0..self.runtime_caches.len() {
4088            let (local_objective, local_score, local_hessian) =
4089                self.adaptive_block_eval(eval, cache_idx, component, derivative)?;
4090            objective += local_objective;
4091            score += &local_score;
4092            hessian += &local_hessian;
4093        }
4094        Ok((objective, score, hessian))
4095    }
4096
4097    fn adaptive_shared_log_epsilon_drift(
4098        &self,
4099        eval: &SpatialAdaptiveExactEvaluation,
4100        component: usize,
4101        direction: &Array1<f64>,
4102    ) -> Result<Array2<f64>, String> {
4103        // Exact shared-log-epsilon Hessian drift:
4104        //
4105        //   T_{eta_p}[u] = sum_m lambda_{m,p} D_beta(U_{m,p,beta beta eta})[u].
4106        let component = AdaptiveComponent::from_index(component)?;
4107        let total_dim = self.design.ncols();
4108        let mut total = Array2::<f64>::zeros((total_dim, total_dim));
4109        for cache_idx in 0..self.runtime_caches.len() {
4110            total += &self.adaptive_block_drift_eval(
4111                eval,
4112                cache_idx,
4113                component,
4114                HyperDriftKind::LogEpsilon,
4115                direction,
4116            )?;
4117        }
4118        Ok(total)
4119    }
4120
4121    fn adaptive_explicit_second_order_parts(
4122        &self,
4123        eval: &SpatialAdaptiveExactEvaluation,
4124        left: SpatialAdaptiveHyperSpec,
4125        right: SpatialAdaptiveHyperSpec,
4126    ) -> Result<(f64, Array1<f64>, Array2<f64>), String> {
4127        // Structural sparsity from the adaptive penalty algebra:
4128        //
4129        //   - alpha_{m,p} / alpha_{n,r} is nonzero only when (m,p) = (n,r),
4130        //   - alpha_{m,p} / eta_r is nonzero only when p = r,
4131        //   - eta_p / eta_r is nonzero only when p = r,
4132        //
4133        // with eta_p contributions summed over all adaptive terms m because the
4134        // three log-epsilon coordinates are shared globally by penalty type.
4135        match left.explicit_second_order_kind(right) {
4136            SpatialAdaptiveExplicitSecondOrderKind::StructuralZero => {
4137                let (score, hessian) = self.zero_hyper_parts();
4138                Ok((0.0, score, hessian))
4139            }
4140            SpatialAdaptiveExplicitSecondOrderKind::LocalAlphaAlpha => self.adaptive_block_eval(
4141                eval,
4142                left.cache_index,
4143                AdaptiveComponent::from_index(left.component_index())?,
4144                HyperDerivativeKind::Rho,
4145            ),
4146            SpatialAdaptiveExplicitSecondOrderKind::LocalAlphaEta => {
4147                let local_alpha = if left.kind.is_log_lambda() {
4148                    left
4149                } else {
4150                    right
4151                };
4152                self.adaptive_block_eval(
4153                    eval,
4154                    local_alpha.cache_index,
4155                    AdaptiveComponent::from_index(local_alpha.component_index())?,
4156                    HyperDerivativeKind::LogEpsilonFirst,
4157                )
4158            }
4159            SpatialAdaptiveExplicitSecondOrderKind::SharedEtaEta => {
4160                self.adaptive_shared_log_epsilon_second_parts(eval, left.component_index())
4161            }
4162        }
4163    }
4164
4165    /// Unified per-block directional-drift assembly. Owns the shared cache /
4166    /// hyperparameter / exact-state lookup, the per-component direction
4167    /// projection through the collocation operators, the operator embedding, and
4168    /// the global embedding via [`Self::embed_local_hyper_hessian`]. The only
4169    /// piece that varies with `drift` is the directional state accessor:
4170    /// [`HyperDriftKind::Rho`] takes the bare directional Hessian drift, while
4171    /// [`HyperDriftKind::LogEpsilon`] takes its `log epsilon` derivative.
4172    fn adaptive_block_drift_eval(
4173        &self,
4174        eval: &SpatialAdaptiveExactEvaluation,
4175        cache_idx: usize,
4176        component: AdaptiveComponent,
4177        drift: HyperDriftKind,
4178        direction: &Array1<f64>,
4179    ) -> Result<Array2<f64>, String> {
4180        let cache = self
4181            .runtime_caches
4182            .get(cache_idx)
4183            .ok_or_else(|| format!("adaptive cache index {} out of bounds", cache_idx))?;
4184        let params = self
4185            .adaptive_params
4186            .get(cache_idx)
4187            .ok_or_else(|| format!("adaptive hyperparameter block {} out of bounds", cache_idx))?;
4188        let state = eval
4189            .adaptive_states
4190            .get(cache_idx)
4191            .ok_or_else(|| format!("adaptive exact state index {} out of bounds", cache_idx))?;
4192        let direction_local = direction.slice(s![cache.coeff_global_range.clone()]);
4193
4194        let local_hessian = match component {
4195            AdaptiveComponent::Magnitude => {
4196                let d0_u = cache.d0.dot(&direction_local);
4197                let mag = &state.magnitude;
4198                let diag = match drift {
4199                    HyperDriftKind::Rho => mag.directionalhessian_diag(&d0_u),
4200                    HyperDriftKind::LogEpsilon => {
4201                        mag.log_epsilon_betahessian_directional_diag(&d0_u)
4202                    }
4203                };
4204                params.lambda[0] * scalar_operatorhessian(&cache.d0, &diag)
4205            }
4206            AdaptiveComponent::Gradient => {
4207                let d1_u = cache.d1.dot(&direction_local);
4208                let direction_blocks = collocationgradient_blocks(&d1_u, cache.dimension)
4209                    .map_err(|e| e.to_string())?;
4210                let grad = &state.gradient;
4211                let blocks = match drift {
4212                    HyperDriftKind::Rho => grad.directionalhessian_blocks(&direction_blocks),
4213                    HyperDriftKind::LogEpsilon => {
4214                        grad.log_epsilon_betahessian_directional_blocks(&direction_blocks)
4215                    }
4216                };
4217                params.lambda[1]
4218                    * grouped_operatorhessian(&cache.d1, cache.dimension, &blocks)
4219                        .map_err(|e| e.to_string())?
4220            }
4221            AdaptiveComponent::Curvature => {
4222                let group = cache.dimension * cache.dimension;
4223                let d2_u = cache.d2.dot(&direction_local);
4224                let direction_blocks =
4225                    collocationhessian_blocks(&d2_u, cache.dimension).map_err(|e| e.to_string())?;
4226                let curv = &state.curvature;
4227                let blocks = match drift {
4228                    HyperDriftKind::Rho => curv.directionalhessian_blocks(&direction_blocks),
4229                    HyperDriftKind::LogEpsilon => {
4230                        curv.log_epsilon_betahessian_directional_blocks(&direction_blocks)
4231                    }
4232                };
4233                params.lambda[2]
4234                    * grouped_operatorhessian(&cache.d2, group, &blocks)
4235                        .map_err(|e| e.to_string())?
4236            }
4237        };
4238
4239        Ok(self.embed_local_hyper_hessian(&cache.coeff_global_range, &local_hessian))
4240    }
4241
4242    fn adaptive_hyper_parts(
4243        &self,
4244        eval: &SpatialAdaptiveExactEvaluation,
4245        hyper: SpatialAdaptiveHyperSpec,
4246    ) -> Result<(f64, Array1<f64>, Array2<f64>), String> {
4247        match hyper.kind {
4248            // Per-term `log lambda` (rho) hyper-derivative: the bare penalty
4249            // pieces for this cache's selected component.
4250            SpatialAdaptiveHyperKind::LogLambdaMagnitude
4251            | SpatialAdaptiveHyperKind::LogLambdaGradient
4252            | SpatialAdaptiveHyperKind::LogLambdaCurvature => self.adaptive_block_eval(
4253                eval,
4254                hyper.cache_index,
4255                AdaptiveComponent::from_index(hyper.component_index())?,
4256                HyperDerivativeKind::Rho,
4257            ),
4258            // Shared `log epsilon` hyper-derivative: summed across all terms.
4259            SpatialAdaptiveHyperKind::LogEpsilonMagnitude
4260            | SpatialAdaptiveHyperKind::LogEpsilonGradient
4261            | SpatialAdaptiveHyperKind::LogEpsilonCurvature => {
4262                self.adaptive_shared_log_epsilon_parts(eval, hyper.component_index())
4263            }
4264        }
4265    }
4266
4267    fn exact_evaluation_uncached(
4268        &self,
4269        beta: &Array1<f64>,
4270    ) -> Result<SpatialAdaptiveExactEvaluation, String> {
4271        let eta = self.total_eta(beta);
4272        let obs = evaluate_standard_familyobservations(
4273            self.family.clone(),
4274            self.latent_cloglog_state.as_ref(),
4275            self.mixture_link_state.as_ref(),
4276            self.sas_link_state.as_ref(),
4277            &self.y,
4278            &self.weights,
4279            &eta,
4280        )
4281        .map_err(|e| e.to_string())?;
4282        let p = beta.len();
4283        let mut penalty_value = 0.0;
4284        let mut penaltygradient = Array1::<f64>::zeros(p);
4285        let mut penaltyhessian = Array2::<f64>::zeros((p, p));
4286        let mut adaptive_states = Vec::with_capacity(self.runtime_caches.len());
4287
4288        for (cache_idx, cache) in self.runtime_caches.iter().enumerate() {
4289            let params = self.adaptive_params.get(cache_idx).ok_or_else(|| {
4290                format!(
4291                    "missing adaptive parameter block for cache {}",
4292                    cache.termname
4293                )
4294            })?;
4295            let beta_local = beta.slice(s![cache.coeff_global_range.clone()]);
4296            let state =
4297                SpatialPenaltyExactState::from_beta_local(beta_local, cache, params.epsilon)
4298                    .map_err(|e| e.to_string())?;
4299
4300            let g0 = scalar_operatorgradient(&cache.d0, &state.magnitude.betagradient_coeff());
4301            let gg = grouped_operatorgradient(
4302                &cache.d1,
4303                cache.dimension,
4304                &state.gradient.betagradient_blocks(),
4305            )
4306            .map_err(|e| e.to_string())?;
4307            let gc = grouped_operatorgradient(
4308                &cache.d2,
4309                cache.dimension * cache.dimension,
4310                &state.curvature.betagradient_blocks(),
4311            )
4312            .map_err(|e| e.to_string())?;
4313            let h0 = scalar_operatorhessian(&cache.d0, &state.magnitude.betahessian_diag());
4314            let hg = grouped_operatorhessian(
4315                &cache.d1,
4316                cache.dimension,
4317                &state.gradient.betahessian_blocks(),
4318            )
4319            .map_err(|e| e.to_string())?;
4320            let hc = grouped_operatorhessian(
4321                &cache.d2,
4322                cache.dimension * cache.dimension,
4323                &state.curvature.betahessian_blocks(),
4324            )
4325            .map_err(|e| e.to_string())?;
4326
4327            let lambda0 = params.lambda[0];
4328            let lambdag = params.lambda[1];
4329            let lambdac = params.lambda[2];
4330
4331            penalty_value += lambda0 * state.magnitude.penalty_value();
4332            penalty_value += lambdag * state.gradient.penalty_value();
4333            penalty_value += lambdac * state.curvature.penalty_value();
4334
4335            let range = cache.coeff_global_range.clone();
4336            {
4337                let mut grad_local = penaltygradient.slice_mut(s![range.clone()]);
4338                grad_local += &(g0.mapv(|v| lambda0 * v));
4339                grad_local += &(gg.mapv(|v| lambdag * v));
4340                grad_local += &(gc.mapv(|v| lambdac * v));
4341            }
4342            {
4343                let mut h_local = penaltyhessian.slice_mut(s![range.clone(), range]);
4344                h_local += &h0.mapv(|v| lambda0 * v);
4345                h_local += &hg.mapv(|v| lambdag * v);
4346                h_local += &hc.mapv(|v| lambdac * v);
4347            }
4348
4349            adaptive_states.push(state);
4350        }
4351
4352        let (fixed_quadraticvalue, fixed_quadraticgradient) = self.fixed_quadratic_terms(beta);
4353        Ok(SpatialAdaptiveExactEvaluation {
4354            obs,
4355            adaptive_states,
4356            adaptive_penalty_value: penalty_value,
4357            adaptive_penaltygradient: penaltygradient,
4358            adaptive_penaltyhessian: penaltyhessian,
4359            fixed_quadraticvalue,
4360            fixed_quadraticgradient,
4361            fixed_quadratichessian: self.fixed_quadratichessian.as_ref().clone(),
4362        })
4363    }
4364
4365    fn exact_evaluation(
4366        &self,
4367        beta: &Array1<f64>,
4368    ) -> Result<Arc<SpatialAdaptiveExactEvaluation>, String> {
4369        {
4370            let cache = self
4371                .exact_eval_cache
4372                .lock()
4373                .map_err(|_| "spatial adaptive exact-evaluation cache lock poisoned".to_string())?;
4374            if let Some(cached) = cache.as_ref()
4375                && cached.beta.len() == beta.len()
4376                && cached
4377                    .beta
4378                    .iter()
4379                    .zip(beta.iter())
4380                    .all(|(&left, &right)| left == right)
4381            {
4382                return Ok(Arc::clone(&cached.eval));
4383            }
4384        }
4385
4386        let eval = Arc::new(self.exact_evaluation_uncached(beta)?);
4387        let mut cache = self
4388            .exact_eval_cache
4389            .lock()
4390            .map_err(|_| "spatial adaptive exact-evaluation cache lock poisoned".to_string())?;
4391        *cache = Some(CachedSpatialAdaptiveExactEvaluation {
4392            beta: beta.clone(),
4393            eval: Arc::clone(&eval),
4394        });
4395        Ok(eval)
4396    }
4397
4398    fn exacthessian_directional_derivative_from_evaluation(
4399        &self,
4400        beta: &Array1<f64>,
4401        eval: &SpatialAdaptiveExactEvaluation,
4402        direction: &Array1<f64>,
4403    ) -> Result<Array2<f64>, String> {
4404        assert_eq!(
4405            beta.len(),
4406            direction.len(),
4407            "beta/direction length mismatch",
4408        );
4409        let d_eta = gam_linalg::faer_ndarray::fast_av(self.design.as_ref(), direction);
4410        let mut total = xt_diag_x_dense(
4411            self.design.view(),
4412            (&eval.obs.neghessian_eta_derivative * &d_eta).view(),
4413        )?;
4414        for (cache_idx, cache) in self.runtime_caches.iter().enumerate() {
4415            let params = self.adaptive_params.get(cache_idx).ok_or_else(|| {
4416                format!(
4417                    "missing adaptive parameter block for cache {}",
4418                    cache.termname
4419                )
4420            })?;
4421            let state = eval
4422                .adaptive_states
4423                .get(cache_idx)
4424                .ok_or_else(|| format!("missing adaptive state for cache {}", cache.termname))?;
4425            let direction_local = direction.slice(s![cache.coeff_global_range.clone()]);
4426            let d0_u = cache.d0.dot(&direction_local);
4427            let d1_u = cache.d1.dot(&direction_local);
4428            let d2_u = cache.d2.dot(&direction_local);
4429            let h0 =
4430                scalar_operatorhessian(&cache.d0, &state.magnitude.directionalhessian_diag(&d0_u))
4431                    .mapv(|v| params.lambda[0] * v);
4432            let hg = grouped_operatorhessian(
4433                &cache.d1,
4434                cache.dimension,
4435                &state.gradient.directionalhessian_blocks(
4436                    &collocationgradient_blocks(&d1_u, cache.dimension)
4437                        .map_err(|e| e.to_string())?,
4438                ),
4439            )
4440            .map_err(|e| e.to_string())?
4441            .mapv(|v| params.lambda[1] * v);
4442            let hc = grouped_operatorhessian(
4443                &cache.d2,
4444                cache.dimension * cache.dimension,
4445                &state.curvature.directionalhessian_blocks(
4446                    &collocationhessian_blocks(&d2_u, cache.dimension)
4447                        .map_err(|e| e.to_string())?,
4448                ),
4449            )
4450            .map_err(|e| e.to_string())?
4451            .mapv(|v| params.lambda[2] * v);
4452            let range = cache.coeff_global_range.clone();
4453            let mut local = total.slice_mut(s![range.clone(), range]);
4454            local += &h0;
4455            local += &hg;
4456            local += &hc;
4457        }
4458        Ok(total)
4459    }
4460
4461    /// Exact second directional derivative `D²_β H[u, v]` of the joint
4462    /// (likelihood + adaptive Charbonnier penalty) Hessian, needed so the outer
4463    /// LAML's joint-Jeffreys curvature drift `D_β H_Φ[β̇]` is exact rather than
4464    /// silently dropped (which leaves the outer hypergradient inconsistent with
4465    /// the `½log|H+H_Φ|` objective it folds `H_Φ` into).
4466    ///
4467    /// The data block contributes `Xᵀ diag(ℓ'''(η_i) (Xu)_i (Xv)_i) X`, where
4468    /// `ℓ'''` is the third derivative of the per-observation log-likelihood in
4469    /// `η`. The observation state exposes the working weight `w=−ℓ''` and its
4470    /// first `η`-derivative `w'` (`neghessian_eta_derivative`) but not `w''`, so
4471    /// the exact data term is available only on the **constant-weight** path
4472    /// (`w' ≡ 0`, e.g. Gaussian identity), where `w'' ≡ 0` and the data block
4473    /// second derivative vanishes. On a varying-weight family we return `None`
4474    /// (the safe, pre-existing behavior: the drift degrades to zero rather than
4475    /// to a wrong value) until the observation contract carries `w''`.
4476    ///
4477    /// The penalty block is always exact: with `λ_m G_mᵀ B_m(G_m β) G_m` the
4478    /// per-component penalty Hessian, `D²_β` is `λ_m Σ_k G_mᵀ N_m,k G_m` using the
4479    /// scalar (`second_directionalhessian_diag`) / grouped
4480    /// (`second_directionalhessian_blocks`) fourth-derivative contractions.
4481    fn exacthessian_second_directional_derivative_from_evaluation(
4482        &self,
4483        eval: &SpatialAdaptiveExactEvaluation,
4484        direction_u: &Array1<f64>,
4485        direction_v: &Array1<f64>,
4486    ) -> Result<Option<Array2<f64>>, String> {
4487        let p = self.design.ncols();
4488        // Data block: exact only when the working weight is constant in η.
4489        if eval.obs.neghessian_eta_derivative.iter().any(|&w| w != 0.0) {
4490            return Ok(None);
4491        }
4492        let mut total = Array2::<f64>::zeros((p, p));
4493        for (cache_idx, cache) in self.runtime_caches.iter().enumerate() {
4494            let params = self.adaptive_params.get(cache_idx).ok_or_else(|| {
4495                format!(
4496                    "missing adaptive parameter block for cache {}",
4497                    cache.termname
4498                )
4499            })?;
4500            let state = eval
4501                .adaptive_states
4502                .get(cache_idx)
4503                .ok_or_else(|| format!("missing adaptive state for cache {}", cache.termname))?;
4504            let u_local = direction_u.slice(s![cache.coeff_global_range.clone()]);
4505            let v_local = direction_v.slice(s![cache.coeff_global_range.clone()]);
4506
4507            // Magnitude (scalar d0).
4508            let q0_u = cache.d0.dot(&u_local);
4509            let q0_v = cache.d0.dot(&v_local);
4510            let h0 = scalar_operatorhessian(
4511                &cache.d0,
4512                &state.magnitude.second_directionalhessian_diag(&q0_u, &q0_v),
4513            )
4514            .mapv(|x| params.lambda[0] * x);
4515
4516            // Gradient (grouped d1, block dim = dimension).
4517            let a1 = collocationgradient_blocks(&cache.d1.dot(&u_local), cache.dimension)
4518                .map_err(|e| e.to_string())?;
4519            let b1 = collocationgradient_blocks(&cache.d1.dot(&v_local), cache.dimension)
4520                .map_err(|e| e.to_string())?;
4521            let hg = grouped_operatorhessian(
4522                &cache.d1,
4523                cache.dimension,
4524                &state.gradient.second_directionalhessian_blocks(&a1, &b1),
4525            )
4526            .map_err(|e| e.to_string())?
4527            .mapv(|x| params.lambda[1] * x);
4528
4529            // Curvature (grouped d2, block dim = dimension²).
4530            let a2 = collocationhessian_blocks(&cache.d2.dot(&u_local), cache.dimension)
4531                .map_err(|e| e.to_string())?;
4532            let b2 = collocationhessian_blocks(&cache.d2.dot(&v_local), cache.dimension)
4533                .map_err(|e| e.to_string())?;
4534            let hc = grouped_operatorhessian(
4535                &cache.d2,
4536                cache.dimension * cache.dimension,
4537                &state.curvature.second_directionalhessian_blocks(&a2, &b2),
4538            )
4539            .map_err(|e| e.to_string())?
4540            .mapv(|x| params.lambda[2] * x);
4541
4542            let range = cache.coeff_global_range.clone();
4543            let mut local = total.slice_mut(s![range.clone(), range]);
4544            local += &h0;
4545            local += &hg;
4546            local += &hc;
4547        }
4548        Ok(Some(total))
4549    }
4550}
4551
4552impl CustomFamily for SpatialAdaptiveExactFamily {
4553    // Preserve the pre-gam#1395 behavior: the trait default flipped to OFF (the
4554    // flat-prior exact-Newton objective carries no Jeffreys term), so families
4555    // that historically armed the term by default opt back in explicitly.
4556    fn joint_jeffreys_term_required(&self) -> bool {
4557        true
4558    }
4559
4560    // Jeffreys/Firth information = the LIKELIHOOD Fisher information only
4561    // (`Xᵀ W X`, `W = −ℓ''(η)`), NOT the penalized joint Newton Hessian
4562    // `Xᵀ W X + ∂²_β penalty` the trait default (`exact_newton_joint_hessian`)
4563    // returns. Two reasons, both load-bearing for the #901 outer-REML
4564    // hypergradient:
4565    //
4566    //   1. CONTRACT. Jeffreys' prior is `Φ = ½ log|I(β)|₊` with `I` the
4567    //      likelihood information; the adaptive Charbonnier term is the PRIOR,
4568    //      not the likelihood, so folding its curvature into `I` is a
4569    //      category error (the trait doc on `joint_jeffreys_information_with_specs`
4570    //      spells this out — "Jeffreys' prior is defined from expected
4571    //      information").
4572    //
4573    //   2. θ-CONSISTENCY. With the full span `Z_J = I`, the reduced
4574    //      information IS `I(β)`. If the penalty Hessian `S_λ,ε(θ)` rode along,
4575    //      `Φ` would depend on the smoothing hyperparameters `θ = (log λ, log ε)`
4576    //      EXPLICITLY through `S_λ,ε`. The outer gradient then needs `−∂_θ Φ`
4577    //      (psi_hyper's `phi_psi`), computed from the EXACT, UNGATED, UNFLOORED
4578    //      `joint_jeffreys_phi_explicit_param_derivative`, whereas the LAML cost
4579    //      folds the GATED + spectrally-FLOORED value `Φ` and a
4580    //      divided-difference `H_Φ` that omits its second-order completion.
4581    //      Those two describe different functions, so the analytic
4582    //      hypergradient disagreed with the central-difference reference by
4583    //      exactly that penalty-driven `∂_θ Φ` — the residual scaling with the
4584    //      Charbonnier group dimension (mass 1, tension 2, curvature 4) the
4585    //      #901 fixture pinned. `Xᵀ W X` carries NO `θ` dependence, so
4586    //      `∂_θ Φ ≡ 0` and the term contributes only its genuine β-mode-response
4587    //      (which the envelope identity already accounts for), restoring
4588    //      analytic-vs-FD agreement to f64 grade.
4589    //
4590    // For Gaussian identity `W ≡ 1`, so this is the constant data Gram `XᵀX`,
4591    // which is also β-independent — its β-directional derivatives below are
4592    // therefore zero, matching the exact Fisher-information geometry. On a
4593    // genuinely near-separating non-Gaussian fit the data information still
4594    // shrinks where the conditioning gate arms, so the self-limiting Firth
4595    // bound is preserved exactly where it is needed.
4596    fn joint_jeffreys_information_with_specs(
4597        &self,
4598        block_states: &[ParameterBlockState],
4599        specs: &[ParameterBlockSpec],
4600    ) -> Result<Option<Array2<f64>>, String> {
4601        let spec = expect_single_blockspec(specs, "spatial adaptive exact family")?;
4602        let beta = &expect_single_block_state(block_states, "spatial adaptive exact family")?.beta;
4603        if spec.design.ncols() != beta.len() {
4604            return Err(SmoothError::dimension_mismatch(format!(
4605                "spatial adaptive Jeffreys information: spec design has {} columns, beta has {}",
4606                spec.design.ncols(),
4607                beta.len()
4608            ))
4609            .into());
4610        }
4611        let eval = self.exact_evaluation(beta)?;
4612        Ok(Some(xt_diag_x_dense(
4613            self.design.view(),
4614            eval.obs.neghessian_eta.view(),
4615        )?))
4616    }
4617
4618    fn joint_jeffreys_information_directional_derivative_with_specs(
4619        &self,
4620        block_states: &[ParameterBlockState],
4621        specs: &[ParameterBlockSpec],
4622        d_beta_flat: &Array1<f64>,
4623    ) -> Result<Option<Array2<f64>>, String> {
4624        // `D_β(Xᵀ W X)[u] = Xᵀ diag(W'(η) (X u)) X`, with `W = −ℓ''(η)` and
4625        // `W' = neghessian_eta_derivative`. Mirrors the data-block term of
4626        // `exacthessian_directional_derivative_from_evaluation`, MINUS the
4627        // penalty contribution (the penalty is not part of the likelihood
4628        // information). Zero for the constant-weight (Gaussian-identity) path.
4629        let spec = expect_single_blockspec(specs, "spatial adaptive exact family")?;
4630        let beta = &expect_single_block_state(block_states, "spatial adaptive exact family")?.beta;
4631        if spec.design.ncols() != d_beta_flat.len() {
4632            return Err(SmoothError::dimension_mismatch(format!(
4633                "spatial adaptive Jeffreys directional derivative: spec design has {} columns, direction has {}",
4634                spec.design.ncols(),
4635                d_beta_flat.len()
4636            ))
4637            .into());
4638        }
4639        let eval = self.exact_evaluation(beta)?;
4640        let d_eta = gam_linalg::faer_ndarray::fast_av(self.design.as_ref(), d_beta_flat);
4641        Ok(Some(xt_diag_x_dense(
4642            self.design.view(),
4643            (&eval.obs.neghessian_eta_derivative * &d_eta).view(),
4644        )?))
4645    }
4646
4647    fn joint_jeffreys_information_second_directional_derivative_with_specs(
4648        &self,
4649        block_states: &[ParameterBlockState],
4650        specs: &[ParameterBlockSpec],
4651        d_beta_u_flat: &Array1<f64>,
4652        d_betav_flat: &Array1<f64>,
4653    ) -> Result<Option<Array2<f64>>, String> {
4654        // `D²_β(Xᵀ W X)[u, v] = Xᵀ diag(W''(η) (X u) (X v)) X`. The observation
4655        // state exposes `W` and `W'` but not `W''`, so this is exact only on the
4656        // constant-weight path (`W' ≡ 0 ⇒ W'' ≡ 0`, the zero matrix), matching
4657        // the guard in `exacthessian_second_directional_derivative_from_evaluation`.
4658        // On a varying-weight family we return `None` so the divided-difference
4659        // completion degrades safely rather than to a wrong value.
4660        let spec = expect_single_blockspec(specs, "spatial adaptive exact family")?;
4661        let beta = &expect_single_block_state(block_states, "spatial adaptive exact family")?.beta;
4662        if spec.design.ncols() != beta.len()
4663            || d_beta_u_flat.len() != beta.len()
4664            || d_betav_flat.len() != beta.len()
4665        {
4666            return Err(SmoothError::dimension_mismatch(format!(
4667                "spatial adaptive Jeffreys second-direction length mismatch: spec cols={}, dirs=({}, {}), expected {}",
4668                spec.design.ncols(),
4669                d_beta_u_flat.len(),
4670                d_betav_flat.len(),
4671                beta.len()
4672            ))
4673            .into());
4674        }
4675        let eval = self.exact_evaluation(beta)?;
4676        if eval.obs.neghessian_eta_derivative.iter().any(|&w| w != 0.0) {
4677            return Ok(None);
4678        }
4679        Ok(Some(Array2::<f64>::zeros((beta.len(), beta.len()))))
4680    }
4681
4682    fn joint_jeffreys_information_matches_observed_hessian(&self) -> bool {
4683        // The Jeffreys information above is the LIKELIHOOD Fisher information,
4684        // which differs from the penalized observed joint Newton Hessian, so the
4685        // observed-Hessian conditioning pre-check must NOT certify a skip from it
4686        // (gam#1020 expected-information caveat).
4687        false
4688    }
4689
4690    fn joint_jeffreys_information_depends_on_psi(&self) -> bool {
4691        // The Jeffreys information is the data Fisher information `Xᵀ W X`, whose
4692        // explicit ψ-dependence is zero: the smoothing hyperparameters
4693        // ψ = (log λ, log ε) act only through the adaptive Charbonnier PENALTY,
4694        // never the design `X`, so `∂_ψ (Xᵀ W X)|_β ≡ 0`. Returning `false`
4695        // suppresses the three explicit-ψ Firth terms the outer engine would
4696        // otherwise form from `∂_ψ(penalty)` (the wrong perturbation), which is
4697        // exactly the spurious hypergradient bias the #901 fixture pinned. The
4698        // implicit β-mode-response of `Φ` is unaffected and still folded.
4699        false
4700    }
4701
4702    fn evaluate(&self, block_states: &[ParameterBlockState]) -> Result<FamilyEvaluation, String> {
4703        let beta = &expect_single_block_state(block_states, "spatial adaptive exact family")?.beta;
4704        let eval = self.exact_evaluation(beta)?;
4705        let mut gradient = fast_atv(&self.design, &eval.obs.score);
4706        gradient -= &eval.total_penaltygradient();
4707        let mut hessian = xt_diag_x_dense(self.design.view(), eval.obs.neghessian_eta.view())?;
4708        hessian += &eval.total_penaltyhessian();
4709        Ok(FamilyEvaluation {
4710            log_likelihood: eval.obs.log_likelihood - eval.total_penalty_value(),
4711            blockworking_sets: vec![BlockWorkingSet::ExactNewton {
4712                gradient,
4713                hessian: SymmetricMatrix::Dense(hessian),
4714            }],
4715        })
4716    }
4717
4718    fn log_likelihood_only(&self, block_states: &[ParameterBlockState]) -> Result<f64, String> {
4719        let state = expect_single_block_state(block_states, "spatial adaptive exact family")?;
4720        let beta = &state.beta;
4721        let obs = evaluate_standard_familyobservations(
4722            self.family.clone(),
4723            self.latent_cloglog_state.as_ref(),
4724            self.mixture_link_state.as_ref(),
4725            self.sas_link_state.as_ref(),
4726            &self.y,
4727            &self.weights,
4728            &state.eta,
4729        )
4730        .map_err(|e| e.to_string())?;
4731        let adaptive_penalty = self.adaptive_penalty_value_only(beta)?;
4732        let (fixed_quadratic, _) = self.fixed_quadratic_terms(beta);
4733        Ok(obs.log_likelihood - adaptive_penalty - fixed_quadratic)
4734    }
4735
4736    fn exact_newton_outerobjective(&self) -> ExactNewtonOuterObjective {
4737        ExactNewtonOuterObjective::StrictPseudoLaplace
4738    }
4739
4740    fn exact_newton_joint_hessian(
4741        &self,
4742        block_states: &[ParameterBlockState],
4743    ) -> Result<Option<Array2<f64>>, String> {
4744        let beta = &expect_single_block_state(block_states, "spatial adaptive exact family")?.beta;
4745        let eval = self.exact_evaluation(beta)?;
4746        Ok(Some(eval.totalobjectivehessian(&self.design)?))
4747    }
4748
4749    fn exact_newton_hessian_directional_derivative(
4750        &self,
4751        block_states: &[ParameterBlockState],
4752        block_idx: usize,
4753        d_beta: &Array1<f64>,
4754    ) -> Result<Option<Array2<f64>>, String> {
4755        expect_block_idx_zero(block_idx, "spatial adaptive exact family", "")?;
4756        self.exact_newton_joint_hessian_directional_derivative(block_states, d_beta)
4757    }
4758
4759    fn exact_newton_joint_hessian_directional_derivative(
4760        &self,
4761        block_states: &[ParameterBlockState],
4762        d_beta_flat: &Array1<f64>,
4763    ) -> Result<Option<Array2<f64>>, String> {
4764        let beta = &expect_single_block_state(block_states, "spatial adaptive exact family")?.beta;
4765        if d_beta_flat.len() != beta.len() {
4766            return Err(SmoothError::dimension_mismatch(format!(
4767                "spatial adaptive exact family direction length mismatch: got {}, expected {}",
4768                d_beta_flat.len(),
4769                beta.len()
4770            ))
4771            .into());
4772        }
4773        let eval = self.exact_evaluation(beta)?;
4774        Ok(Some(
4775            self.exacthessian_directional_derivative_from_evaluation(beta, &eval, d_beta_flat)?,
4776        ))
4777    }
4778
4779    fn exact_newton_joint_hessiansecond_directional_derivative(
4780        &self,
4781        block_states: &[ParameterBlockState],
4782        d_beta_u_flat: &Array1<f64>,
4783        d_betav_flat: &Array1<f64>,
4784    ) -> Result<Option<Array2<f64>>, String> {
4785        let beta = &expect_single_block_state(block_states, "spatial adaptive exact family")?.beta;
4786        if d_beta_u_flat.len() != beta.len() || d_betav_flat.len() != beta.len() {
4787            return Err(SmoothError::dimension_mismatch(format!(
4788                "spatial adaptive exact family second-direction length mismatch: got ({}, {}), expected {}",
4789                d_beta_u_flat.len(),
4790                d_betav_flat.len(),
4791                beta.len()
4792            ))
4793            .into());
4794        }
4795        let eval = self.exact_evaluation(beta)?;
4796        self.exacthessian_second_directional_derivative_from_evaluation(
4797            &eval,
4798            d_beta_u_flat,
4799            d_betav_flat,
4800        )
4801    }
4802
4803    fn block_linear_constraints(
4804        &self,
4805        block_states: &[ParameterBlockState],
4806        block_idx: usize,
4807        block_spec: &ParameterBlockSpec,
4808    ) -> Result<Option<LinearInequalityConstraints>, String> {
4809        assert!(!block_states.is_empty(), "block_states must be non-empty");
4810        assert!(
4811            !block_spec.name.is_empty(),
4812            "block spec name must be non-empty",
4813        );
4814        expect_block_idx_zero(block_idx, "spatial adaptive exact family", "")?;
4815        Ok(self.linear_constraints.clone())
4816    }
4817
4818    fn exact_newton_joint_psi_terms(
4819        &self,
4820        block_states: &[ParameterBlockState],
4821        specs: &[ParameterBlockSpec],
4822        derivative_blocks: &[Vec<CustomFamilyBlockPsiDerivative>],
4823        psi_index: usize,
4824    ) -> Result<Option<ExactNewtonJointPsiTerms>, String> {
4825        if block_states.len() != 1 || specs.len() != 1 || derivative_blocks.len() != 1 {
4826            return Err(SmoothError::dimension_mismatch(format!(
4827                "spatial adaptive exact family expects one block/state/spec/psi payload, got states={} specs={} deriv_blocks={}",
4828                block_states.len(),
4829                specs.len(),
4830                derivative_blocks.len()
4831            ))
4832            .into());
4833        }
4834        derivative_blocks[0]
4835            .get(psi_index)
4836            .ok_or_else(|| format!("adaptive psi index {} out of bounds", psi_index))?;
4837        let hyper = self
4838            .hyperspecs
4839            .get(psi_index)
4840            .ok_or_else(|| format!("adaptive psi index {} out of bounds", psi_index))?;
4841        let beta = &block_states[0].beta;
4842        let eval = self.exact_evaluation(beta)?;
4843        let (direct, beta_mixed, betahessian_explicit) =
4844            self.adaptive_hyper_parts(&eval, *hyper)?;
4845
4846        // Exact pseudo-Laplace psi-gradient.
4847        //
4848        // For one hyperparameter coordinate a we use the exact formula
4849        //
4850        //   d/da L_tilde
4851        //   = J_a + 0.5 tr(H^{-1} Hdot_a),
4852        //
4853        // with
4854        //
4855        //   H u_a   = J_{beta,a},
4856        //   beta_a  = -u_a,
4857        //   Hdot_a  = J_{beta,beta,a} + D_beta(H)[beta_a]
4858        //           = J_{beta,beta,a} - D_beta(H)[u_a].
4859        //
4860        // Here:
4861        //   - `direct` is J_a,
4862        //   - `beta_mixed` is J_{beta,a},
4863        //   - `betahessian_explicit` is J_{beta,beta,a},
4864        //   - `exacthessian_directional_derivative_from_evaluation(..., u)` returns
4865        //     D_beta(H)[u] for the exact likelihood-plus-Charbonnier model.
4866        Ok(Some(ExactNewtonJointPsiTerms {
4867            objective_psi: direct,
4868            score_psi: beta_mixed,
4869            hessian_psi: betahessian_explicit,
4870            hessian_psi_operator: None,
4871        }))
4872    }
4873
4874    fn exact_newton_joint_psisecond_order_terms(
4875        &self,
4876        block_states: &[ParameterBlockState],
4877        specs: &[ParameterBlockSpec],
4878        derivative_blocks: &[Vec<CustomFamilyBlockPsiDerivative>],
4879        psi_i: usize,
4880        psi_j: usize,
4881    ) -> Result<Option<gam_problem::ExactNewtonJointPsiSecondOrderTerms>, String> {
4882        if block_states.len() != 1 || specs.len() != 1 || derivative_blocks.len() != 1 {
4883            return Err(SmoothError::dimension_mismatch(format!(
4884                "spatial adaptive exact family expects one block/state/spec/psi payload, got states={} specs={} deriv_blocks={}",
4885                block_states.len(),
4886                specs.len(),
4887                derivative_blocks.len()
4888            ))
4889            .into());
4890        }
4891        derivative_blocks[0]
4892            .get(psi_i)
4893            .ok_or_else(|| format!("adaptive psi index {} out of bounds", psi_i))?;
4894        derivative_blocks[0]
4895            .get(psi_j)
4896            .ok_or_else(|| format!("adaptive psi index {} out of bounds", psi_j))?;
4897        let hyper_i = self
4898            .hyperspecs
4899            .get(psi_i)
4900            .ok_or_else(|| format!("adaptive psi index {} out of bounds", psi_i))?;
4901        let hyper_j = self
4902            .hyperspecs
4903            .get(psi_j)
4904            .ok_or_else(|| format!("adaptive psi index {} out of bounds", psi_j))?;
4905        let beta = &block_states[0].beta;
4906        let eval = self.exact_evaluation(beta)?;
4907        let (objective_psi_psi, score_psi_psi, hessian_psi_psi) =
4908            self.adaptive_explicit_second_order_parts(&eval, *hyper_i, *hyper_j)?;
4909
4910        Ok(Some(
4911            gam_problem::ExactNewtonJointPsiSecondOrderTerms {
4912                objective_psi_psi,
4913                score_psi_psi,
4914                hessian_psi_psi,
4915                hessian_psi_psi_operator: None,
4916            },
4917        ))
4918    }
4919
4920    fn exact_newton_joint_psihessian_directional_derivative(
4921        &self,
4922        block_states: &[ParameterBlockState],
4923        specs: &[ParameterBlockSpec],
4924        derivative_blocks: &[Vec<CustomFamilyBlockPsiDerivative>],
4925        psi_index: usize,
4926        direction: &Array1<f64>,
4927    ) -> Result<Option<Array2<f64>>, String> {
4928        if block_states.len() != 1 || specs.len() != 1 || derivative_blocks.len() != 1 {
4929            return Err(SmoothError::dimension_mismatch(format!(
4930                "spatial adaptive exact family expects one block/state/spec/psi payload, got states={} specs={} deriv_blocks={}",
4931                block_states.len(),
4932                specs.len(),
4933                derivative_blocks.len()
4934            ))
4935            .into());
4936        }
4937        let beta = &block_states[0].beta;
4938        if direction.len() != beta.len() {
4939            return Err(SmoothError::dimension_mismatch(format!(
4940                "spatial adaptive exact family direction length mismatch: got {}, expected {}",
4941                direction.len(),
4942                beta.len()
4943            ))
4944            .into());
4945        }
4946        derivative_blocks[0]
4947            .get(psi_index)
4948            .ok_or_else(|| format!("adaptive psi index {} out of bounds", psi_index))?;
4949        let hyper = self
4950            .hyperspecs
4951            .get(psi_index)
4952            .ok_or_else(|| format!("adaptive psi index {} out of bounds", psi_index))?;
4953        let eval = self.exact_evaluation(beta)?;
4954        let drift = match hyper.kind {
4955            SpatialAdaptiveHyperKind::LogLambdaMagnitude
4956            | SpatialAdaptiveHyperKind::LogLambdaGradient
4957            | SpatialAdaptiveHyperKind::LogLambdaCurvature => self.adaptive_block_drift_eval(
4958                &eval,
4959                hyper.cache_index,
4960                AdaptiveComponent::from_index(hyper.kind.component_index())?,
4961                HyperDriftKind::Rho,
4962                direction,
4963            )?,
4964            SpatialAdaptiveHyperKind::LogEpsilonMagnitude
4965            | SpatialAdaptiveHyperKind::LogEpsilonGradient
4966            | SpatialAdaptiveHyperKind::LogEpsilonCurvature => self
4967                .adaptive_shared_log_epsilon_drift(
4968                    &eval,
4969                    hyper.kind.component_index(),
4970                    direction,
4971                )?,
4972        };
4973        Ok(Some(drift))
4974    }
4975}
4976
4977fn expect_single_block_state<'a>(
4978    block_states: &'a [ParameterBlockState],
4979    family_name: &str,
4980) -> Result<&'a ParameterBlockState, String> {
4981    crate::block_layout::block_count::validate_block_count::<SmoothError>(
4982        family_name,
4983        1,
4984        block_states.len(),
4985    )?;
4986    Ok(&block_states[0])
4987}
4988
4989fn expect_single_blockspec<'a>(
4990    specs: &'a [ParameterBlockSpec],
4991    family_name: &str,
4992) -> Result<&'a ParameterBlockSpec, String> {
4993    crate::block_layout::block_count::validate_block_count::<SmoothError>(
4994        family_name,
4995        1,
4996        specs.len(),
4997    )?;
4998    Ok(&specs[0])
4999}
5000
5001fn expect_block_idx_zero(block_idx: usize, family_name: &str, context: &str) -> Result<(), String> {
5002    if block_idx != 0 {
5003        return Err(SmoothError::invalid_index(format!(
5004            "{family_name} expects block_idx 0{context}, got {block_idx}"
5005        ))
5006        .into());
5007    }
5008    Ok::<(), _>(())
5009}
5010
5011impl BoundedLinearFamily {
5012    fn bounded_term_derivative_data(
5013        &self,
5014        latent_beta: &Array1<f64>,
5015    ) -> (
5016        Array1<f64>,
5017        Array1<f64>,
5018        Array1<f64>,
5019        Array1<f64>,
5020        Array1<f64>,
5021    ) {
5022        let p = latent_beta.len();
5023        let mut beta_user = latent_beta.clone();
5024        let mut jac_diag = Array1::<f64>::ones(p);
5025        let mut second_diag = Array1::<f64>::zeros(p);
5026        let mut third_diag = Array1::<f64>::zeros(p);
5027        let mut priorthird = Array1::<f64>::zeros(p);
5028        for term in &self.bounded_terms {
5029            let (beta, _, db_dtheta, d2b_dtheta2, d3b_dtheta3) =
5030                bounded_latent_derivatives(latent_beta[term.col_idx], term.min, term.max);
5031            beta_user[term.col_idx] = beta;
5032            jac_diag[term.col_idx] = db_dtheta;
5033            second_diag[term.col_idx] = d2b_dtheta2;
5034            third_diag[term.col_idx] = d3b_dtheta3;
5035            let (_, _, _, prior_neghess_derivative) =
5036                bounded_prior_terms(latent_beta[term.col_idx], &term.prior);
5037            priorthird[term.col_idx] = prior_neghess_derivative;
5038        }
5039        (beta_user, jac_diag, second_diag, third_diag, priorthird)
5040    }
5041
5042    fn user_beta_and_jacobian(&self, latent_beta: &Array1<f64>) -> (Array1<f64>, Array1<f64>) {
5043        let (beta_user, jac_diag, _, _, _) = self.bounded_term_derivative_data(latent_beta);
5044        (beta_user, jac_diag)
5045    }
5046
5047    fn nonlinear_offset_from_latent(&self, latent_beta: &Array1<f64>) -> Array1<f64> {
5048        let mut offset = self.offset.clone();
5049        for term in &self.bounded_terms {
5050            let (beta, _, _) =
5051                bounded_latent_to_user(latent_beta[term.col_idx], term.min, term.max);
5052            offset.scaled_add(beta, &self.design.column(term.col_idx));
5053        }
5054        offset
5055    }
5056
5057    fn effective_design_for_latent(&self, jac_diag: &Array1<f64>) -> Array2<f64> {
5058        let mut x_eff = self.design.clone();
5059        for term in &self.bounded_terms {
5060            x_eff
5061                .column_mut(term.col_idx)
5062                .mapv_inplace(|v| v * jac_diag[term.col_idx]);
5063        }
5064        x_eff
5065    }
5066
5067    fn exacthessian_andgradient(
5068        &self,
5069        latent_beta: &Array1<f64>,
5070    ) -> Result<
5071        (
5072            StandardFamilyObservationState,
5073            Array2<f64>,
5074            Array1<f64>,
5075            f64,
5076            Array1<f64>,
5077            Array1<f64>,
5078            Array1<f64>,
5079        ),
5080        String,
5081    > {
5082        let (_, jac_diag, second_diag, third_diag, priorthird) =
5083            self.bounded_term_derivative_data(latent_beta);
5084        let x_eff = self.effective_design_for_latent(&jac_diag);
5085        let eta =
5086            self.designzeroed.dot(latent_beta) + self.nonlinear_offset_from_latent(latent_beta);
5087        let obs = evaluate_standard_familyobservations(
5088            self.family.clone(),
5089            self.latent_cloglog_state.as_ref(),
5090            self.mixture_link_state.as_ref(),
5091            self.sas_link_state.as_ref(),
5092            &self.y,
5093            &self.weights,
5094            &eta,
5095        )
5096        .map_err(|e| e.to_string())?;
5097
5098        let mut priorgrad = Array1::<f64>::zeros(latent_beta.len());
5099        let mut prior_neghess = Array2::<f64>::zeros((latent_beta.len(), latent_beta.len()));
5100        let mut prior_loglik = 0.0;
5101        for term in &self.bounded_terms {
5102            let (logp, grad, neghess, _) =
5103                bounded_prior_terms(latent_beta[term.col_idx], &term.prior);
5104            prior_loglik += logp;
5105            priorgrad[term.col_idx] += grad;
5106            prior_neghess[[term.col_idx, term.col_idx]] += neghess;
5107        }
5108
5109        let mut hessian = xt_diag_x_dense(x_eff.view(), obs.neghessian_eta.view())?;
5110        let mut gradient = fast_atv(&x_eff, &obs.score);
5111        for term in &self.bounded_terms {
5112            let score_beta = self.design.column(term.col_idx).dot(&obs.score);
5113            hessian[[term.col_idx, term.col_idx]] -= score_beta * second_diag[term.col_idx];
5114        }
5115        hessian += &prior_neghess;
5116        gradient += &priorgrad;
5117
5118        Ok((
5119            obs,
5120            hessian,
5121            gradient,
5122            prior_loglik,
5123            second_diag,
5124            third_diag,
5125            priorthird,
5126        ))
5127    }
5128
5129    fn evaluation_from_latent(
5130        &self,
5131        latent_beta: &Array1<f64>,
5132    ) -> Result<
5133        (
5134            StandardFamilyObservationState,
5135            Array2<f64>,
5136            Array1<f64>,
5137            f64,
5138        ),
5139        String,
5140    > {
5141        let (obs, hessian, gradient, prior_loglik, _, _, _) =
5142            self.exacthessian_andgradient(latent_beta)?;
5143        Ok((obs, hessian, gradient, prior_loglik))
5144    }
5145}
5146
5147impl CustomFamily for BoundedLinearFamily {
5148    // Preserve the pre-gam#1395 behavior: the trait default flipped to OFF (the
5149    // flat-prior exact-Newton objective carries no Jeffreys term), so families
5150    // that historically armed the term by default opt back in explicitly.
5151    fn joint_jeffreys_term_required(&self) -> bool {
5152        true
5153    }
5154
5155    fn evaluate(&self, block_states: &[ParameterBlockState]) -> Result<FamilyEvaluation, String> {
5156        let latent_beta = &expect_single_block_state(block_states, "bounded linear family")?.beta;
5157        let (obs, hessian, gradient, prior_loglik) = self.evaluation_from_latent(latent_beta)?;
5158        Ok(FamilyEvaluation {
5159            log_likelihood: obs.log_likelihood + prior_loglik,
5160            blockworking_sets: vec![BlockWorkingSet::ExactNewton {
5161                gradient,
5162                hessian: SymmetricMatrix::Dense(hessian),
5163            }],
5164        })
5165    }
5166
5167    fn exact_newton_joint_hessian(
5168        &self,
5169        block_states: &[ParameterBlockState],
5170    ) -> Result<Option<Array2<f64>>, String> {
5171        let latent_beta = &expect_single_block_state(block_states, "bounded linear family")?.beta;
5172        let (_, hessian, _, _) = self.evaluation_from_latent(latent_beta)?;
5173        Ok(Some(hessian))
5174    }
5175
5176    fn exact_newton_hessian_directional_derivative(
5177        &self,
5178        block_states: &[ParameterBlockState],
5179        block_idx: usize,
5180        d_beta: &Array1<f64>,
5181    ) -> Result<Option<Array2<f64>>, String> {
5182        expect_block_idx_zero(block_idx, "bounded linear family", "")?;
5183        self.exact_newton_joint_hessian_directional_derivative(block_states, d_beta)
5184    }
5185
5186    fn exact_newton_joint_hessian_directional_derivative(
5187        &self,
5188        block_states: &[ParameterBlockState],
5189        d_beta_flat: &Array1<f64>,
5190    ) -> Result<Option<Array2<f64>>, String> {
5191        let latent_beta = &expect_single_block_state(block_states, "bounded linear family")?.beta;
5192        if d_beta_flat.len() != latent_beta.len() {
5193            return Err(SmoothError::dimension_mismatch(format!(
5194                "bounded linear family directional derivative length mismatch: got {}, expected {}",
5195                d_beta_flat.len(),
5196                latent_beta.len()
5197            ))
5198            .into());
5199        }
5200
5201        let (obs, _, _, _, second_diag, third_diag, priorthird) =
5202            self.exacthessian_andgradient(latent_beta)?;
5203
5204        let (_, jac_diag, _, _, _) = self.bounded_term_derivative_data(latent_beta);
5205        let x_eff = self.effective_design_for_latent(&jac_diag);
5206        let deta = x_eff.dot(d_beta_flat);
5207        let d_neghess_eta = &obs.neghessian_eta_derivative * &deta;
5208
5209        let mut dx_eff = Array2::<f64>::zeros(x_eff.raw_dim());
5210        for term in &self.bounded_terms {
5211            let scale = second_diag[term.col_idx] * d_beta_flat[term.col_idx];
5212            if scale != 0.0 {
5213                let mut col = dx_eff.column_mut(term.col_idx);
5214                col.assign(&self.design.column(term.col_idx));
5215                col.mapv_inplace(|v| v * scale);
5216            }
5217        }
5218
5219        let mut dhessian = xt_diag_x_dense(x_eff.view(), d_neghess_eta.view())?;
5220        let mut wxdx = Array2::<f64>::zeros((x_eff.ncols(), x_eff.ncols()));
5221        for i in 0..x_eff.nrows() {
5222            let wi = obs.neghessian_eta[i];
5223            if wi == 0.0 {
5224                continue;
5225            }
5226            for a in 0..x_eff.ncols() {
5227                let xa = x_eff[[i, a]];
5228                for b in 0..x_eff.ncols() {
5229                    wxdx[[a, b]] += wi * (dx_eff[[i, a]] * x_eff[[i, b]] + xa * dx_eff[[i, b]]);
5230                }
5231            }
5232        }
5233        dhessian += &wxdx;
5234
5235        let d_score = -&obs.neghessian_eta * &deta;
5236        for term in &self.bounded_terms {
5237            let score_beta = self.design.column(term.col_idx).dot(&obs.score);
5238            let d_score_beta = self.design.column(term.col_idx).dot(&d_score);
5239            dhessian[[term.col_idx, term.col_idx]] -= d_score_beta * second_diag[term.col_idx]
5240                + score_beta * third_diag[term.col_idx] * d_beta_flat[term.col_idx];
5241            dhessian[[term.col_idx, term.col_idx]] +=
5242                priorthird[term.col_idx] * d_beta_flat[term.col_idx];
5243        }
5244
5245        Ok(Some(dhessian))
5246    }
5247
5248    fn block_geometry(
5249        &self,
5250        block_states: &[ParameterBlockState],
5251        spec: &ParameterBlockSpec,
5252    ) -> Result<(DesignMatrix, Array1<f64>), String> {
5253        if block_states.is_empty() {
5254            return Ok((
5255                DesignMatrix::Dense(gam_linalg::matrix::DenseDesignMatrix::from(
5256                    self.designzeroed.clone(),
5257                )),
5258                self.offset.clone(),
5259            ));
5260        }
5261        let offset = self.nonlinear_offset_from_latent(
5262            &expect_single_block_state(block_states, "bounded linear family")?.beta,
5263        );
5264        let x = if spec.design.ncols() == self.designzeroed.ncols() {
5265            self.designzeroed.clone()
5266        } else {
5267            return Err(SmoothError::dimension_mismatch(
5268                "bounded linear family design column mismatch",
5269            )
5270            .into());
5271        };
5272        Ok((
5273            DesignMatrix::Dense(gam_linalg::matrix::DenseDesignMatrix::from(x)),
5274            offset,
5275        ))
5276    }
5277
5278    fn block_geometry_is_dynamic(&self) -> bool {
5279        true
5280    }
5281
5282    fn block_geometry_directional_derivative(
5283        &self,
5284        block_states: &[ParameterBlockState],
5285        block_idx: usize,
5286        spec: &ParameterBlockSpec,
5287        d_beta: &Array1<f64>,
5288    ) -> Result<Option<BlockGeometryDirectionalDerivative>, String> {
5289        expect_block_idx_zero(
5290            block_idx,
5291            "bounded linear family",
5292            " for geometry derivative",
5293        )?;
5294        expect_single_block_state(block_states, "bounded linear family")?;
5295        if d_beta.len() != spec.design.ncols() {
5296            return Err(SmoothError::dimension_mismatch(format!(
5297                "bounded linear family geometry derivative direction mismatch: got {}, expected {}",
5298                d_beta.len(),
5299                spec.design.ncols()
5300            ))
5301            .into());
5302        }
5303        let (_, jac_diag, _, _, _) = self.bounded_term_derivative_data(&block_states[0].beta);
5304        let mut d_offset = Array1::<f64>::zeros(self.offset.len());
5305        let has_drift = self
5306            .bounded_terms
5307            .iter()
5308            .any(|term| jac_diag[term.col_idx] != 0.0 && d_beta[term.col_idx] != 0.0);
5309        if !has_drift {
5310            return Ok(Some(BlockGeometryDirectionalDerivative {
5311                d_design: None,
5312                d_offset,
5313            }));
5314        }
5315        for term in &self.bounded_terms {
5316            let col = term.col_idx;
5317            let drift = jac_diag[col] * d_beta[col];
5318            if drift != 0.0 {
5319                d_offset.scaled_add(drift, &self.design.column(col));
5320            }
5321        }
5322        Ok(Some(BlockGeometryDirectionalDerivative {
5323            d_design: None,
5324            d_offset,
5325        }))
5326    }
5327}
5328
5329#[inline]
5330fn dense_diag_gram_chunkrows(p: usize) -> usize {
5331    const MIN_ROWS: usize = 512;
5332    const MAX_ROWS: usize = 2048;
5333    const TARGET_BYTES: usize = 2 * 1024 * 1024;
5334    let bytes_per_row = p.max(1) * std::mem::size_of::<f64>();
5335    (TARGET_BYTES / bytes_per_row).clamp(MIN_ROWS, MAX_ROWS)
5336}
5337
5338fn xt_diag_x_dense(x: ArrayView2<'_, f64>, w: ArrayView1<'_, f64>) -> Result<Array2<f64>, String> {
5339    if x.nrows() != w.len() {
5340        return Err(SmoothError::dimension_mismatch("xt_diag_x_dense row mismatch").into());
5341    }
5342    let (n, p) = x.dim();
5343    if n == 0 || p == 0 {
5344        return Ok(Array2::<f64>::zeros((p, p)));
5345    }
5346
5347    const STREAMING_BYTES_THRESHOLD: usize = 8 * 1024 * 1024;
5348    let dense_work_bytes = n
5349        .checked_mul(p)
5350        .and_then(|cells| cells.checked_mul(std::mem::size_of::<f64>()))
5351        .unwrap_or(usize::MAX);
5352    if dense_work_bytes <= STREAMING_BYTES_THRESHOLD {
5353        let mut weighted = x.to_owned();
5354        ndarray::Zip::from(weighted.rows_mut())
5355            .and(w)
5356            .par_for_each(|mut row, wi| row *= *wi);
5357        return Ok(fast_atb(&x, &weighted));
5358    }
5359
5360    let chunkrows = dense_diag_gram_chunkrows(p).min(n);
5361    let mut weighted_chunk = Array2::<f64>::zeros((chunkrows, p));
5362    let mut out = Array2::<f64>::zeros((p, p));
5363    for row_start in (0..n).step_by(chunkrows) {
5364        let rows = (n - row_start).min(chunkrows);
5365        let x_chunk = x.slice(s![row_start..row_start + rows, ..]);
5366        {
5367            let mut chunk = weighted_chunk.slice_mut(s![0..rows, ..]);
5368            for local_row in 0..rows {
5369                let scale = w[row_start + local_row];
5370                if scale == 0.0 {
5371                    chunk.row_mut(local_row).fill(0.0);
5372                    continue;
5373                }
5374                for col in 0..p {
5375                    chunk[[local_row, col]] = x_chunk[[local_row, col]] * scale;
5376                }
5377            }
5378        }
5379        out += &fast_atb(&x_chunk, &weighted_chunk.slice(s![0..rows, ..]));
5380    }
5381    Ok(out)
5382}
5383
5384fn trace_of_dense_product(a: &Array2<f64>, b: &Array2<f64>) -> Result<f64, String> {
5385    if a.nrows() != a.ncols() || b.nrows() != b.ncols() || a.nrows() != b.nrows() {
5386        return Err(
5387            SmoothError::dimension_mismatch("trace_of_dense_product dimension mismatch").into(),
5388        );
5389    }
5390    let mut trace = 0.0;
5391    for i in 0..a.nrows() {
5392        for j in 0..a.ncols() {
5393            trace += a[[i, j]] * b[[j, i]];
5394        }
5395    }
5396    Ok(trace)
5397}
5398
5399fn exact_bounded_edf(
5400    penalties: &[PenaltySpec],
5401    lambdas: &Array1<f64>,
5402    latent_cov: &Array2<f64>,
5403) -> Result<(Vec<f64>, Vec<f64>, f64), EstimationError> {
5404    if penalties.len() != lambdas.len() {
5405        crate::bail_invalid_estim!(
5406            "bounded EDF penalty/lambda mismatch: {} penalties vs {} lambdas",
5407            penalties.len(),
5408            lambdas.len()
5409        );
5410    }
5411    if latent_cov.nrows() != latent_cov.ncols() {
5412        crate::bail_invalid_estim!("bounded EDF covariance must be square");
5413    }
5414
5415    let p = latent_cov.nrows();
5416    let mut s_lambda = Array2::<f64>::zeros((p, p));
5417    let mut edf_by_block = Vec::with_capacity(penalties.len());
5418    // Raw per-block penalty trace tr_kk = λ_kk·tr(H⁻¹S_kk) (issue #1219).
5419    let mut penalty_block_trace = Vec::with_capacity(penalties.len());
5420    let mut trace_sum = 0.0;
5421
5422    for (k, ps) in penalties.iter().enumerate() {
5423        let lambda_k = lambdas[k];
5424        match ps {
5425            PenaltySpec::Block {
5426                local, col_range, ..
5427            } => {
5428                s_lambda
5429                    .slice_mut(ndarray::s![col_range.clone(), col_range.clone()])
5430                    .scaled_add(lambda_k, local);
5431                // Compute penalty rank from the block-local matrix directly.
5432                let penalty_rank =
5433                    local
5434                        .nrows()
5435                        .saturating_sub(estimate_penalty_nullity(local).map_err(|e| {
5436                            EstimationError::InvalidInput(format!("bounded EDF rank failed: {e}"))
5437                        })?);
5438                // Trace only involves the block slice of latent_cov.
5439                let cov_block = latent_cov.slice(ndarray::s![col_range.clone(), col_range.clone()]);
5440                let trace_k = lambda_k
5441                    * trace_of_dense_product(&cov_block.to_owned(), local)
5442                        .map_err(EstimationError::InvalidInput)?;
5443                trace_sum += trace_k;
5444                penalty_block_trace.push(trace_k);
5445                let p_k = penalty_rank as f64;
5446                edf_by_block.push((p_k - trace_k).clamp(0.0, p_k));
5447            }
5448            PenaltySpec::Dense(m) | PenaltySpec::DenseWithMean { matrix: m, .. } => {
5449                s_lambda.scaled_add(lambda_k, m);
5450                let penalty_rank = p.saturating_sub(estimate_penalty_nullity(m).map_err(|e| {
5451                    EstimationError::InvalidInput(format!("bounded EDF rank failed: {e}"))
5452                })?);
5453                let trace_k = lambda_k
5454                    * trace_of_dense_product(latent_cov, m)
5455                        .map_err(EstimationError::InvalidInput)?;
5456                trace_sum += trace_k;
5457                penalty_block_trace.push(trace_k);
5458                let p_k = penalty_rank as f64;
5459                edf_by_block.push((p_k - trace_k).clamp(0.0, p_k));
5460            }
5461        }
5462    }
5463
5464    let nullity_total = estimate_penalty_nullity(&s_lambda)
5465        .map_err(|e| EstimationError::InvalidInput(format!("bounded EDF nullity failed: {e}")))?
5466        as f64;
5467    let edf_total = (p as f64 - trace_sum).clamp(nullity_total, p as f64);
5468    Ok((edf_by_block, penalty_block_trace, edf_total))
5469}
5470
5471/// Symmetric posterior-precision inverse for the bounded-coefficient path.
5472///
5473/// The penalised Hessian at a strict posterior maximum is SPD, so its inverse
5474/// is the posterior covariance. We eigendecompose the symmetric precision and
5475/// invert the positive-eigenvalue subspace, projecting out the (rare)
5476/// structural null directions a penalised model leaves flat rather than
5477/// δ-ridging them — the same honest pseudo-inverse contract the strict
5478/// pseudo-Laplace covariance uses (gam#748). A genuinely indefinite precision
5479/// (a negative eigenvalue beyond rounding) means the reported mode is not a
5480/// posterior maximum and is surfaced as a fit-quality error rather than
5481/// masked.
5482fn symmetric_positive_definite_inverse_or_pseudo(
5483    precision: &Array2<f64>,
5484) -> Result<Array2<f64>, EstimationError> {
5485    use gam_linalg::faer_ndarray::FaerEigh;
5486    let p = precision.nrows();
5487    if precision.ncols() != p {
5488        crate::bail_invalid_estim!(
5489            "posterior precision inverse requires a square matrix, got {}x{}",
5490            precision.nrows(),
5491            precision.ncols()
5492        );
5493    }
5494    if p == 0 {
5495        return Ok(Array2::<f64>::zeros((0, 0)));
5496    }
5497    let symmetric = (precision + &precision.t().to_owned()) * 0.5;
5498    let (evals, evecs) = symmetric.eigh(faer::Side::Lower).map_err(|e| {
5499        EstimationError::InvalidInput(format!(
5500            "posterior precision eigendecomposition failed: {e}"
5501        ))
5502    })?;
5503    let max_abs_eval = evals.iter().fold(0.0_f64, |acc, &ev| acc.max(ev.abs()));
5504    let tol =
5505        (10.0 * f64::EPSILON * (p as f64) * (p as f64) * max_abs_eval).max(100.0 * f64::EPSILON);
5506    if let Some(&min_eval) = evals
5507        .iter()
5508        .filter(|&&ev| ev < -tol)
5509        .min_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal))
5510    {
5511        crate::bail_invalid_estim!(
5512            "bounded posterior precision is non-PD at the converged optimum (min eigenvalue \
5513             {min_eval:.6e} < -tol={tol:.6e}); the reported mode is not a strict posterior \
5514             maximum, so a covariance would be meaningless"
5515        );
5516    }
5517    // Σ = U diag(1/λ_+) Uᵀ over the positive-eigenvalue subspace.
5518    let mut scaled = evecs.clone();
5519    for (j, &ev) in evals.iter().enumerate() {
5520        let inv = if ev > tol { 1.0 / ev } else { 0.0 };
5521        scaled.column_mut(j).mapv_inplace(|v| v * inv);
5522    }
5523    let cov = scaled.dot(&evecs.t());
5524    Ok((&cov + &cov.t().to_owned()) * 0.5)
5525}
5526
5527fn transform_bounded_latent_precision_to_user_internal(
5528    latent_precision: &Array2<f64>,
5529    jac_diag: &Array1<f64>,
5530) -> Result<Array2<f64>, EstimationError> {
5531    let p = latent_precision.nrows();
5532    if latent_precision.ncols() != p || jac_diag.len() != p {
5533        crate::bail_invalid_estim!(
5534            "bounded precision transform dimension mismatch: precision is {}x{}, jacobian has {} entries",
5535            latent_precision.nrows(),
5536            latent_precision.ncols(),
5537            jac_diag.len()
5538        );
5539    }
5540    let mut out = latent_precision.clone();
5541    for i in 0..p {
5542        let scale = jac_diag[i];
5543        if !scale.is_finite() || scale <= 0.0 {
5544            crate::bail_invalid_estim!(
5545                "bounded precision transform requires a positive finite coefficient jacobian; column {i} has {scale}"
5546            );
5547        }
5548        if scale != 1.0 {
5549            out.row_mut(i).mapv_inplace(|v| v / scale);
5550            out.column_mut(i).mapv_inplace(|v| v / scale);
5551        }
5552    }
5553    Ok(out)
5554}
5555
5556fn fit_bounded_term_collection_with_design(
5557    y: ArrayView1<'_, f64>,
5558    weights: ArrayView1<'_, f64>,
5559    offset: ArrayView1<'_, f64>,
5560    spec: &TermCollectionSpec,
5561    design: &TermCollectionDesign,
5562    heuristic_lambdas: Option<&[f64]>,
5563    family: LikelihoodSpec,
5564    options: &FitOptions,
5565) -> Result<FittedTermCollection, EstimationError> {
5566    let conditioning_cols: Vec<usize> = spec
5567        .linear_terms
5568        .iter()
5569        .enumerate()
5570        .filter_map(|(j, linear)| {
5571            (!linear.double_penalty).then_some(design.intercept_range.end + j)
5572        })
5573        .collect();
5574    let conditioning = LinearFitConditioning::from_columns(design, &conditioning_cols);
5575    let dense_design = design.design.to_dense_cow();
5576    let fit_design = conditioning.apply_to_design(&dense_design);
5577    let fit_penalties = conditioning
5578        .transform_blockwise_penalties_to_internal(&design.penalties, design.design.ncols());
5579    if design.linear_constraints.is_some() {
5580        crate::bail_invalid_estim!(
5581            "bounded() terms are not yet compatible with explicit linear constraints"
5582        );
5583    }
5584    let mut bounded_terms = Vec::<BoundedLinearTermMeta>::new();
5585    for (j, term) in spec.linear_terms.iter().enumerate() {
5586        if term.double_penalty
5587            && matches!(
5588                term.coefficient_geometry,
5589                LinearCoefficientGeometry::Bounded { .. }
5590            )
5591        {
5592            crate::bail_invalid_estim!(
5593                "bounded linear term '{}' cannot also use double_penalty",
5594                term.name
5595            );
5596        }
5597        if let LinearCoefficientGeometry::Bounded { min, max, prior } =
5598            term.coefficient_geometry.clone()
5599        {
5600            let col_idx = design.intercept_range.end + j;
5601            let (min_internal, max_internal) = conditioning.internal_bounds_for(col_idx, min, max);
5602            bounded_terms.push(BoundedLinearTermMeta {
5603                col_idx,
5604                min: min_internal,
5605                max: max_internal,
5606                prior,
5607            });
5608        }
5609    }
5610    if bounded_terms.is_empty() {
5611        crate::bail_invalid_estim!("internal bounded fit path called with no bounded terms");
5612    }
5613
5614    let mut designzeroed = fit_design.clone();
5615    let mut initial_beta = Array1::<f64>::zeros(fit_design.ncols());
5616    for term in &bounded_terms {
5617        designzeroed.column_mut(term.col_idx).fill(0.0);
5618        initial_beta[term.col_idx] = bounded_logit(0.5);
5619    }
5620
5621    let initial_log_lambdas = heuristic_lambdas
5622        .map(|vals| Array1::from_vec(vals.to_vec()))
5623        .unwrap_or_else(|| Array1::zeros(fit_penalties.len()));
5624    if initial_log_lambdas.len() != fit_penalties.len() {
5625        crate::bail_invalid_estim!(
5626            "heuristic lambda length mismatch for bounded model: got {}, expected {}",
5627            initial_log_lambdas.len(),
5628            fit_penalties.len()
5629        );
5630    }
5631
5632    let is_beta_logistic = family.is_binomial_beta_logistic();
5633    let family_adapter = BoundedLinearFamily {
5634        family: family.clone(),
5635        latent_cloglog_state: options.latent_cloglog,
5636        mixture_link_state: options
5637            .mixture_link
5638            .clone()
5639            .as_ref()
5640            .map(state_fromspec)
5641            .transpose()
5642            .map_err(EstimationError::InvalidInput)?,
5643        sas_link_state: options
5644            .sas_link
5645            .map(|spec| {
5646                if is_beta_logistic {
5647                    state_from_beta_logisticspec(spec)
5648                } else {
5649                    state_from_sasspec(spec)
5650                }
5651            })
5652            .transpose()
5653            .map_err(EstimationError::InvalidInput)?,
5654        y: y.to_owned(),
5655        weights: weights.to_owned(),
5656        design: fit_design.clone(),
5657        designzeroed: designzeroed.clone(),
5658        offset: offset.to_owned(),
5659        bounded_terms: bounded_terms.clone(),
5660    };
5661    let blockspec = ParameterBlockSpec {
5662        name: "eta".to_string(),
5663        design: DesignMatrix::Dense(gam_linalg::matrix::DenseDesignMatrix::from(designzeroed)),
5664        offset: offset.to_owned(),
5665        penalties: fit_penalties
5666            .iter()
5667            .map(|ps| match ps {
5668                PenaltySpec::Block {
5669                    local, col_range, ..
5670                } => PenaltyMatrix::Blockwise {
5671                    local: local.clone(),
5672                    col_range: col_range.clone(),
5673                    total_dim: design.design.ncols(),
5674                },
5675                PenaltySpec::Dense(m) | PenaltySpec::DenseWithMean { matrix: m, .. } => {
5676                    PenaltyMatrix::Dense(m.clone())
5677                }
5678            })
5679            .collect(),
5680        nullspace_dims: design.nullspace_dims.clone(),
5681        initial_log_lambdas,
5682        initial_beta: Some(initial_beta),
5683        gauge_priority: 100,
5684        // Report the true β-dependent Jacobian (bounded columns scaled by
5685        // dβ/dθ) to the identifiability audit so it does not mistake the
5686        // deliberately-zeroed placeholder columns for a structural rank
5687        // deficiency. The inner solve still drives η through the family
5688        // adapter, so this does not affect the fit geometry.
5689        jacobian_callback: Some(Arc::new(BoundedEffectiveJacobian {
5690            design: fit_design.clone(),
5691            bounded_terms: bounded_terms.clone(),
5692        })),
5693        stacked_design: None,
5694        stacked_offset: None,
5695    };
5696    let fit = fit_custom_family(
5697        &family_adapter,
5698        &[blockspec],
5699        &BlockwiseFitOptions {
5700            inner_max_cycles: options.max_iter,
5701            inner_tol: options.tol,
5702            outer_max_iter: options.max_iter,
5703            outer_tol: options.tol,
5704            // The bounded path builds its own user-scale covariance below by
5705            // inverting the user-scale penalised Hessian (delta-method through
5706            // the bounded transform's Jacobian + the conditioning map), so it
5707            // does not consume the inner solver's optional canonical-space
5708            // `covariance_conditional`. Inverting the reported precision
5709            // directly guarantees `inv(penalized_hessian) == covariance` and
5710            // works on every bounded fit — including the common no-smoothing
5711            // path where the inner solve surfaces no covariance at all (the
5712            // gam#854 "bounded fit emits no user-scale covariance" symptom).
5713            compute_covariance: false,
5714            ..BlockwiseFitOptions::default()
5715        },
5716    )
5717    .map_err(EstimationError::CustomFamily)?;
5718
5719    let latent_beta = fit.block_states[0].beta.clone();
5720    let (beta_user_internal, jac_diag) = family_adapter.user_beta_and_jacobian(&latent_beta);
5721    let beta_user = conditioning.backtransform_beta(&beta_user_internal);
5722
5723    let (eta_state, h_data, _, _) = family_adapter
5724        .evaluation_from_latent(&latent_beta)
5725        .map_err(EstimationError::InvalidInput)?;
5726    let p_fit = fit_design.ncols();
5727    let mut s_lambda_internal = Array2::<f64>::zeros((p_fit, p_fit));
5728    for (k, penalty) in fit_penalties.iter().enumerate() {
5729        match penalty {
5730            PenaltySpec::Block {
5731                local, col_range, ..
5732            } => {
5733                s_lambda_internal
5734                    .slice_mut(ndarray::s![col_range.clone(), col_range.clone()])
5735                    .scaled_add(fit.lambdas[k], local);
5736            }
5737            PenaltySpec::Dense(m) | PenaltySpec::DenseWithMean { matrix: m, .. } => {
5738                s_lambda_internal.scaled_add(fit.lambdas[k], m);
5739            }
5740        }
5741    }
5742    let mut latent_precision = h_data.clone();
5743    latent_precision += &s_lambda_internal;
5744    let user_precision_internal =
5745        transform_bounded_latent_precision_to_user_internal(&latent_precision, &jac_diag)?;
5746    let penalized_hessian =
5747        conditioning.transform_penalized_hessian_to_original(&user_precision_internal);
5748
5749    // User-scale posterior covariance via the delta method. The reported
5750    // geometry precision `penalized_hessian` is the user-scale penalized
5751    // Hessian `H_user = C⁻ᵀ J⁻¹ (H_latent + S_λ) J⁻¹ C⁻¹` (latent precision
5752    // pushed through the bounded transform's Jacobian `J = diag(dβ_user/dθ)`
5753    // and the conditioning map `C`). Its exact inverse `H_user⁻¹` is the
5754    // delta-method pushforward of the latent posterior precision-inverse
5755    // `(H_latent + S_λ)⁻¹` — but on the UNSCALED (unit-dispersion) scale. For a
5756    // free-dispersion family (profiled Gaussian) the reported coefficient
5757    // covariance is `Vb = φ̂ · H_user⁻¹` with `φ̂ = σ̂²`, so the unscaled inverse
5758    // below is multiplied by the dispersion scale `cov_scale` once `σ̂²` is
5759    // known (after the EDF, which sets the residual d.f.). For fixed-scale
5760    // families (Binomial, `φ ≡ 1`) `cov_scale == 1` and `Vb = H_user⁻¹`
5761    // unchanged. Skipping this scale was gam#1514: an interior, well-identified
5762    // Gaussian bounded slope reported an SE ≈ 1/√Σ(xᵢ−x̄)² instead of
5763    // σ̂/√Σ(xᵢ−x̄)², i.e. ~`1/σ̂` (≈20×) too wide.
5764    //
5765    // Inverting the same matrix the geometry reports keeps
5766    // `inv(penalized_hessian) == cov_scale⁻¹ · covariance` and removes the
5767    // dependency on the inner solver's optional, canonical-space
5768    // `covariance_conditional` (which is `None` whenever the bounded blockspec
5769    // carries no smoothing parameters — the no-rho fit path — leaving a bounded
5770    // fit with a populated precision but no user-scale covariance, the gam#854
5771    // symptom). The latent precision is SPD at a strict posterior maximum; on a
5772    // marginally-indefinite boundary Hessian we invert the positive-eigenvalue
5773    // subspace (the structural null space of a penalised model is a flat
5774    // posterior direction, not something to ridge away), matching the
5775    // strict-pseudo-Laplace covariance contract (gam#748).
5776    let beta_covariance_unscaled = if options.compute_inference {
5777        Some(symmetric_positive_definite_inverse_or_pseudo(
5778            &penalized_hessian,
5779        )?)
5780    } else {
5781        None
5782    };
5783    // EDF `p − Σ_k λ_k tr(H_latent⁻¹ S_k)` is computed in the *latent*
5784    // (untransformed) coordinate system the penalties `fit_penalties` live in,
5785    // so it needs the latent posterior covariance `(H_latent + S_λ)⁻¹`, not the
5786    // user-scale one. Invert the same latent precision that produced the
5787    // reported user precision so the two are an exact transform pair.
5788    let latent_cov = if options.compute_inference {
5789        Some(symmetric_positive_definite_inverse_or_pseudo(
5790            &latent_precision,
5791        )?)
5792    } else {
5793        None
5794    };
5795    let s_lambda_original = weighted_blockwise_penalty_sum(
5796        &design.penalties,
5797        fit.lambdas.as_slice().unwrap(),
5798        design.design.ncols(),
5799    );
5800    let penalty_term = beta_user.dot(&s_lambda_original.dot(&beta_user));
5801    let deviance = if family.is_gaussian_identity() {
5802        y.iter()
5803            .zip(eta_state.mu.iter())
5804            .zip(weights.iter())
5805            .map(|((&yy, &mu), &w)| w.max(0.0) * (yy - mu) * (yy - mu))
5806            .sum()
5807    } else {
5808        -2.0 * eta_state.log_likelihood
5809    };
5810    let (edf_by_block, penalty_block_trace, edf_total) = if let Some(cov) = latent_cov.as_ref() {
5811        exact_bounded_edf(&fit_penalties, &fit.lambdas, cov)?
5812    } else {
5813        (
5814            vec![0.0; fit_penalties.len()],
5815            vec![0.0; fit_penalties.len()],
5816            0.0,
5817        )
5818    };
5819
5820    // Dispersion. The bounded fit's working weight is scale-free for a profiled
5821    // Gaussian (`W = priorweights`), so the unscaled penalized Hessian carries
5822    // unit implicit dispersion and the reported coefficient covariance must be
5823    // restored to `Vb = σ̂²·H_user⁻¹` with the REML residual variance
5824    // `σ̂² = RSS/(n − edf_total)` — identical to the ordinary GAM path
5825    // (`solver/estimate/optimizer.rs`). Fixed-scale families (Binomial here,
5826    // `φ ≡ 1`) keep their full Fisher information in `W`, so `cov_scale == 1`
5827    // and the covariance is `H_user⁻¹` unscaled. The single source of truth for
5828    // the per-family scale is `GlmLikelihoodSpec::coefficient_covariance_scale`
5829    // / `dispersion_from_likelihood`, reused verbatim so the bounded path can
5830    // never drift from the standard contract (gam#1514).
5831    let glm_likelihood = gam_spec::GlmLikelihoodSpec::canonical(family.clone());
5832    let standard_deviation = if family.is_gaussian_identity() {
5833        let denom = if options.compute_inference {
5834            (y.len() as f64 - edf_total).max(1.0)
5835        } else {
5836            (y.len() as f64).max(1.0)
5837        };
5838        (deviance / denom).sqrt()
5839    } else {
5840        1.0
5841    };
5842    let cov_scale = glm_likelihood
5843        .coefficient_covariance_scale(standard_deviation * standard_deviation)
5844        .max(f64::MIN_POSITIVE);
5845    let dispersion = gam_solve::estimate::dispersion_from_likelihood(&glm_likelihood, standard_deviation);
5846    // Apply the dispersion scale to the unscaled inverse, producing the reported
5847    // `Vb = cov_scale · H_user⁻¹` and its diagonal standard errors. The stored
5848    // `penalized_hessian` stays UNSCALED (`H_user`) per the dispersion-ownership
5849    // contract in `inference::dispersion_cov`; the sampler re-applies `√cov_scale`
5850    // when it reconstructs the latent posterior (see `sample_standard_bounded`).
5851    let beta_covariance = beta_covariance_unscaled.map(|mut cov| {
5852        if cov_scale != 1.0 {
5853            cov.mapv_inplace(|v| v * cov_scale);
5854        }
5855        cov
5856    });
5857    let beta_standard_errors = beta_covariance
5858        .as_ref()
5859        .map(|cov| Array1::from_iter((0..cov.nrows()).map(|i| cov[[i, i]].max(0.0).sqrt())));
5860
5861    let geometry = Some(gam_solve::estimate::FitGeometry {
5862        penalized_hessian: penalized_hessian.clone().into(),
5863        working_weights: eta_state.fisherweight.clone(),
5864        working_response: {
5865            let mut working_response = eta_state.eta.clone();
5866            for i in 0..working_response.len() {
5867                let wi = eta_state.fisherweight[i].max(1e-12);
5868                working_response[i] += eta_state.score[i] / wi;
5869            }
5870            working_response
5871        },
5872    });
5873    let max_abs_eta = eta_state
5874        .eta
5875        .iter()
5876        .fold(0.0_f64, |acc, &v| acc.max(v.abs()));
5877    Ok(FittedTermCollection {
5878        fit: {
5879            let log_lambdas = fit.lambdas.mapv(|v| v.max(1e-300).ln());
5880            let inf = FitInference {
5881                edf_by_block,
5882                penalty_block_trace,
5883                edf_total,
5884                smoothing_correction: None,
5885                // Boundary adapter: `penalized_hessian` storage is now
5886                // `UnscaledPrecision`.
5887                penalized_hessian: penalized_hessian.clone().into(),
5888                working_weights: eta_state.fisherweight.clone(),
5889                working_response: {
5890                    let mut working_response = eta_state.eta.clone();
5891                    for i in 0..working_response.len() {
5892                        let wi = eta_state.fisherweight[i].max(1e-12);
5893                        working_response[i] += eta_state.score[i] / wi;
5894                    }
5895                    working_response
5896                },
5897                reparam_qs: None,
5898                dispersion,
5899                beta_covariance: beta_covariance
5900                    .clone()
5901                    .map(gam_problem::dispersion_cov::PhiScaledCovariance::from),
5902                beta_standard_errors,
5903                beta_covariance_corrected: None,
5904                beta_standard_errors_corrected: None,
5905                beta_covariance_frequentist: None,
5906                coefficient_influence: None,
5907                weighted_gram: None,
5908                bias_correction_beta: None,
5909                bias_correction_jacobian: None,
5910            };
5911            let covariance_conditional = beta_covariance;
5912            // Sealed `UnifiedFitResult`: existence certifies inner+outer
5913            // convergence (see `try_from_parts`), so the status is Converged.
5914            let pirls_status_val = gam_solve::pirls::PirlsStatus::Converged;
5915            UnifiedFitResult::try_from_parts(UnifiedFitResultParts {
5916                blocks: vec![gam_solve::estimate::FittedBlock {
5917                    beta: beta_user.clone(),
5918                    role: gam_problem::BlockRole::Mean,
5919                    edf: edf_total,
5920                    lambdas: fit.lambdas.clone(),
5921                }],
5922                log_lambdas,
5923                lambdas: fit.lambdas,
5924                likelihood_scale: family.default_scale_metadata(),
5925                likelihood_family: Some(family),
5926                log_likelihood_normalization:
5927                    gam_spec::LogLikelihoodNormalization::UserProvided,
5928                log_likelihood: eta_state.log_likelihood,
5929                deviance,
5930                reml_score: fit.penalized_objective,
5931                stable_penalty_term: penalty_term,
5932                penalized_objective: fit.penalized_objective,
5933                used_device: false,
5934                outer_iterations: fit.outer_iterations,
5935                // Sealed result ⇒ outer convergence was certified at assembly.
5936                outer_converged: true,
5937                outer_gradient_norm: fit.outer_gradient_norm,
5938                standard_deviation,
5939                covariance_conditional,
5940                covariance_corrected: None,
5941                inference: Some(inf),
5942                fitted_link: gam_solve::estimate::FittedLinkState::Standard(None),
5943                geometry,
5944                block_states: Vec::new(),
5945                pirls_status: pirls_status_val,
5946                max_abs_eta,
5947                constraint_kkt: None,
5948                artifacts: gam_solve::estimate::FitArtifacts {
5949                    pirls: None,
5950                    ..Default::default()
5951                },
5952                inner_cycles: 0,
5953            })?
5954        },
5955        design: design.clone(),
5956        adaptive_diagnostics: None,
5957    })
5958}
5959
5960fn enforce_term_constraint_feasibility(
5961    design: &TermCollectionDesign,
5962    fit: &UnifiedFitResult,
5963) -> Result<(), EstimationError> {
5964    // Geometric (per-row-scaled) tolerance, matching the public contract on
5965    // `ACTIVE_SET_PRIMAL_FEASIBILITY_TOL` and the diagnostic that
5966    // `compute_constraint_kkt_diagnostics` exposes via `fit.constraint_kkt`.
5967    // Lower-bound rows are unit-norm (a_i = e_i) so the scale-invariant and
5968    // raw checks coincide there. Linear-inequality rows generally are NOT
5969    // unit-norm — e.g. a B-spline endpoint-derivative clamp at k = 12 carries
5970    // ‖a_i‖ ≈ 38, so a 1e-6 raw residual is only 2.6e-8 in geometric units.
5971    // Holding this gate to raw 1e-7 while the in-solver acceptance gate
5972    // measures geometric 1e-8 is the inconsistency that made well-conditioned
5973    // clamped fits get rejected after they completed cleanly.
5974    /// Raw (unscaled) constraint-residual tolerance for the post-fit feasibility
5975    /// audit; kept loose enough to be consistent with the geometric in-solver
5976    /// acceptance gate on non-unit-norm linear-inequality rows (see comment).
5977    const CONSTRAINT_FEASIBILITY_RAW_TOL: f64 = 1e-7;
5978    let tol = CONSTRAINT_FEASIBILITY_RAW_TOL;
5979    let smooth_start = design
5980        .design
5981        .ncols()
5982        .saturating_sub(design.smooth.total_smooth_cols());
5983    let mut violations: Vec<String> = Vec::new();
5984    for term in &design.smooth.terms {
5985        let gr = (smooth_start + term.coeff_range.start)..(smooth_start + term.coeff_range.end);
5986        let beta_local = fit.beta.slice(s![gr.clone()]).to_owned();
5987        if let Some(lb) = term.lower_bounds_local.as_ref() {
5988            let mut worst = 0.0_f64;
5989            let mut worst_idx = 0usize;
5990            for i in 0..lb.len().min(beta_local.len()) {
5991                if lb[i].is_finite() {
5992                    let viol = (lb[i] - beta_local[i]).max(0.0);
5993                    if viol > worst {
5994                        worst = viol;
5995                        worst_idx = i;
5996                    }
5997                }
5998            }
5999            if worst > tol {
6000                violations.push(format!(
6001                    "term='{}' kind=lower-bound maxviolation={:.3e} coeff_index={}",
6002                    term.name, worst, worst_idx
6003                ));
6004            }
6005        }
6006        if let Some(lin) = term.linear_constraints_local.as_ref() {
6007            let mut worst = 0.0_f64;
6008            let mut worstrow = 0usize;
6009            for i in 0..lin.a.nrows() {
6010                let norm = lin.a.row(i).dot(&lin.a.row(i)).sqrt();
6011                let inv = if norm > 0.0 { 1.0 / norm } else { 0.0 };
6012                let s = (lin.a.row(i).dot(&beta_local) - lin.b[i]) * inv;
6013                let viol = (-s).max(0.0);
6014                if viol > worst {
6015                    worst = viol;
6016                    worstrow = i;
6017                }
6018            }
6019            if worst > tol {
6020                violations.push(format!(
6021                    "term='{}' kind=linear-inequality maxviolation={:.3e} row={}",
6022                    term.name, worst, worstrow
6023                ));
6024            }
6025        }
6026    }
6027
6028    if !violations.is_empty() {
6029        let mut msg = format!(
6030            "constraint violation after fit ({} violating term constraints): {}",
6031            violations.len(),
6032            violations.join(" | ")
6033        );
6034        if let Some(kkt) = fit.constraint_kkt.as_ref() {
6035            msg.push_str(&format!(
6036                "; KKT[primal={:.3e}, dual={:.3e}, comp={:.3e}, stat={:.3e}]",
6037                kkt.primal_feasibility, kkt.dual_feasibility, kkt.complementarity, kkt.stationarity
6038            ));
6039        }
6040        return Err(EstimationError::ParameterConstraintViolation(msg));
6041    }
6042    Ok(())
6043}
6044
6045fn stratified_spatial_subsample(
6046    data: ArrayView2<'_, f64>,
6047    spec: &TermCollectionSpec,
6048    target_size: usize,
6049) -> Vec<usize> {
6050    use rand::SeedableRng;
6051    use rand::rngs::StdRng;
6052    use rand::seq::SliceRandom;
6053
6054    let n = data.nrows();
6055    if n <= target_size {
6056        return (0..n).collect();
6057    }
6058
6059    let spatial_cols: Option<Vec<usize>> =
6060        spec.smooth_terms.iter().find_map(|term| match &term.basis {
6061            SmoothBasisSpec::ThinPlate { feature_cols, .. }
6062            | SmoothBasisSpec::Matern { feature_cols, .. }
6063            | SmoothBasisSpec::Duchon { feature_cols, .. } => {
6064                if !feature_cols.is_empty() {
6065                    Some(feature_cols.clone())
6066                } else {
6067                    None
6068                }
6069            }
6070            _ => None,
6071        });
6072
6073    let cols = match spatial_cols {
6074        Some(c) if !c.is_empty() => c,
6075        _ => {
6076            let mut rng = StdRng::seed_from_u64(spatial_subsample_seed(data, &[], target_size));
6077            let mut indices: Vec<usize> = (0..n).collect();
6078            indices.shuffle(&mut rng);
6079            indices.truncate(target_size);
6080            indices.sort_unstable();
6081            return indices;
6082        }
6083    };
6084    let mut rng = StdRng::seed_from_u64(spatial_subsample_seed(data, &cols, target_size));
6085
6086    let d = cols.len();
6087    let mut mins = vec![f64::INFINITY; d];
6088    let mut maxs = vec![f64::NEG_INFINITY; d];
6089    for i in 0..n {
6090        for (ax, &col) in cols.iter().enumerate() {
6091            let v = data[[i, col]];
6092            if v < mins[ax] {
6093                mins[ax] = v;
6094            }
6095            if v > maxs[ax] {
6096                maxs[ax] = v;
6097            }
6098        }
6099    }
6100
6101    // Aim for roughly this many sampled points per stratification cell so each
6102    // occupied cell can contribute a representative draw without collapsing the
6103    // grid to one point per cell.
6104    const TARGET_POINTS_PER_CELL: usize = 5;
6105    let total_cells_target = (target_size / TARGET_POINTS_PER_CELL).max(1);
6106    let cells_per_axis = ((total_cells_target as f64).powf(1.0 / d as f64)).ceil() as usize;
6107    let cells_per_axis = cells_per_axis.max(1);
6108
6109    let mut cell_members: std::collections::HashMap<Vec<usize>, Vec<usize>> =
6110        std::collections::HashMap::new();
6111    for i in 0..n {
6112        let mut cell_key = Vec::with_capacity(d);
6113        for (ax, &col) in cols.iter().enumerate() {
6114            let range = maxs[ax] - mins[ax];
6115            let cell = if range <= 0.0 {
6116                0
6117            } else {
6118                let frac = (data[[i, col]] - mins[ax]) / range;
6119                (frac * cells_per_axis as f64).floor() as usize
6120            };
6121            cell_key.push(cell.min(cells_per_axis - 1));
6122        }
6123        cell_members.entry(cell_key).or_default().push(i);
6124    }
6125
6126    let mut selected: Vec<usize> = Vec::with_capacity(target_size);
6127    let mut remaining_budget = target_size;
6128    let mut remaining_population = n;
6129
6130    let mut cells: Vec<(Vec<usize>, Vec<usize>)> = cell_members.into_iter().collect();
6131    cells.sort_by(|a, b| a.0.cmp(&b.0));
6132
6133    for (_, members) in &mut cells {
6134        if remaining_budget == 0 {
6135            break;
6136        }
6137        let alloc = ((members.len() as f64 / remaining_population as f64) * remaining_budget as f64)
6138            .round() as usize;
6139        let alloc = alloc.max(1).min(members.len()).min(remaining_budget);
6140        members.shuffle(&mut rng);
6141        selected.extend_from_slice(&members[..alloc]);
6142        remaining_budget = remaining_budget.saturating_sub(alloc);
6143        remaining_population = remaining_population.saturating_sub(members.len());
6144    }
6145
6146    if selected.len() > target_size {
6147        selected.shuffle(&mut rng);
6148        selected.truncate(target_size);
6149    }
6150
6151    selected.sort_unstable();
6152    selected
6153}
6154
6155fn spatial_subsample_seed(
6156    data: ArrayView2<'_, f64>,
6157    spatial_cols: &[usize],
6158    target_size: usize,
6159) -> u64 {
6160    let mut state = 0x5350_4154_4941_4C53_u64;
6161    spatial_seed_mix(&mut state, data.nrows() as u64);
6162    spatial_seed_mix(&mut state, data.ncols() as u64);
6163    spatial_seed_mix(&mut state, target_size as u64);
6164    spatial_seed_mix(&mut state, spatial_cols.len() as u64);
6165    for &col in spatial_cols {
6166        spatial_seed_mix(&mut state, col as u64);
6167    }
6168
6169    if data.nrows() > 0 {
6170        let mid = data.nrows() / 2;
6171        let last = data.nrows() - 1;
6172        for &row in &[0usize, mid, last] {
6173            for &col in spatial_cols {
6174                let value = data[[row, col]];
6175                spatial_seed_mix(&mut state, value.to_bits());
6176            }
6177        }
6178    }
6179    state
6180}
6181
6182#[inline]
6183fn spatial_seed_mix(state: &mut u64, value: u64) {
6184    // Canonical SplitMix64 step over `value + state` (the step adds G itself),
6185    // then an extra rotate-multiply avalanche unique to the spatial seed mix.
6186    let mut s = value.wrapping_add(*state);
6187    let z = gam_linalg::utils::splitmix64(&mut s);
6188    *state ^= z;
6189    *state = (*state).rotate_left(27).wrapping_mul(0x3C79_AC49_2BA7_B653);
6190}
6191
6192fn sampled_rows(data: ArrayView2<'_, f64>, indices: &[usize]) -> Array2<f64> {
6193    let mut sampled = Array2::<f64>::zeros((indices.len(), data.ncols()));
6194    for (new_row, &orig_row) in indices.iter().enumerate() {
6195        sampled.row_mut(new_row).assign(&data.row(orig_row));
6196    }
6197    sampled
6198}
6199
6200fn spatial_term_user_centers(term: &SmoothTermSpec) -> Option<ArrayView2<'_, f64>> {
6201    match spatial_term_center_strategy(term) {
6202        Some(CenterStrategy::UserProvided(centers)) => Some(centers.view()),
6203        _ => None,
6204    }
6205}
6206
6207fn finite_centered_axis_contrasts(values: &[f64], expected_dim: usize) -> Option<Vec<f64>> {
6208    if values.len() != expected_dim || expected_dim <= 1 {
6209        return None;
6210    }
6211    if values.iter().any(|value| !value.is_finite()) {
6212        return None;
6213    }
6214    Some(center_aniso_log_scales(values))
6215}
6216
6217fn blended_pilot_axis_contrasts(
6218    pilot_data: ArrayView2<'_, f64>,
6219    term: &SmoothTermSpec,
6220    centers: ArrayView2<'_, f64>,
6221) -> Option<Vec<f64>> {
6222    let d = centers.ncols();
6223    if d <= 1 {
6224        return None;
6225    }
6226    let center_eta = initial_aniso_contrasts(centers);
6227    let data_eta = standardized_spatial_term_data(pilot_data, term)
6228        .ok()
6229        .and_then(|x| finite_centered_axis_contrasts(&initial_aniso_contrasts(x.view()), d));
6230    let center_eta = finite_centered_axis_contrasts(&center_eta, d)?;
6231    let blended = match data_eta {
6232        Some(data_eta) => center_eta
6233            .iter()
6234            .zip(data_eta.iter())
6235            .map(|(&from_centers, &from_data)| 0.5 * (from_centers + from_data))
6236            .collect::<Vec<_>>(),
6237        None => center_eta,
6238    };
6239    finite_centered_axis_contrasts(&blended, d)
6240}
6241
6242fn apply_pilot_spatial_psi_reseed(
6243    pilot_data: ArrayView2<'_, f64>,
6244    spec: &TermCollectionSpec,
6245    spatial_terms: &[usize],
6246    kappa_options: &SpatialLengthScaleOptimizationOptions,
6247) -> Result<TermCollectionSpec, EstimationError> {
6248    let dims_per_term = spatial_dims_per_term(spec, spatial_terms);
6249    let use_aniso = has_aniso_terms(spec, spatial_terms);
6250    let log_kappa0 = if use_aniso {
6251        SpatialLogKappaCoords::from_length_scales_aniso(spec, spatial_terms, kappa_options)
6252    } else {
6253        SpatialLogKappaCoords::from_length_scales(spec, spatial_terms, kappa_options)
6254    };
6255    let log_kappa0 = log_kappa0.reseed_from_data(pilot_data, spec, spatial_terms, kappa_options);
6256    let log_kappa_lower = if use_aniso {
6257        SpatialLogKappaCoords::lower_bounds_aniso_from_data(
6258            pilot_data,
6259            spec,
6260            spatial_terms,
6261            &dims_per_term,
6262            kappa_options,
6263        )
6264    } else {
6265        SpatialLogKappaCoords::lower_bounds_from_data(
6266            pilot_data,
6267            spec,
6268            spatial_terms,
6269            kappa_options,
6270        )
6271    };
6272    let log_kappa_upper = if use_aniso {
6273        SpatialLogKappaCoords::upper_bounds_aniso_from_data(
6274            pilot_data,
6275            spec,
6276            spatial_terms,
6277            &dims_per_term,
6278            kappa_options,
6279        )
6280    } else {
6281        SpatialLogKappaCoords::upper_bounds_from_data(
6282            pilot_data,
6283            spec,
6284            spatial_terms,
6285            kappa_options,
6286        )
6287    };
6288    log_kappa0
6289        .clamp_to_bounds(&log_kappa_lower, &log_kappa_upper)
6290        .apply_tospec(spec, spatial_terms)
6291}
6292
6293pub(crate) fn apply_spatial_anisotropy_pilot_initializer(
6294    data: ArrayView2<'_, f64>,
6295    spec: &mut TermCollectionSpec,
6296    spatial_terms: &[usize],
6297    target_size: usize,
6298    kappa_options: &SpatialLengthScaleOptimizationOptions,
6299) -> usize {
6300    if target_size == 0 || data.nrows() <= target_size.saturating_mul(2) || spatial_terms.is_empty()
6301    {
6302        return 0;
6303    }
6304    if !has_aniso_terms(spec, spatial_terms) {
6305        return 0;
6306    }
6307    let indices = stratified_spatial_subsample(data, spec, target_size);
6308    let pilot_data = sampled_rows(data, &indices);
6309    let mut working = spec.clone();
6310    let mut updated_terms = 0usize;
6311    const GEOMETRY_UPDATES: usize = 2;
6312
6313    for pass in 0..GEOMETRY_UPDATES {
6314        let planned_terms = match plan_joint_spatial_centers_for_term_blocks(
6315            pilot_data.view(),
6316            &[working.smooth_terms.clone()],
6317        )
6318        .and_then(|mut blocks| {
6319            blocks.pop().ok_or_else(|| {
6320                BasisError::InvalidInput(
6321                    "pilot geometry initializer produced no smooth-term block".to_string(),
6322                )
6323            })
6324        }) {
6325            Ok(terms) => terms,
6326            Err(err) => {
6327                log::warn!(
6328                    "[spatial-kappa] pilot geometry initializer skipped after center planning failed: {err}"
6329                );
6330                return updated_terms;
6331            }
6332        };
6333
6334        for &term_idx in spatial_terms {
6335            let Some(current_eta) = get_spatial_aniso_log_scales(&working, term_idx) else {
6336                continue;
6337            };
6338            let Some(d) = get_spatial_feature_dim(&working, term_idx) else {
6339                continue;
6340            };
6341            if d <= 1 || current_eta.len() != d {
6342                continue;
6343            }
6344            let Some(planned_term) = planned_terms.get(term_idx) else {
6345                continue;
6346            };
6347            let Some(centers) = spatial_term_user_centers(planned_term) else {
6348                continue;
6349            };
6350            let Some(eta) = blended_pilot_axis_contrasts(pilot_data.view(), planned_term, centers)
6351            else {
6352                continue;
6353            };
6354            if set_spatial_aniso_log_scales(&mut working, term_idx, eta).is_ok() {
6355                updated_terms += usize::from(pass == 0);
6356            }
6357        }
6358
6359        match apply_pilot_spatial_psi_reseed(
6360            pilot_data.view(),
6361            &working,
6362            spatial_terms,
6363            kappa_options,
6364        ) {
6365            Ok(updated) => {
6366                working = updated;
6367            }
6368            Err(err) => {
6369                log::warn!(
6370                    "[spatial-kappa] pilot geometry ψ reseed skipped after deterministic initializer error: {err}"
6371                );
6372                break;
6373            }
6374        }
6375    }
6376
6377    if updated_terms > 0 {
6378        log::info!(
6379            "[spatial-kappa] initialized anisotropy from {}-row pilot geometry for {} spatial term(s); proceeding to full-data optimization",
6380            indices.len(),
6381            updated_terms
6382        );
6383        *spec = working;
6384    }
6385    updated_terms
6386}
6387
6388pub(crate) fn spatial_length_scale_term_indices(spec: &TermCollectionSpec) -> Vec<usize> {
6389    spec.smooth_terms
6390        .iter()
6391        .enumerate()
6392        .filter_map(|(idx, _)| spatial_term_supports_hyper_optimization(spec, idx).then_some(idx))
6393        .collect()
6394}
6395
6396/// Returns `true` when every spatial term in `spec` has a locked kernel
6397/// scale (explicit `length_scale=X` without anisotropy) and therefore
6398/// contributes no outer ψ/κ optimization axis. Empty term collections
6399/// also return `true` — there are no kappas to optimize.
6400///
6401/// Used by family entry points that want to honor a user-supplied scalar
6402/// length scale exactly: when all spatial terms are locked the n-block
6403/// joint-spatial outer solver has nothing to optimize, and routing
6404/// through it merely spends ~80 outer iters chasing a stalled ARC at the
6405/// user's chosen ρ. Skipping straight to the rho-only path avoids that
6406/// waste and respects the user's explicit kernel-scale input.
6407fn fit_score(fit: &UnifiedFitResult) -> f64 {
6408    if fit.reml_score.is_finite() {
6409        return fit.reml_score;
6410    }
6411    let score = 0.5 * fit.deviance + 0.5 * fit.stable_penalty_term;
6412    if score.is_finite() {
6413        score
6414    } else {
6415        f64::INFINITY
6416    }
6417}
6418
6419/// Classify an outer-evaluation error as a *recoverable trial-point
6420/// infeasibility* versus a genuine fatal failure.
6421///
6422/// The spatial-κ / anisotropy outer optimizer probes a sequence of trial
6423/// hyperparameters. At an extreme trial point the realized kernel design or
6424/// its ψ-derivatives may simply be non-constructible — e.g. a learned
6425/// per-axis log-scale stretches the anisotropic distance `r = |Λh|` until the
6426/// Duchon polyharmonic blocks `r^(2m−d)` overflow, or a degenerate metric
6427/// collapses two centers onto a non-C² collision. Those points lie outside
6428/// the model's feasible domain; the principled response is to treat them like
6429/// the cost-only path already does (objective `+∞`) so the line-search /
6430/// trust-region solver retreats, rather than aborting the entire REML fit.
6431///
6432/// A `BasisError` is exactly this class: it means "the basis/design cannot be
6433/// built at this hyperparameter". The same retreat semantics also apply when a
6434/// trial reaches the inner solve but produces a singular/unstable curvature:
6435/// those cases are reported by the shared inner-solve retreat classifier, or
6436/// by the final fit validator when an inference-only matrix derived from
6437/// `H⁻¹` (not the fitted mean coefficients themselves) becomes non-finite.
6438/// Everything else (layout/topology invariants, over-parameterization, and
6439/// arbitrary invalid inputs) stays fatal so genuine bugs are never masked.
6440fn is_recoverable_trial_point_error(err: &EstimationError) -> bool {
6441    matches!(err, EstimationError::BasisError(_))
6442        || err.is_inner_solve_retreat()
6443        || is_recoverable_fit_inference_finiteness_error(err)
6444}
6445
6446fn is_recoverable_fit_inference_finiteness_error(err: &EstimationError) -> bool {
6447    let EstimationError::InvalidInput(message) = err else {
6448        return false;
6449    };
6450
6451    message.contains("must be finite")
6452        && [
6453            "fit_result.beta_covariance_frequentist",
6454            "fit_result.coefficient_influence",
6455            "fit_result.weighted_gram",
6456        ]
6457        .iter()
6458        .any(|field| message.contains(field))
6459}
6460
6461#[cfg(test)]
6462mod spatial_trial_recovery_tests {
6463    use super::*;
6464
6465    #[test]
6466    fn nonfinite_frequentist_covariance_is_recoverable_trial_point() {
6467        let err = EstimationError::InvalidInput(
6468            "fit_result.beta_covariance_frequentist[0] must be finite, got NaN".to_string(),
6469        );
6470
6471        assert!(
6472            is_recoverable_trial_point_error(&err),
6473            "singular trial-point curvature should make spatial κ retreat, not abort"
6474        );
6475    }
6476
6477    #[test]
6478    fn arbitrary_invalid_input_remains_fatal_trial_point_error() {
6479        let err = EstimationError::InvalidInput("outer rho bounds are invalid".to_string());
6480
6481        assert!(
6482            !is_recoverable_trial_point_error(&err),
6483            "the spatial κ recovery gate must not mask unrelated invalid inputs"
6484        );
6485    }
6486}
6487
6488fn require_successful_spatial_optimization_result<T>(
6489    initial_score: f64,
6490    result: Result<Option<(T, f64)>, EstimationError>,
6491) -> Result<T, EstimationError> {
6492    match result {
6493        Ok(Some((value, exact_score))) => {
6494            // Allow rounding-level worsening: REML scores accumulate
6495            // log-determinant terms whose finite-precision re-evaluation
6496            // can drift well past 1e-10 absolute near a converged optimum
6497            // (we have seen ~1e-6 between two evaluations whose printed
6498            // values round to identical 6-digit scientific). Reject genuine
6499            // worsenings (>1 unit) but admit anything within ~1e-6
6500            // absolute / 1e-8 relative — meaningful REML gains are
6501            // orders of magnitude larger.
6502            const SCORE_DRIFT_ABS_TOL: f64 = 1e-6;
6503            const SCORE_DRIFT_REL_TOL: f64 = 1e-8;
6504            let tol = SCORE_DRIFT_ABS_TOL.max(initial_score.abs() * SCORE_DRIFT_REL_TOL);
6505            if exact_score <= initial_score + tol {
6506                Ok(value)
6507            } else {
6508                Err(EstimationError::RemlOptimizationFailed(format!(
6509                    "spatial kappa optimization made REML score worse ({initial_score:.6e} -> {exact_score:.6e})"
6510                )))
6511            }
6512        }
6513        Ok(None) => Err(EstimationError::RemlOptimizationFailed(
6514            "spatial kappa optimization is unavailable for one or more eligible spatial terms"
6515                .to_string(),
6516        )),
6517        Err(err) => Err(EstimationError::RemlOptimizationFailed(format!(
6518            "spatial kappa optimization failed: {err}"
6519        ))),
6520    }
6521}
6522
6523fn external_opts_for_design(
6524    family: &LikelihoodSpec,
6525    design: &TermCollectionDesign,
6526    options: &FitOptions,
6527) -> ExternalOptimOptions {
6528    ExternalOptimOptions {
6529        family: family.clone(),
6530        latent_cloglog: options.latent_cloglog,
6531        mixture_link: options.mixture_link.clone(),
6532        optimize_mixture: options.optimize_mixture,
6533        sas_link: options.sas_link,
6534        optimize_sas: options.optimize_sas,
6535        compute_inference: options.compute_inference,
6536        skip_rho_posterior_inference: options.skip_rho_posterior_inference,
6537        max_iter: options.max_iter,
6538        tol: options.tol,
6539        nullspace_dims: design.nullspace_dims.clone(),
6540        linear_constraints: design.linear_constraints.clone(),
6541        firth_bias_reduction: Some(options.firth_bias_reduction),
6542        penalty_shrinkage_floor: options.penalty_shrinkage_floor,
6543        rho_prior: options.rho_prior.clone(),
6544        // Propagate Kronecker structure so the joint optimizer minimizes the
6545        // same REML surface as the baseline/refit (adaptive_fit_options_base).
6546        kronecker_penalty_system: design.kronecker_penalty_system(),
6547        kronecker_factored: design
6548            .smooth
6549            .terms
6550            .iter()
6551            .find_map(|t| t.kronecker_factored.clone()),
6552        persist_warm_start_disk: options.persist_warm_start_disk,
6553    }
6554}
6555
6556/// Evaluate the joint REML cost, gradient, and Hessian result at a given θ = [ρ, ψ]
6557/// for a single-block term collection with spatial hyperparameters.
6558///
6559/// This provides a direct evaluation of the profiled REML objective using the
6560/// external-caller interface, which exposes exact cost/gradient/Hessian without
6561/// running the full outer smoothing loop. The returned tuple is
6562/// `(cost, gradient, hessian)` in the joint [ρ, ψ] space.
6563fn evaluate_joint_reml_outer_eval_at_theta(
6564    evaluator: &mut gam_solve::estimate::ExternalJointHyperEvaluator<'_>,
6565    design: &TermCollectionDesign,
6566    theta: &Array1<f64>,
6567    rho_dim: usize,
6568    hyper_dirs: Vec<gam_solve::estimate::reml::DirectionalHyperParam>,
6569    warm_start_beta: Option<ArrayView1<'_, f64>>,
6570    order: gam_solve::rho_optimizer::OuterEvalOrder,
6571    design_revision: Option<u64>,
6572) -> Result<
6573    (
6574        f64,
6575        Array1<f64>,
6576        gam_problem::HessianValue,
6577    ),
6578    EstimationError,
6579> {
6580    evaluator.evaluate_with_order(
6581        &design.design,
6582        &design.penalties,
6583        &design.nullspace_dims,
6584        design.linear_constraints.clone(),
6585        theta,
6586        rho_dim,
6587        hyper_dirs,
6588        warm_start_beta,
6589        "evaluate_joint_reml_outer_eval_at_theta",
6590        order,
6591        design_revision,
6592    )
6593}
6594
6595fn evaluate_joint_reml_efs_at_theta(
6596    evaluator: &mut gam_solve::estimate::ExternalJointHyperEvaluator<'_>,
6597    design: &TermCollectionDesign,
6598    theta: &Array1<f64>,
6599    rho_dim: usize,
6600    hyper_dirs: Vec<gam_solve::estimate::reml::DirectionalHyperParam>,
6601    warm_start_beta: Option<ArrayView1<'_, f64>>,
6602    design_revision: Option<u64>,
6603) -> Result<gam_problem::EfsEval, EstimationError> {
6604    evaluator.evaluate_efs(
6605        &design.design,
6606        &design.penalties,
6607        &design.nullspace_dims,
6608        design.linear_constraints.clone(),
6609        theta,
6610        rho_dim,
6611        hyper_dirs,
6612        warm_start_beta,
6613        "evaluate_joint_reml_efs_at_theta",
6614        design_revision,
6615    )
6616}
6617
6618fn exact_joint_spatial_outer_hessian_available(
6619    family: &LikelihoodSpec,
6620    design: &TermCollectionDesign,
6621) -> bool {
6622    // Every `LikelihoodSpec` variant (Gaussian, Binomial-*, Poisson, Gamma,
6623    // Royston-Parmar) routes through the unified evaluator's outer-Hessian
6624    // path: Gaussian Identity uses the no-correction dense form, all GLM
6625    // variants supply scalar-GLM derivative ingredients consumed by
6626    // `compute_outer_hessian` / `build_outer_hessian_operator`, and the
6627    // (n, p, K) crossover in `prefer_outer_hessian_operator` chooses the
6628    // matrix-free `HessianValue::Operator` representation at large scale
6629    // for dense-lazy designs.  The previous `Identity || sparse_design`
6630    // gate predates that operator routing and forced binomial+logit+Matern
6631    // (and any other non-Gaussian dense-lazy spatial design) onto the
6632    // gradient-only BFGS path even though analytic Hessian is fully
6633    // available — capability check, not cost.  Match every variant
6634    // explicitly so any future family addition (which may not yet provide
6635    // outer-Hessian ingredients) forces an authoring decision here rather
6636    // than silently inheriting `true`.
6637    // Every supported response (Gaussian, Binomial-*, Poisson, Tweedie,
6638    // NegativeBinomial, Beta, Gamma, Royston-Parmar) routes through the
6639    // unified evaluator's outer-Hessian path; the spec-level capability
6640    // check therefore always succeeds. Match every response explicitly so
6641    // any future family addition (which may not yet provide outer-Hessian
6642    // ingredients) forces an authoring decision here rather than silently
6643    // inheriting `true`.
6644    let family_supported = match &family.response {
6645        ResponseFamily::Gaussian
6646        | ResponseFamily::Binomial
6647        | ResponseFamily::Poisson
6648        | ResponseFamily::Tweedie { .. }
6649        | ResponseFamily::NegativeBinomial { .. }
6650        | ResponseFamily::Beta { .. }
6651        | ResponseFamily::Gamma
6652        | ResponseFamily::RoystonParmar => true,
6653    };
6654    // A design with zero columns has no joint outer-Hessian to compute;
6655    // the analytic path is only meaningful for non-empty parameter blocks.
6656    family_supported && design.design.ncols() > 0
6657}
6658
6659fn smooth_term_penalty_index(
6660    spec: &TermCollectionSpec,
6661    design: &TermCollectionDesign,
6662    term_idx: usize,
6663) -> Option<usize> {
6664    if term_idx >= design.smooth.terms.len() || term_idx >= spec.smooth_terms.len() {
6665        return None;
6666    }
6667    if design.smooth.terms[term_idx].penalties_local.is_empty() {
6668        return None;
6669    }
6670    let linear_penalties = spec
6671        .linear_terms
6672        .iter()
6673        .filter(|t| t.double_penalty)
6674        .count()
6675        * 2;
6676    let random_penalties = design
6677        .random_effect_ranges
6678        .iter()
6679        .filter(|(_, range)| !range.is_empty())
6680        .count();
6681    let smooth_offset = linear_penalties + random_penalties;
6682    let local_offset = design
6683        .smooth
6684        .terms
6685        .iter()
6686        .take(term_idx)
6687        .map(|term| term.penalties_local.len())
6688        .sum::<usize>();
6689    Some(smooth_offset + local_offset)
6690}
6691
6692fn try_build_spatial_term_log_kappa_derivativeinfo(
6693    data: ArrayView2<'_, f64>,
6694    resolvedspec: &TermCollectionSpec,
6695    design: &TermCollectionDesign,
6696    term_idx: usize,
6697) -> Result<Option<SpatialPsiDerivative>, EstimationError> {
6698    let Some((
6699        global_range,
6700        total_p,
6701        x_psi_local,
6702        s_psi_local_check,
6703        x_psi_psi_local,
6704        s_psi_psi_local,
6705        s_psi_components_local,
6706        s_psi_psi_components_local,
6707        implicit_operator,
6708    )) = try_build_spatial_term_log_kappa_derivative(data, resolvedspec, design, term_idx)?
6709    else {
6710        return Ok(None);
6711    };
6712    let Some(penalty_start) = smooth_term_penalty_index(resolvedspec, design, term_idx) else {
6713        return Ok(None);
6714    };
6715    if s_psi_components_local.is_empty() || s_psi_psi_components_local.is_empty() {
6716        return Ok(None);
6717    }
6718    if s_psi_components_local.len() != s_psi_psi_components_local.len() {
6719        return Ok(None);
6720    }
6721    let penalty_indices = (0..s_psi_components_local.len())
6722        .map(|j| penalty_start + j)
6723        .collect::<Vec<_>>();
6724    let penalty_index = penalty_indices[0];
6725    if s_psi_local_check.nrows() == 0 || s_psi_psi_local.nrows() == 0 {
6726        return Ok(None);
6727    }
6728    Ok(Some(SpatialPsiDerivative {
6729        penalty_index,
6730        penalty_indices,
6731        global_range,
6732        total_p,
6733        x_psi_local,
6734        s_psi_components_local,
6735        x_psi_psi_local,
6736        s_psi_psi_components_local,
6737        aniso_group_id: None,
6738        aniso_cross_designs: None,
6739        aniso_cross_penalty_provider: None,
6740        implicit_operator,
6741        implicit_axis: 0,
6742    }))
6743}
6744
6745pub(crate) fn try_build_spatial_log_kappa_derivativeinfo_list(
6746    data: ArrayView2<'_, f64>,
6747    resolvedspec: &TermCollectionSpec,
6748    design: &TermCollectionDesign,
6749    spatial_terms: &[usize],
6750) -> Result<Option<Vec<SpatialPsiDerivative>>, EstimationError> {
6751    let mut out = Vec::new();
6752    let mut aniso_gid = 0usize;
6753    for &term_idx in spatial_terms {
6754        if spatial_term_uses_per_axis_psi(resolvedspec, term_idx) {
6755            if let Some(entries) = try_build_spatial_term_log_kappa_aniso_derivativeinfos(
6756                data,
6757                resolvedspec,
6758                design,
6759                term_idx,
6760                aniso_gid,
6761            )? {
6762                aniso_gid += 1;
6763                out.extend(entries);
6764                continue;
6765            } else {
6766                return Ok(None);
6767            }
6768        }
6769        let Some(info) =
6770            try_build_spatial_term_log_kappa_derivativeinfo(data, resolvedspec, design, term_idx)?
6771        else {
6772            return Ok(None);
6773        };
6774        out.push(info);
6775    }
6776    Ok(Some(out))
6777}
6778
6779/// For an aniso term with d axes, produce d `SpatialPsiDerivative` entries.
6780fn try_build_spatial_term_log_kappa_aniso_derivativeinfos(
6781    data: ArrayView2<'_, f64>,
6782    resolvedspec: &TermCollectionSpec,
6783    design: &TermCollectionDesign,
6784    term_idx: usize,
6785    aniso_group_id: usize,
6786) -> Result<Option<Vec<SpatialPsiDerivative>>, EstimationError> {
6787    let Some(smooth_term) = design.smooth.terms.get(term_idx) else {
6788        return Ok(None);
6789    };
6790    let Some(termspec) = resolvedspec.smooth_terms.get(term_idx) else {
6791        return Ok(None);
6792    };
6793    let mut aniso_result = match &termspec.basis {
6794        SmoothBasisSpec::Sphere { .. } => return Ok(None),
6795        SmoothBasisSpec::Matern {
6796            feature_cols,
6797            spec,
6798            input_scales,
6799        } => {
6800            let mut x = select_columns(data, feature_cols).map_err(EstimationError::from)?;
6801            if let Some(s) = input_scales {
6802                apply_input_standardization(&mut x, s);
6803            }
6804            // #1122: the realized Matérn design always carries the operator
6805            // {mass, tension, stiffness} penalty triplet (`build_term` overrides
6806            // the `double_penalty` kernel penalty via
6807            // `matern_operator_penalty_triplet_from_metadata`). The per-axis
6808            // κ-gradient must differentiate that SAME triplet, not the kernel
6809            // double-penalty blocks, or the analytic `tr(S⁺ Ṡ)` desyncs from the
6810            // FD of the criterion's operator-triplet `log|Sλ|₊` (the iso-axis
6811            // analogue is handled in `try_build_spatial_term_log_kappa_derivative`).
6812            let mut spec_operator = spec.clone();
6813            spec_operator.double_penalty = false;
6814            build_matern_basis_log_kappa_aniso_derivatives(x.view(), &spec_operator)
6815                .map_err(EstimationError::from)?
6816        }
6817        // Measure-jet: the grouped dial coordinates ride the same per-axis
6818        // carrier. The producer runs on the FROZEN spec (the driver runs
6819        // post-freeze), so per-trial rebuilds move only the dials; the
6820        // coordinate layout, zero design drift, and shared candidate
6821        // normalization are owned by `build_measure_jet_basis_psi_derivatives`.
6822        SmoothBasisSpec::MeasureJet {
6823            feature_cols,
6824            spec,
6825            input_scales,
6826        } => {
6827            let mut x = select_columns(data, feature_cols).map_err(EstimationError::from)?;
6828            if let Some(s) = input_scales {
6829                apply_input_standardization(&mut x, s);
6830            }
6831            build_measure_jet_basis_psi_derivatives(x.view(), spec)
6832                .map_err(EstimationError::from)?
6833        }
6834        _ => return Ok(None),
6835    };
6836    // Get number of axes from the shared operator when available; otherwise
6837    // fall back to the dense design list.
6838    let d = if let Some(ref op) = aniso_result.implicit_operator {
6839        op.n_axes()
6840    } else if !aniso_result.design_first.is_empty() {
6841        aniso_result.design_first.len()
6842    } else {
6843        0
6844    };
6845    if d == 0 {
6846        return Ok(None);
6847    }
6848    let Some(penalty_start) = smooth_term_penalty_index(resolvedspec, design, term_idx) else {
6849        return Ok(None);
6850    };
6851    let p_total = design.design.ncols();
6852    let smooth_start = p_total.saturating_sub(design.smooth.total_smooth_cols());
6853    let global_range = (smooth_start + smooth_term.coeff_range.start)
6854        ..(smooth_start + smooth_term.coeff_range.end);
6855    let num_penalties = aniso_result.penalties_first[0].len();
6856    let penalty_indices: Vec<usize> = (0..num_penalties).map(|j| penalty_start + j).collect();
6857    let penalties_cross_provider = aniso_result.penalties_cross_provider.clone();
6858
6859    // Dense first/diagonal-second matrices may be present even when the shared
6860    // operator is available. The operator remains the canonical source for
6861    // exact cross-axis second derivatives.
6862    let use_implicit_design = aniso_result.design_first.is_empty();
6863    let implicit_op_arc = aniso_result
6864        .implicit_operator
6865        .as_ref()
6866        .map(|op| std::sync::Arc::new(op.clone()));
6867
6868    let mut entries = Vec::with_capacity(d);
6869    for a in 0..d {
6870        let (x_psi_local, x_psi_psi_local) = if use_implicit_design {
6871            // Implicit path: design-derivative matvecs will be dispatched through
6872            // the ImplicitDerivativeOp inside HyperDesignDerivative, so we do NOT
6873            // need to materialize the dense (n x p) matrices here.  Store empty
6874            // placeholders — they are never read when the implicit operator is
6875            // present (spatial_log_kappa_hyper_dirs_frominfo_list uses from_implicit).
6876            (Array2::<f64>::zeros((0, 0)), Array2::<f64>::zeros((0, 0)))
6877        } else {
6878            // Move the dense (n × p) matrices out of aniso_result instead of
6879            // cloning. Each axis index `a` is read exactly once across the
6880            // loop, and aniso_result is dropped at function exit, so leaving
6881            // empty placeholders behind in those vec slots is safe.
6882            let x_first = std::mem::take(&mut aniso_result.design_first[a]);
6883            let x_second = std::mem::take(&mut aniso_result.design_second_diag[a]);
6884            if x_first.ncols() != smooth_term.coeff_range.len() {
6885                return Ok(None);
6886            }
6887            (x_first, x_second)
6888        };
6889        let s_psi_components = std::mem::take(&mut aniso_result.penalties_first[a]);
6890        let s_psi_psi_components = std::mem::take(&mut aniso_result.penalties_second_diag[a]);
6891        // Build cross-design entries for other axes b != a in this group.
6892        // These will be indexed by (b, cross_matrix) where b is the axis
6893        // offset within the d-entry block.
6894        // Cross-axis second derivatives are sourced from the shared operator,
6895        // so we only need placeholder entries to preserve the axis layout.
6896        let cross_designs = if implicit_op_arc.is_some() {
6897            let mut cd = Vec::with_capacity(d - 1);
6898            for b in 0..d {
6899                if b == a {
6900                    continue;
6901                }
6902                cd.push((b, Array2::<f64>::zeros((0, 0))));
6903            }
6904            cd
6905        } else if !aniso_result.design_second_cross.is_empty() {
6906            let mut cd = Vec::new();
6907            for (cross_idx, &(pa, pb)) in aniso_result.design_second_cross_pairs.iter().enumerate()
6908            {
6909                if pa == a {
6910                    cd.push((pb, aniso_result.design_second_cross[cross_idx].clone()));
6911                } else if pb == a {
6912                    cd.push((pa, aniso_result.design_second_cross[cross_idx].clone()));
6913                }
6914            }
6915            cd
6916        } else {
6917            Vec::new()
6918        };
6919        let cross_penalty_provider = if d > 1 {
6920            let penalties_cross_provider = penalties_cross_provider.clone();
6921            Some(std::sync::Arc::new(
6922                move |b_axis: usize| -> Result<Vec<Array2<f64>>, EstimationError> {
6923                    if b_axis == a {
6924                        return Ok(Vec::new());
6925                    }
6926                    let (axis_lo, axis_hi) = if a < b_axis { (a, b_axis) } else { (b_axis, a) };
6927                    if let Some(provider) = penalties_cross_provider.as_ref() {
6928                        provider
6929                            .evaluate(axis_lo, axis_hi)
6930                            .map_err(EstimationError::from)
6931                    } else {
6932                        // No provider: either the pair is unregistered, or it
6933                        // was registered without data (early-return raw-operator
6934                        // paths). Both cases contribute no cross penalties.
6935                        Ok(Vec::new())
6936                    }
6937                },
6938            )
6939                as std::sync::Arc<
6940                    dyn Fn(usize) -> Result<Vec<Array2<f64>>, EstimationError>
6941                        + Send
6942                        + Sync
6943                        + 'static,
6944                >)
6945        } else {
6946            None
6947        };
6948
6949        entries.push(SpatialPsiDerivative {
6950            penalty_index: penalty_indices[0],
6951            penalty_indices: penalty_indices.clone(),
6952            global_range: global_range.clone(),
6953            total_p: p_total,
6954            x_psi_local,
6955            s_psi_components_local: s_psi_components,
6956            x_psi_psi_local,
6957            s_psi_psi_components_local: s_psi_psi_components,
6958            aniso_group_id: Some(aniso_group_id),
6959            aniso_cross_designs: if cross_designs.is_empty() {
6960                None
6961            } else {
6962                Some(cross_designs)
6963            },
6964            aniso_cross_penalty_provider: cross_penalty_provider,
6965            implicit_operator: implicit_op_arc.clone(),
6966            implicit_axis: a,
6967        });
6968    }
6969    Ok(Some(entries))
6970}
6971
6972#[cfg(test)]
6973mod glm_eta_observation_fd_tests {
6974    //! #1615/#1616: the non-Gaussian GLM arms of `evaluate_standard_familyobservations`
6975    //! (Poisson / Gamma / NegativeBinomial / Tweedie) must have a self-consistent
6976    //! derivative tower: `score = ∂ℓ/∂η`, `neghessian_eta = −∂(score)/∂η`, and
6977    //! `neghessian_eta_derivative = ∂(neghessian_eta)/∂η`. Pin each against central
6978    //! finite differences of the assembled log-likelihood / score.
6979    use super::*;
6980
6981    fn one_obs(spec: &LikelihoodSpec, y: f64, eta: f64) -> StandardFamilyObservationState {
6982        let yv = Array1::from_vec(vec![y]);
6983        let wv = Array1::from_vec(vec![1.0]);
6984        let ev = Array1::from_vec(vec![eta]);
6985        evaluate_standard_familyobservations(spec.clone(), None, None, None, &yv, &wv, &ev)
6986            .expect("standard family observation state assembles")
6987    }
6988
6989    fn check_fd(label: &str, spec: &LikelihoodSpec, y: f64, eta: f64) {
6990        let h = 1e-5;
6991        let s0 = one_obs(spec, y, eta);
6992        let sp = one_obs(spec, y, eta + h);
6993        let sm = one_obs(spec, y, eta - h);
6994
6995        // score = d(log_likelihood)/d(eta)
6996        let score_fd = (sp.log_likelihood - sm.log_likelihood) / (2.0 * h);
6997        let score = s0.score[0];
6998        assert!(
6999            (score - score_fd).abs() <= 1e-4 * (1.0 + score.abs()),
7000            "{label}: score {score} vs FD {score_fd}"
7001        );
7002
7003        // neghessian_eta = -d(score)/d(eta)
7004        let neghess_fd = -(sp.score[0] - sm.score[0]) / (2.0 * h);
7005        let neghess = s0.neghessian_eta[0];
7006        assert!(
7007            (neghess - neghess_fd).abs() <= 1e-3 * (1.0 + neghess.abs()),
7008            "{label}: neghessian_eta {neghess} vs FD {neghess_fd}"
7009        );
7010
7011        // neghessian_eta_derivative = d(neghessian_eta)/d(eta)
7012        let nhd_fd = (sp.neghessian_eta[0] - sm.neghessian_eta[0]) / (2.0 * h);
7013        let nhd = s0.neghessian_eta_derivative[0];
7014        assert!(
7015            (nhd - nhd_fd).abs() <= 1e-2 * (1.0 + nhd.abs()),
7016            "{label}: neghessian_eta_derivative {nhd} vs FD {nhd_fd}"
7017        );
7018    }
7019
7020    #[test]
7021    fn poisson_gamma_nb_tweedie_arms_match_finite_differences_1615_1616() {
7022        let log = InverseLink::Standard(StandardLink::Log);
7023        let poisson = LikelihoodSpec {
7024            response: ResponseFamily::Poisson,
7025            link: log.clone(),
7026        };
7027        check_fd("poisson y=3", &poisson, 3.0, 0.4);
7028        check_fd("poisson y=0", &poisson, 0.0, -0.2);
7029
7030        let gamma = LikelihoodSpec {
7031            response: ResponseFamily::Gamma,
7032            link: log.clone(),
7033        };
7034        check_fd("gamma y=2.5", &gamma, 2.5, 0.3);
7035        check_fd("gamma y=0.7", &gamma, 0.7, -0.1);
7036
7037        let nb = LikelihoodSpec {
7038            response: ResponseFamily::NegativeBinomial {
7039                theta: 1.5,
7040                theta_fixed: true,
7041            },
7042            link: log.clone(),
7043        };
7044        check_fd("negbin y=4", &nb, 4.0, 0.5);
7045        check_fd("negbin y=0", &nb, 0.0, -0.3);
7046
7047        let tweedie = LikelihoodSpec {
7048            response: ResponseFamily::Tweedie { p: 1.5 },
7049            link: log.clone(),
7050        };
7051        check_fd("tweedie y=2", &tweedie, 2.0, 0.25);
7052        check_fd("tweedie y=0.5", &tweedie, 0.5, -0.15);
7053    }
7054}