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