Skip to main content

gam_models/fit_orchestration/drivers/
spatial_optimization.rs

1fn try_build_spatial_term_log_kappa_derivative(
2    data: ArrayView2<'_, f64>,
3    resolvedspec: &TermCollectionSpec,
4    design: &TermCollectionDesign,
5    term_idx: usize,
6) -> Result<
7    Option<(
8        Range<usize>,
9        usize,
10        Array2<f64>,
11        Array2<f64>,
12        Array2<f64>,
13        Array2<f64>,
14        Vec<Array2<f64>>,
15        Vec<Array2<f64>>,
16        Option<std::sync::Arc<gam_terms::basis::ImplicitDesignPsiDerivative>>,
17    )>,
18    EstimationError,
19> {
20    let Some(smooth_term) = design.smooth.terms.get(term_idx) else {
21        return Ok(None);
22    };
23    let Some(termspec) = resolvedspec.smooth_terms.get(term_idx) else {
24        return Ok(None);
25    };
26
27    let derivative_bundle = match &termspec.basis {
28        SmoothBasisSpec::ThinPlate {
29            feature_cols,
30            spec,
31            input_scale,
32        } => {
33            let mut x = select_columns(data, feature_cols).map_err(EstimationError::from)?;
34            let mut spec_local = spec.clone();
35            if let Some(scale) = input_scale {
36                scale.standardize(&mut x);
37                spec_local.length_scale = scale.to_standardized_units(spec.length_scale);
38            }
39            build_thin_plate_basis_log_kappa_derivatives(x.view(), &spec_local)
40                .map_err(EstimationError::from)?
41        }
42        SmoothBasisSpec::Sphere { .. } => return Ok(None),
43        // Constant-curvature smooths expose κ as one signed, design-moving
44        // outer ψ-coordinate (#944 stage 3 final wiring). Unlike the Matérn /
45        // Duchon / TPS kernels — whose ψ-coordinate is `log κ = −log ℓ` — the
46        // constant-curvature ψ-coordinate is the **raw curvature κ itself**, so
47        // κ = 0 stays an interior point of the `S^d ← ℝ^d → H^d` family. The
48        // bundle therefore carries `∂·/∂κ` / `∂²·/∂κ²` directly, and the chart
49        // coordinates are consumed verbatim (no input standardization — the
50        // gauge `1 + κ‖x‖²` defines what κ means; see the basis builder).
51        SmoothBasisSpec::ConstantCurvature { feature_cols, spec } => {
52            let x = select_columns(data, feature_cols).map_err(EstimationError::from)?;
53            build_constant_curvature_basis_kappa_derivatives(x.view(), spec)
54                .map_err(EstimationError::from)?
55        }
56        // Measure-jet routes through the GROUPED dial builder
57        // (`try_build_spatial_term_log_kappa_aniso_derivativeinfos`):
58        // `spatial_term_uses_per_axis_psi` is true for every enrolled
59        // measure-jet term, so this isotropic path only sees unenrolled
60        // terms (`measure_jet_enrolls_psi` = false), which expose no ψ bundle.
61        SmoothBasisSpec::MeasureJet { .. } => return Ok(None),
62        SmoothBasisSpec::Matern {
63            feature_cols,
64            spec,
65            input_scale,
66        } => {
67            let mut x = select_columns(data, feature_cols).map_err(EstimationError::from)?;
68            let mut spec_local = spec.clone();
69            if let Some(scale) = input_scale {
70                scale.standardize(&mut x);
71                let length_scale = spec.length_scale.resolved().ok_or_else(|| {
72                    EstimationError::InvalidInput(
73                        "Matérn Auto length_scale reached derivative construction unresolved"
74                            .to_string(),
75                    )
76                })?;
77                spec_local
78                    .length_scale
79                    .set_resolved(scale.to_standardized_units(length_scale));
80            }
81            // The realized Matérn DESIGN penalty is ALWAYS the operator-collocation
82            // {mass, tension, stiffness} triplet — the term-collection assembler
83            // overrides whatever `double_penalty` produced at the basis level with
84            // `matern_operator_penalty_triplet_from_metadata` (see
85            // `gam_terms::smooth::term_specs`, "The Matérn design ALWAYS uses the
86            // operator-collocation … triplet"; #1074/#1270). The ψ=log κ outer
87            // gradient must differentiate the SAME penalty the REML cost is built
88            // on, so the derivative is forced onto the operator-triplet path here.
89            // Honoring `double_penalty: true` instead returned the kernel-Gram
90            // double-penalty ψ-derivatives — a penalty the design does NOT carry —
91            // which desynced the analytic iso-κ gradient from the cost's FD and
92            // stalled the κ-optimizer at its iteration cap with a large residual
93            // gradient (#1122). `double_penalty: false` reproduces the operator
94            // triplet exactly (verified: the 2-D iso-κ FD matches to ~1e-9).
95            spec_local.double_penalty = false;
96            build_matern_basis_log_kappa_derivatives(x.view(), &spec_local)
97                .map_err(EstimationError::from)?
98        }
99        SmoothBasisSpec::Duchon {
100            feature_cols,
101            spec,
102            input_scale,
103        } => {
104            let mut x = select_columns(data, feature_cols).map_err(EstimationError::from)?;
105            let mut spec_local = spec.clone();
106            if let Some(scale) = input_scale {
107                scale.standardize(&mut x);
108                spec_local.length_scale = spec
109                    .length_scale
110                    .map(|length| scale.to_standardized_units(length));
111            }
112            let BasisMetadata::Duchon {
113                centers,
114                identifiability_transform,
115                operator_collocation_points,
116                radial_reparam,
117                ..
118            } = &smooth_term.metadata
119            else {
120                return Ok(None);
121            };
122            // #1355: replay the frozen data-metric reparam into the derivative
123            // spec so the ψ-derivative arms assemble in the rotated radial basis.
124            if spec_local.radial_reparam.is_none() {
125                spec_local.radial_reparam = radial_reparam.clone();
126            }
127            gam_terms::basis::build_duchon_basis_log_kappa_derivativeswith_collocationwithworkspace(
128                x.view(),
129                &spec_local,
130                centers.view(),
131                identifiability_transform.as_ref(),
132                operator_collocation_points
133                    .as_ref()
134                    .map(|points| points.view()),
135                &mut BasisWorkspace::default(),
136            )
137            .map_err(EstimationError::from)?
138        }
139        SmoothBasisSpec::BSpline1D { .. }
140        | SmoothBasisSpec::TensorBSpline { .. }
141        | SmoothBasisSpec::ByVariable { .. }
142        | SmoothBasisSpec::FactorSumToZero { .. }
143        | SmoothBasisSpec::BySmooth { .. }
144        | SmoothBasisSpec::FactorSmooth { .. }
145        | SmoothBasisSpec::Pca { .. } => {
146            return Ok(None);
147        }
148    };
149    let mut implicit_operator = derivative_bundle.implicit_operator;
150    let BasisPsiDerivativeResult {
151        design_derivative: mut local_x_psi,
152        penalties_derivative: mut local_s_psi,
153        implicit_operator: local_implicit_first_unused,
154    } = derivative_bundle.first;
155    let BasisPsiSecondDerivativeResult {
156        designsecond_derivative: mut local_x_psi_psi,
157        penaltiessecond_derivative: mut local_s_psi_psi,
158        implicit_operator: local_implicit_second_unused,
159    } = derivative_bundle.second;
160    assert!(local_implicit_first_unused.is_none());
161    assert!(local_implicit_second_unused.is_none());
162
163    if let Some(rotation) = smooth_term.joint_null_rotation.as_ref() {
164        let q = &rotation.rotation;
165        if let Some(op) = implicit_operator.take() {
166            implicit_operator = Some(op.append_full_transform(q).map_err(EstimationError::from)?);
167        } else {
168            if local_x_psi.ncols() != q.nrows() || local_x_psi_psi.ncols() != q.nrows() {
169                return Ok(None);
170            }
171            local_x_psi = fast_ab(&local_x_psi, q);
172            local_x_psi_psi = fast_ab(&local_x_psi_psi, q);
173        }
174        let rotate_penalty = |s_local: Array2<f64>| -> Option<Array2<f64>> {
175            if s_local.nrows() != q.nrows() || s_local.ncols() != q.nrows() {
176                return None;
177            }
178            let qt_s = gam_linalg::faer_ndarray::fast_atb(q, &s_local);
179            Some(gam_linalg::faer_ndarray::fast_ab(&qt_s, q))
180        };
181        let Some(rotated_s_psi) = local_s_psi
182            .into_iter()
183            .map(|s| rotate_penalty(s))
184            .collect::<Option<Vec<_>>>()
185        else {
186            return Ok(None);
187        };
188        local_s_psi = rotated_s_psi;
189        let Some(rotated_s_psi_psi) = local_s_psi_psi
190            .into_iter()
191            .map(|s| rotate_penalty(s))
192            .collect::<Option<Vec<_>>>()
193        else {
194            return Ok(None);
195        };
196        local_s_psi_psi = rotated_s_psi_psi;
197    }
198    let implicit_operator = implicit_operator.map(std::sync::Arc::new);
199
200    if let Some(ref op) = implicit_operator {
201        if op.p_out() != smooth_term.coeff_range.len() {
202            return Ok(None);
203        }
204    } else {
205        if local_x_psi.ncols() != smooth_term.coeff_range.len() {
206            return Ok(None);
207        }
208        if local_x_psi_psi.ncols() != smooth_term.coeff_range.len() {
209            return Ok(None);
210        }
211    }
212    if local_s_psi.is_empty() || local_s_psi.len() != local_s_psi_psi.len() {
213        return Ok(None);
214    }
215    if local_s_psi.iter().any(|s| {
216        s.nrows() != smooth_term.coeff_range.len() || s.ncols() != smooth_term.coeff_range.len()
217    }) {
218        return Ok(None);
219    }
220    if local_s_psi_psi.iter().any(|s| {
221        s.nrows() != smooth_term.coeff_range.len() || s.ncols() != smooth_term.coeff_range.len()
222    }) {
223        return Ok(None);
224    }
225
226    let p_total = design.design.ncols();
227    let smooth_start = p_total.saturating_sub(design.smooth.total_smooth_cols());
228    let global_range = (smooth_start + smooth_term.coeff_range.start)
229        ..(smooth_start + smooth_term.coeff_range.end);
230
231    Ok(Some((
232        global_range,
233        p_total,
234        local_x_psi,
235        local_s_psi.iter().fold(
236            Array2::<f64>::zeros((smooth_term.coeff_range.len(), smooth_term.coeff_range.len())),
237            |acc, m| acc + m,
238        ),
239        local_x_psi_psi,
240        local_s_psi_psi.iter().fold(
241            Array2::<f64>::zeros((smooth_term.coeff_range.len(), smooth_term.coeff_range.len())),
242            |acc, m| acc + m,
243        ),
244        local_s_psi,
245        local_s_psi_psi,
246        implicit_operator,
247    )))
248}
249
250fn try_build_spatial_log_kappa_hyper_dirs(
251    data: ArrayView2<'_, f64>,
252    resolvedspec: &TermCollectionSpec,
253    design: &TermCollectionDesign,
254    spatial_terms: &[usize],
255) -> Result<Option<Vec<DirectionalHyperParam>>, EstimationError> {
256    // Each spatial term contributes one continuous scale hyperparameter
257    //   psi = log(kappa) = -log(length_scale),
258    // while rho = log(lambda) still indexes the smoothing parameters of the
259    // three operator penalties. The joint outer vector is therefore
260    //   theta = (rho_0, ..., rho_{K-1}, psi_1, ..., psi_q)
261    // for q spatial terms participating in exact joint optimization.
262    let Some(info_list) =
263        try_build_spatial_log_kappa_derivativeinfo_list(data, resolvedspec, design, spatial_terms)?
264    else {
265        return Ok(None);
266    };
267    Ok(Some(spatial_log_kappa_hyper_dirs_frominfo_list(info_list)?))
268}
269
270pub(crate) fn try_build_latent_coord_hyper_dirs(
271    latent: std::sync::Arc<gam_terms::latent::LatentCoordValues>,
272    resolvedspec: &TermCollectionSpec,
273    design: &TermCollectionDesign,
274    latent_terms: &[gam_problem::types::SmoothTermIdx],
275    analytic_rho_count: usize,
276) -> Result<Option<Vec<DirectionalHyperParam>>, EstimationError> {
277    if latent_terms.is_empty() || latent.is_empty() {
278        return Ok(None);
279    }
280    if latent_terms.len() != 1 {
281        crate::bail_invalid_estim!(
282            "LatentCoord standard-fit hyper_dirs currently require exactly one latent smooth term"
283                .to_string(),
284        );
285    }
286    let term_idx = latent_terms[0];
287    let smooth_term = design.smooth.terms.get(term_idx.get()).ok_or_else(|| {
288        EstimationError::InvalidInput(format!(
289            "LatentCoord term index {term_idx} out of bounds for realized smooth design"
290        ))
291    })?;
292    let termspec = resolvedspec
293        .smooth_terms
294        .get(term_idx.get())
295        .ok_or_else(|| {
296            EstimationError::InvalidInput(format!(
297                "LatentCoord term index {term_idx} out of bounds for resolved smooth spec"
298            ))
299        })?;
300    let p_total = design.design.ncols();
301    let smooth_start = p_total.saturating_sub(design.smooth.total_smooth_cols());
302    let global_range = (smooth_start + smooth_term.coeff_range.start)
303        ..(smooth_start + smooth_term.coeff_range.end);
304
305    // Spline bases do not add a separate continuous basis-scale ψ coordinate
306    // here. When they are latent-coordinate terms, their ψ directions are the
307    // latent-coordinate axes below, using the same DirectionalHyperParam layout
308    // as Matérn and Duchon.
309    let operator = match (&termspec.basis, &smooth_term.metadata) {
310        (
311            SmoothBasisSpec::Matern { .. },
312            BasisMetadata::Matern {
313                centers,
314                length_scale,
315                nu,
316                include_intercept,
317                identifiability_transform,
318                ..
319            },
320        ) => gam_terms::basis::LatentCoordDesignDerivative::new_matern(
321            latent.clone(),
322            std::sync::Arc::new(centers.clone()),
323            *length_scale,
324            *nu,
325            *include_intercept,
326            identifiability_transform.clone(),
327        )
328        .map_err(EstimationError::from)?,
329        (
330            SmoothBasisSpec::Duchon { .. },
331            BasisMetadata::Duchon {
332                centers,
333                length_scale,
334                power,
335                nullspace_order,
336                identifiability_transform,
337                ..
338            },
339        ) => gam_terms::basis::LatentCoordDesignDerivative::new_duchon(
340            latent.clone(),
341            std::sync::Arc::new(centers.clone()),
342            *length_scale,
343            *power,
344            *nullspace_order,
345            identifiability_transform.clone(),
346        )
347        .map_err(EstimationError::from)?,
348        (
349            SmoothBasisSpec::Sphere { .. },
350            BasisMetadata::Sphere {
351                centers,
352                penalty_order,
353                method,
354                constraint_transform,
355                ..
356            },
357        ) if matches!(*method, gam_terms::basis::SphereMethod::Wahba) => {
358            gam_terms::basis::LatentCoordDesignDerivative::new_sphere(
359                latent.clone(),
360                std::sync::Arc::new(centers.clone()),
361                *penalty_order,
362                constraint_transform.clone(),
363            )
364            .map_err(EstimationError::from)?
365        }
366        (
367            SmoothBasisSpec::BSpline1D { spec, .. },
368            BasisMetadata::BSpline1D {
369                knots,
370                identifiability_transform,
371                periodic,
372                degree: meta_degree,
373                ..
374            },
375        ) => {
376            // Issue #340: use the metadata-recorded effective degree so the
377            // latent-design Jacobian matches what `build_bspline_basis_1d`
378            // actually built at fit time after auto-shrink.
379            let effective_degree = meta_degree.unwrap_or(spec.degree);
380            if let Some((domain_start, period, num_basis)) = periodic {
381                gam_terms::basis::LatentCoordDesignDerivative::new_periodic_bspline(
382                    latent.clone(),
383                    (*domain_start, *domain_start + *period),
384                    effective_degree,
385                    *num_basis,
386                    identifiability_transform.clone(),
387                )
388                .map_err(EstimationError::from)?
389            } else {
390                gam_terms::basis::LatentCoordDesignDerivative::new_tensor_bspline(
391                    latent.clone(),
392                    vec![knots.clone()],
393                    vec![effective_degree],
394                    identifiability_transform.clone(),
395                )
396                .map_err(EstimationError::from)?
397            }
398        }
399        (
400            SmoothBasisSpec::TensorBSpline { .. },
401            BasisMetadata::TensorBSpline {
402                knots,
403                degrees,
404                identifiability_transform,
405                ..
406            },
407        ) => gam_terms::basis::LatentCoordDesignDerivative::new_tensor_bspline(
408            latent.clone(),
409            knots.clone(),
410            degrees.clone(),
411            identifiability_transform.clone(),
412        )
413        .map_err(EstimationError::from)?,
414        (SmoothBasisSpec::Pca { .. }, BasisMetadata::Pca { basis_matrix, .. }) => {
415            gam_terms::basis::LatentCoordDesignDerivative::new_pca(
416                latent.clone(),
417                std::sync::Arc::new(basis_matrix.clone()),
418            )
419            .map_err(EstimationError::from)?
420        }
421        _ => return Ok(None),
422    };
423    if operator.p_out() != global_range.len() {
424        crate::bail_invalid_estim!(
425            "LatentCoord derivative width mismatch for term '{}': operator p={}, coeff range={}",
426            smooth_term.name,
427            operator.p_out(),
428            global_range.len()
429        );
430    }
431    let operator = std::sync::Arc::new(operator);
432    let mut hyper_dirs = Vec::with_capacity(operator.n_axes());
433    for flat_axis in 0..operator.n_axes() {
434        let dir = DirectionalHyperParam::new_compact(
435            gam_solve::estimate::reml::HyperDesignDerivative::from_latent_coord(
436                operator.clone(),
437                flat_axis,
438                global_range.clone(),
439                p_total,
440            ),
441            Vec::new(),
442            None,
443            None,
444        )?
445        .not_penalty_like();
446        hyper_dirs.push(dir);
447    }
448    let direct_dim = latent_coord_direct_hyper_count(latent.id_mode(), latent.latent_dim());
449    if analytic_rho_count + direct_dim > 0 {
450        let zero_x = gam_solve::estimate::reml::HyperDesignDerivative::from(Array2::<f64>::zeros(
451            (design.design.nrows(), p_total),
452        ));
453        for _ in 0..analytic_rho_count {
454            hyper_dirs.push(
455                DirectionalHyperParam::new_compact(zero_x.clone(), Vec::new(), None, None)?
456                    .not_penalty_like(),
457            );
458        }
459        for _ in 0..direct_dim {
460            hyper_dirs.push(
461                DirectionalHyperParam::new_compact(zero_x.clone(), Vec::new(), None, None)?
462                    .not_penalty_like(),
463            );
464        }
465    }
466    Ok(Some(hyper_dirs))
467}
468
469fn latent_coord_direct_hyper_count(
470    id_mode: &gam_terms::latent::LatentIdMode,
471    latent_dim: usize,
472) -> usize {
473    use gam_terms::latent::{AuxPriorStrength, LatentIdMode};
474    match id_mode {
475        LatentIdMode::AuxPrior { strength, .. } => match strength {
476            AuxPriorStrength::Auto => 1,
477            AuxPriorStrength::Fixed(_) => 0,
478        },
479        LatentIdMode::AuxPriorDimSelection { strength, .. } => {
480            latent_dim
481                + match strength {
482                    AuxPriorStrength::Auto => 1,
483                    AuxPriorStrength::Fixed(_) => 0,
484                }
485        }
486        LatentIdMode::DimSelection { .. } => latent_dim,
487        // A fixed-reference anchor carries at most the REML-selectable log-`μ`
488        // (one direct hyper when `Auto`, none when `Fixed`), like `AuxPrior`.
489        LatentIdMode::IsometryToReference { strength, .. } => match strength {
490            AuxPriorStrength::Auto => 1,
491            AuxPriorStrength::Fixed(_) => 0,
492        },
493        // The behavioral head appends one (1 + d) coefficient block per
494        // η-channel, plus the composed per-axis ARD log-precisions.
495        LatentIdMode::AuxOutcome { head, .. } => head.n_coeffs(latent_dim) + latent_dim,
496        LatentIdMode::None => 0,
497    }
498}
499
500fn latent_coord_initial_direct_hypers(
501    id_mode: &gam_terms::latent::LatentIdMode,
502    latent_dim: usize,
503) -> Result<Array1<f64>, EstimationError> {
504    use gam_terms::latent::{AuxPriorStrength, LatentIdMode};
505    let mut values = Vec::with_capacity(latent_coord_direct_hyper_count(id_mode, latent_dim));
506    match id_mode {
507        LatentIdMode::AuxPrior { strength, .. } => {
508            if matches!(strength, AuxPriorStrength::Auto) {
509                values.push(0.0);
510            }
511        }
512        LatentIdMode::AuxPriorDimSelection {
513            strength,
514            init_log_precision,
515            ..
516        } => {
517            if matches!(strength, AuxPriorStrength::Auto) {
518                values.push(0.0);
519            }
520            append_latent_ard_seed(&mut values, init_log_precision.as_ref(), latent_dim)?;
521        }
522        LatentIdMode::DimSelection { init_log_precision } => {
523            append_latent_ard_seed(&mut values, init_log_precision.as_ref(), latent_dim)?;
524        }
525        LatentIdMode::IsometryToReference { strength, .. } => {
526            if matches!(strength, AuxPriorStrength::Auto) {
527                values.push(0.0);
528            }
529        }
530        LatentIdMode::AuxOutcome {
531            head,
532            init_log_precision,
533        } => {
534            // Head coefficients seed at zero: intercept 0 ⇒ baseline rate, all
535            // loadings 0 ⇒ no behavioral anchoring at start (REML/Newton move
536            // them). One (1 + d) block per η-channel.
537            values.extend(std::iter::repeat_n(0.0, head.n_coeffs(latent_dim)));
538            append_latent_ard_seed(&mut values, init_log_precision.as_ref(), latent_dim)?;
539        }
540        LatentIdMode::None => {}
541    }
542    Ok(Array1::from_vec(values))
543}
544
545fn append_latent_ard_seed(
546    values: &mut Vec<f64>,
547    init: Option<&Array1<f64>>,
548    latent_dim: usize,
549) -> Result<(), EstimationError> {
550    if let Some(init) = init {
551        if init.len() != latent_dim {
552            crate::bail_invalid_estim!(
553                "latent dim_selection init_log_precision length mismatch: got {}, expected {}",
554                init.len(),
555                latent_dim
556            );
557        }
558        values.extend(init.iter().copied());
559    } else {
560        values.extend(std::iter::repeat_n(0.0, latent_dim));
561    }
562    Ok(())
563}
564
565struct LatentIdObjectiveContribution {
566    cost: f64,
567    gradient: Array1<f64>,
568}
569
570fn latent_id_objective_contribution(
571    theta: &Array1<f64>,
572    rho_dim: usize,
573    analytic_rho_count: usize,
574    latent: &gam_terms::latent::LatentCoordValues,
575) -> Result<LatentIdObjectiveContribution, EstimationError> {
576    use gam_terms::latent::{AuxPriorStrength, LatentIdMode, aux_prior_targets};
577    let n_obs = latent.n_obs();
578    let latent_dim = latent.latent_dim();
579    let flat_len = latent.len();
580    let mut gradient = Array1::<f64>::zeros(theta.len());
581    let t_start = rho_dim;
582    let direct_start = t_start + flat_len + analytic_rho_count;
583    if theta.len() < direct_start {
584        crate::bail_invalid_estim!(
585            "latent-coordinate theta too short for id objective: got {}, need at least {}",
586            theta.len(),
587            direct_start
588        );
589    }
590    let t = latent.as_matrix();
591    let mut cost = 0.0;
592    let mut cursor = direct_start;
593
594    match latent.id_mode() {
595        LatentIdMode::AuxPrior {
596            u,
597            family,
598            strength,
599        }
600        | LatentIdMode::AuxPriorDimSelection {
601            u,
602            family,
603            strength,
604            ..
605        } => {
606            let (log_mu, mu) = match strength {
607                AuxPriorStrength::Fixed(mu) => (
608                    gam_problem::checked_log_strength(*mu).map_err(|error| {
609                        EstimationError::InvalidInput(format!(
610                            "fixed latent auxiliary-prior precision is outside the canonical physical-strength domain: {error}"
611                        ))
612                    })?,
613                    *mu,
614                ),
615                AuxPriorStrength::Auto => {
616                    let log_mu = *theta.get(cursor).ok_or_else(|| {
617                        EstimationError::InvalidInput(format!(
618                            "latent auxiliary-prior precision coordinate {cursor} is missing from theta length {}",
619                            theta.len(),
620                        ))
621                    })?;
622                    cursor += 1;
623                    let mu = gam_problem::checked_exp_log_strength(log_mu).map_err(|error| {
624                        EstimationError::InvalidInput(format!(
625                            "latent auxiliary-prior log precision is outside the canonical log-strength domain: {error}"
626                        ))
627                    })?;
628                    (log_mu, mu)
629                }
630            };
631            let targets = aux_prior_targets(t.view(), u.view(), *family)
632                .map_err(EstimationError::InvalidInput)?;
633            let residual = &t - &targets;
634            let q = residual.iter().map(|v| v * v).sum::<f64>();
635            // The single shared precision `mu` governs every one of the
636            // `n_obs · latent_dim` scalar latent coordinates, so the prior
637            // log-determinant normalizer `−0.5·log det₊(mu · I_K)` counts
638            // `K = n_obs · latent_dim`. (The per-axis ARD path below emits
639            // `−0.5·n_obs·ln(α)` for each of `latent_dim` axes; one shared `mu`
640            // must equal that sum.)
641            let k = (n_obs * latent_dim) as f64;
642            cost += 0.5 * mu * q - 0.5 * k * log_mu;
643
644            let projected_residual = aux_prior_targets(residual.view(), u.view(), *family)
645                .map_err(EstimationError::InvalidInput)?;
646            let grad_base = residual - projected_residual;
647            for n in 0..n_obs {
648                for axis in 0..latent_dim {
649                    gradient[t_start + n * latent_dim + axis] += mu * grad_base[[n, axis]];
650                }
651            }
652            if matches!(strength, AuxPriorStrength::Auto) {
653                gradient[direct_start] += 0.5 * mu * q - 0.5 * k;
654            }
655        }
656        LatentIdMode::IsometryToReference {
657            reference,
658            strength,
659        } => {
660            // Fixed-reference anchor `½ μ ‖t − reference‖²` with REML-selectable
661            // `μ`. Identical structure to `AuxPrior` except the target is a
662            // constant configuration (independent of `t`), so the latent
663            // gradient is the plain `μ · (t − reference)` with no projection
664            // term (`AuxPrior` subtracts the projected residual only because its
665            // target `ĥ(u)` depends on `t` through the internal ridge fit).
666            if reference.dim() != (n_obs, latent_dim) {
667                crate::bail_invalid_estim!(
668                    "IsometryToReference reference shape {:?} must equal (n_obs, latent_dim) = ({}, {})",
669                    reference.dim(),
670                    n_obs,
671                    latent_dim
672                );
673            }
674            let mu_slot = cursor;
675            let (log_mu, mu) = match strength {
676                AuxPriorStrength::Fixed(mu) => (
677                    gam_problem::checked_log_strength(*mu).map_err(|error| {
678                        EstimationError::InvalidInput(format!(
679                            "fixed latent isometry precision is outside the canonical physical-strength domain: {error}"
680                        ))
681                    })?,
682                    *mu,
683                ),
684                AuxPriorStrength::Auto => {
685                    let log_mu = *theta.get(cursor).ok_or_else(|| {
686                        EstimationError::InvalidInput(format!(
687                            "latent isometry precision coordinate {cursor} is missing from theta length {}",
688                            theta.len(),
689                        ))
690                    })?;
691                    cursor += 1;
692                    let mu = gam_problem::checked_exp_log_strength(log_mu).map_err(|error| {
693                        EstimationError::InvalidInput(format!(
694                            "latent isometry log precision is outside the canonical log-strength domain: {error}"
695                        ))
696                    })?;
697                    (log_mu, mu)
698                }
699            };
700            let residual = &t - reference;
701            let q = residual.iter().map(|v| v * v).sum::<f64>();
702            // Shared precision `mu` over all `K = n_obs · latent_dim` scalar
703            // coordinates: the normalizer `−0.5·log det₊(mu · I_K)` counts `K`,
704            // matching the AuxPrior arm and the ARD path's per-axis sum.
705            let k = (n_obs * latent_dim) as f64;
706            cost += 0.5 * mu * q - 0.5 * k * log_mu;
707            for n in 0..n_obs {
708                for axis in 0..latent_dim {
709                    gradient[t_start + n * latent_dim + axis] += mu * residual[[n, axis]];
710                }
711            }
712            if matches!(strength, AuxPriorStrength::Auto) {
713                gradient[mu_slot] += 0.5 * mu * q - 0.5 * k;
714            }
715        }
716        LatentIdMode::AuxOutcome { head, .. } => {
717            // Behavioral head likelihood channel: the head's design columns are
718            // the live latent codes, so its NLL enters the SAME joint objective
719            // as the reconstruction term and REML balances the two channels.
720            // The head coefficients occupy `head.n_coeffs(d)` direct-hyper slots
721            // starting at `cursor`; their gradient drives the β-tier update and
722            // the head's latent-code gradient flows into the `t` block (the
723            // arrow-Schur cross-channel coupling).
724            let n_coeffs = head.n_coeffs(latent_dim);
725            if cursor + n_coeffs > theta.len() {
726                crate::bail_invalid_estim!(
727                    "latent auxiliary-outcome coefficient block overruns theta: start={cursor}, width={n_coeffs}, theta_len={}",
728                    theta.len(),
729                );
730            }
731            let coeffs = theta
732                .slice(ndarray::s![cursor..cursor + n_coeffs])
733                .to_owned();
734            let (head_nll, grad_coeffs, grad_t) = head
735                .neg_loglik_and_grad(t.view(), coeffs.view())
736                .map_err(EstimationError::InvalidInput)?;
737            cost += head_nll;
738            for (offset, &g) in grad_coeffs.iter().enumerate() {
739                gradient[cursor + offset] += g;
740            }
741            for n in 0..n_obs {
742                for axis in 0..latent_dim {
743                    gradient[t_start + n * latent_dim + axis] += grad_t[[n, axis]];
744                }
745            }
746            cursor += n_coeffs;
747        }
748        LatentIdMode::DimSelection { .. } | LatentIdMode::None => {}
749    }
750
751    match latent.id_mode() {
752        LatentIdMode::AuxPriorDimSelection { .. }
753        | LatentIdMode::DimSelection { .. }
754        | LatentIdMode::AuxOutcome { .. } => {
755            if cursor + latent_dim > theta.len() {
756                crate::bail_invalid_estim!(
757                    "latent dimension-selection precision block overruns theta: start={cursor}, width={latent_dim}, theta_len={}",
758                    theta.len(),
759                );
760            }
761            let alphas = gam_problem::checked_exp_log_strengths(
762                theta.slice(s![cursor..cursor + latent_dim]).iter().copied(),
763            )
764            .map_err(|error| {
765                EstimationError::InvalidInput(format!(
766                    "latent dimension-selection log precision is outside the canonical log-strength domain: {error}"
767                ))
768            })?;
769            for axis in 0..latent_dim {
770                let log_alpha = theta[cursor + axis];
771                let alpha = alphas[axis];
772                let mut q_axis = 0.0;
773                for n in 0..n_obs {
774                    let flat_idx = n * latent_dim + axis;
775                    let value = latent.as_flat()[flat_idx];
776                    q_axis += value * value;
777                    gradient[t_start + flat_idx] += alpha * value;
778                }
779                cost += 0.5 * alpha * q_axis - 0.5 * n_obs as f64 * log_alpha;
780                gradient[cursor + axis] += 0.5 * alpha * q_axis - 0.5 * n_obs as f64;
781            }
782            cursor += latent_dim;
783        }
784        LatentIdMode::AuxPrior { .. }
785        | LatentIdMode::IsometryToReference { .. }
786        | LatentIdMode::None => {}
787    }
788
789    if cursor != theta.len() {
790        crate::bail_invalid_estim!(
791            "latent-coordinate direct hyperparameter length mismatch: consumed {}, theta len {}",
792            cursor,
793            theta.len()
794        );
795    }
796    Ok(LatentIdObjectiveContribution { cost, gradient })
797}
798
799fn add_latent_id_objective_to_eval(
800    theta: &Array1<f64>,
801    rho_dim: usize,
802    analytic_rho_count: usize,
803    latent: &gam_terms::latent::LatentCoordValues,
804    eval: &mut (f64, Array1<f64>, gam_problem::HessianValue),
805) -> Result<(), EstimationError> {
806    let contribution =
807        latent_id_objective_contribution(theta, rho_dim, analytic_rho_count, latent)?;
808    eval.0 += contribution.cost;
809    if eval.1.len() != contribution.gradient.len() {
810        crate::bail_invalid_estim!(
811            "latent-coordinate REML gradient length mismatch: base={}, id={}",
812            eval.1.len(),
813            contribution.gradient.len()
814        );
815    }
816    eval.1 += &contribution.gradient;
817    if eval.2.is_analytic() {
818        eval.2 = gam_problem::HessianValue::Unavailable;
819    }
820    Ok(())
821}
822
823fn analytic_penalty_objective_contribution(
824    theta: &Array1<f64>,
825    rho_dim: usize,
826    latent: &gam_terms::latent::LatentCoordValues,
827    registry: &gam_terms::AnalyticPenaltyRegistry,
828) -> Result<LatentIdObjectiveContribution, EstimationError> {
829    let flat_len = latent.len();
830    let t_start = rho_dim;
831    let t_end = t_start + flat_len;
832    let rho_start = t_end;
833    let rho_end = rho_start + registry.total_rho_count();
834    if theta.len() < rho_end {
835        crate::bail_invalid_estim!(
836            "latent-coordinate theta too short for analytic penalties: got {}, need at least {}",
837            theta.len(),
838            rho_end
839        );
840    }
841    let target_t = theta.slice(s![t_start..t_end]);
842    let rho = theta.slice(s![rho_start..rho_end]);
843    registry
844        .validate_rho(rho)
845        .map_err(EstimationError::InvalidInput)?;
846    let mut cost = 0.0_f64;
847    let mut gradient = Array1::<f64>::zeros(theta.len());
848    for (penalty, (rho_slice, tier, name)) in registry.penalties.iter().zip(registry.rho_layout()) {
849        let rho_local = rho.slice(s![rho_slice.clone()]);
850        match tier {
851            gam_terms::PenaltyTier::Psi => {
852                cost += penalty.value(target_t.view(), rho_local);
853                let grad = penalty.grad_target(target_t.view(), rho_local);
854                if grad.len() != flat_len {
855                    crate::bail_invalid_estim!(
856                        "analytic penalty {name:?} gradient length mismatch: got {}, expected {}",
857                        grad.len(),
858                        flat_len
859                    );
860                }
861                for i in 0..flat_len {
862                    gradient[t_start + i] += grad[i];
863                }
864                let grad_rho_local = penalty.grad_rho(target_t.view(), rho_local);
865                if grad_rho_local.len() != rho_slice.len() {
866                    crate::bail_invalid_estim!(
867                        "analytic penalty {name:?} rho-gradient length mismatch: got {}, expected {}",
868                        grad_rho_local.len(),
869                        rho_slice.len()
870                    );
871                }
872                for local_idx in 0..grad_rho_local.len() {
873                    gradient[rho_start + rho_slice.start + local_idx] += grad_rho_local[local_idx];
874                }
875            }
876            gam_terms::PenaltyTier::Beta => {}
877            gam_terms::PenaltyTier::Rho => {}
878        }
879    }
880    Ok(LatentIdObjectiveContribution { cost, gradient })
881}
882
883fn add_analytic_penalty_hessian_to_eval(
884    theta: &Array1<f64>,
885    rho_dim: usize,
886    latent: &gam_terms::latent::LatentCoordValues,
887    registry: &gam_terms::AnalyticPenaltyRegistry,
888    eval: &mut (f64, Array1<f64>, gam_problem::HessianValue),
889) -> Result<(), EstimationError> {
890    let flat_len = latent.len();
891    let t_start = rho_dim;
892    let t_end = t_start + flat_len;
893    let rho_start = t_end;
894    let rho_end = rho_start + registry.total_rho_count();
895    if theta.len() < rho_end {
896        crate::bail_invalid_estim!(
897            "latent-coordinate theta too short for analytic penalty Hessian: got {}, need at least {}",
898            theta.len(),
899            rho_end
900        );
901    }
902    let gam_problem::HessianValue::Dense(hessian) = &mut eval.2 else {
903        if eval.2.is_analytic() {
904            eval.2 = gam_problem::HessianValue::Unavailable;
905        }
906        return Ok(());
907    };
908    if hessian.dim() != (theta.len(), theta.len()) {
909        crate::bail_invalid_estim!(
910            "analytic penalty Hessian target shape mismatch: got {}x{}, expected {}x{}",
911            hessian.nrows(),
912            hessian.ncols(),
913            theta.len(),
914            theta.len()
915        );
916    }
917    let target_t = theta.slice(s![t_start..t_end]);
918    let rho = theta.slice(s![rho_start..rho_end]);
919    registry
920        .validate_rho(rho)
921        .map_err(EstimationError::InvalidInput)?;
922    for (penalty, (rho_slice, tier, _name)) in registry.penalties.iter().zip(registry.rho_layout())
923    {
924        let rho_local = rho.slice(s![rho_slice]);
925        if !matches!(tier, gam_terms::PenaltyTier::Psi) {
926            continue;
927        }
928        if let Some(diag) = penalty.hessian_diag(target_t.view(), rho_local) {
929            if diag.len() != flat_len {
930                crate::bail_invalid_estim!(
931                    "analytic penalty Hessian diagonal length mismatch: got {}, expected {}",
932                    diag.len(),
933                    flat_len
934                );
935            }
936            for i in 0..flat_len {
937                hessian[[t_start + i, t_start + i]] += diag[i];
938            }
939            continue;
940        }
941        let mut probe = Array1::<f64>::zeros(flat_len);
942        for col in 0..flat_len {
943            probe[col] = 1.0;
944            let hv = penalty.hvp(target_t.view(), rho_local, probe.view());
945            if hv.len() != flat_len {
946                crate::bail_invalid_estim!(
947                    "analytic penalty Hessian-vector length mismatch: got {}, expected {}",
948                    hv.len(),
949                    flat_len
950                );
951            }
952            for row in 0..flat_len {
953                hessian[[t_start + row, t_start + col]] += hv[row];
954            }
955            probe[col] = 0.0;
956        }
957    }
958    Ok(())
959}
960
961fn add_analytic_penalty_objective_to_eval(
962    theta: &Array1<f64>,
963    rho_dim: usize,
964    latent: &gam_terms::latent::LatentCoordValues,
965    registry: &gam_terms::AnalyticPenaltyRegistry,
966    eval: &mut (f64, Array1<f64>, gam_problem::HessianValue),
967) -> Result<(), EstimationError> {
968    let contribution = analytic_penalty_objective_contribution(theta, rho_dim, latent, registry)?;
969    eval.0 += contribution.cost;
970    if eval.1.len() != contribution.gradient.len() {
971        crate::bail_invalid_estim!(
972            "latent-coordinate REML gradient length mismatch: base={}, analytic_penalty={}",
973            eval.1.len(),
974            contribution.gradient.len()
975        );
976    }
977    eval.1 += &contribution.gradient;
978    add_analytic_penalty_hessian_to_eval(theta, rho_dim, latent, registry, eval)?;
979    Ok(())
980}
981
982fn spatial_log_kappa_hyper_dirs_frominfo_list(
983    info_list: Vec<SpatialPsiDerivative>,
984) -> Result<Vec<DirectionalHyperParam>, EstimationError> {
985    use gam_solve::estimate::reml::ImplicitDerivLevel;
986    use std::collections::HashMap;
987
988    let log_kappa_dim = info_list.len();
989    // Layout-only metadata (group_id per axis) is cheap to snapshot up front so
990    // the consumption loop below can MOVE the dense (n × p) derivative arrays
991    // out of each entry instead of cloning. At large scale (n≈3×10⁵, 16-axis
992    // CTN) the prior `.clone()` sites doubled peak working memory for the
993    // psi-derivative pass through several GiB.
994    let group_ids: Vec<Option<usize>> = info_list.iter().map(|e| e.aniso_group_id).collect();
995    let mut group_indices_map: HashMap<usize, Vec<usize>> = HashMap::new();
996    for (idx, gid) in group_ids.iter().enumerate() {
997        if let Some(g) = gid {
998            group_indices_map.entry(*g).or_default().push(idx);
999        }
1000    }
1001
1002    let mut hyper_dirs = Vec::with_capacity(log_kappa_dim);
1003    for (i, info) in info_list.into_iter().enumerate() {
1004        let SpatialPsiDerivative {
1005            penalty_index: _,
1006            penalty_indices,
1007            global_range,
1008            total_p,
1009            x_psi_local,
1010            s_psi_components_local,
1011            x_psi_psi_local,
1012            s_psi_psi_components_local,
1013            aniso_group_id,
1014            aniso_cross_designs,
1015            aniso_cross_penalty_provider,
1016            implicit_operator,
1017            implicit_axis,
1018        } = info;
1019
1020        let mut xsecond = vec![None; log_kappa_dim];
1021        // Diagonal second derivative (same axis).
1022        xsecond[i] = Some(if let Some(ref op) = implicit_operator {
1023            gam_solve::estimate::reml::HyperDesignDerivative::from_implicit(
1024                op.clone(),
1025                ImplicitDerivLevel::SecondDiag(implicit_axis),
1026                global_range.clone(),
1027                total_p,
1028            )
1029        } else {
1030            gam_solve::estimate::reml::HyperDesignDerivative::from_embedded(
1031                x_psi_psi_local,
1032                global_range.clone(),
1033                total_p,
1034            )
1035        });
1036        // Cross second derivatives for axes in the same aniso group.
1037        if let Some(cross_designs) = aniso_cross_designs {
1038            // Use the base index of this aniso group in the original info_list.
1039            // Entries for the same group are contiguous: the first index in the
1040            // group gives the base, and axis b is at base+b.
1041            if let Some(gid) = aniso_group_id {
1042                let base = group_indices_map
1043                    .get(&gid)
1044                    .and_then(|v| v.first().copied())
1045                    .unwrap_or(i);
1046                for (b_axis, cross_mat) in cross_designs.into_iter() {
1047                    let j = base + b_axis;
1048                    if j < log_kappa_dim {
1049                        xsecond[j] = Some(if let Some(ref op) = implicit_operator {
1050                            gam_solve::estimate::reml::HyperDesignDerivative::from_implicit(
1051                                op.clone(),
1052                                ImplicitDerivLevel::SecondCross(implicit_axis, b_axis),
1053                                global_range.clone(),
1054                                total_p,
1055                            )
1056                        } else {
1057                            gam_solve::estimate::reml::HyperDesignDerivative::from_embedded(
1058                                cross_mat,
1059                                global_range.clone(),
1060                                total_p,
1061                            )
1062                        });
1063                    }
1064                }
1065            }
1066        }
1067        let s_components = penalty_indices
1068            .iter()
1069            .copied()
1070            .zip(s_psi_components_local.into_iter().map(|local| {
1071                gam_solve::estimate::reml::HyperPenaltyDerivative::from_embedded(
1072                    local,
1073                    global_range.clone(),
1074                    total_p,
1075                )
1076            }))
1077            .collect::<Vec<_>>();
1078        let s2_components = penalty_indices
1079            .iter()
1080            .copied()
1081            .zip(s_psi_psi_components_local.into_iter().map(|local| {
1082                gam_solve::estimate::reml::HyperPenaltyDerivative::from_embedded(
1083                    local,
1084                    global_range.clone(),
1085                    total_p,
1086                )
1087            }))
1088            .collect::<Vec<_>>();
1089        let mut ssecond_components = vec![None; log_kappa_dim];
1090        ssecond_components[i] = Some(s2_components);
1091        let mut penaltysecond_partner_indices: Option<Vec<usize>> = None;
1092        let penaltysecond_component_provider =
1093            if let (Some(provider), Some(gid)) = (aniso_cross_penalty_provider, aniso_group_id) {
1094                let group_indices = group_indices_map.get(&gid).cloned().unwrap_or_default();
1095                let axis_in_group =
1096                    group_indices
1097                        .iter()
1098                        .position(|&idx| idx == i)
1099                        .ok_or_else(|| {
1100                            EstimationError::InvalidInput(format!(
1101                                "missing spatial hyper axis {} in anisotropy group {}",
1102                                i, gid
1103                            ))
1104                        })?;
1105                penaltysecond_partner_indices = Some(
1106                    group_indices
1107                        .iter()
1108                        .copied()
1109                        .filter(|&idx| idx != i)
1110                        .collect(),
1111                );
1112                let penalty_indices_inner = penalty_indices.clone();
1113                let global_range_inner = global_range.clone();
1114                let total_p_inner = total_p;
1115                let group_indices_inner = group_indices;
1116                Some(std::sync::Arc::new(
1117                    move |j: usize| -> Result<
1118                        Option<Vec<gam_solve::estimate::reml::PenaltyDerivativeComponent>>,
1119                        EstimationError,
1120                    > {
1121                        let Some(other_axis_in_group) =
1122                            group_indices_inner.iter().position(|&idx| idx == j)
1123                        else {
1124                            return Ok(None);
1125                        };
1126                        if other_axis_in_group == axis_in_group {
1127                            return Ok(None);
1128                        }
1129                        let cross_pens = provider(other_axis_in_group)?;
1130                        if cross_pens.is_empty() {
1131                            return Ok(None);
1132                        }
1133                        Ok(Some(
1134                        penalty_indices_inner
1135                            .iter()
1136                            .copied()
1137                            .zip(cross_pens.into_iter().map(|local| {
1138                                gam_solve::estimate::reml::HyperPenaltyDerivative::from_embedded(
1139                                    local,
1140                                    global_range_inner.clone(),
1141                                    total_p_inner,
1142                                )
1143                            }))
1144                            .map(|(penalty_index, matrix)| {
1145                                gam_solve::estimate::reml::PenaltyDerivativeComponent {
1146                                    penalty_index,
1147                                    matrix,
1148                                }
1149                            })
1150                            .collect(),
1151                    ))
1152                    },
1153                )
1154                    as std::sync::Arc<
1155                        dyn Fn(
1156                                usize,
1157                            ) -> Result<
1158                                Option<Vec<gam_solve::estimate::reml::PenaltyDerivativeComponent>>,
1159                                EstimationError,
1160                            > + Send
1161                            + Sync
1162                            + 'static,
1163                    >)
1164            } else {
1165                None
1166            };
1167        // First derivative: use implicit operator when available to avoid
1168        // storing dense (n x p) matrices for all D axes simultaneously.
1169        let x_first_hyper = if let Some(ref op) = implicit_operator {
1170            gam_solve::estimate::reml::HyperDesignDerivative::from_implicit(
1171                op.clone(),
1172                ImplicitDerivLevel::First(implicit_axis),
1173                global_range.clone(),
1174                total_p,
1175            )
1176        } else {
1177            gam_solve::estimate::reml::HyperDesignDerivative::from_embedded(
1178                x_psi_local,
1179                global_range.clone(),
1180                total_p,
1181            )
1182        };
1183        let mut dir = DirectionalHyperParam::new_compact(
1184            x_first_hyper,
1185            s_components,
1186            Some(xsecond),
1187            Some(ssecond_components),
1188        )?
1189        .not_penalty_like();
1190        if let Some(provider) = penaltysecond_component_provider {
1191            dir = dir.with_penaltysecond_component_provider(provider);
1192        }
1193        if let Some(partner_indices) = penaltysecond_partner_indices {
1194            dir = dir.with_penaltysecond_partner_indices(partner_indices);
1195        }
1196        hyper_dirs.push(dir);
1197    }
1198    Ok(hyper_dirs)
1199}
1200
1201/// Compute `dims_per_term` for a list of spatial term indices.
1202///
1203/// Returns a vector where entry i is the number of stored ψ values for
1204/// spatial term i: `d` for terms that enroll per-axis anisotropy in the
1205/// REML joint vector (`spatial_term_uses_per_axis_psi`), `1` otherwise.
1206pub(crate) fn spatial_dims_per_term(
1207    resolvedspec: &TermCollectionSpec,
1208    spatial_terms: &[usize],
1209) -> Vec<usize> {
1210    spatial_terms
1211        .iter()
1212        .map(|&term_idx| {
1213            if let Some(mj) = measure_jet_term_spec(resolvedspec, term_idx) {
1214                // Dial group, not per-axis anisotropy; layout owned by
1215                // `measure_jet_psi_dim`.
1216                measure_jet_psi_dim(mj)
1217            } else if spatial_term_uses_per_axis_psi(resolvedspec, term_idx) {
1218                get_spatial_feature_dim(resolvedspec, term_idx).unwrap_or(1)
1219            } else {
1220                1
1221            }
1222        })
1223        .collect()
1224}
1225
1226/// Check whether any spatial terms enroll per-axis anisotropic ψ in the joint
1227/// outer vector. Mirrors the hyper_dirs builder's enrollment predicate so the
1228/// outer θ-layout cannot drift from the inner evaluator's ψ count.
1229fn has_aniso_terms(resolvedspec: &TermCollectionSpec, spatial_terms: &[usize]) -> bool {
1230    spatial_terms
1231        .iter()
1232        .any(|&term_idx| spatial_term_uses_per_axis_psi(resolvedspec, term_idx))
1233}
1234
1235/// Emits the `theta`-keyed memoization accessors shared verbatim by the
1236/// single-block and n-block exact-joint design caches. Both carry the same
1237/// `current_theta` / `last_cost` / `last_eval` fields, so the cost/eval
1238/// lookups and the `store_eval` writer are identical; this macro is the single
1239/// source so the two inherent impls cannot drift.
1240macro_rules! impl_exact_joint_theta_memo {
1241    () => {
1242        fn memoized_cost(&self, theta: &Array1<f64>) -> Option<f64> {
1243            if self
1244                .current_theta
1245                .as_ref()
1246                .is_some_and(|cached| theta_values_match(cached, theta))
1247            {
1248                self.last_eval
1249                    .as_ref()
1250                    .map(|cached| cached.0)
1251                    .or(self.last_cost)
1252            } else {
1253                None
1254            }
1255        }
1256
1257        fn memoized_eval(
1258            &self,
1259            theta: &Array1<f64>,
1260        ) -> Option<(f64, Array1<f64>, gam_problem::HessianValue)> {
1261            if self
1262                .current_theta
1263                .as_ref()
1264                .is_some_and(|cached| theta_values_match(cached, theta))
1265            {
1266                self.last_eval.clone()
1267            } else {
1268                None
1269            }
1270        }
1271
1272        fn store_eval(&mut self, eval: (f64, Array1<f64>, gam_problem::HessianValue)) {
1273            self.last_cost = Some(eval.0);
1274            self.last_eval = Some(eval);
1275        }
1276    };
1277}
1278
1279struct SingleBlockExactJointDesignCache<'d> {
1280    realizer: FrozenTermCollectionIncrementalRealizer<'d>,
1281    current_theta: Option<Array1<f64>>,
1282    // Memo key for `last_cost`/`last_eval`. Distinct from `current_theta` (which
1283    // tracks the θ the n×k design is REALIZED at): on the #1033 certified
1284    // Gaussian path `eval_full` evaluates a trial ψ WITHOUT re-realizing the
1285    // design (the tensor serves value+gradient n-free), so the eval θ and the
1286    // realized-design θ diverge. Keying the memo on a dedicated field keeps a
1287    // ψ-skip from ever mis-associating one ψ's cost/eval with another ψ's key.
1288    last_eval_theta: Option<Array1<f64>>,
1289    last_cost: Option<f64>,
1290    last_eval: Option<(f64, Array1<f64>, gam_problem::HessianValue)>,
1291    // #1033: ψ-invariant hyper-direction slab cache. The κ hyper_dirs (the n×k
1292    // ∂X/∂ψ design-derivative slabs + their k×k penalty derivatives) are a pure
1293    // function of (data, frozen spec, REALIZED column layout) — they do NOT
1294    // depend on the trial ψ once the design is fixed. On the certified Gaussian
1295    // n-free path `eval_full` evaluates trial ψ WITHOUT re-realizing the design,
1296    // so the realized layout (and hence the hyper_dirs) is identical across an
1297    // entire run of skip-path trials. Rebuilding them each trial re-runs the
1298    // basis ψ-derivative over all n rows + an O(n·k²) `fast_ab` rotation — the
1299    // last per-trial O(n) pass in the κ loop. Cache them keyed by the realizer
1300    // `design_revision`: a skip-path trial (revision unchanged) reuses the
1301    // build; a slow-path trial (revision advanced) rebuilds and re-keys.
1302    cached_hyper_dirs: Option<(u64, Vec<DirectionalHyperParam>)>,
1303    spatial_terms: Vec<usize>,
1304    rho_dim: usize,
1305    dims_per_term: Vec<usize>,
1306}
1307
1308impl<'d> SingleBlockExactJointDesignCache<'d> {
1309    fn new_with_policy(
1310        data: ArrayView2<'d, f64>,
1311        spec: TermCollectionSpec,
1312        design: TermCollectionDesign,
1313        spatial_terms: Vec<usize>,
1314        rho_dim: usize,
1315        dims_per_term: Vec<usize>,
1316        policy: &gam_runtime::resource::ResourcePolicy,
1317    ) -> Result<Self, String> {
1318        Ok(Self {
1319            realizer: FrozenTermCollectionIncrementalRealizer::new_with_policy(
1320                data, spec, design, policy,
1321            )?,
1322            current_theta: None,
1323            last_eval_theta: None,
1324            last_cost: None,
1325            last_eval: None,
1326            cached_hyper_dirs: None,
1327            spatial_terms,
1328            rho_dim,
1329            dims_per_term,
1330        })
1331    }
1332
1333    fn design_revision(&self) -> u64 {
1334        self.realizer.design_revision()
1335    }
1336
1337    /// Build the κ hyper-directions for the CURRENT realized design, reusing the
1338    /// `cached_hyper_dirs` slab when the realizer revision has not advanced since
1339    /// the last build (#1033). The slab is ψ-invariant at a fixed realized
1340    /// layout, so a skip-path trial (which does not re-realize the design) gets a
1341    /// bit-identical clone instead of re-running the per-row basis ψ-derivative +
1342    /// O(n·k²) rotation. A revision change (slow-path re-realization) rebuilds and
1343    /// re-keys. The clone is an O(n·k) memcpy — far cheaper than the O(n·k²)
1344    /// rebuild, and the conditioning pass it feeds is itself skipped on the
1345    /// certified path (see `prepare_eval_state`'s fast path).
1346    fn hyper_dirs_for_current_design(
1347        &mut self,
1348        data: ArrayView2<'_, f64>,
1349        kind: SpatialHyperKind,
1350    ) -> Result<Vec<DirectionalHyperParam>, EstimationError> {
1351        let revision = self.realizer.design_revision();
1352        if let Some((cached_rev, dirs)) = self.cached_hyper_dirs.as_ref()
1353            && *cached_rev == revision
1354        {
1355            return Ok(dirs.clone());
1356        }
1357        let dirs = try_build_spatial_log_kappa_hyper_dirs(
1358            data,
1359            self.realizer.spec(),
1360            self.realizer.design(),
1361            &self.spatial_terms,
1362        )?
1363        .ok_or_else(|| {
1364            EstimationError::InvalidInput(format!(
1365                "failed to build {} hyper_dirs at current {}",
1366                kind.adjective(),
1367                kind.coord_name(),
1368            ))
1369        })?;
1370        self.cached_hyper_dirs = Some((revision, dirs.clone()));
1371        Ok(dirs)
1372    }
1373
1374    fn nfree_tensor_gradient_hyper_dirs(
1375        &mut self,
1376        theta: &Array1<f64>,
1377    ) -> Result<Vec<DirectionalHyperParam>, EstimationError> {
1378        let psi = &theta.as_slice().ok_or_else(|| {
1379            EstimationError::InvalidInput(
1380                "nfree_tensor_gradient_hyper_dirs: theta is not contiguous".to_string(),
1381            )
1382        })?[self.rho_dim..];
1383        let (global_range, p_total, s_psi_components) = self
1384            .realizer
1385            .canonical_penalty_derivatives_at_psi(&self.spatial_terms, psi)
1386            .map_err(EstimationError::InvalidInput)?;
1387        let zero_x = gam_solve::estimate::reml::HyperDesignDerivative::zero(
1388            self.realizer.design().design.nrows(),
1389            p_total,
1390        );
1391        let components = s_psi_components
1392            .into_iter()
1393            .enumerate()
1394            .map(|(penalty_index, local)| {
1395                (
1396                    penalty_index,
1397                    gam_solve::estimate::reml::HyperPenaltyDerivative::from_embedded(
1398                        local,
1399                        global_range.clone(),
1400                        p_total,
1401                    ),
1402                )
1403            })
1404            .collect::<Vec<_>>();
1405        Ok(DirectionalHyperParam::new_compact(zero_x, components, None, None)?.not_penalty_like())
1406            .map(|dir| vec![dir])
1407    }
1408
1409    fn ensure_theta(&mut self, theta: &Array1<f64>) -> Result<(), String> {
1410        if self
1411            .current_theta
1412            .as_ref()
1413            .is_some_and(|cached| theta_values_match(cached, theta))
1414        {
1415            return Ok(());
1416        }
1417        let t_ensure = std::time::Instant::now();
1418        let log_kappa = SpatialLogKappaCoords::from_theta_tail_with_dims(
1419            theta,
1420            self.rho_dim,
1421            self.dims_per_term.clone(),
1422        );
1423        self.realizer
1424            .apply_log_kappa(&log_kappa, &self.spatial_terms)?;
1425        log::info!(
1426            "[STAGE] ensure_theta (apply_log_kappa, {} terms): {:.3}s",
1427            self.spatial_terms.len(),
1428            t_ensure.elapsed().as_secs_f64(),
1429        );
1430        self.current_theta = Some(theta.clone());
1431        self.last_eval_theta = None;
1432        self.last_cost = None;
1433        self.last_eval = None;
1434        Ok(())
1435    }
1436
1437    // Memo methods keyed on `last_eval_theta` (NOT `current_theta`): the #1033
1438    // certified Gaussian path evaluates a trial ψ without re-realizing the
1439    // design, so the eval θ and the realized-design θ can differ. Keying the
1440    // memo on the eval θ keeps a ψ-skip from mis-associating one ψ's result
1441    // with another ψ's key. The other exact-joint caches still use the shared
1442    // `impl_exact_joint_theta_memo!` macro (they always realize before eval).
1443    fn memoized_cost(&self, theta: &Array1<f64>) -> Option<f64> {
1444        if self
1445            .last_eval_theta
1446            .as_ref()
1447            .is_some_and(|cached| theta_values_match(cached, theta))
1448        {
1449            self.last_eval
1450                .as_ref()
1451                .map(|cached| cached.0)
1452                .or(self.last_cost)
1453        } else {
1454            None
1455        }
1456    }
1457
1458    fn memoized_eval(
1459        &self,
1460        theta: &Array1<f64>,
1461    ) -> Option<(f64, Array1<f64>, gam_problem::HessianValue)> {
1462        if self
1463            .last_eval_theta
1464            .as_ref()
1465            .is_some_and(|cached| theta_values_match(cached, theta))
1466        {
1467            self.last_eval.clone()
1468        } else {
1469            None
1470        }
1471    }
1472
1473    /// Record an eval result keyed to the θ it was computed at. Used in place of
1474    /// the macro's `store_eval` so the memo key reflects the EVAL θ even when the
1475    /// design was not re-realized at that θ (#1033 certified skip).
1476    fn store_eval_at(
1477        &mut self,
1478        theta: &Array1<f64>,
1479        eval: (f64, Array1<f64>, gam_problem::HessianValue),
1480    ) {
1481        self.last_eval_theta = Some(theta.clone());
1482        self.last_cost = Some(eval.0);
1483        self.last_eval = Some(eval);
1484    }
1485
1486    /// Record a cost-only result keyed to the θ it was computed at, so
1487    /// `memoized_cost` keys on the EVAL θ (matching `store_eval_at`).
1488    fn store_cost_at(&mut self, theta: &Array1<f64>, cost: f64) {
1489        self.last_eval_theta = Some(theta.clone());
1490        self.last_cost = Some(cost);
1491        // A cost-only probe carries no gradient/Hessian, so drop any prior
1492        // full eval: `memoized_cost` prefers `last_eval.0`, and a stale
1493        // `last_eval` from a different θ must never answer for this θ.
1494        self.last_eval = None;
1495    }
1496
1497    fn spec(&self) -> &TermCollectionSpec {
1498        self.realizer.spec()
1499    }
1500
1501    fn design(&self) -> &TermCollectionDesign {
1502        self.realizer.design()
1503    }
1504
1505    /// True when the single spatial term's frozen geometry admits an EXACT,
1506    /// n-free penalty re-key at a new length-scale (#1033). The κ-loop fast path
1507    /// gates its design-realization skip on this (replacing the old certified
1508    /// `psi_penalty_tensor_covers` gate): the skip leaves `reset_surface`
1509    /// un-run, so it is sound only when `S(ψ_new)` can be rebuilt n-free.
1510    fn supports_nfree_penalty_rekey(&self) -> bool {
1511        self.realizer
1512            .supports_nfree_penalty_rekey(&self.spatial_terms)
1513    }
1514
1515    fn supports_nfree_gradient_only_routing(&self) -> bool {
1516        self.realizer
1517            .supports_nfree_gradient_only_routing(&self.spatial_terms)
1518    }
1519
1520    /// Build the EXACT canonical penalty surface `S(ψ)` at the length-scale
1521    /// implied by `theta`'s ψ tail, entirely n-free (#1033). Maps ψ→length-scale
1522    /// with the IDENTICAL `spatial_term_psi_to_length_scale_and_aniso` the slow
1523    /// path uses, reuses the frozen basis geometry, and runs the SAME
1524    /// `canonicalize_penalty_specs` pipeline `reset_surface` runs — so the
1525    /// returned canonical list is the one the kept reference surface must be
1526    /// re-keyed with on the design-revision fast path. The caller (which holds
1527    /// `cache`) computes this and hands the owned result to the evaluator via
1528    /// `stage_fast_path_penalty`, avoiding a `&mut cache` borrow alias.
1529    fn canonical_penalties_at(
1530        &mut self,
1531        theta: &Array1<f64>,
1532    ) -> Result<(Vec<gam_terms::construction::CanonicalPenalty>, Vec<usize>), String> {
1533        let psi = &theta
1534            .as_slice()
1535            .ok_or_else(|| "canonical_penalties_at: theta is not contiguous".to_string())?
1536            [self.rho_dim..];
1537        self.realizer
1538            .canonical_penalties_at_psi(&self.spatial_terms, psi)
1539    }
1540}
1541
1542struct SingleBlockLatentCoordDesignCache {
1543    data: Array2<f64>,
1544    spec: TermCollectionSpec,
1545    design: TermCollectionDesign,
1546    current_theta: Option<Array1<f64>>,
1547    current_latent: Option<std::sync::Arc<gam_terms::latent::LatentCoordValues>>,
1548    current_hyper_dirs: Option<Vec<gam_solve::estimate::reml::DirectionalHyperParam>>,
1549    current_design_cache_id: Option<u64>,
1550    latent_design_cache: gam_solve::latent_cache::LatentDesignCache,
1551    last_cost: Option<f64>,
1552    last_eval: Option<(f64, Array1<f64>, gam_problem::HessianValue)>,
1553    term_index: gam_problem::types::SmoothTermIdx,
1554    feature_cols: Vec<usize>,
1555    rho_dim: usize,
1556    n_obs: usize,
1557    latent_dim: usize,
1558    id_mode: gam_terms::latent::LatentIdMode,
1559    manifold: gam_terms::latent::LatentManifold,
1560    retraction_registry: gam_solve::latent_cache::LatentRetractionRegistry,
1561    latent_id: u64,
1562    analytic_penalties: Option<std::sync::Arc<gam_terms::AnalyticPenaltyRegistry>>,
1563    analytic_rho_count: usize,
1564    design_revision: u64,
1565    // Stamp the outer-iter the cached cost/eval was computed under; analytic
1566    // penalty weight schedules advance with this counter, so a stale stamp
1567    // invalidates the memo even at unchanged θ.
1568    last_outer_iter: Option<u64>,
1569}
1570
1571impl SingleBlockLatentCoordDesignCache {
1572    fn new(
1573        data: Array2<f64>,
1574        spec: TermCollectionSpec,
1575        design: TermCollectionDesign,
1576        latent: &StandardLatentCoordConfig,
1577        rho_dim: usize,
1578    ) -> Result<Self, String> {
1579        if latent.term_index.get() >= spec.smooth_terms.len() {
1580            return Err(SmoothError::dimension_mismatch(format!(
1581                "latent-coordinate term index {} out of bounds for {} smooth terms",
1582                latent.term_index,
1583                spec.smooth_terms.len()
1584            ))
1585            .into());
1586        }
1587        if latent.feature_cols.len() != latent.values.latent_dim() {
1588            return Err(SmoothError::dimension_mismatch(format!(
1589                "latent-coordinate feature width mismatch: feature_cols={}, latent_dim={}",
1590                latent.feature_cols.len(),
1591                latent.values.latent_dim()
1592            ))
1593            .into());
1594        }
1595        if latent.values.n_obs() != data.nrows() {
1596            return Err(SmoothError::dimension_mismatch(format!(
1597                "latent-coordinate row mismatch: latent n={}, data n={}",
1598                latent.values.n_obs(),
1599                data.nrows()
1600            ))
1601            .into());
1602        }
1603        let analytic_rho_count = latent
1604            .analytic_penalties
1605            .as_ref()
1606            .map_or(0, |registry| registry.total_rho_count());
1607        Ok(Self {
1608            data,
1609            spec,
1610            design,
1611            current_theta: None,
1612            current_latent: None,
1613            current_hyper_dirs: None,
1614            current_design_cache_id: None,
1615            latent_design_cache: gam_solve::latent_cache::LatentDesignCache::default(),
1616            last_cost: None,
1617            last_eval: None,
1618            term_index: latent.term_index,
1619            feature_cols: latent.feature_cols.clone(),
1620            rho_dim,
1621            n_obs: latent.values.n_obs(),
1622            latent_dim: latent.values.latent_dim(),
1623            id_mode: latent.values.id_mode().clone(),
1624            manifold: latent.values.manifold().clone(),
1625            retraction_registry: latent.values.retraction_registry().clone(),
1626            latent_id: latent.values.latent_id(),
1627            analytic_penalties: latent.analytic_penalties.clone(),
1628            analytic_rho_count,
1629            design_revision: 0,
1630            last_outer_iter: None,
1631        })
1632    }
1633
1634    fn design_revision(&self) -> u64 {
1635        self.design_revision
1636    }
1637
1638    fn design(&self) -> &TermCollectionDesign {
1639        &self.design
1640    }
1641
1642    fn latent(&self) -> Result<std::sync::Arc<gam_terms::latent::LatentCoordValues>, String> {
1643        self.current_latent
1644            .as_ref()
1645            .cloned()
1646            .ok_or_else(|| "latent-coordinate cache has not been realized".to_string())
1647    }
1648
1649    fn analytic_penalties(&self) -> Option<std::sync::Arc<gam_terms::AnalyticPenaltyRegistry>> {
1650        self.analytic_penalties.clone()
1651    }
1652
1653    fn analytic_penalty_rho_count(&self) -> usize {
1654        self.analytic_rho_count
1655    }
1656
1657    fn hyper_dirs(&self) -> Result<Vec<gam_solve::estimate::reml::DirectionalHyperParam>, String> {
1658        self.current_hyper_dirs
1659            .as_ref()
1660            .cloned()
1661            .ok_or_else(|| "latent-coordinate hyper_dirs cache has not been realized".to_string())
1662    }
1663
1664    fn latent_basis_kind(&self) -> Result<gam_solve::latent_cache::LatentBasisKind, String> {
1665        let smooth_term = self
1666            .design
1667            .smooth
1668            .terms
1669            .get(self.term_index.get())
1670            .ok_or_else(|| {
1671                SmoothError::dimension_mismatch(format!(
1672                    "LatentCoord term index {} out of bounds for realized smooth design",
1673                    self.term_index
1674                ))
1675            })?;
1676        let termspec = self
1677            .spec
1678            .smooth_terms
1679            .get(self.term_index.get())
1680            .ok_or_else(|| {
1681                SmoothError::dimension_mismatch(format!(
1682                    "LatentCoord term index {} out of bounds for resolved smooth spec",
1683                    self.term_index
1684                ))
1685            })?;
1686        match (&termspec.basis, &smooth_term.metadata) {
1687            (
1688                SmoothBasisSpec::Matern { .. },
1689                BasisMetadata::Matern {
1690                    centers,
1691                    length_scale,
1692                    nu,
1693                    aniso_log_scales,
1694                    ..
1695                },
1696            ) => Ok(gam_solve::latent_cache::LatentBasisKind::Matern {
1697                centers: centers.clone(),
1698                length_scale: *length_scale,
1699                nu: *nu,
1700                aniso_log_scales: aniso_log_scales
1701                    .clone()
1702                    .unwrap_or_else(|| vec![0.0; centers.ncols()]),
1703                chunk_size: gam_terms::basis::auto_streaming_chunk_size_for_dense(
1704                    self.n_obs,
1705                    centers.nrows(),
1706                ),
1707            }),
1708            (
1709                SmoothBasisSpec::Duchon { .. },
1710                BasisMetadata::Duchon {
1711                    centers,
1712                    length_scale,
1713                    power,
1714                    nullspace_order,
1715                    aniso_log_scales,
1716                    ..
1717                },
1718            ) => Ok(gam_solve::latent_cache::LatentBasisKind::Duchon {
1719                centers: centers.clone(),
1720                length_scale: *length_scale,
1721                power: *power,
1722                nullspace_order: *nullspace_order,
1723                aniso_log_scales: aniso_log_scales
1724                    .clone()
1725                    .unwrap_or_else(|| vec![0.0; centers.ncols()]),
1726            }),
1727            (
1728                SmoothBasisSpec::Sphere { .. },
1729                BasisMetadata::Sphere {
1730                    centers,
1731                    penalty_order,
1732                    method,
1733                    ..
1734                },
1735            ) if matches!(*method, gam_terms::basis::SphereMethod::Wahba) => {
1736                Ok(gam_solve::latent_cache::LatentBasisKind::Sphere {
1737                    centers: centers.clone(),
1738                    penalty_order: *penalty_order,
1739                    chunk_size: gam_terms::basis::auto_streaming_chunk_size_for_dense(
1740                        self.n_obs,
1741                        centers.nrows(),
1742                    ),
1743                })
1744            }
1745            (
1746                SmoothBasisSpec::BSpline1D { spec, .. },
1747                BasisMetadata::BSpline1D {
1748                    knots,
1749                    periodic,
1750                    degree: meta_degree,
1751                    ..
1752                },
1753            ) => {
1754                // Issue #340: prefer the metadata-recorded effective degree
1755                // (which reflects fit-time auto-shrink) over the upstream
1756                // user-requested `spec.degree`.
1757                let effective_degree = meta_degree.unwrap_or(spec.degree);
1758                if let Some((domain_start, period, num_basis)) = periodic {
1759                    Ok(gam_solve::latent_cache::LatentBasisKind::PeriodicBspline {
1760                        domain_start: *domain_start,
1761                        period: *period,
1762                        degree: effective_degree,
1763                        num_basis: *num_basis,
1764                        chunk_size: gam_terms::basis::auto_streaming_chunk_size_for_dense(
1765                            self.n_obs, *num_basis,
1766                        ),
1767                    })
1768                } else {
1769                    let num_basis_est = knots.len().saturating_sub(effective_degree + 1);
1770                    Ok(gam_solve::latent_cache::LatentBasisKind::TensorBspline {
1771                        knots: vec![knots.clone()],
1772                        degrees: vec![effective_degree],
1773                        chunk_size: gam_terms::basis::auto_streaming_chunk_size_for_dense(
1774                            self.n_obs,
1775                            num_basis_est,
1776                        ),
1777                    })
1778                }
1779            }
1780            (
1781                SmoothBasisSpec::TensorBSpline { .. },
1782                BasisMetadata::TensorBSpline { knots, degrees, .. },
1783            ) => Ok(gam_solve::latent_cache::LatentBasisKind::TensorBspline {
1784                knots: knots.clone(),
1785                degrees: degrees.clone(),
1786                chunk_size: None,
1787            }),
1788            (
1789                SmoothBasisSpec::Pca { .. },
1790                BasisMetadata::Pca {
1791                    basis_matrix,
1792                    centered,
1793                    smooth_penalty,
1794                    center_mean,
1795                    pca_basis_path,
1796                    chunk_size,
1797                    ..
1798                },
1799            ) => {
1800                let center_mean_fingerprint = if *centered && pca_basis_path.is_none() {
1801                    let mean = center_mean.as_ref().ok_or_else(|| {
1802                        SmoothError::invalid_config(
1803                            "latent-coordinate Pca cache key requires center_mean when centered",
1804                        )
1805                    })?;
1806                    Some(gam_solve::latent_cache::pca_center_mean_fingerprint(mean))
1807                } else {
1808                    None
1809                };
1810                Ok(gam_solve::latent_cache::LatentBasisKind::Pca {
1811                    basis_matrix: basis_matrix.clone(),
1812                    centered: *centered,
1813                    center_mean_fingerprint,
1814                    smooth_penalty: *smooth_penalty,
1815                    pca_basis_path: pca_basis_path.clone(),
1816                    chunk_size: *chunk_size,
1817                })
1818            }
1819            _ => Err(SmoothError::invalid_config(
1820                "latent-coordinate design cache could not key the realized latent smooth basis"
1821                    .to_string(),
1822            )
1823            .into()),
1824        }
1825    }
1826
1827    fn ensure_theta(&mut self, theta: &Array1<f64>) -> Result<(), String> {
1828        if self
1829            .current_theta
1830            .as_ref()
1831            .is_some_and(|cached| theta_values_match(cached, theta))
1832        {
1833            return Ok(());
1834        }
1835        let latent_flat_len = self.n_obs * self.latent_dim;
1836        let direct_hyper_count = latent_coord_direct_hyper_count(&self.id_mode, self.latent_dim);
1837        let expected =
1838            self.rho_dim + latent_flat_len + self.analytic_rho_count + direct_hyper_count;
1839        if theta.len() != expected {
1840            return Err(SmoothError::dimension_mismatch(format!(
1841                "latent-coordinate theta length mismatch: got {}, expected {} (rho_dim={}, n={}, d={}, analytic_rhos={}, direct_hypers={})",
1842                theta.len(),
1843                expected,
1844                self.rho_dim,
1845                self.n_obs,
1846                self.latent_dim,
1847                self.analytic_rho_count,
1848                direct_hyper_count
1849            ))
1850            .into());
1851        }
1852        let flat = theta
1853            .slice(s![self.rho_dim..self.rho_dim + latent_flat_len])
1854            .to_owned();
1855        let latent = std::sync::Arc::new(
1856            gam_terms::latent::LatentCoordValues::from_flat_with_manifold_and_retraction_and_id(
1857                flat,
1858                self.n_obs,
1859                self.latent_dim,
1860                self.id_mode.clone(),
1861                self.manifold.clone(),
1862                self.retraction_registry.clone(),
1863                self.latent_id,
1864            ),
1865        );
1866        let latent_values_changed = self
1867            .current_latent
1868            .as_ref()
1869            .map(|cached| !latent_values_match(cached.as_flat(), latent.as_flat()))
1870            .unwrap_or(true);
1871        if latent_values_changed {
1872            self.latent_design_cache.invalidate_all();
1873            self.current_design_cache_id = None;
1874            self.design_revision = self.design_revision.wrapping_add(1);
1875        }
1876        for n in 0..self.n_obs {
1877            for axis in 0..self.latent_dim {
1878                let col = self.feature_cols[axis];
1879                self.data[[n, col]] = latent.as_flat()[n * self.latent_dim + axis];
1880            }
1881        }
1882
1883        let basis_kind = self.latent_basis_kind()?;
1884        let rebuilt_width = self.design.design.ncols();
1885        let spec = self.spec.clone();
1886        let term_index = self.term_index;
1887        let analytic_rho_count = self.analytic_rho_count;
1888        let data = self.data.view();
1889        let design_context_digest = gam_solve::latent_cache::latent_design_context_cache_digest(
1890            data,
1891            &spec,
1892            term_index,
1893            analytic_rho_count,
1894            &self.feature_cols,
1895        )
1896        .map_err(|e| e.to_string())?;
1897        let lookup = self
1898            .latent_design_cache
1899            .lookup_or_compute(latent.clone(), basis_kind, design_context_digest, || {
1900                let rebuilt = build_term_collection_design(data, &spec).map_err(|e| {
1901                    EstimationError::InvalidInput(format!(
1902                        "failed to rebuild latent-coordinate design: {e}"
1903                    ))
1904                })?;
1905                if rebuilt.design.ncols() != rebuilt_width {
1906                    crate::bail_invalid_estim!(
1907                        "latent-coordinate design topology changed: rebuilt p={}, cached p={}",
1908                        rebuilt.design.ncols(),
1909                        rebuilt_width
1910                    );
1911                }
1912                let hyper_dirs = try_build_latent_coord_hyper_dirs(
1913                    latent.clone(),
1914                    &spec,
1915                    &rebuilt,
1916                    &[term_index],
1917                    analytic_rho_count,
1918                )?
1919                .ok_or_else(|| {
1920                    EstimationError::InvalidInput(
1921                        "failed to build latent-coordinate hyper_dirs".to_string(),
1922                    )
1923                })?;
1924                Ok(gam_solve::latent_cache::ComputedLatentDesign {
1925                    design: rebuilt,
1926                    hyper_dirs,
1927                })
1928            })
1929            .map_err(|e| e.to_string())?;
1930        if lookup.cached.design.design.ncols() != self.design.design.ncols() {
1931            return Err(SmoothError::dimension_mismatch(format!(
1932                "latent-coordinate design topology changed: rebuilt p={}, cached p={}",
1933                lookup.cached.design.design.ncols(),
1934                self.design.design.ncols()
1935            ))
1936            .into());
1937        }
1938        self.design = lookup.cached.design.clone();
1939        self.current_hyper_dirs = Some(lookup.cached.hyper_dirs.clone());
1940        self.current_latent = Some(latent);
1941        self.current_theta = Some(theta.clone());
1942        self.last_cost = None;
1943        self.last_eval = None;
1944        self.last_outer_iter = None;
1945        if !latent_values_changed && self.current_design_cache_id != Some(lookup.entry_id) {
1946            self.design_revision = self.design_revision.wrapping_add(1);
1947        }
1948        self.current_design_cache_id = Some(lookup.entry_id);
1949        Ok(())
1950    }
1951
1952    fn memoized_cost(&self, theta: &Array1<f64>) -> Option<f64> {
1953        if self
1954            .current_theta
1955            .as_ref()
1956            .is_some_and(|cached| theta_values_match(cached, theta))
1957            && self.last_outer_iter
1958                == Some(gam_solve::estimate::reml::outer_eval::current_outer_iter())
1959        {
1960            self.last_eval
1961                .as_ref()
1962                .map(|cached| cached.0)
1963                .or(self.last_cost)
1964        } else {
1965            None
1966        }
1967    }
1968
1969    fn memoized_eval(
1970        &self,
1971        theta: &Array1<f64>,
1972    ) -> Option<(f64, Array1<f64>, gam_problem::HessianValue)> {
1973        if self
1974            .current_theta
1975            .as_ref()
1976            .is_some_and(|cached| theta_values_match(cached, theta))
1977            && self.last_outer_iter
1978                == Some(gam_solve::estimate::reml::outer_eval::current_outer_iter())
1979        {
1980            self.last_eval.clone()
1981        } else {
1982            None
1983        }
1984    }
1985
1986    fn store_eval(&mut self, eval: (f64, Array1<f64>, gam_problem::HessianValue)) {
1987        self.last_cost = Some(eval.0);
1988        self.last_eval = Some(eval);
1989        self.last_outer_iter = Some(gam_solve::estimate::reml::outer_eval::current_outer_iter());
1990    }
1991
1992    fn store_cost(&mut self, cost: f64) {
1993        self.last_cost = Some(cost);
1994        self.last_outer_iter = Some(gam_solve::estimate::reml::outer_eval::current_outer_iter());
1995    }
1996
1997    fn reset(&mut self) {
1998        self.current_theta = None;
1999        self.current_latent = None;
2000        self.current_hyper_dirs = None;
2001        self.current_design_cache_id = None;
2002        self.latent_design_cache.invalidate();
2003        self.last_cost = None;
2004        self.last_eval = None;
2005        self.last_outer_iter = None;
2006    }
2007}
2008
2009/// #1464: the fixed-κ profiled-REML score `V_p(κ)` for a single constant-curvature
2010/// term — pin κ on the term, fit with κ-optimisation DISABLED so only the
2011/// smoothing parameters ρ are profiled, and return the resulting REML/LAML
2012/// negative-log-evidence (the value the outer loop minimises). This is exactly
2013/// the criterion the `curvature_inference_forspec` CI oracle evaluates; factoring
2014/// it here lets the production joint-fit path reuse the SAME sign-correct profiled
2015/// criterion to pick the κ-sign basin before the joint [ρ, ψ] solve, instead of
2016/// letting the joint optimiser descend from a single κ seed into the spurious +κ
2017/// collapsed-kernel corner (the headline #1464 sign-blindness).
2018///
2019/// `pub` so a regression test can evaluate the EXACT production criterion at two
2020/// pinned κ (e.g. +κ vs −κ on a hyperbolic dataset) and settle solver-vs-criterion:
2021/// if `V_p(+κ) < V_p(−κ)` for hyperbolic data, the criterion itself prefers the
2022/// collapsed +κ corner and the bug is in the constant-curvature REML/Occam term,
2023/// not the optimiser.
2024pub fn fixed_kappa_profiled_reml_score(
2025    data: ArrayView2<'_, f64>,
2026    y: ArrayView1<'_, f64>,
2027    weights: ArrayView1<'_, f64>,
2028    offset: ArrayView1<'_, f64>,
2029    resolvedspec: &TermCollectionSpec,
2030    term_idx: usize,
2031    kappa: f64,
2032    family: LikelihoodSpec,
2033    options: &FitOptions,
2034) -> Result<f64, EstimationError> {
2035    if !kappa.is_finite() {
2036        crate::bail_invalid_estim!("fixed-κ profiled score probed a non-finite κ = {kappa}");
2037    }
2038    // Resolve the constant-curvature term's feature columns and base spec so the
2039    // criterion is probed on the production constant-curvature design.
2040    let (feature_cols, mut probe_basis) =
2041        match resolvedspec.smooth_terms.get(term_idx).map(|t| &t.basis) {
2042            Some(SmoothBasisSpec::ConstantCurvature {
2043                feature_cols, spec, ..
2044            }) => (feature_cols.clone(), spec.clone()),
2045            _ => {
2046                crate::bail_invalid_estim!(
2047                    "fixed-κ profiled score: term {term_idx} is not a constant-curvature smooth"
2048                )
2049            }
2050        };
2051    probe_basis.kappa = kappa;
2052
2053    // #1464: the curvature κ criterion the CI/flatness oracle walks (and the
2054    // `constant_curvature_profiled_reml_scores` export reports) is the HONEST
2055    // fixed-κ profiled REML of the realized constant-curvature design —
2056    // `dof·log(rss/dof) + log|H| − log|λS|₊` profiled over λ on `[1|K_κ·z]`
2057    // (`constant_curvature_honest_profiled_reml_score`). NOT the production
2058    // full-fit `reml_score`: that score heavily SMOOTHS this RKHS kernel, and under
2059    // heavy smoothing the +κ chart's geodesic-distance compression makes the
2060    // collapsed kernel a uniformly better fit of the over-smoothed target for ANY
2061    // data, so it is MONOTONE toward the +chart bound regardless of the true
2062    // curvature sign (the #1464 sign-blindness — `bug_hunt_1464_criterion_vs_solver`
2063    // shows V_p(+2) < V_p(0) < V_p(−2) on hyperbolic data with the raw score). The
2064    // honest profiled REML keeps the curvature-shape signal in the data fit, so its
2065    // argmin tracks the planted sign, and as a proper profiled-REML deviance the
2066    // CI/flatness LR thresholds stay χ²-calibrated; on constant-mean data it is
2067    // ~flat in κ, giving the flatness test correct size. Gaussian-identity is the
2068    // only family the curvature-as-estimand path serves; a weighted response, a
2069    // non-zero offset, or a non-Gaussian link routes to the production fixed-κ fit
2070    // (those configurations are not exercised by curvature inference, and the
2071    // fallback keeps their behaviour byte-identical).
2072    let is_unweighted = weights.iter().all(|&w| (w - 1.0).abs() <= 1e-12);
2073    let is_zero_offset = offset.iter().all(|&o| o.abs() <= 1e-12);
2074    if family == LikelihoodSpec::gaussian_identity() && is_unweighted && is_zero_offset {
2075        let x_term = select_columns(data, &feature_cols).map_err(EstimationError::from)?;
2076        let score = gam_terms::basis::constant_curvature_honest_profiled_reml_score(
2077            x_term.view(),
2078            y,
2079            &probe_basis,
2080        )
2081        .map_err(|e| {
2082            EstimationError::InvalidInput(format!(
2083                "fixed-κ honest profiled-REML score at κ={kappa} failed: {e}"
2084            ))
2085        })?;
2086        if !score.is_finite() {
2087            crate::bail_invalid_estim!(
2088                "fixed-κ honest profiled-REML score at κ={kappa} is non-finite"
2089            );
2090        }
2091        return Ok(score);
2092    }
2093
2094    // Fallback (weighted / offset / non-Gaussian): the production fixed-κ fit.
2095    let mut probe_spec = resolvedspec.clone();
2096    match probe_spec
2097        .smooth_terms
2098        .get_mut(term_idx)
2099        .map(|t| &mut t.basis)
2100    {
2101        Some(SmoothBasisSpec::ConstantCurvature { spec, .. }) => spec.kappa = kappa,
2102        _ => {
2103            crate::bail_invalid_estim!(
2104                "fixed-κ profiled score: term {term_idx} is not a constant-curvature smooth"
2105            )
2106        }
2107    }
2108    let fixed_kappa_options = SpatialLengthScaleOptimizationOptions {
2109        enabled: false,
2110        ..SpatialLengthScaleOptimizationOptions::default()
2111    };
2112    let fit = fit_term_collectionwith_spatial_length_scale_optimization(
2113        data,
2114        y.to_owned(),
2115        weights.to_owned(),
2116        offset.to_owned(),
2117        &probe_spec,
2118        family,
2119        options,
2120        &fixed_kappa_options,
2121    )?;
2122    let score = fit_score(&fit.fit);
2123    if !score.is_finite() {
2124        crate::bail_invalid_estim!("fixed-κ profiled fit at κ={kappa} returned a non-finite score");
2125    }
2126    Ok(score)
2127}
2128
2129/// Evaluate a Gaussian REML profile and its exact envelope derivative along one
2130/// signed-curvature direction. The smoothing parameter is selected by the
2131/// closed-form stationary-root enumerator; no lattice participates in either
2132/// the value or the derivative.
2133fn profiled_gaussian_reml_value_kappa_gradient(
2134    design: &Array2<f64>,
2135    design_kappa: &Array2<f64>,
2136    penalty: &Array2<f64>,
2137    penalty_kappa: &Array2<f64>,
2138    response: ArrayView1<'_, f64>,
2139) -> Result<(f64, f64), EstimationError> {
2140    if design.dim() != design_kappa.dim()
2141        || penalty.dim() != penalty_kappa.dim()
2142        || penalty.dim() != (design.ncols(), design.ncols())
2143        || response.len() != design.nrows()
2144    {
2145        crate::bail_invalid_estim!("constant-curvature profile value/gradient shape mismatch");
2146    }
2147
2148    let response_2d = response.insert_axis(ndarray::Axis(1));
2149    let fit = gam_solve::gaussian_reml::gaussian_reml_multi_closed_form(
2150        design.view(),
2151        response_2d.view(),
2152        penalty.view(),
2153        None,
2154        None,
2155    )?;
2156    let backward = gam_solve::gaussian_reml::gaussian_reml_multi_closed_form_backward_from_fit(
2157        design.view(),
2158        response_2d.view(),
2159        penalty.view(),
2160        None,
2161        &fit,
2162        0.0,
2163        None,
2164        None,
2165        1.0,
2166        0.0,
2167    )?;
2168    let derivative = backward
2169        .grad_x
2170        .iter()
2171        .zip(design_kappa.iter())
2172        .map(|(&adjoint, &direction)| adjoint * direction)
2173        .sum::<f64>()
2174        + backward
2175            .grad_penalty
2176            .iter()
2177            .zip(penalty_kappa.iter())
2178            .map(|(&adjoint, &direction)| adjoint * direction)
2179            .sum::<f64>();
2180    if !(fit.reml_score.is_finite() && derivative.is_finite()) {
2181        crate::bail_invalid_estim!(
2182            "constant-curvature analytic profile returned a non-finite value or derivative"
2183        );
2184    }
2185    Ok((fit.reml_score, derivative))
2186}
2187
2188/// Curvature-neutral radial reference used by the fair profile. Equal-width
2189/// Euclidean-radius bins use Sturges' data-size rule, so the reference resolution
2190/// is determined by the sample rather than a hand-picked bin count.
2191fn constant_curvature_radial_reference(
2192    data: ArrayView2<'_, f64>,
2193    y: ArrayView1<'_, f64>,
2194) -> Result<Array1<f64>, EstimationError> {
2195    if y.len() != data.nrows() || y.is_empty() {
2196        crate::bail_invalid_estim!(
2197            "constant-curvature radial reference needs one non-empty response per row"
2198        );
2199    }
2200    let radii: Array1<f64> = data.outer_iter().map(|row| row.dot(&row).sqrt()).collect();
2201    let r_max = radii.iter().copied().fold(0.0_f64, f64::max);
2202    if r_max <= f64::MIN_POSITIVE {
2203        let mean = y.sum() / y.len() as f64;
2204        return Ok(Array1::from_elem(y.len(), mean));
2205    }
2206
2207    let bin_count = (data.nrows() as f64).log2().ceil() as usize + 1;
2208    let bin_of = |radius: f64| -> usize {
2209        ((radius / r_max * bin_count as f64) as usize).min(bin_count - 1)
2210    };
2211    let mut sums = vec![0.0; bin_count];
2212    let mut counts = vec![0usize; bin_count];
2213    for (row, &radius) in radii.iter().enumerate() {
2214        let bin = bin_of(radius);
2215        sums[bin] += y[row];
2216        counts[bin] += 1;
2217    }
2218    let means: Vec<f64> = sums
2219        .into_iter()
2220        .zip(counts)
2221        .map(
2222            |(sum, count)| {
2223                if count == 0 { 0.0 } else { sum / count as f64 }
2224            },
2225        )
2226        .collect();
2227    Ok(radii.mapv(|radius| means[bin_of(radius)]))
2228}
2229
2230/// Value and exact first derivative of the continuously smoothing-profiled
2231/// curvature-fair negative log evidence. Both the observed response and its
2232/// curvature-neutral reference are profiled independently, then differenced on
2233/// the scale consumed by likelihood-ratio inference.
2234fn constant_curvature_kappa_fair_profile_value_gradient(
2235    data: ArrayView2<'_, f64>,
2236    y: ArrayView1<'_, f64>,
2237    y_ref: ArrayView1<'_, f64>,
2238    spec: &gam_terms::basis::ConstantCurvatureBasisSpec,
2239) -> Result<(f64, f64), EstimationError> {
2240    if y.len() != data.nrows() || y_ref.len() != data.nrows() {
2241        crate::bail_invalid_estim!(
2242            "constant-curvature fair profile row mismatch: data={}, response={}, reference={}",
2243            data.nrows(),
2244            y.len(),
2245            y_ref.len(),
2246        );
2247    }
2248
2249    let mut profile_spec = spec.clone();
2250    profile_spec.double_penalty = false;
2251    let basis = gam_terms::basis::build_constant_curvature_basis(data, &profile_spec)
2252        .map_err(EstimationError::from)?;
2253    let derivatives =
2254        gam_terms::basis::build_constant_curvature_basis_kappa_derivatives(data, &profile_spec)
2255            .map_err(EstimationError::from)?;
2256    if basis.active_penalties.len() != 1 || derivatives.first.penalties_derivative.len() != 1 {
2257        crate::bail_invalid_estim!(
2258            "constant-curvature fair profile expected one primary penalty; value blocks={}, derivative blocks={}",
2259            basis.active_penalties.len(),
2260            derivatives.first.penalties_derivative.len(),
2261        );
2262    }
2263
2264    let smooth_design = basis.design.to_dense();
2265    let smooth_design_kappa = &derivatives.first.design_derivative;
2266    let smooth_penalty = &basis.active_penalties[0].matrix;
2267    let smooth_penalty_kappa = &derivatives.first.penalties_derivative[0];
2268    let n = smooth_design.nrows();
2269    let p = smooth_design.ncols();
2270    if smooth_design_kappa.dim() != (n, p)
2271        || smooth_penalty.dim() != (p, p)
2272        || smooth_penalty_kappa.dim() != (p, p)
2273    {
2274        crate::bail_invalid_estim!(
2275            "constant-curvature kappa derivative bundle does not match its value basis"
2276        );
2277    }
2278
2279    let mut design = Array2::<f64>::ones((n, p + 1));
2280    design.slice_mut(s![.., 1..]).assign(&smooth_design);
2281    let mut design_kappa = Array2::<f64>::zeros((n, p + 1));
2282    design_kappa
2283        .slice_mut(s![.., 1..])
2284        .assign(smooth_design_kappa);
2285    let mut penalty = Array2::<f64>::zeros((p + 1, p + 1));
2286    penalty.slice_mut(s![1.., 1..]).assign(smooth_penalty);
2287    let mut penalty_kappa = Array2::<f64>::zeros((p + 1, p + 1));
2288    penalty_kappa
2289        .slice_mut(s![1.., 1..])
2290        .assign(smooth_penalty_kappa);
2291
2292    let (value_y, derivative_y) = profiled_gaussian_reml_value_kappa_gradient(
2293        &design,
2294        &design_kappa,
2295        &penalty,
2296        &penalty_kappa,
2297        y,
2298    )?;
2299    let (value_ref, derivative_ref) = profiled_gaussian_reml_value_kappa_gradient(
2300        &design,
2301        &design_kappa,
2302        &penalty,
2303        &penalty_kappa,
2304        y_ref,
2305    )?;
2306    Ok((value_y - value_ref, derivative_y - derivative_ref))
2307}
2308
2309struct ConstantCurvatureFairProfile<'a> {
2310    data: ArrayView2<'a, f64>,
2311    response: ArrayView1<'a, f64>,
2312    radial_reference: Array1<f64>,
2313    spec: gam_terms::basis::ConstantCurvatureBasisSpec,
2314    cache: std::cell::RefCell<std::collections::HashMap<u64, (f64, f64)>>,
2315}
2316
2317impl ConstantCurvatureFairProfile<'_> {
2318    fn evaluate(&self, kappa: f64) -> Result<(f64, f64), EstimationError> {
2319        if !kappa.is_finite() {
2320            crate::bail_invalid_estim!("constant-curvature fair profile probed a non-finite kappa");
2321        }
2322        let key = kappa.to_bits();
2323        if let Some(&cached) = self.cache.borrow().get(&key) {
2324            return Ok(cached);
2325        }
2326        let mut probe_spec = self.spec.clone();
2327        probe_spec.kappa = kappa;
2328        let sample = constant_curvature_kappa_fair_profile_value_gradient(
2329            self.data,
2330            self.response,
2331            self.radial_reference.view(),
2332            &probe_spec,
2333        )?;
2334        self.cache.borrow_mut().insert(key, sample);
2335        Ok(sample)
2336    }
2337}
2338
2339fn validate_constant_curvature_fair_profile_inputs(
2340    weights: ArrayView1<'_, f64>,
2341    offset: ArrayView1<'_, f64>,
2342    family: &LikelihoodSpec,
2343) -> Result<(), EstimationError> {
2344    if *family != LikelihoodSpec::gaussian_identity() {
2345        crate::bail_invalid_estim!(
2346            "curvature-as-an-estimand profile currently requires Gaussian identity likelihood"
2347        );
2348    }
2349    let input_tolerance = f64::EPSILON.sqrt();
2350    if weights
2351        .iter()
2352        .any(|&weight| (weight - 1.0).abs() > input_tolerance)
2353        || offset.iter().any(|&value| value.abs() > input_tolerance)
2354    {
2355        crate::bail_invalid_estim!(
2356            "curvature-as-an-estimand profile requires unit weights and zero offset"
2357        );
2358    }
2359    Ok(())
2360}
2361
2362/// Minimize the continuously smoothing-profiled curvature-fair evidence on the
2363/// chart-valid interval with the shared bounded analytic outer solver. The
2364/// curvature coordinate is the sole auxiliary coordinate, so every accepted
2365/// result has passed the solver's final box-KKT projected-gradient certificate.
2366/// No sampled point is ever returned as the estimate: samples are only line-
2367/// search probes for the continuous BFGS solve.
2368fn constant_curvature_kappa_fair_optimum(
2369    data: ArrayView2<'_, f64>,
2370    y: ArrayView1<'_, f64>,
2371    resolvedspec: &TermCollectionSpec,
2372    term_idx: usize,
2373    options: &FitOptions,
2374) -> Result<f64, EstimationError> {
2375    let (kappa_min, kappa_max) = constant_curvature_kappa_bounds(data, resolvedspec, term_idx);
2376    if !(kappa_min.is_finite() && kappa_max.is_finite() && kappa_max > kappa_min) {
2377        crate::bail_invalid_estim!(
2378            "constant-curvature term {term_idx} has invalid kappa bounds [{kappa_min}, {kappa_max}]"
2379        );
2380    }
2381    let (feature_cols, base_spec) = match resolvedspec
2382        .smooth_terms
2383        .get(term_idx)
2384        .map(|term| &term.basis)
2385    {
2386        Some(SmoothBasisSpec::ConstantCurvature {
2387            feature_cols, spec, ..
2388        }) => (feature_cols, spec.clone()),
2389        _ => {
2390            crate::bail_invalid_estim!(
2391                "constant-curvature optimum requested for non-curvature term {term_idx}"
2392            )
2393        }
2394    };
2395    let x_term = select_columns(data, feature_cols).map_err(EstimationError::from)?;
2396    let y_ref = constant_curvature_radial_reference(x_term.view(), y)?;
2397    let profile = ConstantCurvatureFairProfile {
2398        data: x_term.view(),
2399        response: y,
2400        radial_reference: y_ref,
2401        spec: base_spec,
2402        cache: std::cell::RefCell::new(std::collections::HashMap::new()),
2403    };
2404    let mut seed_config = gam_problem::SeedConfig::default();
2405    seed_config.bounds = (kappa_min, kappa_max);
2406    seed_config.max_seeds = 1;
2407    seed_config.seed_budget = 1;
2408    seed_config.risk_profile = gam_problem::SeedRiskProfile::Gaussian;
2409    seed_config.num_auxiliary_trailing = 1;
2410    seed_config.over_smoothing_probe_rho = None;
2411    let initial_kappa = profile.spec.kappa.clamp(kappa_min, kappa_max);
2412    let problem = gam_solve::rho_optimizer::OuterProblem::new(1)
2413        .with_gradient(gam_problem::Derivative::Analytic)
2414        .with_hessian(gam_problem::DeclaredHessianForm::Unavailable)
2415        .with_prefer_gradient_only(true)
2416        .with_disable_fixed_point(true)
2417        .with_fallback_policy(gam_solve::rho_optimizer::FallbackPolicy::Disabled)
2418        .with_psi_dim(1)
2419        .with_tolerance(options.tol.max(f64::EPSILON.sqrt()))
2420        .with_max_iter(options.max_iter.max(1))
2421        .with_bounds(
2422            Array1::from_vec(vec![kappa_min]),
2423            Array1::from_vec(vec![kappa_max]),
2424        )
2425        .with_initial_rho(Array1::from_vec(vec![initial_kappa]))
2426        .with_seed_config(seed_config);
2427    let mut objective = problem.build_objective(
2428        profile,
2429        |profile: &mut ConstantCurvatureFairProfile<'_>, theta: &Array1<f64>| {
2430            profile.evaluate(theta[0]).map(|(value, _)| value)
2431        },
2432        |profile: &mut ConstantCurvatureFairProfile<'_>, theta: &Array1<f64>| {
2433            let (cost, derivative) = profile.evaluate(theta[0])?;
2434            Ok(gam_problem::OuterEval {
2435                cost,
2436                gradient: Array1::from_vec(vec![derivative]),
2437                hessian: gam_problem::HessianValue::Unavailable,
2438                inner_beta_hint: None,
2439            })
2440        },
2441        None::<fn(&mut ConstantCurvatureFairProfile<'_>)>,
2442        None::<
2443            fn(
2444                &mut ConstantCurvatureFairProfile<'_>,
2445                &Array1<f64>,
2446            ) -> Result<gam_problem::EfsEval, EstimationError>,
2447        >,
2448    );
2449    let result = problem.run(
2450        &mut objective,
2451        &format!("constant-curvature fair profile term {term_idx}"),
2452    )?;
2453    if !result.converged {
2454        crate::bail_invalid_estim!(
2455            "constant-curvature fair-profile κ optimization did not converge for term {} after {} iterations (negative_log_evidence={:.6e}, final_grad_norm={})",
2456            term_idx,
2457            result.iterations,
2458            result.final_value,
2459            result.final_grad_norm_report(),
2460        );
2461    }
2462    let kappa_hat = result.rho[0];
2463    log::info!(
2464        "[spatial-kappa] continuous fair-profile optimum kappa_hat={:.6} \
2465         (negative_log_evidence={:.6e}, projected_gradient={}) for term {term_idx}",
2466        kappa_hat,
2467        result.final_value,
2468        result.final_grad_norm_report(),
2469    );
2470    Ok(kappa_hat)
2471}
2472
2473fn try_exact_joint_spatial_length_scale_optimization(
2474    data: ArrayView2<'_, f64>,
2475    y: ArrayView1<'_, f64>,
2476    weights: ArrayView1<'_, f64>,
2477    offset: ArrayView1<'_, f64>,
2478    resolvedspec: &TermCollectionSpec,
2479    best: &FittedTermCollection,
2480    family: LikelihoodSpec,
2481    options: &FitOptions,
2482    kappa_options: &SpatialLengthScaleOptimizationOptions,
2483    spatial_terms: &[usize],
2484) -> Result<Option<FittedTermCollectionWithSpec>, EstimationError> {
2485    if spatial_terms.is_empty() {
2486        return Ok(None);
2487    }
2488    // Fail loud on nonsensical κ options rather than letting them propagate
2489    // silent NaNs (e.g. inverted min/max inverts the BFGS window, negative
2490    // scales produce NaN logs). This is the first function on every outer-κ
2491    // path; downstream paths assume validated options.
2492    kappa_options
2493        .validate()
2494        .map_err(EstimationError::InvalidInput)?;
2495
2496    if try_build_spatial_log_kappa_hyper_dirs(data, resolvedspec, &best.design, spatial_terms)?
2497        .is_none()
2498    {
2499        if !constant_curvature_term_indices(resolvedspec).is_empty() {
2500            log::info!(
2501                "[#1464-trace] try_exact_joint RETURNED None (hyper_dirs unavailable); \
2502                 κ̂ comes from a NON-joint path"
2503            );
2504        }
2505        return Ok(None);
2506    }
2507    if !constant_curvature_term_indices(resolvedspec).is_empty() {
2508        log::info!(
2509            "[#1464-trace] try_exact_joint ENTERED for {} spatial term(s); CC present",
2510            spatial_terms.len()
2511        );
2512    }
2513
2514    const JOINT_RHO_BOUND: f64 = 12.0;
2515    let rho_dim = best.fit.lambdas.len();
2516
2517    // #1464: a constant-curvature `curv()` term's geodesic-exponential kernel
2518    // COLLAPSES toward the constant function as κ grows positive (sphere
2519    // distances compress), so its global REML optimum at the +κ side is a LARGE
2520    // smoothing λ — often ρ > +JOINT_RHO_BOUND. With the symmetric ±12 box the
2521    // joint [ρ,ψ] optimizer is structurally clamped into the shallow
2522    // under-smoothing basin whose spuriously-low deviance rails κ̂ to the +chart
2523    // bound for any curved data (hyperbolic truth mis-recovered as spherical).
2524    // When a constant-curvature term is present, widen ONLY the over-smoothing
2525    // (upper) ρ bound to the standard `RHO_BOUND`, leaving the lower bound at
2526    // −JOINT_RHO_BOUND so an overfit origin is never reachable — the same
2527    // asymmetric-bound rationale the standard scalar-ρ path uses for the
2528    // gam#1266 high-λ basin. Every other spatial/Matérn/Duchon/sphere joint fit
2529    // keeps the historical ±12 box byte-for-byte.
2530    let has_constant_curvature_term = !constant_curvature_term_indices(resolvedspec).is_empty();
2531    let rho_upper_bound = if has_constant_curvature_term {
2532        gam_solve::estimate::RHO_BOUND
2533    } else {
2534        JOINT_RHO_BOUND
2535    };
2536
2537    // Compute per-term dimensionality for anisotropic terms.
2538    let dims_per_term = spatial_dims_per_term(resolvedspec, spatial_terms);
2539    let use_aniso = has_aniso_terms(resolvedspec, spatial_terms);
2540
2541    // Build initial ψ values and bounds, using aniso-aware constructors
2542    // when any term has d > 1 axes. Bounds are tied to each term's center
2543    // geometry (r_min, r_max) so κ cannot saturate at an upper bound that
2544    // has no relationship to the data's distance scale.
2545    let log_kappa0 = if use_aniso {
2546        SpatialLogKappaCoords::from_length_scales_aniso(resolvedspec, spatial_terms, kappa_options)
2547    } else {
2548        SpatialLogKappaCoords::from_length_scales(resolvedspec, spatial_terms, kappa_options)
2549    };
2550    // If the user/spec did not set a length_scale, re-seed ψ at the midpoint
2551    // of the data-derived window instead of the arbitrary options fallback.
2552    let mut log_kappa0 = log_kappa0
2553        .reseed_from_data(data, resolvedspec, spatial_terms, kappa_options)
2554        .map_err(EstimationError::BasisError)?;
2555    // Constant curvature is selected once, continuously, before the baseline
2556    // fit. The full joint solve therefore profiles only nuisance ρ (and any
2557    // non-curvature spatial coordinates) at that certified κ. User-pinned and
2558    // estimated values share the same fixed-coordinate treatment, including κ=0.
2559    let mut cc_profiled_values: Vec<(usize, f64)> = Vec::new();
2560    if has_constant_curvature_term {
2561        for (slot, &term_idx) in spatial_terms.iter().enumerate() {
2562            if constant_curvature_term_spec(resolvedspec, term_idx).is_none() {
2563                continue;
2564            }
2565            let kappa = get_constant_curvature_kappa(resolvedspec, term_idx)
2566                .expect("constant-curvature term exposes its kappa");
2567            log_kappa0.set_scalar_slot(slot, kappa);
2568            cc_profiled_values.push((slot, kappa));
2569        }
2570    }
2571    let log_kappa_lower = if use_aniso {
2572        SpatialLogKappaCoords::lower_bounds_aniso_from_data(
2573            data,
2574            resolvedspec,
2575            spatial_terms,
2576            &dims_per_term,
2577            kappa_options,
2578        )
2579    } else {
2580        SpatialLogKappaCoords::lower_bounds_from_data(
2581            data,
2582            resolvedspec,
2583            spatial_terms,
2584            kappa_options,
2585        )
2586    }
2587    .map_err(EstimationError::BasisError)?;
2588    let log_kappa_upper = if use_aniso {
2589        SpatialLogKappaCoords::upper_bounds_aniso_from_data(
2590            data,
2591            resolvedspec,
2592            spatial_terms,
2593            &dims_per_term,
2594            kappa_options,
2595        )
2596    } else {
2597        SpatialLogKappaCoords::upper_bounds_from_data(
2598            data,
2599            resolvedspec,
2600            spatial_terms,
2601            kappa_options,
2602        )
2603    }
2604    .map_err(EstimationError::BasisError)?;
2605    let mut log_kappa_lower = log_kappa_lower;
2606    let mut log_kappa_upper = log_kappa_upper;
2607    for &(slot, kappa) in &cc_profiled_values {
2608        log_kappa_lower.set_scalar_slot(slot, kappa);
2609        log_kappa_upper.set_scalar_slot(slot, kappa);
2610        log::info!("[spatial-kappa] slot {slot}: profiling rho at certified kappa={kappa}");
2611    }
2612    // Project seed onto data-derived bounds; spec.length_scale is a hint,
2613    // not a hard constraint. BFGS requires theta0 ∈ [lower, upper].
2614    let log_kappa0 = log_kappa0.clamp_to_bounds(&log_kappa_lower, &log_kappa_upper);
2615    let setup = ExactJointHyperSetup::new(
2616        best.fit.lambdas.mapv(f64::ln),
2617        Array1::<f64>::from_elem(rho_dim, -JOINT_RHO_BOUND),
2618        Array1::<f64>::from_elem(rho_dim, rho_upper_bound),
2619        log_kappa0,
2620        log_kappa_lower,
2621        log_kappa_upper,
2622    );
2623
2624    let theta0 = setup.theta0();
2625    let lower = setup.lower();
2626    let upper = setup.upper();
2627
2628    // ───────────────────────────────────────────────────────────────────────
2629    //  Both coordinate kinds drive the SAME exact joint optimizer
2630    //  (`run_exact_joint_spatial_optimization`): the unified REML evaluator with
2631    //  ext_coords for joint [ρ, ψ] optimization, with analytic gradient +
2632    //  Hessian flowing through the
2633    //  AnisoBasisPsiDerivatives / SpatialPsiDerivative → DirectionalHyperParam →
2634    //  HyperCoord pipeline for Newton/BFGS quadratic convergence. The only
2635    //  difference is the coordinate kind: anisotropic carries one ψ per axis per
2636    //  term, isotropic one log-κ per term. `outer_strategy` handles the
2637    //  centralized degradation path when the analytic Hessian is unavailable.
2638    // ───────────────────────────────────────────────────────────────────────
2639    let kind = if use_aniso {
2640        SpatialHyperKind::Anisotropic
2641    } else {
2642        SpatialHyperKind::Isotropic
2643    };
2644    let (theta_star, joint_final_value, kappa_timing) = run_exact_joint_spatial_optimization(
2645        kind,
2646        data,
2647        y,
2648        weights,
2649        offset,
2650        resolvedspec,
2651        &best.design,
2652        family.clone(),
2653        options,
2654        spatial_terms,
2655        &dims_per_term,
2656        &theta0,
2657        &lower,
2658        &upper,
2659        rho_dim,
2660        kappa_options,
2661    )?;
2662
2663    let baseline_score = fit_score(&best.fit);
2664
2665    // Compare the joint optimizer's certified cost (final_value at theta*)
2666    // against the baseline. Tolerance ≥ options.tol because both endpoints
2667    // are outer-BFGS approximations accurate to options.tol; a tighter
2668    // gate would reject true improvements due to floating-point noise.
2669    let accept_tol = options.tol.max(1e-8 * baseline_score.abs()).max(1e-12);
2670    if joint_final_value > baseline_score + accept_tol {
2671        return Err(EstimationError::RemlOptimizationFailed(format!(
2672            "exact joint spatial optimization failed its objective-monotonicity certificate: \
2673             initial={baseline_score:.6e}, final={joint_final_value:.6e}, \
2674             acceptance_tolerance={accept_tol:.3e}, theta_checkpoint={:?}",
2675            theta_star.to_vec(),
2676        )));
2677    }
2678
2679    let selected_lambdas = Array1::from_vec(
2680        gam_problem::checked_exp_log_strengths(
2681            theta_star.slice(s![..rho_dim]).iter().copied(),
2682        )
2683        .map_err(|error| {
2684            EstimationError::InvalidInput(format!(
2685                "selected joint spatial smoothing coordinate is outside the canonical log-strength domain: {error}"
2686            ))
2687        })?,
2688    );
2689    let log_kappa_star =
2690        SpatialLogKappaCoords::from_theta_tail_with_dims(&theta_star, rho_dim, dims_per_term);
2691    // #1464 diagnostic (ban-clean): the joint solver's CONVERGED ψ-tail κ for each
2692    // CC term — the value BEFORE any spec write-back / freeze / readback. If this
2693    // is negative for the hyperbolic dataset but `get_constant_curvature_kappa`
2694    // later returns +1.08, the railing is a POST-SOLVE clamp/readback, not the
2695    // optimiser. If this is itself +1.08, the joint solver railed past the pin.
2696    if has_constant_curvature_term {
2697        let star = log_kappa_star.as_array();
2698        let dims = log_kappa_star.dims_per_term();
2699        for (slot, &term_idx) in spatial_terms.iter().enumerate() {
2700            if constant_curvature_term_spec(resolvedspec, term_idx).is_some() {
2701                let off: usize = dims[..slot].iter().sum();
2702                log::info!(
2703                    "[#1464-trace] term {term_idx}: joint solver CONVERGED ψ-tail κ = {} \
2704                     (this is the optimised candidate; joint_final_value={joint_final_value})",
2705                    star[off]
2706                );
2707            }
2708        }
2709    }
2710    let optimized_spec = log_kappa_star.apply_tospec(resolvedspec, spatial_terms)?;
2711    let optimized = fit_term_collection_forspecwith_heuristic_lambdas(
2712        data,
2713        y,
2714        weights,
2715        offset,
2716        &optimized_spec,
2717        selected_lambdas.as_slice(),
2718        family.clone(),
2719        options,
2720    )?;
2721
2722    // Stamp reml_score with joint_final_value so downstream consumers see a
2723    // score consistent with the gate decision; the refit serves as a
2724    // β/inference harvester at the certified (ρ*, ψ*).
2725    let mut fit = optimized.fit;
2726    fit.reml_score = joint_final_value;
2727    let optimized_result = FittedTermCollectionWithSpec {
2728        fit,
2729        design: optimized.design,
2730        resolvedspec: optimized_spec,
2731        adaptive_diagnostics: optimized.adaptive_diagnostics,
2732        kappa_timing: Some(kappa_timing),
2733    };
2734
2735    Ok(Some(optimized_result))
2736}
2737
2738/// Coordinate kind for the exact joint spatial hyperparameter optimizer.
2739///
2740/// Anisotropic and isotropic spatial terms drive the *same* joint `[ρ, ψ]`
2741/// optimizer: identical outer-Hessian policy, identical
2742/// `ExternalJointHyperEvaluator` wiring, identical convergence processing, and
2743/// an identical `eval_full / eval_efs / eval_cost`
2744/// inner loop that routes ψ through `try_build_spatial_log_kappa_hyper_dirs`.
2745/// The coordinate *kind* distinguishes per-axis log scales (ψ_a) from one
2746/// log-κ per term and selects diagnostic labels. It also tells the startup
2747/// policy when an isotropic Matérn point has already won the explicit certified
2748/// endpoint comparison, in which case that point owns the sole joint start.
2749#[derive(Clone, Copy, PartialEq, Eq, Debug)]
2750enum SpatialHyperKind {
2751    Anisotropic,
2752    Isotropic,
2753}
2754
2755impl SpatialHyperKind {
2756    /// Stable diagnostic prefix used in every `log::*` line and as the
2757    /// `ExternalJointHyperEvaluator` / cost-only label root.
2758    fn label(self) -> &'static str {
2759        match self {
2760            SpatialHyperKind::Anisotropic => "spatial-aniso-joint",
2761            SpatialHyperKind::Isotropic => "spatial-iso-joint",
2762        }
2763    }
2764
2765    /// Human-readable adjective for error strings ("anisotropic" / "isotropic").
2766    fn adjective(self) -> &'static str {
2767        match self {
2768            SpatialHyperKind::Anisotropic => "anisotropic",
2769            SpatialHyperKind::Isotropic => "isotropic",
2770        }
2771    }
2772
2773    /// Name of the directional coordinate being optimized ("psi" / "kappa"),
2774    /// used only in hyper-direction construction error messages.
2775    fn coord_name(self) -> &'static str {
2776        match self {
2777            SpatialHyperKind::Anisotropic => "psi",
2778            SpatialHyperKind::Isotropic => "kappa",
2779        }
2780    }
2781}
2782
2783/// Shared context for the exact joint spatial optimizer's closures. Holds the
2784/// realized-design cache and the joint REML evaluator, plus the coordinate
2785/// `kind` whose only effect is the diagnostic label routed into the cost-only
2786/// evaluation path. The `eval_full / eval_efs / eval_cost` methods are the
2787/// single source of truth for both anisotropic and isotropic spatial terms.
2788struct SpatialFrozenGlmInputs {
2789    y: Array1<f64>,
2790    weights: Array1<f64>,
2791    offset: Array1<f64>,
2792    family: LikelihoodSpec,
2793}
2794
2795/// True when the frozen-weight GLM ψ-tensor (#1111 / #1033 mechanism (c)) is a
2796/// faithful first-Fisher-step provider for this family.
2797///
2798/// The mechanism freezes the working weight `w = w(η_warm)` and working response
2799/// `z = z(η_warm)` once per outer ψ-sweep, so it is exact for ANY family whose
2800/// per-iteration PIRLS reduces to a Gaussian working model with a SINGLE
2801/// canonical Fisher weight at a FIXED dispersion — i.e. the one-parameter
2802/// exponential families Binomial, Poisson, Gamma, and Negative-Binomial (the
2803/// θ-fixed running-seed weight `W = μθ/(θ+μ)` is a clean per-row Fisher weight).
2804/// These are precisely the "Poisson/Binomial/etc" families the issue names.
2805///
2806/// Tweedie and Beta jointly estimate an extra dispersion parameter that moves
2807/// the working weight outside the frozen snapshot, so the frozen-W stand-in is
2808/// not faithful for them and they keep the exact per-trial PIRLS rebuild.
2809/// Gaussian-identity is served by the (exact, converged) `PsiGramTensor` lane,
2810/// and Royston-Parmar is the survival path, neither of which routes here.
2811fn frozen_glm_tensor_eligible_family(family: &LikelihoodSpec) -> bool {
2812    !family.is_gaussian_identity()
2813        && matches!(
2814            &family.response,
2815            ResponseFamily::Binomial
2816                | ResponseFamily::Poisson
2817                | ResponseFamily::Gamma
2818                | ResponseFamily::NegativeBinomial { .. }
2819        )
2820}
2821
2822struct SpatialJointContext<'d> {
2823    data: ArrayView2<'d, f64>,
2824    rho_dim: usize,
2825    kind: SpatialHyperKind,
2826    cache: SingleBlockExactJointDesignCache<'d>,
2827    evaluator: gam_solve::estimate::ExternalJointHyperEvaluator<'d>,
2828    frozen_glm_inputs: Option<SpatialFrozenGlmInputs>,
2829    frozen_glm_psi_bounds: Option<(f64, f64)>,
2830    frozen_glm_tensor: Option<gam_solve::glm_sufficient_lane::FrozenWeightGramTensor>,
2831    frozen_glm_tensor_attempted: bool,
2832    /// #1033: memo of the frozen-W trial Fisher weights keyed on the warm β that
2833    /// produced them. `stage_frozen_glm_trial_statistics` runs on EVERY κ trial
2834    /// (every cost / gradient probe), and the only β-dependent quantity it needs
2835    /// is the current Fisher weight vector `W(η)` (η = Xβ + offset) for the
2836    /// drift check and the n-free gradient soundness gate. Computing `W` is an
2837    /// O(n·p) GEMV + O(n) family evaluation; β only changes when the inner solve
2838    /// re-converges (after an accepted outer step), so recomputing it on every
2839    /// same-β probe was a redundant per-trial n-touch. Cache `(β, W)` and reuse
2840    /// `W` whenever β is unchanged — the GEMV runs once per distinct β, i.e.
2841    /// O(outer steps), not O(trials). `None` until the first compute / when no
2842    /// frozen-W inputs are installed.
2843    frozen_glm_weight_memo: Option<(Array1<f64>, Array1<f64>)>,
2844}
2845
2846#[derive(Clone, Copy, Debug, Default)]
2847struct NfreeSkipGateStatus {
2848    shape: bool,
2849    value: bool,
2850    gradient: bool,
2851    penalty: bool,
2852    revision: bool,
2853    second_order: bool,
2854}
2855
2856impl NfreeSkipGateStatus {
2857    fn would_skip(self, require_gradient: bool) -> bool {
2858        self.shape
2859            && self.value
2860            && (!require_gradient || self.gradient)
2861            && self.penalty
2862            && self.revision
2863            && !self.second_order
2864    }
2865}
2866
2867fn nfree_skip_gate_status_from_parts(
2868    shape: bool,
2869    covers_value: bool,
2870    covers_skip: bool,
2871    covers_gradient: bool,
2872    penalty: bool,
2873    revision: bool,
2874    allow_second_order: bool,
2875    require_gradient: bool,
2876) -> NfreeSkipGateStatus {
2877    NfreeSkipGateStatus {
2878        shape,
2879        // A value-only cost probe consumes only the Chebyshev Gram value; it
2880        // does not expose a beta/row-space object, so the #1264 reduced-basis
2881        // skip witness is not part of the value soundness certificate. Requiring
2882        // `covers_skip` here forces harmless cost probes across basis-rotation
2883        // seams onto `reset_surface`, reintroducing an O(n) pass into the κ
2884        // trial loop. Gradient probes still require the skip witness because
2885        // they return a stationary beta/gradient in the frozen reduced basis.
2886        value: shape && covers_value && (!require_gradient || covers_skip),
2887        gradient: shape && (!require_gradient || covers_gradient),
2888        penalty,
2889        revision,
2890        second_order: allow_second_order,
2891    }
2892}
2893
2894impl<'d> SpatialJointContext<'d> {
2895    fn nfree_skip_gate_status(
2896        &self,
2897        theta: &Array1<f64>,
2898        allow_second_order: bool,
2899        require_gradient: bool,
2900    ) -> NfreeSkipGateStatus {
2901        let shape = theta.len() == self.rho_dim + 1;
2902        let (covers_value, covers_skip, covers_gradient) = if shape {
2903            let psi = theta[self.rho_dim];
2904            (
2905                self.evaluator.psi_gram_tensor_covers(psi),
2906                self.evaluator.psi_gram_tensor_covers_skip(psi),
2907                self.evaluator.psi_gram_tensor_covers_gradient(psi),
2908            )
2909        } else {
2910            (false, false, false)
2911        };
2912        nfree_skip_gate_status_from_parts(
2913            shape,
2914            covers_value,
2915            covers_skip,
2916            covers_gradient,
2917            self.evaluator.supports_nfree_penalty_rekey(),
2918            self.evaluator.nfree_fast_path_revision().is_some(),
2919            allow_second_order,
2920            require_gradient,
2921        )
2922    }
2923
2924    fn frozen_glm_working_state(
2925        &self,
2926        beta: &Array1<f64>,
2927    ) -> Result<Option<(Array1<f64>, Array1<f64>)>, EstimationError> {
2928        let Some(inputs) = self.frozen_glm_inputs.as_ref() else {
2929            return Ok(None);
2930        };
2931        if beta.len() != self.cache.design().design.ncols() {
2932            return Ok(None);
2933        }
2934        let mut eta = self.cache.design().design.matrixvectormultiply(beta);
2935        if eta.len() != inputs.offset.len() {
2936            crate::bail_invalid_estim!(
2937                "frozen GLM tensor warm-state row mismatch: eta={}, offset={}",
2938                eta.len(),
2939                inputs.offset.len()
2940            );
2941        }
2942        eta += &inputs.offset;
2943        let obs = evaluate_standard_familyobservations(
2944            inputs.family.clone(),
2945            None,
2946            None,
2947            None,
2948            &inputs.y,
2949            &inputs.weights,
2950            &eta,
2951        )?;
2952        let mut working_response = obs.eta.clone();
2953        for i in 0..working_response.len() {
2954            let wi = obs.fisherweight[i].max(1e-12);
2955            working_response[i] += obs.score[i] / wi;
2956        }
2957        Ok(Some((obs.fisherweight, working_response)))
2958    }
2959
2960    /// #1033: the trial Fisher weight vector `W(η)` for `beta`, memoized on
2961    /// `beta`. `stage_frozen_glm_trial_statistics` consults `W` on EVERY κ trial
2962    /// (drift check + n-free gradient soundness gate) but `W` is a deterministic
2963    /// function of β (η = Xβ + offset), and β only changes when the inner solve
2964    /// re-converges — many cost / gradient probes share one β. Recompute the
2965    /// O(n·p) working state only when β differs from the memoized key; otherwise
2966    /// return the cached weights. Returns `None` exactly when
2967    /// `frozen_glm_working_state` does (no frozen-W inputs / β shape mismatch).
2968    fn frozen_glm_trial_weights(
2969        &mut self,
2970        beta: &Array1<f64>,
2971    ) -> Result<Option<Array1<f64>>, EstimationError> {
2972        if let Some((memo_beta, memo_w)) = self.frozen_glm_weight_memo.as_ref()
2973            && memo_beta.len() == beta.len()
2974            && memo_beta
2975                .iter()
2976                .zip(beta.iter())
2977                .all(|(a, b)| a.to_bits() == b.to_bits())
2978        {
2979            return Ok(Some(memo_w.clone()));
2980        }
2981        match self.frozen_glm_working_state(beta)? {
2982            Some((current_w, _)) => {
2983                self.frozen_glm_weight_memo = Some((beta.clone(), current_w.clone()));
2984                Ok(Some(current_w))
2985            }
2986            None => Ok(None),
2987        }
2988    }
2989
2990    fn ensure_frozen_glm_tensor(
2991        &mut self,
2992        theta: &Array1<f64>,
2993        warm_beta: Option<&Array1<f64>>,
2994    ) -> Result<(), EstimationError> {
2995        if self.frozen_glm_tensor.is_some() || self.frozen_glm_tensor_attempted {
2996            return Ok(());
2997        }
2998        let Some((psi_lo, psi_hi)) = self.frozen_glm_psi_bounds else {
2999            return Ok(());
3000        };
3001        if theta.len() != self.rho_dim + 1 {
3002            self.frozen_glm_tensor_attempted = true;
3003            return Ok(());
3004        }
3005        let Some(beta) = warm_beta else {
3006            return Ok(());
3007        };
3008        let Some((frozen_w, working_z)) = self.frozen_glm_working_state(beta)? else {
3009            self.frozen_glm_tensor_attempted = true;
3010            return Ok(());
3011        };
3012        let theta_probe_base = theta.clone();
3013        let rho_dim = self.rho_dim;
3014        // Build through the evaluator so the frozen-W Gram is assembled in the
3015        // SAME conditioned `x_fit` column frame the inner PIRLS solve uses
3016        // (the evaluator owns the ψ-invariant parametric conditioning). Disjoint
3017        // mutable borrows of `cache` (in the realizer) and `evaluator` (the
3018        // build host) — both fields of `self` — exactly as the Gaussian
3019        // `build_and_set_psi_gram_tensor` site does.
3020        let Self {
3021            cache, evaluator, ..
3022        } = self;
3023        let tensor = evaluator.build_frozen_glm_gram_tensor(
3024            |psi| {
3025                let mut theta_probe = theta_probe_base.clone();
3026                theta_probe[rho_dim] = psi;
3027                cache.ensure_theta(&theta_probe)?;
3028                Ok(cache.design().design.clone())
3029            },
3030            frozen_w.view(),
3031            working_z.view(),
3032            psi_lo,
3033            psi_hi,
3034        );
3035        self.cache
3036            .ensure_theta(theta)
3037            .map_err(EstimationError::InvalidInput)?;
3038        self.frozen_glm_tensor_attempted = true;
3039        if let Some(tensor) = tensor {
3040            self.frozen_glm_tensor = Some(tensor);
3041            log::info!(
3042                "[STAGE] {} certified frozen-W GLM ψ tensor over [{psi_lo:.3}, {psi_hi:.3}]",
3043                self.kind.label(),
3044            );
3045        } else {
3046            log::info!(
3047                "[STAGE] {} frozen-W GLM ψ tensor did not certify over [{psi_lo:.3}, {psi_hi:.3}]",
3048                self.kind.label(),
3049            );
3050        }
3051        Ok(())
3052    }
3053
3054    fn stage_frozen_glm_trial_statistics(
3055        &mut self,
3056        theta: &Array1<f64>,
3057        warm_beta: Option<&Array1<f64>>,
3058        allow_gradient: bool,
3059    ) -> Result<(), EstimationError> {
3060        let kind = self.kind;
3061        let mut staged_gram: Option<Array2<f64>> = None;
3062        let mut staged_deriv: Option<(Array2<f64>, Array1<f64>)> = None;
3063        if theta.len() == self.rho_dim + 1 {
3064            let psi = theta[self.rho_dim];
3065            // Compute the β-memoized trial Fisher weights up front (mutable
3066            // self borrow) so the immutable `self.frozen_glm_tensor` borrow
3067            // below does not alias it. `frozen_glm_trial_weights` recomputes the
3068            // O(n·p) working state only on a β change, so a same-β probe pays
3069            // nothing here (#1033). Only proceed when a tensor is installed and
3070            // covers this ψ — otherwise skip the weight compute entirely.
3071            let tensor_covers = self
3072                .frozen_glm_tensor
3073                .as_ref()
3074                .is_some_and(|t| t.contains(psi));
3075            let current_w = if tensor_covers {
3076                match warm_beta {
3077                    Some(beta) => self.frozen_glm_trial_weights(beta)?,
3078                    None => None,
3079                }
3080            } else {
3081                None
3082            };
3083            if let (Some(tensor), Some(current_w)) =
3084                (self.frozen_glm_tensor.as_ref(), current_w.as_ref())
3085            {
3086                const FROZEN_GLM_WEIGHT_DRIFT_RTOL: f64 = 1e-3;
3087                if tensor.weight_drift_within(current_w.view(), FROZEN_GLM_WEIGHT_DRIFT_RTOL) {
3088                    staged_gram = Some(tensor.gram_at(psi));
3089                    log::debug!(
3090                        "[STAGE] {} trial at psi={psi:.6}: serving frozen-W GLM \
3091                         first-Fisher-step XᵀWX n-free (weight drift within tol)",
3092                        kind.label(),
3093                    );
3094                }
3095                if allow_gradient
3096                    && tensor.contains_for_gradient(psi)
3097                    && let Some((dgram_dpsi, drhs_dpsi)) =
3098                        tensor.gradient_pair_if_sound(psi, current_w.view())
3099                {
3100                    staged_deriv = Some((dgram_dpsi, drhs_dpsi));
3101                    log::debug!(
3102                        "[STAGE] {} trial at psi={psi:.6}: serving frozen-W GLM \
3103                         ψ-gradient (∂G/∂ψ, ∂b/∂ψ) n-free (gradient weight drift within \
3104                         tight tol); B_j stays exact",
3105                        kind.label(),
3106                    );
3107                }
3108            }
3109        }
3110        self.evaluator.stage_glm_first_step_gram(staged_gram);
3111        self.evaluator.stage_glm_psi_gram_deriv(staged_deriv);
3112        Ok(())
3113    }
3114
3115    /// Full evaluation on the current realized design + hyper_dirs.
3116    fn eval_full(
3117        &mut self,
3118        theta: &Array1<f64>,
3119        order: gam_solve::rho_optimizer::OuterEvalOrder,
3120        analytic_outer_hessian_available: bool,
3121    ) -> Result<(f64, Array1<f64>, gam_problem::HessianValue), EstimationError> {
3122        use gam_solve::rho_optimizer::OuterEvalOrder;
3123        let allow_second_order = matches!(order, OuterEvalOrder::ValueGradientHessian)
3124            && analytic_outer_hessian_available;
3125        if let Some(eval) = self.cache.memoized_eval(theta) {
3126            let cached_satisfies_order = !allow_second_order || eval.2.is_analytic();
3127            if cached_satisfies_order {
3128                return Ok(eval);
3129            }
3130        }
3131        let kind = self.kind;
3132        // #1033: the per-trial n×k design re-realization (`ensure_theta` →
3133        // `apply_log_kappa`) plus the downstream n-row reconditioning
3134        // (`reset_surface`) are the LAST n-passes in the certified κ loop. They
3135        // are redundant on the Gaussian-identity certified path: the inner
3136        // Gaussian PLS reads its `XᵀWX(ψ)/XᵀW(y−offset)(ψ)` entirely from the
3137        // ψ-keyed `GaussianFixedCache` the certified tensor installs (zero row
3138        // access), and the ψ-gradient HyperCoord is served from the k-space
3139        // `(∂G/∂ψ, ∂b/∂ψ)` tensor derivatives — never the n×k ∂X/∂ψ slab. So when
3140        //   (a) this is the single design-moving ψ coordinate (`rho_dim + 1`),
3141        //   (b) the certified ψ-Gram tensor covers ψ for BOTH the value lane
3142        //       (`psi_gram_tensor_covers`) AND the gradient window
3143        //       (`psi_gram_tensor_covers_gradient`) — so neither channel reads
3144        //       the realized rows,
3145        //   (c) this eval is gradient-only (`!allow_second_order`) — the exact
3146        //       outer-Hessian `B_j` path DOES read the slab, so a Hessian trial
3147        //       must keep a faithful (freshly realized) design, and
3148        //   (d) the evaluator has a pinned canonical slow-path revision — i.e.
3149        //       a prior slow-path eval already built a faithful reference surface,
3150        //       which `prepare_eval_state` will reuse while re-installing the
3151        //       ψ-keyed cache,
3152        // we SKIP `ensure_theta`. The realizer revision then does not advance, so
3153        // `prepare_eval_state` takes its design-revision fast path by receiving
3154        // that pinned revision back: it skips `reset_surface` + the n×k
3155        // `apply_to_design`, keeps the reference surface, and re-keys the
3156        // `GaussianFixedCache` to this ψ. The hyper_dirs built below are a pure
3157        // function of (data, frozen spec, column layout) — ψ-invariant — so they
3158        // are bit-identical whether or not the design was re-realized, and the
3159        // tensor branch never reads their n×k slab anyway. Net: criterion +
3160        // gradient + inner solve come from k-space statistics only, with no
3161        // per-trial O(n·k) pass.
3162        //
3163        // When ANY gate clause fails (non-Gaussian, off-window, off the gradient
3164        // sub-window, a Hessian eval, or no pinned canonical surface yet) we
3165        // realize the design as before so the slow path rebuilds a faithful
3166        // surface — the existing exact lane runs unchanged.
3167        let nfree_fast_path_revision = self.evaluator.nfree_fast_path_revision();
3168        let skip_design_realization = !allow_second_order && theta.len() == self.rho_dim + 1 && {
3169            let psi = theta[self.rho_dim];
3170            self.evaluator.psi_gram_tensor_covers(psi)
3171                    // #1033 gradient coverage: the skip serves the ψ-gradient n-free
3172                    // only where the analytic Chebyshev derivative is CERTIFIED.
3173                    // The kappa sufficient-statistic outer loop is routed here only
3174                    // when the certified gradient window spans the entire optimizer
3175                    // bounds, so a measured trial cannot pay an edge streamed
3176                    // ∂X/∂ψ pass after the initial priming eval.
3177                    && self.evaluator.psi_gram_tensor_covers_gradient(psi)
3178                    // #1264 (RESTORED) reduced-basis-rotation soundness precondition.
3179                    // The Gaussian inner penalized solve `(QsᵀGQs+S)β=b` runs in the
3180                    // CONDITIONED reduced basis. On the near-singular production
3181                    // Duchon Gram (κ(G)≈9.5e14) that basis ROTATES with ψ, and the
3182                    // skip installs the Chebyshev-interpolated `gram_at(ψ)` (≤1e-10
3183                    // vs streamed exact). When the trial-ψ basis differs from the
3184                    // reference surface's, the κ-amplified round-off moves β̂ by
3185                    // ~1.7e-5 — 17× the issue's 1e-6 bar — EVEN at a ψ the n-free
3186                    // VALUE window admits (cluster: β̂rel=1.749e-5 at ψ=2.803). The
3187                    // "stale-penalty-not-stale-basis" theory that dropped this gate
3188                    // was empirically refuted. So the skip is β̂-sound ONLY where the
3189                    // gauge-invariant range projector is unchanged vs the pinning ψ:
3190                    // `reduced_basis_equal(psi_ref, psi)`. Value coverage is NOT
3191                    // sufficient. This forces the exact O(n) `reset_surface` fallback
3192                    // across a basis rotation — correctness over n-independence
3193                    // (#1033 is frontier-blocked on rotating Duchon geometry).
3194                    && self.evaluator.psi_gram_tensor_covers_skip(psi)
3195                    // #1033 penalty lane: ψ moves S(ψ) too, and the skip leaves
3196                    // `reset_surface` un-run; only skip when the penalty can be
3197                    // rebuilt EXACTLY and n-free on the fast path, else the inner
3198                    // solve would pair XᵀWX(ψ_new) with the stale S(ψ_old).
3199                    && self.evaluator.supports_nfree_penalty_rekey()
3200                    && nfree_fast_path_revision.is_some()
3201        };
3202        // #1868: the #1033 n-free design-realization skip is armed above. A prior
3203        // debug override (`TEMP-SKIPOFF-1122`) hard-forced `skip_design_realization`
3204        // to `false` here to test whether the n-free ψ-Gram Chebyshev interpolant
3205        // was the source of the #1122 H-side FD-vs-analytic gap. That override was
3206        // never removed, so every in-window κ `eval_full` trial fell through to the
3207        // O(n) `ensure_theta` → `apply_log_kappa` + `reset_surface` lane — the O(n)
3208        // per-callback regression #1868 reports. The skip is already gated on
3209        // `!allow_second_order`, so it never fires on the H (Hessian) trials the
3210        // #1122 diagnostic was probing; the override only ever suppressed the
3211        // n-free gradient/value lane. Removing it routes the gradient eval through
3212        // the k-space `GaussianFixedCache` + ψ-derivative tensor as intended.
3213        if skip_design_realization {
3214            log::debug!(
3215                "[STAGE] {} eval_full at psi={:.6}: skipping n×k design re-realization \
3216                 + reconditioning — criterion/gradient/inner-solve served n-free from \
3217                 the certified ψ-gram tensor (GaussianFixedCache + k-space ψ-derivatives)",
3218                kind.label(),
3219                theta[self.rho_dim],
3220            );
3221        } else {
3222            self.cache
3223                .ensure_theta(theta)
3224                .map_err(EstimationError::InvalidInput)?;
3225        }
3226        let warm_beta = self.evaluator.current_beta();
3227        self.ensure_frozen_glm_tensor(theta, warm_beta.as_ref())?;
3228        // #1033 / #1111: stage the GLM frozen-W first-step Gram and conditioned
3229        // ψ-gradient whenever the certified frozen-weight tensor covers this
3230        // trial's ψ. The provider applies its drift guards, so misses clear the
3231        // staged slots and the exact streamed path runs.
3232        //
3233        // Stage through a shared helper because cost-only line-search probes use
3234        // the same first-Fisher-step Gram; they simply pass `allow_gradient=false`.
3235        self.stage_frozen_glm_trial_statistics(theta, warm_beta.as_ref(), !allow_second_order)?;
3236        // #1033: on the certified Gaussian skip path the value and ψ-gradient
3237        // are both served by k-space tensor statistics, so the row-wise X_ψ slab
3238        // is dead. Build only the exact n-free S_ψ components from frozen
3239        // geometry and attach a zero-storage design derivative placeholder.
3240        // Edge-gradient/Hessian/non-certified trials keep the exact row-wise
3241        // builder, because those lanes genuinely consume X_ψ.
3242        let hyper_dirs = if skip_design_realization {
3243            self.cache.nfree_tensor_gradient_hyper_dirs(theta)?
3244        } else {
3245            self.cache.hyper_dirs_for_current_design(self.data, kind)?
3246        };
3247
3248        let design_revision = if skip_design_realization {
3249            nfree_fast_path_revision
3250        } else {
3251            Some(self.cache.design_revision())
3252        };
3253        // #1033 penalty lane: stage the EXACT n-free `S(ψ)` for this trial so the
3254        // evaluator's design-revision fast path can re-key the kept reference
3255        // surface without `reset_surface`. Built from the FROZEN basis geometry
3256        // (centers + identifiability transform + operator collocation points) at
3257        // the trial length-scale — no data rows — so it is valid even on the
3258        // design-realization skip path (where the design was not re-realized). The
3259        // caller (holding `cache`) computes it and hands the owned result to the
3260        // evaluator, sidestepping a `&mut cache` borrow alias. On the slow path
3261        // the evaluator ignores + clears the staged value (it rebuilds S from the
3262        // realized design). A build error here clears the stage; if the skip
3263        // already fired (fast path), the evaluator then hard-errors rather than
3264        // pairing a stale S — the safe outcome, since a rebuild from frozen
3265        // geometry should never fail in practice.
3266        if self.evaluator.supports_nfree_penalty_rekey() {
3267            match self.cache.canonical_penalties_at(theta) {
3268                Ok(penalty) => self.evaluator.stage_fast_path_penalty(Some(penalty)),
3269                Err(e) => {
3270                    log::warn!(
3271                        "[STAGE] {} eval_full at psi={:.6}: exact n-free S(ψ) rebuild failed \
3272                         ({e}); clearing stage (eval falls to slow path)",
3273                        kind.label(),
3274                        theta[self.rho_dim],
3275                    );
3276                    self.evaluator.stage_fast_path_penalty(None);
3277                }
3278            }
3279        }
3280        // Warm-start PIRLS from the previous outer step's converged β. This is
3281        // especially impactful for GLM families (Poisson, NB, Binomial) that
3282        // cannot use the Gaussian Gram tensor n-free shortcut: without the warm
3283        // β every outer step cold-solves a full PIRLS from β=0, paying the full
3284        // O(n·p²) cost × PIRLS-iters × outer-iters budget. With the warm β the
3285        // inner solve typically converges in 1-2 Newton steps instead of 4-8.
3286        let eval = evaluate_joint_reml_outer_eval_at_theta(
3287            &mut self.evaluator,
3288            self.cache.design(),
3289            theta,
3290            self.rho_dim,
3291            hyper_dirs,
3292            warm_beta.as_ref().map(|b: &Array1<f64>| b.view()),
3293            if allow_second_order {
3294                order
3295            } else {
3296                OuterEvalOrder::ValueAndGradient
3297            },
3298            design_revision,
3299        );
3300        if let Ok(ref value) = eval {
3301            self.cache.store_eval_at(theta, value.clone());
3302        }
3303        eval
3304    }
3305
3306    fn eval_efs(&mut self, theta: &Array1<f64>) -> Result<gam_problem::EfsEval, EstimationError> {
3307        self.cache
3308            .ensure_theta(theta)
3309            .map_err(EstimationError::InvalidInput)?;
3310        let kind = self.kind;
3311        let hyper_dirs = try_build_spatial_log_kappa_hyper_dirs(
3312            self.data,
3313            self.cache.spec(),
3314            self.cache.design(),
3315            &self.cache.spatial_terms,
3316        )?
3317        .ok_or_else(|| {
3318            EstimationError::InvalidInput(format!(
3319                "failed to build {} hyper_dirs for exact-joint EFS",
3320                kind.adjective(),
3321            ))
3322        })?;
3323        let design_revision = Some(self.cache.design_revision());
3324        let warm_beta = self.evaluator.current_beta();
3325        evaluate_joint_reml_efs_at_theta(
3326            &mut self.evaluator,
3327            self.cache.design(),
3328            theta,
3329            self.rho_dim,
3330            hyper_dirs,
3331            warm_beta.as_ref().map(|b: &Array1<f64>| b.view()),
3332            design_revision,
3333        )
3334    }
3335
3336    /// Cost-only evaluation. BFGS line-search probes route through the
3337    /// evaluator's true value-only path so they neither construct
3338    /// `try_build_spatial_log_kappa_hyper_dirs` nor assemble a gradient that
3339    /// the line search will discard. Split-borrow on `self.cache` +
3340    /// `self.evaluator` matches the pattern already used by `eval_full`.
3341    fn eval_cost(&mut self, theta: &Array1<f64>) -> f64 {
3342        if let Some(cost) = self.cache.memoized_cost(theta) {
3343            return cost;
3344        }
3345        // #1029: a BFGS line-search VALUE probe. It converges the inner PIRLS to
3346        // the SAME tolerance the accepted-point full eval uses (NOT a capped
3347        // surrogate — a cap returns ∞ for a feasible point and re-imports the
3348        // #787/#808 outer stall), so probe and incumbent values live in ONE
3349        // refinement regime (measure-consistent Armijo). It is cheaper only
3350        // because it skips the gradient / hyper-dir assembly. Time the inner
3351        // cost-only solve and report it alongside the trial-θ distance from the
3352        // last evaluated point so this convergence-critical regression class is
3353        // visible in the STAGE trace (the spatial REML lane has no PROGRESS-
3354        // EXTENDED refine multiplier — that knob is SAE-only — so there is no
3355        // extended polish to strip from a probe here).
3356        //
3357        // Capture the previous evaluated θ BEFORE `ensure_theta` overwrites it,
3358        // so the logged distance reflects the backtracking step rather than 0.
3359        let probe_start = std::time::Instant::now();
3360        let psi_distance = self
3361            .cache
3362            .current_theta
3363            .as_ref()
3364            .filter(|reference| reference.len() == theta.len())
3365            .map(|reference| {
3366                reference
3367                    .iter()
3368                    .zip(theta.iter())
3369                    .map(|(a, b)| (a - b) * (a - b))
3370                    .sum::<f64>()
3371                    .sqrt()
3372            })
3373            .unwrap_or(f64::NAN);
3374        // #1033: a VALUE-only line-search probe needs only the certified ψ-Gram
3375        // tensor's value lane (`XᵀWX(ψ)/XᵀW(y−offset)(ψ)`), which the inner
3376        // Gaussian PLS reads n-free from the ψ-keyed `GaussianFixedCache`. So when
3377        // the single design-moving ψ is covered for the VALUE lane and the
3378        // evaluator has a pinned canonical slow-path revision, skip the n×k
3379        // design re-realization: `evaluate_cost_only` receives that pinned
3380        // revision, takes its `prepare_eval_state_cost_only` fast path (which
3381        // skips `reset_surface` + the n×k `apply_to_design` and re-keys the cache
3382        // to this probe's ψ), and the probe cost comes from k-space statistics
3383        // only. Line-search probes are the bulk of the κ-loop per-trial work, so
3384        // this is the dominant n-flat lever. Any miss (non-Gaussian, off-window,
3385        // missing penalty re-key support, or no pinned surface yet) realizes the
3386        // design and runs the exact streamed probe unchanged.
3387        let nfree_fast_path_revision = self.evaluator.nfree_fast_path_revision();
3388        let skip_value_realization = theta.len() == self.rho_dim + 1 && {
3389            let psi = theta[self.rho_dim];
3390            self.evaluator.psi_gram_tensor_covers(psi)
3391                    // #1868: a VALUE-only line-search probe does NOT need the
3392                    // #1264 `reduced_basis_equal` (`covers_skip`) soundness gate the
3393                    // ACCEPTED gradient eval (`eval_full`, still gated) requires. That
3394                    // gate exists because the design-realization skip freezes the
3395                    // conditioned reduced basis at the pinning ψ, and on the near-
3396                    // singular Duchon Gram (κ(G)≈9.5e14) a ψ-rotation makes the
3397                    // frozen basis interpolate β̂ with a κ-amplified round-off of
3398                    // β̂rel≈1.7e-5 — which matters for the RETURNED coefficients/
3399                    // gradient. A cost probe returns only the scalar REML criterion
3400                    // for the line search, and that criterion is STATIONARY in β̂ at
3401                    // the inner minimizer (envelope theorem): a β̂ perturbation δβ
3402                    // moves the data-fit+penalty term by O(δβ²) and leaves the
3403                    // `log|H|` term (built from the EXACT tensor Gram G(ψ), not β̂)
3404                    // untouched, so the RELATIVE cost error is ~δβ² ≈ 3e-10 — orders
3405                    // below the line search's 1e-5 Armijo tolerance. So the probe
3406                    // cannot be mis-ranked, and the converged κ/β̂ (pinned by the
3407                    // covers_skip-gated `eval_full` at accepted iterates) is
3408                    // unchanged. Gating the probe on `covers_skip` instead forced the
3409                    // O(n) `reset_surface` lane for every line-search step that
3410                    // overshoots the (n-drifting) reduced-basis-stable band — the
3411                    // #1868 per-callback reset climb: the band's rotation rate dP/dψ
3412                    // grows with n (sample-std standardization), so more probes fall
3413                    // just past PSI_GRAM_SKIP_PROJ_ATOL as n grows, defeating the
3414                    // n-independence the tensor lane was built for. The evaluator's
3415                    // own value-probe fast path (`prepare_eval_state_cost_only`) is
3416                    // gated on VALUE coverage exactly for this reason; aligning the
3417                    // driver here lets the probe cost come from the n-free k-space
3418                    // Gram/penalty statistics across the rotation.
3419                    //
3420                    // #1033 penalty lane: the value-probe fast path also skips
3421                    // `reset_surface`, so the probe must be able to re-key S(ψ)
3422                    // EXACTLY and n-free; otherwise its cost would use the stale
3423                    // S(ψ_old) and mis-rank the line search.
3424                    && self.evaluator.supports_nfree_penalty_rekey()
3425                    && nfree_fast_path_revision.is_some()
3426        };
3427        if theta.len() == self.rho_dim + 1
3428            && self.evaluator.has_psi_gram_tensor()
3429            && !self.evaluator.psi_gram_tensor_covers(theta[self.rho_dim])
3430        {
3431            self.cache.store_cost_at(theta, f64::INFINITY);
3432            return f64::INFINITY;
3433        }
3434        if !skip_value_realization && self.cache.ensure_theta(theta).is_err() {
3435            return f64::INFINITY;
3436        }
3437        // #1033 penalty lane: stage the EXACT n-free `S(ψ)` for this probe's ψ so
3438        // the cost-only fast path re-keys the kept surface without `reset_surface`
3439        // (built from frozen geometry — valid even when the design was not
3440        // re-realized). The slow path clears it. A rebuild failure clears the
3441        // stage; the evaluator then takes the slow path or hard-errors (safe).
3442        if self.evaluator.supports_nfree_penalty_rekey() {
3443            match self.cache.canonical_penalties_at(theta) {
3444                Ok(penalty) => self.evaluator.stage_fast_path_penalty(Some(penalty)),
3445                Err(_) => self.evaluator.stage_fast_path_penalty(None),
3446            }
3447        }
3448        let warm_beta = self.evaluator.current_beta();
3449        if let Err(err) = self.ensure_frozen_glm_tensor(theta, warm_beta.as_ref()) {
3450            log::warn!(
3451                "[STAGE] {} value-probe at psi={:.6}: frozen-W GLM tensor setup failed ({err}); \
3452                 falling back to exact streamed Gram",
3453                self.kind.label(),
3454                if theta.len() > self.rho_dim {
3455                    theta[self.rho_dim]
3456                } else {
3457                    f64::NAN
3458                },
3459            );
3460            self.evaluator.stage_glm_first_step_gram(None);
3461            self.evaluator.stage_glm_psi_gram_deriv(None);
3462        } else if let Err(err) =
3463            self.stage_frozen_glm_trial_statistics(theta, warm_beta.as_ref(), false)
3464        {
3465            log::warn!(
3466                "[STAGE] {} value-probe at psi={:.6}: frozen-W GLM staging failed ({err}); \
3467                 falling back to exact streamed Gram",
3468                self.kind.label(),
3469                if theta.len() > self.rho_dim {
3470                    theta[self.rho_dim]
3471                } else {
3472                    f64::NAN
3473                },
3474            );
3475            self.evaluator.stage_glm_first_step_gram(None);
3476            self.evaluator.stage_glm_psi_gram_deriv(None);
3477        }
3478        let design_revision = if skip_value_realization {
3479            nfree_fast_path_revision
3480        } else {
3481            Some(self.cache.design_revision())
3482        };
3483        let cost_label = self.kind.label();
3484        let result = {
3485            let design = self.cache.design();
3486            self.evaluator.evaluate_cost_only(
3487                &design.design,
3488                &design.penalties,
3489                &design.nullspace_dims,
3490                design.linear_constraints.clone(),
3491                theta,
3492                self.rho_dim,
3493                warm_beta.as_ref().map(|b: &Array1<f64>| b.view()),
3494                cost_label,
3495                design_revision,
3496            )
3497        };
3498        match result {
3499            Ok(cost) => {
3500                log::debug!(
3501                    "[STAGE] {cost_label} value-probe (order=Value): elapsed={:.3}s \
3502                     cost={cost:.6e} trial_theta_distance={psi_distance:.3e}",
3503                    probe_start.elapsed().as_secs_f64(),
3504                );
3505                self.cache.store_cost_at(theta, cost);
3506                cost
3507            }
3508            Err(_) => f64::INFINITY,
3509        }
3510    }
3511
3512    fn reset(&mut self) {
3513        self.cache.current_theta = None;
3514        self.cache.last_eval_theta = None;
3515        self.cache.last_cost = None;
3516        self.cache.last_eval = None;
3517    }
3518}
3519
3520/// Exact joint `[ρ, ψ]` optimization for spatial terms using analytic
3521/// derivatives through the unified REML evaluator. This is the single shared
3522/// engine for both the anisotropic and isotropic coordinate kinds (selected by
3523/// `kind`).
3524///
3525/// At each outer iteration, the frozen term topology is reused and only the
3526/// spatial realized blocks affected by the current ψ are refreshed before the
3527/// unified evaluator returns cost + gradient + Hessian for the full
3528/// θ = [ρ, ψ] vector. The ψ derivatives flow through:
3529///
3530///   `AnisoBasisPsiDerivatives` / `SpatialPsiDerivative` → `DirectionalHyperParam`
3531///     → `build_tau_unified_objects` → `HyperCoord` ext_coords → unified evaluator
3532///
3533/// This gives Newton/BFGS quadratic convergence on the length-scale /
3534/// anisotropy parameters while jointly optimizing the smoothing parameters.
3535///
3536/// The ψ coordinates are parameterized as unconstrained log-scales. For the
3537/// anisotropic kind the decomposition into isotropic scale (ψ̄ = mean(ψ_a)) and
3538/// anisotropy (η_a = ψ_a − ψ̄, with Ση_a = 0) happens only on writeback via
3539/// `SpatialLogKappaCoords::apply_tospec`; the all-ones direction in ψ-space is
3540/// NOT a gauge direction — it controls the identifiable isotropic scale
3541/// κ = exp(ψ̄). The isotropic kind carries one log-κ coordinate per term. In
3542/// neither case is a sum-to-zero constraint enforced during optimization.
3543fn kphase_log_norms(theta: &Array1<f64>, rho_dim: usize) -> (f64, f64) {
3544    let theta_norm = theta.iter().map(|v| v * v).sum::<f64>().sqrt();
3545    let log_kappa_norm = theta
3546        .iter()
3547        .skip(rho_dim)
3548        .map(|v| v * v)
3549        .sum::<f64>()
3550        .sqrt();
3551    (theta_norm, log_kappa_norm)
3552}
3553
3554fn run_exact_joint_spatial_optimization(
3555    kind: SpatialHyperKind,
3556    data: ArrayView2<'_, f64>,
3557    y: ArrayView1<'_, f64>,
3558    weights: ArrayView1<'_, f64>,
3559    offset: ArrayView1<'_, f64>,
3560    resolvedspec: &TermCollectionSpec,
3561    baseline_design: &TermCollectionDesign,
3562    family: LikelihoodSpec,
3563    options: &FitOptions,
3564    spatial_terms: &[usize],
3565    dims_per_term: &[usize],
3566    theta0: &Array1<f64>,
3567    lower: &Array1<f64>,
3568    upper: &Array1<f64>,
3569    rho_dim: usize,
3570    kappa_options: &SpatialLengthScaleOptimizationOptions,
3571) -> Result<(Array1<f64>, f64, SpatialLengthScaleOptimizationTiming), EstimationError> {
3572    let label = kind.label();
3573    let effective_offset = baseline_design
3574        .compose_offset(offset, "spatial joint fit")
3575        .map_err(EstimationError::BasisError)?;
3576    let offset = effective_offset.view();
3577    // Use bounds and design metadata for validation.
3578    assert!(
3579        lower.len() == theta0.len() && upper.len() == theta0.len(),
3580        "spatial hyperparameter bounds must match theta length: lower_len={}, upper_len={}, theta_len={}",
3581        lower.len(),
3582        upper.len(),
3583        theta0.len()
3584    );
3585    assert!(
3586        baseline_design.smooth.terms.len() >= spatial_terms.len(),
3587        "baseline design must have at least one smooth term per spatial term: baseline_terms={}, spatial_terms={}",
3588        baseline_design.smooth.terms.len(),
3589        spatial_terms.len()
3590    );
3591    use gam_problem::{DeclaredHessianForm, Derivative, OuterEval};
3592    use gam_solve::rho_optimizer::OuterEvalOrder;
3593
3594    let theta_dim = theta0.len();
3595    // Directional-coordinate dimension: psi-per-axis (anisotropic) or
3596    // kappa-per-term (isotropic). The numerics below are identical either way.
3597    let coord_dim = theta_dim - rho_dim;
3598    // Capability is declared solely from derivative coverage, not from
3599    // problem size. The unified REML evaluator now exposes exact matrix-free
3600    // outer Hessian operators for the costly third/fourth-derivative
3601    // contractions used by spatial ψ coordinates; its internal
3602    // `(n, p, K)` work model chooses `HessianValue::Operator` at large-scale
3603    // scale and the dense analytic matrix only below that crossover. Keeping
3604    // `Derivative::Analytic` here preserves ARC / trust-region-CG second-order
3605    // optimization for `n > 50_000` and `coord_dim > 30` instead of forcing the
3606    // obsolete HybridEFS compatibility path.
3607    let analytic_outer_hessian_available =
3608        exact_joint_spatial_outer_hessian_available(&family, baseline_design);
3609    if !analytic_outer_hessian_available {
3610        log::info!(
3611            "[{label}] analytic outer Hessian unavailable for family/design; routing without second-order geometry (coord_dim={coord_dim})"
3612        );
3613    }
3614    // Cost-aware second-order routing, mirroring the n-block path's
3615    // work-budget policy: past the pair budget gradient-only quasi-Newton
3616    // converges to the same optimum strictly cheaper per eval; below it,
3617    // exact second-order keeps the ARC/TR-CG geometry. The budget's
3618    // derivation is owned by `EXACT_JOINT_SECOND_ORDER_THETA_CAP`.
3619    let mut prefer_gradient_only = theta_dim > EXACT_JOINT_SECOND_ORDER_THETA_CAP;
3620    if prefer_gradient_only {
3621        log::info!(
3622            "[{label}] joint θ-dim {theta_dim} exceeds the exact pair-Hessian budget \
3623             ({EXACT_JOINT_SECOND_ORDER_THETA_CAP}); routing gradient-only quasi-Newton"
3624        );
3625    }
3626    // #1033: set when the n-free Gaussian ψ-lane arms below. It must SUPPRESS the
3627    // declared analytic outer Hessian (force `Unavailable`), not merely prefer
3628    // gradient-only: the planner keeps the second-order ARC solver whenever an
3629    // analytic Hessian is declared `Either`, even under `prefer_gradient_only`
3630    // (see `plan_prefer_gradient_only_does_not_hide_analytic_hessian`). A
3631    // `ValueGradientHessian` eval forces the O(n) design re-realization because
3632    // the outer Hessian curvature slab `B_j` is irreducibly n-dependent, so only
3633    // routing to a gradient-only solver (BFGS) keeps every in-window κ-trial on
3634    // the n-free `ValueAndGradient` skip.
3635    let mut suppress_outer_hessian_for_nfree = false;
3636
3637    log::trace!(
3638        "[{}] starting analytic optimization: rho_dim={}, coord_dim={}, dims_per_term={:?}",
3639        label,
3640        rho_dim,
3641        coord_dim,
3642        dims_per_term,
3643    );
3644
3645    let mut ctx = SpatialJointContext {
3646        data,
3647        rho_dim,
3648        kind,
3649        cache: SingleBlockExactJointDesignCache::new_with_policy(
3650            data,
3651            resolvedspec.clone(),
3652            baseline_design.clone(),
3653            spatial_terms.to_vec(),
3654            rho_dim,
3655            dims_per_term.to_vec(),
3656            &options.resource_policy,
3657        )
3658        .map_err(EstimationError::InvalidInput)?,
3659        evaluator: gam_solve::estimate::ExternalJointHyperEvaluator::new(
3660            y,
3661            weights,
3662            &baseline_design.design,
3663            offset,
3664            &baseline_design.penalties,
3665            &external_opts_for_design(&family, baseline_design, options),
3666            label,
3667        )?,
3668        frozen_glm_inputs: if coord_dim == 1 && frozen_glm_tensor_eligible_family(&family) {
3669            Some(SpatialFrozenGlmInputs {
3670                y: y.to_owned(),
3671                weights: weights.to_owned(),
3672                offset: offset.to_owned(),
3673                family: family.clone(),
3674            })
3675        } else {
3676            None
3677        },
3678        frozen_glm_psi_bounds: if coord_dim == 1 && frozen_glm_tensor_eligible_family(&family) {
3679            Some((lower[rho_dim], upper[rho_dim]))
3680        } else {
3681            None
3682        },
3683        frozen_glm_tensor: None,
3684        frozen_glm_tensor_attempted: false,
3685        frozen_glm_weight_memo: None,
3686    };
3687
3688    // #1033b: single isotropic design-moving coordinate on a Gaussian-identity
3689    // fit — build the certified Chebyshev-in-ψ Gram tensor ONCE over the
3690    // optimizer's ψ window and hand it to the evaluator. Every in-window trial
3691    // then receives its Gaussian sufficient statistics (XᵀWX(ψ), XᵀW(y−offset),
3692    // (y−offset)ᵀW(y−offset)) assembled n-free instead of paying the per-trial
3693    // O(n·p²) Gram re-stream after the design rebuild. The realizer closure
3694    // returns the RAW realized design; the evaluator threads it through its
3695    // own (fixed, ψ-invariant) parametric column conditioning so the tensor
3696    // lives in the same frame as the streamed Gram. Certification failure,
3697    // off-window trials, or any other ineligibility silently keep the exact
3698    // streamed path (same numbers, the tensor is certified to
3699    // PSI_GRAM_SPOT_RTOL against the exact rebuild).
3700    // #1033 (rank-stable κ-floor): set to the lowest ψ at which the certified
3701    // tensor's conditioned Gram holds maximal numerical rank. Below it the
3702    // reduced basis collapses/rotates and the design-realization skip is SOUNDLY
3703    // refused (→ O(n) reset_surface); the κ window floor `ln(2/r_max)` lands
3704    // inside that degenerate sliver and DRIFTS with n through the sample-std
3705    // standardization, so n=2000's line search re-enters the slow lane while
3706    // n=1000's does not. Lifting the optimizer's lower bound to this n-FREE
3707    // (k-space) floor keeps every in-window trial on the fast path for all n,
3708    // and only excludes over-smoothed length scales the `2/r_max` geometry floor
3709    // already meant to exclude (the κ-optimum lives well above it).
3710    let mut psi_rank_stable_floor: Option<f64> = None;
3711    // #1033 (rank-stable κ-ceiling): symmetric twin of the floor. The conditioned
3712    // Gram is rank-deficient at the HIGH window edge too (the longest-frequency
3713    // radial mode goes collinear), so a line-search overshoot above the maximal-
3714    // rank band soundly refuses the design-realization skip → O(n) reset_surface,
3715    // and the deficient pinning ψ it records makes the NEXT in-band trial reset a
3716    // second time. Clamping the optimizer's UPPER bound to this n-free k-space
3717    // ceiling keeps every trial inside the band. The κ-optimum lives well inside
3718    // it, so the clamp only excludes over-fit (too-short) length scales.
3719    let mut psi_rank_stable_ceiling: Option<f64> = None;
3720    let nfree_penalty_capable =
3721        coord_dim == 1 && family.is_gaussian_identity() && ctx.cache.supports_nfree_penalty_rekey();
3722    if nfree_penalty_capable {
3723        let psi_lo = lower[rho_dim];
3724        let psi_hi = upper[rho_dim];
3725        let z = Array1::from_iter(y.iter().zip(offset.iter()).map(|(yi, oi)| yi - oi));
3726        let theta_probe_base = theta0.clone();
3727        // Disjoint mutable borrows of `cache` (in the realizer) and
3728        // `evaluator` (the build target) — both fields of `ctx`.
3729        let SpatialJointContext {
3730            cache, evaluator, ..
3731        } = &mut ctx;
3732        let attached = evaluator.build_and_set_psi_gram_tensor(
3733            |psi| {
3734                let mut theta_probe = theta_probe_base.clone();
3735                theta_probe[rho_dim] = psi;
3736                cache.ensure_theta(&theta_probe)?;
3737                Ok(cache.design().design.clone())
3738            },
3739            weights,
3740            z.view(),
3741            psi_lo,
3742            psi_hi,
3743        );
3744        if attached {
3745            log::info!(
3746                "[{label}] certified ψ-gram tensor over [{psi_lo:.3}, {psi_hi:.3}]: \
3747                 in-window trials assemble Gaussian sufficient statistics n-free"
3748            );
3749            // #1033: read the n-free rank-stable κ-floor off the k-space tensor.
3750            // Only lift INTO the window (never below psi_lo, never above the seed
3751            // ψ — the seed is the geometric-mean midpoint and is well clear of the
3752            // degenerate band), so the optimizer never starts outside its bounds.
3753            let psi_anchor = theta0[rho_dim];
3754            psi_rank_stable_floor = evaluator
3755                .psi_gram_rank_stable_floor(psi_anchor)
3756                .filter(|&f| f.is_finite() && f > psi_lo && f < psi_anchor);
3757            log::info!(
3758                "[KAPPA-PHASE-FLOOR] n_rows={} psi_lo={psi_lo:.6} psi_anchor={psi_anchor:.6} \
3759                 rank_stable_floor={:?} lifted={}",
3760                data.nrows(),
3761                evaluator.psi_gram_rank_stable_floor(psi_anchor),
3762                psi_rank_stable_floor.is_some(),
3763            );
3764            if let Some(floor) = psi_rank_stable_floor {
3765                log::info!(
3766                    "[{label}] rank-stable κ-floor ψ_floor={floor:.6} > window floor \
3767                     ψ_lo={psi_lo:.6}: lifting the optimizer lower bound to keep every \
3768                     in-window trial on the n-free design-realization skip (#1033). The \
3769                     conditioned Gram is rank-deficient below ψ_floor (longest-length-scale \
3770                     radial mode collapses into the nullspace), where the skip is soundly \
3771                     refused; that band drifts with n via the sample-std standardization, \
3772                     so this n-free k-space floor is the n-independent fix."
3773                );
3774            }
3775            // #1033: read the n-free rank-stable κ-CEILING (symmetric twin of the
3776            // floor). Only clamp INTO the window (strictly below psi_hi, strictly
3777            // above the seed ψ — the seed is the geometric-mean midpoint, well
3778            // inside the maximal-rank band), so the optimizer never starts outside
3779            // its bounds. This is the fix for the n=16000 fast-ladder resets: the
3780            // line search overshot to ψ≈1.0 (rank 11→10 at the high edge), tripping
3781            // two O(n) reset_surface calls; clamping the upper bound keeps the
3782            // search inside the band where the n-free skip stays sound.
3783            psi_rank_stable_ceiling = evaluator
3784                .psi_gram_rank_stable_ceiling(psi_anchor)
3785                .filter(|&c| c.is_finite() && c < psi_hi && c > psi_anchor);
3786            log::info!(
3787                "[KAPPA-PHASE-CEIL] n_rows={} psi_hi={psi_hi:.6} psi_anchor={psi_anchor:.6} \
3788                 rank_stable_ceiling={:?} clamped={}",
3789                data.nrows(),
3790                evaluator.psi_gram_rank_stable_ceiling(psi_anchor),
3791                psi_rank_stable_ceiling.is_some(),
3792            );
3793            if let Some(ceiling) = psi_rank_stable_ceiling {
3794                log::info!(
3795                    "[{label}] rank-stable κ-ceiling ψ_ceil={ceiling:.6} < window ceiling \
3796                     ψ_hi={psi_hi:.6}: clamping the optimizer upper bound to keep every \
3797                     in-window trial on the n-free design-realization skip (#1033). The \
3798                     conditioned Gram is rank-deficient above ψ_ceil (longest-frequency \
3799                     radial mode goes collinear), where the skip is soundly refused; a \
3800                     line-search overshoot there trips the O(n) reset_surface lane (and the \
3801                     deficient pinning ψ it records resets the next in-band trial too)."
3802                );
3803            }
3804            let gradient_covers_full_window = evaluator.psi_gram_tensor_covers_gradient(psi_lo)
3805                && evaluator.psi_gram_tensor_covers_gradient(psi_hi);
3806            if gradient_covers_full_window {
3807                log::info!(
3808                    "[{label}] certified ψ-gram tensor gradient lane covers the full \
3809                     optimizer window [{psi_lo:.3}, {psi_hi:.3}]"
3810                );
3811            } else {
3812                log::info!(
3813                    "[{label}] ψ-gram tensor value lane certified, but the gradient lane \
3814                     does not cover the full optimizer window [{psi_lo:.3}, {psi_hi:.3}]; \
3815                     keeping exact streamed kappa routing"
3816                );
3817            }
3818            // #1033 penalty lane: ψ also moves the penalty `S(ψ)` (the
3819            // Duchon/ThinPlate Hilbert scale is an analytic function of the
3820            // length-scale, built from the FROZEN basis CENTERS — not the data
3821            // rows). The design-revision fast path that the Gram tensor enables
3822            // SKIPS `reset_surface`, the only place the canonical penalty surface
3823            // is rebuilt; without re-keying, the inner solve would pair
3824            // `XᵀWX(ψ_new)` with the stale `S(ψ_old)` and converge to the wrong
3825            // β̂ / κ-optimum. Rather than interpolate `S(ψ)`, the fast path rebuilds
3826            // it EXACTLY and n-free per trial from the frozen geometry via
3827            // `cache.canonical_penalties_at(theta)` (the SAME
3828            // `canonicalize_penalty_specs` pipeline the slow `reset_surface` runs).
3829            // Here we only DECLARE the capability to the evaluator; the per-trial
3830            // staging happens in `eval_full` / `eval_cost`. The skip is enabled
3831            // exactly when the single spatial term's frozen metadata
3832            // (Duchon/ThinPlate) admits the exact rebuild. Matérn deliberately
3833            // does not enter this block: mixing tensor value probes with exact
3834            // streamed gradients/Hessians changed its selected κ enough to miss
3835            // the truth-recovery quality gate, so Matérn stays on one exact
3836            // streamed objective for value, gradient, and Hessian.
3837            evaluator.set_supports_nfree_penalty_rekey(true);
3838            log::info!(
3839                "[{label}] exact n-free ψ-penalty re-key enabled over [{psi_lo:.3}, \
3840                 {psi_hi:.3}]: in-window fast-path trials rebuild S(ψ) n-free from frozen \
3841                 geometry (no reset_surface)"
3842            );
3843        } else {
3844            log::info!(
3845                "[{label}] ψ-gram tensor did not certify over [{psi_lo:.3}, {psi_hi:.3}]; \
3846                 keeping the exact per-trial path"
3847            );
3848        }
3849        // #1033 (n-independent outer loop): with the n-free Gaussian lane fully
3850        // armed (Gram tensor attached + exact n-free penalty re-key), the design-
3851        // realization skip serves the criterion AND the ψ-gradient `(a_j, g_j)`
3852        // n-free for every in-window trial — but ONLY a `ValueAndGradient` eval
3853        // takes that skip. A `ValueGradientHessian` eval sets `allow_second_order`,
3854        // which forces `ensure_theta` → `reset_surface` (the O(n) design re-
3855        // realization) because the outer Hessian curvature `B_j` is the exact
3856        // n-dependent slab. So second-order outer steps are the LAST O(n) per-trial
3857        // cost in the κ search, and they make the outer loop scale with n. Route
3858        // gradient-only here: the spatial length-scale objective is smooth and the
3859        // budget policy already establishes that gradient-only quasi-Newton
3860        // converges to the same optimum strictly cheaper per eval past the pair-
3861        // Hessian budget — and with the tensor, the realized Hessian is the only
3862        // remaining expensive operation, so the same argument applies for ANY n
3863        // once the lane is armed. This keeps every in-window κ-trial on the n-free
3864        // `ValueAndGradient` skip, delivering the n-independent outer loop. The
3865        // exact second-order geometry is preserved whenever the lane is NOT armed
3866        // for gradient-only routing (non-Gaussian, multi-term, Matérn, or an
3867        // uncertified window), where it still pays O(n) per Hessian but keeps the
3868        // quality-sensitive exact second-order path.
3869        if attached
3870            && evaluator.psi_gram_tensor_covers_gradient(psi_lo)
3871            && evaluator.psi_gram_tensor_covers_gradient(psi_hi)
3872            && evaluator.supports_nfree_penalty_rekey()
3873            && cache.supports_nfree_gradient_only_routing()
3874        {
3875            suppress_outer_hessian_for_nfree = true;
3876            prefer_gradient_only = true;
3877            log::info!(
3878                "[{label}] n-free Gaussian ψ-lane armed; suppressing the analytic outer \
3879                 Hessian and routing gradient-only (BFGS) so the κ outer loop never realizes \
3880                the O(n) second-order slab — n-independent outer loop (#1033)"
3881            );
3882        }
3883    } else if coord_dim == 1 && family.is_gaussian_identity() {
3884        log::info!(
3885            "[{label}] exact n-free ψ-penalty re-key unavailable; skipping ψ-gram tensor \
3886             attachment so value, gradient, and Hessian remain on the same exact streamed \
3887             objective"
3888        );
3889    }
3890
3891    let kphase_prime_order =
3892        if analytic_outer_hessian_available && !suppress_outer_hessian_for_nfree {
3893            OuterEvalOrder::ValueGradientHessian
3894        } else {
3895            OuterEvalOrder::ValueAndGradient
3896        };
3897    let kphase_prime_start = std::time::Instant::now();
3898    drop(ctx.eval_full(theta0, kphase_prime_order, analytic_outer_hessian_available)?);
3899    log::info!(
3900        "[KAPPA-PHASE-PRIME] n_rows={} order={:?} elapsed_s={:.4} slow_path_resets_total={} design_revision={}",
3901        data.nrows(),
3902        kphase_prime_order,
3903        kphase_prime_start.elapsed().as_secs_f64(),
3904        ctx.evaluator.slow_path_reset_count(),
3905        ctx.cache.design_revision(),
3906    );
3907
3908    let kphase_cost_calls = std::cell::Cell::new(0usize);
3909    let kphase_eval_calls = std::cell::Cell::new(0usize);
3910    let kphase_efs_calls = std::cell::Cell::new(0usize);
3911    let kphase_cost_total_s = std::cell::Cell::new(0.0);
3912    let kphase_eval_total_s = std::cell::Cell::new(0.0);
3913    let kphase_efs_total_s = std::cell::Cell::new(0.0);
3914    let kphase_nfree_miss_shape = std::cell::Cell::new(0u64);
3915    let kphase_nfree_miss_value = std::cell::Cell::new(0u64);
3916    let kphase_nfree_miss_gradient = std::cell::Cell::new(0u64);
3917    let kphase_nfree_miss_penalty = std::cell::Cell::new(0u64);
3918    let kphase_nfree_miss_revision = std::cell::Cell::new(0u64);
3919    let kphase_nfree_miss_second_order = std::cell::Cell::new(0u64);
3920    let kphase_nfree_miss_other = std::cell::Cell::new(0u64);
3921    let kphase_optim_start = std::time::Instant::now();
3922    let kphase_log_kappa_dim = coord_dim;
3923    let kphase_slow_resets_start = ctx.evaluator.slow_path_reset_count();
3924    let kphase_design_revision_start = ctx.cache.design_revision();
3925    // #1868: snapshot the deterministic n-free skip-path row-touch accumulator
3926    // AFTER the one-time priming eval above, so the reported delta measures only
3927    // the per-trial inner-synthesis row work across the κ-optimisation phase.
3928    let kphase_nfree_skip_touches_start = gam_solve::pirls::nfree_skip_row_element_touches();
3929
3930    // #1033: lift the ψ (log-κ) lower bound to the n-free rank-stable floor so the
3931    // optimizer never line-searches into the rank-deficient sliver where the
3932    // design-realization skip is soundly refused (→ O(n) reset_surface). The lift
3933    // touches ONLY the single design-moving ψ coordinate at `rho_dim`; all ρ
3934    // bounds are untouched. `psi_rank_stable_floor` is already constrained to lie
3935    // strictly inside `(psi_lo, theta0[rho_dim])`, so theta0 stays feasible.
3936    let lower_effective: std::borrow::Cow<'_, Array1<f64>> = match psi_rank_stable_floor {
3937        Some(floor) if coord_dim == 1 && floor > lower[rho_dim] => {
3938            let mut lifted = lower.clone();
3939            lifted[rho_dim] = floor;
3940            std::borrow::Cow::Owned(lifted)
3941        }
3942        _ => std::borrow::Cow::Borrowed(lower),
3943    };
3944    let lower = lower_effective.as_ref();
3945
3946    // #1033: clamp the ψ (log-κ) upper bound DOWN to the n-free rank-stable ceiling
3947    // so the optimizer never line-searches into the high-edge rank-deficient sliver
3948    // where the design-realization skip is soundly refused (→ O(n) reset_surface,
3949    // plus a second reset from the deficient pinning ψ). Touches ONLY the single
3950    // design-moving ψ coordinate at `rho_dim`; all ρ bounds are untouched.
3951    // `psi_rank_stable_ceiling` is already constrained to lie strictly inside
3952    // `(theta0[rho_dim], psi_hi)`, so theta0 stays feasible.
3953    let upper_effective: std::borrow::Cow<'_, Array1<f64>> = match psi_rank_stable_ceiling {
3954        Some(ceiling) if coord_dim == 1 && ceiling < upper[rho_dim] => {
3955            let mut clamped = upper.clone();
3956            clamped[rho_dim] = ceiling;
3957            std::borrow::Cow::Owned(clamped)
3958        }
3959        _ => std::borrow::Cow::Borrowed(upper),
3960    };
3961    let upper = upper_effective.as_ref();
3962
3963    let problem = exact_joint_multistart_outer_problem(
3964        theta0,
3965        lower,
3966        upper,
3967        rho_dim,
3968        coord_dim,
3969        theta_dim,
3970        Derivative::Analytic,
3971        if analytic_outer_hessian_available && !suppress_outer_hessian_for_nfree {
3972            DeclaredHessianForm::Either
3973        } else {
3974            // `Unavailable` when the n-free Gaussian ψ-lane is armed (#1033): the
3975            // planner then selects BFGS instead of ARC, so the κ loop issues only
3976            // `ValueAndGradient` evals and every in-window trial takes the n-free
3977            // design-realization skip.
3978            DeclaredHessianForm::Unavailable
3979        },
3980        prefer_gradient_only,
3981        // Single-block spatial path: penalty-like rho + spatial psi.
3982        // EFS/HybridEFS remain eligible (the Wood-Fasiolo PSD structure holds
3983        // for single-block families with β-independent joint H_L) UNLESS the
3984        // n-free Gaussian ψ-lane is armed (#1033): HybridEFS forms the trace Gram
3985        // `tr(H⁻¹ B_d H⁻¹ B_e)` from the n-dependent curvature slab `B_d`, so it
3986        // realizes O(n) per step exactly like a Hessian eval. Disabling the
3987        // fixed-point lane there forces the planner to BFGS (`(Analytic,
3988        // Unavailable)` → `S::Bfgs`), keeping every in-window κ-trial on the
3989        // n-free `ValueAndGradient` skip even when `n_params` exceeds the small-
3990        // BFGS threshold (aniso / multi-ψ).
3991        suppress_outer_hessian_for_nfree,
3992        seed_risk_profile_for_likelihood_family(&family),
3993        kappa_options.rel_tol.max(1e-6),
3994        kappa_options.max_outer_iter.max(1),
3995        // Rho-axis BFGS cap: log-λ's natural step is ≈ 5 per
3996        // `first_order_bfgs_loglambda_step_cap`. Anything tighter throttles
3997        // BFGS on flat REML valleys.
3998        Some(5.0),
3999        // Psi-axis BFGS cap: kappa / aniso-log-scale needs ~ln 2 per iter.
4000        Some(kappa_options.log_step.clamp(0.25, 1.0)),
4001        None,
4002        // Calibrate the outer to the n-scaled profiled REML/LAML objective for
4003        // every family — the iso-κ non-convergence cure (#1053 1-D Matérn,
4004        // #1066 2-D binomial geo, #1069 GP/kriging). p = baseline design column
4005        // count.
4006        Some((data.nrows(), baseline_design.design.ncols())),
4007        // #1464: widen the over-smoothing ρ ceiling + seed a high-λ probe when a
4008        // constant-curvature term is present (collapsing +κ kernel needs a large
4009        // smoothing λ beyond the historical ±12 box).
4010        !constant_curvature_term_indices(resolvedspec).is_empty(),
4011        // The scalar Matérn endpoint comparison has already selected and
4012        // certified the range basin. Give its explicit theta0 the only joint
4013        // start; anisotropic and non-Matérn paths keep their established seed
4014        // policy.
4015        kind == SpatialHyperKind::Isotropic
4016            && constant_curvature_term_indices(resolvedspec).is_empty()
4017            && spatial_terms.iter().any(|&term_idx| {
4018                matches!(
4019                    resolvedspec
4020                        .smooth_terms
4021                        .get(term_idx)
4022                        .map(|term| &term.basis),
4023                    Some(SmoothBasisSpec::Matern { .. })
4024                )
4025            }),
4026    )?;
4027
4028    let eval_outer = |ctx: &mut &mut SpatialJointContext<'_>,
4029                      theta: &Array1<f64>,
4030                      order: OuterEvalOrder|
4031     -> Result<OuterEval, EstimationError> {
4032        let t0 = std::time::Instant::now();
4033        let allow_second_order_for_call = matches!(order, OuterEvalOrder::ValueGradientHessian)
4034            && analytic_outer_hessian_available;
4035        let gate = ctx.nfree_skip_gate_status(theta, allow_second_order_for_call, true);
4036        let resets_before = ctx.evaluator.slow_path_reset_count();
4037        let raw = ctx.eval_full(theta, order, analytic_outer_hessian_available);
4038        let reset_delta = ctx
4039            .evaluator
4040            .slow_path_reset_count()
4041            .saturating_sub(resets_before);
4042        if reset_delta > 0 {
4043            if !gate.shape {
4044                kphase_nfree_miss_shape.set(kphase_nfree_miss_shape.get() + reset_delta);
4045            }
4046            if gate.shape && !gate.value {
4047                kphase_nfree_miss_value.set(kphase_nfree_miss_value.get() + reset_delta);
4048            }
4049            if gate.shape && gate.value && !gate.gradient {
4050                kphase_nfree_miss_gradient.set(kphase_nfree_miss_gradient.get() + reset_delta);
4051            }
4052            if gate.shape && gate.value && gate.gradient && !gate.penalty {
4053                kphase_nfree_miss_penalty.set(kphase_nfree_miss_penalty.get() + reset_delta);
4054            }
4055            if gate.shape && gate.value && gate.gradient && gate.penalty && !gate.revision {
4056                kphase_nfree_miss_revision.set(kphase_nfree_miss_revision.get() + reset_delta);
4057            }
4058            if gate.shape
4059                && gate.value
4060                && gate.gradient
4061                && gate.penalty
4062                && gate.revision
4063                && gate.second_order
4064            {
4065                kphase_nfree_miss_second_order
4066                    .set(kphase_nfree_miss_second_order.get() + reset_delta);
4067            }
4068            if gate.would_skip(true) {
4069                kphase_nfree_miss_other.set(kphase_nfree_miss_other.get() + reset_delta);
4070            }
4071        }
4072        let elapsed_s = t0.elapsed().as_secs_f64();
4073        kphase_eval_calls.set(kphase_eval_calls.get() + 1);
4074        kphase_eval_total_s.set(kphase_eval_total_s.get() + elapsed_s);
4075        let (theta_norm, log_kappa_norm) = kphase_log_norms(theta, rho_dim);
4076        log::info!(
4077            "[KAPPA-PHASE] phase=eval_outer call={} order={:?} design_revision={:?} theta_norm={:.4e} log_kappa_norm={:.4e} elapsed_s={:.4}",
4078            kphase_eval_calls.get(),
4079            order,
4080            Some(ctx.cache.design_revision()),
4081            theta_norm,
4082            log_kappa_norm,
4083            elapsed_s,
4084        );
4085        match raw {
4086            Ok((cost, grad, hess)) => Ok(OuterEval {
4087                cost,
4088                gradient: grad,
4089                hessian: hess,
4090                inner_beta_hint: None,
4091            }),
4092            // A trial hyperparameter at which the spatial kernel design /
4093            // ψ-derivatives are non-constructible is an infeasible point, not
4094            // a fatal error: the gradient/Hessian path must retreat exactly as
4095            // the cost-only path (which already returns +∞) does. Returning
4096            // `OuterEval::infeasible` keeps the two paths symmetric so a single
4097            // bad probe — e.g. an anisotropy that overflows the Duchon radial
4098            // kernel — no longer aborts the whole REML optimization.
4099            Err(err) if is_recoverable_trial_point_error(&err) => {
4100                log::debug!(
4101                    "[{label}] trial point infeasible (kernel design \
4102                     not constructible at theta={theta:?}): {err}; retreating",
4103                );
4104                Ok(OuterEval::infeasible(theta_dim))
4105            }
4106            Err(err) => Err(err),
4107        }
4108    };
4109
4110    let mut obj = problem.build_objective_with_eval_order(
4111        &mut ctx,
4112        |ctx: &mut &mut SpatialJointContext<'_>, theta: &Array1<f64>| {
4113            let t0 = std::time::Instant::now();
4114            let gate = ctx.nfree_skip_gate_status(theta, false, false);
4115            let resets_before = ctx.evaluator.slow_path_reset_count();
4116            let cost = ctx.eval_cost(theta);
4117            let reset_delta = ctx
4118                .evaluator
4119                .slow_path_reset_count()
4120                .saturating_sub(resets_before);
4121            if reset_delta > 0 {
4122                if !gate.shape {
4123                    kphase_nfree_miss_shape.set(kphase_nfree_miss_shape.get() + reset_delta);
4124                }
4125                if gate.shape && !gate.value {
4126                    kphase_nfree_miss_value.set(kphase_nfree_miss_value.get() + reset_delta);
4127                }
4128                if gate.shape && gate.value && !gate.penalty {
4129                    kphase_nfree_miss_penalty.set(kphase_nfree_miss_penalty.get() + reset_delta);
4130                }
4131                if gate.shape && gate.value && gate.penalty && !gate.revision {
4132                    kphase_nfree_miss_revision.set(kphase_nfree_miss_revision.get() + reset_delta);
4133                }
4134                if gate.would_skip(false) {
4135                    kphase_nfree_miss_other.set(kphase_nfree_miss_other.get() + reset_delta);
4136                }
4137            }
4138            let elapsed_s = t0.elapsed().as_secs_f64();
4139            kphase_cost_calls.set(kphase_cost_calls.get() + 1);
4140            kphase_cost_total_s.set(kphase_cost_total_s.get() + elapsed_s);
4141            let (theta_norm, log_kappa_norm) = kphase_log_norms(theta, rho_dim);
4142            log::info!(
4143                "[KAPPA-PHASE] phase=cost call={} design_revision={:?} theta_norm={:.4e} log_kappa_norm={:.4e} elapsed_s={:.4}",
4144                kphase_cost_calls.get(),
4145                Some(ctx.cache.design_revision()),
4146                theta_norm,
4147                log_kappa_norm,
4148                elapsed_s,
4149            );
4150            Ok(cost)
4151        },
4152        |ctx: &mut &mut SpatialJointContext<'_>, theta: &Array1<f64>| {
4153            eval_outer(
4154                ctx,
4155                theta,
4156                // #1033: when the n-free Gaussian ψ-lane is armed we suppress the
4157                // outer Hessian and route BFGS — so this default gradient eval MUST
4158                // request `ValueAndGradient`, not `ValueGradientHessian`. A
4159                // second-order order sets `allow_second_order`, which forces
4160                // `ensure_theta` → the O(n) design re-realization (the Hessian slab
4161                // is irreducibly n-dependent), DISARMING the design-revision fast
4162                // path for every trial — exactly the O(n) κ-loop this lane exists to
4163                // remove. Gating only the planner's solver (Unavailable→BFGS)
4164                // without gating this eval-order left every trial second-order.
4165                if analytic_outer_hessian_available && !suppress_outer_hessian_for_nfree {
4166                    OuterEvalOrder::ValueGradientHessian
4167                } else {
4168                    OuterEvalOrder::ValueAndGradient
4169                },
4170            )
4171        },
4172        |ctx: &mut &mut SpatialJointContext<'_>, theta: &Array1<f64>, order: OuterEvalOrder| {
4173            eval_outer(ctx, theta, order)
4174        },
4175        Some(|ctx: &mut &mut SpatialJointContext<'_>| {
4176            ctx.reset();
4177        }),
4178        Some(|ctx: &mut &mut SpatialJointContext<'_>, theta: &Array1<f64>| {
4179            let t0 = std::time::Instant::now();
4180            let eval = ctx.eval_efs(theta);
4181            let elapsed_s = t0.elapsed().as_secs_f64();
4182            kphase_efs_calls.set(kphase_efs_calls.get() + 1);
4183            kphase_efs_total_s.set(kphase_efs_total_s.get() + elapsed_s);
4184            let (theta_norm, log_kappa_norm) = kphase_log_norms(theta, rho_dim);
4185            log::info!(
4186                "[KAPPA-PHASE] phase=efs call={} design_revision={:?} theta_norm={:.4e} log_kappa_norm={:.4e} elapsed_s={:.4}",
4187                kphase_efs_calls.get(),
4188                Some(ctx.cache.design_revision()),
4189                theta_norm,
4190                log_kappa_norm,
4191                elapsed_s,
4192            );
4193            eval
4194        }),
4195    );
4196
4197    let run_label = match kind {
4198        SpatialHyperKind::Anisotropic => "aniso-psi joint REML",
4199        SpatialHyperKind::Isotropic => "iso-kappa joint REML",
4200    };
4201    let result = problem.run(&mut obj, run_label)?;
4202    if !result.converged {
4203        crate::bail_invalid_estim!(
4204            "{} did not converge after {} iterations (final_objective={:.6e}, final_grad_norm={})",
4205            run_label,
4206            result.iterations,
4207            result.final_value,
4208            result.final_grad_norm_report(),
4209        );
4210    }
4211    drop(obj);
4212    let kphase_total_s = kphase_optim_start.elapsed().as_secs_f64();
4213    let kphase_slow_resets = ctx
4214        .evaluator
4215        .slow_path_reset_count()
4216        .saturating_sub(kphase_slow_resets_start);
4217    let kphase_design_revision_delta = ctx
4218        .cache
4219        .design_revision()
4220        .saturating_sub(kphase_design_revision_start);
4221    let kphase_nfree_skip_touches = gam_solve::pirls::nfree_skip_row_element_touches()
4222        .saturating_sub(kphase_nfree_skip_touches_start);
4223    log::info!(
4224        "[KAPPA-PHASE-SUMMARY] n_rows={} log_kappa_dim={} n_cost={} cost_total_s={:.4} n_eval={} eval_total_s={:.4} n_efs={} efs_total_s={:.4} slow_path_resets={} design_revision_delta={} nfree_skip_row_touches={} nfree_miss_shape={} nfree_miss_value={} nfree_miss_gradient={} nfree_miss_penalty={} nfree_miss_revision={} nfree_miss_second_order={} nfree_miss_other={} optim_total_s={:.4}",
4225        data.nrows(),
4226        kphase_log_kappa_dim,
4227        kphase_cost_calls.get(),
4228        kphase_cost_total_s.get(),
4229        kphase_eval_calls.get(),
4230        kphase_eval_total_s.get(),
4231        kphase_efs_calls.get(),
4232        kphase_efs_total_s.get(),
4233        kphase_slow_resets,
4234        kphase_design_revision_delta,
4235        kphase_nfree_skip_touches,
4236        kphase_nfree_miss_shape.get(),
4237        kphase_nfree_miss_value.get(),
4238        kphase_nfree_miss_gradient.get(),
4239        kphase_nfree_miss_penalty.get(),
4240        kphase_nfree_miss_revision.get(),
4241        kphase_nfree_miss_second_order.get(),
4242        kphase_nfree_miss_other.get(),
4243        kphase_total_s,
4244    );
4245    let timing = SpatialLengthScaleOptimizationTiming {
4246        log_kappa_dim: kphase_log_kappa_dim,
4247        cost_calls: kphase_cost_calls.get(),
4248        cost_total_s: kphase_cost_total_s.get(),
4249        eval_calls: kphase_eval_calls.get(),
4250        eval_total_s: kphase_eval_total_s.get(),
4251        efs_calls: kphase_efs_calls.get(),
4252        efs_total_s: kphase_efs_total_s.get(),
4253        slow_path_resets: kphase_slow_resets,
4254        design_revision_delta: kphase_design_revision_delta,
4255        nfree_skip_row_touches: kphase_nfree_skip_touches,
4256        nfree_miss_shape: kphase_nfree_miss_shape.get(),
4257        nfree_miss_value: kphase_nfree_miss_value.get(),
4258        nfree_miss_gradient: kphase_nfree_miss_gradient.get(),
4259        nfree_miss_penalty: kphase_nfree_miss_penalty.get(),
4260        nfree_miss_revision: kphase_nfree_miss_revision.get(),
4261        nfree_miss_second_order: kphase_nfree_miss_second_order.get(),
4262        nfree_miss_other: kphase_nfree_miss_other.get(),
4263        optim_total_s: kphase_total_s,
4264    };
4265    log::trace!(
4266        "[{}] converged in {} iterations, final_value={:.6e}, grad_norm={}",
4267        label,
4268        result.iterations,
4269        result.final_value,
4270        result.final_grad_norm_report(),
4271    );
4272    // No sum-to-zero enforcement needed: ψ coordinates are unconstrained during
4273    // optimization. For the anisotropic kind the decomposition into (ψ̄, η)
4274    // happens later in apply_tospec.
4275    let theta_star = result.rho;
4276    Ok((theta_star, result.final_value, timing))
4277}
4278
4279/// Apply a length scale to a single `SmoothTermSpec` (independent of any
4280/// outer `TermCollectionSpec`). Mirrors `set_spatial_length_scale` but on a
4281/// term in isolation; used by the incremental realizer's cached planned spec.
4282fn set_single_term_spatial_length_scale(
4283    term: &mut SmoothTermSpec,
4284    length_scale: f64,
4285) -> Result<(), EstimationError> {
4286    match &mut term.basis {
4287        SmoothBasisSpec::ThinPlate { spec, .. } => {
4288            spec.length_scale = length_scale;
4289            Ok(())
4290        }
4291        SmoothBasisSpec::Matern { spec, .. } => {
4292            spec.length_scale.set_resolved(length_scale);
4293            Ok(())
4294        }
4295        SmoothBasisSpec::Duchon { spec, .. } => {
4296            spec.length_scale = Some(length_scale);
4297            Ok(())
4298        }
4299        _ => Err(EstimationError::InvalidInput(format!(
4300            "term '{}' does not expose a spatial length scale",
4301            term.name
4302        ))),
4303    }
4304}
4305
4306/// Apply anisotropy contrasts to a single `SmoothTermSpec`. Mirrors
4307/// `set_spatial_aniso_log_scales` but on a term in isolation; used by the
4308/// incremental realizer's cached planned spec.
4309fn set_single_term_spatial_aniso_log_scales(
4310    term: &mut SmoothTermSpec,
4311    eta: Vec<f64>,
4312) -> Result<(), EstimationError> {
4313    let eta = center_aniso_log_scales(&eta);
4314    match &mut term.basis {
4315        SmoothBasisSpec::Matern { spec, .. } => {
4316            spec.aniso_log_scales = Some(eta);
4317            Ok(())
4318        }
4319        SmoothBasisSpec::Duchon { spec, .. } => {
4320            spec.aniso_log_scales = Some(eta);
4321            Ok(())
4322        }
4323        _ => Err(EstimationError::InvalidInput(format!(
4324            "term '{}' does not support aniso_log_scales",
4325            term.name
4326        ))),
4327    }
4328}
4329
4330/// Freeze the design-moving representer length-scale dial on every measure-jet
4331/// term in `spec` (sets `learn_length_scale = false`), so ℓ stays at its
4332/// realized auto value with no outer REML enrollment.
4333///
4334/// Used by COUPLED-block families (bernoulli marginal-slope: a shared mjs
4335/// surface feeds both the marginal mean and the log-slope). In that coupling a
4336/// design-moving kernel-scale dial on the shared covariates is an
4337/// identifiability hazard: the outer search can reach a sharp ℓ at which a
4338/// marginal smooth direction trades off against the log-slope into a
4339/// separation-scale runaway (#1116). A single Gaussian surface has no such
4340/// coupling and keeps ℓ learnable. Returns the number of terms frozen.
4341/// The signed sectional curvature κ of a constant-curvature smooth at
4342/// `term_idx`, or `None` if that term is not a `curv(...)` smooth. After a fit
4343/// with κ-optimization enabled this reads the **fitted κ̂** out of the resolved
4344/// spec (`freeze_term_collection_from_design` writes the optimized κ back into
4345/// the spec, and `BasisMetadata::ConstantCurvature.kappa` carries the same
4346/// value). This is the headline #944 estimand accessor — the κ̂ in
4347/// "κ̂ = −1.8 (95% CI …)". Mirrors [`get_spatial_length_scale`].
4348pub fn get_constant_curvature_kappa(spec: &TermCollectionSpec, term_idx: usize) -> Option<f64> {
4349    constant_curvature_term_spec(spec, term_idx).map(|cc| cc.kappa)
4350}
4351
4352/// `true` when `term_idx` is a `curv(...)` smooth whose user PINNED the
4353/// sectional curvature with an explicit `kappa=` (the mgcv-`sp=` convention,
4354/// gam#2152). A pinned κ is a fixed geometry: the outer loop must hold it
4355/// constant and never run the continuous curvature-fair profile optimizer on
4356/// that term. Non-CC terms and CC terms whose `kappa=` was omitted (κ free,
4357/// #944/#1464 estimation) return `false`.
4358pub fn constant_curvature_kappa_is_fixed(spec: &TermCollectionSpec, term_idx: usize) -> bool {
4359    constant_curvature_term_spec(spec, term_idx).is_some_and(|cc| cc.kappa_fixed)
4360}
4361
4362/// Indices of every constant-curvature (`curv(...)`) smooth term in `spec`.
4363pub fn constant_curvature_term_indices(spec: &TermCollectionSpec) -> Vec<usize> {
4364    (0..spec.smooth_terms.len())
4365        .filter(|&idx| constant_curvature_term_spec(spec, idx).is_some())
4366        .collect()
4367}
4368
4369#[derive(Debug, Clone)]
4370struct SingleSmoothTermRealization {
4371    design_local: DesignMatrix,
4372    term: SmoothTerm,
4373    dropped_penaltyinfo: Vec<DroppedPenaltyBlockInfo>,
4374}
4375
4376impl SingleSmoothTermRealization {
4377    fn active_penalty_count(&self) -> usize {
4378        self.term.active_penalties.len()
4379    }
4380}
4381
4382fn build_single_smooth_term_realization_with_policy(
4383    data: ArrayView2<'_, f64>,
4384    termspec: &SmoothTermSpec,
4385    policy: &gam_runtime::resource::ResourcePolicy,
4386) -> Result<SingleSmoothTermRealization, BasisError> {
4387    let mut workspace = gam_terms::basis::BasisWorkspace::with_policy(policy.clone());
4388    let raw =
4389        build_smooth_design_withworkspace(data, std::slice::from_ref(termspec), &mut workspace)?;
4390    finish_single_smooth_term_realization(raw)
4391}
4392
4393fn finish_single_smooth_term_realization(
4394    raw: RawSmoothDesign,
4395) -> Result<SingleSmoothTermRealization, BasisError> {
4396    let RawSmoothDesign {
4397        term_designs,
4398        dropped_penaltyinfo,
4399        terms,
4400        ..
4401    } = raw;
4402    let term = terms.into_iter().next().ok_or_else(|| {
4403        BasisError::InvalidInput("single-term smooth build returned no term".to_string())
4404    })?;
4405    let design = term_designs.into_iter().next().ok_or_else(|| {
4406        BasisError::InvalidInput("single-term smooth build returned no term design".to_string())
4407    })?;
4408
4409    Ok(SingleSmoothTermRealization {
4410        design_local: design,
4411        term,
4412        dropped_penaltyinfo,
4413    })
4414}
4415
4416/// Wrap a fresh `LocalSmoothTermBuild` (produced by `build_single_local_smooth_term`)
4417/// into a `SingleSmoothTermRealization`. Mirrors the single-term portion of
4418/// `build_smooth_design_withworkspace_unvalidated`, but skips the joint center
4419/// planner and per-term workspace fork — the realizer drives κ-only rebuilds
4420/// directly with its persistent workspace so basis caches survive across BFGS
4421/// κ proposals.
4422fn wrap_local_build_as_realization(
4423    mut local: LocalSmoothTermBuild,
4424    termspec: &SmoothTermSpec,
4425) -> Result<SingleSmoothTermRealization, String> {
4426    let p_local = local.dim;
4427    let lb_local = if local.box_reparam {
4428        shape_lower_bounds_local(termspec.shape, p_local)
4429    } else {
4430        None
4431    };
4432
4433    let dropped_penaltyinfo = local
4434        .dropped_penalties
4435        .iter()
4436        .map(|info| DroppedPenaltyBlockInfo {
4437            termname: Some(termspec.name.clone()),
4438            penalty: info.clone(),
4439        })
4440        .collect();
4441
4442    // Stage-2 joint-null absorption rotation, same logic as the main
4443    // aggregation loop in `build_smooth_design_withworkspace_unvalidated`:
4444    // apply Q when Some AND the smooth has no shape constraints.
4445    let applied_rotation: Option<gam_terms::basis::JointNullRotation> = match (
4446        local.joint_null_rotation.take(),
4447        lb_local.is_some(),
4448        local.linear_constraints.is_some(),
4449    ) {
4450        (Some(rot), false, false) => {
4451            let q = &rot.rotation;
4452            local.design =
4453                apply_smooth_transform_to_design(local.design.clone(), q, &termspec.name).map_err(
4454                    |e| {
4455                        format!(
4456                            "joint-null absorption rotation failed for term '{}': {}",
4457                            termspec.name, e
4458                        )
4459                    },
4460                )?;
4461            for penalty in &mut local.active_penalties {
4462                let qt_s = gam_linalg::faer_ndarray::fast_atb(q, &penalty.matrix);
4463                penalty.matrix = gam_linalg::faer_ndarray::fast_ab(&qt_s, q);
4464                penalty.null_eigenvectors = penalty
4465                    .null_eigenvectors
4466                    .as_ref()
4467                    .map(|basis| gam_linalg::faer_ndarray::fast_atb(q, basis));
4468                penalty.op = None;
4469                penalty.info.kronecker_factors = None;
4470            }
4471            local.kronecker_factored = None;
4472            Some(rot)
4473        }
4474        (Some(_), _, _) => None,
4475        (None, _, _) => None,
4476    };
4477
4478    let smooth_term = SmoothTerm {
4479        name: termspec.name.clone(),
4480        coeff_range: 0..p_local,
4481        shape: termspec.shape,
4482        active_penalties: local.active_penalties.clone(),
4483        dropped_penalties: local.dropped_penalties.clone(),
4484        metadata: local.metadata.clone(),
4485        lower_bounds_local: lb_local,
4486        linear_constraints_local: local.linear_constraints.clone(),
4487        kronecker_factored: local.kronecker_factored.take(),
4488        joint_null_rotation: applied_rotation,
4489        // Single-term realizations never run the global ownership pass, so
4490        // there is no overlap residualization to export here (#978).
4491        unabsorbed_global_orthogonality: None,
4492    };
4493
4494    Ok(SingleSmoothTermRealization {
4495        design_local: local.design,
4496        term: smooth_term,
4497        dropped_penaltyinfo,
4498    })
4499}
4500
4501/// Extract the κ-invariant pieces of a freshly-built spatial basis — center
4502/// cloud (in standardized coords) and `input_scale` — and bake them into a
4503/// `SmoothTermSpec` whose `center_strategy` becomes `UserProvided` and whose
4504/// `input_scale` is `Some`. Subsequent rebuilds driven from this cached spec
4505/// will short-circuit `select_centers_by_strategy` (KMeans / FarthestPoint /
4506/// EqualMass cluster searches over n×d data) and isotropic scale estimation,
4507/// leaving only the κ-dependent kernel
4508/// values and basis assembly. Returns `None` for non-spatial families or when
4509/// the metadata does not yet expose the required pieces (for instance when a
4510/// ThinPlate request was auto-promoted to Duchon during the build).
4511fn freeze_geometry_from_metadata(
4512    termspec: &SmoothTermSpec,
4513    metadata: &BasisMetadata,
4514) -> Option<SmoothTermSpec> {
4515    let mut frozen = termspec.clone();
4516    match (&mut frozen.basis, metadata) {
4517        (
4518            SmoothBasisSpec::Matern {
4519                spec,
4520                input_scale: spec_scale,
4521                ..
4522            },
4523            BasisMetadata::Matern {
4524                centers,
4525                input_scale: metadata_scale,
4526                identifiability_transform,
4527                ..
4528            },
4529        ) => {
4530            spec.center_strategy = CenterStrategy::UserProvided(centers.clone());
4531            *spec_scale = Some(*metadata_scale);
4532            // Freeze the cold-build coefficient chart. Double-penalty topology
4533            // is structural (the explicit intercept only), so no numerical
4534            // nullspace decision needs to be carried across κ trials.
4535            if let Some(transform) = identifiability_transform.clone() {
4536                spec.identifiability = MaternIdentifiability::FrozenTransform { transform };
4537            }
4538            Some(frozen)
4539        }
4540        (
4541            SmoothBasisSpec::Duchon {
4542                spec,
4543                input_scale: spec_scale,
4544                ..
4545            },
4546            BasisMetadata::Duchon {
4547                centers,
4548                input_scale: metadata_scale,
4549                ..
4550            },
4551        ) => {
4552            spec.center_strategy = CenterStrategy::UserProvided(centers.clone());
4553            *spec_scale = Some(*metadata_scale);
4554            Some(frozen)
4555        }
4556        (
4557            SmoothBasisSpec::ThinPlate {
4558                spec,
4559                input_scale: spec_scale,
4560                ..
4561            },
4562            BasisMetadata::ThinPlate {
4563                centers,
4564                input_scale: metadata_scale,
4565                ..
4566            },
4567        ) => {
4568            spec.center_strategy = CenterStrategy::UserProvided(centers.clone());
4569            *spec_scale = Some(*metadata_scale);
4570            Some(frozen)
4571        }
4572        // Family mismatch (e.g. ThinPlate auto-promotion to Duchon) leaves the
4573        // cache empty; we'll retry materialization on the next κ apply.
4574        _ => None,
4575    }
4576}
4577
4578fn rebuild_smooth_auxiliary_state(
4579    smooth: &mut SmoothDesign,
4580    dropped_penaltyinfo_by_term: &[Vec<DroppedPenaltyBlockInfo>],
4581) -> Result<(), String> {
4582    if dropped_penaltyinfo_by_term.len() != smooth.terms.len() {
4583        return Err(SmoothError::dimension_mismatch(format!(
4584            "smooth dropped-penalty cache mismatch: terms={}, dropped_sets={}",
4585            smooth.terms.len(),
4586            dropped_penaltyinfo_by_term.len()
4587        ))
4588        .into());
4589    }
4590
4591    let total_p = smooth.total_smooth_cols();
4592    let mut coefficient_lower_bounds = Array1::<f64>::from_elem(total_p, f64::NEG_INFINITY);
4593    let mut any_bounds = false;
4594    let mut linear_constraintrows: Vec<Array1<f64>> = Vec::new();
4595    let mut linear_constraint_b: Vec<f64> = Vec::new();
4596
4597    for term in &smooth.terms {
4598        let range = term.coeff_range.clone();
4599        if let Some(lb_local) = term.lower_bounds_local.as_ref() {
4600            if lb_local.len() != range.len() {
4601                return Err(SmoothError::dimension_mismatch(format!(
4602                    "smooth lower-bound cache mismatch for term '{}': bounds={}, coeffs={}",
4603                    term.name,
4604                    lb_local.len(),
4605                    range.len()
4606                ))
4607                .into());
4608            }
4609            coefficient_lower_bounds
4610                .slice_mut(s![range.clone()])
4611                .assign(lb_local);
4612            any_bounds = true;
4613        }
4614        if let Some(lin_local) = term.linear_constraints_local.as_ref() {
4615            if lin_local.a.ncols() != range.len() {
4616                return Err(SmoothError::dimension_mismatch(format!(
4617                    "smooth linear-constraint cache mismatch for term '{}': cols={}, coeffs={}",
4618                    term.name,
4619                    lin_local.a.ncols(),
4620                    range.len()
4621                ))
4622                .into());
4623            }
4624            for r in 0..lin_local.a.nrows() {
4625                let mut row = Array1::<f64>::zeros(total_p);
4626                row.slice_mut(s![range.clone()]).assign(&lin_local.a.row(r));
4627                linear_constraintrows.push(row);
4628                linear_constraint_b.push(lin_local.b[r]);
4629            }
4630        }
4631    }
4632
4633    smooth.coefficient_lower_bounds = if any_bounds {
4634        Some(coefficient_lower_bounds)
4635    } else {
4636        None
4637    };
4638    smooth.linear_constraints = if linear_constraintrows.is_empty() {
4639        None
4640    } else {
4641        let mut a = Array2::<f64>::zeros((linear_constraintrows.len(), total_p));
4642        for (i, row) in linear_constraintrows.iter().enumerate() {
4643            a.row_mut(i).assign(row);
4644        }
4645        Some(LinearInequalityConstraints {
4646            a,
4647            b: Array1::from_vec(linear_constraint_b),
4648        })
4649    };
4650    smooth.dropped_penaltyinfo = dropped_penaltyinfo_by_term
4651        .iter()
4652        .flat_map(|infos| infos.iter().cloned())
4653        .collect();
4654    Ok(())
4655}
4656
4657fn rebuild_term_collection_auxiliary_state(
4658    spec: &TermCollectionSpec,
4659    design: &mut TermCollectionDesign,
4660) -> Result<(), String> {
4661    if spec.linear_terms.len() != design.linear_ranges.len() {
4662        return Err(SmoothError::dimension_mismatch(format!(
4663            "term-collection linear bookkeeping mismatch: spec_terms={}, design_ranges={}",
4664            spec.linear_terms.len(),
4665            design.linear_ranges.len()
4666        ))
4667        .into());
4668    }
4669
4670    let p_total = design.design.ncols();
4671    let smooth_start = p_total.saturating_sub(design.smooth.total_smooth_cols());
4672    let mut coefficient_lower_bounds = Array1::<f64>::from_elem(p_total, f64::NEG_INFINITY);
4673    let mut any_bounds = false;
4674    let mut linear_constraintrows: Vec<Array1<f64>> = Vec::new();
4675    let mut linear_constraint_b: Vec<f64> = Vec::new();
4676
4677    for (linear, (_, range)) in spec.linear_terms.iter().zip(design.linear_ranges.iter()) {
4678        if range.len() != 1 {
4679            return Err(SmoothError::dimension_mismatch(format!(
4680                "linear term '{}' expected one coefficient column, found {}",
4681                linear.name,
4682                range.len()
4683            ))
4684            .into());
4685        }
4686        let col = range.start;
4687        if let Some(lb) = linear.coefficient_min {
4688            let mut row = Array1::<f64>::zeros(p_total);
4689            row[col] = 1.0;
4690            linear_constraintrows.push(row);
4691            linear_constraint_b.push(lb);
4692        }
4693        if let Some(ub) = linear.coefficient_max {
4694            let mut row = Array1::<f64>::zeros(p_total);
4695            row[col] = -1.0;
4696            linear_constraintrows.push(row);
4697            linear_constraint_b.push(-ub);
4698        }
4699    }
4700
4701    if let Some(lb_smooth) = design.smooth.coefficient_lower_bounds.as_ref() {
4702        if lb_smooth.len() != design.smooth.total_smooth_cols() {
4703            return Err(SmoothError::dimension_mismatch(format!(
4704                "smooth lower-bound width mismatch: bounds={}, smooth_cols={}",
4705                lb_smooth.len(),
4706                design.smooth.total_smooth_cols()
4707            ))
4708            .into());
4709        }
4710        coefficient_lower_bounds
4711            .slice_mut(s![
4712                smooth_start..(smooth_start + design.smooth.total_smooth_cols())
4713            ])
4714            .assign(lb_smooth);
4715        any_bounds = true;
4716    }
4717    if let Some(lin_smooth) = design.smooth.linear_constraints.as_ref() {
4718        if lin_smooth.a.ncols() != design.smooth.total_smooth_cols() {
4719            return Err(SmoothError::dimension_mismatch(format!(
4720                "smooth linear-constraint width mismatch: cols={}, smooth_cols={}",
4721                lin_smooth.a.ncols(),
4722                design.smooth.total_smooth_cols()
4723            ))
4724            .into());
4725        }
4726        let mut a_global = Array2::<f64>::zeros((lin_smooth.a.nrows(), p_total));
4727        a_global
4728            .slice_mut(s![
4729                ..,
4730                smooth_start..(smooth_start + design.smooth.total_smooth_cols())
4731            ])
4732            .assign(&lin_smooth.a);
4733        for r in 0..a_global.nrows() {
4734            linear_constraintrows.push(a_global.row(r).to_owned());
4735            linear_constraint_b.push(lin_smooth.b[r]);
4736        }
4737    }
4738
4739    let lower_bound_constraints = if any_bounds {
4740        linear_constraints_from_lower_bounds_global(&coefficient_lower_bounds)
4741    } else {
4742        None
4743    };
4744    let explicit_linear_constraints = if linear_constraintrows.is_empty() {
4745        None
4746    } else {
4747        let mut a = Array2::<f64>::zeros((linear_constraintrows.len(), p_total));
4748        for (i, row) in linear_constraintrows.iter().enumerate() {
4749            a.row_mut(i).assign(row);
4750        }
4751        Some(LinearInequalityConstraints {
4752            a,
4753            b: Array1::from_vec(linear_constraint_b),
4754        })
4755    };
4756
4757    design.coefficient_lower_bounds = if any_bounds {
4758        Some(coefficient_lower_bounds)
4759    } else {
4760        None
4761    };
4762    design.linear_constraints =
4763        merge_linear_constraints_global(explicit_linear_constraints, lower_bound_constraints)
4764            .map_err(|error| error.to_string())?;
4765    design.dropped_penaltyinfo = design.smooth.dropped_penaltyinfo.clone();
4766    Ok(())
4767}
4768
4769fn theta_values_match(left: &Array1<f64>, right: &Array1<f64>) -> bool {
4770    left.len() == right.len()
4771        && left
4772            .iter()
4773            .zip(right.iter())
4774            .all(|(&l, &r)| l.to_bits() == r.to_bits())
4775}
4776
4777fn latent_values_match(left: &Array1<f64>, right: &Array1<f64>) -> bool {
4778    theta_values_match(left, right)
4779}
4780
4781fn spatial_aniso_matches(left: Option<&[f64]>, right: Option<&[f64]>) -> bool {
4782    match (left, right) {
4783        (None, None) => true,
4784        (Some(a), Some(b)) => {
4785            a.len() == b.len()
4786                && a.iter()
4787                    .zip(b.iter())
4788                    .all(|(&x, &y)| x.to_bits() == y.to_bits())
4789        }
4790        _ => false,
4791    }
4792}
4793
4794fn spatial_length_scale_matches(left: Option<f64>, right: Option<f64>) -> bool {
4795    match (left, right) {
4796        (None, None) => true,
4797        (Some(a), Some(b)) => a.to_bits() == b.to_bits(),
4798        _ => false,
4799    }
4800}
4801
4802struct FrozenTermCollectionIncrementalRealizer<'d> {
4803    data: ArrayView2<'d, f64>,
4804    spec: TermCollectionSpec,
4805    design: TermCollectionDesign,
4806    fixed_blocks: Vec<DesignBlock>,
4807    dropped_penaltyinfo_by_term: Vec<Vec<DroppedPenaltyBlockInfo>>,
4808    smooth_penalty_ranges: Vec<Range<usize>>,
4809    full_penalty_ranges: Vec<Range<usize>>,
4810    /// Persistent workspace for basis cache reuse across κ proposals.
4811    /// Distance matrices are cached here so they're computed once and
4812    /// reused across repeated `apply_log_kappa_to_term` calls.
4813    basisworkspace: gam_terms::basis::BasisWorkspace,
4814    /// Per-term cached realization geometry for incremental κ updates.
4815    ///
4816    /// On the first κ-driven rebuild of term `i`, this slot is populated with a
4817    /// `SmoothTermSpec` whose κ-invariant geometry — center cloud (as
4818    /// `CenterStrategy::UserProvided`) and `input_scale` — has been frozen
4819    /// out of the realized basis metadata. Subsequent
4820    /// `apply_log_kappa_to_term` calls reuse this spec, mutating only the
4821    /// κ / aniso fields. This short-circuits `select_centers_by_strategy`
4822    /// (KMeans / FarthestPoint / EqualMass cluster searches over the n×d data
4823    /// matrix) and isotropic scale estimation over n rows on every BFGS
4824    /// κ-eval, leaving the kernel-value pass and
4825    /// basis assembly as the only work.
4826    spatial_realization_geometry: Vec<Option<SmoothTermSpec>>,
4827    /// Monotonic counter incremented every time `apply_log_kappa` actually
4828    /// rebuilds the realized design / smooth penalties. Read by the
4829    /// design-revision-counter fast path in `ExternalJointHyperEvaluator`
4830    /// to skip redundant canonical-penalty rebuilds and cache wipes when
4831    /// the outer BFGS loop probes the same ψ twice in a row.
4832    design_revision: u64,
4833}
4834
4835impl<'d> std::fmt::Debug for FrozenTermCollectionIncrementalRealizer<'d> {
4836    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
4837        f.debug_struct("FrozenTermCollectionIncrementalRealizer")
4838            .field("data_shape", &(self.data.nrows(), self.data.ncols()))
4839            .field("fixed_blocks", &self.fixed_blocks.len())
4840            .finish_non_exhaustive()
4841    }
4842}
4843
4844/// Translate the authoritative emitted global penalty layout into the two
4845/// coordinate systems the incremental realizer updates.
4846///
4847/// The model-global ranges come directly from `TermCollectionDesign`; the
4848/// smooth-local ranges are their exact translation past the recorded leading
4849/// penalty prefix. Keeping this outside the constructor makes the layout
4850/// invariant independently testable without constructing any κ-specific
4851/// spatial caches.
4852fn emitted_smooth_penalty_ranges(
4853    design: &TermCollectionDesign,
4854) -> Result<(Vec<Range<usize>>, Vec<Range<usize>>), String> {
4855    let leading = design.leading_penalty_blocks_before_smooth();
4856    let mut smooth_penalty_ranges = Vec::with_capacity(design.smooth.terms.len());
4857    let mut full_penalty_ranges = Vec::with_capacity(design.smooth.terms.len());
4858    let mut smooth_cursor = 0usize;
4859    for term_idx in 0..design.smooth.terms.len() {
4860        let full_range = design.smooth_term_penalty_range(term_idx)?;
4861        match full_range {
4862            Some(full_range) => {
4863                let local_start = full_range.start.checked_sub(leading).ok_or_else(|| {
4864                    "incremental realizer smooth penalty range precedes the emitted smooth prefix"
4865                        .to_string()
4866                })?;
4867                let local_end = full_range.end.checked_sub(leading).ok_or_else(|| {
4868                    "incremental realizer smooth penalty range precedes the emitted smooth prefix"
4869                        .to_string()
4870                })?;
4871                if local_start != smooth_cursor {
4872                    return Err(format!(
4873                        "incremental realizer non-contiguous emitted smooth layout at term {term_idx}: expected local start {smooth_cursor}, got {local_start}"
4874                    ));
4875                }
4876                smooth_cursor = local_end;
4877                smooth_penalty_ranges.push(local_start..local_end);
4878                full_penalty_ranges.push(full_range);
4879            }
4880            None => {
4881                smooth_penalty_ranges.push(smooth_cursor..smooth_cursor);
4882                let global_cursor = leading.checked_add(smooth_cursor).ok_or_else(|| {
4883                    "incremental realizer empty smooth penalty range overflow".to_string()
4884                })?;
4885                full_penalty_ranges.push(global_cursor..global_cursor);
4886            }
4887        }
4888    }
4889    if smooth_cursor != design.smooth.penalties.len() {
4890        return Err(format!(
4891            "incremental realizer smooth penalty mismatch: ranged={}, actual={}",
4892            smooth_cursor,
4893            design.smooth.penalties.len()
4894        ));
4895    }
4896    Ok((smooth_penalty_ranges, full_penalty_ranges))
4897}
4898
4899impl<'d> FrozenTermCollectionIncrementalRealizer<'d> {
4900    fn new(
4901        data: ArrayView2<'d, f64>,
4902        spec: TermCollectionSpec,
4903        design: TermCollectionDesign,
4904    ) -> Result<Self, String> {
4905        let policy = gam_runtime::resource::ResourcePolicy::default_library();
4906        Self::new_with_policy(data, spec, design, &policy)
4907    }
4908
4909    fn new_with_policy(
4910        data: ArrayView2<'d, f64>,
4911        spec: TermCollectionSpec,
4912        design: TermCollectionDesign,
4913        policy: &gam_runtime::resource::ResourcePolicy,
4914    ) -> Result<Self, String> {
4915        if spec.smooth_terms.len() != design.smooth.terms.len() {
4916            return Err(SmoothError::dimension_mismatch(format!(
4917                "incremental realizer smooth term mismatch: spec_terms={}, design_terms={}",
4918                spec.smooth_terms.len(),
4919                design.smooth.terms.len()
4920            ))
4921            .into());
4922        }
4923
4924        // Cache the exact ranges reported by the emitted global layout. Do not
4925        // reconstruct a second global offset from term specs or coefficient
4926        // blocks: unpenalized fixed/random effects own columns but emit no
4927        // penalty, and multi-penalty smooths own more than one coordinate.
4928        let (smooth_penalty_ranges, full_penalty_ranges) = emitted_smooth_penalty_ranges(&design)?;
4929        let fixed_blocks = build_term_collection_fixed_blocks(data, &spec)
4930            .map_err(|e| format!("failed to cache fixed term-collection blocks: {e}"))?;
4931
4932        let mut dropped_penaltyinfo_by_term = Vec::with_capacity(spec.smooth_terms.len());
4933        for (term_idx, termspec) in spec.smooth_terms.iter().enumerate() {
4934            let realization = build_single_smooth_term_realization_with_policy(
4935                data, termspec, policy,
4936            )
4937            .map_err(|e| {
4938                format!(
4939                    "failed to build cached realization for smooth term '{}' (index {}): {e}",
4940                    termspec.name, term_idx
4941                )
4942            })?;
4943            let expected_cols = design.smooth.terms[term_idx].coeff_range.len();
4944            if realization.design_local.ncols() != expected_cols {
4945                return Err(SmoothError::dimension_mismatch(format!(
4946                    "cached realization width mismatch for term '{}': cached_cols={}, design_cols={}",
4947                    termspec.name,
4948                    realization.design_local.ncols(),
4949                    expected_cols
4950                ))
4951                .into());
4952            }
4953            if realization.active_penalty_count()
4954                != design.smooth.terms[term_idx].active_penalties.len()
4955            {
4956                return Err(SmoothError::dimension_mismatch(format!(
4957                    "cached realization penalty mismatch for term '{}': cached_penalties={}, design_penalties={}",
4958                    termspec.name,
4959                    realization.active_penalty_count(),
4960                    design.smooth.terms[term_idx].active_penalties.len()
4961                ))
4962                .into());
4963            }
4964            dropped_penaltyinfo_by_term.push(realization.dropped_penaltyinfo);
4965        }
4966
4967        let geometry_slots = spec.smooth_terms.len();
4968        Ok(Self {
4969            data,
4970            spec,
4971            design,
4972            fixed_blocks,
4973            dropped_penaltyinfo_by_term,
4974            smooth_penalty_ranges,
4975            full_penalty_ranges,
4976            basisworkspace: gam_terms::basis::BasisWorkspace::with_policy(policy.clone()),
4977            spatial_realization_geometry: vec![None; geometry_slots],
4978            design_revision: 0,
4979        })
4980    }
4981
4982    fn design_revision(&self) -> u64 {
4983        self.design_revision
4984    }
4985
4986    fn spec(&self) -> &TermCollectionSpec {
4987        &self.spec
4988    }
4989
4990    fn design(&self) -> &TermCollectionDesign {
4991        &self.design
4992    }
4993
4994    /// True when this realizer carries exactly ONE spatial smooth term whose
4995    /// frozen basis geometry (`BasisMetadata::Duchon`/`ThinPlate`)
4996    /// admits an EXACT, n-free penalty rebuild at a new length-scale (#1033).
4997    /// The κ-loop fast path gates its design-realization skip on this: the skip
4998    /// leaves `reset_surface` un-run, so it is only sound when `S(ψ_new)` can be
4999    /// re-keyed n-free from the frozen geometry (centers + identifiability
5000    /// transform + operator collocation points), never from the data rows, AND
5001    /// the re-keyed penalty's block topology is IDENTICAL to the one the frozen
5002    /// design carries.
5003    ///
5004    /// Matérn stays on the exact slow re-key path here, but NOT for the reason
5005    /// #1270 originally pinned. The operator-triplet penalty re-key (#1274) IS
5006    /// fully landed: `canonical_penalties_at_psi` and
5007    /// `canonical_penalty_derivatives_at_psi` both rebuild the realized Matérn
5008    /// `{mass, tension, stiffness}` triplet (and its analytic ψ-derivative)
5009    /// n-free from the frozen collocation geometry, routed through the SAME
5010    /// shared `matern_operator_penalty_triplet_at_length_scale` builder the
5011    /// design uses — so the block topology is ψ-stable by construction and the
5012    /// surface is byte-identical to the slow path across the ψ window (pinned
5013    /// to <1e-10 by `matern_nfree_rekey_topology_tests`). The historical
5014    /// "the re-key cannot reproduce the operator triplet" rationale is resolved.
5015    ///
5016    /// Re-admission is nonetheless withheld because it is net-negative on the
5017    /// CURRENT architecture, for two independent reasons the #1274 acceptance
5018    /// gates surface:
5019    ///   1. NO SPEED WIN. Even with the penalty re-keyed, the #1264
5020    ///      reduced-basis-rotation soundness gate (`psi_gram_tensor_covers_skip`)
5021    ///      refuses Matérn's rotating collocation geometry, so the design-
5022    ///      realization skip still falls to the exact O(n) `reset_surface`
5023    ///      re-realization every trial — admitting the penalty rekey alone buys
5024    ///      no n-independence. Closing this needs an n-free re-key of the Matérn
5025    ///      *design* (Chebyshev-in-ψ Gram over the rotating basis), which is the
5026    ///      remaining design-scope work, not a flag flip.
5027    ///   2. QUALITY REGRESSION. Re-admitting Matérn (as #1033 `6a5a2e1` did,
5028    ///      reverted by `feb0eb5`) perturbs the selected fit enough to miss the
5029    ///      mgcv/GP truth-recovery bar (`matern_nu_sweep_*`) — slower AND worse.
5030    ///
5031    /// So Matérn is deliberately "slow-but-right". Duchon/ThinPlate are the
5032    /// #1033 acceptance lane. `matern_nfree_rekey_topology_tests` test (b) pins
5033    /// this negative admission contract: a flip must first re-clear both gates.
5034    fn supports_nfree_penalty_rekey(&self, spatial_terms: &[usize]) -> bool {
5035        if spatial_terms.len() != 1 {
5036            return false;
5037        }
5038        let term_idx = spatial_terms[0];
5039        matches!(
5040            self.design.smooth.terms.get(term_idx).map(|t| &t.metadata),
5041            Some(BasisMetadata::Duchon { .. } | BasisMetadata::ThinPlate { .. })
5042        )
5043    }
5044
5045    /// True when the armed n-free Gaussian lane should suppress exact outer
5046    /// Hessians and route κ search through gradient-only BFGS.
5047    ///
5048    /// This is deliberately narrower than [`Self::supports_nfree_penalty_rekey`]:
5049    /// Matérn has an exact n-free operator-triplet `S(ψ)` re-key (#1274), but its
5050    /// quality gate still depends on the exact second-order outer route. Duchon
5051    /// and ThinPlate are the #1033 n-independent acceptance lane where the exact
5052    /// Hessian slab is the remaining O(n) per-trial cost.
5053    fn supports_nfree_gradient_only_routing(&self, spatial_terms: &[usize]) -> bool {
5054        if spatial_terms.len() != 1 {
5055            return false;
5056        }
5057        let term_idx = spatial_terms[0];
5058        matches!(
5059            self.design.smooth.terms.get(term_idx).map(|t| &t.metadata),
5060            Some(BasisMetadata::Duchon { .. } | BasisMetadata::ThinPlate { .. })
5061        )
5062    }
5063
5064    /// Rebuild the EXACT canonical penalty surface `S(ψ)` at the length-scale
5065    /// implied by `psi`, entirely n-free (#1033). Reuses the FROZEN basis
5066    /// geometry from the single spatial term's `BasisMetadata` (centers,
5067    /// identifiability transform, operator collocation points — all `k × d`, no
5068    /// data rows) and the spec's `(power, nullspace_order, operator_penalties,
5069    /// nu, …)`; only the length-scale moves. The reconstructed term-local
5070    /// penalty matrices replace the `local` of the FROZEN
5071    /// `design.penalties` templates (whose `col_range` / `prior_mean` /
5072    /// `structure_hint` / `op` are ψ-invariant), so the resulting
5073    /// `PenaltySpec`s are bit-identical in topology to the slow path's; running
5074    /// them through the SAME `canonicalize_penalty_specs` pipeline yields the
5075    /// canonical list the kept reference surface must be re-keyed with.
5076    fn canonical_penalties_at_psi(
5077        &mut self,
5078        spatial_terms: &[usize],
5079        psi: &[f64],
5080    ) -> Result<(Vec<gam_terms::construction::CanonicalPenalty>, Vec<usize>), String> {
5081        if spatial_terms.len() != 1 {
5082            return Err(format!(
5083                "n-free penalty re-key requires exactly one spatial term, found {}",
5084                spatial_terms.len()
5085            ));
5086        }
5087        let term_idx = spatial_terms[0];
5088        // Decode ψ with the same chart used by the slow rebuild path. For
5089        // Matérn, per-axis ψ entries are REML hyper-coordinates, so the n-free
5090        // penalty rebuild must consume the trial η contrasts as well as the
5091        // scalar length scale. Duchon keeps η as fixed geometry and continues
5092        // to use frozen metadata below.
5093        let (ls_opt, aniso_from_psi) = spatial_term_psi_to_length_scale_and_aniso(psi);
5094        // Pull the spec-level penalty configuration (which operator orders are
5095        // active / double_penalty) — ψ-invariant, frozen at construction.
5096        let termspec =
5097            self.spec.smooth_terms.get(term_idx).ok_or_else(|| {
5098                format!("spatial term {term_idx} out of range for n-free penalty")
5099            })?;
5100        let term = self
5101            .design
5102            .smooth
5103            .terms
5104            .get(term_idx)
5105            .ok_or_else(|| format!("realized smooth term {term_idx} out of range"))?;
5106        // The per-term penalties live contiguously in the collection penalty
5107        // list at the term's `coeff_range` (single-spatial-term collection).
5108        let p_total = self.design.design.ncols();
5109        let (locals, nullspace_dims): (Vec<Array2<f64>>, Vec<usize>) = match &term.metadata {
5110            BasisMetadata::Duchon {
5111                centers,
5112                identifiability_transform,
5113                operator_collocation_points,
5114                power,
5115                nullspace_order,
5116                aniso_log_scales,
5117                input_scale,
5118                radial_reparam,
5119                ..
5120            } => {
5121                let operator_penalties = match &termspec.basis {
5122                    SmoothBasisSpec::Duchon { spec, .. } => spec.operator_penalties.clone(),
5123                    _ => gam_terms::basis::DuchonOperatorPenaltySpec::default(),
5124                };
5125                // Slow-path Duchon realization stores centers/collocation points
5126                // in standardized coordinates and compensates the user-facing
5127                // length_scale by the scalar input frame before building penalties. The n-free
5128                // re-key must use the same effective length scale, or the fast
5129                // path pairs G(ψ_new) with an S(ψ_new) from a different
5130                // coordinate scale.
5131                let effective_ls =
5132                    ls_opt.map(|length| input_scale.to_standardized_units(length));
5133                gam_terms::basis::duchon_penalties_at_length_scale(
5134                    centers.view(),
5135                    identifiability_transform.as_ref(),
5136                    operator_collocation_points.as_ref().map(|p| p.view()),
5137                    &operator_penalties,
5138                    *power,
5139                    *nullspace_order,
5140                    aniso_log_scales.as_deref(),
5141                    radial_reparam.as_ref(),
5142                    effective_ls,
5143                    &mut self.basisworkspace,
5144                )
5145                .map_err(|e| e.to_string())?
5146            }
5147            BasisMetadata::Matern {
5148                centers,
5149                periodic,
5150                nu,
5151                include_intercept,
5152                identifiability_transform,
5153                aniso_log_scales,
5154                input_scale,
5155                ..
5156            } => {
5157                // `spatial_term_psi_to_length_scale_and_aniso` decodes ψ to a
5158                // length scale in ORIGINAL data coordinates — exactly what the
5159                // slow-path rebuild writes into `spec.length_scale` before
5160                // `matern_operator_penalty_triplet_from_metadata` compensates it
5161                // by the scalar input frame. Compensate identically here so the n-free re-key
5162                // reproduces the slow-path penalty surface byte-for-byte (#706).
5163                let ls = ls_opt.ok_or_else(|| {
5164                    "Matérn n-free penalty re-key requires a finite length-scale".to_string()
5165                })?;
5166                let effective_ls = input_scale.to_standardized_units(ls);
5167                let aniso_for_penalty = aniso_from_psi.as_deref().or(aniso_log_scales.as_deref());
5168                // Route through the SAME canonical operator-triplet builder the
5169                // realized design uses (`matern_operator_penalty_triplet_from_
5170                // metadata`). The Matérn design ALWAYS uses this {mass, tension,
5171                // stiffness} triplet (see the Matérn penalty selection in
5172                // term_specs.rs; #1074 confirmed by MSI measurement that the RKHS
5173                // kernel penalty does not improve recovery and regresses the
5174                // high-frequency guard), so re-keying via the kernel path would
5175                // produce a 1-block surface against a 3-block frozen design — the
5176                // topology desync #1270 hard-errored on. Sharing the builder
5177                // makes the block count ψ-stable by construction.
5178                let filtered = matern_operator_penalty_triplet_at_length_scale(
5179                    centers.view(),
5180                    periodic.as_deref(),
5181                    identifiability_transform.as_ref(),
5182                    *nu,
5183                    *include_intercept,
5184                    aniso_for_penalty,
5185                    effective_ls,
5186                )
5187                .map_err(|e| e.to_string())?;
5188                let locals = filtered
5189                    .active
5190                    .iter()
5191                    .map(|penalty| penalty.matrix.clone())
5192                    .collect();
5193                let nullspace_dims = filtered
5194                    .active
5195                    .iter()
5196                    .map(|penalty| penalty.nullity)
5197                    .collect();
5198                (locals, nullspace_dims)
5199            }
5200            BasisMetadata::ThinPlate {
5201                centers,
5202                identifiability_transform,
5203                radial_reparam,
5204                ..
5205            } => {
5206                let ls = ls_opt.ok_or_else(|| {
5207                    "thin-plate n-free penalty re-key requires a finite length-scale".to_string()
5208                })?;
5209                let double_penalty = match &termspec.basis {
5210                    SmoothBasisSpec::ThinPlate { spec, .. } => spec.double_penalty,
5211                    _ => false,
5212                };
5213                gam_terms::basis::thin_plate_penalties_at_length_scale(
5214                    centers.view(),
5215                    identifiability_transform.as_ref(),
5216                    radial_reparam.as_ref(),
5217                    ls,
5218                    double_penalty,
5219                    &mut self.basisworkspace,
5220                )
5221                .map_err(|e| e.to_string())?
5222            }
5223            other => {
5224                return Err(format!(
5225                    "n-free penalty re-key unsupported for basis metadata {:?}",
5226                    std::mem::discriminant(other)
5227                ));
5228            }
5229        };
5230        // The frozen collection penalties for THIS term are the templates whose
5231        // ψ-invariant structure (col_range / prior_mean / structure_hint / op)
5232        // we keep, swapping only the numeric `local`. For a single-spatial-term
5233        // collection the term owns the whole penalty list.
5234        let templates = &self.design.penalties;
5235        if templates.len() != locals.len() {
5236            return Err(format!(
5237                "n-free penalty re-key produced {} blocks but the frozen design carries {} \
5238                 — penalty topology is not ψ-stable",
5239                locals.len(),
5240                templates.len()
5241            ));
5242        }
5243        let specs: Vec<gam_solve::estimate::PenaltySpec> = templates
5244            .iter()
5245            .zip(locals.into_iter())
5246            .map(|(tmpl, local)| gam_solve::estimate::PenaltySpec::Block {
5247                local,
5248                col_range: tmpl.col_range.clone(),
5249                prior_mean: tmpl.prior_mean.clone(),
5250                structure_hint: tmpl.structure_hint.clone(),
5251                op: tmpl.op.clone(),
5252            })
5253            .collect();
5254        gam_terms::construction::canonicalize_penalty_specs(
5255            &specs,
5256            &nullspace_dims,
5257            p_total,
5258            "nfree-psi-penalty",
5259        )
5260        .map_err(|e| e.to_string())
5261    }
5262
5263    fn canonical_penalty_derivatives_at_psi(
5264        &mut self,
5265        spatial_terms: &[usize],
5266        psi: &[f64],
5267    ) -> Result<(Range<usize>, usize, Vec<Array2<f64>>), String> {
5268        if spatial_terms.len() != 1 {
5269            return Err(format!(
5270                "n-free penalty derivative re-key requires exactly one spatial term, found {}",
5271                spatial_terms.len()
5272            ));
5273        }
5274        let term_idx = spatial_terms[0];
5275        let (ls_opt, aniso_from_psi) = spatial_term_psi_to_length_scale_and_aniso(psi);
5276        let termspec = self.spec.smooth_terms.get(term_idx).ok_or_else(|| {
5277            format!("spatial term {term_idx} out of range for n-free penalty derivative")
5278        })?;
5279        let term = self
5280            .design
5281            .smooth
5282            .terms
5283            .get(term_idx)
5284            .ok_or_else(|| format!("realized smooth term {term_idx} out of range"))?;
5285        let p_total = self.design.design.ncols();
5286        let smooth_start = p_total.saturating_sub(self.design.smooth.total_smooth_cols());
5287        let global_range =
5288            (smooth_start + term.coeff_range.start)..(smooth_start + term.coeff_range.end);
5289
5290        let locals = match &term.metadata {
5291            BasisMetadata::Duchon {
5292                centers,
5293                identifiability_transform,
5294                operator_collocation_points,
5295                power,
5296                nullspace_order,
5297                aniso_log_scales,
5298                input_scale,
5299                radial_reparam,
5300                ..
5301            } => {
5302                let mut spec = match &termspec.basis {
5303                    SmoothBasisSpec::Duchon { spec, .. } => spec.clone(),
5304                    _ => {
5305                        return Err(
5306                            "Duchon n-free penalty derivative requires a Duchon term spec"
5307                                .to_string(),
5308                        );
5309                    }
5310                };
5311                let effective_ls =
5312                    ls_opt.map(|length| input_scale.to_standardized_units(length));
5313                spec.length_scale = effective_ls;
5314                spec.power = *power;
5315                spec.nullspace_order = *nullspace_order;
5316                spec.aniso_log_scales = aniso_log_scales.clone();
5317                // #1355: replay the frozen data-metric reparam so the n-free
5318                // penalty ψ-derivative matches the rotated forward penalty.
5319                spec.radial_reparam = radial_reparam.clone();
5320                if spec.length_scale.is_none() {
5321                    return Err(
5322                        "Duchon n-free penalty derivative requires a hybrid length-scale"
5323                            .to_string(),
5324                    );
5325                }
5326                let collocation = operator_collocation_points
5327                    .as_ref()
5328                    .map(|points| points.view())
5329                    .unwrap_or_else(|| centers.view());
5330                let (_native_sources, mut first, _native_second) =
5331                    gam_terms::basis::build_duchon_native_penalty_psi_derivatives(
5332                        centers.view(),
5333                        &spec,
5334                        identifiability_transform.as_ref(),
5335                        &mut self.basisworkspace,
5336                    )
5337                    .map_err(|e| e.to_string())?;
5338                let (_operator_sources, operator_first, _operator_second) =
5339                    gam_terms::basis::build_duchon_operator_penalty_psi_derivatives(
5340                        collocation,
5341                        centers.view(),
5342                        &spec,
5343                        identifiability_transform.as_ref(),
5344                        &mut self.basisworkspace,
5345                    )
5346                    .map_err(|e| e.to_string())?;
5347                first.extend(operator_first);
5348                first
5349            }
5350            BasisMetadata::Matern {
5351                centers,
5352                periodic,
5353                nu,
5354                include_intercept,
5355                identifiability_transform,
5356                aniso_log_scales,
5357                input_scale,
5358                ..
5359            } => {
5360                let ls = ls_opt.ok_or_else(|| {
5361                    "Matérn n-free penalty derivative requires a finite length-scale".to_string()
5362                })?;
5363                let effective_ls = input_scale.to_standardized_units(ls);
5364                let penalty_centers = gam_terms::basis::expand_periodic_centers(
5365                    &centers.to_owned(),
5366                    periodic.as_deref(),
5367                )
5368                .map_err(|e| e.to_string())?;
5369                let aniso_for_penalty = aniso_from_psi.as_deref().or(aniso_log_scales.as_deref());
5370                let (first, _second) =
5371                    gam_terms::basis::build_matern_operator_penalty_psi_derivatives(
5372                        penalty_centers.view(),
5373                        effective_ls,
5374                        *nu,
5375                        *include_intercept,
5376                        identifiability_transform.as_ref(),
5377                        aniso_for_penalty,
5378                    )
5379                    .map_err(|e| e.to_string())?;
5380                first
5381            }
5382            BasisMetadata::ThinPlate {
5383                centers,
5384                identifiability_transform,
5385                radial_reparam,
5386                ..
5387            } => {
5388                let ls = ls_opt.ok_or_else(|| {
5389                    "thin-plate n-free penalty derivative requires a finite length-scale"
5390                        .to_string()
5391                })?;
5392                let mut spec = match &termspec.basis {
5393                    SmoothBasisSpec::ThinPlate { spec, .. } => spec.clone(),
5394                    _ => {
5395                        return Err(
5396                            "thin-plate n-free penalty derivative requires a ThinPlate term spec"
5397                                .to_string(),
5398                        );
5399                    }
5400                };
5401                spec.length_scale = ls;
5402                if spec.radial_reparam.is_none() {
5403                    spec.radial_reparam = radial_reparam.clone();
5404                }
5405                let (primary, _primary_second, nullspace, _nullspace_second) =
5406                    gam_terms::basis::build_thin_plate_penalty_psi_derivativeswithworkspace(
5407                        centers.view(),
5408                        &spec,
5409                        identifiability_transform.as_ref(),
5410                        &mut self.basisworkspace,
5411                    )
5412                    .map_err(|e| e.to_string())?;
5413                if self.design.penalties.len() > 1 {
5414                    vec![primary, nullspace]
5415                } else {
5416                    vec![primary]
5417                }
5418            }
5419            other => {
5420                return Err(format!(
5421                    "n-free penalty derivative re-key unsupported for basis metadata {:?}",
5422                    std::mem::discriminant(other)
5423                ));
5424            }
5425        };
5426        if locals.len() != self.design.penalties.len() {
5427            return Err(format!(
5428                "n-free penalty derivative re-key produced {} blocks but the frozen design carries {} \
5429                 — penalty topology is not ψ-stable",
5430                locals.len(),
5431                self.design.penalties.len()
5432            ));
5433        }
5434        Ok((global_range, p_total, locals))
5435    }
5436
5437    fn apply_log_kappa(
5438        &mut self,
5439        log_kappa: &SpatialLogKappaCoords,
5440        term_indices: &[usize],
5441    ) -> Result<(), String> {
5442        if term_indices.len() != log_kappa.dims_per_term().len() {
5443            return Err(SmoothError::dimension_mismatch(format!(
5444                "incremental realizer log-kappa term mismatch: term_indices={}, dims_per_term={}",
5445                term_indices.len(),
5446                log_kappa.dims_per_term().len()
5447            ))
5448            .into());
5449        }
5450
5451        let mut any_changed = false;
5452        for (slot, &term_idx) in term_indices.iter().enumerate() {
5453            any_changed |= self.apply_log_kappa_to_term(term_idx, log_kappa.term_slice(slot))?;
5454        }
5455
5456        if any_changed {
5457            self.refresh_full_design_operator()?;
5458            rebuild_smooth_auxiliary_state(
5459                &mut self.design.smooth,
5460                &self.dropped_penaltyinfo_by_term,
5461            )?;
5462            rebuild_term_collection_auxiliary_state(&self.spec, &mut self.design)?;
5463            self.design_revision = self.design_revision.wrapping_add(1);
5464        }
5465        Ok(())
5466    }
5467
5468    fn apply_log_kappa_to_term(&mut self, term_idx: usize, psi: &[f64]) -> Result<bool, String> {
5469        if !spatial_term_supports_hyper_optimization(&self.spec, term_idx) {
5470            return Err(SmoothError::invalid_config(format!(
5471                "incremental realizer term {term_idx} does not expose spatial hyperparameters"
5472            ))
5473            .into());
5474        }
5475        // Measure-jet ψ slots are dial coordinates, not log-κ (dial docs:
5476        // the MEASURE_JET_PSI_* bounds block); route through the dial setter
5477        // so the κ-translation below never misreads them as log-scales.
5478        let measure_jet_term = measure_jet_term_spec(&self.spec, term_idx).is_some();
5479        // Constant-curvature ψ is the raw signed curvature κ, NOT a log-scale;
5480        // route through the κ setter so `spatial_term_psi_to_length_scale_and_aniso`
5481        // never misreads it (and never hits the "no length scale" rejection).
5482        let constant_curvature_term = constant_curvature_term_spec(&self.spec, term_idx).is_some();
5483        let mut next_length_scale = None;
5484        let mut next_aniso: Option<Vec<f64>> = None;
5485        if measure_jet_term {
5486            if !set_measure_jet_psi_dials(&mut self.spec, term_idx, psi)
5487                .map_err(|e| e.to_string())?
5488            {
5489                return Ok(false);
5490            }
5491        } else if constant_curvature_term {
5492            if !set_constant_curvature_kappa(&mut self.spec, term_idx, psi)
5493                .map_err(|e| e.to_string())?
5494            {
5495                return Ok(false);
5496            }
5497        } else {
5498            let current_length_scale = get_spatial_length_scale(&self.spec, term_idx);
5499            let current_aniso = get_spatial_aniso_log_scales(&self.spec, term_idx);
5500            let (ls, eta) = spatial_term_psi_to_length_scale_and_aniso(psi);
5501            next_length_scale = ls;
5502            next_aniso = eta;
5503            let same_length = spatial_length_scale_matches(current_length_scale, next_length_scale);
5504            let same_aniso = spatial_aniso_matches(current_aniso.as_deref(), next_aniso.as_deref());
5505            if same_length && same_aniso {
5506                return Ok(false);
5507            }
5508            if let Some(length_scale) = next_length_scale {
5509                set_spatial_length_scale(&mut self.spec, term_idx, length_scale)
5510                    .map_err(|e| e.to_string())?;
5511            }
5512            if let Some(eta) = next_aniso.clone() {
5513                set_spatial_aniso_log_scales(&mut self.spec, term_idx, eta)
5514                    .map_err(|e| e.to_string())?;
5515            }
5516        }
5517
5518        // Pick the spec to drive the rebuild. If the per-term geometry cache
5519        // is populated, it carries already-resolved centers
5520        // (`CenterStrategy::UserProvided`) and frozen `input_scale`; reusing
5521        // it short-circuits `select_centers_by_strategy` (KMeans /
5522        // FarthestPoint / EqualMass cluster searches) and
5523        // isotropic scale estimation over n rows in
5524        // the family builders. Centers in the cached spec live in
5525        // standardized coordinates (matching the cached `input_scale`), so
5526        // the same standardization + kernel path runs without recomputation
5527        // of the geometry.
5528        let geometry_slot = self
5529            .spatial_realization_geometry
5530            .get(term_idx)
5531            .ok_or_else(|| format!("incremental realizer geometry slot {term_idx} out of range"))?;
5532        let mut build_spec = match geometry_slot {
5533            Some(cached) => cached.clone(),
5534            None => self
5535                .spec
5536                .smooth_terms
5537                .get(term_idx)
5538                .ok_or_else(|| format!("incremental realizer smooth term {term_idx} out of range"))?
5539                .clone(),
5540        };
5541        if measure_jet_term {
5542            // The cached build spec carries the frozen geometry (UserProvided
5543            // barycenter nodes, frozen quadrature + transform); only the
5544            // dials move per trial.
5545            set_single_term_measure_jet_psi_dials(&mut build_spec, psi)
5546                .map_err(|e| e.to_string())?;
5547        } else if constant_curvature_term {
5548            // The cached build spec carries the κ-fixed geometry (UserProvided
5549            // centers, frozen ℓ and constraint transform); only κ moves per
5550            // trial, written through the raw-κ setter to match the collection
5551            // write-back above.
5552            set_single_term_constant_curvature_kappa(&mut build_spec, psi)
5553                .map_err(|e| e.to_string())?;
5554        } else {
5555            if let Some(length_scale) = next_length_scale {
5556                set_single_term_spatial_length_scale(&mut build_spec, length_scale)
5557                    .map_err(|e| e.to_string())?;
5558            }
5559            if let Some(eta) = next_aniso {
5560                set_single_term_spatial_aniso_log_scales(&mut build_spec, eta)
5561                    .map_err(|e| e.to_string())?;
5562            }
5563        }
5564
5565        let termname = build_spec.name.clone();
5566        let local = build_single_local_smooth_term(
5567            self.data,
5568            &build_spec,
5569            &mut self.basisworkspace,
5570        )
5571        .map_err(|e| {
5572            format!(
5573                "failed to rebuild smooth term '{termname}' during incremental κ realization: {e}"
5574            )
5575        })?;
5576
5577        // Populate the geometry cache from the realized metadata on first use.
5578        // Family auto-promotion (ThinPlate -> Duchon) is detected as a basis /
5579        // metadata mismatch in `freeze_geometry_from_metadata` and leaves the
5580        // cache empty so the next call re-tries with the (now stable) family.
5581        if self.spatial_realization_geometry[term_idx].is_none()
5582            && let Some(frozen) = freeze_geometry_from_metadata(&build_spec, &local.metadata)
5583        {
5584            // Mirror the frozen identifiability (pinned `Z` + double-penalty
5585            // nullspace-shrinkage decision, #787/#860/#1122) back onto the
5586            // collection spec the analytic ψ-gradient reads
5587            // (`try_build_spatial_log_kappa_hyper_dirs(self.spec(), …)`). The
5588            // value rebuild consumes the cached `build_spec`, so without this
5589            // copy the gradient would keep re-running the κ-DEPENDENT spectral
5590            // test on the un-frozen collection spec while the value uses the
5591            // frozen decision — re-introducing the very objective↔gradient
5592            // desync the freeze removes. Pinning both to the same frozen
5593            // transform keeps the per-trial value and its analytic gradient on
5594            // one fixed `Z` and one fixed null dimension `r`.
5595            if let (
5596                SmoothBasisSpec::Matern {
5597                    spec: frozen_spec, ..
5598                },
5599                Some(SmoothBasisSpec::Matern {
5600                    spec: live_spec, ..
5601                }),
5602            ) = (
5603                &frozen.basis,
5604                self.spec
5605                    .smooth_terms
5606                    .get_mut(term_idx)
5607                    .map(|t| &mut t.basis),
5608            ) {
5609                live_spec.identifiability = frozen_spec.identifiability.clone();
5610                live_spec.center_strategy = frozen_spec.center_strategy.clone();
5611            }
5612            self.spatial_realization_geometry[term_idx] = Some(frozen);
5613        }
5614
5615        let realization = wrap_local_build_as_realization(local, &build_spec)?;
5616        self.replace_term_realization(term_idx, realization)?;
5617        Ok(true)
5618    }
5619
5620    fn replace_term_realization(
5621        &mut self,
5622        term_idx: usize,
5623        realization: SingleSmoothTermRealization,
5624    ) -> Result<(), String> {
5625        let t_replace = std::time::Instant::now();
5626        let SingleSmoothTermRealization {
5627            design_local,
5628            term,
5629            dropped_penaltyinfo,
5630        } = realization;
5631        let SmoothTerm {
5632            name,
5633            active_penalties,
5634            dropped_penalties,
5635            metadata,
5636            lower_bounds_local,
5637            linear_constraints_local,
5638            joint_null_rotation,
5639            ..
5640        } = term;
5641        let coeff_range = self
5642            .design
5643            .smooth
5644            .terms
5645            .get(term_idx)
5646            .ok_or_else(|| format!("incremental realizer smooth term {term_idx} out of range"))?
5647            .coeff_range
5648            .clone();
5649        if design_local.ncols() != coeff_range.len() {
5650            return Err(SmoothError::dimension_mismatch(format!(
5651                "incremental realizer width mismatch for term {}: rebuilt_cols={}, cached_cols={}",
5652                term_idx,
5653                design_local.ncols(),
5654                coeff_range.len()
5655            ))
5656            .into());
5657        }
5658        if design_local.nrows() != self.design.design.nrows() {
5659            return Err(SmoothError::dimension_mismatch(format!(
5660                "incremental realizer row mismatch for term {}: rebuilt_rows={}, design_rows={}",
5661                term_idx,
5662                design_local.nrows(),
5663                self.design.design.nrows()
5664            ))
5665            .into());
5666        }
5667
5668        let smooth_penalty_range = self
5669            .smooth_penalty_ranges
5670            .get(term_idx)
5671            .ok_or_else(|| {
5672                format!("incremental realizer missing smooth penalty range for term {term_idx}")
5673            })?
5674            .clone();
5675        let full_penalty_range = self
5676            .full_penalty_ranges
5677            .get(term_idx)
5678            .ok_or_else(|| {
5679                format!("incremental realizer missing full penalty range for term {term_idx}")
5680            })?
5681            .clone();
5682        if active_penalties.len() != smooth_penalty_range.len() {
5683            return Err(SmoothError::dimension_mismatch(format!(
5684                "incremental realizer topology changed for term '{}': active_penalties={}, cached_penalties={}",
5685                name,
5686                active_penalties.len(),
5687                smooth_penalty_range.len()
5688            ))
5689            .into());
5690        }
5691
5692        self.design.smooth.term_designs[term_idx] = design_local;
5693
5694        for (offset, active_penalty) in active_penalties.iter().enumerate() {
5695            let smooth_penalty_idx = smooth_penalty_range.start + offset;
5696            let full_penalty_idx = full_penalty_range.start + offset;
5697            let penalty_local = &active_penalty.matrix;
5698
5699            if penalty_local.nrows() != coeff_range.len()
5700                || penalty_local.ncols() != coeff_range.len()
5701            {
5702                return Err(SmoothError::dimension_mismatch(format!(
5703                    "incremental realizer penalty shape mismatch for term '{}' penalty {}: \
5704                     penalty is {}x{} but coeff_range has {} columns",
5705                    name,
5706                    offset,
5707                    penalty_local.nrows(),
5708                    penalty_local.ncols(),
5709                    coeff_range.len()
5710                ))
5711                .into());
5712            }
5713
5714            let smooth_penalty = self
5715                .design
5716                .smooth
5717                .penalties
5718                .get_mut(smooth_penalty_idx)
5719                .ok_or_else(|| {
5720                    format!(
5721                        "incremental realizer smooth penalty {} out of range for term {}",
5722                        smooth_penalty_idx, term_idx
5723                    )
5724                })?;
5725            // With per-term block-local penalties, col_range already targets
5726            // this specific term, so .local is p_k × p_k.
5727            smooth_penalty.local.assign(penalty_local);
5728            smooth_penalty.op = active_penalty.op.clone();
5729
5730            let full_bp = self
5731                .design
5732                .penalties
5733                .get_mut(full_penalty_idx)
5734                .ok_or_else(|| {
5735                    format!(
5736                        "incremental realizer full penalty {} out of range for term {}",
5737                        full_penalty_idx, term_idx
5738                    )
5739                })?;
5740            // With per-term block-local penalties, col_range already targets
5741            // this specific term, so .local is p_k × p_k.
5742            full_bp.local.assign(penalty_local);
5743            full_bp.op = active_penalty.op.clone();
5744
5745            self.design.smooth.nullspace_dims[smooth_penalty_idx] = active_penalty.nullity;
5746            self.design.nullspace_dims[full_penalty_idx] = active_penalty.nullity;
5747
5748            self.design.smooth.penaltyinfo[smooth_penalty_idx].global_index = smooth_penalty_idx;
5749            self.design.smooth.penaltyinfo[smooth_penalty_idx].termname = Some(name.clone());
5750            self.design.smooth.penaltyinfo[smooth_penalty_idx].penalty =
5751                active_penalty.info.clone();
5752
5753            self.design.penaltyinfo[full_penalty_idx].global_index = full_penalty_idx;
5754            self.design.penaltyinfo[full_penalty_idx].termname = Some(name.clone());
5755            self.design.penaltyinfo[full_penalty_idx].penalty = active_penalty.info.clone();
5756        }
5757
5758        let target_term = self.design.smooth.terms.get_mut(term_idx).ok_or_else(|| {
5759            format!("incremental realizer smooth term {term_idx} disappeared during replacement")
5760        })?;
5761        target_term.active_penalties = active_penalties;
5762        target_term.dropped_penalties = dropped_penalties;
5763        target_term.metadata = metadata;
5764        target_term.lower_bounds_local = lower_bounds_local;
5765        target_term.linear_constraints_local = linear_constraints_local;
5766        target_term.joint_null_rotation = joint_null_rotation;
5767        self.dropped_penaltyinfo_by_term[term_idx] = dropped_penaltyinfo;
5768        log::info!(
5769            "[STAGE] smooth basis rebuild (term {}, '{}', cols={}): {:.3}s",
5770            term_idx,
5771            target_term.name,
5772            coeff_range.len(),
5773            t_replace.elapsed().as_secs_f64(),
5774        );
5775        Ok(())
5776    }
5777
5778    fn refresh_full_design_operator(&mut self) -> Result<(), String> {
5779        let mut blocks = Vec::<DesignBlock>::with_capacity(
5780            self.fixed_blocks.len() + self.design.smooth.term_designs.len(),
5781        );
5782        blocks.extend(self.fixed_blocks.iter().cloned());
5783        for term_design in &self.design.smooth.term_designs {
5784            blocks.push(DesignBlock::from(term_design));
5785        }
5786        self.design.design = assemble_term_collection_design_matrix(blocks)
5787            .map_err(|e| format!("failed to refresh term-collection design: {e}"))?;
5788        Ok(())
5789    }
5790}
5791
5792fn build_term_collection_fixed_blocks(
5793    data: ArrayView2<'_, f64>,
5794    spec: &TermCollectionSpec,
5795) -> Result<Vec<DesignBlock>, BasisError> {
5796    let mut blocks = Vec::<DesignBlock>::new();
5797    if !term_collection_has_anchored_bspline(spec) {
5798        blocks.push(DesignBlock::Intercept(data.nrows()));
5799    }
5800
5801    if !spec.linear_terms.is_empty() {
5802        let mut linear_block = Array2::<f64>::zeros((data.nrows(), spec.linear_terms.len()));
5803        for (j, linear) in spec.linear_terms.iter().enumerate() {
5804            // Single shared realizer: numeric product gated by any
5805            // categorical-level indicators (factor-aware `:` interaction),
5806            // mirroring `build_term_collection_design_inner`.
5807            let column = linear
5808                .realized_design_column(data)
5809                .map_err(BasisError::InvalidInput)?;
5810            linear_block.column_mut(j).assign(&column);
5811        }
5812        blocks.push(DesignBlock::Dense(
5813            gam_linalg::matrix::DenseDesignMatrix::from(linear_block),
5814        ));
5815    }
5816
5817    for term in &spec.random_effect_terms {
5818        let block = build_random_effect_block(data, term)?;
5819        let re_op = RandomEffectOperator::new(block.group_ids, block.num_groups);
5820        blocks.push(DesignBlock::RandomEffect(Arc::new(re_op)));
5821    }
5822
5823    Ok(blocks)
5824}
5825
5826// ---------------------------------------------------------------------------
5827// N-block spatial length-scale optimizer.
5828// ---------------------------------------------------------------------------
5829
5830pub struct SpatialLengthScaleOptimizationResult<FitOut> {
5831    pub resolved_specs: Vec<TermCollectionSpec>,
5832    pub designs: Vec<TermCollectionDesign>,
5833    pub fit: FitOut,
5834    pub certified_outer: Option<gam_solve::rho_optimizer::CertifiedOuterResult>,
5835    pub timing: Option<SpatialLengthScaleOptimizationTiming>,
5836}
5837
5838/// One exact outer-objective evaluation together with the owned coefficient
5839/// mode that produced it.
5840///
5841/// `mode` is deliberately generic and move-only.  The spatial driver never
5842/// interprets or clones it; it retains the carrier that belongs to the latest
5843/// successful evaluation and transfers that exact ownership into final fit
5844/// assembly after the outer certificate has been issued.
5845pub struct ExactJointEvaluation<M> {
5846    pub objective: f64,
5847    pub gradient: Array1<f64>,
5848    pub hessian: gam_problem::HessianValue,
5849    pub mode: M,
5850}
5851
5852/// One exact fixed-point evaluation and the owned coefficient mode that
5853/// produced its value and update equations.
5854pub struct ExactJointEfsEvaluation<M> {
5855    pub evaluation: gam_problem::EfsEval,
5856    pub mode: M,
5857}
5858
5859pub enum SpatialFitProvenance<'a, M> {
5860    NoOuterOptimization,
5861    Certified {
5862        outer: &'a gam_solve::rho_optimizer::CertifiedOuterResult,
5863        mode: M,
5864    },
5865}
5866
5867/// Exact-joint hyper-parameter setup for N-block spatial length-scale optimization.
5868#[derive(Debug, Clone)]
5869pub struct ExactJointHyperSetup {
5870    rho0: Array1<f64>,
5871    rho_lower: Array1<f64>,
5872    rho_upper: Array1<f64>,
5873    log_kappa0: SpatialLogKappaCoords,
5874    log_kappa_lower: SpatialLogKappaCoords,
5875    log_kappa_upper: SpatialLogKappaCoords,
5876    auxiliary0: Array1<f64>,
5877    auxiliary_lower: Array1<f64>,
5878    auxiliary_upper: Array1<f64>,
5879}
5880
5881impl ExactJointHyperSetup {
5882    fn sanitize_rho_seed(
5883        rho0: Array1<f64>,
5884        rho_lower: &Array1<f64>,
5885        rho_upper: &Array1<f64>,
5886    ) -> Array1<f64> {
5887        Array1::from_iter(rho0.iter().enumerate().map(|(idx, &value)| {
5888            let lo = rho_lower[idx];
5889            let hi = rho_upper[idx];
5890            let fallback = 0.0_f64.clamp(lo, hi);
5891            if value.is_finite() {
5892                value.clamp(lo, hi)
5893            } else {
5894                fallback
5895            }
5896        }))
5897    }
5898
5899    pub(crate) fn new(
5900        rho0: Array1<f64>,
5901        rho_lower: Array1<f64>,
5902        rho_upper: Array1<f64>,
5903        log_kappa0: SpatialLogKappaCoords,
5904        log_kappa_lower: SpatialLogKappaCoords,
5905        log_kappa_upper: SpatialLogKappaCoords,
5906    ) -> Self {
5907        let rho0 = Self::sanitize_rho_seed(rho0, &rho_lower, &rho_upper);
5908        Self {
5909            rho0,
5910            rho_lower,
5911            rho_upper,
5912            log_kappa0,
5913            log_kappa_lower,
5914            log_kappa_upper,
5915            auxiliary0: Array1::zeros(0),
5916            auxiliary_lower: Array1::zeros(0),
5917            auxiliary_upper: Array1::zeros(0),
5918        }
5919    }
5920
5921    pub(crate) fn with_auxiliary(
5922        mut self,
5923        auxiliary0: Array1<f64>,
5924        auxiliary_lower: Array1<f64>,
5925        auxiliary_upper: Array1<f64>,
5926    ) -> Self {
5927        assert_eq!(
5928            auxiliary0.len(),
5929            auxiliary_lower.len(),
5930            "auxiliary lower bound length mismatch"
5931        );
5932        assert_eq!(
5933            auxiliary0.len(),
5934            auxiliary_upper.len(),
5935            "auxiliary upper bound length mismatch"
5936        );
5937        self.auxiliary0 = Self::sanitize_rho_seed(auxiliary0, &auxiliary_lower, &auxiliary_upper);
5938        self.auxiliary_lower = auxiliary_lower;
5939        self.auxiliary_upper = auxiliary_upper;
5940        self
5941    }
5942
5943    pub(crate) fn rho_dim(&self) -> usize {
5944        self.rho0.len()
5945    }
5946
5947    pub(crate) fn log_kappa_dim(&self) -> usize {
5948        self.log_kappa0.len()
5949    }
5950
5951    pub(crate) fn auxiliary_dim(&self) -> usize {
5952        self.auxiliary0.len()
5953    }
5954
5955    pub(crate) fn theta0(&self) -> Array1<f64> {
5956        let mut out =
5957            Array1::<f64>::zeros(self.rho_dim() + self.log_kappa_dim() + self.auxiliary_dim());
5958        out.slice_mut(s![..self.rho_dim()]).assign(&self.rho0);
5959        out.slice_mut(s![self.rho_dim()..self.rho_dim() + self.log_kappa_dim()])
5960            .assign(self.log_kappa0.as_array());
5961        out.slice_mut(s![self.rho_dim() + self.log_kappa_dim()..])
5962            .assign(&self.auxiliary0);
5963        out
5964    }
5965
5966    pub(crate) fn lower(&self) -> Array1<f64> {
5967        let mut out =
5968            Array1::<f64>::zeros(self.rho_dim() + self.log_kappa_dim() + self.auxiliary_dim());
5969        out.slice_mut(s![..self.rho_dim()]).assign(&self.rho_lower);
5970        out.slice_mut(s![self.rho_dim()..self.rho_dim() + self.log_kappa_dim()])
5971            .assign(self.log_kappa_lower.as_array());
5972        out.slice_mut(s![self.rho_dim() + self.log_kappa_dim()..])
5973            .assign(&self.auxiliary_lower);
5974        out
5975    }
5976
5977    pub(crate) fn upper(&self) -> Array1<f64> {
5978        let mut out =
5979            Array1::<f64>::zeros(self.rho_dim() + self.log_kappa_dim() + self.auxiliary_dim());
5980        out.slice_mut(s![..self.rho_dim()]).assign(&self.rho_upper);
5981        out.slice_mut(s![self.rho_dim()..self.rho_dim() + self.log_kappa_dim()])
5982            .assign(self.log_kappa_upper.as_array());
5983        out.slice_mut(s![self.rho_dim() + self.log_kappa_dim()..])
5984            .assign(&self.auxiliary_upper);
5985        out
5986    }
5987
5988    /// Per-term dimensionality layout for the psi block.
5989    pub(crate) fn log_kappa_dims_per_term(&self) -> Vec<usize> {
5990        self.log_kappa0.dims_per_term().to_vec()
5991    }
5992}
5993
5994/// N-block design cache for exact-joint spatial length-scale optimization.
5995///
5996/// Each block owns a `FrozenTermCollectionIncrementalRealizer` and a list of
5997/// spatial term indices within that block's spec. The cache splits the
5998/// combined psi vector into per-block slices using precomputed offsets.
5999struct ExactJointDesignCache<'d> {
6000    realizers: Vec<FrozenTermCollectionIncrementalRealizer<'d>>,
6001    block_term_indices: Vec<Vec<usize>>,
6002    current_theta: Option<Array1<f64>>,
6003    last_cost: Option<f64>,
6004    last_eval: Option<(f64, Array1<f64>, gam_problem::HessianValue)>,
6005    rho_dim: usize,
6006    all_dims: Vec<usize>,
6007    log_kappa_dim: usize,
6008    block_term_counts: Vec<usize>,
6009}
6010
6011impl<'d> ExactJointDesignCache<'d> {
6012    fn new(
6013        data: ArrayView2<'d, f64>,
6014        blocks: Vec<(TermCollectionSpec, TermCollectionDesign, Vec<usize>)>,
6015        rho_dim: usize,
6016        all_dims: Vec<usize>,
6017    ) -> Result<Self, String> {
6018        let n_blocks = blocks.len();
6019        let mut realizers = Vec::with_capacity(n_blocks);
6020        let mut block_term_indices = Vec::with_capacity(n_blocks);
6021        let mut block_term_counts = Vec::with_capacity(n_blocks);
6022
6023        for (spec, design, terms) in blocks {
6024            block_term_counts.push(terms.len());
6025            block_term_indices.push(terms);
6026            realizers.push(FrozenTermCollectionIncrementalRealizer::new(
6027                data, spec, design,
6028            )?);
6029        }
6030
6031        Ok(Self {
6032            realizers,
6033            block_term_indices,
6034            current_theta: None,
6035            last_cost: None,
6036            last_eval: None,
6037            rho_dim,
6038            log_kappa_dim: all_dims.iter().sum(),
6039            all_dims,
6040            block_term_counts,
6041        })
6042    }
6043
6044    fn ensure_theta(&mut self, theta: &Array1<f64>) -> Result<(), String> {
6045        if self
6046            .current_theta
6047            .as_ref()
6048            .is_some_and(|cached| theta_values_match(cached, theta))
6049        {
6050            return Ok(());
6051        }
6052
6053        let t_ensure = std::time::Instant::now();
6054        let kappa_theta_len = self.rho_dim + self.log_kappa_dim;
6055        if theta.len() < kappa_theta_len {
6056            return Err(SmoothError::dimension_mismatch(format!(
6057                "exact-joint theta length mismatch: got {}, expected at least {} (rho_dim={}, log_kappa_dim={})",
6058                theta.len(),
6059                kappa_theta_len,
6060                self.rho_dim,
6061                self.log_kappa_dim
6062            ))
6063            .into());
6064        }
6065        let theta_kappa = theta.slice(s![..kappa_theta_len]).to_owned();
6066        let full_log_kappa = SpatialLogKappaCoords::from_theta_tail_with_dims(
6067            &theta_kappa,
6068            self.rho_dim,
6069            self.all_dims.clone(),
6070        );
6071
6072        // Split the full log_kappa into per-block sub-coords using split_at.
6073        // We split from the front iteratively: after extracting block 0..N-2,
6074        // the remainder is the last block.
6075        let n = self.realizers.len();
6076        let mut remaining = full_log_kappa;
6077        for block_idx in 0..n {
6078            let count = self.block_term_counts[block_idx];
6079            if block_idx < n - 1 {
6080                let (block_lk, rest) = remaining.split_at(count);
6081                self.realizers[block_idx]
6082                    .apply_log_kappa(&block_lk, &self.block_term_indices[block_idx])?;
6083                remaining = rest;
6084            } else {
6085                // Last block gets the remainder.
6086                self.realizers[block_idx]
6087                    .apply_log_kappa(&remaining, &self.block_term_indices[block_idx])?;
6088            }
6089        }
6090
6091        log::info!(
6092            "[STAGE] ensure_theta (n-block, {} blocks, {} realizers): {:.3}s",
6093            n,
6094            self.realizers.len(),
6095            t_ensure.elapsed().as_secs_f64(),
6096        );
6097        self.current_theta = Some(theta.clone());
6098        self.last_cost = None;
6099        self.last_eval = None;
6100        Ok(())
6101    }
6102
6103    impl_exact_joint_theta_memo!();
6104
6105    /// Cache a cost-only result. Called after `ensure_theta(theta)` for
6106    /// literal-seed and line-search cost probes. We
6107    /// intentionally do not populate `last_eval` because no gradient was
6108    /// computed; the next outer evaluation at this θ will recompute
6109    /// (V, ∇V) via `evaluate_with_order` if the optimizer asks for it.
6110    fn store_cost_only(&mut self, theta: &Array1<f64>, cost: f64) {
6111        if self
6112            .current_theta
6113            .as_ref()
6114            .is_some_and(|cached| theta_values_match(cached, theta))
6115        {
6116            self.last_cost = Some(cost);
6117        }
6118    }
6119
6120    /// Revoke objective values when the row measure changes while retaining
6121    /// the realized design at the current theta.
6122    fn invalidate_objective_memo(&mut self) {
6123        self.last_cost = None;
6124        self.last_eval = None;
6125    }
6126
6127    fn specs(&self) -> Vec<&TermCollectionSpec> {
6128        self.realizers.iter().map(|r| r.spec()).collect()
6129    }
6130
6131    fn designs(&self) -> Vec<&TermCollectionDesign> {
6132        self.realizers.iter().map(|r| r.design()).collect()
6133    }
6134
6135    /// Combined monotonic design revision across all per-block realizers.
6136    ///
6137    /// Mirrors `SingleBlockExactJointDesignCache::design_revision` for the
6138    /// n-block exact-joint path. Each realizer's `design_revision` counter
6139    /// advances iff `apply_log_kappa` actually rebuilt that block's realized
6140    /// design / smooth penalties; the wrapping sum therefore changes iff
6141    /// *any* block rebuilt. Equal values across two calls imply no realizer
6142    /// has been rebuilt in between, which is the invariant the
6143    /// `ExternalJointHyperEvaluator` canonical-penalty fast path needs.
6144    fn design_revision(&self) -> u64 {
6145        self.realizers
6146            .iter()
6147            .fold(0u64, |acc, r| acc.wrapping_add(r.design_revision()))
6148    }
6149}
6150
6151pub(crate) fn seed_risk_profile_for_likelihood_family(
6152    family: &LikelihoodSpec,
6153) -> gam_problem::SeedRiskProfile {
6154    match &family.response {
6155        ResponseFamily::Gaussian => gam_problem::SeedRiskProfile::Gaussian,
6156        ResponseFamily::RoystonParmar => gam_problem::SeedRiskProfile::Survival,
6157        ResponseFamily::Binomial
6158        | ResponseFamily::Poisson
6159        | ResponseFamily::Tweedie { .. }
6160        | ResponseFamily::NegativeBinomial { .. }
6161        | ResponseFamily::Beta { .. }
6162        | ResponseFamily::Gamma => gam_problem::SeedRiskProfile::GeneralizedLinear,
6163    }
6164}
6165
6166/// Joint-θ dimension above which the single-block exact-joint driver routes
6167/// gradient-only (this doc owns the derivation; the routing site only
6168/// compares against it). The exact outer Hessian builds θ(θ+1)/2 pairwise
6169/// hyper operators, so per-eval cost grows quadratically in θ-dim —
6170/// profiled: `TauTauPairHyperOperator::mul_vec` dominates wall-clock at
6171/// spectral-mode measure-jet candidate counts (θ ≈ 9–11), while θ ≤ 8
6172/// (classic Matérn κ/η fits) keeps cheap exact second-order geometry.
6173const EXACT_JOINT_SECOND_ORDER_THETA_CAP: usize = 8;
6174
6175fn exact_joint_seed_config(
6176    risk_profile: gam_problem::SeedRiskProfile,
6177    auxiliary_dim: usize,
6178    initial_seed_only: bool,
6179) -> gam_problem::SeedConfig {
6180    let mut config = gam_problem::SeedConfig {
6181        risk_profile,
6182        num_auxiliary_trailing: auxiliary_dim,
6183        ..Default::default()
6184    };
6185    match risk_profile {
6186        gam_problem::SeedRiskProfile::Gaussian
6187        | gam_problem::SeedRiskProfile::GaussianLocationScale => {
6188            config.max_seeds = 4;
6189            config.seed_budget = 2;
6190        }
6191        gam_problem::SeedRiskProfile::GeneralizedLinear => {
6192            // Bernoulli marginal-slope Matérn fits use the exact-joint spatial
6193            // driver rather than the family-local BMS outer. Mirror BMS proper:
6194            // screen one principled heuristic seed deeply enough to reach the
6195            // KKT basin instead of spending minutes screening equivalent starts.
6196            config.max_seeds = 1;
6197            config.seed_budget = 1;
6198            config.screen_max_inner_iterations = 8;
6199        }
6200        gam_problem::SeedRiskProfile::Survival => {
6201            // Survival marginal-slope has an additional time/hazard block and
6202            // is the most sensitive Matérn startup regime. Keep more of the
6203            // coherent SPDE candidate manifold alive through truncation and
6204            // validate enough starts that one bad transient does not report
6205            // "no candidate seeds" before reaching a viable basin.
6206            config.max_seeds = 8;
6207            config.seed_budget = 4;
6208            config.screen_max_inner_iterations = 8;
6209        }
6210    }
6211    if initial_seed_only {
6212        // The isotropic Matérn path has already compared and fully profiled its
6213        // two geometry-derived range basins. Its winning [rho, psi] point is an
6214        // explicit certified initial point, so launching another heuristic seed
6215        // would repeat basin selection inside the local joint solve. A budget of
6216        // one gives that explicit initial point sole ownership of the run-plan
6217        // slot (run_plan inserts it at slot zero and skips seed screening).
6218        config.max_seeds = 1;
6219        config.seed_budget = 1;
6220        config.over_smoothing_probe_rho = None;
6221    }
6222    config
6223}
6224
6225#[cfg(test)]
6226mod exact_joint_seed_config_tests {
6227    use super::*;
6228
6229    #[test]
6230    fn exact_joint_marginal_slope_profiles_get_deeper_startup_validation() {
6231        let bms =
6232            exact_joint_seed_config(gam_problem::SeedRiskProfile::GeneralizedLinear, 2, false);
6233        assert_eq!(bms.max_seeds, 1);
6234        assert_eq!(bms.seed_budget, 1);
6235        assert_eq!(bms.screen_max_inner_iterations, 8);
6236        assert_eq!(bms.num_auxiliary_trailing, 2);
6237
6238        let survival = exact_joint_seed_config(gam_problem::SeedRiskProfile::Survival, 3, false);
6239        assert_eq!(survival.max_seeds, 8);
6240        assert_eq!(survival.seed_budget, 4);
6241        assert_eq!(survival.screen_max_inner_iterations, 8);
6242        assert_eq!(survival.num_auxiliary_trailing, 3);
6243    }
6244
6245    #[test]
6246    fn exact_joint_gaussian_keeps_tight_historical_multistart_budget() {
6247        let gaussian = exact_joint_seed_config(gam_problem::SeedRiskProfile::Gaussian, 1, false);
6248        assert_eq!(gaussian.max_seeds, 4);
6249        assert_eq!(gaussian.seed_budget, 2);
6250        assert_eq!(
6251            gaussian.screen_max_inner_iterations,
6252            gam_problem::SeedConfig::default().screen_max_inner_iterations
6253        );
6254        assert_eq!(gaussian.num_auxiliary_trailing, 1);
6255    }
6256
6257    #[test]
6258    fn certified_matern_basin_owns_the_only_joint_start() {
6259        let gaussian = exact_joint_seed_config(gam_problem::SeedRiskProfile::Gaussian, 1, true);
6260        assert_eq!(gaussian.max_seeds, 1);
6261        assert_eq!(gaussian.seed_budget, 1);
6262        assert_eq!(gaussian.over_smoothing_probe_rho, None);
6263        assert_eq!(gaussian.num_auxiliary_trailing, 1);
6264    }
6265}
6266
6267#[cfg(test)]
6268mod wood_reference_df_tests {
6269    use super::*;
6270
6271    // The whole-term LR reference d.f. (#1766). `wood_reference_df` returns
6272    // Wood's smoothing-selection-corrected `edf1 = 2·tr(F) − tr(F²)` on the
6273    // coefficient-influence block, NOT the earlier Satterthwaite ratio
6274    // `tr(F)²/tr(F²)` that collapsed toward 0 for a boundary-shrunk term.
6275
6276    #[test]
6277    fn edf1_equals_two_trace_minus_trace_of_square() {
6278        // A symmetric smoother block with eigenvalues {0.9, 0.4}: a partially
6279        // shrunk penalized term. edf = tr = 1.3; tr(F²) = 0.81 + 0.16 = 0.97;
6280        // edf1 = 2·1.3 − 0.97 = 1.63. (Diagonal ⇒ block F² trace = Σ λ².)
6281        let f = ndarray::array![[0.9_f64, 0.0], [0.0, 0.4]];
6282        let got = wood_reference_df(Some(&f), &(0..2)).unwrap();
6283        assert!(
6284            (got - 1.63).abs() < 1e-12,
6285            "edf1 should be 2*tr - tr(F^2) = 1.63, got {got}"
6286        );
6287        // And it must dominate the raw edf (edf1 >= edf) — the invariant the LR
6288        // reference relies on.
6289        let edf = 1.3;
6290        assert!(got >= edf - 1e-12, "edf1 {got} must be >= edf {edf}");
6291    }
6292
6293    #[test]
6294    fn edf1_never_collapses_below_edf_when_offdiagonals_blow_up() {
6295        // The #1766 degeneracy: a NON-symmetric influence block whose
6296        // off-diagonal coupling makes tr(F²) = Σ_ij F_ij F_ji run away, so the
6297        // old Satterthwaite ratio tr(F)²/tr(F²) crashed toward 0. Here tr = 1.0
6298        // but tr(F²) = 1·1 + 50·(-50) ... take F with large opposite-sign
6299        // off-diagonals so tr(F²) is huge: edf1 = 2·tr − tr(F²) would go very
6300        // negative, and the `.max(tr)` guard must floor it back at edf = tr.
6301        let f = ndarray::array![[0.5_f64, 40.0], [40.0, 0.5]];
6302        let tr = 1.0_f64;
6303        let got = wood_reference_df(Some(&f), &(0..2)).unwrap();
6304        assert!(
6305            got >= tr - 1e-12,
6306            "edf1 must be floored at edf (=tr={tr}) even when tr(F^2) explodes, got {got}"
6307        );
6308        assert!(
6309            got.is_finite() && got > 0.0,
6310            "edf1 must stay finite/positive"
6311        );
6312    }
6313
6314    #[test]
6315    fn returns_none_on_nonpositive_or_missing_trace() {
6316        // No influence matrix at all → None (caller falls back to the
6317        // max(edf, null_dim, 1) floor).
6318        assert!(wood_reference_df(None, &(0..2)).is_none());
6319        // A fully-shrunk block with a non-positive trace → None.
6320        let zero = ndarray::array![[0.0_f64, 0.0], [0.0, 0.0]];
6321        assert!(wood_reference_df(Some(&zero), &(0..2)).is_none());
6322        // An out-of-bounds range → None, never a panic.
6323        let f = ndarray::array![[0.5_f64, 0.0], [0.0, 0.5]];
6324        assert!(wood_reference_df(Some(&f), &(0..5)).is_none());
6325    }
6326}
6327
6328pub(crate) fn exact_joint_multistart_outer_problem(
6329    theta0: &Array1<f64>,
6330    lower: &Array1<f64>,
6331    upper: &Array1<f64>,
6332    rho_dim: usize,
6333    auxiliary_dim: usize,
6334    n_params: usize,
6335    gradient: gam_problem::Derivative,
6336    hessian: gam_problem::DeclaredHessianForm,
6337    prefer_gradient_only: bool,
6338    disable_fixed_point: bool,
6339    risk_profile: gam_problem::SeedRiskProfile,
6340    tolerance: f64,
6341    max_iter: usize,
6342    // BFGS step caps split by parameter type. `bfgs_step_cap` (rho-axis cap)
6343    // bounds first-trial moves on log-λ; documented natural step is ≈ 5.
6344    // `bfgs_step_cap_psi` bounds moves on the trailing `auxiliary_dim`
6345    // psi-axes (kappa / aniso-log-scales), where ≈ ln 2 keeps the kernel
6346    // scale from oscillating across orders of magnitude per iter. Using a
6347    // single uniform cap (the old API) starved rho on the survival-marg-slope
6348    // joint solver because the psi-calibrated value (`ln 2 ≈ 0.69`) was
6349    // applied to log-λ, where |d|≈5 is the natural quasi-Newton magnitude.
6350    bfgs_step_cap: Option<f64>,
6351    bfgs_step_cap_psi: Option<f64>,
6352    screening_cap: Option<Arc<AtomicUsize>>,
6353    // `Some((n_obs, p_cols))` calibrates the outer solver to the n-scaled
6354    // profiled REML/LAML criterion exactly as the primary REML outer
6355    // (`solver/estimate.rs`) does. The profiled criterion is a sum over the n
6356    // observations, so its magnitude is O(n) (|f| ~ thousands at n ~ 10³) for
6357    // EVERY family — Gaussian, binomial, GP/kriging alike. A scale-blind outer
6358    // takes the bare `tolerance` (≈1e-6) as the *absolute* projected-gradient
6359    // floor, which is hopelessly tight against an n-scaled gradient: in-basin
6360    // iterates (e.g. ‖g‖≈7e-2 at |f|≈17, or single-digit ‖g‖ at |f|≈1.3e3)
6361    // never clear it and the fit bails at the iteration cap. Worse, ARC's
6362    // trust-region reduction ratios and default initial regularization are
6363    // referenced against the wrong curvature magnitude, so the first step can
6364    // overshoot and diverge (the ‖g‖≈½|f| blow-ups in #1053/#1066). Threading
6365    // the scale (→ absolute floor = max(tol, n·1e-9)) plus a warm ARC
6366    // regularization (σ₀ = 0.25) and operator trust radius (4.0) makes the
6367    // spatial exact-joint outer converge as robustly as the primary REML outer
6368    // across 1-D Matérn (#1053), 2-D binomial geo (#1066), and GP/kriging
6369    // (#1069). This is NOT a loosening of the `τ·(1+|f|)` REML acceptance gate
6370    // — that relative-to-cost criterion is unchanged; only the nonsensical
6371    // scale-free *absolute* floor and the solver's curvature reference are
6372    // corrected. `None` preserves the prior scale-free calibration.
6373    profiled_objective_size: Option<(usize, usize)>,
6374    // #1464: `true` when the fit carries a constant-curvature `curv()` term. Its
6375    // geodesic-exponential kernel collapses toward the constant function on the
6376    // +κ side, so the joint REML optimum there is a LARGE smoothing λ beyond the
6377    // historical ±12 ρ box. For that case the over-smoothing ρ ceiling is widened
6378    // to `RHO_BOUND` and an explicit high-ρ over-smoothing multistart probe is
6379    // seeded so the joint ARC can reach that basin. `false` keeps the historical
6380    // ±12 box and seed grid byte-for-byte for every other spatial/Matérn/Duchon/
6381    // sphere/survival joint fit.
6382    has_constant_curvature: bool,
6383    // `true` only after the isotropic Matérn endpoint profiler has certified a
6384    // winning range basin. The explicit theta0 then owns the sole joint-start
6385    // budget; generic multi-block and latent-coordinate callers retain their
6386    // family-specific multistart policies.
6387    initial_seed_only: bool,
6388) -> Result<gam_solve::rho_optimizer::OuterProblem, EstimationError> {
6389    if rho_dim > theta0.len() {
6390        crate::bail_invalid_estim!(
6391            "exact joint outer problem declares {rho_dim} smoothing coordinates for theta length {}",
6392            theta0.len(),
6393        );
6394    }
6395    let mut seed_heuristic = theta0.to_vec();
6396    let initial_lambdas = gam_problem::checked_exp_log_strengths(
6397        theta0.iter().take(rho_dim).copied(),
6398    )
6399    .map_err(|error| {
6400        EstimationError::InvalidInput(format!(
6401            "exact joint initial smoothing coordinate is outside the canonical log-strength domain: {error}"
6402        ))
6403    })?;
6404    for (value, lambda) in seed_heuristic[..rho_dim].iter_mut().zip(initial_lambdas) {
6405        *value = lambda;
6406    }
6407    // Over-smoothing ρ ceiling: widened only for a constant-curvature fit (see
6408    // the `has_constant_curvature` param doc). Drives both the scalar saturation
6409    // reference and the seed-grid clamp; the actual box is the per-dim
6410    // `lower`/`upper` arrays passed in.
6411    let rho_ceiling = if has_constant_curvature {
6412        gam_solve::estimate::RHO_BOUND
6413    } else {
6414        12.0
6415    };
6416    let mut problem = gam_solve::rho_optimizer::OuterProblem::new(n_params)
6417        .with_gradient(gradient)
6418        .with_hessian(hessian)
6419        .with_prefer_gradient_only(prefer_gradient_only)
6420        .with_disable_fixed_point(disable_fixed_point)
6421        // Re-enable the automatic fallback ladder for exact joint spatial
6422        // problems. It was previously `Disabled` to suppress a geo-bench
6423        // fallback bug where HybridEFS ψ stagnation degraded silently to
6424        // BfgsApprox on a Charbonnier surface. With the ψ-stagnation guard
6425        // in OuterFixedPointBridge (`MAX_CONSECUTIVE_PSI_STAGNATION`) the
6426        // bridge now surfaces `EFS_FIRST_ORDER_FALLBACK_MARKER` when ψ
6427        // stationarity cannot be enforced, so the ladder routes correctly
6428        // to a joint gradient-based solver instead of grinding HybridEFS
6429        // for thousands of iterations.
6430        .with_fallback_policy(gam_solve::rho_optimizer::FallbackPolicy::Automatic)
6431        .with_psi_dim(auxiliary_dim)
6432        .with_tolerance(tolerance)
6433        .with_max_iter(max_iter)
6434        .with_bounds(lower.clone(), upper.clone())
6435        .with_initial_rho(theta0.clone())
6436        .with_bfgs_step_cap(bfgs_step_cap)
6437        .with_bfgs_step_cap_psi(bfgs_step_cap_psi)
6438        .with_seed_config({
6439            let mut sc = exact_joint_seed_config(risk_profile, auxiliary_dim, initial_seed_only);
6440            if has_constant_curvature {
6441                // Let the seed grid reach the widened over-smoothing ceiling so a
6442                // smooth whose true REML optimum genuinely lives at large λ can be
6443                // discovered (#1464).
6444                sc.bounds = (sc.bounds.0, rho_ceiling);
6445                // gam#1464: do NOT inject an explicit over-smoothing probe at
6446                // ρ ≈ +15 for constant-curvature terms. The probe seeds the joint
6447                // [ρ, ψ] solve at the collapsed-kernel corner where the geodesic
6448                // exponential exp(−d_κ/L) degenerates to a near-constant. There the
6449                // criterion is flat in κ (the kernel no longer resolves curvature)
6450                // and reduces to the monotone log-det Occam term, so keep-best
6451                // adopts the low-Occam collapsed null regardless of the true κ sign
6452                // — the bit-identical κ̂ → +chart-bound rail for both ±κ datasets
6453                // (the headline #1464 sign-blindness). Curvature is instead chosen
6454                // once by the sign-correct continuous fair-profile solve before
6455                // this joint nuisance optimization, and its coordinate is pinned
6456                // here. The widened ρ ceiling is retained: legitimate
6457                // over-smoothing remains reachable by the analytic gradient solve
6458                // without pre-pinning a start at the collapsed corner.
6459            }
6460            sc
6461        })
6462        .with_rho_bound(rho_ceiling)
6463        .with_heuristic_lambdas(seed_heuristic);
6464    if let Some((n_obs, p_cols)) = profiled_objective_size {
6465        // Calibrate to the n-scaled profiled criterion (see the param doc):
6466        // n-aware objective scale → sane absolute gradient floor + correct ARC
6467        // reduction-ratio reference, plus a warm ARC regularization / operator
6468        // trust radius that prevents the first-step overshoot. These are the
6469        // knobs the spatial exact-joint path was missing relative to the
6470        // primary REML outer; without them the iso-κ length-scale fit stalls or
6471        // diverges as |f| grows with n (#1053 / #1066 / #1069).
6472        problem = problem
6473            .with_objective_scale(Some(n_obs as f64))
6474            .with_problem_size(n_obs, p_cols)
6475            .with_arc_initial_regularization(Some(0.25))
6476            .with_operator_initial_trust_radius(Some(4.0));
6477    }
6478    if let Some(screening_cap) = screening_cap {
6479        problem = problem
6480            .with_screening_cap(screening_cap)
6481            .with_screen_initial_rho(true);
6482    }
6483    Ok(problem)
6484}
6485
6486pub fn optimize_spatial_length_scale_exact_joint<FitOut, Mode, FitFn, ExactFn, ExactEfsFn, SeedFn>(
6487    data: ArrayView2<'_, f64>,
6488    block_specs: &[TermCollectionSpec],
6489    block_term_indices: &[Vec<usize>],
6490    kappa_options: &SpatialLengthScaleOptimizationOptions,
6491    joint_setup: &ExactJointHyperSetup,
6492    seed_risk_profile: gam_problem::SeedRiskProfile,
6493    analytic_joint_gradient_available: bool,
6494    analytic_joint_hessian_available: bool,
6495    disable_fixed_point: bool,
6496    screening_cap: Option<Arc<AtomicUsize>>,
6497    outer_derivative_policy: gam_model_api::families::custom_family::OuterDerivativePolicy,
6498    mut fit_fn: FitFn,
6499    mut exact_fn: ExactFn,
6500    mut exact_efs_fn: ExactEfsFn,
6501    mut seed_inner_beta_fn: SeedFn,
6502) -> Result<SpatialLengthScaleOptimizationResult<FitOut>, String>
6503where
6504    FitFn: FnMut(
6505        &Array1<f64>,
6506        &[TermCollectionSpec],
6507        &[TermCollectionDesign],
6508        SpatialFitProvenance<'_, Mode>,
6509    ) -> Result<FitOut, String>,
6510    ExactFn: FnMut(
6511        &Array1<f64>,
6512        &[TermCollectionSpec],
6513        &[TermCollectionDesign],
6514        gam_solve::estimate::reml::reml_outer_engine::EvalMode,
6515        &gam_problem::outer_subsample::RowSet,
6516    ) -> Result<ExactJointEvaluation<Mode>, String>,
6517    ExactEfsFn: FnMut(
6518        &Array1<f64>,
6519        &[TermCollectionSpec],
6520        &[TermCollectionDesign],
6521        &gam_problem::outer_subsample::RowSet,
6522    ) -> Result<ExactJointEfsEvaluation<Mode>, String>,
6523    SeedFn: FnMut(&Array1<f64>) -> Result<gam_solve::rho_optimizer::SeedOutcome, EstimationError>,
6524{
6525    let n_blocks = block_specs.len();
6526    if block_term_indices.len() != n_blocks {
6527        return Err(SmoothError::dimension_mismatch(format!(
6528            "block_specs ({}) and block_term_indices ({}) length mismatch",
6529            n_blocks,
6530            block_term_indices.len()
6531        ))
6532        .into());
6533    }
6534
6535    let log_kappa_dim = joint_setup.log_kappa_dim();
6536
6537    log::trace!(
6538        "[spatial-exact-joint] driver entry: aux_dim={} log_kappa_dim={} kappa_enabled={} rho_dim={} theta0_len={}",
6539        joint_setup.auxiliary_dim(),
6540        log_kappa_dim,
6541        kappa_options.enabled,
6542        joint_setup.rho_dim(),
6543        joint_setup.theta0().len()
6544    );
6545
6546    // -----------------------------------------------------------------------
6547    // Fast path: kappa disabled or no spatial terms — build designs once.
6548    // -----------------------------------------------------------------------
6549    if joint_setup.auxiliary_dim() == 0 && (!kappa_options.enabled || log_kappa_dim == 0) {
6550        log::trace!(
6551            "[spatial-exact-joint] taking fast path (no outer theta optimization in this driver)"
6552        );
6553        let (designs, resolved_specs) = build_term_collection_designs_and_freeze_joint(
6554            data, block_specs,
6555        )
6556        .map_err(|e| {
6557            format!("failed to build and freeze joint block designs during exact joint kappa optimization: {e}")
6558        })?;
6559        let theta0 = joint_setup.theta0();
6560
6561        // Build temporary owned slices for the closure call.
6562        let spec_refs: Vec<TermCollectionSpec> = resolved_specs.clone();
6563        let design_refs: Vec<TermCollectionDesign> = designs.clone();
6564        let fit = fit_fn(
6565            &theta0,
6566            &spec_refs,
6567            &design_refs,
6568            SpatialFitProvenance::NoOuterOptimization,
6569        )?;
6570        return Ok(SpatialLengthScaleOptimizationResult {
6571            resolved_specs,
6572            designs,
6573            fit,
6574            certified_outer: None,
6575            timing: None,
6576        });
6577    }
6578
6579    // -----------------------------------------------------------------------
6580    // Full optimization path.
6581    // -----------------------------------------------------------------------
6582    let theta0 = joint_setup.theta0();
6583    let lower = joint_setup.lower();
6584    let upper = joint_setup.upper();
6585    if theta0.len() < log_kappa_dim || lower.len() != theta0.len() || upper.len() != theta0.len() {
6586        return Err(SmoothError::dimension_mismatch(format!(
6587            "invalid exact joint theta setup: theta0={}, lower={}, upper={}, required_log_kappa_dim={}",
6588            theta0.len(),
6589            lower.len(),
6590            upper.len(),
6591            log_kappa_dim
6592        ))
6593        .into());
6594    }
6595    let rho_dim = joint_setup.rho_dim();
6596    let all_dims = joint_setup.log_kappa_dims_per_term();
6597
6598    // Build bootstrap designs and frozen specs for each block.
6599    let (boot_designs, best_specs) = build_term_collection_designs_and_freeze_joint(
6600        data,
6601        block_specs,
6602    )
6603    .map_err(|e| {
6604        format!(
6605            "failed to build and freeze joint block designs during exact joint kappa bootstrap: {e}"
6606        )
6607    })?;
6608    // Capability vs realized policy: the family may *advertise* an exact
6609    // analytic outer Hessian, but at this realized (n, psi_dim, rho_dim,
6610    // p_total) the predicted per-eval cost can still exceed the universal
6611    // outer-Hessian work budget. In that regime we route the outer optimizer
6612    // through gradient-only BFGS / L-BFGS, which is **convergent** to the
6613    // exact MLE — it just takes more line-search iterations. This is **not**
6614    // a feature drop: quasi-Newton picks up curvature from successive
6615    // analytic gradients, and the per-eval cost saving (`O(p)` instead of
6616    // `O(p²)`) more than pays for the iteration overhead at large scale.
6617    let policy_hessian_form = outer_derivative_policy.declared_hessian_form();
6618    let analytic_outer_hessian_available = analytic_joint_hessian_available
6619        && matches!(
6620            policy_hessian_form,
6621            gam_problem::DeclaredHessianForm::Either
6622                | gam_problem::DeclaredHessianForm::Dense
6623                | gam_problem::DeclaredHessianForm::Operator { .. }
6624        );
6625    let prefer_gradient_only = !analytic_outer_hessian_available;
6626
6627    let theta_dim = theta0.len();
6628    let psi_dim = theta_dim - rho_dim;
6629
6630    // Build the cache with one realizer per block.
6631    let cache_blocks: Vec<(TermCollectionSpec, TermCollectionDesign, Vec<usize>)> = best_specs
6632        .iter()
6633        .zip(boot_designs.iter())
6634        .zip(block_term_indices.iter())
6635        .map(|((spec, design), terms)| (spec.clone(), design.clone(), terms.clone()))
6636        .collect();
6637
6638    struct NBlockExactJointState<'d, M> {
6639        cache: ExactJointDesignCache<'d>,
6640        row_set: gam_problem::outer_subsample::RowSet,
6641        staged_pilot_active: bool,
6642        terminal_mode: Option<(Array1<f64>, f64, M)>,
6643    }
6644
6645    impl<M> NBlockExactJointState<'_, M> {
6646        fn ensure_theta(&mut self, theta: &Array1<f64>) -> Result<(), String> {
6647            let theta_changed = !self
6648                .cache
6649                .current_theta
6650                .as_ref()
6651                .is_some_and(|current| theta_values_match(current, theta));
6652            if theta_changed {
6653                self.terminal_mode = None;
6654            }
6655            self.cache.ensure_theta(theta)
6656        }
6657
6658        fn install_terminal_mode(&mut self, theta: &Array1<f64>, objective: f64, mode: M) {
6659            self.terminal_mode = Some((theta.clone(), objective, mode));
6660        }
6661
6662        fn terminal_mode_matches(&self, theta: &Array1<f64>, objective: f64) -> bool {
6663            self.terminal_mode
6664                .as_ref()
6665                .is_some_and(|(mode_theta, mode_objective, _)| {
6666                    theta_values_match(mode_theta, theta)
6667                        && mode_objective.to_bits() == objective.to_bits()
6668                })
6669        }
6670    }
6671
6672    let mut state = NBlockExactJointState {
6673        cache: ExactJointDesignCache::new(data, cache_blocks, rho_dim, all_dims.clone())?,
6674        row_set: gam_problem::outer_subsample::RowSet::All,
6675        staged_pilot_active: false,
6676        terminal_mode: None,
6677    };
6678
6679    // ── P7: staged-κ schedule ────────────────────────────────────────────
6680    //
6681    // The κ MLE for a stationary spatial process is asymptotically
6682    // *invariant* in `n` once `n` is past the Monte-Carlo resolution of
6683    // the cell-moment kernel. At large scale (`n ≥ STAGED_KAPPA_*`) the
6684    // Monte-Carlo error of a `K = 5_000`-row pilot is ≪ the κ posterior
6685    // width, so estimating θ on a stratified `K`-row pilot returns
6686    // statistically the *same* estimate as the full-data fit at a
6687    // fraction of the wall-clock cost. The shared outer runner then continues
6688    // from that checkpoint on the exact full-data measure and issues its
6689    // mandatory analytic certificate only after the transition.
6690    //
6691    // This is **not a heuristic shortcut**. It is the textbook
6692    // pilot-then-refine schedule for stationary-process likelihoods,
6693    // chosen here because the per-eval cost of the κ gradient grows
6694    // linearly in `n` and the pilot subsample reduces that cost by a
6695    // factor of `n / K`. The exact full-data refinement starts literally at
6696    // the pilot checkpoint and retains the learned trust radius and Hessian;
6697    // it costs one terminal
6698    // full-data evaluation when the pilot point already certifies and keeps
6699    // optimizing when it does not.
6700    //
6701    // At `n < STAGED_KAPPA_TRIGGER_N` the schedule collapses to one
6702    // full-data stage — identical to the pre-P7 behaviour.
6703    // Note: the n≥30_000 pilot trigger lives in
6704    // `outer_derivative_policy.should_use_staged_kappa(n_total)`; this fn
6705    // only carries the constants it consumes directly.
6706    const KAPPA_PILOT_K: usize = 5_000;
6707
6708    let n_total = data.nrows();
6709    let use_staged_kappa = outer_derivative_policy.should_use_staged_kappa(n_total);
6710    if use_staged_kappa {
6711        log::info!(
6712            "[KAPPA-STAGED] auto-engaging pilot+exact schedule: n={} pilot_k={}",
6713            n_total,
6714            KAPPA_PILOT_K,
6715        );
6716    }
6717
6718    // Build the initial row mask for the κ optimization.
6719    //
6720    // * `use_staged_kappa = false`: full data (`RowSet::All`). The
6721    //   schedule collapses to the historical single-stage path.
6722    // * `use_staged_kappa = true`: deterministic uniform pilot of size
6723    //   `min(KAPPA_PILOT_K, n_total)`, wrapped as a `RowSet::Subsample`
6724    //   with per-row HT weight `n_total / k_pilot`. The uniform pick is
6725    //   a valid unbiased estimator on its own; the stratified
6726    //   per-decile picker
6727    //   (`marginal_slope_shared::auto_outer_score_subsample`) requires
6728    //   the response vector `z`, which only the family evaluator can
6729    //   produce. **Agent C replaces this with the stratified pick once
6730    //   `exact_fn` exposes the per-row score.**
6731    //
6732    // Sampling RNG is seeded from `n_total` so the pilot is
6733    // deterministic across reruns at fixed `n`.
6734    fn build_uniform_pilot_subsample(
6735        n_total: usize,
6736        k_target: usize,
6737        seed: u64,
6738    ) -> gam_problem::outer_subsample::OuterScoreSubsample {
6739        use gam_problem::outer_subsample::OuterScoreSubsample;
6740        let k = k_target.min(n_total);
6741        if k == 0 || n_total == 0 {
6742            return OuterScoreSubsample::from_uniform_inclusion_mask(Vec::new(), n_total, seed);
6743        }
6744        // Reservoir-free deterministic pick: linear congruential walk
6745        // over a shuffled index set; for the pilot, a fast Floyd-style
6746        // sample is sufficient.
6747        let mut mask: Vec<usize> = Vec::with_capacity(k);
6748        // Splitmix64-driven Floyd's sampler.
6749        let mut state = seed.wrapping_add(0x9E3779B97F4A7C15);
6750        let splitmix = |s: &mut u64| -> u64 { gam_linalg::utils::splitmix64(s) };
6751        let mut taken = std::collections::HashSet::with_capacity(k);
6752        for j in (n_total - k)..n_total {
6753            let r = (splitmix(&mut state) % (j as u64 + 1)) as usize;
6754            if !taken.insert(r) {
6755                taken.insert(j);
6756                mask.push(j);
6757            } else {
6758                mask.push(r);
6759            }
6760        }
6761        mask.sort_unstable();
6762        mask.dedup();
6763        OuterScoreSubsample::from_uniform_inclusion_mask(mask, n_total, seed)
6764    }
6765
6766    if use_staged_kappa {
6767        let pilot = build_uniform_pilot_subsample(n_total, KAPPA_PILOT_K, n_total as u64);
6768        state.row_set = gam_problem::outer_subsample::RowSet::Subsample {
6769            rows: std::sync::Arc::clone(&pilot.rows),
6770            n_full: n_total,
6771        };
6772        state.staged_pilot_active = true;
6773    }
6774
6775    let exact_fn_cell = std::cell::RefCell::new(&mut exact_fn);
6776    let exact_efs_fn_cell = std::cell::RefCell::new(&mut exact_efs_fn);
6777
6778    // ── κ-optimization scaling instrumentation ──
6779    //
6780    // Per-phase wall-clock counters for the three kinds of evaluator
6781    // invocation the κ outer drives: cost-only line-search probes,
6782    // value-and-gradient(/Hessian) evaluations at accepted iterates, and
6783    // EFS fixed-point evaluations. Each invocation emits one
6784    // `[KAPPA-PHASE]` log line with a per-call elapsed time, plus the
6785    // running call counter and a summary `theta_norm` /
6786    // `log_kappa_norm` so the bench runner can attribute cost to
6787    // particular trajectory regions. A single `[KAPPA-PHASE-SUMMARY]`
6788    // line is emitted on optimization exit. Grepping these is the
6789    // production-fit κ-scaling probe (task #32) — measurement happens
6790    // in real large-scale fits rather than a synthetic harness, so the
6791    // scaling law reflects the actual workload.
6792    use std::cell::Cell;
6793    let kphase_cost_calls: Cell<usize> = Cell::new(0);
6794    let kphase_cost_total_s: Cell<f64> = Cell::new(0.0);
6795    let kphase_eval_calls: Cell<usize> = Cell::new(0);
6796    let kphase_eval_total_s: Cell<f64> = Cell::new(0.0);
6797    let kphase_efs_calls: Cell<usize> = Cell::new(0);
6798    let kphase_efs_total_s: Cell<f64> = Cell::new(0.0);
6799    let kphase_optim_start = std::time::Instant::now();
6800    let kphase_log_kappa_dim = log_kappa_dim;
6801    let kphase_log_norms = |theta: &Array1<f64>| -> (f64, f64) {
6802        let theta_norm = theta.iter().map(|v| v * v).sum::<f64>().sqrt();
6803        let log_kappa_norm = if kphase_log_kappa_dim > 0 && theta.len() >= kphase_log_kappa_dim {
6804            let start = theta.len() - kphase_log_kappa_dim;
6805            theta.iter().skip(start).map(|v| v * v).sum::<f64>().sqrt()
6806        } else {
6807            0.0
6808        };
6809        (theta_norm, log_kappa_norm)
6810    };
6811
6812    use gam_problem::{DeclaredHessianForm, Derivative, OuterEval};
6813    use gam_solve::rho_optimizer::OuterEvalOrder;
6814
6815    // Joint design width across blocks → the `p` reported to the outer solver's
6816    // operator-vs-dense Hessian crossover. `n_total` is the load-bearing
6817    // profiled-objective scale (see `exact_joint_multistart_outer_problem`).
6818    let joint_p_cols: usize = boot_designs
6819        .iter()
6820        .map(|d| d.design.ncols())
6821        .sum::<usize>()
6822        .max(1);
6823
6824    let problem = exact_joint_multistart_outer_problem(
6825        &theta0,
6826        &lower,
6827        &upper,
6828        rho_dim,
6829        psi_dim,
6830        theta_dim,
6831        if analytic_joint_gradient_available {
6832            Derivative::Analytic
6833        } else {
6834            Derivative::Unavailable
6835        },
6836        if analytic_outer_hessian_available {
6837            DeclaredHessianForm::Either
6838        } else {
6839            DeclaredHessianForm::Unavailable
6840        },
6841        prefer_gradient_only,
6842        disable_fixed_point,
6843        seed_risk_profile,
6844        kappa_options.rel_tol.max(1e-6),
6845        kappa_options.max_outer_iter.max(1),
6846        // Rho-axis cap: log-λ natural step ≈ 5.
6847        Some(5.0),
6848        // Psi-axis cap: kappa scale needs ~ln 2 per iter.
6849        Some(kappa_options.log_step.clamp(0.25, 1.0)),
6850        screening_cap.clone(),
6851        // n-scaled profiled-criterion calibration for every family (#1053 /
6852        // #1066 / #1069 iso-κ non-convergence cure).
6853        Some((n_total, joint_p_cols)),
6854        // #1464: widen the over-smoothing ρ ceiling + seed a high-λ probe when
6855        // any block carries a constant-curvature term.
6856        block_specs
6857            .iter()
6858            .any(|s| !constant_curvature_term_indices(s).is_empty()),
6859        // Multi-block optimization has no preceding scalar Matérn endpoint
6860        // certificate, so retain its family-specific seed cascade.
6861        false,
6862    )
6863    .map_err(|e| e.to_string())?;
6864
6865    // Helper: collect specs and designs from cache into owned Vecs for closure calls.
6866    fn collect_specs(cache: &ExactJointDesignCache<'_>) -> Vec<TermCollectionSpec> {
6867        cache.specs().into_iter().cloned().collect()
6868    }
6869    fn collect_designs(cache: &ExactJointDesignCache<'_>) -> Vec<TermCollectionDesign> {
6870        cache.designs().into_iter().cloned().collect()
6871    }
6872
6873    let result = {
6874        let eval_outer = |ctx: &mut &mut NBlockExactJointState<'_, Mode>,
6875                          theta: &Array1<f64>,
6876                          order: OuterEvalOrder|
6877         -> Result<OuterEval, EstimationError> {
6878            if let Some((cost, grad, hess)) = ctx.cache.memoized_eval(theta)
6879                && ctx.terminal_mode_matches(theta, cost)
6880            {
6881                let cached_satisfies_order = match order {
6882                    OuterEvalOrder::Value => true,
6883                    OuterEvalOrder::ValueAndGradient => grad.len() == theta.len(),
6884                    OuterEvalOrder::ValueGradientHessian => {
6885                        grad.len() == theta.len() && hess.is_analytic()
6886                    }
6887                };
6888                if cached_satisfies_order {
6889                    if !cost.is_finite() {
6890                        return Ok(OuterEval::infeasible(theta.len()));
6891                    }
6892                    // Symmetric with the non-finite-cost guard above: a non-finite
6893                    // gradient marks this θ as infeasible just as a non-finite cost
6894                    // does (e.g. degenerate tied / zero-gap survival times drive the
6895                    // analytic exact-joint gradient channel to NaN/Inf). Return the
6896                    // bounded infeasible sentinel so the outer optimizer rejects the
6897                    // step and shrinks its trust region — instead of hard-failing the
6898                    // entire REML fit and handing the driver an unbroken stream of
6899                    // objective failures whose recovery path deepens once per outer
6900                    // step until the worker stack overflows (the survival
6901                    // location-scale path is the one that routes through this analytic
6902                    // gradient, which is why it crashed where the cost-only paths only
6903                    // stall).
6904                    if grad.iter().any(|v| !v.is_finite()) {
6905                        return Ok(OuterEval::infeasible(theta.len()));
6906                    }
6907                    return Ok(OuterEval {
6908                        cost,
6909                        gradient: grad,
6910                        hessian: hess,
6911                        inner_beta_hint: None,
6912                    });
6913                }
6914            }
6915            ctx.ensure_theta(theta).map_err(|err| {
6916                EstimationError::InvalidInput(format!(
6917                    "n-block exact-joint spatial design realization failed: {err}"
6918                ))
6919            })?;
6920            let design_revision = Some(ctx.cache.design_revision());
6921            let specs = collect_specs(&ctx.cache);
6922            let designs = collect_designs(&ctx.cache);
6923            // Clamp the requested order against the realized outer
6924            // derivative policy. The capability-aware
6925            // `analytic_outer_hessian_available` already encodes the
6926            // policy gate; re-checking through `order_for_evaluation`
6927            // here keeps the per-eval branch in lockstep with the
6928            // top-of-function declaration so the optimizer and the
6929            // evaluator never disagree on what was requested.
6930            let clamped = outer_derivative_policy.order_for_evaluation(order);
6931            let value_only = matches!(clamped, OuterEvalOrder::Value);
6932            let need_hessian = matches!(clamped, OuterEvalOrder::ValueGradientHessian)
6933                && analytic_outer_hessian_available;
6934            let eval_mode = if value_only {
6935                gam_solve::estimate::reml::reml_outer_engine::EvalMode::ValueOnly
6936            } else if need_hessian {
6937                gam_solve::estimate::reml::reml_outer_engine::EvalMode::ValueGradientHessian
6938            } else {
6939                gam_solve::estimate::reml::reml_outer_engine::EvalMode::ValueAndGradient
6940            };
6941            let t0 = std::time::Instant::now();
6942            let result =
6943                (*exact_fn_cell.borrow_mut())(theta, &specs, &designs, eval_mode, &ctx.row_set);
6944            let elapsed_s = t0.elapsed().as_secs_f64();
6945            kphase_eval_calls.set(kphase_eval_calls.get() + 1);
6946            kphase_eval_total_s.set(kphase_eval_total_s.get() + elapsed_s);
6947            let (theta_norm, log_kappa_norm) = kphase_log_norms(theta);
6948            log::info!(
6949                "[KAPPA-PHASE] phase=eval_outer call={} order={:?} design_revision={:?} theta_norm={:.4e} log_kappa_norm={:.4e} elapsed_s={:.4}",
6950                kphase_eval_calls.get(),
6951                order,
6952                design_revision,
6953                theta_norm,
6954                log_kappa_norm,
6955                elapsed_s,
6956            );
6957            match result {
6958                Ok(ExactJointEvaluation {
6959                    objective: cost,
6960                    gradient: grad,
6961                    hessian: hess,
6962                    mode,
6963                }) => {
6964                    ctx.install_terminal_mode(theta, cost, mode);
6965                    if value_only {
6966                        ctx.cache.store_cost_only(theta, cost);
6967                    } else {
6968                        ctx.cache.store_eval((cost, grad.clone(), hess.clone()));
6969                    }
6970                    if !cost.is_finite() {
6971                        return Ok(OuterEval::infeasible(theta.len()));
6972                    }
6973                    // Symmetric with the non-finite-cost guard above: a non-finite
6974                    // gradient marks this θ as infeasible just as a non-finite cost
6975                    // does (e.g. degenerate tied / zero-gap survival times drive the
6976                    // analytic exact-joint gradient channel to NaN/Inf). Return the
6977                    // bounded infeasible sentinel so the outer optimizer rejects the
6978                    // step and shrinks its trust region — instead of hard-failing the
6979                    // entire REML fit and handing the driver an unbroken stream of
6980                    // objective failures whose recovery path deepens once per outer
6981                    // step until the worker stack overflows (the survival
6982                    // location-scale path is the one that routes through this analytic
6983                    // gradient, which is why it crashed where the cost-only paths only
6984                    // stall).
6985                    if grad.iter().any(|v| !v.is_finite()) {
6986                        return Ok(OuterEval::infeasible(theta.len()));
6987                    }
6988                    Ok(OuterEval {
6989                        cost,
6990                        gradient: grad,
6991                        hessian: hess,
6992                        inner_beta_hint: None,
6993                    })
6994                }
6995                Err(err) => Err(EstimationError::RemlOptimizationFailed(format!(
6996                    "n-block exact-joint spatial evaluation failed: {err}"
6997                ))),
6998            }
6999        };
7000
7001        let obj = problem.build_objective_with_eval_order(
7002            &mut state,
7003            |ctx: &mut &mut NBlockExactJointState<'_, Mode>, theta: &Array1<f64>| {
7004                if let Some(cost) = ctx.cache.memoized_cost(theta)
7005                    && ctx.terminal_mode_matches(theta, cost)
7006                {
7007                    return Ok(cost);
7008                }
7009                ctx.ensure_theta(theta).map_err(|err| {
7010                    EstimationError::InvalidInput(format!(
7011                        "n-block exact-joint spatial design realization failed: {err}"
7012                    ))
7013                })?;
7014                let design_revision = Some(ctx.cache.design_revision());
7015                let specs = collect_specs(&ctx.cache);
7016                let designs = collect_designs(&ctx.cache);
7017                // Cost-only line-search probe: pass `ValueOnly` so the closure
7018                // skips gradient and Hessian assembly. This is the principled
7019                // fix for the N-block joint optimization V+G-per-probe waste —
7020                // gradient construction (≈ 6.5·10⁹ FLOPs per CTN step at
7021                // n=320 000, n_grid=293, p_resp=32, p_cov=23) is now paid only
7022                // when the outer evaluator actually requests it.
7023                let t0 = std::time::Instant::now();
7024                let result = (*exact_fn_cell.borrow_mut())(
7025                    theta,
7026                    &specs,
7027                    &designs,
7028                    gam_solve::estimate::reml::reml_outer_engine::EvalMode::ValueOnly,
7029                    &ctx.row_set,
7030                );
7031                let elapsed_s = t0.elapsed().as_secs_f64();
7032                kphase_cost_calls.set(kphase_cost_calls.get() + 1);
7033                kphase_cost_total_s.set(kphase_cost_total_s.get() + elapsed_s);
7034                let (theta_norm, log_kappa_norm) = kphase_log_norms(theta);
7035                log::info!(
7036                    "[KAPPA-PHASE] phase=cost call={} design_revision={:?} theta_norm={:.4e} log_kappa_norm={:.4e} elapsed_s={:.4}",
7037                    kphase_cost_calls.get(),
7038                    design_revision,
7039                    theta_norm,
7040                    log_kappa_norm,
7041                    elapsed_s,
7042                );
7043                match result {
7044                    Ok(ExactJointEvaluation {
7045                        objective: cost,
7046                        mode,
7047                        ..
7048                    }) => {
7049                        ctx.install_terminal_mode(theta, cost, mode);
7050                        // Don't `store_eval`: that path is only valid when the
7051                        // closure produced a real gradient. The next outer-eval
7052                        // call will recompute (V, ∇V) at this θ if needed; the
7053                        // memoized_cost path covers the common case where the
7054                        // line search returns to an accepted iterate.
7055                        ctx.cache.store_cost_only(theta, cost);
7056                        Ok(cost)
7057                    }
7058                    Err(err) => Err(EstimationError::RemlOptimizationFailed(format!(
7059                        "n-block exact-joint spatial cost evaluation failed: {err}"
7060                    ))),
7061                }
7062            },
7063            |ctx: &mut &mut NBlockExactJointState<'_, Mode>, theta: &Array1<f64>| {
7064                eval_outer(
7065                    ctx,
7066                    theta,
7067                    if analytic_outer_hessian_available {
7068                        OuterEvalOrder::ValueGradientHessian
7069                    } else {
7070                        OuterEvalOrder::ValueAndGradient
7071                    },
7072                )
7073            },
7074            |ctx: &mut &mut NBlockExactJointState<'_, Mode>,
7075             theta: &Array1<f64>,
7076             order: OuterEvalOrder| { eval_outer(ctx, theta, order) },
7077            None::<fn(&mut &mut NBlockExactJointState<'_, Mode>)>,
7078            Some(
7079                |ctx: &mut &mut NBlockExactJointState<'_, Mode>, theta: &Array1<f64>| {
7080                    ctx
7081                        .ensure_theta(theta)
7082                        .map_err(EstimationError::InvalidInput)?;
7083                    let design_revision = Some(ctx.cache.design_revision());
7084                    let specs = collect_specs(&ctx.cache);
7085                    let designs = collect_designs(&ctx.cache);
7086                    let t0 = std::time::Instant::now();
7087                    let eval_result = (*exact_efs_fn_cell.borrow_mut())(
7088                        theta,
7089                        &specs,
7090                        &designs,
7091                        &ctx.row_set,
7092                    );
7093                    let elapsed_s = t0.elapsed().as_secs_f64();
7094                    kphase_efs_calls.set(kphase_efs_calls.get() + 1);
7095                    kphase_efs_total_s.set(kphase_efs_total_s.get() + elapsed_s);
7096                    let (theta_norm, log_kappa_norm) = kphase_log_norms(theta);
7097                    log::info!(
7098                        "[KAPPA-PHASE] phase=efs call={} design_revision={:?} theta_norm={:.4e} log_kappa_norm={:.4e} elapsed_s={:.4}",
7099                        kphase_efs_calls.get(),
7100                        design_revision,
7101                        theta_norm,
7102                        log_kappa_norm,
7103                        elapsed_s,
7104                    );
7105                    let ExactJointEfsEvaluation { evaluation, mode } =
7106                        eval_result.map_err(EstimationError::RemlOptimizationFailed)?;
7107                    // An EFS solve can select a different coefficient mode at
7108                    // the same theta.  Revoke any derivative memo assembled
7109                    // from the previous mode before installing this carrier;
7110                    // a later analytic certification must then re-evaluate and
7111                    // replace both the derivative payload and the owned mode
7112                    // atomically.
7113                    ctx.cache.invalidate_objective_memo();
7114                    ctx.cache.store_cost_only(theta, evaluation.cost);
7115                    ctx.install_terminal_mode(theta, evaluation.cost, mode);
7116                    Ok(evaluation)
7117                },
7118            ),
7119        );
7120        let mut obj = obj
7121            .with_seed_inner_state(
7122                move |_ctx: &mut &mut NBlockExactJointState<'_, Mode>, beta: &Array1<f64>| {
7123                    (seed_inner_beta_fn)(beta)
7124                },
7125            )
7126            .with_exact_polish(|ctx: &mut &mut NBlockExactJointState<'_, Mode>| {
7127                if !ctx.staged_pilot_active {
7128                    return false;
7129                }
7130                // Objective memoization is theta-only, so a pilot value at the
7131                // warm checkpoint must not alias the exact full-data value.
7132                // Keep the realized design and warm coefficient state: only the
7133                // score measure changes here.
7134                ctx.cache.invalidate_objective_memo();
7135                ctx.terminal_mode = None;
7136                ctx.row_set = gam_problem::outer_subsample::RowSet::All;
7137                ctx.staged_pilot_active = false;
7138                true
7139            });
7140
7141        problem
7142            .run_certified(&mut obj, "n-block exact-joint spatial")
7143            .map_err(|error| error.to_string())?
7144    }; // obj dropped here, releasing mutable borrow on state
7145
7146    // ── κ-optimization scaling summary ──
7147    //
7148    // Single line summarizing all per-call wall-clock counters
7149    // accumulated above. The bench runner / scaling-law analyzer
7150    // can pivot on this directly without parsing the per-call
7151    // [KAPPA-PHASE] markers (which remain available for
7152    // attribution).
7153    let kphase_total_s = kphase_optim_start.elapsed().as_secs_f64();
7154    log::info!(
7155        "[KAPPA-PHASE-SUMMARY] log_kappa_dim={} n_cost={} cost_total_s={:.4} n_eval={} eval_total_s={:.4} n_efs={} efs_total_s={:.4} optim_total_s={:.4}",
7156        kphase_log_kappa_dim,
7157        kphase_cost_calls.get(),
7158        kphase_cost_total_s.get(),
7159        kphase_eval_calls.get(),
7160        kphase_eval_total_s.get(),
7161        kphase_efs_calls.get(),
7162        kphase_efs_total_s.get(),
7163        kphase_total_s,
7164    );
7165    let timing = SpatialLengthScaleOptimizationTiming {
7166        log_kappa_dim: kphase_log_kappa_dim,
7167        cost_calls: kphase_cost_calls.get(),
7168        cost_total_s: kphase_cost_total_s.get(),
7169        eval_calls: kphase_eval_calls.get(),
7170        eval_total_s: kphase_eval_total_s.get(),
7171        efs_calls: kphase_efs_calls.get(),
7172        efs_total_s: kphase_efs_total_s.get(),
7173        slow_path_resets: 0,
7174        design_revision_delta: 0,
7175        nfree_skip_row_touches: 0,
7176        nfree_miss_shape: 0,
7177        nfree_miss_value: 0,
7178        nfree_miss_gradient: 0,
7179        nfree_miss_penalty: 0,
7180        nfree_miss_revision: 0,
7181        nfree_miss_second_order: 0,
7182        nfree_miss_other: 0,
7183        optim_total_s: kphase_total_s,
7184    };
7185
7186    if !matches!(state.row_set, gam_problem::outer_subsample::RowSet::All) {
7187        return Err(
7188            "n-block exact-joint spatial optimization returned before its exact full-data transition"
7189                .to_string(),
7190        );
7191    }
7192    let certified_outer = result;
7193    let theta_star = certified_outer.rho().clone();
7194
7195    // ── P7 stage rotation ────────────────────────────────────────────────
7196    // The returned theta and certificate now belong to the exact full-data
7197    // refinement. No separate probe may mutate that certified identity before
7198    // the final coefficient fit.
7199    state.ensure_theta(&theta_star)?;
7200    let (mode_theta, mode_objective, mode) = state.terminal_mode.take().ok_or_else(|| {
7201        "n-block exact-joint spatial optimization produced a certificate without retaining the owned terminal coefficient mode"
7202            .to_string()
7203    })?;
7204    if !theta_values_match(&mode_theta, &theta_star) {
7205        return Err(
7206            "n-block exact-joint spatial terminal coefficient mode does not bitwise match the certified hyperparameter vector"
7207                .to_string(),
7208        );
7209    }
7210    if mode_objective.to_bits() != certified_outer.final_value().to_bits() {
7211        return Err(format!(
7212            "n-block exact-joint spatial terminal coefficient mode objective does not bitwise match the certified objective: mode={mode_objective:.17e}, certified={:.17e}",
7213            certified_outer.final_value(),
7214        ));
7215    }
7216
7217    let resolved_specs: Vec<TermCollectionSpec> = collect_specs(&state.cache);
7218    let designs: Vec<TermCollectionDesign> = collect_designs(&state.cache);
7219
7220    let fit = fit_fn(
7221        &theta_star,
7222        &resolved_specs,
7223        &designs,
7224        SpatialFitProvenance::Certified {
7225            outer: &certified_outer,
7226            mode,
7227        },
7228    )?;
7229
7230    for spec in &resolved_specs {
7231        log_spatial_aniso_scales(spec);
7232    }
7233
7234    Ok(SpatialLengthScaleOptimizationResult {
7235        resolved_specs,
7236        designs,
7237        fit,
7238        certified_outer: Some(certified_outer),
7239        timing: Some(timing),
7240    })
7241}
7242
7243fn try_exact_joint_latent_coord_optimization(
7244    data: ArrayView2<'_, f64>,
7245    y: ArrayView1<'_, f64>,
7246    weights: ArrayView1<'_, f64>,
7247    offset: ArrayView1<'_, f64>,
7248    resolvedspec: &TermCollectionSpec,
7249    best: &FittedTermCollection,
7250    family: LikelihoodSpec,
7251    options: &FitOptions,
7252    latent: &StandardLatentCoordConfig,
7253) -> Result<FittedTermCollectionWithSpec, EstimationError> {
7254    use gam_problem::{DeclaredHessianForm, Derivative, OuterEval};
7255    use gam_solve::rho_optimizer::OuterEvalOrder;
7256
7257    let rho_dim = best.fit.lambdas.len();
7258    let latent_flat_dim = latent.values.len();
7259    if latent_flat_dim == 0 {
7260        crate::bail_invalid_estim!(
7261            "latent-coordinate optimization requires a non-empty latent block"
7262        );
7263    }
7264    let direct_hypers =
7265        latent_coord_initial_direct_hypers(latent.values.id_mode(), latent.values.latent_dim())?;
7266    let analytic_rho_count = latent
7267        .analytic_penalties
7268        .as_ref()
7269        .map_or(0, |registry| registry.total_rho_count());
7270    let latent_coord_ext_dim = latent_flat_dim + analytic_rho_count + direct_hypers.len();
7271
7272    let mut theta0 = Array1::<f64>::zeros(rho_dim + latent_coord_ext_dim);
7273    theta0
7274        .slice_mut(s![..rho_dim])
7275        .assign(&best.fit.lambdas.mapv(f64::ln));
7276    theta0
7277        .slice_mut(s![rho_dim..rho_dim + latent_flat_dim])
7278        .assign(latent.values.as_flat());
7279    if !direct_hypers.is_empty() {
7280        let direct_start = rho_dim + latent_flat_dim + analytic_rho_count;
7281        theta0
7282            .slice_mut(s![direct_start..direct_start + direct_hypers.len()])
7283            .assign(&direct_hypers);
7284    }
7285
7286    let mut lower = Array1::<f64>::from_elem(theta0.len(), -12.0);
7287    let mut upper = Array1::<f64>::from_elem(theta0.len(), 12.0);
7288    let latent_bound = latent
7289        .values
7290        .as_flat()
7291        .iter()
7292        .fold(1.0_f64, |acc, &v| acc.max(v.abs()))
7293        + 10.0;
7294    for axis in rho_dim..rho_dim + latent_flat_dim {
7295        lower[axis] = -latent_bound;
7296        upper[axis] = latent_bound;
7297    }
7298    if let Some(registry) = latent.analytic_penalties.as_ref() {
7299        let (domain_lower, domain_upper) = registry
7300            .rho_domain_bounds()
7301            .map_err(EstimationError::InvalidInput)?;
7302        let start = rho_dim + latent_flat_dim;
7303        for local in 0..analytic_rho_count {
7304            lower[start + local] = lower[start + local].max(domain_lower[local]);
7305            upper[start + local] = upper[start + local].min(domain_upper[local]);
7306            if lower[start + local] >= upper[start + local] {
7307                return Err(EstimationError::InvalidInput(format!(
7308                    "analytic-penalty rho domain has no searchable interval at coordinate {local}: lower={}, upper={}",
7309                    lower[start + local],
7310                    upper[start + local]
7311                )));
7312            }
7313        }
7314    }
7315
7316    struct LatentJointContext<'d> {
7317        rho_dim: usize,
7318        cache: SingleBlockLatentCoordDesignCache,
7319        evaluator: gam_solve::estimate::ExternalJointHyperEvaluator<'d>,
7320    }
7321
7322    impl<'d> LatentJointContext<'d> {
7323        fn eval_full(
7324            &mut self,
7325            theta: &Array1<f64>,
7326            order: OuterEvalOrder,
7327        ) -> Result<(f64, Array1<f64>, gam_problem::HessianValue), EstimationError> {
7328            if let Some(eval) = self.cache.memoized_eval(theta) {
7329                return Ok(eval);
7330            }
7331            self.cache
7332                .ensure_theta(theta)
7333                .map_err(EstimationError::InvalidInput)?;
7334            let hyper_dirs = self
7335                .cache
7336                .hyper_dirs()
7337                .map_err(EstimationError::InvalidInput)?;
7338            let design_revision = Some(self.cache.design_revision());
7339            let registry_for_key = self.cache.analytic_penalties();
7340            self.evaluator
7341                .set_analytic_penalty_registry(registry_for_key.as_deref());
7342            let mut eval = evaluate_joint_reml_outer_eval_at_theta(
7343                &mut self.evaluator,
7344                self.cache.design(),
7345                theta,
7346                self.rho_dim,
7347                hyper_dirs,
7348                None,
7349                order,
7350                design_revision,
7351            )?;
7352            let latent = self.cache.latent().map_err(EstimationError::InvalidInput)?;
7353            if let Some(registry) = registry_for_key {
7354                let mut registry = registry.as_ref().clone();
7355                registry.apply_weight_schedules(
7356                    gam_solve::estimate::reml::outer_eval::current_outer_iter() as usize,
7357                );
7358                add_analytic_penalty_objective_to_eval(
7359                    theta,
7360                    self.rho_dim,
7361                    latent.as_ref(),
7362                    &registry,
7363                    &mut eval,
7364                )?;
7365            }
7366            add_latent_id_objective_to_eval(
7367                theta,
7368                self.rho_dim,
7369                self.cache.analytic_penalty_rho_count(),
7370                latent.as_ref(),
7371                &mut eval,
7372            )?;
7373            self.cache.store_eval(eval.clone());
7374            Ok(eval)
7375        }
7376
7377        fn eval_efs(
7378            &mut self,
7379            theta: &Array1<f64>,
7380        ) -> Result<gam_problem::EfsEval, EstimationError> {
7381            self.cache
7382                .ensure_theta(theta)
7383                .map_err(EstimationError::InvalidInput)?;
7384            let hyper_dirs = self
7385                .cache
7386                .hyper_dirs()
7387                .map_err(EstimationError::InvalidInput)?;
7388            let registry_for_key = self.cache.analytic_penalties();
7389            self.evaluator
7390                .set_analytic_penalty_registry(registry_for_key.as_deref());
7391            let mut efs = evaluate_joint_reml_efs_at_theta(
7392                &mut self.evaluator,
7393                self.cache.design(),
7394                theta,
7395                self.rho_dim,
7396                hyper_dirs,
7397                None,
7398                Some(self.cache.design_revision()),
7399            )?;
7400            if let Some(registry) = registry_for_key {
7401                let mut registry = registry.as_ref().clone();
7402                registry.apply_weight_schedules(
7403                    gam_solve::estimate::reml::outer_eval::current_outer_iter() as usize,
7404                );
7405                let latent = self.cache.latent().map_err(EstimationError::InvalidInput)?;
7406                let contribution = analytic_penalty_objective_contribution(
7407                    theta,
7408                    self.rho_dim,
7409                    latent.as_ref(),
7410                    &registry,
7411                )?;
7412                efs.cost += contribution.cost;
7413                if let (Some(psi_gradient), Some(psi_indices)) =
7414                    (efs.psi_gradient.as_mut(), efs.psi_indices.as_ref())
7415                {
7416                    if psi_gradient.len() != psi_indices.len() {
7417                        crate::bail_invalid_estim!(
7418                            "latent-coordinate analytic penalty EFS psi gradient length mismatch: gradient={}, indices={}",
7419                            psi_gradient.len(),
7420                            psi_indices.len()
7421                        );
7422                    }
7423                    for (local_idx, &theta_idx) in psi_indices.iter().enumerate() {
7424                        psi_gradient[local_idx] += contribution.gradient[theta_idx];
7425                    }
7426                }
7427            }
7428            Ok(efs)
7429        }
7430
7431        fn eval_cost(&mut self, theta: &Array1<f64>) -> f64 {
7432            if let Some(cost) = self.cache.memoized_cost(theta) {
7433                return cost;
7434            }
7435            if self.cache.ensure_theta(theta).is_err() {
7436                return f64::INFINITY;
7437            }
7438            let design_revision = Some(self.cache.design_revision());
7439            let registry_for_key = self.cache.analytic_penalties();
7440            self.evaluator
7441                .set_analytic_penalty_registry(registry_for_key.as_deref());
7442            let result = {
7443                let design = self.cache.design();
7444                self.evaluator.evaluate_cost_only(
7445                    &design.design,
7446                    &design.penalties,
7447                    &design.nullspace_dims,
7448                    design.linear_constraints.clone(),
7449                    theta,
7450                    self.rho_dim,
7451                    None,
7452                    "latent-coordinate-joint cost-only",
7453                    design_revision,
7454                )
7455            };
7456            match result {
7457                Ok(cost) => {
7458                    let latent = match self.cache.latent() {
7459                        Ok(latent) => latent,
7460                        Err(_) => return f64::INFINITY,
7461                    };
7462                    let contribution = match latent_id_objective_contribution(
7463                        theta,
7464                        self.rho_dim,
7465                        self.cache.analytic_penalty_rho_count(),
7466                        latent.as_ref(),
7467                    ) {
7468                        Ok(contribution) => contribution,
7469                        Err(_) => return f64::INFINITY,
7470                    };
7471                    let cost = cost + contribution.cost;
7472                    let cost = if let Some(registry) = registry_for_key {
7473                        let mut registry = registry.as_ref().clone();
7474                        registry.apply_weight_schedules(
7475                            gam_solve::estimate::reml::outer_eval::current_outer_iter() as usize,
7476                        );
7477                        match analytic_penalty_objective_contribution(
7478                            theta,
7479                            self.rho_dim,
7480                            latent.as_ref(),
7481                            &registry,
7482                        ) {
7483                            Ok(contribution) => cost + contribution.cost,
7484                            Err(_) => return f64::INFINITY,
7485                        }
7486                    } else {
7487                        cost
7488                    };
7489                    self.cache.store_cost(cost);
7490                    cost
7491                }
7492                Err(_) => f64::INFINITY,
7493            }
7494        }
7495    }
7496
7497    let effective_offset = best
7498        .design
7499        .compose_offset(offset, "latent-coordinate joint fit")
7500        .map_err(EstimationError::BasisError)?;
7501    let mut ctx = LatentJointContext {
7502        rho_dim,
7503        cache: SingleBlockLatentCoordDesignCache::new(
7504            data.to_owned(),
7505            resolvedspec.clone(),
7506            best.design.clone(),
7507            latent,
7508            rho_dim,
7509        )
7510        .map_err(EstimationError::InvalidInput)?,
7511        evaluator: gam_solve::estimate::ExternalJointHyperEvaluator::new(
7512            y,
7513            weights,
7514            &best.design.design,
7515            effective_offset.view(),
7516            &best.design.penalties,
7517            &external_opts_for_design(&family, &best.design, options),
7518            "latent-coordinate-joint",
7519        )?,
7520    };
7521    let registry_for_key = ctx.cache.analytic_penalties();
7522    ctx.evaluator
7523        .set_analytic_penalty_registry(registry_for_key.as_deref());
7524    ctx.evaluator
7525        .set_persistent_latent_values_fingerprint(latent.values.id_mode());
7526    if let Some(cached_t) = ctx
7527        .evaluator
7528        .load_persistent_latent_values(latent.values.n_obs(), latent.values.latent_dim())
7529    {
7530        let cached_t: Array2<f64> = cached_t;
7531        for (dst, src) in theta0
7532            .slice_mut(s![rho_dim..rho_dim + latent_flat_dim])
7533            .iter_mut()
7534            .zip(cached_t.iter())
7535        {
7536            *dst = *src;
7537        }
7538    }
7539
7540    let problem = exact_joint_multistart_outer_problem(
7541        &theta0,
7542        &lower,
7543        &upper,
7544        rho_dim,
7545        latent_coord_ext_dim,
7546        theta0.len(),
7547        Derivative::Analytic,
7548        DeclaredHessianForm::Unavailable,
7549        false,
7550        false,
7551        seed_risk_profile_for_likelihood_family(&family),
7552        options.tol,
7553        options.max_iter.max(1),
7554        Some(5.0),
7555        Some(0.5),
7556        None,
7557        // n-scaled profiled-criterion calibration (same absolute-gradient-floor
7558        // correction as the spatial paths; #1053 / #1066 / #1069).
7559        Some((data.nrows(), best.design.design.ncols().max(1))),
7560        // #1464: widen the over-smoothing ρ ceiling and seed the high-ρ probe
7561        // only when a constant-curvature curv() term is present in this fit.
7562        !constant_curvature_term_indices(resolvedspec).is_empty(),
7563        // Latent-coordinate optimization is not a profiled Matérn range solve.
7564        false,
7565    )?;
7566
7567    let eval_outer = |ctx: &mut &mut LatentJointContext<'_>,
7568                      theta: &Array1<f64>,
7569                      order: OuterEvalOrder|
7570     -> Result<OuterEval, EstimationError> {
7571        let (cost, gradient, hessian) = ctx.eval_full(theta, order)?;
7572        Ok(OuterEval {
7573            cost,
7574            gradient,
7575            hessian,
7576            inner_beta_hint: None,
7577        })
7578    };
7579
7580    let result = {
7581        let mut obj = problem.build_objective_with_eval_order(
7582            &mut ctx,
7583            |ctx: &mut &mut LatentJointContext<'_>, theta: &Array1<f64>| Ok(ctx.eval_cost(theta)),
7584            |ctx: &mut &mut LatentJointContext<'_>, theta: &Array1<f64>| {
7585                eval_outer(ctx, theta, OuterEvalOrder::ValueAndGradient)
7586            },
7587            |ctx: &mut &mut LatentJointContext<'_>, theta: &Array1<f64>, order: OuterEvalOrder| {
7588                eval_outer(ctx, theta, order)
7589            },
7590            Some(|ctx: &mut &mut LatentJointContext<'_>| {
7591                ctx.cache.reset();
7592            }),
7593            Some(|ctx: &mut &mut LatentJointContext<'_>, theta: &Array1<f64>| ctx.eval_efs(theta)),
7594        );
7595
7596        problem
7597            .run(&mut obj, "latent-coordinate joint REML")
7598            .map_err(|e| {
7599                EstimationError::InvalidInput(format!(
7600                    "latent-coordinate joint optimization failed after exhausting strategy fallbacks: {e}"
7601                ))
7602            })?
7603    };
7604    if !result.converged {
7605        crate::bail_invalid_estim!(
7606            "latent-coordinate joint optimization did not converge after {} iterations (final_objective={:.6e}, final_grad_norm={})",
7607            result.iterations,
7608            result.final_value,
7609            result.final_grad_norm_report(),
7610        );
7611    }
7612
7613    let theta_star = result.rho;
7614    let selected_lambdas = Array1::from_vec(
7615        gam_problem::checked_exp_log_strengths(
7616            theta_star.slice(s![..rho_dim]).iter().copied(),
7617        )
7618        .map_err(|error| {
7619            EstimationError::InvalidInput(format!(
7620                "selected latent-coordinate smoothing coordinate is outside the canonical log-strength domain: {error}"
7621            ))
7622        })?,
7623    );
7624    let mut final_data = data.to_owned();
7625    let flat_t = theta_star
7626        .slice(s![rho_dim..rho_dim + latent_flat_dim])
7627        .to_owned();
7628    let mut fitted_latent_values =
7629        Array2::<f64>::zeros((latent.values.n_obs(), latent.values.latent_dim()));
7630    for n in 0..latent.values.n_obs() {
7631        for axis in 0..latent.values.latent_dim() {
7632            let value = flat_t[n * latent.values.latent_dim() + axis];
7633            fitted_latent_values[[n, axis]] = value;
7634            final_data[[n, latent.feature_cols[axis]]] = value;
7635        }
7636    }
7637    let optimized = fit_term_collection_forspecwith_heuristic_lambdas(
7638        final_data.view(),
7639        y,
7640        weights,
7641        offset,
7642        resolvedspec,
7643        selected_lambdas.as_slice(),
7644        family,
7645        options,
7646    )?;
7647    ctx.evaluator
7648        .store_persistent_latent_values(&fitted_latent_values);
7649    let mut fit = optimized.fit;
7650    fit.reml_score = result.final_value;
7651    fit.penalized_objective = result.final_value;
7652    Ok(FittedTermCollectionWithSpec {
7653        fit,
7654        design: optimized.design,
7655        resolvedspec: resolvedspec.clone(),
7656        adaptive_diagnostics: optimized.adaptive_diagnostics,
7657        kappa_timing: None,
7658    })
7659}
7660
7661pub fn fit_term_collectionwith_latent_coord_optimization(
7662    data: ArrayView2<'_, f64>,
7663    y: Array1<f64>,
7664    weights: Array1<f64>,
7665    offset: Array1<f64>,
7666    spec: &TermCollectionSpec,
7667    latent: &StandardLatentCoordConfig,
7668    family: LikelihoodSpec,
7669    options: &FitOptions,
7670) -> Result<FittedTermCollectionWithSpec, EstimationError> {
7671    let n = data.nrows();
7672    if !(y.len() == n && weights.len() == n && offset.len() == n) {
7673        crate::bail_invalid_estim!(
7674            "fit_term_collectionwith_latent_coord_optimization row mismatch: n={}, y={}, weights={}, offset={}",
7675            n,
7676            y.len(),
7677            weights.len(),
7678            offset.len()
7679        );
7680    }
7681    let best = fit_term_collection_forspec(
7682        data,
7683        y.view(),
7684        weights.view(),
7685        offset.view(),
7686        spec,
7687        family.clone(),
7688        options,
7689    )?;
7690    let resolvedspec = freeze_term_collection_from_design(spec, &best.design)?;
7691    try_exact_joint_latent_coord_optimization(
7692        data,
7693        y.view(),
7694        weights.view(),
7695        offset.view(),
7696        &resolvedspec,
7697        &best,
7698        family,
7699        options,
7700        latent,
7701    )
7702}
7703
7704/// Resolve the two physically distinct isotropic Matérn range basins before the
7705/// local joint `[rho, psi]` solve.
7706///
7707/// The short/rich basin is represented by the ordinary cold fit at the
7708/// observation-density seed. The competing long-range basin has one canonical,
7709/// data-derived representative: the rotation-invariant fill distance of the
7710/// reduced-rank center set, `extent_rot / sqrt(k)`. Profile all smoothing
7711/// parameters at that endpoint once, and retain it only when its certified REML
7712/// objective is strictly lower. The subsequent joint optimizer therefore runs
7713/// exactly once, from the winning basin.
7714///
7715/// This is a closed endpoint comparison, not a lattice, sweep, or collection of
7716/// joint restarts. Both profiles pass through `fit_term_collection_forspec`, so
7717/// either produces a fully certified fit or the error is surfaced; there is no
7718/// best-effort fallback. Pairwise-distance bounds and the incumbent Matérn seed
7719/// are Euclidean invariants, so the decision is unchanged by rigid rotations.
7720fn select_isotropic_matern_range_basin(
7721    data: ArrayView2<'_, f64>,
7722    y: ArrayView1<'_, f64>,
7723    weights: ArrayView1<'_, f64>,
7724    offset: ArrayView1<'_, f64>,
7725    mut resolvedspec: TermCollectionSpec,
7726    mut best: FittedTermCollection,
7727    family: &LikelihoodSpec,
7728    options: &FitOptions,
7729    kappa_options: &SpatialLengthScaleOptimizationOptions,
7730    spatial_terms: &[usize],
7731) -> Result<(TermCollectionSpec, FittedTermCollection), EstimationError> {
7732    // Per-axis anisotropy and signed curvature have dedicated geometry
7733    // estimators. Their outer coordinates are not a scalar Matérn range and
7734    // therefore do not participate in this two-basin decision.
7735    if has_aniso_terms(&resolvedspec, spatial_terms)
7736        || !constant_curvature_term_indices(&resolvedspec).is_empty()
7737    {
7738        return Ok((resolvedspec, best));
7739    }
7740
7741    let mut best_score = fit_score(&best.fit);
7742    if !best_score.is_finite() {
7743        crate::bail_invalid_estim!(
7744            "isotropic Matérn basin selection received a non-finite incumbent profile"
7745        );
7746    }
7747
7748    for &term_idx in spatial_terms {
7749        let Some(SmoothBasisSpec::Matern {
7750            feature_cols,
7751            spec: matern,
7752            ..
7753        }) = resolvedspec
7754            .smooth_terms
7755            .get(term_idx)
7756            .map(|term| &term.basis)
7757        else {
7758            continue;
7759        };
7760        let num_centers = gam_terms::basis::center_strategy_num_centers(&matern.center_strategy)
7761            .ok_or_else(|| {
7762                EstimationError::InvalidInput(format!(
7763                    "resolved isotropic Matérn term {term_idx} has no finite center count"
7764                ))
7765            })?;
7766        let companion_length_scale = matern_low_rank_center_resolution_length_scale(
7767            data,
7768            feature_cols,
7769            num_centers,
7770        )
7771        .ok_or_else(|| {
7772            EstimationError::InvalidInput(format!(
7773                "resolved isotropic Matérn term {term_idx} has no finite center-resolution range"
7774            ))
7775        })?;
7776        let (psi_long_bound, psi_short_bound) =
7777            spatial_term_psi_bounds(data, &resolvedspec, term_idx, kappa_options)
7778                .map_err(EstimationError::BasisError)?;
7779        let psi_long = (-companion_length_scale.ln()).clamp(psi_long_bound, psi_short_bound);
7780        let long_length_scale = (-psi_long).exp();
7781        if !(long_length_scale.is_finite() && long_length_scale > 0.0) {
7782            crate::bail_invalid_estim!(
7783                "isotropic Matérn term {term_idx} produced an invalid long-range endpoint from psi={psi_long}"
7784            );
7785        }
7786        if get_spatial_length_scale(&resolvedspec, term_idx)
7787            .is_some_and(|current| current == long_length_scale)
7788        {
7789            continue;
7790        }
7791
7792        let mut endpoint_spec = resolvedspec.clone();
7793        set_spatial_length_scale(&mut endpoint_spec, term_idx, long_length_scale)?;
7794        // Profile rho at the competing geometry by starting the ordinary outer
7795        // optimizer literally at the already certified incumbent rho. This is
7796        // still a full standard REML
7797        // solve (including its ordinary seed certification), but it avoids
7798        // throwing away the exact smoothing optimum immediately before a
7799        // deliberately coarser center-resolution geometry move. The incumbent
7800        // lambdas provide the well-scaled starting chart needed for that profile
7801        // to reach its KKT certificate rather than exhausting its startup plans
7802        // a few ulps above stationarity.
7803        let endpoint = fit_term_collection_forspecwith_heuristic_lambdas(
7804            data,
7805            y,
7806            weights,
7807            offset,
7808            &endpoint_spec,
7809            best.fit.lambdas.as_slice(),
7810            family.clone(),
7811            options,
7812        )?;
7813        let endpoint_score = fit_score(&endpoint.fit);
7814        if !endpoint_score.is_finite() {
7815            crate::bail_invalid_estim!(
7816                "isotropic Matérn term {term_idx} long-range endpoint returned a non-finite profiled REML score"
7817            );
7818        }
7819
7820        if endpoint_score < best_score {
7821            log::info!(
7822                "[spatial-kappa] term {term_idx} selected certified long-range basin: \
7823                 length_scale={long_length_scale:.6}, profiled REML {endpoint_score:.6} \
7824                 < short-basin {best_score:.6}"
7825            );
7826            resolvedspec = freeze_term_collection_from_design(&endpoint_spec, &endpoint.design)?;
7827            best = endpoint;
7828            best_score = endpoint_score;
7829        } else {
7830            log::info!(
7831                "[spatial-kappa] term {term_idx} retained certified short-range basin: \
7832                 profiled REML {best_score:.6} <= long-endpoint {endpoint_score:.6} \
7833                 at length_scale={long_length_scale:.6}"
7834            );
7835        }
7836    }
7837
7838    Ok((resolvedspec, best))
7839}
7840
7841pub fn fit_term_collectionwith_spatial_length_scale_optimization(
7842    data: ArrayView2<'_, f64>,
7843    y: Array1<f64>,
7844    weights: Array1<f64>,
7845    offset: Array1<f64>,
7846    spec: &TermCollectionSpec,
7847    family: LikelihoodSpec,
7848    options: &FitOptions,
7849    kappa_options: &SpatialLengthScaleOptimizationOptions,
7850) -> Result<FittedTermCollectionWithSpec, EstimationError> {
7851    // Spatial hyperparameters change kernel geometry nonlinearly, so each
7852    // proposal rebuilds the spatial basis. Hybrid/isotropic terms expose a
7853    // scalar κ (= 1/length_scale); pure Duchon anisotropy exposes only
7854    // per-axis shape coordinates.
7855    //
7856    // When exact derivative information is available for the rebuilt basis and
7857    // penalty, kappa is promoted to a first-class outer hyperparameter beside
7858    // rho = log(lambda). In that mode this routine runs a joint outer solve in
7859    // theta = [rho, psi], where psi = log(kappa) = -log(length_scale), and the
7860    // optimizer is expected to consume a real joint Hessian. ARC is not meant
7861    // to run on a gradient-only surrogate here.
7862    //
7863    // Any eligible spatial smooth participates in this outer solve. If an
7864    // eligible spatial basis does not expose derivative information, that is
7865    // now a hard error.
7866    let mut resolvedspec = spec.clone();
7867    let spatial_terms = spatial_length_scale_term_indices(&resolvedspec);
7868    let n = data.nrows();
7869    if !(y.len() == n && weights.len() == n && offset.len() == n) {
7870        crate::bail_invalid_estim!(
7871            "fit_term_collectionwith_spatial_length_scale_optimization row mismatch: n={}, y={}, weights={}, offset={}",
7872            n,
7873            y.len(),
7874            weights.len(),
7875            offset.len()
7876        );
7877    }
7878    if !kappa_options.enabled || spatial_terms.is_empty() {
7879        let out = fit_term_collection_forspec(
7880            data,
7881            y.view(),
7882            weights.view(),
7883            offset.view(),
7884            &resolvedspec,
7885            family,
7886            options,
7887        )?;
7888        let resolvedspec = freeze_term_collection_from_design(&resolvedspec, &out.design)?;
7889        return Ok(FittedTermCollectionWithSpec {
7890            fit: out.fit,
7891            design: out.design,
7892            resolvedspec,
7893            adaptive_diagnostics: out.adaptive_diagnostics,
7894            kappa_timing: None,
7895        });
7896    }
7897    if kappa_options.max_outer_iter == 0 {
7898        crate::bail_invalid_estim!("spatial kappa optimization requires max_outer_iter >= 1");
7899    }
7900    if !(kappa_options.log_step.is_finite() && kappa_options.log_step > 0.0) {
7901        crate::bail_invalid_estim!("spatial kappa optimization requires log_step > 0");
7902    }
7903    if !(kappa_options.min_length_scale.is_finite()
7904        && kappa_options.max_length_scale.is_finite()
7905        && kappa_options.min_length_scale > 0.0
7906        && kappa_options.max_length_scale >= kappa_options.min_length_scale)
7907    {
7908        crate::bail_invalid_estim!(
7909            "spatial kappa optimization requires valid positive length_scale bounds"
7910        );
7911    }
7912
7913    let pilot_threshold = kappa_options.pilot_subsample_threshold;
7914    if pilot_threshold > 0 && n > pilot_threshold * 2 {
7915        log::info!(
7916            "[spatial-kappa] n={n} exceeds pilot threshold {}; using pilot geometry only for deterministic anisotropy initialization",
7917            pilot_threshold * 2,
7918        );
7919        apply_spatial_anisotropy_pilot_initializer(
7920            data,
7921            &mut resolvedspec,
7922            &spatial_terms,
7923            pilot_threshold,
7924            kappa_options,
7925        )?;
7926    }
7927
7928    // #1376: the geometry-only anisotropy seed (`initial_aniso_contrasts`, from
7929    // per-axis knot-coordinate spread) is blind to the response, so a signal
7930    // axis and a nuisance axis with equal coordinate spread both seed to ~0 and
7931    // the κ optimizer can stall at the symmetric point (it found a weak/flat
7932    // antisymmetric gradient, amplified by double-penalty nullspace shrinkage).
7933    // Add a bounded, response-aware per-axis nudge so the optimizer starts in
7934    // the correct basin. This runs whether or not the pilot initializer fired
7935    // (the pilot path is gated on a large-n threshold).
7936    apply_response_aware_anisotropy_seed(data, y.view(), &mut resolvedspec, &spatial_terms);
7937
7938    // Select every free constant-curvature coordinate once from its continuous,
7939    // analytically differentiated fair profile before fitting the baseline. The
7940    // subsequent joint solve profiles rho at this certified curvature.
7941    let free_curvature_terms: Vec<usize> = constant_curvature_term_indices(&resolvedspec)
7942        .into_iter()
7943        .filter(|&term_idx| !constant_curvature_kappa_is_fixed(&resolvedspec, term_idx))
7944        .collect();
7945    if !free_curvature_terms.is_empty() {
7946        validate_constant_curvature_fair_profile_inputs(weights.view(), offset.view(), &family)?;
7947    }
7948    for term_idx in free_curvature_terms {
7949        let kappa_hat = constant_curvature_kappa_fair_optimum(
7950            data,
7951            y.view(),
7952            &resolvedspec,
7953            term_idx,
7954            options,
7955        )?;
7956        if let Some(SmoothBasisSpec::ConstantCurvature { spec: cc, .. }) = resolvedspec
7957            .smooth_terms
7958            .get_mut(term_idx)
7959            .map(|term| &mut term.basis)
7960        {
7961            cc.kappa = kappa_hat;
7962        }
7963    }
7964
7965    let baseline_options = superseded_fit_options(options);
7966    let best = fit_term_collection_forspec(
7967        data,
7968        y.view(),
7969        weights.view(),
7970        offset.view(),
7971        &resolvedspec,
7972        family.clone(),
7973        &baseline_options,
7974    )?;
7975    resolvedspec = freeze_term_collection_from_design(&resolvedspec, &best.design)?;
7976    // The freeze step can rewrite a term's basis variant — most notably when
7977    // `build_thin_plate_basis_with_workspace` auto-promotes an infeasible
7978    // canonical-TPS request to a pure Duchon spline (length_scale = None,
7979    // no anisotropy). The pre-fit eligibility list was computed against the
7980    // ThinPlate spec, which has length_scale set, so it included that term.
7981    // After the rewrite the same term is a *pure* Duchon basis with no free
7982    // length-scale parameter to optimize, and the downstream kappa solver
7983    // (which assumes hybrid Duchon for log-κ derivatives) errors out. Refresh
7984    // the index list so it reflects the post-freeze spec.
7985    let spatial_terms = spatial_length_scale_term_indices(&resolvedspec);
7986    let (next_spec, best) = select_isotropic_matern_range_basin(
7987        data,
7988        y.view(),
7989        weights.view(),
7990        offset.view(),
7991        resolvedspec,
7992        best,
7993        &family,
7994        &baseline_options,
7995        kappa_options,
7996        &spatial_terms,
7997    )?;
7998    resolvedspec = next_spec;
7999    // Sync knot-cloud-derived aniso contrasts from the basis metadata back
8000    // into the spec so the optimizer starts from the geometry-informed η values
8001    // rather than the zero sentinel from --scale-dimensions.
8002    sync_aniso_contrasts_from_metadata(&mut resolvedspec, &best.design.smooth);
8003    if spatial_terms.is_empty() {
8004        let fitted = fit_term_collection_forspecwith_heuristic_lambdas(
8005            data,
8006            y.view(),
8007            weights.view(),
8008            offset.view(),
8009            &resolvedspec,
8010            best.fit.lambdas.as_slice(),
8011            family,
8012            options,
8013        )?;
8014        return Ok(FittedTermCollectionWithSpec {
8015            fit: fitted.fit,
8016            design: fitted.design,
8017            resolvedspec,
8018            adaptive_diagnostics: fitted.adaptive_diagnostics,
8019            kappa_timing: None,
8020        });
8021    }
8022    let initial_score = fit_score(&best.fit);
8023    if !initial_score.is_finite() {
8024        crate::bail_invalid_estim!(
8025            "spatial kappa optimization received a non-finite initial profiled score"
8026        );
8027    }
8028    let exact_joint = try_exact_joint_spatial_length_scale_optimization(
8029        data,
8030        y.view(),
8031        weights.view(),
8032        offset.view(),
8033        &resolvedspec,
8034        &best,
8035        family.clone(),
8036        options,
8037        kappa_options,
8038        &spatial_terms,
8039    )?
8040    .ok_or_else(|| {
8041        EstimationError::RemlOptimizationFailed(
8042            "spatial kappa optimization is unavailable for one or more eligible spatial terms"
8043                .to_string(),
8044        )
8045    })?;
8046    let exact_score = fit_score(&exact_joint.fit);
8047    let exact_joint = require_successful_spatial_optimization_result(
8048        initial_score,
8049        Ok(Some((exact_joint, exact_score))),
8050    )?;
8051
8052    log_spatial_aniso_scales(&exact_joint.resolvedspec);
8053    Ok(exact_joint)
8054}
8055
8056/// The end-to-end curvature-as-an-estimand report for one `curv(...)` smooth:
8057/// the fitted κ̂, its profile-likelihood confidence interval, the interior
8058/// κ = 0 likelihood-ratio flatness test, and the topology-free geometry
8059/// verdict. This is the #944 headline — it turns "we chose hyperbolic space"
8060/// into "κ̂ = −1.8 (95% CI −2.6, −1.1), flat rejected at p = …".
8061#[derive(Clone, Debug)]
8062pub struct CurvatureInference {
8063    /// Smooth-term index of the `curv(...)` term this report is about.
8064    pub term_idx: usize,
8065    /// The fitted signed sectional curvature κ̂ (the bounded analytic
8066    /// curvature-fair profile optimum).
8067    pub kappa_hat: f64,
8068    /// Profile-likelihood CI for κ and the geometry verdict from its sign.
8069    pub ci: gam_geometry::curvature_estimand::KappaProfileCi,
8070    /// Interior-point κ = 0 likelihood-ratio flatness test (full χ²₁, no
8071    /// half-χ² boundary correction — κ = 0 is an interior point of the
8072    /// `S^d ← ℝ^d → H^d` family).
8073    pub flatness: gam_geometry::curvature_estimand::FlatnessTest,
8074}
8075
8076/// Compute the #944 curvature inference for the constant-curvature smooth at
8077/// `term_idx`, given the already-fitted resolved spec (carrying κ̂) and the same
8078/// fit inputs used to produce it.
8079///
8080/// The point estimate and inference share the same continuously smoothing-
8081/// profiled curvature-fair evidence and its analytic profile score. Each CI
8082/// endpoint solves the Wilks likelihood-ratio equation directly inside the
8083/// chart-bound bracket with safeguarded Newton steps; bisection is the
8084/// guaranteed-progress fallback. A bound is reported as open only when the
8085/// analytic score certifies that the connected likelihood set containing κ̂
8086/// remains monotone all the way to that bound.
8087fn curvature_profile_lr_endpoint<F>(
8088    profile: &mut F,
8089    kappa_hat: f64,
8090    value_hat: f64,
8091    bound: f64,
8092    half_threshold: f64,
8093    x_tolerance: f64,
8094    score_tolerance: f64,
8095) -> Result<(f64, bool), String>
8096where
8097    F: FnMut(f64) -> Result<(f64, f64), String>,
8098{
8099    let direction = (bound - kappa_hat).signum();
8100    let span = (bound - kappa_hat).abs();
8101    if direction == 0.0 || span <= x_tolerance {
8102        return Ok((bound, true));
8103    }
8104
8105    let (bound_value, bound_score) = profile(bound)?;
8106    let outward_score = direction * bound_score;
8107    if outward_score < -score_tolerance {
8108        return Err(format!(
8109            "curvature profile is not outward-monotone at chart bound {bound}: \
8110             outward score {outward_score:.6e} is below tolerance {score_tolerance:.6e}"
8111        ));
8112    }
8113    let value_tolerance = score_tolerance * span;
8114    if bound_value < value_hat - value_tolerance {
8115        return Err(format!(
8116            "fitted curvature is not the minimum of its inference profile: \
8117             V(bound={bound})={bound_value:.6e} < V(kappa_hat)={value_hat:.6e}"
8118        ));
8119    }
8120    let bound_residual = bound_value - value_hat - half_threshold;
8121    if bound_residual < 0.0 {
8122        return Ok((bound, true));
8123    }
8124    if bound_residual == 0.0 {
8125        return Ok((bound, false));
8126    }
8127
8128    // `inside` is in the connected likelihood set and `outside` is beyond its
8129    // first threshold crossing. Newton uses the exact profile score. It is
8130    // accepted only in the central half of the current bracket, so every other
8131    // iteration is a bisection-quality contraction even on a nearly flat score.
8132    let mut inside_x = kappa_hat;
8133    let mut outside_x = bound;
8134    let mut outside_residual = bound_residual;
8135    let mut outside_score = bound_score;
8136    while (outside_x - inside_x).abs() > x_tolerance {
8137        let lo = inside_x.min(outside_x);
8138        let hi = inside_x.max(outside_x);
8139        let width = hi - lo;
8140        let central_lo = lo + 0.25 * width;
8141        let central_hi = hi - 0.25 * width;
8142        let newton = outside_x - outside_residual / outside_score;
8143        let probe = if newton.is_finite() && newton > central_lo && newton < central_hi {
8144            newton
8145        } else {
8146            lo + 0.5 * width
8147        };
8148        if !(probe > lo && probe < hi) {
8149            break;
8150        }
8151        let (value, score) = profile(probe)?;
8152        let outward_score = direction * score;
8153        if outward_score < -score_tolerance {
8154            return Err(format!(
8155                "curvature profile changed direction before its likelihood crossing at \
8156                 kappa={probe}: outward score {outward_score:.6e} is below tolerance \
8157                 {score_tolerance:.6e}"
8158            ));
8159        }
8160        let residual = value - value_hat - half_threshold;
8161        if residual >= 0.0 {
8162            outside_x = probe;
8163            outside_residual = residual;
8164            outside_score = score;
8165        } else {
8166            inside_x = probe;
8167        }
8168    }
8169    Ok((inside_x + 0.5 * (outside_x - inside_x), false))
8170}
8171
8172fn curvature_profile_ci_from_analytic_score<F>(
8173    profile: &mut F,
8174    kappa_hat: f64,
8175    kappa_min: f64,
8176    kappa_max: f64,
8177    level: f64,
8178    relative_tolerance: f64,
8179) -> Result<gam_geometry::curvature_estimand::KappaProfileCi, String>
8180where
8181    F: FnMut(f64) -> Result<(f64, f64), String>,
8182{
8183    if !(kappa_min < kappa_max && kappa_hat >= kappa_min && kappa_hat <= kappa_max) {
8184        return Err("curvature profile requires kappa_hat inside valid chart bounds".to_string());
8185    }
8186    if !(level > 0.0 && level < 1.0) {
8187        return Err("curvature profile level must lie in (0, 1)".to_string());
8188    }
8189    let z = gam_geometry::curvature_estimand::wald_half_width(1.0, level)
8190        .ok_or_else(|| "curvature profile threshold is not finite".to_string())?;
8191    let half_threshold = 0.5 * z * z;
8192    let (value_hat, score_hat) = profile(kappa_hat)?;
8193    let relative_tolerance = relative_tolerance.max(f64::EPSILON.sqrt());
8194    let x_tolerance = relative_tolerance * (1.0 + kappa_min.abs().max(kappa_max.abs()));
8195    let score_tolerance = relative_tolerance * (1.0 + value_hat.abs());
8196    let at_lower = (kappa_hat - kappa_min).abs() <= x_tolerance;
8197    let at_upper = (kappa_hat - kappa_max).abs() <= x_tolerance;
8198    let stationary = if at_lower {
8199        score_hat >= -score_tolerance
8200    } else if at_upper {
8201        score_hat <= score_tolerance
8202    } else {
8203        score_hat.abs() <= score_tolerance
8204    };
8205    if !stationary {
8206        return Err(format!(
8207            "curvature inference rejected a non-stationary point estimate: \
8208             kappa_hat={kappa_hat}, score={score_hat:.6e}, \
8209             stationarity_bound={score_tolerance:.6e}"
8210        ));
8211    }
8212
8213    let (ci_lo, lo_at_bound) = curvature_profile_lr_endpoint(
8214        profile,
8215        kappa_hat,
8216        value_hat,
8217        kappa_min,
8218        half_threshold,
8219        x_tolerance,
8220        score_tolerance,
8221    )?;
8222    let (ci_hi, hi_at_bound) = curvature_profile_lr_endpoint(
8223        profile,
8224        kappa_hat,
8225        value_hat,
8226        kappa_max,
8227        half_threshold,
8228        x_tolerance,
8229        score_tolerance,
8230    )?;
8231    let verdict = if ci_lo > 0.0 {
8232        gam_geometry::curvature_estimand::CurvatureVerdict::Spherical
8233    } else if ci_hi < 0.0 {
8234        gam_geometry::curvature_estimand::CurvatureVerdict::Hyperbolic
8235    } else {
8236        gam_geometry::curvature_estimand::CurvatureVerdict::Flat
8237    };
8238    Ok(gam_geometry::curvature_estimand::KappaProfileCi {
8239        kappa_hat,
8240        ci_lo,
8241        ci_hi,
8242        lo_at_bound,
8243        hi_at_bound,
8244        verdict,
8245    })
8246}
8247
8248pub fn curvature_inference_forspec(
8249    data: ArrayView2<'_, f64>,
8250    y: ArrayView1<'_, f64>,
8251    weights: ArrayView1<'_, f64>,
8252    offset: ArrayView1<'_, f64>,
8253    resolvedspec: &TermCollectionSpec,
8254    term_idx: usize,
8255    family: LikelihoodSpec,
8256    options: &FitOptions,
8257    level: f64,
8258) -> Result<CurvatureInference, EstimationError> {
8259    let kappa_hat = get_constant_curvature_kappa(resolvedspec, term_idx).ok_or_else(|| {
8260        EstimationError::InvalidInput(format!(
8261            "curvature_inference_forspec: term {term_idx} is not a constant-curvature smooth"
8262        ))
8263    })?;
8264    if constant_curvature_kappa_is_fixed(resolvedspec, term_idx) {
8265        crate::bail_invalid_estim!(
8266            "curvature inference requires an estimated curvature; term {term_idx} has user-pinned kappa={kappa_hat}"
8267        );
8268    }
8269    if y.len() != data.nrows() || weights.len() != data.nrows() || offset.len() != data.nrows() {
8270        crate::bail_invalid_estim!(
8271            "curvature inference row mismatch: data={}, y={}, weights={}, offset={}",
8272            data.nrows(),
8273            y.len(),
8274            weights.len(),
8275            offset.len(),
8276        );
8277    }
8278    validate_constant_curvature_fair_profile_inputs(weights, offset, &family)?;
8279    let (kappa_min, kappa_max) = constant_curvature_kappa_bounds(data, resolvedspec, term_idx);
8280    let (feature_cols, base_spec) = match resolvedspec
8281        .smooth_terms
8282        .get(term_idx)
8283        .map(|term| &term.basis)
8284    {
8285        Some(SmoothBasisSpec::ConstantCurvature {
8286            feature_cols, spec, ..
8287        }) => (feature_cols, spec.clone()),
8288        _ => {
8289            return Err(EstimationError::InvalidInput(format!(
8290                "constant-curvature κ profile: smooth term {term_idx} is not a \
8291                 constant-curvature basis"
8292            )));
8293        }
8294    };
8295    let x_term = select_columns(data, feature_cols).map_err(EstimationError::from)?;
8296    let radial_reference = constant_curvature_radial_reference(x_term.view(), y)?;
8297    let fair_profile = ConstantCurvatureFairProfile {
8298        data: x_term.view(),
8299        response: y,
8300        radial_reference,
8301        spec: base_spec,
8302        cache: std::cell::RefCell::new(std::collections::HashMap::new()),
8303    };
8304
8305    // CI and flatness revisit κ̂ and κ=0. The shared profile caches each joint
8306    // value/analytic-score pair so every statistic consumes the same evaluation.
8307    let mut v_p = |kappa: f64| -> Result<(f64, f64), String> {
8308        if !kappa.is_finite() {
8309            return Err(format!("V_p probed a non-finite κ = {kappa}"));
8310        }
8311        let sample = fair_profile.evaluate(kappa).map_err(|error| {
8312            format!("analytic curvature profile at kappa={kappa} failed: {error}")
8313        })?;
8314        Ok(sample)
8315    };
8316    let ci = curvature_profile_ci_from_analytic_score(
8317        &mut v_p,
8318        kappa_hat,
8319        kappa_min,
8320        kappa_max,
8321        level,
8322        options.tol,
8323    )
8324    .map_err(EstimationError::RemlOptimizationFailed)?;
8325    let flatness = gam_geometry::curvature_estimand::flatness_lr_test(
8326        |kappa| v_p(kappa).map(|(value, _)| value),
8327        kappa_hat,
8328    )
8329    .map_err(EstimationError::RemlOptimizationFailed)?;
8330
8331    Ok(CurvatureInference {
8332        term_idx,
8333        kappa_hat,
8334        ci,
8335        flatness,
8336    })
8337}
8338
8339#[cfg(test)]
8340mod curvature_profile_score_tests {
8341    use super::*;
8342
8343    #[test]
8344    fn analytic_profile_score_finds_exact_quadratic_lr_crossings() {
8345        let kappa_hat = -0.37;
8346        let curvature = 16.0;
8347        let level = 0.95;
8348        let mut profile = |kappa: f64| -> Result<(f64, f64), String> {
8349            let displacement = kappa - kappa_hat;
8350            Ok((
8351                7.0 + 0.5 * curvature * displacement * displacement,
8352                curvature * displacement,
8353            ))
8354        };
8355        let ci = curvature_profile_ci_from_analytic_score(
8356            &mut profile,
8357            kappa_hat,
8358            -3.0,
8359            3.0,
8360            level,
8361            1.0e-10,
8362        )
8363        .expect("analytic quadratic profile CI");
8364        let z = gam_geometry::curvature_estimand::wald_half_width(1.0, level)
8365            .expect("valid normal quantile");
8366        let expected_half_width = z / curvature.sqrt();
8367        assert!((ci.ci_lo - (kappa_hat - expected_half_width)).abs() <= 1.0e-8);
8368        assert!((ci.ci_hi - (kappa_hat + expected_half_width)).abs() <= 1.0e-8);
8369        assert!(!ci.lo_at_bound && !ci.hi_at_bound);
8370    }
8371
8372    #[test]
8373    fn analytic_profile_marks_chart_bound_when_wilks_set_never_crosses() {
8374        let mut profile =
8375            |kappa: f64| -> Result<(f64, f64), String> { Ok((0.5 * kappa * kappa, kappa)) };
8376        let ci =
8377            curvature_profile_ci_from_analytic_score(&mut profile, 0.0, -0.1, 0.1, 0.95, 1.0e-10)
8378                .expect("open bounded profile CI");
8379        assert_eq!(ci.ci_lo, -0.1);
8380        assert_eq!(ci.ci_hi, 0.1);
8381        assert!(ci.lo_at_bound && ci.hi_at_bound);
8382    }
8383}
8384
8385/// Provenance tag for the smooth-term significance correction (#1063): which
8386/// statistic the reported p-value is built from.
8387#[derive(Clone, Copy, Debug, PartialEq, Eq)]
8388pub enum SmoothLrCorrection {
8389    /// A per-term LR statistic corrected by the full estimated-λ Lawley factor,
8390    /// including the ρ̂-sampling-variation contribution from the regularized
8391    /// inverse REML/LAML outer Hessian.
8392    LawleyLrEstimatedLambda,
8393    /// A per-term likelihood-ratio statistic `W = 2(ℓ_full − ℓ_null)` that has
8394    /// been Bartlett-corrected with the fixed-λ Lawley factor `c = E[W|λ]/d`
8395    /// (`W* = W/c`, referenced against `χ²_d`). This is used only when the
8396    /// estimated-λ handoff is unavailable.
8397    LawleyLrFixedLambda,
8398    /// No second-order correction was applied — either the family has no
8399    /// closed-form Lawley cumulant jets or the null refit did not converge — so
8400    /// the uncorrected `χ²_d` of the raw LR statistic stands.
8401    None,
8402}
8403
8404impl SmoothLrCorrection {
8405    /// The serialized provenance label surfaced in the summary table.
8406    pub fn label(self) -> &'static str {
8407        match self {
8408            SmoothLrCorrection::LawleyLrEstimatedLambda => "lawley_lr_estimated_lambda",
8409            SmoothLrCorrection::LawleyLrFixedLambda => "lawley_lr_fixed_lambda",
8410            SmoothLrCorrection::None => "none",
8411        }
8412    }
8413}
8414
8415/// The Bartlett-corrected per-term significance report for one penalized smooth
8416/// term (#1063). Unlike the summary table's Wood rank-truncated **Wald**
8417/// statistic, this is a genuine **likelihood-ratio** statistic from a
8418/// constrained refit (the smooth dropped), so the exact Lawley LR Bartlett
8419/// factor corrects the right quantity.
8420#[derive(Clone, Debug)]
8421pub struct SmoothTermLrInference {
8422    /// Smooth-term name (matches the summary row).
8423    pub name: String,
8424    /// Smooth-term index within `resolvedspec.smooth_terms`.
8425    pub term_idx: usize,
8426    /// The uncorrected likelihood-ratio statistic `W = 2(ℓ_full − ℓ_null)`,
8427    /// floored at zero (a non-negative LR by construction).
8428    pub statistic_lr: f64,
8429    /// Reference degrees of freedom `d` (the Wood truncation `tr(F)²/tr(F²)` on
8430    /// the term's influence block, falling back to the term EDF).
8431    pub ref_df: f64,
8432    /// Lawley LR Bartlett factor `c = E[W]/d = 1 + Δε/d` when computable, else
8433    /// `1.0` (no correction).
8434    pub bartlett_factor: f64,
8435    /// Fixed-λ conditional factor `c_cond = 1 + Δε(ρ̂)/d` when the estimated-λ
8436    /// correction was applied. `None` means the applied factor was either the
8437    /// fixed-λ factor itself or no Lawley correction was available.
8438    pub bartlett_factor_conditional: Option<f64>,
8439    /// Increment in Lawley's LR mean shift due solely to ρ̂ sampling variation,
8440    /// `0.5 * tr(H_Δε Cov(ρ̂))`, when estimated-λ correction was applied.
8441    pub rho_variation_shift: Option<f64>,
8442    /// Bartlett-corrected statistic `W* = W / c`.
8443    pub statistic_corrected: f64,
8444    /// Uncorrected p-value `P(χ²_d > W)`.
8445    pub p_value_uncorrected: f64,
8446    /// Corrected p-value `P(χ²_d > W*)`; equals the uncorrected value when no
8447    /// correction was applied.
8448    pub p_value_corrected: f64,
8449    /// Whether the second-order correction is **material** (#939 deliverable 4):
8450    /// the per-test diagnostic "is `n` too small for first-order inference
8451    /// *here*?". `true` when a correction was applied and it moves the result by
8452    /// more than [`SMOOTH_LR_MATERIAL_THRESHOLD`] — measured as the larger of the
8453    /// relative Bartlett-factor distance from one `|c − 1|` and the relative
8454    /// p-value change `|p* − p| / max(p, p*, ε)`. `false` when `correction` is
8455    /// [`SmoothLrCorrection::None`] (no correction was applied).
8456    pub material: bool,
8457    /// Which statistic the corrected p-value is built from.
8458    pub correction: SmoothLrCorrection,
8459}
8460
8461/// The materiality threshold for [`SmoothTermLrInference::material`] (#939
8462/// deliverable 4): a correction is flagged material when it changes the result
8463/// by more than 10%.
8464pub const SMOOTH_LR_MATERIAL_THRESHOLD: f64 = 0.10;
8465
8466/// Build `S_b = lambda_b * S_b^unit` as global `p_total x p_total` matrices in
8467/// exactly the fitted rho/lambda ordering. This is the narrow handoff the
8468/// estimated-lambda Lawley correction needs: the same `design.penalties` order
8469/// already paired with `fit.lambdas`, without changing #740's outer-Hessian
8470/// algebra or the production penalty assembly.
8471fn fitted_rho_penalty_components(
8472    penalties: &[BlockwisePenalty],
8473    lambdas: &[f64],
8474    p_total: usize,
8475) -> Result<Vec<gam_terms::inference::lawley::RhoPenaltyComponent>, EstimationError> {
8476    if penalties.len() != lambdas.len() {
8477        return Err(EstimationError::InvalidInput(format!(
8478            "smooth_term_lr_inference: penalty/lambda count mismatch ({} penalties, {} lambdas)",
8479            penalties.len(),
8480            lambdas.len()
8481        )));
8482    }
8483    let mut components = Vec::with_capacity(penalties.len());
8484    for (idx, (penalty, &lambda)) in penalties.iter().zip(lambdas.iter()).enumerate() {
8485        if !(lambda.is_finite() && lambda >= 0.0) {
8486            return Err(EstimationError::InvalidInput(format!(
8487                "smooth_term_lr_inference: lambda[{idx}] is invalid: {lambda}"
8488            )));
8489        }
8490        let r = &penalty.col_range;
8491        if r.end > p_total {
8492            return Err(EstimationError::InvalidInput(format!(
8493                "smooth_term_lr_inference: penalty[{idx}] range {:?} exceeds coefficient dimension {p_total}",
8494                r
8495            )));
8496        }
8497        let mut s_component = Array2::<f64>::zeros((p_total, p_total));
8498        s_component
8499            .slice_mut(s![r.start..r.end, r.start..r.end])
8500            .scaled_add(lambda, &penalty.local);
8501        components.push(gam_terms::inference::lawley::RhoPenaltyComponent { s_component });
8502    }
8503    Ok(components)
8504}
8505
8506/// The end-to-end per-term likelihood-ratio significance report for every
8507/// penalized (shape-unconstrained) smooth term in a fitted model, magically
8508/// Bartlett-corrected when the family carries closed-form Lawley cumulant jets
8509/// (#1063, follow-up to #939).
8510///
8511/// # Why an LR statistic (not the summary Wald)
8512///
8513/// The summary table's `wood_smooth_test` is Wood's rank-truncated **Wald**
8514/// statistic `T = β̂'Σ̂⁻β̂`. Lawley's ε corrects the **likelihood-ratio**
8515/// statistic, and under penalization the Wald form is already a weighted χ²
8516/// whose second-order mean is *not* `d + Δε` — dividing `T` by the LR factor
8517/// would correct the wrong statistic. The principled route (#1063 Option 1) is
8518/// to compute a real per-term LR statistic by a constrained refit and correct
8519/// *that*:
8520///
8521/// ```text
8522/// W = 2(ℓ_full − ℓ_null),   W* = W / c,   c = 1 + Δε/d,   p = P(χ²_d > W*).
8523/// ```
8524///
8525/// # Method
8526///
8527/// 1. Fit the full model and read `ℓ_full` and the per-term coefficient ranges /
8528///    EDF / influence block. The full design's column layout fixes the tested
8529///    block for the Lawley factor.
8530/// 2. For each penalized smooth term, refit a null model with that term dropped
8531///    from the spec; `W = max(2(ℓ_full − ℓ_null), 0)`.
8532/// 3. The reference d.f. `d` is the Wood truncation `tr(F)²/tr(F²)` on the
8533///    term's influence block (the same `ref_df` the summary Wald row reports),
8534///    floored at `max(edf, null_dim, 1)`: this LR test drops the whole term, so
8535///    `d` is at least the dimension the term spans when present (its null-space
8536///    dimension, never below 1). The non-symmetric `tr(F²)` can collapse toward
8537///    0 at a shrunk-to-null fit and violate that bound — see the inline note at
8538///    the `ref_df` binding.
8539/// 4. When the family has closed-form cumulant jets, evaluate Lawley's ε at the
8540///    **null** linear predictor (an expectation evaluated at the null fit), fold
8541///    the full λ-scaled penalty `S_λ` into the information, and Bartlett-correct
8542///    `W` with [`gam_terms::inference::lawley::lawley_lr_bartlett_factor`]. The
8543///    null annihilates the tested block's penalty (`S_λ β₀ = 0` on that block),
8544///    so the penalized Lawley expansion applies verbatim.
8545/// 5. Otherwise (no closed-form jets, or a null refit that did not converge) the
8546///    uncorrected `χ²_d` stands with provenance `none` — never weakened.
8547///
8548/// Random-effect smooths and shape-constrained smooths are skipped (their tests
8549/// are not a central-χ² LR), matching the summary table's policy.
8550pub fn smooth_term_lr_inference_forspec(
8551    data: ArrayView2<'_, f64>,
8552    y: ArrayView1<'_, f64>,
8553    weights: ArrayView1<'_, f64>,
8554    offset: ArrayView1<'_, f64>,
8555    resolvedspec: &TermCollectionSpec,
8556    family: LikelihoodSpec,
8557    options: &FitOptions,
8558) -> Result<Vec<SmoothTermLrInference>, EstimationError> {
8559    use gam_terms::inference::lawley::{
8560        LAWLEY_PAIR_MATRIX_MAX_ROWS, known_scale_expected_jets_with_dispersion,
8561        lawley_lr_bartlett_factor, lawley_lr_mean_shift_with_rho_variation,
8562    };
8563
8564    let n = data.nrows();
8565    // Full fit: ℓ_full, the per-term coefficient ranges/EDF/influence, and the
8566    // full design whose column layout fixes each tested block for Lawley.
8567    let full = fit_term_collection_forspec(
8568        data,
8569        y,
8570        weights,
8571        offset,
8572        resolvedspec,
8573        family.clone(),
8574        options,
8575    )?;
8576    let ll_full = full.fit.log_likelihood;
8577    let p_total = full.design.design.ncols();
8578    let lambdas = full.fit.lambdas.as_slice().ok_or_else(|| {
8579        EstimationError::InvalidInput(
8580            "smooth_term_lr_inference: non-contiguous lambda vector".to_string(),
8581        )
8582    })?;
8583    let s_lambda = weighted_blockwise_penalty_sum(&full.design.penalties, lambdas, p_total);
8584    let rho_penalty_components =
8585        fitted_rho_penalty_components(&full.design.penalties, lambdas, p_total)?;
8586    let rho_covariance = full.fit.artifacts.rho_covariance.as_ref().filter(|cov| {
8587        cov.nrows() == rho_penalty_components.len() && cov.ncols() == rho_penalty_components.len()
8588    });
8589    // Full design as a dense n×p array for the Lawley pair-matrix reduction.
8590    let full_design_dense = full.design.design.to_dense();
8591    let influence = full.fit.coefficient_influence();
8592    let fitted_likelihood = resolved_likelihood_for_fit(&full.fit)?;
8593    let family_disp = lawley_dispersion_for_family(&fitted_likelihood, &full.fit)?;
8594    let coefficient_covariance_scale = fitted_likelihood
8595        .coefficient_covariance_scale(family_disp)
8596        .map_err(|error| EstimationError::InvalidInput(error.to_string()))?;
8597
8598    let mut out = Vec::<SmoothTermLrInference>::new();
8599    for (term_idx, design_term) in full.design.smooth.terms.iter().enumerate() {
8600        let penalty_range = full
8601            .design
8602            .smooth_term_penalty_range(term_idx)
8603            .map_err(EstimationError::InvalidInput)?;
8604        let (block_start, k) = penalty_range
8605            .map(|range| (range.start, range.len()))
8606            .unwrap_or((0, 0));
8607        // Shape-constrained smooths get no central-χ² LR (cone-projected
8608        // boundary test); the summary table skips them too.
8609        if design_term.shape != ShapeConstraint::None {
8610            continue;
8611        }
8612        let coeff_range = design_term.coeff_range.clone();
8613        if coeff_range.start >= coeff_range.end || coeff_range.end > p_total {
8614            continue;
8615        }
8616        // Per-term EDF for the χ² reference df FALLBACK (used only when the
8617        // influence matrix `F` is unavailable). Route through `per_term_edf`,
8618        // which uses the ADDITIVE per-block trace channel
8619        // (`|coeff_range| − Σ_{kk∈term} tr_kk`) and caps at the model total,
8620        // rather than the raw `edf_by_block` block-sum `Σ_{kk}(rank_kk − tr_kk)`.
8621        // For a multi-penalty term (te/ti/double-penalty) the penalties share one
8622        // coefficient range, so the rank-based block-sum OVER-COUNTS the term EDF
8623        // (Σ rank_kk > |coeff_range|) and would inflate the LR reference df,
8624        // biasing the smooth-term test conservative on large/sparse fits where `F`
8625        // is not materialised. (Same per-block over-count class as the multinomial
8626        // `edf_per_class` fix.)
8627        let edf = full.fit.per_term_edf(coeff_range.clone(), block_start, k);
8628        // The term's **joint** unpenalized null-space dimension: the coefficient
8629        // directions penalized by *no* active penalty — the polynomial part a
8630        // penalized smooth always carries when present, which no penalty can
8631        // shrink. This is `dim(∩_k null(S_k)) = p_local − rank(Σ_k S_k)`, the
8632        // INTERSECTION of the per-penalty null spaces, computed by
8633        // `wald_unpenalized_dim()` — the very same scalar the summary Wald test
8634        // (`wood_smooth_test`) floors its reference d.f. at, so the LR and Wald
8635        // tests reference a consistent d.f.
8636        //
8637        // It must NOT be `nullspace_dims.iter().sum()`: that *unions* the null
8638        // spaces (the #1360 defect — see `joint_unpenalized_dim`'s docs). A
8639        // double-penalty smooth carries a bending penalty (null space = its
8640        // polynomial part) plus a complementary null-space ridge (which penalizes
8641        // exactly that polynomial part), so the two null spaces are disjoint and
8642        // the joint null space is EMPTY (dim 0) — yet the per-penalty dims sum to
8643        // ~`p_local`. Flooring `ref_df` at that sum pins it to the full basis
8644        // dimension for every fit (e.g. 11 for a k=12 s(x)), making the LR test
8645        // badly conservative for genuine moderate signals while only accidentally
8646        // masking the collapse.
8647        let null_dim = design_term.wald_unpenalized_dim();
8648        // χ² reference d.f. for the whole-term LR test. The statistic W tests the
8649        // term present vs entirely absent, so the reference d.f. must be at least
8650        // the dimension the term spans when present — its joint null-space
8651        // dimension (`null_dim`), the effective d.f. it uses in the fit (`edf`),
8652        // and never below 1 (you cannot test "is this function present" with
8653        // fewer than one degree of freedom). The primary reference is Wood's
8654        // smoothing-selection-corrected `edf1 = 2·tr(F) − tr(F²)`
8655        // (`wood_reference_df`), which dominates in calibrated and high-power
8656        // fits and removes the post-selection left-tail anti-conservatism the raw
8657        // `edf` reference leaves behind. The floor guards the DEGENERATE collapse:
8658        // as REML shrinks a term fully onto its null space `edf1 → edf → ~0` and
8659        // the non-symmetric `F = H⁻¹X'WX` can make the block `tr(F²)` numerically
8660        // unreliable, so `wood_reference_df` returns `None` and the floor supplies
8661        // `max(edf, null_dim, 1)`. Without a floor, a positive `W` referenced
8662        // against `χ²_{~0}` would report a flat, shrunk-to-null term as MAXIMALLY
8663        // significant (`p ~ 1e-12`) — a Type-I error decided by a degenerate
8664        // reference d.f., not the data (#1766). The floor binds ONLY on that
8665        // collapse; `edf` and `null_dim` never exceed `edf1` in a healthy fit.
8666        // This mirrors the summary Wald path, which floors its reference at the
8667        // statistic's own rank for the same reason.
8668        let rho_uncertainty_df = match wps_block_uncertainty_df(
8669            full.fit.weighted_gram(),
8670            full.fit.smoothing_correction(),
8671            &coeff_range,
8672            coefficient_covariance_scale,
8673        )? {
8674            // Missing correction artifacts explicitly mean that this optional
8675            // WPS term was not computed; the base Wood reference df remains.
8676            Some(extra_df) => extra_df,
8677            None => 0.0,
8678        };
8679        let ref_df = (wood_reference_df(influence, &coeff_range)
8680            .unwrap_or(0.0)
8681            .max(edf)
8682            + rho_uncertainty_df)
8683            .max(null_dim as f64)
8684            .max(1.0);
8685        if !(ref_df.is_finite() && ref_df > 0.0) {
8686            continue;
8687        }
8688
8689        // Null model: drop this smooth term from the spec and refit. The term's
8690        // name pins which spec entry to remove (design and spec share names).
8691        let mut null_spec = resolvedspec.clone();
8692        let Some(spec_pos) = null_spec
8693            .smooth_terms
8694            .iter()
8695            .position(|t| t.name == design_term.name)
8696        else {
8697            continue;
8698        };
8699        null_spec.smooth_terms.remove(spec_pos);
8700        let null_fit = fit_term_collection_forspec(
8701            data,
8702            y,
8703            weights,
8704            offset,
8705            &null_spec,
8706            family.clone(),
8707            options,
8708        );
8709        let (statistic_lr, eta_null) = match null_fit {
8710            Ok(null) if null.fit.log_likelihood.is_finite() => {
8711                let w = (2.0 * (ll_full - null.fit.log_likelihood)).max(0.0);
8712                // η at the null fit: X_null β_null + affine_offset + offset
8713                // (per-row linear predictor; design-layout independent — Lawley
8714                // reads it on the full design rows). `compose_offset` folds the
8715                // design's fixed affine channel (non-zero endpoint anchor,
8716                // #2297) into the user offset.
8717                let null_offset = null
8718                    .design
8719                    .compose_offset(offset, "smooth likelihood-ratio null model")
8720                    .map_err(|error| EstimationError::InvalidInput(error.to_string()))?;
8721                let mut eta = null.design.design.dot(&null.fit.beta);
8722                eta += &null_offset;
8723                (w, Some(eta))
8724            }
8725            _ => (f64::NAN, None),
8726        };
8727
8728        let chi2 = statrs::distribution::ChiSquared::new(ref_df).ok();
8729        let p_uncorrected = match (chi2.as_ref(), statistic_lr.is_finite()) {
8730            (Some(dist), true) => {
8731                use statrs::distribution::ContinuousCDF;
8732                (1.0 - dist.cdf(statistic_lr)).clamp(0.0, 1.0)
8733            }
8734            _ => f64::NAN,
8735        };
8736
8737        // Magic Bartlett correction: only when the LR statistic is finite, the
8738        // family has closed-form jets, n is in the resolvable regime, and the
8739        // factor is computable. Otherwise the uncorrected χ² stands.
8740        let mut bartlett_factor = 1.0;
8741        let mut bartlett_factor_conditional = None;
8742        let mut rho_variation_shift = None;
8743        let mut statistic_corrected = statistic_lr;
8744        let mut p_corrected = p_uncorrected;
8745        let mut correction = SmoothLrCorrection::None;
8746        if let (Some(eta), true, true) = (
8747            eta_null.as_ref(),
8748            statistic_lr.is_finite(),
8749            n <= LAWLEY_PAIR_MATRIX_MAX_ROWS,
8750        ) {
8751            let kappas: Option<Vec<_>> = (0..n)
8752                .map(|i| {
8753                    known_scale_expected_jets_with_dispersion(
8754                        &fitted_likelihood.spec,
8755                        eta[i],
8756                        family_disp,
8757                    )
8758                    .and_then(|jets| jets.kappas().ok())
8759                })
8760                .collect();
8761            if let (Some(kappas), Some(dist)) = (kappas, chi2.as_ref()) {
8762                let fixed_factor = lawley_lr_bartlett_factor(
8763                    full_design_dense.view(),
8764                    &kappas,
8765                    Some(s_lambda.view()),
8766                    coeff_range.clone(),
8767                    ref_df,
8768                );
8769                if let Ok(c_cond) = fixed_factor
8770                    && c_cond.is_finite()
8771                    && c_cond > 0.0
8772                {
8773                    let mut c_applied = c_cond;
8774                    correction = SmoothLrCorrection::LawleyLrFixedLambda;
8775                    if let Some(cov) = rho_covariance
8776                        && let Ok(total_shift) = lawley_lr_mean_shift_with_rho_variation(
8777                            full_design_dense.view(),
8778                            &kappas,
8779                            s_lambda.view(),
8780                            coeff_range.clone(),
8781                            &rho_penalty_components,
8782                            cov.view(),
8783                        )
8784                    {
8785                        let mean_w = ref_df + total_shift;
8786                        if let Some(c_est) =
8787                            gam_terms::inference::higher_order::bartlett_factor_from_mean(
8788                                mean_w, ref_df,
8789                            )
8790                            && c_est.is_finite()
8791                            && c_est > 0.0
8792                        {
8793                            let conditional_shift = (c_cond - 1.0) * ref_df;
8794                            c_applied = c_est;
8795                            bartlett_factor_conditional = Some(c_cond);
8796                            rho_variation_shift = Some(total_shift - conditional_shift);
8797                            correction = SmoothLrCorrection::LawleyLrEstimatedLambda;
8798                        }
8799                    }
8800                    use statrs::distribution::ContinuousCDF;
8801                    bartlett_factor = c_applied;
8802                    statistic_corrected = statistic_lr / c_applied;
8803                    p_corrected = (1.0 - dist.cdf(statistic_corrected)).clamp(0.0, 1.0);
8804                }
8805            }
8806        }
8807
8808        // Materiality (#939 deliverable 4): only when a correction was actually
8809        // applied, flagged when it moves the result by more than the 10%
8810        // threshold — by the Bartlett factor's distance from one OR the relative
8811        // p-value shift, whichever is larger (a factor near one can still flip a
8812        // p-value sitting on the α boundary, and vice versa).
8813        let material = match correction {
8814            SmoothLrCorrection::LawleyLrEstimatedLambda
8815            | SmoothLrCorrection::LawleyLrFixedLambda => {
8816                let factor_move = (bartlett_factor - 1.0).abs();
8817                let p_denom = p_uncorrected.max(p_corrected).max(f64::MIN_POSITIVE);
8818                let p_move = if p_uncorrected.is_finite() && p_corrected.is_finite() {
8819                    (p_corrected - p_uncorrected).abs() / p_denom
8820                } else {
8821                    0.0
8822                };
8823                factor_move > SMOOTH_LR_MATERIAL_THRESHOLD || p_move > SMOOTH_LR_MATERIAL_THRESHOLD
8824            }
8825            SmoothLrCorrection::None => false,
8826        };
8827
8828        out.push(SmoothTermLrInference {
8829            name: design_term.name.clone(),
8830            term_idx,
8831            statistic_lr,
8832            ref_df,
8833            bartlett_factor,
8834            bartlett_factor_conditional,
8835            rho_variation_shift,
8836            statistic_corrected,
8837            p_value_uncorrected: p_uncorrected,
8838            p_value_corrected: p_corrected,
8839            material,
8840            correction,
8841        });
8842    }
8843    Ok(out)
8844}
8845
8846fn resolved_likelihood_for_fit(
8847    fit: &UnifiedFitResult,
8848) -> Result<gam_spec::GlmLikelihoodSpec, EstimationError> {
8849    let spec = fit.likelihood_family.as_ref().ok_or_else(|| {
8850        EstimationError::InvalidInput(
8851            "smooth-term LR inference requires an engine-level GLM likelihood".to_string(),
8852        )
8853    })?;
8854    gam_spec::GlmLikelihoodSpec::try_new(spec.clone(), fit.likelihood_scale.clone())
8855        .map_err(|error| EstimationError::InvalidInput(error.to_string()))
8856}
8857
8858/// The response dispersion `phi` Lawley needs for cumulant scaling. This is
8859/// deliberately distinct from the coefficient-covariance multiplier used by
8860/// the WPS trace below: Gamma Lawley uses `1 / shape`, while its PIRLS Hessian
8861/// already carries `shape` and therefore has covariance multiplier one.
8862fn lawley_dispersion_for_family(
8863    likelihood: &gam_spec::GlmLikelihoodSpec,
8864    fit: &UnifiedFitResult,
8865) -> Result<f64, EstimationError> {
8866    let profiled_standard_deviation = matches!(
8867        likelihood
8868            .resolved_scale()
8869            .map_err(|error| EstimationError::InvalidInput(error.to_string()))?,
8870        gam_spec::ResolvedLikelihoodScale::ProfiledGaussian
8871    )
8872    .then_some(fit.standard_deviation);
8873    gam_solve::estimate::dispersion_from_likelihood(likelihood, profiled_standard_deviation)
8874        .map(|dispersion| dispersion.phi())
8875}
8876
8877fn wps_block_uncertainty_df(
8878    weighted_gram: Option<&Array2<f64>>,
8879    smoothing_correction: Option<&Array2<f64>>,
8880    coeff_range: &Range<usize>,
8881    coefficient_covariance_scale: f64,
8882) -> Result<Option<f64>, EstimationError> {
8883    let (Some(xwx), Some(corr)) = (weighted_gram, smoothing_correction) else {
8884        return Ok(None);
8885    };
8886    let (start, end) = (coeff_range.start, coeff_range.end);
8887    if start >= end {
8888        return Err(EstimationError::InvalidInput(format!(
8889            "WPS coefficient block must be non-empty, got {coeff_range:?}"
8890        )));
8891    }
8892    if xwx.nrows() != xwx.ncols() || corr.nrows() != corr.ncols() {
8893        return Err(EstimationError::InvalidInput(format!(
8894            "WPS matrices must be square, got X'WX={}x{} and correction={}x{}",
8895            xwx.nrows(),
8896            xwx.ncols(),
8897            corr.nrows(),
8898            corr.ncols()
8899        )));
8900    }
8901    if xwx.dim() != corr.dim() || end > xwx.nrows() {
8902        return Err(EstimationError::InvalidInput(format!(
8903            "WPS block {coeff_range:?} is incompatible with X'WX={:?} and correction={:?}",
8904            xwx.dim(),
8905            corr.dim()
8906        )));
8907    }
8908    if !(coefficient_covariance_scale.is_finite() && coefficient_covariance_scale > 0.0) {
8909        return Err(EstimationError::InvalidInput(format!(
8910            "WPS coefficient-covariance scale must be finite and strictly positive, got {coefficient_covariance_scale:?}"
8911        )));
8912    }
8913
8914    let mut trace = gam_linalg::utils::KahanSum::default();
8915    for i in start..end {
8916        for j in start..end {
8917            let gram_value = xwx[[i, j]];
8918            let correction_value = corr[[j, i]];
8919            if !gram_value.is_finite() || !correction_value.is_finite() {
8920                return Err(EstimationError::InvalidInput(format!(
8921                    "WPS trace has non-finite matrix entry at ({i}, {j}): X'WX={gram_value:?}, correction-transpose={correction_value:?}"
8922                )));
8923            }
8924            let product = gram_value * correction_value;
8925            if !product.is_finite() {
8926                return Err(EstimationError::InvalidInput(format!(
8927                    "WPS trace product is not representable at ({i}, {j}): {gram_value:?} * {correction_value:?}"
8928                )));
8929            }
8930            trace.add(product);
8931        }
8932    }
8933    let trace = trace.sum() / coefficient_covariance_scale;
8934    if !trace.is_finite() {
8935        return Err(EstimationError::InvalidInput(format!(
8936            "WPS corrected-EDF trace is not representable after coefficient scale {coefficient_covariance_scale:?}: {trace:?}"
8937        )));
8938    }
8939    if trace < 0.0 {
8940        return Err(EstimationError::InvalidInput(format!(
8941            "WPS corrected-EDF trace must be non-negative, got {trace:?}"
8942        )));
8943    }
8944    Ok(Some(trace))
8945}
8946
8947/// Wood's smoothing-selection-corrected reference d.f. `edf1 = 2·tr(F_jj) −
8948/// tr(F_jj²)` on the coefficient-influence block `F = H⁻¹ X'WX` restricted to
8949/// `coeff_range` (mgcv's `edf1`). This is the reference degrees of freedom
8950/// Wood (2013) recommends for a smooth-component significance test when the
8951/// smoothing parameter is *estimated* (as it always is here): it inflates the
8952/// raw effective d.f. `edf = tr(F)` by the selection-bias term
8953/// `tr(F) − tr(F²) = Σ_i λ_i(1 − λ_i) ≥ 0` (the F-block eigenvalues `λ_i` lie in
8954/// `[0, 1]`), which is exactly the excess variance the fit spends locating a
8955/// term on noise. Referencing the whole-term LR statistic against this larger
8956/// `edf1` instead of the raw `edf` removes the post-selection left-tail
8957/// anti-conservatism (a REML fit that turns a smooth *on* against pure noise no
8958/// longer over-rejects): with the raw `edf` reference the mean null p-value is a
8959/// calibrated ~0.5 but the 5%-level false-positive rate runs ~0.15; the `edf1`
8960/// correction pulls it toward the nominal α while leaving the collapsed-term
8961/// verdict (`W ≈ 0 ⇒ p ≈ 1`) untouched.
8962///
8963/// `edf1 ∈ [edf, 2·edf]` analytically, so it never collapses toward 0 the way
8964/// the earlier Satterthwaite ratio `tr(F)²/tr(F²)` did when the non-symmetric
8965/// block's `tr(F²)` ran away (the #1766 degeneracy). The `.max(tr)` guard keeps
8966/// it ≥ `edf` even if a corrupted block drives `tr(F²)` above `2·tr(F)`. Returns
8967/// `None` when the influence block is unavailable or its trace is non-finite /
8968/// non-positive (a fully shrunk term), in which case the caller falls back to
8969/// the `max(edf, null_dim, 1)` floor.
8970fn wood_reference_df(influence: Option<&Array2<f64>>, coeff_range: &Range<usize>) -> Option<f64> {
8971    let f = influence?;
8972    let (start, end) = (coeff_range.start, coeff_range.end);
8973    if start >= end || end > f.nrows() || end > f.ncols() {
8974        return None;
8975    }
8976    let block = f.slice(s![start..end, start..end]);
8977    let tr = (0..block.nrows()).map(|i| block[[i, i]]).sum::<f64>();
8978    let tr2 = block.dot(&block).diag().sum();
8979    (tr.is_finite() && tr2.is_finite() && tr > 0.0).then(|| (2.0 * tr - tr2).max(tr).max(1e-12))
8980}
8981
8982#[cfg(test)]
8983mod likelihood_scale_wps_tests {
8984    use super::wps_block_uncertainty_df;
8985    use ndarray::array;
8986
8987    #[test]
8988    fn wps_trace_uses_coefficient_covariance_scale() {
8989        let xwx = array![[1.0, 0.0], [0.0, 1.0]];
8990        let correction = array![[1.0, 0.0], [0.0, 1.0]];
8991        let extra_df = wps_block_uncertainty_df(Some(&xwx), Some(&correction), &(0..2), 4.0)
8992            .expect("valid WPS geometry")
8993            .expect("correction artifacts are present");
8994        assert_eq!(extra_df, 0.5);
8995    }
8996
8997    #[test]
8998    fn wps_absence_is_distinct_from_invalid_geometry() {
8999        let xwx = array![[1.0]];
9000        assert_eq!(
9001            wps_block_uncertainty_df(Some(&xwx), None, &(0..1), 1.0)
9002                .expect("missing optional artifact is not malformed geometry"),
9003            None
9004        );
9005
9006        let negative_correction = array![[-1.0]];
9007        let error = wps_block_uncertainty_df(Some(&xwx), Some(&negative_correction), &(0..1), 1.0)
9008            .expect_err("negative corrected EDF must not be silently zeroed");
9009        assert!(error.to_string().contains("must be non-negative"));
9010    }
9011}
9012
9013#[cfg(test)]
9014mod nfree_gate_tests {
9015    use super::nfree_skip_gate_status_from_parts;
9016
9017    #[test]
9018    fn value_only_nfree_gate_does_not_require_basis_skip_witness() {
9019        let gate = nfree_skip_gate_status_from_parts(
9020            true,  // shape
9021            true,  // Chebyshev Gram value covers this ψ
9022            false, // reduced-basis skip witness absent across a rotation seam
9023            false, // gradient coverage irrelevant for a value-only cost probe
9024            true,  // penalty can be re-keyed without rows
9025            true,  // design revision is pinned
9026            false, // no Hessian request
9027            false, // value-only cost probe
9028        );
9029        assert!(
9030            gate.would_skip(false),
9031            "value-only κ cost probes must stay n-free when the Gram value is certified; \
9032             the reduced-basis skip witness is required only for beta/gradient probes"
9033        );
9034    }
9035
9036    #[test]
9037    fn gradient_nfree_gate_still_requires_basis_skip_witness() {
9038        let gate =
9039            nfree_skip_gate_status_from_parts(true, true, false, true, true, true, false, true);
9040        assert!(
9041            !gate.would_skip(true),
9042            "gradient probes return beta/gradient objects in a reduced basis and must not \
9043             skip the row lane without the reduced-basis witness"
9044        );
9045    }
9046}