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