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