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