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