1fn try_build_spatial_term_log_kappa_derivative(
2 data: ArrayView2<'_, f64>,
3 resolvedspec: &TermCollectionSpec,
4 design: &TermCollectionDesign,
5 term_idx: usize,
6) -> Result<
7 Option<(
8 Range<usize>,
9 usize,
10 Array2<f64>,
11 Array2<f64>,
12 Array2<f64>,
13 Array2<f64>,
14 Vec<Array2<f64>>,
15 Vec<Array2<f64>>,
16 Option<std::sync::Arc<gam_terms::basis::ImplicitDesignPsiDerivative>>,
17 )>,
18 EstimationError,
19> {
20 let Some(smooth_term) = design.smooth.terms.get(term_idx) else {
21 return Ok(None);
22 };
23 let Some(termspec) = resolvedspec.smooth_terms.get(term_idx) else {
24 return Ok(None);
25 };
26
27 let derivative_bundle = match &termspec.basis {
28 SmoothBasisSpec::ThinPlate {
29 feature_cols,
30 spec,
31 input_scale,
32 } => {
33 let mut x = select_columns(data, feature_cols).map_err(EstimationError::from)?;
34 let mut spec_local = spec.clone();
35 if let Some(scale) = input_scale {
36 scale.standardize(&mut x);
37 spec_local.length_scale = scale.to_standardized_units(spec.length_scale);
38 }
39 build_thin_plate_basis_log_kappa_derivatives(x.view(), &spec_local)
40 .map_err(EstimationError::from)?
41 }
42 SmoothBasisSpec::Sphere { .. } => return Ok(None),
43 SmoothBasisSpec::ConstantCurvature { feature_cols, spec } => {
52 let x = select_columns(data, feature_cols).map_err(EstimationError::from)?;
53 build_constant_curvature_basis_kappa_derivatives(x.view(), spec)
54 .map_err(EstimationError::from)?
55 }
56 SmoothBasisSpec::MeasureJet { .. } => return Ok(None),
62 SmoothBasisSpec::Matern {
63 feature_cols,
64 spec,
65 input_scale,
66 } => {
67 let mut x = select_columns(data, feature_cols).map_err(EstimationError::from)?;
68 let mut spec_local = spec.clone();
69 if let Some(scale) = input_scale {
70 scale.standardize(&mut x);
71 let length_scale = spec.length_scale.resolved().ok_or_else(|| {
72 EstimationError::InvalidInput(
73 "Matérn Auto length_scale reached derivative construction unresolved"
74 .to_string(),
75 )
76 })?;
77 spec_local
78 .length_scale
79 .set_resolved(scale.to_standardized_units(length_scale));
80 }
81 spec_local.double_penalty = false;
96 build_matern_basis_log_kappa_derivatives(x.view(), &spec_local)
97 .map_err(EstimationError::from)?
98 }
99 SmoothBasisSpec::Duchon {
100 feature_cols,
101 spec,
102 input_scale,
103 } => {
104 let mut x = select_columns(data, feature_cols).map_err(EstimationError::from)?;
105 let mut spec_local = spec.clone();
106 if let Some(scale) = input_scale {
107 scale.standardize(&mut x);
108 spec_local.length_scale = spec
109 .length_scale
110 .map(|length| scale.to_standardized_units(length));
111 }
112 let BasisMetadata::Duchon {
113 centers,
114 identifiability_transform,
115 operator_collocation_points,
116 radial_reparam,
117 ..
118 } = &smooth_term.metadata
119 else {
120 return Ok(None);
121 };
122 if spec_local.radial_reparam.is_none() {
125 spec_local.radial_reparam = radial_reparam.clone();
126 }
127 gam_terms::basis::build_duchon_basis_log_kappa_derivativeswith_collocationwithworkspace(
128 x.view(),
129 &spec_local,
130 centers.view(),
131 identifiability_transform.as_ref(),
132 operator_collocation_points
133 .as_ref()
134 .map(|points| points.view()),
135 &mut BasisWorkspace::default(),
136 )
137 .map_err(EstimationError::from)?
138 }
139 SmoothBasisSpec::BSpline1D { .. }
140 | SmoothBasisSpec::TensorBSpline { .. }
141 | SmoothBasisSpec::ByVariable { .. }
142 | SmoothBasisSpec::FactorSumToZero { .. }
143 | SmoothBasisSpec::BySmooth { .. }
144 | SmoothBasisSpec::FactorSmooth { .. }
145 | SmoothBasisSpec::Pca { .. } => {
146 return Ok(None);
147 }
148 };
149 let mut implicit_operator = derivative_bundle.implicit_operator;
150 let BasisPsiDerivativeResult {
151 design_derivative: mut local_x_psi,
152 penalties_derivative: mut local_s_psi,
153 implicit_operator: local_implicit_first_unused,
154 } = derivative_bundle.first;
155 let BasisPsiSecondDerivativeResult {
156 designsecond_derivative: mut local_x_psi_psi,
157 penaltiessecond_derivative: mut local_s_psi_psi,
158 implicit_operator: local_implicit_second_unused,
159 } = derivative_bundle.second;
160 assert!(local_implicit_first_unused.is_none());
161 assert!(local_implicit_second_unused.is_none());
162
163 if let Some(rotation) = smooth_term.joint_null_rotation.as_ref() {
164 let q = &rotation.rotation;
165 if let Some(op) = implicit_operator.take() {
166 implicit_operator = Some(op.append_full_transform(q).map_err(EstimationError::from)?);
167 } else {
168 if local_x_psi.ncols() != q.nrows() || local_x_psi_psi.ncols() != q.nrows() {
169 return Ok(None);
170 }
171 local_x_psi = fast_ab(&local_x_psi, q);
172 local_x_psi_psi = fast_ab(&local_x_psi_psi, q);
173 }
174 let rotate_penalty = |s_local: Array2<f64>| -> Option<Array2<f64>> {
175 if s_local.nrows() != q.nrows() || s_local.ncols() != q.nrows() {
176 return None;
177 }
178 let qt_s = gam_linalg::faer_ndarray::fast_atb(q, &s_local);
179 Some(gam_linalg::faer_ndarray::fast_ab(&qt_s, q))
180 };
181 let Some(rotated_s_psi) = local_s_psi
182 .into_iter()
183 .map(|s| rotate_penalty(s))
184 .collect::<Option<Vec<_>>>()
185 else {
186 return Ok(None);
187 };
188 local_s_psi = rotated_s_psi;
189 let Some(rotated_s_psi_psi) = local_s_psi_psi
190 .into_iter()
191 .map(|s| rotate_penalty(s))
192 .collect::<Option<Vec<_>>>()
193 else {
194 return Ok(None);
195 };
196 local_s_psi_psi = rotated_s_psi_psi;
197 }
198 let implicit_operator = implicit_operator.map(std::sync::Arc::new);
199
200 if let Some(ref op) = implicit_operator {
201 if op.p_out() != smooth_term.coeff_range.len() {
202 return Ok(None);
203 }
204 } else {
205 if local_x_psi.ncols() != smooth_term.coeff_range.len() {
206 return Ok(None);
207 }
208 if local_x_psi_psi.ncols() != smooth_term.coeff_range.len() {
209 return Ok(None);
210 }
211 }
212 if local_s_psi.is_empty() || local_s_psi.len() != local_s_psi_psi.len() {
213 return Ok(None);
214 }
215 if local_s_psi.iter().any(|s| {
216 s.nrows() != smooth_term.coeff_range.len() || s.ncols() != smooth_term.coeff_range.len()
217 }) {
218 return Ok(None);
219 }
220 if local_s_psi_psi.iter().any(|s| {
221 s.nrows() != smooth_term.coeff_range.len() || s.ncols() != smooth_term.coeff_range.len()
222 }) {
223 return Ok(None);
224 }
225
226 let p_total = design.design.ncols();
227 let smooth_start = p_total.saturating_sub(design.smooth.total_smooth_cols());
228 let global_range = (smooth_start + smooth_term.coeff_range.start)
229 ..(smooth_start + smooth_term.coeff_range.end);
230
231 Ok(Some((
232 global_range,
233 p_total,
234 local_x_psi,
235 local_s_psi.iter().fold(
236 Array2::<f64>::zeros((smooth_term.coeff_range.len(), smooth_term.coeff_range.len())),
237 |acc, m| acc + m,
238 ),
239 local_x_psi_psi,
240 local_s_psi_psi.iter().fold(
241 Array2::<f64>::zeros((smooth_term.coeff_range.len(), smooth_term.coeff_range.len())),
242 |acc, m| acc + m,
243 ),
244 local_s_psi,
245 local_s_psi_psi,
246 implicit_operator,
247 )))
248}
249
250fn try_build_spatial_log_kappa_hyper_dirs(
251 data: ArrayView2<'_, f64>,
252 resolvedspec: &TermCollectionSpec,
253 design: &TermCollectionDesign,
254 spatial_terms: &[usize],
255) -> Result<Option<Vec<DirectionalHyperParam>>, EstimationError> {
256 let Some(info_list) =
263 try_build_spatial_log_kappa_derivativeinfo_list(data, resolvedspec, design, spatial_terms)?
264 else {
265 return Ok(None);
266 };
267 Ok(Some(spatial_log_kappa_hyper_dirs_frominfo_list(info_list)?))
268}
269
270pub(crate) fn try_build_latent_coord_hyper_dirs(
271 latent: std::sync::Arc<gam_terms::latent::LatentCoordValues>,
272 resolvedspec: &TermCollectionSpec,
273 design: &TermCollectionDesign,
274 latent_terms: &[gam_problem::types::SmoothTermIdx],
275 analytic_rho_count: usize,
276) -> Result<Option<Vec<DirectionalHyperParam>>, EstimationError> {
277 if latent_terms.is_empty() || latent.is_empty() {
278 return Ok(None);
279 }
280 if latent_terms.len() != 1 {
281 crate::bail_invalid_estim!(
282 "LatentCoord standard-fit hyper_dirs currently require exactly one latent smooth term"
283 .to_string(),
284 );
285 }
286 let term_idx = latent_terms[0];
287 let smooth_term = design.smooth.terms.get(term_idx.get()).ok_or_else(|| {
288 EstimationError::InvalidInput(format!(
289 "LatentCoord term index {term_idx} out of bounds for realized smooth design"
290 ))
291 })?;
292 let termspec = resolvedspec
293 .smooth_terms
294 .get(term_idx.get())
295 .ok_or_else(|| {
296 EstimationError::InvalidInput(format!(
297 "LatentCoord term index {term_idx} out of bounds for resolved smooth spec"
298 ))
299 })?;
300 let p_total = design.design.ncols();
301 let smooth_start = p_total.saturating_sub(design.smooth.total_smooth_cols());
302 let global_range = (smooth_start + smooth_term.coeff_range.start)
303 ..(smooth_start + smooth_term.coeff_range.end);
304
305 let operator = match (&termspec.basis, &smooth_term.metadata) {
310 (
311 SmoothBasisSpec::Matern { .. },
312 BasisMetadata::Matern {
313 centers,
314 length_scale,
315 nu,
316 include_intercept,
317 identifiability_transform,
318 ..
319 },
320 ) => gam_terms::basis::LatentCoordDesignDerivative::new_matern(
321 latent.clone(),
322 std::sync::Arc::new(centers.clone()),
323 *length_scale,
324 *nu,
325 *include_intercept,
326 identifiability_transform.clone(),
327 )
328 .map_err(EstimationError::from)?,
329 (
330 SmoothBasisSpec::Duchon { .. },
331 BasisMetadata::Duchon {
332 centers,
333 length_scale,
334 power,
335 nullspace_order,
336 identifiability_transform,
337 ..
338 },
339 ) => gam_terms::basis::LatentCoordDesignDerivative::new_duchon(
340 latent.clone(),
341 std::sync::Arc::new(centers.clone()),
342 *length_scale,
343 *power,
344 *nullspace_order,
345 identifiability_transform.clone(),
346 )
347 .map_err(EstimationError::from)?,
348 (
349 SmoothBasisSpec::Sphere { .. },
350 BasisMetadata::Sphere {
351 centers,
352 penalty_order,
353 method,
354 constraint_transform,
355 ..
356 },
357 ) if matches!(*method, gam_terms::basis::SphereMethod::Wahba) => {
358 gam_terms::basis::LatentCoordDesignDerivative::new_sphere(
359 latent.clone(),
360 std::sync::Arc::new(centers.clone()),
361 *penalty_order,
362 constraint_transform.clone(),
363 )
364 .map_err(EstimationError::from)?
365 }
366 (
367 SmoothBasisSpec::BSpline1D { spec, .. },
368 BasisMetadata::BSpline1D {
369 knots,
370 identifiability_transform,
371 periodic,
372 degree: meta_degree,
373 ..
374 },
375 ) => {
376 let effective_degree = meta_degree.unwrap_or(spec.degree);
380 if let Some((domain_start, period, num_basis)) = periodic {
381 gam_terms::basis::LatentCoordDesignDerivative::new_periodic_bspline(
382 latent.clone(),
383 (*domain_start, *domain_start + *period),
384 effective_degree,
385 *num_basis,
386 identifiability_transform.clone(),
387 )
388 .map_err(EstimationError::from)?
389 } else {
390 gam_terms::basis::LatentCoordDesignDerivative::new_tensor_bspline(
391 latent.clone(),
392 vec![knots.clone()],
393 vec![effective_degree],
394 identifiability_transform.clone(),
395 )
396 .map_err(EstimationError::from)?
397 }
398 }
399 (
400 SmoothBasisSpec::TensorBSpline { .. },
401 BasisMetadata::TensorBSpline {
402 knots,
403 degrees,
404 identifiability_transform,
405 ..
406 },
407 ) => gam_terms::basis::LatentCoordDesignDerivative::new_tensor_bspline(
408 latent.clone(),
409 knots.clone(),
410 degrees.clone(),
411 identifiability_transform.clone(),
412 )
413 .map_err(EstimationError::from)?,
414 (SmoothBasisSpec::Pca { .. }, BasisMetadata::Pca { basis_matrix, .. }) => {
415 gam_terms::basis::LatentCoordDesignDerivative::new_pca(
416 latent.clone(),
417 std::sync::Arc::new(basis_matrix.clone()),
418 )
419 .map_err(EstimationError::from)?
420 }
421 _ => return Ok(None),
422 };
423 if operator.p_out() != global_range.len() {
424 crate::bail_invalid_estim!(
425 "LatentCoord derivative width mismatch for term '{}': operator p={}, coeff range={}",
426 smooth_term.name,
427 operator.p_out(),
428 global_range.len()
429 );
430 }
431 let operator = std::sync::Arc::new(operator);
432 let mut hyper_dirs = Vec::with_capacity(operator.n_axes());
433 for flat_axis in 0..operator.n_axes() {
434 let dir = DirectionalHyperParam::new_compact(
435 gam_solve::estimate::reml::HyperDesignDerivative::from_latent_coord(
436 operator.clone(),
437 flat_axis,
438 global_range.clone(),
439 p_total,
440 ),
441 Vec::new(),
442 None,
443 None,
444 )?
445 .not_penalty_like();
446 hyper_dirs.push(dir);
447 }
448 let direct_dim = latent_coord_direct_hyper_count(latent.id_mode(), latent.latent_dim());
449 if analytic_rho_count + direct_dim > 0 {
450 let zero_x = gam_solve::estimate::reml::HyperDesignDerivative::from(Array2::<f64>::zeros(
451 (design.design.nrows(), p_total),
452 ));
453 for _ in 0..analytic_rho_count {
454 hyper_dirs.push(
455 DirectionalHyperParam::new_compact(zero_x.clone(), Vec::new(), None, None)?
456 .not_penalty_like(),
457 );
458 }
459 for _ in 0..direct_dim {
460 hyper_dirs.push(
461 DirectionalHyperParam::new_compact(zero_x.clone(), Vec::new(), None, None)?
462 .not_penalty_like(),
463 );
464 }
465 }
466 Ok(Some(hyper_dirs))
467}
468
469fn latent_coord_direct_hyper_count(
470 id_mode: &gam_terms::latent::LatentIdMode,
471 latent_dim: usize,
472) -> usize {
473 use gam_terms::latent::{AuxPriorStrength, LatentIdMode};
474 match id_mode {
475 LatentIdMode::AuxPrior { strength, .. } => match strength {
476 AuxPriorStrength::Auto => 1,
477 AuxPriorStrength::Fixed(_) => 0,
478 },
479 LatentIdMode::AuxPriorDimSelection { strength, .. } => {
480 latent_dim
481 + match strength {
482 AuxPriorStrength::Auto => 1,
483 AuxPriorStrength::Fixed(_) => 0,
484 }
485 }
486 LatentIdMode::DimSelection { .. } => latent_dim,
487 LatentIdMode::IsometryToReference { strength, .. } => match strength {
490 AuxPriorStrength::Auto => 1,
491 AuxPriorStrength::Fixed(_) => 0,
492 },
493 LatentIdMode::AuxOutcome { head, .. } => head.n_coeffs(latent_dim) + latent_dim,
496 LatentIdMode::None => 0,
497 }
498}
499
500fn latent_coord_initial_direct_hypers(
501 id_mode: &gam_terms::latent::LatentIdMode,
502 latent_dim: usize,
503) -> Result<Array1<f64>, EstimationError> {
504 use gam_terms::latent::{AuxPriorStrength, LatentIdMode};
505 let mut values = Vec::with_capacity(latent_coord_direct_hyper_count(id_mode, latent_dim));
506 match id_mode {
507 LatentIdMode::AuxPrior { strength, .. } => {
508 if matches!(strength, AuxPriorStrength::Auto) {
509 values.push(0.0);
510 }
511 }
512 LatentIdMode::AuxPriorDimSelection {
513 strength,
514 init_log_precision,
515 ..
516 } => {
517 if matches!(strength, AuxPriorStrength::Auto) {
518 values.push(0.0);
519 }
520 append_latent_ard_seed(&mut values, init_log_precision.as_ref(), latent_dim)?;
521 }
522 LatentIdMode::DimSelection { init_log_precision } => {
523 append_latent_ard_seed(&mut values, init_log_precision.as_ref(), latent_dim)?;
524 }
525 LatentIdMode::IsometryToReference { strength, .. } => {
526 if matches!(strength, AuxPriorStrength::Auto) {
527 values.push(0.0);
528 }
529 }
530 LatentIdMode::AuxOutcome {
531 head,
532 init_log_precision,
533 } => {
534 values.extend(std::iter::repeat_n(0.0, head.n_coeffs(latent_dim)));
538 append_latent_ard_seed(&mut values, init_log_precision.as_ref(), latent_dim)?;
539 }
540 LatentIdMode::None => {}
541 }
542 Ok(Array1::from_vec(values))
543}
544
545fn append_latent_ard_seed(
546 values: &mut Vec<f64>,
547 init: Option<&Array1<f64>>,
548 latent_dim: usize,
549) -> Result<(), EstimationError> {
550 if let Some(init) = init {
551 if init.len() != latent_dim {
552 crate::bail_invalid_estim!(
553 "latent dim_selection init_log_precision length mismatch: got {}, expected {}",
554 init.len(),
555 latent_dim
556 );
557 }
558 values.extend(init.iter().copied());
559 } else {
560 values.extend(std::iter::repeat_n(0.0, latent_dim));
561 }
562 Ok(())
563}
564
565struct LatentIdObjectiveContribution {
566 cost: f64,
567 gradient: Array1<f64>,
568}
569
570fn latent_id_objective_contribution(
571 theta: &Array1<f64>,
572 rho_dim: usize,
573 analytic_rho_count: usize,
574 latent: &gam_terms::latent::LatentCoordValues,
575) -> Result<LatentIdObjectiveContribution, EstimationError> {
576 use gam_terms::latent::{AuxPriorStrength, LatentIdMode, aux_prior_targets};
577 let n_obs = latent.n_obs();
578 let latent_dim = latent.latent_dim();
579 let flat_len = latent.len();
580 let mut gradient = Array1::<f64>::zeros(theta.len());
581 let t_start = rho_dim;
582 let direct_start = t_start + flat_len + analytic_rho_count;
583 if theta.len() < direct_start {
584 crate::bail_invalid_estim!(
585 "latent-coordinate theta too short for id objective: got {}, need at least {}",
586 theta.len(),
587 direct_start
588 );
589 }
590 let t = latent.as_matrix();
591 let mut cost = 0.0;
592 let mut cursor = direct_start;
593
594 match latent.id_mode() {
595 LatentIdMode::AuxPrior {
596 u,
597 family,
598 strength,
599 }
600 | LatentIdMode::AuxPriorDimSelection {
601 u,
602 family,
603 strength,
604 ..
605 } => {
606 let (log_mu, mu) = match strength {
607 AuxPriorStrength::Fixed(mu) => (
608 gam_problem::checked_log_strength(*mu).map_err(|error| {
609 EstimationError::InvalidInput(format!(
610 "fixed latent auxiliary-prior precision is outside the canonical physical-strength domain: {error}"
611 ))
612 })?,
613 *mu,
614 ),
615 AuxPriorStrength::Auto => {
616 let log_mu = *theta.get(cursor).ok_or_else(|| {
617 EstimationError::InvalidInput(format!(
618 "latent auxiliary-prior precision coordinate {cursor} is missing from theta length {}",
619 theta.len(),
620 ))
621 })?;
622 cursor += 1;
623 let mu = gam_problem::checked_exp_log_strength(log_mu).map_err(|error| {
624 EstimationError::InvalidInput(format!(
625 "latent auxiliary-prior log precision is outside the canonical log-strength domain: {error}"
626 ))
627 })?;
628 (log_mu, mu)
629 }
630 };
631 let targets = aux_prior_targets(t.view(), u.view(), *family)
632 .map_err(EstimationError::InvalidInput)?;
633 let residual = &t - &targets;
634 let q = residual.iter().map(|v| v * v).sum::<f64>();
635 let k = (n_obs * latent_dim) as f64;
642 cost += 0.5 * mu * q - 0.5 * k * log_mu;
643
644 let projected_residual = aux_prior_targets(residual.view(), u.view(), *family)
645 .map_err(EstimationError::InvalidInput)?;
646 let grad_base = residual - projected_residual;
647 for n in 0..n_obs {
648 for axis in 0..latent_dim {
649 gradient[t_start + n * latent_dim + axis] += mu * grad_base[[n, axis]];
650 }
651 }
652 if matches!(strength, AuxPriorStrength::Auto) {
653 gradient[direct_start] += 0.5 * mu * q - 0.5 * k;
654 }
655 }
656 LatentIdMode::IsometryToReference {
657 reference,
658 strength,
659 } => {
660 if reference.dim() != (n_obs, latent_dim) {
667 crate::bail_invalid_estim!(
668 "IsometryToReference reference shape {:?} must equal (n_obs, latent_dim) = ({}, {})",
669 reference.dim(),
670 n_obs,
671 latent_dim
672 );
673 }
674 let mu_slot = cursor;
675 let (log_mu, mu) = match strength {
676 AuxPriorStrength::Fixed(mu) => (
677 gam_problem::checked_log_strength(*mu).map_err(|error| {
678 EstimationError::InvalidInput(format!(
679 "fixed latent isometry precision is outside the canonical physical-strength domain: {error}"
680 ))
681 })?,
682 *mu,
683 ),
684 AuxPriorStrength::Auto => {
685 let log_mu = *theta.get(cursor).ok_or_else(|| {
686 EstimationError::InvalidInput(format!(
687 "latent isometry precision coordinate {cursor} is missing from theta length {}",
688 theta.len(),
689 ))
690 })?;
691 cursor += 1;
692 let mu = gam_problem::checked_exp_log_strength(log_mu).map_err(|error| {
693 EstimationError::InvalidInput(format!(
694 "latent isometry log precision is outside the canonical log-strength domain: {error}"
695 ))
696 })?;
697 (log_mu, mu)
698 }
699 };
700 let residual = &t - reference;
701 let q = residual.iter().map(|v| v * v).sum::<f64>();
702 let k = (n_obs * latent_dim) as f64;
706 cost += 0.5 * mu * q - 0.5 * k * log_mu;
707 for n in 0..n_obs {
708 for axis in 0..latent_dim {
709 gradient[t_start + n * latent_dim + axis] += mu * residual[[n, axis]];
710 }
711 }
712 if matches!(strength, AuxPriorStrength::Auto) {
713 gradient[mu_slot] += 0.5 * mu * q - 0.5 * k;
714 }
715 }
716 LatentIdMode::AuxOutcome { head, .. } => {
717 let n_coeffs = head.n_coeffs(latent_dim);
725 if cursor + n_coeffs > theta.len() {
726 crate::bail_invalid_estim!(
727 "latent auxiliary-outcome coefficient block overruns theta: start={cursor}, width={n_coeffs}, theta_len={}",
728 theta.len(),
729 );
730 }
731 let coeffs = theta
732 .slice(ndarray::s![cursor..cursor + n_coeffs])
733 .to_owned();
734 let (head_nll, grad_coeffs, grad_t) = head
735 .neg_loglik_and_grad(t.view(), coeffs.view())
736 .map_err(EstimationError::InvalidInput)?;
737 cost += head_nll;
738 for (offset, &g) in grad_coeffs.iter().enumerate() {
739 gradient[cursor + offset] += g;
740 }
741 for n in 0..n_obs {
742 for axis in 0..latent_dim {
743 gradient[t_start + n * latent_dim + axis] += grad_t[[n, axis]];
744 }
745 }
746 cursor += n_coeffs;
747 }
748 LatentIdMode::DimSelection { .. } | LatentIdMode::None => {}
749 }
750
751 match latent.id_mode() {
752 LatentIdMode::AuxPriorDimSelection { .. }
753 | LatentIdMode::DimSelection { .. }
754 | LatentIdMode::AuxOutcome { .. } => {
755 if cursor + latent_dim > theta.len() {
756 crate::bail_invalid_estim!(
757 "latent dimension-selection precision block overruns theta: start={cursor}, width={latent_dim}, theta_len={}",
758 theta.len(),
759 );
760 }
761 let alphas = gam_problem::checked_exp_log_strengths(
762 theta.slice(s![cursor..cursor + latent_dim]).iter().copied(),
763 )
764 .map_err(|error| {
765 EstimationError::InvalidInput(format!(
766 "latent dimension-selection log precision is outside the canonical log-strength domain: {error}"
767 ))
768 })?;
769 for axis in 0..latent_dim {
770 let log_alpha = theta[cursor + axis];
771 let alpha = alphas[axis];
772 let mut q_axis = 0.0;
773 for n in 0..n_obs {
774 let flat_idx = n * latent_dim + axis;
775 let value = latent.as_flat()[flat_idx];
776 q_axis += value * value;
777 gradient[t_start + flat_idx] += alpha * value;
778 }
779 cost += 0.5 * alpha * q_axis - 0.5 * n_obs as f64 * log_alpha;
780 gradient[cursor + axis] += 0.5 * alpha * q_axis - 0.5 * n_obs as f64;
781 }
782 cursor += latent_dim;
783 }
784 LatentIdMode::AuxPrior { .. }
785 | LatentIdMode::IsometryToReference { .. }
786 | LatentIdMode::None => {}
787 }
788
789 if cursor != theta.len() {
790 crate::bail_invalid_estim!(
791 "latent-coordinate direct hyperparameter length mismatch: consumed {}, theta len {}",
792 cursor,
793 theta.len()
794 );
795 }
796 Ok(LatentIdObjectiveContribution { cost, gradient })
797}
798
799fn add_latent_id_objective_to_eval(
800 theta: &Array1<f64>,
801 rho_dim: usize,
802 analytic_rho_count: usize,
803 latent: &gam_terms::latent::LatentCoordValues,
804 eval: &mut (f64, Array1<f64>, gam_problem::HessianValue),
805) -> Result<(), EstimationError> {
806 let contribution =
807 latent_id_objective_contribution(theta, rho_dim, analytic_rho_count, latent)?;
808 eval.0 += contribution.cost;
809 if eval.1.len() != contribution.gradient.len() {
810 crate::bail_invalid_estim!(
811 "latent-coordinate REML gradient length mismatch: base={}, id={}",
812 eval.1.len(),
813 contribution.gradient.len()
814 );
815 }
816 eval.1 += &contribution.gradient;
817 if eval.2.is_analytic() {
818 eval.2 = gam_problem::HessianValue::Unavailable;
819 }
820 Ok(())
821}
822
823fn analytic_penalty_objective_contribution(
824 theta: &Array1<f64>,
825 rho_dim: usize,
826 latent: &gam_terms::latent::LatentCoordValues,
827 registry: &gam_terms::AnalyticPenaltyRegistry,
828) -> Result<LatentIdObjectiveContribution, EstimationError> {
829 let flat_len = latent.len();
830 let t_start = rho_dim;
831 let t_end = t_start + flat_len;
832 let rho_start = t_end;
833 let rho_end = rho_start + registry.total_rho_count();
834 if theta.len() < rho_end {
835 crate::bail_invalid_estim!(
836 "latent-coordinate theta too short for analytic penalties: got {}, need at least {}",
837 theta.len(),
838 rho_end
839 );
840 }
841 let target_t = theta.slice(s![t_start..t_end]);
842 let rho = theta.slice(s![rho_start..rho_end]);
843 registry
844 .validate_rho(rho)
845 .map_err(EstimationError::InvalidInput)?;
846 let mut cost = 0.0_f64;
847 let mut gradient = Array1::<f64>::zeros(theta.len());
848 for (penalty, (rho_slice, tier, name)) in registry.penalties.iter().zip(registry.rho_layout()) {
849 let rho_local = rho.slice(s![rho_slice.clone()]);
850 match tier {
851 gam_terms::PenaltyTier::Psi => {
852 cost += penalty.value(target_t.view(), rho_local);
853 let grad = penalty.grad_target(target_t.view(), rho_local);
854 if grad.len() != flat_len {
855 crate::bail_invalid_estim!(
856 "analytic penalty {name:?} gradient length mismatch: got {}, expected {}",
857 grad.len(),
858 flat_len
859 );
860 }
861 for i in 0..flat_len {
862 gradient[t_start + i] += grad[i];
863 }
864 let grad_rho_local = penalty.grad_rho(target_t.view(), rho_local);
865 if grad_rho_local.len() != rho_slice.len() {
866 crate::bail_invalid_estim!(
867 "analytic penalty {name:?} rho-gradient length mismatch: got {}, expected {}",
868 grad_rho_local.len(),
869 rho_slice.len()
870 );
871 }
872 for local_idx in 0..grad_rho_local.len() {
873 gradient[rho_start + rho_slice.start + local_idx] += grad_rho_local[local_idx];
874 }
875 }
876 gam_terms::PenaltyTier::Beta => {}
877 gam_terms::PenaltyTier::Rho => {}
878 }
879 }
880 Ok(LatentIdObjectiveContribution { cost, gradient })
881}
882
883fn add_analytic_penalty_hessian_to_eval(
884 theta: &Array1<f64>,
885 rho_dim: usize,
886 latent: &gam_terms::latent::LatentCoordValues,
887 registry: &gam_terms::AnalyticPenaltyRegistry,
888 eval: &mut (f64, Array1<f64>, gam_problem::HessianValue),
889) -> Result<(), EstimationError> {
890 let flat_len = latent.len();
891 let t_start = rho_dim;
892 let t_end = t_start + flat_len;
893 let rho_start = t_end;
894 let rho_end = rho_start + registry.total_rho_count();
895 if theta.len() < rho_end {
896 crate::bail_invalid_estim!(
897 "latent-coordinate theta too short for analytic penalty Hessian: got {}, need at least {}",
898 theta.len(),
899 rho_end
900 );
901 }
902 let gam_problem::HessianValue::Dense(hessian) = &mut eval.2 else {
903 if eval.2.is_analytic() {
904 eval.2 = gam_problem::HessianValue::Unavailable;
905 }
906 return Ok(());
907 };
908 if hessian.dim() != (theta.len(), theta.len()) {
909 crate::bail_invalid_estim!(
910 "analytic penalty Hessian target shape mismatch: got {}x{}, expected {}x{}",
911 hessian.nrows(),
912 hessian.ncols(),
913 theta.len(),
914 theta.len()
915 );
916 }
917 let target_t = theta.slice(s![t_start..t_end]);
918 let rho = theta.slice(s![rho_start..rho_end]);
919 registry
920 .validate_rho(rho)
921 .map_err(EstimationError::InvalidInput)?;
922 for (penalty, (rho_slice, tier, _name)) in registry.penalties.iter().zip(registry.rho_layout())
923 {
924 let rho_local = rho.slice(s![rho_slice]);
925 if !matches!(tier, gam_terms::PenaltyTier::Psi) {
926 continue;
927 }
928 if let Some(diag) = penalty.hessian_diag(target_t.view(), rho_local) {
929 if diag.len() != flat_len {
930 crate::bail_invalid_estim!(
931 "analytic penalty Hessian diagonal length mismatch: got {}, expected {}",
932 diag.len(),
933 flat_len
934 );
935 }
936 for i in 0..flat_len {
937 hessian[[t_start + i, t_start + i]] += diag[i];
938 }
939 continue;
940 }
941 let mut probe = Array1::<f64>::zeros(flat_len);
942 for col in 0..flat_len {
943 probe[col] = 1.0;
944 let hv = penalty.hvp(target_t.view(), rho_local, probe.view());
945 if hv.len() != flat_len {
946 crate::bail_invalid_estim!(
947 "analytic penalty Hessian-vector length mismatch: got {}, expected {}",
948 hv.len(),
949 flat_len
950 );
951 }
952 for row in 0..flat_len {
953 hessian[[t_start + row, t_start + col]] += hv[row];
954 }
955 probe[col] = 0.0;
956 }
957 }
958 Ok(())
959}
960
961fn add_analytic_penalty_objective_to_eval(
962 theta: &Array1<f64>,
963 rho_dim: usize,
964 latent: &gam_terms::latent::LatentCoordValues,
965 registry: &gam_terms::AnalyticPenaltyRegistry,
966 eval: &mut (f64, Array1<f64>, gam_problem::HessianValue),
967) -> Result<(), EstimationError> {
968 let contribution = analytic_penalty_objective_contribution(theta, rho_dim, latent, registry)?;
969 eval.0 += contribution.cost;
970 if eval.1.len() != contribution.gradient.len() {
971 crate::bail_invalid_estim!(
972 "latent-coordinate REML gradient length mismatch: base={}, analytic_penalty={}",
973 eval.1.len(),
974 contribution.gradient.len()
975 );
976 }
977 eval.1 += &contribution.gradient;
978 add_analytic_penalty_hessian_to_eval(theta, rho_dim, latent, registry, eval)?;
979 Ok(())
980}
981
982fn spatial_log_kappa_hyper_dirs_frominfo_list(
983 info_list: Vec<SpatialPsiDerivative>,
984) -> Result<Vec<DirectionalHyperParam>, EstimationError> {
985 use gam_solve::estimate::reml::ImplicitDerivLevel;
986 use std::collections::HashMap;
987
988 let log_kappa_dim = info_list.len();
989 let group_ids: Vec<Option<usize>> = info_list.iter().map(|e| e.aniso_group_id).collect();
995 let mut group_indices_map: HashMap<usize, Vec<usize>> = HashMap::new();
996 for (idx, gid) in group_ids.iter().enumerate() {
997 if let Some(g) = gid {
998 group_indices_map.entry(*g).or_default().push(idx);
999 }
1000 }
1001
1002 let mut hyper_dirs = Vec::with_capacity(log_kappa_dim);
1003 for (i, info) in info_list.into_iter().enumerate() {
1004 let SpatialPsiDerivative {
1005 penalty_index: _,
1006 penalty_indices,
1007 global_range,
1008 total_p,
1009 x_psi_local,
1010 s_psi_components_local,
1011 x_psi_psi_local,
1012 s_psi_psi_components_local,
1013 aniso_group_id,
1014 aniso_cross_designs,
1015 aniso_cross_penalty_provider,
1016 implicit_operator,
1017 implicit_axis,
1018 } = info;
1019
1020 let mut xsecond = vec![None; log_kappa_dim];
1021 xsecond[i] = Some(if let Some(ref op) = implicit_operator {
1023 gam_solve::estimate::reml::HyperDesignDerivative::from_implicit(
1024 op.clone(),
1025 ImplicitDerivLevel::SecondDiag(implicit_axis),
1026 global_range.clone(),
1027 total_p,
1028 )
1029 } else {
1030 gam_solve::estimate::reml::HyperDesignDerivative::from_embedded(
1031 x_psi_psi_local,
1032 global_range.clone(),
1033 total_p,
1034 )
1035 });
1036 if let Some(cross_designs) = aniso_cross_designs {
1038 if let Some(gid) = aniso_group_id {
1042 let base = group_indices_map
1043 .get(&gid)
1044 .and_then(|v| v.first().copied())
1045 .unwrap_or(i);
1046 for (b_axis, cross_mat) in cross_designs.into_iter() {
1047 let j = base + b_axis;
1048 if j < log_kappa_dim {
1049 xsecond[j] = Some(if let Some(ref op) = implicit_operator {
1050 gam_solve::estimate::reml::HyperDesignDerivative::from_implicit(
1051 op.clone(),
1052 ImplicitDerivLevel::SecondCross(implicit_axis, b_axis),
1053 global_range.clone(),
1054 total_p,
1055 )
1056 } else {
1057 gam_solve::estimate::reml::HyperDesignDerivative::from_embedded(
1058 cross_mat,
1059 global_range.clone(),
1060 total_p,
1061 )
1062 });
1063 }
1064 }
1065 }
1066 }
1067 let s_components = penalty_indices
1068 .iter()
1069 .copied()
1070 .zip(s_psi_components_local.into_iter().map(|local| {
1071 gam_solve::estimate::reml::HyperPenaltyDerivative::from_embedded(
1072 local,
1073 global_range.clone(),
1074 total_p,
1075 )
1076 }))
1077 .collect::<Vec<_>>();
1078 let s2_components = penalty_indices
1079 .iter()
1080 .copied()
1081 .zip(s_psi_psi_components_local.into_iter().map(|local| {
1082 gam_solve::estimate::reml::HyperPenaltyDerivative::from_embedded(
1083 local,
1084 global_range.clone(),
1085 total_p,
1086 )
1087 }))
1088 .collect::<Vec<_>>();
1089 let mut ssecond_components = vec![None; log_kappa_dim];
1090 ssecond_components[i] = Some(s2_components);
1091 let mut penaltysecond_partner_indices: Option<Vec<usize>> = None;
1092 let penaltysecond_component_provider =
1093 if let (Some(provider), Some(gid)) = (aniso_cross_penalty_provider, aniso_group_id) {
1094 let group_indices = group_indices_map.get(&gid).cloned().unwrap_or_default();
1095 let axis_in_group =
1096 group_indices
1097 .iter()
1098 .position(|&idx| idx == i)
1099 .ok_or_else(|| {
1100 EstimationError::InvalidInput(format!(
1101 "missing spatial hyper axis {} in anisotropy group {}",
1102 i, gid
1103 ))
1104 })?;
1105 penaltysecond_partner_indices = Some(
1106 group_indices
1107 .iter()
1108 .copied()
1109 .filter(|&idx| idx != i)
1110 .collect(),
1111 );
1112 let penalty_indices_inner = penalty_indices.clone();
1113 let global_range_inner = global_range.clone();
1114 let total_p_inner = total_p;
1115 let group_indices_inner = group_indices;
1116 Some(std::sync::Arc::new(
1117 move |j: usize| -> Result<
1118 Option<Vec<gam_solve::estimate::reml::PenaltyDerivativeComponent>>,
1119 EstimationError,
1120 > {
1121 let Some(other_axis_in_group) =
1122 group_indices_inner.iter().position(|&idx| idx == j)
1123 else {
1124 return Ok(None);
1125 };
1126 if other_axis_in_group == axis_in_group {
1127 return Ok(None);
1128 }
1129 let cross_pens = provider(other_axis_in_group)?;
1130 if cross_pens.is_empty() {
1131 return Ok(None);
1132 }
1133 Ok(Some(
1134 penalty_indices_inner
1135 .iter()
1136 .copied()
1137 .zip(cross_pens.into_iter().map(|local| {
1138 gam_solve::estimate::reml::HyperPenaltyDerivative::from_embedded(
1139 local,
1140 global_range_inner.clone(),
1141 total_p_inner,
1142 )
1143 }))
1144 .map(|(penalty_index, matrix)| {
1145 gam_solve::estimate::reml::PenaltyDerivativeComponent {
1146 penalty_index,
1147 matrix,
1148 }
1149 })
1150 .collect(),
1151 ))
1152 },
1153 )
1154 as std::sync::Arc<
1155 dyn Fn(
1156 usize,
1157 ) -> Result<
1158 Option<Vec<gam_solve::estimate::reml::PenaltyDerivativeComponent>>,
1159 EstimationError,
1160 > + Send
1161 + Sync
1162 + 'static,
1163 >)
1164 } else {
1165 None
1166 };
1167 let x_first_hyper = if let Some(ref op) = implicit_operator {
1170 gam_solve::estimate::reml::HyperDesignDerivative::from_implicit(
1171 op.clone(),
1172 ImplicitDerivLevel::First(implicit_axis),
1173 global_range.clone(),
1174 total_p,
1175 )
1176 } else {
1177 gam_solve::estimate::reml::HyperDesignDerivative::from_embedded(
1178 x_psi_local,
1179 global_range.clone(),
1180 total_p,
1181 )
1182 };
1183 let mut dir = DirectionalHyperParam::new_compact(
1184 x_first_hyper,
1185 s_components,
1186 Some(xsecond),
1187 Some(ssecond_components),
1188 )?
1189 .not_penalty_like();
1190 if let Some(provider) = penaltysecond_component_provider {
1191 dir = dir.with_penaltysecond_component_provider(provider);
1192 }
1193 if let Some(partner_indices) = penaltysecond_partner_indices {
1194 dir = dir.with_penaltysecond_partner_indices(partner_indices);
1195 }
1196 hyper_dirs.push(dir);
1197 }
1198 Ok(hyper_dirs)
1199}
1200
1201pub(crate) fn spatial_dims_per_term(
1207 resolvedspec: &TermCollectionSpec,
1208 spatial_terms: &[usize],
1209) -> Vec<usize> {
1210 spatial_terms
1211 .iter()
1212 .map(|&term_idx| {
1213 if let Some(mj) = measure_jet_term_spec(resolvedspec, term_idx) {
1214 measure_jet_psi_dim(mj)
1217 } else if spatial_term_uses_per_axis_psi(resolvedspec, term_idx) {
1218 get_spatial_feature_dim(resolvedspec, term_idx).unwrap_or(1)
1219 } else {
1220 1
1221 }
1222 })
1223 .collect()
1224}
1225
1226fn has_aniso_terms(resolvedspec: &TermCollectionSpec, spatial_terms: &[usize]) -> bool {
1230 spatial_terms
1231 .iter()
1232 .any(|&term_idx| spatial_term_uses_per_axis_psi(resolvedspec, term_idx))
1233}
1234
1235macro_rules! impl_exact_joint_theta_memo {
1241 () => {
1242 fn memoized_cost(&self, theta: &Array1<f64>) -> Option<f64> {
1243 if self
1244 .current_theta
1245 .as_ref()
1246 .is_some_and(|cached| theta_values_match(cached, theta))
1247 {
1248 self.last_eval
1249 .as_ref()
1250 .map(|cached| cached.0)
1251 .or(self.last_cost)
1252 } else {
1253 None
1254 }
1255 }
1256
1257 fn memoized_eval(
1258 &self,
1259 theta: &Array1<f64>,
1260 ) -> Option<(f64, Array1<f64>, gam_problem::HessianValue)> {
1261 if self
1262 .current_theta
1263 .as_ref()
1264 .is_some_and(|cached| theta_values_match(cached, theta))
1265 {
1266 self.last_eval.clone()
1267 } else {
1268 None
1269 }
1270 }
1271
1272 fn store_eval(&mut self, eval: (f64, Array1<f64>, gam_problem::HessianValue)) {
1273 self.last_cost = Some(eval.0);
1274 self.last_eval = Some(eval);
1275 }
1276 };
1277}
1278
1279struct SingleBlockExactJointDesignCache<'d> {
1280 realizer: FrozenTermCollectionIncrementalRealizer<'d>,
1281 current_theta: Option<Array1<f64>>,
1282 last_eval_theta: Option<Array1<f64>>,
1289 last_cost: Option<f64>,
1290 last_eval: Option<(f64, Array1<f64>, gam_problem::HessianValue)>,
1291 cached_hyper_dirs: Option<(u64, Vec<DirectionalHyperParam>)>,
1303 spatial_terms: Vec<usize>,
1304 rho_dim: usize,
1305 dims_per_term: Vec<usize>,
1306}
1307
1308impl<'d> SingleBlockExactJointDesignCache<'d> {
1309 fn new_with_policy(
1310 data: ArrayView2<'d, f64>,
1311 spec: TermCollectionSpec,
1312 design: TermCollectionDesign,
1313 spatial_terms: Vec<usize>,
1314 rho_dim: usize,
1315 dims_per_term: Vec<usize>,
1316 policy: &gam_runtime::resource::ResourcePolicy,
1317 ) -> Result<Self, String> {
1318 Ok(Self {
1319 realizer: FrozenTermCollectionIncrementalRealizer::new_with_policy(
1320 data, spec, design, policy,
1321 )?,
1322 current_theta: None,
1323 last_eval_theta: None,
1324 last_cost: None,
1325 last_eval: None,
1326 cached_hyper_dirs: None,
1327 spatial_terms,
1328 rho_dim,
1329 dims_per_term,
1330 })
1331 }
1332
1333 fn design_revision(&self) -> u64 {
1334 self.realizer.design_revision()
1335 }
1336
1337 fn hyper_dirs_for_current_design(
1347 &mut self,
1348 data: ArrayView2<'_, f64>,
1349 kind: SpatialHyperKind,
1350 ) -> Result<Vec<DirectionalHyperParam>, EstimationError> {
1351 let revision = self.realizer.design_revision();
1352 if let Some((cached_rev, dirs)) = self.cached_hyper_dirs.as_ref()
1353 && *cached_rev == revision
1354 {
1355 return Ok(dirs.clone());
1356 }
1357 let dirs = try_build_spatial_log_kappa_hyper_dirs(
1358 data,
1359 self.realizer.spec(),
1360 self.realizer.design(),
1361 &self.spatial_terms,
1362 )?
1363 .ok_or_else(|| {
1364 EstimationError::InvalidInput(format!(
1365 "failed to build {} hyper_dirs at current {}",
1366 kind.adjective(),
1367 kind.coord_name(),
1368 ))
1369 })?;
1370 self.cached_hyper_dirs = Some((revision, dirs.clone()));
1371 Ok(dirs)
1372 }
1373
1374 fn nfree_tensor_gradient_hyper_dirs(
1375 &mut self,
1376 theta: &Array1<f64>,
1377 ) -> Result<Vec<DirectionalHyperParam>, EstimationError> {
1378 let psi = &theta.as_slice().ok_or_else(|| {
1379 EstimationError::InvalidInput(
1380 "nfree_tensor_gradient_hyper_dirs: theta is not contiguous".to_string(),
1381 )
1382 })?[self.rho_dim..];
1383 let (global_range, p_total, s_psi_components) = self
1384 .realizer
1385 .canonical_penalty_derivatives_at_psi(&self.spatial_terms, psi)
1386 .map_err(EstimationError::InvalidInput)?;
1387 let zero_x = gam_solve::estimate::reml::HyperDesignDerivative::zero(
1388 self.realizer.design().design.nrows(),
1389 p_total,
1390 );
1391 let components = s_psi_components
1392 .into_iter()
1393 .enumerate()
1394 .map(|(penalty_index, local)| {
1395 (
1396 penalty_index,
1397 gam_solve::estimate::reml::HyperPenaltyDerivative::from_embedded(
1398 local,
1399 global_range.clone(),
1400 p_total,
1401 ),
1402 )
1403 })
1404 .collect::<Vec<_>>();
1405 Ok(DirectionalHyperParam::new_compact(zero_x, components, None, None)?.not_penalty_like())
1406 .map(|dir| vec![dir])
1407 }
1408
1409 fn ensure_theta(&mut self, theta: &Array1<f64>) -> Result<(), String> {
1410 if self
1411 .current_theta
1412 .as_ref()
1413 .is_some_and(|cached| theta_values_match(cached, theta))
1414 {
1415 return Ok(());
1416 }
1417 let t_ensure = std::time::Instant::now();
1418 let log_kappa = SpatialLogKappaCoords::from_theta_tail_with_dims(
1419 theta,
1420 self.rho_dim,
1421 self.dims_per_term.clone(),
1422 );
1423 self.realizer
1424 .apply_log_kappa(&log_kappa, &self.spatial_terms)?;
1425 log::info!(
1426 "[STAGE] ensure_theta (apply_log_kappa, {} terms): {:.3}s",
1427 self.spatial_terms.len(),
1428 t_ensure.elapsed().as_secs_f64(),
1429 );
1430 self.current_theta = Some(theta.clone());
1431 self.last_eval_theta = None;
1432 self.last_cost = None;
1433 self.last_eval = None;
1434 Ok(())
1435 }
1436
1437 fn memoized_cost(&self, theta: &Array1<f64>) -> Option<f64> {
1444 if self
1445 .last_eval_theta
1446 .as_ref()
1447 .is_some_and(|cached| theta_values_match(cached, theta))
1448 {
1449 self.last_eval
1450 .as_ref()
1451 .map(|cached| cached.0)
1452 .or(self.last_cost)
1453 } else {
1454 None
1455 }
1456 }
1457
1458 fn memoized_eval(
1459 &self,
1460 theta: &Array1<f64>,
1461 ) -> Option<(f64, Array1<f64>, gam_problem::HessianValue)> {
1462 if self
1463 .last_eval_theta
1464 .as_ref()
1465 .is_some_and(|cached| theta_values_match(cached, theta))
1466 {
1467 self.last_eval.clone()
1468 } else {
1469 None
1470 }
1471 }
1472
1473 fn store_eval_at(
1477 &mut self,
1478 theta: &Array1<f64>,
1479 eval: (f64, Array1<f64>, gam_problem::HessianValue),
1480 ) {
1481 self.last_eval_theta = Some(theta.clone());
1482 self.last_cost = Some(eval.0);
1483 self.last_eval = Some(eval);
1484 }
1485
1486 fn store_cost_at(&mut self, theta: &Array1<f64>, cost: f64) {
1489 self.last_eval_theta = Some(theta.clone());
1490 self.last_cost = Some(cost);
1491 self.last_eval = None;
1495 }
1496
1497 fn spec(&self) -> &TermCollectionSpec {
1498 self.realizer.spec()
1499 }
1500
1501 fn design(&self) -> &TermCollectionDesign {
1502 self.realizer.design()
1503 }
1504
1505 fn supports_nfree_penalty_rekey(&self) -> bool {
1511 self.realizer
1512 .supports_nfree_penalty_rekey(&self.spatial_terms)
1513 }
1514
1515 fn supports_nfree_gradient_only_routing(&self) -> bool {
1516 self.realizer
1517 .supports_nfree_gradient_only_routing(&self.spatial_terms)
1518 }
1519
1520 fn canonical_penalties_at(
1530 &mut self,
1531 theta: &Array1<f64>,
1532 ) -> Result<(Vec<gam_terms::construction::CanonicalPenalty>, Vec<usize>), String> {
1533 let psi = &theta
1534 .as_slice()
1535 .ok_or_else(|| "canonical_penalties_at: theta is not contiguous".to_string())?
1536 [self.rho_dim..];
1537 self.realizer
1538 .canonical_penalties_at_psi(&self.spatial_terms, psi)
1539 }
1540}
1541
1542struct SingleBlockLatentCoordDesignCache {
1543 data: Array2<f64>,
1544 spec: TermCollectionSpec,
1545 design: TermCollectionDesign,
1546 current_theta: Option<Array1<f64>>,
1547 current_latent: Option<std::sync::Arc<gam_terms::latent::LatentCoordValues>>,
1548 current_hyper_dirs: Option<Vec<gam_solve::estimate::reml::DirectionalHyperParam>>,
1549 current_design_cache_id: Option<u64>,
1550 latent_design_cache: gam_solve::latent_cache::LatentDesignCache,
1551 last_cost: Option<f64>,
1552 last_eval: Option<(f64, Array1<f64>, gam_problem::HessianValue)>,
1553 term_index: gam_problem::types::SmoothTermIdx,
1554 feature_cols: Vec<usize>,
1555 rho_dim: usize,
1556 n_obs: usize,
1557 latent_dim: usize,
1558 id_mode: gam_terms::latent::LatentIdMode,
1559 manifold: gam_terms::latent::LatentManifold,
1560 retraction_registry: gam_solve::latent_cache::LatentRetractionRegistry,
1561 latent_id: u64,
1562 analytic_penalties: Option<std::sync::Arc<gam_terms::AnalyticPenaltyRegistry>>,
1563 analytic_rho_count: usize,
1564 design_revision: u64,
1565 last_outer_iter: Option<u64>,
1569}
1570
1571impl SingleBlockLatentCoordDesignCache {
1572 fn new(
1573 data: Array2<f64>,
1574 spec: TermCollectionSpec,
1575 design: TermCollectionDesign,
1576 latent: &StandardLatentCoordConfig,
1577 rho_dim: usize,
1578 ) -> Result<Self, String> {
1579 if latent.term_index.get() >= spec.smooth_terms.len() {
1580 return Err(SmoothError::dimension_mismatch(format!(
1581 "latent-coordinate term index {} out of bounds for {} smooth terms",
1582 latent.term_index,
1583 spec.smooth_terms.len()
1584 ))
1585 .into());
1586 }
1587 if latent.feature_cols.len() != latent.values.latent_dim() {
1588 return Err(SmoothError::dimension_mismatch(format!(
1589 "latent-coordinate feature width mismatch: feature_cols={}, latent_dim={}",
1590 latent.feature_cols.len(),
1591 latent.values.latent_dim()
1592 ))
1593 .into());
1594 }
1595 if latent.values.n_obs() != data.nrows() {
1596 return Err(SmoothError::dimension_mismatch(format!(
1597 "latent-coordinate row mismatch: latent n={}, data n={}",
1598 latent.values.n_obs(),
1599 data.nrows()
1600 ))
1601 .into());
1602 }
1603 let analytic_rho_count = latent
1604 .analytic_penalties
1605 .as_ref()
1606 .map_or(0, |registry| registry.total_rho_count());
1607 Ok(Self {
1608 data,
1609 spec,
1610 design,
1611 current_theta: None,
1612 current_latent: None,
1613 current_hyper_dirs: None,
1614 current_design_cache_id: None,
1615 latent_design_cache: gam_solve::latent_cache::LatentDesignCache::default(),
1616 last_cost: None,
1617 last_eval: None,
1618 term_index: latent.term_index,
1619 feature_cols: latent.feature_cols.clone(),
1620 rho_dim,
1621 n_obs: latent.values.n_obs(),
1622 latent_dim: latent.values.latent_dim(),
1623 id_mode: latent.values.id_mode().clone(),
1624 manifold: latent.values.manifold().clone(),
1625 retraction_registry: latent.values.retraction_registry().clone(),
1626 latent_id: latent.values.latent_id(),
1627 analytic_penalties: latent.analytic_penalties.clone(),
1628 analytic_rho_count,
1629 design_revision: 0,
1630 last_outer_iter: None,
1631 })
1632 }
1633
1634 fn design_revision(&self) -> u64 {
1635 self.design_revision
1636 }
1637
1638 fn design(&self) -> &TermCollectionDesign {
1639 &self.design
1640 }
1641
1642 fn latent(&self) -> Result<std::sync::Arc<gam_terms::latent::LatentCoordValues>, String> {
1643 self.current_latent
1644 .as_ref()
1645 .cloned()
1646 .ok_or_else(|| "latent-coordinate cache has not been realized".to_string())
1647 }
1648
1649 fn analytic_penalties(&self) -> Option<std::sync::Arc<gam_terms::AnalyticPenaltyRegistry>> {
1650 self.analytic_penalties.clone()
1651 }
1652
1653 fn analytic_penalty_rho_count(&self) -> usize {
1654 self.analytic_rho_count
1655 }
1656
1657 fn hyper_dirs(&self) -> Result<Vec<gam_solve::estimate::reml::DirectionalHyperParam>, String> {
1658 self.current_hyper_dirs
1659 .as_ref()
1660 .cloned()
1661 .ok_or_else(|| "latent-coordinate hyper_dirs cache has not been realized".to_string())
1662 }
1663
1664 fn latent_basis_kind(&self) -> Result<gam_solve::latent_cache::LatentBasisKind, String> {
1665 let smooth_term = self
1666 .design
1667 .smooth
1668 .terms
1669 .get(self.term_index.get())
1670 .ok_or_else(|| {
1671 SmoothError::dimension_mismatch(format!(
1672 "LatentCoord term index {} out of bounds for realized smooth design",
1673 self.term_index
1674 ))
1675 })?;
1676 let termspec = self
1677 .spec
1678 .smooth_terms
1679 .get(self.term_index.get())
1680 .ok_or_else(|| {
1681 SmoothError::dimension_mismatch(format!(
1682 "LatentCoord term index {} out of bounds for resolved smooth spec",
1683 self.term_index
1684 ))
1685 })?;
1686 match (&termspec.basis, &smooth_term.metadata) {
1687 (
1688 SmoothBasisSpec::Matern { .. },
1689 BasisMetadata::Matern {
1690 centers,
1691 length_scale,
1692 nu,
1693 aniso_log_scales,
1694 ..
1695 },
1696 ) => Ok(gam_solve::latent_cache::LatentBasisKind::Matern {
1697 centers: centers.clone(),
1698 length_scale: *length_scale,
1699 nu: *nu,
1700 aniso_log_scales: aniso_log_scales
1701 .clone()
1702 .unwrap_or_else(|| vec![0.0; centers.ncols()]),
1703 chunk_size: gam_terms::basis::auto_streaming_chunk_size_for_dense(
1704 self.n_obs,
1705 centers.nrows(),
1706 ),
1707 }),
1708 (
1709 SmoothBasisSpec::Duchon { .. },
1710 BasisMetadata::Duchon {
1711 centers,
1712 length_scale,
1713 power,
1714 nullspace_order,
1715 aniso_log_scales,
1716 ..
1717 },
1718 ) => Ok(gam_solve::latent_cache::LatentBasisKind::Duchon {
1719 centers: centers.clone(),
1720 length_scale: *length_scale,
1721 power: *power,
1722 nullspace_order: *nullspace_order,
1723 aniso_log_scales: aniso_log_scales
1724 .clone()
1725 .unwrap_or_else(|| vec![0.0; centers.ncols()]),
1726 }),
1727 (
1728 SmoothBasisSpec::Sphere { .. },
1729 BasisMetadata::Sphere {
1730 centers,
1731 penalty_order,
1732 method,
1733 ..
1734 },
1735 ) if matches!(*method, gam_terms::basis::SphereMethod::Wahba) => {
1736 Ok(gam_solve::latent_cache::LatentBasisKind::Sphere {
1737 centers: centers.clone(),
1738 penalty_order: *penalty_order,
1739 chunk_size: gam_terms::basis::auto_streaming_chunk_size_for_dense(
1740 self.n_obs,
1741 centers.nrows(),
1742 ),
1743 })
1744 }
1745 (
1746 SmoothBasisSpec::BSpline1D { spec, .. },
1747 BasisMetadata::BSpline1D {
1748 knots,
1749 periodic,
1750 degree: meta_degree,
1751 ..
1752 },
1753 ) => {
1754 let effective_degree = meta_degree.unwrap_or(spec.degree);
1758 if let Some((domain_start, period, num_basis)) = periodic {
1759 Ok(gam_solve::latent_cache::LatentBasisKind::PeriodicBspline {
1760 domain_start: *domain_start,
1761 period: *period,
1762 degree: effective_degree,
1763 num_basis: *num_basis,
1764 chunk_size: gam_terms::basis::auto_streaming_chunk_size_for_dense(
1765 self.n_obs, *num_basis,
1766 ),
1767 })
1768 } else {
1769 let num_basis_est = knots.len().saturating_sub(effective_degree + 1);
1770 Ok(gam_solve::latent_cache::LatentBasisKind::TensorBspline {
1771 knots: vec![knots.clone()],
1772 degrees: vec![effective_degree],
1773 chunk_size: gam_terms::basis::auto_streaming_chunk_size_for_dense(
1774 self.n_obs,
1775 num_basis_est,
1776 ),
1777 })
1778 }
1779 }
1780 (
1781 SmoothBasisSpec::TensorBSpline { .. },
1782 BasisMetadata::TensorBSpline { knots, degrees, .. },
1783 ) => Ok(gam_solve::latent_cache::LatentBasisKind::TensorBspline {
1784 knots: knots.clone(),
1785 degrees: degrees.clone(),
1786 chunk_size: None,
1787 }),
1788 (
1789 SmoothBasisSpec::Pca { .. },
1790 BasisMetadata::Pca {
1791 basis_matrix,
1792 centered,
1793 smooth_penalty,
1794 center_mean,
1795 pca_basis_path,
1796 chunk_size,
1797 ..
1798 },
1799 ) => {
1800 let center_mean_fingerprint = if *centered && pca_basis_path.is_none() {
1801 let mean = center_mean.as_ref().ok_or_else(|| {
1802 SmoothError::invalid_config(
1803 "latent-coordinate Pca cache key requires center_mean when centered",
1804 )
1805 })?;
1806 Some(gam_solve::latent_cache::pca_center_mean_fingerprint(mean))
1807 } else {
1808 None
1809 };
1810 Ok(gam_solve::latent_cache::LatentBasisKind::Pca {
1811 basis_matrix: basis_matrix.clone(),
1812 centered: *centered,
1813 center_mean_fingerprint,
1814 smooth_penalty: *smooth_penalty,
1815 pca_basis_path: pca_basis_path.clone(),
1816 chunk_size: *chunk_size,
1817 })
1818 }
1819 _ => Err(SmoothError::invalid_config(
1820 "latent-coordinate design cache could not key the realized latent smooth basis"
1821 .to_string(),
1822 )
1823 .into()),
1824 }
1825 }
1826
1827 fn ensure_theta(&mut self, theta: &Array1<f64>) -> Result<(), String> {
1828 if self
1829 .current_theta
1830 .as_ref()
1831 .is_some_and(|cached| theta_values_match(cached, theta))
1832 {
1833 return Ok(());
1834 }
1835 let latent_flat_len = self.n_obs * self.latent_dim;
1836 let direct_hyper_count = latent_coord_direct_hyper_count(&self.id_mode, self.latent_dim);
1837 let expected =
1838 self.rho_dim + latent_flat_len + self.analytic_rho_count + direct_hyper_count;
1839 if theta.len() != expected {
1840 return Err(SmoothError::dimension_mismatch(format!(
1841 "latent-coordinate theta length mismatch: got {}, expected {} (rho_dim={}, n={}, d={}, analytic_rhos={}, direct_hypers={})",
1842 theta.len(),
1843 expected,
1844 self.rho_dim,
1845 self.n_obs,
1846 self.latent_dim,
1847 self.analytic_rho_count,
1848 direct_hyper_count
1849 ))
1850 .into());
1851 }
1852 let flat = theta
1853 .slice(s![self.rho_dim..self.rho_dim + latent_flat_len])
1854 .to_owned();
1855 let latent = std::sync::Arc::new(
1856 gam_terms::latent::LatentCoordValues::from_flat_with_manifold_and_retraction_and_id(
1857 flat,
1858 self.n_obs,
1859 self.latent_dim,
1860 self.id_mode.clone(),
1861 self.manifold.clone(),
1862 self.retraction_registry.clone(),
1863 self.latent_id,
1864 ),
1865 );
1866 let latent_values_changed = self
1867 .current_latent
1868 .as_ref()
1869 .map(|cached| !latent_values_match(cached.as_flat(), latent.as_flat()))
1870 .unwrap_or(true);
1871 if latent_values_changed {
1872 self.latent_design_cache.invalidate_all();
1873 self.current_design_cache_id = None;
1874 self.design_revision = self.design_revision.wrapping_add(1);
1875 }
1876 for n in 0..self.n_obs {
1877 for axis in 0..self.latent_dim {
1878 let col = self.feature_cols[axis];
1879 self.data[[n, col]] = latent.as_flat()[n * self.latent_dim + axis];
1880 }
1881 }
1882
1883 let basis_kind = self.latent_basis_kind()?;
1884 let rebuilt_width = self.design.design.ncols();
1885 let spec = self.spec.clone();
1886 let term_index = self.term_index;
1887 let analytic_rho_count = self.analytic_rho_count;
1888 let data = self.data.view();
1889 let design_context_digest = gam_solve::latent_cache::latent_design_context_cache_digest(
1890 data,
1891 &spec,
1892 term_index,
1893 analytic_rho_count,
1894 &self.feature_cols,
1895 )
1896 .map_err(|e| e.to_string())?;
1897 let lookup = self
1898 .latent_design_cache
1899 .lookup_or_compute(latent.clone(), basis_kind, design_context_digest, || {
1900 let rebuilt = build_term_collection_design(data, &spec).map_err(|e| {
1901 EstimationError::InvalidInput(format!(
1902 "failed to rebuild latent-coordinate design: {e}"
1903 ))
1904 })?;
1905 if rebuilt.design.ncols() != rebuilt_width {
1906 crate::bail_invalid_estim!(
1907 "latent-coordinate design topology changed: rebuilt p={}, cached p={}",
1908 rebuilt.design.ncols(),
1909 rebuilt_width
1910 );
1911 }
1912 let hyper_dirs = try_build_latent_coord_hyper_dirs(
1913 latent.clone(),
1914 &spec,
1915 &rebuilt,
1916 &[term_index],
1917 analytic_rho_count,
1918 )?
1919 .ok_or_else(|| {
1920 EstimationError::InvalidInput(
1921 "failed to build latent-coordinate hyper_dirs".to_string(),
1922 )
1923 })?;
1924 Ok(gam_solve::latent_cache::ComputedLatentDesign {
1925 design: rebuilt,
1926 hyper_dirs,
1927 })
1928 })
1929 .map_err(|e| e.to_string())?;
1930 if lookup.cached.design.design.ncols() != self.design.design.ncols() {
1931 return Err(SmoothError::dimension_mismatch(format!(
1932 "latent-coordinate design topology changed: rebuilt p={}, cached p={}",
1933 lookup.cached.design.design.ncols(),
1934 self.design.design.ncols()
1935 ))
1936 .into());
1937 }
1938 self.design = lookup.cached.design.clone();
1939 self.current_hyper_dirs = Some(lookup.cached.hyper_dirs.clone());
1940 self.current_latent = Some(latent);
1941 self.current_theta = Some(theta.clone());
1942 self.last_cost = None;
1943 self.last_eval = None;
1944 self.last_outer_iter = None;
1945 if !latent_values_changed && self.current_design_cache_id != Some(lookup.entry_id) {
1946 self.design_revision = self.design_revision.wrapping_add(1);
1947 }
1948 self.current_design_cache_id = Some(lookup.entry_id);
1949 Ok(())
1950 }
1951
1952 fn memoized_cost(&self, theta: &Array1<f64>) -> Option<f64> {
1953 if self
1954 .current_theta
1955 .as_ref()
1956 .is_some_and(|cached| theta_values_match(cached, theta))
1957 && self.last_outer_iter
1958 == Some(gam_solve::estimate::reml::outer_eval::current_outer_iter())
1959 {
1960 self.last_eval
1961 .as_ref()
1962 .map(|cached| cached.0)
1963 .or(self.last_cost)
1964 } else {
1965 None
1966 }
1967 }
1968
1969 fn memoized_eval(
1970 &self,
1971 theta: &Array1<f64>,
1972 ) -> Option<(f64, Array1<f64>, gam_problem::HessianValue)> {
1973 if self
1974 .current_theta
1975 .as_ref()
1976 .is_some_and(|cached| theta_values_match(cached, theta))
1977 && self.last_outer_iter
1978 == Some(gam_solve::estimate::reml::outer_eval::current_outer_iter())
1979 {
1980 self.last_eval.clone()
1981 } else {
1982 None
1983 }
1984 }
1985
1986 fn store_eval(&mut self, eval: (f64, Array1<f64>, gam_problem::HessianValue)) {
1987 self.last_cost = Some(eval.0);
1988 self.last_eval = Some(eval);
1989 self.last_outer_iter = Some(gam_solve::estimate::reml::outer_eval::current_outer_iter());
1990 }
1991
1992 fn store_cost(&mut self, cost: f64) {
1993 self.last_cost = Some(cost);
1994 self.last_outer_iter = Some(gam_solve::estimate::reml::outer_eval::current_outer_iter());
1995 }
1996
1997 fn reset(&mut self) {
1998 self.current_theta = None;
1999 self.current_latent = None;
2000 self.current_hyper_dirs = None;
2001 self.current_design_cache_id = None;
2002 self.latent_design_cache.invalidate();
2003 self.last_cost = None;
2004 self.last_eval = None;
2005 self.last_outer_iter = None;
2006 }
2007}
2008
2009pub fn fixed_kappa_profiled_reml_score(
2025 data: ArrayView2<'_, f64>,
2026 y: ArrayView1<'_, f64>,
2027 weights: ArrayView1<'_, f64>,
2028 offset: ArrayView1<'_, f64>,
2029 resolvedspec: &TermCollectionSpec,
2030 term_idx: usize,
2031 kappa: f64,
2032 family: LikelihoodSpec,
2033 options: &FitOptions,
2034) -> Result<f64, EstimationError> {
2035 if !kappa.is_finite() {
2036 crate::bail_invalid_estim!("fixed-κ profiled score probed a non-finite κ = {kappa}");
2037 }
2038 let (feature_cols, mut probe_basis) =
2041 match resolvedspec.smooth_terms.get(term_idx).map(|t| &t.basis) {
2042 Some(SmoothBasisSpec::ConstantCurvature {
2043 feature_cols, spec, ..
2044 }) => (feature_cols.clone(), spec.clone()),
2045 _ => {
2046 crate::bail_invalid_estim!(
2047 "fixed-κ profiled score: term {term_idx} is not a constant-curvature smooth"
2048 )
2049 }
2050 };
2051 probe_basis.kappa = kappa;
2052
2053 let is_unweighted = weights.iter().all(|&w| (w - 1.0).abs() <= 1e-12);
2073 let is_zero_offset = offset.iter().all(|&o| o.abs() <= 1e-12);
2074 if family == LikelihoodSpec::gaussian_identity() && is_unweighted && is_zero_offset {
2075 let x_term = select_columns(data, &feature_cols).map_err(EstimationError::from)?;
2076 let score = gam_terms::basis::constant_curvature_honest_profiled_reml_score(
2077 x_term.view(),
2078 y,
2079 &probe_basis,
2080 )
2081 .map_err(|e| {
2082 EstimationError::InvalidInput(format!(
2083 "fixed-κ honest profiled-REML score at κ={kappa} failed: {e}"
2084 ))
2085 })?;
2086 if !score.is_finite() {
2087 crate::bail_invalid_estim!(
2088 "fixed-κ honest profiled-REML score at κ={kappa} is non-finite"
2089 );
2090 }
2091 return Ok(score);
2092 }
2093
2094 let mut probe_spec = resolvedspec.clone();
2096 match probe_spec
2097 .smooth_terms
2098 .get_mut(term_idx)
2099 .map(|t| &mut t.basis)
2100 {
2101 Some(SmoothBasisSpec::ConstantCurvature { spec, .. }) => spec.kappa = kappa,
2102 _ => {
2103 crate::bail_invalid_estim!(
2104 "fixed-κ profiled score: term {term_idx} is not a constant-curvature smooth"
2105 )
2106 }
2107 }
2108 let fixed_kappa_options = SpatialLengthScaleOptimizationOptions {
2109 enabled: false,
2110 ..SpatialLengthScaleOptimizationOptions::default()
2111 };
2112 let fit = fit_term_collectionwith_spatial_length_scale_optimization(
2113 data,
2114 y.to_owned(),
2115 weights.to_owned(),
2116 offset.to_owned(),
2117 &probe_spec,
2118 family,
2119 options,
2120 &fixed_kappa_options,
2121 )?;
2122 let score = fit_score(&fit.fit);
2123 if !score.is_finite() {
2124 crate::bail_invalid_estim!("fixed-κ profiled fit at κ={kappa} returned a non-finite score");
2125 }
2126 Ok(score)
2127}
2128
2129fn profiled_gaussian_reml_value_kappa_gradient(
2134 design: &Array2<f64>,
2135 design_kappa: &Array2<f64>,
2136 penalty: &Array2<f64>,
2137 penalty_kappa: &Array2<f64>,
2138 response: ArrayView1<'_, f64>,
2139) -> Result<(f64, f64), EstimationError> {
2140 if design.dim() != design_kappa.dim()
2141 || penalty.dim() != penalty_kappa.dim()
2142 || penalty.dim() != (design.ncols(), design.ncols())
2143 || response.len() != design.nrows()
2144 {
2145 crate::bail_invalid_estim!("constant-curvature profile value/gradient shape mismatch");
2146 }
2147
2148 let response_2d = response.insert_axis(ndarray::Axis(1));
2149 let fit = gam_solve::gaussian_reml::gaussian_reml_multi_closed_form(
2150 design.view(),
2151 response_2d.view(),
2152 penalty.view(),
2153 None,
2154 None,
2155 )?;
2156 let backward = gam_solve::gaussian_reml::gaussian_reml_multi_closed_form_backward_from_fit(
2157 design.view(),
2158 response_2d.view(),
2159 penalty.view(),
2160 None,
2161 &fit,
2162 0.0,
2163 None,
2164 None,
2165 1.0,
2166 0.0,
2167 )?;
2168 let derivative = backward
2169 .grad_x
2170 .iter()
2171 .zip(design_kappa.iter())
2172 .map(|(&adjoint, &direction)| adjoint * direction)
2173 .sum::<f64>()
2174 + backward
2175 .grad_penalty
2176 .iter()
2177 .zip(penalty_kappa.iter())
2178 .map(|(&adjoint, &direction)| adjoint * direction)
2179 .sum::<f64>();
2180 if !(fit.reml_score.is_finite() && derivative.is_finite()) {
2181 crate::bail_invalid_estim!(
2182 "constant-curvature analytic profile returned a non-finite value or derivative"
2183 );
2184 }
2185 Ok((fit.reml_score, derivative))
2186}
2187
2188fn constant_curvature_radial_reference(
2192 data: ArrayView2<'_, f64>,
2193 y: ArrayView1<'_, f64>,
2194) -> Result<Array1<f64>, EstimationError> {
2195 if y.len() != data.nrows() || y.is_empty() {
2196 crate::bail_invalid_estim!(
2197 "constant-curvature radial reference needs one non-empty response per row"
2198 );
2199 }
2200 let radii: Array1<f64> = data.outer_iter().map(|row| row.dot(&row).sqrt()).collect();
2201 let r_max = radii.iter().copied().fold(0.0_f64, f64::max);
2202 if r_max <= f64::MIN_POSITIVE {
2203 let mean = y.sum() / y.len() as f64;
2204 return Ok(Array1::from_elem(y.len(), mean));
2205 }
2206
2207 let bin_count = (data.nrows() as f64).log2().ceil() as usize + 1;
2208 let bin_of = |radius: f64| -> usize {
2209 ((radius / r_max * bin_count as f64) as usize).min(bin_count - 1)
2210 };
2211 let mut sums = vec![0.0; bin_count];
2212 let mut counts = vec![0usize; bin_count];
2213 for (row, &radius) in radii.iter().enumerate() {
2214 let bin = bin_of(radius);
2215 sums[bin] += y[row];
2216 counts[bin] += 1;
2217 }
2218 let means: Vec<f64> = sums
2219 .into_iter()
2220 .zip(counts)
2221 .map(
2222 |(sum, count)| {
2223 if count == 0 { 0.0 } else { sum / count as f64 }
2224 },
2225 )
2226 .collect();
2227 Ok(radii.mapv(|radius| means[bin_of(radius)]))
2228}
2229
2230fn constant_curvature_kappa_fair_profile_value_gradient(
2235 data: ArrayView2<'_, f64>,
2236 y: ArrayView1<'_, f64>,
2237 y_ref: ArrayView1<'_, f64>,
2238 spec: &gam_terms::basis::ConstantCurvatureBasisSpec,
2239) -> Result<(f64, f64), EstimationError> {
2240 if y.len() != data.nrows() || y_ref.len() != data.nrows() {
2241 crate::bail_invalid_estim!(
2242 "constant-curvature fair profile row mismatch: data={}, response={}, reference={}",
2243 data.nrows(),
2244 y.len(),
2245 y_ref.len(),
2246 );
2247 }
2248
2249 let mut profile_spec = spec.clone();
2250 profile_spec.double_penalty = false;
2251 let basis = gam_terms::basis::build_constant_curvature_basis(data, &profile_spec)
2252 .map_err(EstimationError::from)?;
2253 let derivatives =
2254 gam_terms::basis::build_constant_curvature_basis_kappa_derivatives(data, &profile_spec)
2255 .map_err(EstimationError::from)?;
2256 if basis.active_penalties.len() != 1 || derivatives.first.penalties_derivative.len() != 1 {
2257 crate::bail_invalid_estim!(
2258 "constant-curvature fair profile expected one primary penalty; value blocks={}, derivative blocks={}",
2259 basis.active_penalties.len(),
2260 derivatives.first.penalties_derivative.len(),
2261 );
2262 }
2263
2264 let smooth_design = basis.design.to_dense();
2265 let smooth_design_kappa = &derivatives.first.design_derivative;
2266 let smooth_penalty = &basis.active_penalties[0].matrix;
2267 let smooth_penalty_kappa = &derivatives.first.penalties_derivative[0];
2268 let n = smooth_design.nrows();
2269 let p = smooth_design.ncols();
2270 if smooth_design_kappa.dim() != (n, p)
2271 || smooth_penalty.dim() != (p, p)
2272 || smooth_penalty_kappa.dim() != (p, p)
2273 {
2274 crate::bail_invalid_estim!(
2275 "constant-curvature kappa derivative bundle does not match its value basis"
2276 );
2277 }
2278
2279 let mut design = Array2::<f64>::ones((n, p + 1));
2280 design.slice_mut(s![.., 1..]).assign(&smooth_design);
2281 let mut design_kappa = Array2::<f64>::zeros((n, p + 1));
2282 design_kappa
2283 .slice_mut(s![.., 1..])
2284 .assign(smooth_design_kappa);
2285 let mut penalty = Array2::<f64>::zeros((p + 1, p + 1));
2286 penalty.slice_mut(s![1.., 1..]).assign(smooth_penalty);
2287 let mut penalty_kappa = Array2::<f64>::zeros((p + 1, p + 1));
2288 penalty_kappa
2289 .slice_mut(s![1.., 1..])
2290 .assign(smooth_penalty_kappa);
2291
2292 let (value_y, derivative_y) = profiled_gaussian_reml_value_kappa_gradient(
2293 &design,
2294 &design_kappa,
2295 &penalty,
2296 &penalty_kappa,
2297 y,
2298 )?;
2299 let (value_ref, derivative_ref) = profiled_gaussian_reml_value_kappa_gradient(
2300 &design,
2301 &design_kappa,
2302 &penalty,
2303 &penalty_kappa,
2304 y_ref,
2305 )?;
2306 Ok((value_y - value_ref, derivative_y - derivative_ref))
2307}
2308
2309struct ConstantCurvatureFairProfile<'a> {
2310 data: ArrayView2<'a, f64>,
2311 response: ArrayView1<'a, f64>,
2312 radial_reference: Array1<f64>,
2313 spec: gam_terms::basis::ConstantCurvatureBasisSpec,
2314 cache: std::cell::RefCell<std::collections::HashMap<u64, (f64, f64)>>,
2315}
2316
2317impl ConstantCurvatureFairProfile<'_> {
2318 fn evaluate(&self, kappa: f64) -> Result<(f64, f64), EstimationError> {
2319 if !kappa.is_finite() {
2320 crate::bail_invalid_estim!("constant-curvature fair profile probed a non-finite kappa");
2321 }
2322 let key = kappa.to_bits();
2323 if let Some(&cached) = self.cache.borrow().get(&key) {
2324 return Ok(cached);
2325 }
2326 let mut probe_spec = self.spec.clone();
2327 probe_spec.kappa = kappa;
2328 let sample = constant_curvature_kappa_fair_profile_value_gradient(
2329 self.data,
2330 self.response,
2331 self.radial_reference.view(),
2332 &probe_spec,
2333 )?;
2334 self.cache.borrow_mut().insert(key, sample);
2335 Ok(sample)
2336 }
2337}
2338
2339fn validate_constant_curvature_fair_profile_inputs(
2340 weights: ArrayView1<'_, f64>,
2341 offset: ArrayView1<'_, f64>,
2342 family: &LikelihoodSpec,
2343) -> Result<(), EstimationError> {
2344 if *family != LikelihoodSpec::gaussian_identity() {
2345 crate::bail_invalid_estim!(
2346 "curvature-as-an-estimand profile currently requires Gaussian identity likelihood"
2347 );
2348 }
2349 let input_tolerance = f64::EPSILON.sqrt();
2350 if weights
2351 .iter()
2352 .any(|&weight| (weight - 1.0).abs() > input_tolerance)
2353 || offset.iter().any(|&value| value.abs() > input_tolerance)
2354 {
2355 crate::bail_invalid_estim!(
2356 "curvature-as-an-estimand profile requires unit weights and zero offset"
2357 );
2358 }
2359 Ok(())
2360}
2361
2362fn constant_curvature_kappa_fair_optimum(
2369 data: ArrayView2<'_, f64>,
2370 y: ArrayView1<'_, f64>,
2371 resolvedspec: &TermCollectionSpec,
2372 term_idx: usize,
2373 options: &FitOptions,
2374) -> Result<f64, EstimationError> {
2375 let (kappa_min, kappa_max) = constant_curvature_kappa_bounds(data, resolvedspec, term_idx);
2376 if !(kappa_min.is_finite() && kappa_max.is_finite() && kappa_max > kappa_min) {
2377 crate::bail_invalid_estim!(
2378 "constant-curvature term {term_idx} has invalid kappa bounds [{kappa_min}, {kappa_max}]"
2379 );
2380 }
2381 let (feature_cols, base_spec) = match resolvedspec
2382 .smooth_terms
2383 .get(term_idx)
2384 .map(|term| &term.basis)
2385 {
2386 Some(SmoothBasisSpec::ConstantCurvature {
2387 feature_cols, spec, ..
2388 }) => (feature_cols, spec.clone()),
2389 _ => {
2390 crate::bail_invalid_estim!(
2391 "constant-curvature optimum requested for non-curvature term {term_idx}"
2392 )
2393 }
2394 };
2395 let x_term = select_columns(data, feature_cols).map_err(EstimationError::from)?;
2396 let y_ref = constant_curvature_radial_reference(x_term.view(), y)?;
2397 let profile = ConstantCurvatureFairProfile {
2398 data: x_term.view(),
2399 response: y,
2400 radial_reference: y_ref,
2401 spec: base_spec,
2402 cache: std::cell::RefCell::new(std::collections::HashMap::new()),
2403 };
2404 let mut seed_config = gam_problem::SeedConfig::default();
2405 seed_config.bounds = (kappa_min, kappa_max);
2406 seed_config.max_seeds = 1;
2407 seed_config.seed_budget = 1;
2408 seed_config.risk_profile = gam_problem::SeedRiskProfile::Gaussian;
2409 seed_config.num_auxiliary_trailing = 1;
2410 seed_config.over_smoothing_probe_rho = None;
2411 let initial_kappa = profile.spec.kappa.clamp(kappa_min, kappa_max);
2412 let problem = gam_solve::rho_optimizer::OuterProblem::new(1)
2413 .with_gradient(gam_problem::Derivative::Analytic)
2414 .with_hessian(gam_problem::DeclaredHessianForm::Unavailable)
2415 .with_prefer_gradient_only(true)
2416 .with_disable_fixed_point(true)
2417 .with_fallback_policy(gam_solve::rho_optimizer::FallbackPolicy::Disabled)
2418 .with_psi_dim(1)
2419 .with_tolerance(options.tol.max(f64::EPSILON.sqrt()))
2420 .with_max_iter(options.max_iter.max(1))
2421 .with_bounds(
2422 Array1::from_vec(vec![kappa_min]),
2423 Array1::from_vec(vec![kappa_max]),
2424 )
2425 .with_initial_rho(Array1::from_vec(vec![initial_kappa]))
2426 .with_seed_config(seed_config);
2427 let mut objective = problem.build_objective(
2428 profile,
2429 |profile: &mut ConstantCurvatureFairProfile<'_>, theta: &Array1<f64>| {
2430 profile.evaluate(theta[0]).map(|(value, _)| value)
2431 },
2432 |profile: &mut ConstantCurvatureFairProfile<'_>, theta: &Array1<f64>| {
2433 let (cost, derivative) = profile.evaluate(theta[0])?;
2434 Ok(gam_problem::OuterEval {
2435 cost,
2436 gradient: Array1::from_vec(vec![derivative]),
2437 hessian: gam_problem::HessianValue::Unavailable,
2438 inner_beta_hint: None,
2439 })
2440 },
2441 None::<fn(&mut ConstantCurvatureFairProfile<'_>)>,
2442 None::<
2443 fn(
2444 &mut ConstantCurvatureFairProfile<'_>,
2445 &Array1<f64>,
2446 ) -> Result<gam_problem::EfsEval, EstimationError>,
2447 >,
2448 );
2449 let result = problem.run(
2450 &mut objective,
2451 &format!("constant-curvature fair profile term {term_idx}"),
2452 )?;
2453 if !result.converged {
2454 crate::bail_invalid_estim!(
2455 "constant-curvature fair-profile κ optimization did not converge for term {} after {} iterations (negative_log_evidence={:.6e}, final_grad_norm={})",
2456 term_idx,
2457 result.iterations,
2458 result.final_value,
2459 result.final_grad_norm_report(),
2460 );
2461 }
2462 let kappa_hat = result.rho[0];
2463 log::info!(
2464 "[spatial-kappa] continuous fair-profile optimum kappa_hat={:.6} \
2465 (negative_log_evidence={:.6e}, projected_gradient={}) for term {term_idx}",
2466 kappa_hat,
2467 result.final_value,
2468 result.final_grad_norm_report(),
2469 );
2470 Ok(kappa_hat)
2471}
2472
2473fn try_exact_joint_spatial_length_scale_optimization(
2474 data: ArrayView2<'_, f64>,
2475 y: ArrayView1<'_, f64>,
2476 weights: ArrayView1<'_, f64>,
2477 offset: ArrayView1<'_, f64>,
2478 resolvedspec: &TermCollectionSpec,
2479 best: &FittedTermCollection,
2480 family: LikelihoodSpec,
2481 options: &FitOptions,
2482 kappa_options: &SpatialLengthScaleOptimizationOptions,
2483 spatial_terms: &[usize],
2484) -> Result<Option<FittedTermCollectionWithSpec>, EstimationError> {
2485 if spatial_terms.is_empty() {
2486 return Ok(None);
2487 }
2488 kappa_options
2493 .validate()
2494 .map_err(EstimationError::InvalidInput)?;
2495
2496 if try_build_spatial_log_kappa_hyper_dirs(data, resolvedspec, &best.design, spatial_terms)?
2497 .is_none()
2498 {
2499 if !constant_curvature_term_indices(resolvedspec).is_empty() {
2500 log::info!(
2501 "[#1464-trace] try_exact_joint RETURNED None (hyper_dirs unavailable); \
2502 κ̂ comes from a NON-joint path"
2503 );
2504 }
2505 return Ok(None);
2506 }
2507 if !constant_curvature_term_indices(resolvedspec).is_empty() {
2508 log::info!(
2509 "[#1464-trace] try_exact_joint ENTERED for {} spatial term(s); CC present",
2510 spatial_terms.len()
2511 );
2512 }
2513
2514 const JOINT_RHO_BOUND: f64 = 12.0;
2515 let rho_dim = best.fit.lambdas.len();
2516
2517 let has_constant_curvature_term = !constant_curvature_term_indices(resolvedspec).is_empty();
2531 let rho_upper_bound = if has_constant_curvature_term {
2532 gam_solve::estimate::RHO_BOUND
2533 } else {
2534 JOINT_RHO_BOUND
2535 };
2536
2537 let dims_per_term = spatial_dims_per_term(resolvedspec, spatial_terms);
2539 let use_aniso = has_aniso_terms(resolvedspec, spatial_terms);
2540
2541 let log_kappa0 = if use_aniso {
2546 SpatialLogKappaCoords::from_length_scales_aniso(resolvedspec, spatial_terms, kappa_options)
2547 } else {
2548 SpatialLogKappaCoords::from_length_scales(resolvedspec, spatial_terms, kappa_options)
2549 };
2550 let mut log_kappa0 = log_kappa0
2553 .reseed_from_data(data, resolvedspec, spatial_terms, kappa_options)
2554 .map_err(EstimationError::BasisError)?;
2555 let mut cc_profiled_values: Vec<(usize, f64)> = Vec::new();
2560 if has_constant_curvature_term {
2561 for (slot, &term_idx) in spatial_terms.iter().enumerate() {
2562 if constant_curvature_term_spec(resolvedspec, term_idx).is_none() {
2563 continue;
2564 }
2565 let kappa = get_constant_curvature_kappa(resolvedspec, term_idx)
2566 .expect("constant-curvature term exposes its kappa");
2567 log_kappa0.set_scalar_slot(slot, kappa);
2568 cc_profiled_values.push((slot, kappa));
2569 }
2570 }
2571 let log_kappa_lower = if use_aniso {
2572 SpatialLogKappaCoords::lower_bounds_aniso_from_data(
2573 data,
2574 resolvedspec,
2575 spatial_terms,
2576 &dims_per_term,
2577 kappa_options,
2578 )
2579 } else {
2580 SpatialLogKappaCoords::lower_bounds_from_data(
2581 data,
2582 resolvedspec,
2583 spatial_terms,
2584 kappa_options,
2585 )
2586 }
2587 .map_err(EstimationError::BasisError)?;
2588 let log_kappa_upper = if use_aniso {
2589 SpatialLogKappaCoords::upper_bounds_aniso_from_data(
2590 data,
2591 resolvedspec,
2592 spatial_terms,
2593 &dims_per_term,
2594 kappa_options,
2595 )
2596 } else {
2597 SpatialLogKappaCoords::upper_bounds_from_data(
2598 data,
2599 resolvedspec,
2600 spatial_terms,
2601 kappa_options,
2602 )
2603 }
2604 .map_err(EstimationError::BasisError)?;
2605 let mut log_kappa_lower = log_kappa_lower;
2606 let mut log_kappa_upper = log_kappa_upper;
2607 for &(slot, kappa) in &cc_profiled_values {
2608 log_kappa_lower.set_scalar_slot(slot, kappa);
2609 log_kappa_upper.set_scalar_slot(slot, kappa);
2610 log::info!("[spatial-kappa] slot {slot}: profiling rho at certified kappa={kappa}");
2611 }
2612 let log_kappa0 = log_kappa0.clamp_to_bounds(&log_kappa_lower, &log_kappa_upper);
2615 let setup = ExactJointHyperSetup::new(
2616 best.fit.lambdas.mapv(f64::ln),
2617 Array1::<f64>::from_elem(rho_dim, -JOINT_RHO_BOUND),
2618 Array1::<f64>::from_elem(rho_dim, rho_upper_bound),
2619 log_kappa0,
2620 log_kappa_lower,
2621 log_kappa_upper,
2622 );
2623
2624 let theta0 = setup.theta0();
2625 let lower = setup.lower();
2626 let upper = setup.upper();
2627
2628 let kind = if use_aniso {
2640 SpatialHyperKind::Anisotropic
2641 } else {
2642 SpatialHyperKind::Isotropic
2643 };
2644 let (theta_star, joint_final_value, kappa_timing) = run_exact_joint_spatial_optimization(
2645 kind,
2646 data,
2647 y,
2648 weights,
2649 offset,
2650 resolvedspec,
2651 &best.design,
2652 family.clone(),
2653 options,
2654 spatial_terms,
2655 &dims_per_term,
2656 &theta0,
2657 &lower,
2658 &upper,
2659 rho_dim,
2660 kappa_options,
2661 )?;
2662
2663 let baseline_score = fit_score(&best.fit);
2664
2665 let accept_tol = options.tol.max(1e-8 * baseline_score.abs()).max(1e-12);
2670 if joint_final_value > baseline_score + accept_tol {
2671 return Err(EstimationError::RemlOptimizationFailed(format!(
2672 "exact joint spatial optimization failed its objective-monotonicity certificate: \
2673 initial={baseline_score:.6e}, final={joint_final_value:.6e}, \
2674 acceptance_tolerance={accept_tol:.3e}, theta_checkpoint={:?}",
2675 theta_star.to_vec(),
2676 )));
2677 }
2678
2679 let selected_lambdas = Array1::from_vec(
2680 gam_problem::checked_exp_log_strengths(
2681 theta_star.slice(s![..rho_dim]).iter().copied(),
2682 )
2683 .map_err(|error| {
2684 EstimationError::InvalidInput(format!(
2685 "selected joint spatial smoothing coordinate is outside the canonical log-strength domain: {error}"
2686 ))
2687 })?,
2688 );
2689 let log_kappa_star =
2690 SpatialLogKappaCoords::from_theta_tail_with_dims(&theta_star, rho_dim, dims_per_term);
2691 if has_constant_curvature_term {
2697 let star = log_kappa_star.as_array();
2698 let dims = log_kappa_star.dims_per_term();
2699 for (slot, &term_idx) in spatial_terms.iter().enumerate() {
2700 if constant_curvature_term_spec(resolvedspec, term_idx).is_some() {
2701 let off: usize = dims[..slot].iter().sum();
2702 log::info!(
2703 "[#1464-trace] term {term_idx}: joint solver CONVERGED ψ-tail κ = {} \
2704 (this is the optimised candidate; joint_final_value={joint_final_value})",
2705 star[off]
2706 );
2707 }
2708 }
2709 }
2710 let optimized_spec = log_kappa_star.apply_tospec(resolvedspec, spatial_terms)?;
2711 let optimized = fit_term_collection_forspecwith_heuristic_lambdas(
2712 data,
2713 y,
2714 weights,
2715 offset,
2716 &optimized_spec,
2717 selected_lambdas.as_slice(),
2718 family.clone(),
2719 options,
2720 )?;
2721
2722 let mut fit = optimized.fit;
2726 fit.reml_score = joint_final_value;
2727 let optimized_result = FittedTermCollectionWithSpec {
2728 fit,
2729 design: optimized.design,
2730 resolvedspec: optimized_spec,
2731 adaptive_diagnostics: optimized.adaptive_diagnostics,
2732 kappa_timing: Some(kappa_timing),
2733 };
2734
2735 Ok(Some(optimized_result))
2736}
2737
2738#[derive(Clone, Copy, PartialEq, Eq, Debug)]
2750enum SpatialHyperKind {
2751 Anisotropic,
2752 Isotropic,
2753}
2754
2755impl SpatialHyperKind {
2756 fn label(self) -> &'static str {
2759 match self {
2760 SpatialHyperKind::Anisotropic => "spatial-aniso-joint",
2761 SpatialHyperKind::Isotropic => "spatial-iso-joint",
2762 }
2763 }
2764
2765 fn adjective(self) -> &'static str {
2767 match self {
2768 SpatialHyperKind::Anisotropic => "anisotropic",
2769 SpatialHyperKind::Isotropic => "isotropic",
2770 }
2771 }
2772
2773 fn coord_name(self) -> &'static str {
2776 match self {
2777 SpatialHyperKind::Anisotropic => "psi",
2778 SpatialHyperKind::Isotropic => "kappa",
2779 }
2780 }
2781}
2782
2783struct SpatialFrozenGlmInputs {
2789 y: Array1<f64>,
2790 weights: Array1<f64>,
2791 offset: Array1<f64>,
2792 family: LikelihoodSpec,
2793}
2794
2795fn frozen_glm_tensor_eligible_family(family: &LikelihoodSpec) -> bool {
2812 !family.is_gaussian_identity()
2813 && matches!(
2814 &family.response,
2815 ResponseFamily::Binomial
2816 | ResponseFamily::Poisson
2817 | ResponseFamily::Gamma
2818 | ResponseFamily::NegativeBinomial { .. }
2819 )
2820}
2821
2822struct SpatialJointContext<'d> {
2823 data: ArrayView2<'d, f64>,
2824 rho_dim: usize,
2825 kind: SpatialHyperKind,
2826 cache: SingleBlockExactJointDesignCache<'d>,
2827 evaluator: gam_solve::estimate::ExternalJointHyperEvaluator<'d>,
2828 frozen_glm_inputs: Option<SpatialFrozenGlmInputs>,
2829 frozen_glm_psi_bounds: Option<(f64, f64)>,
2830 frozen_glm_tensor: Option<gam_solve::glm_sufficient_lane::FrozenWeightGramTensor>,
2831 frozen_glm_tensor_attempted: bool,
2832 frozen_glm_weight_memo: Option<(Array1<f64>, Array1<f64>)>,
2844}
2845
2846#[derive(Clone, Copy, Debug, Default)]
2847struct NfreeSkipGateStatus {
2848 shape: bool,
2849 value: bool,
2850 gradient: bool,
2851 penalty: bool,
2852 revision: bool,
2853 second_order: bool,
2854}
2855
2856impl NfreeSkipGateStatus {
2857 fn would_skip(self, require_gradient: bool) -> bool {
2858 self.shape
2859 && self.value
2860 && (!require_gradient || self.gradient)
2861 && self.penalty
2862 && self.revision
2863 && !self.second_order
2864 }
2865}
2866
2867fn nfree_skip_gate_status_from_parts(
2868 shape: bool,
2869 covers_value: bool,
2870 covers_skip: bool,
2871 covers_gradient: bool,
2872 penalty: bool,
2873 revision: bool,
2874 allow_second_order: bool,
2875 require_gradient: bool,
2876) -> NfreeSkipGateStatus {
2877 NfreeSkipGateStatus {
2878 shape,
2879 value: shape && covers_value && (!require_gradient || covers_skip),
2887 gradient: shape && (!require_gradient || covers_gradient),
2888 penalty,
2889 revision,
2890 second_order: allow_second_order,
2891 }
2892}
2893
2894impl<'d> SpatialJointContext<'d> {
2895 fn nfree_skip_gate_status(
2896 &self,
2897 theta: &Array1<f64>,
2898 allow_second_order: bool,
2899 require_gradient: bool,
2900 ) -> NfreeSkipGateStatus {
2901 let shape = theta.len() == self.rho_dim + 1;
2902 let (covers_value, covers_skip, covers_gradient) = if shape {
2903 let psi = theta[self.rho_dim];
2904 (
2905 self.evaluator.psi_gram_tensor_covers(psi),
2906 self.evaluator.psi_gram_tensor_covers_skip(psi),
2907 self.evaluator.psi_gram_tensor_covers_gradient(psi),
2908 )
2909 } else {
2910 (false, false, false)
2911 };
2912 nfree_skip_gate_status_from_parts(
2913 shape,
2914 covers_value,
2915 covers_skip,
2916 covers_gradient,
2917 self.evaluator.supports_nfree_penalty_rekey(),
2918 self.evaluator.nfree_fast_path_revision().is_some(),
2919 allow_second_order,
2920 require_gradient,
2921 )
2922 }
2923
2924 fn frozen_glm_working_state(
2925 &self,
2926 beta: &Array1<f64>,
2927 ) -> Result<Option<(Array1<f64>, Array1<f64>)>, EstimationError> {
2928 let Some(inputs) = self.frozen_glm_inputs.as_ref() else {
2929 return Ok(None);
2930 };
2931 if beta.len() != self.cache.design().design.ncols() {
2932 return Ok(None);
2933 }
2934 let mut eta = self.cache.design().design.matrixvectormultiply(beta);
2935 if eta.len() != inputs.offset.len() {
2936 crate::bail_invalid_estim!(
2937 "frozen GLM tensor warm-state row mismatch: eta={}, offset={}",
2938 eta.len(),
2939 inputs.offset.len()
2940 );
2941 }
2942 eta += &inputs.offset;
2943 let obs = evaluate_standard_familyobservations(
2944 inputs.family.clone(),
2945 None,
2946 None,
2947 None,
2948 &inputs.y,
2949 &inputs.weights,
2950 &eta,
2951 )?;
2952 let mut working_response = obs.eta.clone();
2953 for i in 0..working_response.len() {
2954 let wi = obs.fisherweight[i].max(1e-12);
2955 working_response[i] += obs.score[i] / wi;
2956 }
2957 Ok(Some((obs.fisherweight, working_response)))
2958 }
2959
2960 fn frozen_glm_trial_weights(
2969 &mut self,
2970 beta: &Array1<f64>,
2971 ) -> Result<Option<Array1<f64>>, EstimationError> {
2972 if let Some((memo_beta, memo_w)) = self.frozen_glm_weight_memo.as_ref()
2973 && memo_beta.len() == beta.len()
2974 && memo_beta
2975 .iter()
2976 .zip(beta.iter())
2977 .all(|(a, b)| a.to_bits() == b.to_bits())
2978 {
2979 return Ok(Some(memo_w.clone()));
2980 }
2981 match self.frozen_glm_working_state(beta)? {
2982 Some((current_w, _)) => {
2983 self.frozen_glm_weight_memo = Some((beta.clone(), current_w.clone()));
2984 Ok(Some(current_w))
2985 }
2986 None => Ok(None),
2987 }
2988 }
2989
2990 fn ensure_frozen_glm_tensor(
2991 &mut self,
2992 theta: &Array1<f64>,
2993 warm_beta: Option<&Array1<f64>>,
2994 ) -> Result<(), EstimationError> {
2995 if self.frozen_glm_tensor.is_some() || self.frozen_glm_tensor_attempted {
2996 return Ok(());
2997 }
2998 let Some((psi_lo, psi_hi)) = self.frozen_glm_psi_bounds else {
2999 return Ok(());
3000 };
3001 if theta.len() != self.rho_dim + 1 {
3002 self.frozen_glm_tensor_attempted = true;
3003 return Ok(());
3004 }
3005 let Some(beta) = warm_beta else {
3006 return Ok(());
3007 };
3008 let Some((frozen_w, working_z)) = self.frozen_glm_working_state(beta)? else {
3009 self.frozen_glm_tensor_attempted = true;
3010 return Ok(());
3011 };
3012 let theta_probe_base = theta.clone();
3013 let rho_dim = self.rho_dim;
3014 let Self {
3021 cache, evaluator, ..
3022 } = self;
3023 let tensor = evaluator.build_frozen_glm_gram_tensor(
3024 |psi| {
3025 let mut theta_probe = theta_probe_base.clone();
3026 theta_probe[rho_dim] = psi;
3027 cache.ensure_theta(&theta_probe)?;
3028 Ok(cache.design().design.clone())
3029 },
3030 frozen_w.view(),
3031 working_z.view(),
3032 psi_lo,
3033 psi_hi,
3034 );
3035 self.cache
3036 .ensure_theta(theta)
3037 .map_err(EstimationError::InvalidInput)?;
3038 self.frozen_glm_tensor_attempted = true;
3039 if let Some(tensor) = tensor {
3040 self.frozen_glm_tensor = Some(tensor);
3041 log::info!(
3042 "[STAGE] {} certified frozen-W GLM ψ tensor over [{psi_lo:.3}, {psi_hi:.3}]",
3043 self.kind.label(),
3044 );
3045 } else {
3046 log::info!(
3047 "[STAGE] {} frozen-W GLM ψ tensor did not certify over [{psi_lo:.3}, {psi_hi:.3}]",
3048 self.kind.label(),
3049 );
3050 }
3051 Ok(())
3052 }
3053
3054 fn stage_frozen_glm_trial_statistics(
3055 &mut self,
3056 theta: &Array1<f64>,
3057 warm_beta: Option<&Array1<f64>>,
3058 allow_gradient: bool,
3059 ) -> Result<(), EstimationError> {
3060 let kind = self.kind;
3061 let mut staged_gram: Option<Array2<f64>> = None;
3062 let mut staged_deriv: Option<(Array2<f64>, Array1<f64>)> = None;
3063 if theta.len() == self.rho_dim + 1 {
3064 let psi = theta[self.rho_dim];
3065 let tensor_covers = self
3072 .frozen_glm_tensor
3073 .as_ref()
3074 .is_some_and(|t| t.contains(psi));
3075 let current_w = if tensor_covers {
3076 match warm_beta {
3077 Some(beta) => self.frozen_glm_trial_weights(beta)?,
3078 None => None,
3079 }
3080 } else {
3081 None
3082 };
3083 if let (Some(tensor), Some(current_w)) =
3084 (self.frozen_glm_tensor.as_ref(), current_w.as_ref())
3085 {
3086 const FROZEN_GLM_WEIGHT_DRIFT_RTOL: f64 = 1e-3;
3087 if tensor.weight_drift_within(current_w.view(), FROZEN_GLM_WEIGHT_DRIFT_RTOL) {
3088 staged_gram = Some(tensor.gram_at(psi));
3089 log::debug!(
3090 "[STAGE] {} trial at psi={psi:.6}: serving frozen-W GLM \
3091 first-Fisher-step XᵀWX n-free (weight drift within tol)",
3092 kind.label(),
3093 );
3094 }
3095 if allow_gradient
3096 && tensor.contains_for_gradient(psi)
3097 && let Some((dgram_dpsi, drhs_dpsi)) =
3098 tensor.gradient_pair_if_sound(psi, current_w.view())
3099 {
3100 staged_deriv = Some((dgram_dpsi, drhs_dpsi));
3101 log::debug!(
3102 "[STAGE] {} trial at psi={psi:.6}: serving frozen-W GLM \
3103 ψ-gradient (∂G/∂ψ, ∂b/∂ψ) n-free (gradient weight drift within \
3104 tight tol); B_j stays exact",
3105 kind.label(),
3106 );
3107 }
3108 }
3109 }
3110 self.evaluator.stage_glm_first_step_gram(staged_gram);
3111 self.evaluator.stage_glm_psi_gram_deriv(staged_deriv);
3112 Ok(())
3113 }
3114
3115 fn eval_full(
3117 &mut self,
3118 theta: &Array1<f64>,
3119 order: gam_solve::rho_optimizer::OuterEvalOrder,
3120 analytic_outer_hessian_available: bool,
3121 ) -> Result<(f64, Array1<f64>, gam_problem::HessianValue), EstimationError> {
3122 use gam_solve::rho_optimizer::OuterEvalOrder;
3123 let allow_second_order = matches!(order, OuterEvalOrder::ValueGradientHessian)
3124 && analytic_outer_hessian_available;
3125 if let Some(eval) = self.cache.memoized_eval(theta) {
3126 let cached_satisfies_order = !allow_second_order || eval.2.is_analytic();
3127 if cached_satisfies_order {
3128 return Ok(eval);
3129 }
3130 }
3131 let kind = self.kind;
3132 let nfree_fast_path_revision = self.evaluator.nfree_fast_path_revision();
3168 let skip_design_realization = !allow_second_order && theta.len() == self.rho_dim + 1 && {
3169 let psi = theta[self.rho_dim];
3170 self.evaluator.psi_gram_tensor_covers(psi)
3171 && self.evaluator.psi_gram_tensor_covers_gradient(psi)
3178 && self.evaluator.psi_gram_tensor_covers_skip(psi)
3195 && self.evaluator.supports_nfree_penalty_rekey()
3200 && nfree_fast_path_revision.is_some()
3201 };
3202 if skip_design_realization {
3214 log::debug!(
3215 "[STAGE] {} eval_full at psi={:.6}: skipping n×k design re-realization \
3216 + reconditioning — criterion/gradient/inner-solve served n-free from \
3217 the certified ψ-gram tensor (GaussianFixedCache + k-space ψ-derivatives)",
3218 kind.label(),
3219 theta[self.rho_dim],
3220 );
3221 } else {
3222 self.cache
3223 .ensure_theta(theta)
3224 .map_err(EstimationError::InvalidInput)?;
3225 }
3226 let warm_beta = self.evaluator.current_beta();
3227 self.ensure_frozen_glm_tensor(theta, warm_beta.as_ref())?;
3228 self.stage_frozen_glm_trial_statistics(theta, warm_beta.as_ref(), !allow_second_order)?;
3236 let hyper_dirs = if skip_design_realization {
3243 self.cache.nfree_tensor_gradient_hyper_dirs(theta)?
3244 } else {
3245 self.cache.hyper_dirs_for_current_design(self.data, kind)?
3246 };
3247
3248 let design_revision = if skip_design_realization {
3249 nfree_fast_path_revision
3250 } else {
3251 Some(self.cache.design_revision())
3252 };
3253 if self.evaluator.supports_nfree_penalty_rekey() {
3267 match self.cache.canonical_penalties_at(theta) {
3268 Ok(penalty) => self.evaluator.stage_fast_path_penalty(Some(penalty)),
3269 Err(e) => {
3270 log::warn!(
3271 "[STAGE] {} eval_full at psi={:.6}: exact n-free S(ψ) rebuild failed \
3272 ({e}); clearing stage (eval falls to slow path)",
3273 kind.label(),
3274 theta[self.rho_dim],
3275 );
3276 self.evaluator.stage_fast_path_penalty(None);
3277 }
3278 }
3279 }
3280 let eval = evaluate_joint_reml_outer_eval_at_theta(
3287 &mut self.evaluator,
3288 self.cache.design(),
3289 theta,
3290 self.rho_dim,
3291 hyper_dirs,
3292 warm_beta.as_ref().map(|b: &Array1<f64>| b.view()),
3293 if allow_second_order {
3294 order
3295 } else {
3296 OuterEvalOrder::ValueAndGradient
3297 },
3298 design_revision,
3299 );
3300 if let Ok(ref value) = eval {
3301 self.cache.store_eval_at(theta, value.clone());
3302 }
3303 eval
3304 }
3305
3306 fn eval_efs(&mut self, theta: &Array1<f64>) -> Result<gam_problem::EfsEval, EstimationError> {
3307 self.cache
3308 .ensure_theta(theta)
3309 .map_err(EstimationError::InvalidInput)?;
3310 let kind = self.kind;
3311 let hyper_dirs = try_build_spatial_log_kappa_hyper_dirs(
3312 self.data,
3313 self.cache.spec(),
3314 self.cache.design(),
3315 &self.cache.spatial_terms,
3316 )?
3317 .ok_or_else(|| {
3318 EstimationError::InvalidInput(format!(
3319 "failed to build {} hyper_dirs for exact-joint EFS",
3320 kind.adjective(),
3321 ))
3322 })?;
3323 let design_revision = Some(self.cache.design_revision());
3324 let warm_beta = self.evaluator.current_beta();
3325 evaluate_joint_reml_efs_at_theta(
3326 &mut self.evaluator,
3327 self.cache.design(),
3328 theta,
3329 self.rho_dim,
3330 hyper_dirs,
3331 warm_beta.as_ref().map(|b: &Array1<f64>| b.view()),
3332 design_revision,
3333 )
3334 }
3335
3336 fn eval_cost(&mut self, theta: &Array1<f64>) -> f64 {
3342 if let Some(cost) = self.cache.memoized_cost(theta) {
3343 return cost;
3344 }
3345 let probe_start = std::time::Instant::now();
3360 let psi_distance = self
3361 .cache
3362 .current_theta
3363 .as_ref()
3364 .filter(|reference| reference.len() == theta.len())
3365 .map(|reference| {
3366 reference
3367 .iter()
3368 .zip(theta.iter())
3369 .map(|(a, b)| (a - b) * (a - b))
3370 .sum::<f64>()
3371 .sqrt()
3372 })
3373 .unwrap_or(f64::NAN);
3374 let nfree_fast_path_revision = self.evaluator.nfree_fast_path_revision();
3388 let skip_value_realization = theta.len() == self.rho_dim + 1 && {
3389 let psi = theta[self.rho_dim];
3390 self.evaluator.psi_gram_tensor_covers(psi)
3391 && self.evaluator.supports_nfree_penalty_rekey()
3425 && nfree_fast_path_revision.is_some()
3426 };
3427 if theta.len() == self.rho_dim + 1
3428 && self.evaluator.has_psi_gram_tensor()
3429 && !self.evaluator.psi_gram_tensor_covers(theta[self.rho_dim])
3430 {
3431 self.cache.store_cost_at(theta, f64::INFINITY);
3432 return f64::INFINITY;
3433 }
3434 if !skip_value_realization && self.cache.ensure_theta(theta).is_err() {
3435 return f64::INFINITY;
3436 }
3437 if self.evaluator.supports_nfree_penalty_rekey() {
3443 match self.cache.canonical_penalties_at(theta) {
3444 Ok(penalty) => self.evaluator.stage_fast_path_penalty(Some(penalty)),
3445 Err(_) => self.evaluator.stage_fast_path_penalty(None),
3446 }
3447 }
3448 let warm_beta = self.evaluator.current_beta();
3449 if let Err(err) = self.ensure_frozen_glm_tensor(theta, warm_beta.as_ref()) {
3450 log::warn!(
3451 "[STAGE] {} value-probe at psi={:.6}: frozen-W GLM tensor setup failed ({err}); \
3452 falling back to exact streamed Gram",
3453 self.kind.label(),
3454 if theta.len() > self.rho_dim {
3455 theta[self.rho_dim]
3456 } else {
3457 f64::NAN
3458 },
3459 );
3460 self.evaluator.stage_glm_first_step_gram(None);
3461 self.evaluator.stage_glm_psi_gram_deriv(None);
3462 } else if let Err(err) =
3463 self.stage_frozen_glm_trial_statistics(theta, warm_beta.as_ref(), false)
3464 {
3465 log::warn!(
3466 "[STAGE] {} value-probe at psi={:.6}: frozen-W GLM staging failed ({err}); \
3467 falling back to exact streamed Gram",
3468 self.kind.label(),
3469 if theta.len() > self.rho_dim {
3470 theta[self.rho_dim]
3471 } else {
3472 f64::NAN
3473 },
3474 );
3475 self.evaluator.stage_glm_first_step_gram(None);
3476 self.evaluator.stage_glm_psi_gram_deriv(None);
3477 }
3478 let design_revision = if skip_value_realization {
3479 nfree_fast_path_revision
3480 } else {
3481 Some(self.cache.design_revision())
3482 };
3483 let cost_label = self.kind.label();
3484 let result = {
3485 let design = self.cache.design();
3486 self.evaluator.evaluate_cost_only(
3487 &design.design,
3488 &design.penalties,
3489 &design.nullspace_dims,
3490 design.linear_constraints.clone(),
3491 theta,
3492 self.rho_dim,
3493 warm_beta.as_ref().map(|b: &Array1<f64>| b.view()),
3494 cost_label,
3495 design_revision,
3496 )
3497 };
3498 match result {
3499 Ok(cost) => {
3500 log::debug!(
3501 "[STAGE] {cost_label} value-probe (order=Value): elapsed={:.3}s \
3502 cost={cost:.6e} trial_theta_distance={psi_distance:.3e}",
3503 probe_start.elapsed().as_secs_f64(),
3504 );
3505 self.cache.store_cost_at(theta, cost);
3506 cost
3507 }
3508 Err(_) => f64::INFINITY,
3509 }
3510 }
3511
3512 fn reset(&mut self) {
3513 self.cache.current_theta = None;
3514 self.cache.last_eval_theta = None;
3515 self.cache.last_cost = None;
3516 self.cache.last_eval = None;
3517 }
3518}
3519
3520fn kphase_log_norms(theta: &Array1<f64>, rho_dim: usize) -> (f64, f64) {
3544 let theta_norm = theta.iter().map(|v| v * v).sum::<f64>().sqrt();
3545 let log_kappa_norm = theta
3546 .iter()
3547 .skip(rho_dim)
3548 .map(|v| v * v)
3549 .sum::<f64>()
3550 .sqrt();
3551 (theta_norm, log_kappa_norm)
3552}
3553
3554fn run_exact_joint_spatial_optimization(
3555 kind: SpatialHyperKind,
3556 data: ArrayView2<'_, f64>,
3557 y: ArrayView1<'_, f64>,
3558 weights: ArrayView1<'_, f64>,
3559 offset: ArrayView1<'_, f64>,
3560 resolvedspec: &TermCollectionSpec,
3561 baseline_design: &TermCollectionDesign,
3562 family: LikelihoodSpec,
3563 options: &FitOptions,
3564 spatial_terms: &[usize],
3565 dims_per_term: &[usize],
3566 theta0: &Array1<f64>,
3567 lower: &Array1<f64>,
3568 upper: &Array1<f64>,
3569 rho_dim: usize,
3570 kappa_options: &SpatialLengthScaleOptimizationOptions,
3571) -> Result<(Array1<f64>, f64, SpatialLengthScaleOptimizationTiming), EstimationError> {
3572 let label = kind.label();
3573 let effective_offset = baseline_design
3574 .compose_offset(offset, "spatial joint fit")
3575 .map_err(EstimationError::BasisError)?;
3576 let offset = effective_offset.view();
3577 assert!(
3579 lower.len() == theta0.len() && upper.len() == theta0.len(),
3580 "spatial hyperparameter bounds must match theta length: lower_len={}, upper_len={}, theta_len={}",
3581 lower.len(),
3582 upper.len(),
3583 theta0.len()
3584 );
3585 assert!(
3586 baseline_design.smooth.terms.len() >= spatial_terms.len(),
3587 "baseline design must have at least one smooth term per spatial term: baseline_terms={}, spatial_terms={}",
3588 baseline_design.smooth.terms.len(),
3589 spatial_terms.len()
3590 );
3591 use gam_problem::{DeclaredHessianForm, Derivative, OuterEval};
3592 use gam_solve::rho_optimizer::OuterEvalOrder;
3593
3594 let theta_dim = theta0.len();
3595 let coord_dim = theta_dim - rho_dim;
3598 let analytic_outer_hessian_available =
3608 exact_joint_spatial_outer_hessian_available(&family, baseline_design);
3609 if !analytic_outer_hessian_available {
3610 log::info!(
3611 "[{label}] analytic outer Hessian unavailable for family/design; routing without second-order geometry (coord_dim={coord_dim})"
3612 );
3613 }
3614 let mut prefer_gradient_only = theta_dim > EXACT_JOINT_SECOND_ORDER_THETA_CAP;
3620 if prefer_gradient_only {
3621 log::info!(
3622 "[{label}] joint θ-dim {theta_dim} exceeds the exact pair-Hessian budget \
3623 ({EXACT_JOINT_SECOND_ORDER_THETA_CAP}); routing gradient-only quasi-Newton"
3624 );
3625 }
3626 let mut suppress_outer_hessian_for_nfree = false;
3636
3637 log::trace!(
3638 "[{}] starting analytic optimization: rho_dim={}, coord_dim={}, dims_per_term={:?}",
3639 label,
3640 rho_dim,
3641 coord_dim,
3642 dims_per_term,
3643 );
3644
3645 let mut ctx = SpatialJointContext {
3646 data,
3647 rho_dim,
3648 kind,
3649 cache: SingleBlockExactJointDesignCache::new_with_policy(
3650 data,
3651 resolvedspec.clone(),
3652 baseline_design.clone(),
3653 spatial_terms.to_vec(),
3654 rho_dim,
3655 dims_per_term.to_vec(),
3656 &options.resource_policy,
3657 )
3658 .map_err(EstimationError::InvalidInput)?,
3659 evaluator: gam_solve::estimate::ExternalJointHyperEvaluator::new(
3660 y,
3661 weights,
3662 &baseline_design.design,
3663 offset,
3664 &baseline_design.penalties,
3665 &external_opts_for_design(&family, baseline_design, options),
3666 label,
3667 )?,
3668 frozen_glm_inputs: if coord_dim == 1 && frozen_glm_tensor_eligible_family(&family) {
3669 Some(SpatialFrozenGlmInputs {
3670 y: y.to_owned(),
3671 weights: weights.to_owned(),
3672 offset: offset.to_owned(),
3673 family: family.clone(),
3674 })
3675 } else {
3676 None
3677 },
3678 frozen_glm_psi_bounds: if coord_dim == 1 && frozen_glm_tensor_eligible_family(&family) {
3679 Some((lower[rho_dim], upper[rho_dim]))
3680 } else {
3681 None
3682 },
3683 frozen_glm_tensor: None,
3684 frozen_glm_tensor_attempted: false,
3685 frozen_glm_weight_memo: None,
3686 };
3687
3688 let mut psi_rank_stable_floor: Option<f64> = None;
3711 let mut psi_rank_stable_ceiling: Option<f64> = None;
3720 let nfree_penalty_capable =
3721 coord_dim == 1 && family.is_gaussian_identity() && ctx.cache.supports_nfree_penalty_rekey();
3722 if nfree_penalty_capable {
3723 let psi_lo = lower[rho_dim];
3724 let psi_hi = upper[rho_dim];
3725 let z = Array1::from_iter(y.iter().zip(offset.iter()).map(|(yi, oi)| yi - oi));
3726 let theta_probe_base = theta0.clone();
3727 let SpatialJointContext {
3730 cache, evaluator, ..
3731 } = &mut ctx;
3732 let attached = evaluator.build_and_set_psi_gram_tensor(
3733 |psi| {
3734 let mut theta_probe = theta_probe_base.clone();
3735 theta_probe[rho_dim] = psi;
3736 cache.ensure_theta(&theta_probe)?;
3737 Ok(cache.design().design.clone())
3738 },
3739 weights,
3740 z.view(),
3741 psi_lo,
3742 psi_hi,
3743 );
3744 if attached {
3745 log::info!(
3746 "[{label}] certified ψ-gram tensor over [{psi_lo:.3}, {psi_hi:.3}]: \
3747 in-window trials assemble Gaussian sufficient statistics n-free"
3748 );
3749 let psi_anchor = theta0[rho_dim];
3754 psi_rank_stable_floor = evaluator
3755 .psi_gram_rank_stable_floor(psi_anchor)
3756 .filter(|&f| f.is_finite() && f > psi_lo && f < psi_anchor);
3757 log::info!(
3758 "[KAPPA-PHASE-FLOOR] n_rows={} psi_lo={psi_lo:.6} psi_anchor={psi_anchor:.6} \
3759 rank_stable_floor={:?} lifted={}",
3760 data.nrows(),
3761 evaluator.psi_gram_rank_stable_floor(psi_anchor),
3762 psi_rank_stable_floor.is_some(),
3763 );
3764 if let Some(floor) = psi_rank_stable_floor {
3765 log::info!(
3766 "[{label}] rank-stable κ-floor ψ_floor={floor:.6} > window floor \
3767 ψ_lo={psi_lo:.6}: lifting the optimizer lower bound to keep every \
3768 in-window trial on the n-free design-realization skip (#1033). The \
3769 conditioned Gram is rank-deficient below ψ_floor (longest-length-scale \
3770 radial mode collapses into the nullspace), where the skip is soundly \
3771 refused; that band drifts with n via the sample-std standardization, \
3772 so this n-free k-space floor is the n-independent fix."
3773 );
3774 }
3775 psi_rank_stable_ceiling = evaluator
3784 .psi_gram_rank_stable_ceiling(psi_anchor)
3785 .filter(|&c| c.is_finite() && c < psi_hi && c > psi_anchor);
3786 log::info!(
3787 "[KAPPA-PHASE-CEIL] n_rows={} psi_hi={psi_hi:.6} psi_anchor={psi_anchor:.6} \
3788 rank_stable_ceiling={:?} clamped={}",
3789 data.nrows(),
3790 evaluator.psi_gram_rank_stable_ceiling(psi_anchor),
3791 psi_rank_stable_ceiling.is_some(),
3792 );
3793 if let Some(ceiling) = psi_rank_stable_ceiling {
3794 log::info!(
3795 "[{label}] rank-stable κ-ceiling ψ_ceil={ceiling:.6} < window ceiling \
3796 ψ_hi={psi_hi:.6}: clamping the optimizer upper bound to keep every \
3797 in-window trial on the n-free design-realization skip (#1033). The \
3798 conditioned Gram is rank-deficient above ψ_ceil (longest-frequency \
3799 radial mode goes collinear), where the skip is soundly refused; a \
3800 line-search overshoot there trips the O(n) reset_surface lane (and the \
3801 deficient pinning ψ it records resets the next in-band trial too)."
3802 );
3803 }
3804 let gradient_covers_full_window = evaluator.psi_gram_tensor_covers_gradient(psi_lo)
3805 && evaluator.psi_gram_tensor_covers_gradient(psi_hi);
3806 if gradient_covers_full_window {
3807 log::info!(
3808 "[{label}] certified ψ-gram tensor gradient lane covers the full \
3809 optimizer window [{psi_lo:.3}, {psi_hi:.3}]"
3810 );
3811 } else {
3812 log::info!(
3813 "[{label}] ψ-gram tensor value lane certified, but the gradient lane \
3814 does not cover the full optimizer window [{psi_lo:.3}, {psi_hi:.3}]; \
3815 keeping exact streamed kappa routing"
3816 );
3817 }
3818 evaluator.set_supports_nfree_penalty_rekey(true);
3838 log::info!(
3839 "[{label}] exact n-free ψ-penalty re-key enabled over [{psi_lo:.3}, \
3840 {psi_hi:.3}]: in-window fast-path trials rebuild S(ψ) n-free from frozen \
3841 geometry (no reset_surface)"
3842 );
3843 } else {
3844 log::info!(
3845 "[{label}] ψ-gram tensor did not certify over [{psi_lo:.3}, {psi_hi:.3}]; \
3846 keeping the exact per-trial path"
3847 );
3848 }
3849 if attached
3870 && evaluator.psi_gram_tensor_covers_gradient(psi_lo)
3871 && evaluator.psi_gram_tensor_covers_gradient(psi_hi)
3872 && evaluator.supports_nfree_penalty_rekey()
3873 && cache.supports_nfree_gradient_only_routing()
3874 {
3875 suppress_outer_hessian_for_nfree = true;
3876 prefer_gradient_only = true;
3877 log::info!(
3878 "[{label}] n-free Gaussian ψ-lane armed; suppressing the analytic outer \
3879 Hessian and routing gradient-only (BFGS) so the κ outer loop never realizes \
3880 the O(n) second-order slab — n-independent outer loop (#1033)"
3881 );
3882 }
3883 } else if coord_dim == 1 && family.is_gaussian_identity() {
3884 log::info!(
3885 "[{label}] exact n-free ψ-penalty re-key unavailable; skipping ψ-gram tensor \
3886 attachment so value, gradient, and Hessian remain on the same exact streamed \
3887 objective"
3888 );
3889 }
3890
3891 let kphase_prime_order =
3892 if analytic_outer_hessian_available && !suppress_outer_hessian_for_nfree {
3893 OuterEvalOrder::ValueGradientHessian
3894 } else {
3895 OuterEvalOrder::ValueAndGradient
3896 };
3897 let kphase_prime_start = std::time::Instant::now();
3898 drop(ctx.eval_full(theta0, kphase_prime_order, analytic_outer_hessian_available)?);
3899 log::info!(
3900 "[KAPPA-PHASE-PRIME] n_rows={} order={:?} elapsed_s={:.4} slow_path_resets_total={} design_revision={}",
3901 data.nrows(),
3902 kphase_prime_order,
3903 kphase_prime_start.elapsed().as_secs_f64(),
3904 ctx.evaluator.slow_path_reset_count(),
3905 ctx.cache.design_revision(),
3906 );
3907
3908 let kphase_cost_calls = std::cell::Cell::new(0usize);
3909 let kphase_eval_calls = std::cell::Cell::new(0usize);
3910 let kphase_efs_calls = std::cell::Cell::new(0usize);
3911 let kphase_cost_total_s = std::cell::Cell::new(0.0);
3912 let kphase_eval_total_s = std::cell::Cell::new(0.0);
3913 let kphase_efs_total_s = std::cell::Cell::new(0.0);
3914 let kphase_nfree_miss_shape = std::cell::Cell::new(0u64);
3915 let kphase_nfree_miss_value = std::cell::Cell::new(0u64);
3916 let kphase_nfree_miss_gradient = std::cell::Cell::new(0u64);
3917 let kphase_nfree_miss_penalty = std::cell::Cell::new(0u64);
3918 let kphase_nfree_miss_revision = std::cell::Cell::new(0u64);
3919 let kphase_nfree_miss_second_order = std::cell::Cell::new(0u64);
3920 let kphase_nfree_miss_other = std::cell::Cell::new(0u64);
3921 let kphase_optim_start = std::time::Instant::now();
3922 let kphase_log_kappa_dim = coord_dim;
3923 let kphase_slow_resets_start = ctx.evaluator.slow_path_reset_count();
3924 let kphase_design_revision_start = ctx.cache.design_revision();
3925 let kphase_nfree_skip_touches_start = gam_solve::pirls::nfree_skip_row_element_touches();
3929
3930 let lower_effective: std::borrow::Cow<'_, Array1<f64>> = match psi_rank_stable_floor {
3937 Some(floor) if coord_dim == 1 && floor > lower[rho_dim] => {
3938 let mut lifted = lower.clone();
3939 lifted[rho_dim] = floor;
3940 std::borrow::Cow::Owned(lifted)
3941 }
3942 _ => std::borrow::Cow::Borrowed(lower),
3943 };
3944 let lower = lower_effective.as_ref();
3945
3946 let upper_effective: std::borrow::Cow<'_, Array1<f64>> = match psi_rank_stable_ceiling {
3954 Some(ceiling) if coord_dim == 1 && ceiling < upper[rho_dim] => {
3955 let mut clamped = upper.clone();
3956 clamped[rho_dim] = ceiling;
3957 std::borrow::Cow::Owned(clamped)
3958 }
3959 _ => std::borrow::Cow::Borrowed(upper),
3960 };
3961 let upper = upper_effective.as_ref();
3962
3963 let problem = exact_joint_multistart_outer_problem(
3964 theta0,
3965 lower,
3966 upper,
3967 rho_dim,
3968 coord_dim,
3969 theta_dim,
3970 Derivative::Analytic,
3971 if analytic_outer_hessian_available && !suppress_outer_hessian_for_nfree {
3972 DeclaredHessianForm::Either
3973 } else {
3974 DeclaredHessianForm::Unavailable
3979 },
3980 prefer_gradient_only,
3981 suppress_outer_hessian_for_nfree,
3992 seed_risk_profile_for_likelihood_family(&family),
3993 kappa_options.rel_tol.max(1e-6),
3994 kappa_options.max_outer_iter.max(1),
3995 Some(5.0),
3999 Some(kappa_options.log_step.clamp(0.25, 1.0)),
4001 None,
4002 Some((data.nrows(), baseline_design.design.ncols())),
4007 !constant_curvature_term_indices(resolvedspec).is_empty(),
4011 kind == SpatialHyperKind::Isotropic
4016 && constant_curvature_term_indices(resolvedspec).is_empty()
4017 && spatial_terms.iter().any(|&term_idx| {
4018 matches!(
4019 resolvedspec
4020 .smooth_terms
4021 .get(term_idx)
4022 .map(|term| &term.basis),
4023 Some(SmoothBasisSpec::Matern { .. })
4024 )
4025 }),
4026 )?;
4027
4028 let eval_outer = |ctx: &mut &mut SpatialJointContext<'_>,
4029 theta: &Array1<f64>,
4030 order: OuterEvalOrder|
4031 -> Result<OuterEval, EstimationError> {
4032 let t0 = std::time::Instant::now();
4033 let allow_second_order_for_call = matches!(order, OuterEvalOrder::ValueGradientHessian)
4034 && analytic_outer_hessian_available;
4035 let gate = ctx.nfree_skip_gate_status(theta, allow_second_order_for_call, true);
4036 let resets_before = ctx.evaluator.slow_path_reset_count();
4037 let raw = ctx.eval_full(theta, order, analytic_outer_hessian_available);
4038 let reset_delta = ctx
4039 .evaluator
4040 .slow_path_reset_count()
4041 .saturating_sub(resets_before);
4042 if reset_delta > 0 {
4043 if !gate.shape {
4044 kphase_nfree_miss_shape.set(kphase_nfree_miss_shape.get() + reset_delta);
4045 }
4046 if gate.shape && !gate.value {
4047 kphase_nfree_miss_value.set(kphase_nfree_miss_value.get() + reset_delta);
4048 }
4049 if gate.shape && gate.value && !gate.gradient {
4050 kphase_nfree_miss_gradient.set(kphase_nfree_miss_gradient.get() + reset_delta);
4051 }
4052 if gate.shape && gate.value && gate.gradient && !gate.penalty {
4053 kphase_nfree_miss_penalty.set(kphase_nfree_miss_penalty.get() + reset_delta);
4054 }
4055 if gate.shape && gate.value && gate.gradient && gate.penalty && !gate.revision {
4056 kphase_nfree_miss_revision.set(kphase_nfree_miss_revision.get() + reset_delta);
4057 }
4058 if gate.shape
4059 && gate.value
4060 && gate.gradient
4061 && gate.penalty
4062 && gate.revision
4063 && gate.second_order
4064 {
4065 kphase_nfree_miss_second_order
4066 .set(kphase_nfree_miss_second_order.get() + reset_delta);
4067 }
4068 if gate.would_skip(true) {
4069 kphase_nfree_miss_other.set(kphase_nfree_miss_other.get() + reset_delta);
4070 }
4071 }
4072 let elapsed_s = t0.elapsed().as_secs_f64();
4073 kphase_eval_calls.set(kphase_eval_calls.get() + 1);
4074 kphase_eval_total_s.set(kphase_eval_total_s.get() + elapsed_s);
4075 let (theta_norm, log_kappa_norm) = kphase_log_norms(theta, rho_dim);
4076 log::info!(
4077 "[KAPPA-PHASE] phase=eval_outer call={} order={:?} design_revision={:?} theta_norm={:.4e} log_kappa_norm={:.4e} elapsed_s={:.4}",
4078 kphase_eval_calls.get(),
4079 order,
4080 Some(ctx.cache.design_revision()),
4081 theta_norm,
4082 log_kappa_norm,
4083 elapsed_s,
4084 );
4085 match raw {
4086 Ok((cost, grad, hess)) => Ok(OuterEval {
4087 cost,
4088 gradient: grad,
4089 hessian: hess,
4090 inner_beta_hint: None,
4091 }),
4092 Err(err) if is_recoverable_trial_point_error(&err) => {
4100 log::debug!(
4101 "[{label}] trial point infeasible (kernel design \
4102 not constructible at theta={theta:?}): {err}; retreating",
4103 );
4104 Ok(OuterEval::infeasible(theta_dim))
4105 }
4106 Err(err) => Err(err),
4107 }
4108 };
4109
4110 let mut obj = problem.build_objective_with_eval_order(
4111 &mut ctx,
4112 |ctx: &mut &mut SpatialJointContext<'_>, theta: &Array1<f64>| {
4113 let t0 = std::time::Instant::now();
4114 let gate = ctx.nfree_skip_gate_status(theta, false, false);
4115 let resets_before = ctx.evaluator.slow_path_reset_count();
4116 let cost = ctx.eval_cost(theta);
4117 let reset_delta = ctx
4118 .evaluator
4119 .slow_path_reset_count()
4120 .saturating_sub(resets_before);
4121 if reset_delta > 0 {
4122 if !gate.shape {
4123 kphase_nfree_miss_shape.set(kphase_nfree_miss_shape.get() + reset_delta);
4124 }
4125 if gate.shape && !gate.value {
4126 kphase_nfree_miss_value.set(kphase_nfree_miss_value.get() + reset_delta);
4127 }
4128 if gate.shape && gate.value && !gate.penalty {
4129 kphase_nfree_miss_penalty.set(kphase_nfree_miss_penalty.get() + reset_delta);
4130 }
4131 if gate.shape && gate.value && gate.penalty && !gate.revision {
4132 kphase_nfree_miss_revision.set(kphase_nfree_miss_revision.get() + reset_delta);
4133 }
4134 if gate.would_skip(false) {
4135 kphase_nfree_miss_other.set(kphase_nfree_miss_other.get() + reset_delta);
4136 }
4137 }
4138 let elapsed_s = t0.elapsed().as_secs_f64();
4139 kphase_cost_calls.set(kphase_cost_calls.get() + 1);
4140 kphase_cost_total_s.set(kphase_cost_total_s.get() + elapsed_s);
4141 let (theta_norm, log_kappa_norm) = kphase_log_norms(theta, rho_dim);
4142 log::info!(
4143 "[KAPPA-PHASE] phase=cost call={} design_revision={:?} theta_norm={:.4e} log_kappa_norm={:.4e} elapsed_s={:.4}",
4144 kphase_cost_calls.get(),
4145 Some(ctx.cache.design_revision()),
4146 theta_norm,
4147 log_kappa_norm,
4148 elapsed_s,
4149 );
4150 Ok(cost)
4151 },
4152 |ctx: &mut &mut SpatialJointContext<'_>, theta: &Array1<f64>| {
4153 eval_outer(
4154 ctx,
4155 theta,
4156 if analytic_outer_hessian_available && !suppress_outer_hessian_for_nfree {
4166 OuterEvalOrder::ValueGradientHessian
4167 } else {
4168 OuterEvalOrder::ValueAndGradient
4169 },
4170 )
4171 },
4172 |ctx: &mut &mut SpatialJointContext<'_>, theta: &Array1<f64>, order: OuterEvalOrder| {
4173 eval_outer(ctx, theta, order)
4174 },
4175 Some(|ctx: &mut &mut SpatialJointContext<'_>| {
4176 ctx.reset();
4177 }),
4178 Some(|ctx: &mut &mut SpatialJointContext<'_>, theta: &Array1<f64>| {
4179 let t0 = std::time::Instant::now();
4180 let eval = ctx.eval_efs(theta);
4181 let elapsed_s = t0.elapsed().as_secs_f64();
4182 kphase_efs_calls.set(kphase_efs_calls.get() + 1);
4183 kphase_efs_total_s.set(kphase_efs_total_s.get() + elapsed_s);
4184 let (theta_norm, log_kappa_norm) = kphase_log_norms(theta, rho_dim);
4185 log::info!(
4186 "[KAPPA-PHASE] phase=efs call={} design_revision={:?} theta_norm={:.4e} log_kappa_norm={:.4e} elapsed_s={:.4}",
4187 kphase_efs_calls.get(),
4188 Some(ctx.cache.design_revision()),
4189 theta_norm,
4190 log_kappa_norm,
4191 elapsed_s,
4192 );
4193 eval
4194 }),
4195 );
4196
4197 let run_label = match kind {
4198 SpatialHyperKind::Anisotropic => "aniso-psi joint REML",
4199 SpatialHyperKind::Isotropic => "iso-kappa joint REML",
4200 };
4201 let result = problem.run(&mut obj, run_label)?;
4202 if !result.converged {
4203 crate::bail_invalid_estim!(
4204 "{} did not converge after {} iterations (final_objective={:.6e}, final_grad_norm={})",
4205 run_label,
4206 result.iterations,
4207 result.final_value,
4208 result.final_grad_norm_report(),
4209 );
4210 }
4211 drop(obj);
4212 let kphase_total_s = kphase_optim_start.elapsed().as_secs_f64();
4213 let kphase_slow_resets = ctx
4214 .evaluator
4215 .slow_path_reset_count()
4216 .saturating_sub(kphase_slow_resets_start);
4217 let kphase_design_revision_delta = ctx
4218 .cache
4219 .design_revision()
4220 .saturating_sub(kphase_design_revision_start);
4221 let kphase_nfree_skip_touches = gam_solve::pirls::nfree_skip_row_element_touches()
4222 .saturating_sub(kphase_nfree_skip_touches_start);
4223 log::info!(
4224 "[KAPPA-PHASE-SUMMARY] n_rows={} log_kappa_dim={} n_cost={} cost_total_s={:.4} n_eval={} eval_total_s={:.4} n_efs={} efs_total_s={:.4} slow_path_resets={} design_revision_delta={} nfree_skip_row_touches={} nfree_miss_shape={} nfree_miss_value={} nfree_miss_gradient={} nfree_miss_penalty={} nfree_miss_revision={} nfree_miss_second_order={} nfree_miss_other={} optim_total_s={:.4}",
4225 data.nrows(),
4226 kphase_log_kappa_dim,
4227 kphase_cost_calls.get(),
4228 kphase_cost_total_s.get(),
4229 kphase_eval_calls.get(),
4230 kphase_eval_total_s.get(),
4231 kphase_efs_calls.get(),
4232 kphase_efs_total_s.get(),
4233 kphase_slow_resets,
4234 kphase_design_revision_delta,
4235 kphase_nfree_skip_touches,
4236 kphase_nfree_miss_shape.get(),
4237 kphase_nfree_miss_value.get(),
4238 kphase_nfree_miss_gradient.get(),
4239 kphase_nfree_miss_penalty.get(),
4240 kphase_nfree_miss_revision.get(),
4241 kphase_nfree_miss_second_order.get(),
4242 kphase_nfree_miss_other.get(),
4243 kphase_total_s,
4244 );
4245 let timing = SpatialLengthScaleOptimizationTiming {
4246 log_kappa_dim: kphase_log_kappa_dim,
4247 cost_calls: kphase_cost_calls.get(),
4248 cost_total_s: kphase_cost_total_s.get(),
4249 eval_calls: kphase_eval_calls.get(),
4250 eval_total_s: kphase_eval_total_s.get(),
4251 efs_calls: kphase_efs_calls.get(),
4252 efs_total_s: kphase_efs_total_s.get(),
4253 slow_path_resets: kphase_slow_resets,
4254 design_revision_delta: kphase_design_revision_delta,
4255 nfree_skip_row_touches: kphase_nfree_skip_touches,
4256 nfree_miss_shape: kphase_nfree_miss_shape.get(),
4257 nfree_miss_value: kphase_nfree_miss_value.get(),
4258 nfree_miss_gradient: kphase_nfree_miss_gradient.get(),
4259 nfree_miss_penalty: kphase_nfree_miss_penalty.get(),
4260 nfree_miss_revision: kphase_nfree_miss_revision.get(),
4261 nfree_miss_second_order: kphase_nfree_miss_second_order.get(),
4262 nfree_miss_other: kphase_nfree_miss_other.get(),
4263 optim_total_s: kphase_total_s,
4264 };
4265 log::trace!(
4266 "[{}] converged in {} iterations, final_value={:.6e}, grad_norm={}",
4267 label,
4268 result.iterations,
4269 result.final_value,
4270 result.final_grad_norm_report(),
4271 );
4272 let theta_star = result.rho;
4276 Ok((theta_star, result.final_value, timing))
4277}
4278
4279fn set_single_term_spatial_length_scale(
4283 term: &mut SmoothTermSpec,
4284 length_scale: f64,
4285) -> Result<(), EstimationError> {
4286 match &mut term.basis {
4287 SmoothBasisSpec::ThinPlate { spec, .. } => {
4288 spec.length_scale = length_scale;
4289 Ok(())
4290 }
4291 SmoothBasisSpec::Matern { spec, .. } => {
4292 spec.length_scale.set_resolved(length_scale);
4293 Ok(())
4294 }
4295 SmoothBasisSpec::Duchon { spec, .. } => {
4296 spec.length_scale = Some(length_scale);
4297 Ok(())
4298 }
4299 _ => Err(EstimationError::InvalidInput(format!(
4300 "term '{}' does not expose a spatial length scale",
4301 term.name
4302 ))),
4303 }
4304}
4305
4306fn set_single_term_spatial_aniso_log_scales(
4310 term: &mut SmoothTermSpec,
4311 eta: Vec<f64>,
4312) -> Result<(), EstimationError> {
4313 let eta = center_aniso_log_scales(&eta);
4314 match &mut term.basis {
4315 SmoothBasisSpec::Matern { spec, .. } => {
4316 spec.aniso_log_scales = Some(eta);
4317 Ok(())
4318 }
4319 SmoothBasisSpec::Duchon { spec, .. } => {
4320 spec.aniso_log_scales = Some(eta);
4321 Ok(())
4322 }
4323 _ => Err(EstimationError::InvalidInput(format!(
4324 "term '{}' does not support aniso_log_scales",
4325 term.name
4326 ))),
4327 }
4328}
4329
4330pub fn get_constant_curvature_kappa(spec: &TermCollectionSpec, term_idx: usize) -> Option<f64> {
4349 constant_curvature_term_spec(spec, term_idx).map(|cc| cc.kappa)
4350}
4351
4352pub fn constant_curvature_kappa_is_fixed(spec: &TermCollectionSpec, term_idx: usize) -> bool {
4359 constant_curvature_term_spec(spec, term_idx).is_some_and(|cc| cc.kappa_fixed)
4360}
4361
4362pub fn constant_curvature_term_indices(spec: &TermCollectionSpec) -> Vec<usize> {
4364 (0..spec.smooth_terms.len())
4365 .filter(|&idx| constant_curvature_term_spec(spec, idx).is_some())
4366 .collect()
4367}
4368
4369#[derive(Debug, Clone)]
4370struct SingleSmoothTermRealization {
4371 design_local: DesignMatrix,
4372 term: SmoothTerm,
4373 dropped_penaltyinfo: Vec<DroppedPenaltyBlockInfo>,
4374}
4375
4376impl SingleSmoothTermRealization {
4377 fn active_penalty_count(&self) -> usize {
4378 self.term.active_penalties.len()
4379 }
4380}
4381
4382fn build_single_smooth_term_realization_with_policy(
4383 data: ArrayView2<'_, f64>,
4384 termspec: &SmoothTermSpec,
4385 policy: &gam_runtime::resource::ResourcePolicy,
4386) -> Result<SingleSmoothTermRealization, BasisError> {
4387 let mut workspace = gam_terms::basis::BasisWorkspace::with_policy(policy.clone());
4388 let raw =
4389 build_smooth_design_withworkspace(data, std::slice::from_ref(termspec), &mut workspace)?;
4390 finish_single_smooth_term_realization(raw)
4391}
4392
4393fn finish_single_smooth_term_realization(
4394 raw: RawSmoothDesign,
4395) -> Result<SingleSmoothTermRealization, BasisError> {
4396 let RawSmoothDesign {
4397 term_designs,
4398 dropped_penaltyinfo,
4399 terms,
4400 ..
4401 } = raw;
4402 let term = terms.into_iter().next().ok_or_else(|| {
4403 BasisError::InvalidInput("single-term smooth build returned no term".to_string())
4404 })?;
4405 let design = term_designs.into_iter().next().ok_or_else(|| {
4406 BasisError::InvalidInput("single-term smooth build returned no term design".to_string())
4407 })?;
4408
4409 Ok(SingleSmoothTermRealization {
4410 design_local: design,
4411 term,
4412 dropped_penaltyinfo,
4413 })
4414}
4415
4416fn wrap_local_build_as_realization(
4423 mut local: LocalSmoothTermBuild,
4424 termspec: &SmoothTermSpec,
4425) -> Result<SingleSmoothTermRealization, String> {
4426 let p_local = local.dim;
4427 let lb_local = if local.box_reparam {
4428 shape_lower_bounds_local(termspec.shape, p_local)
4429 } else {
4430 None
4431 };
4432
4433 let dropped_penaltyinfo = local
4434 .dropped_penalties
4435 .iter()
4436 .map(|info| DroppedPenaltyBlockInfo {
4437 termname: Some(termspec.name.clone()),
4438 penalty: info.clone(),
4439 })
4440 .collect();
4441
4442 let applied_rotation: Option<gam_terms::basis::JointNullRotation> = match (
4446 local.joint_null_rotation.take(),
4447 lb_local.is_some(),
4448 local.linear_constraints.is_some(),
4449 ) {
4450 (Some(rot), false, false) => {
4451 let q = &rot.rotation;
4452 local.design =
4453 apply_smooth_transform_to_design(local.design.clone(), q, &termspec.name).map_err(
4454 |e| {
4455 format!(
4456 "joint-null absorption rotation failed for term '{}': {}",
4457 termspec.name, e
4458 )
4459 },
4460 )?;
4461 for penalty in &mut local.active_penalties {
4462 let qt_s = gam_linalg::faer_ndarray::fast_atb(q, &penalty.matrix);
4463 penalty.matrix = gam_linalg::faer_ndarray::fast_ab(&qt_s, q);
4464 penalty.null_eigenvectors = penalty
4465 .null_eigenvectors
4466 .as_ref()
4467 .map(|basis| gam_linalg::faer_ndarray::fast_atb(q, basis));
4468 penalty.op = None;
4469 penalty.info.kronecker_factors = None;
4470 }
4471 local.kronecker_factored = None;
4472 Some(rot)
4473 }
4474 (Some(_), _, _) => None,
4475 (None, _, _) => None,
4476 };
4477
4478 let smooth_term = SmoothTerm {
4479 name: termspec.name.clone(),
4480 coeff_range: 0..p_local,
4481 shape: termspec.shape,
4482 active_penalties: local.active_penalties.clone(),
4483 dropped_penalties: local.dropped_penalties.clone(),
4484 metadata: local.metadata.clone(),
4485 lower_bounds_local: lb_local,
4486 linear_constraints_local: local.linear_constraints.clone(),
4487 kronecker_factored: local.kronecker_factored.take(),
4488 joint_null_rotation: applied_rotation,
4489 unabsorbed_global_orthogonality: None,
4492 };
4493
4494 Ok(SingleSmoothTermRealization {
4495 design_local: local.design,
4496 term: smooth_term,
4497 dropped_penaltyinfo,
4498 })
4499}
4500
4501fn freeze_geometry_from_metadata(
4512 termspec: &SmoothTermSpec,
4513 metadata: &BasisMetadata,
4514) -> Option<SmoothTermSpec> {
4515 let mut frozen = termspec.clone();
4516 match (&mut frozen.basis, metadata) {
4517 (
4518 SmoothBasisSpec::Matern {
4519 spec,
4520 input_scale: spec_scale,
4521 ..
4522 },
4523 BasisMetadata::Matern {
4524 centers,
4525 input_scale: metadata_scale,
4526 identifiability_transform,
4527 ..
4528 },
4529 ) => {
4530 spec.center_strategy = CenterStrategy::UserProvided(centers.clone());
4531 *spec_scale = Some(*metadata_scale);
4532 if let Some(transform) = identifiability_transform.clone() {
4536 spec.identifiability = MaternIdentifiability::FrozenTransform { transform };
4537 }
4538 Some(frozen)
4539 }
4540 (
4541 SmoothBasisSpec::Duchon {
4542 spec,
4543 input_scale: spec_scale,
4544 ..
4545 },
4546 BasisMetadata::Duchon {
4547 centers,
4548 input_scale: metadata_scale,
4549 ..
4550 },
4551 ) => {
4552 spec.center_strategy = CenterStrategy::UserProvided(centers.clone());
4553 *spec_scale = Some(*metadata_scale);
4554 Some(frozen)
4555 }
4556 (
4557 SmoothBasisSpec::ThinPlate {
4558 spec,
4559 input_scale: spec_scale,
4560 ..
4561 },
4562 BasisMetadata::ThinPlate {
4563 centers,
4564 input_scale: metadata_scale,
4565 ..
4566 },
4567 ) => {
4568 spec.center_strategy = CenterStrategy::UserProvided(centers.clone());
4569 *spec_scale = Some(*metadata_scale);
4570 Some(frozen)
4571 }
4572 _ => None,
4575 }
4576}
4577
4578fn rebuild_smooth_auxiliary_state(
4579 smooth: &mut SmoothDesign,
4580 dropped_penaltyinfo_by_term: &[Vec<DroppedPenaltyBlockInfo>],
4581) -> Result<(), String> {
4582 if dropped_penaltyinfo_by_term.len() != smooth.terms.len() {
4583 return Err(SmoothError::dimension_mismatch(format!(
4584 "smooth dropped-penalty cache mismatch: terms={}, dropped_sets={}",
4585 smooth.terms.len(),
4586 dropped_penaltyinfo_by_term.len()
4587 ))
4588 .into());
4589 }
4590
4591 let total_p = smooth.total_smooth_cols();
4592 let mut coefficient_lower_bounds = Array1::<f64>::from_elem(total_p, f64::NEG_INFINITY);
4593 let mut any_bounds = false;
4594 let mut linear_constraintrows: Vec<Array1<f64>> = Vec::new();
4595 let mut linear_constraint_b: Vec<f64> = Vec::new();
4596
4597 for term in &smooth.terms {
4598 let range = term.coeff_range.clone();
4599 if let Some(lb_local) = term.lower_bounds_local.as_ref() {
4600 if lb_local.len() != range.len() {
4601 return Err(SmoothError::dimension_mismatch(format!(
4602 "smooth lower-bound cache mismatch for term '{}': bounds={}, coeffs={}",
4603 term.name,
4604 lb_local.len(),
4605 range.len()
4606 ))
4607 .into());
4608 }
4609 coefficient_lower_bounds
4610 .slice_mut(s![range.clone()])
4611 .assign(lb_local);
4612 any_bounds = true;
4613 }
4614 if let Some(lin_local) = term.linear_constraints_local.as_ref() {
4615 if lin_local.a.ncols() != range.len() {
4616 return Err(SmoothError::dimension_mismatch(format!(
4617 "smooth linear-constraint cache mismatch for term '{}': cols={}, coeffs={}",
4618 term.name,
4619 lin_local.a.ncols(),
4620 range.len()
4621 ))
4622 .into());
4623 }
4624 for r in 0..lin_local.a.nrows() {
4625 let mut row = Array1::<f64>::zeros(total_p);
4626 row.slice_mut(s![range.clone()]).assign(&lin_local.a.row(r));
4627 linear_constraintrows.push(row);
4628 linear_constraint_b.push(lin_local.b[r]);
4629 }
4630 }
4631 }
4632
4633 smooth.coefficient_lower_bounds = if any_bounds {
4634 Some(coefficient_lower_bounds)
4635 } else {
4636 None
4637 };
4638 smooth.linear_constraints = if linear_constraintrows.is_empty() {
4639 None
4640 } else {
4641 let mut a = Array2::<f64>::zeros((linear_constraintrows.len(), total_p));
4642 for (i, row) in linear_constraintrows.iter().enumerate() {
4643 a.row_mut(i).assign(row);
4644 }
4645 Some(LinearInequalityConstraints {
4646 a,
4647 b: Array1::from_vec(linear_constraint_b),
4648 })
4649 };
4650 smooth.dropped_penaltyinfo = dropped_penaltyinfo_by_term
4651 .iter()
4652 .flat_map(|infos| infos.iter().cloned())
4653 .collect();
4654 Ok(())
4655}
4656
4657fn rebuild_term_collection_auxiliary_state(
4658 spec: &TermCollectionSpec,
4659 design: &mut TermCollectionDesign,
4660) -> Result<(), String> {
4661 if spec.linear_terms.len() != design.linear_ranges.len() {
4662 return Err(SmoothError::dimension_mismatch(format!(
4663 "term-collection linear bookkeeping mismatch: spec_terms={}, design_ranges={}",
4664 spec.linear_terms.len(),
4665 design.linear_ranges.len()
4666 ))
4667 .into());
4668 }
4669
4670 let p_total = design.design.ncols();
4671 let smooth_start = p_total.saturating_sub(design.smooth.total_smooth_cols());
4672 let mut coefficient_lower_bounds = Array1::<f64>::from_elem(p_total, f64::NEG_INFINITY);
4673 let mut any_bounds = false;
4674 let mut linear_constraintrows: Vec<Array1<f64>> = Vec::new();
4675 let mut linear_constraint_b: Vec<f64> = Vec::new();
4676
4677 for (linear, (_, range)) in spec.linear_terms.iter().zip(design.linear_ranges.iter()) {
4678 if range.len() != 1 {
4679 return Err(SmoothError::dimension_mismatch(format!(
4680 "linear term '{}' expected one coefficient column, found {}",
4681 linear.name,
4682 range.len()
4683 ))
4684 .into());
4685 }
4686 let col = range.start;
4687 if let Some(lb) = linear.coefficient_min {
4688 let mut row = Array1::<f64>::zeros(p_total);
4689 row[col] = 1.0;
4690 linear_constraintrows.push(row);
4691 linear_constraint_b.push(lb);
4692 }
4693 if let Some(ub) = linear.coefficient_max {
4694 let mut row = Array1::<f64>::zeros(p_total);
4695 row[col] = -1.0;
4696 linear_constraintrows.push(row);
4697 linear_constraint_b.push(-ub);
4698 }
4699 }
4700
4701 if let Some(lb_smooth) = design.smooth.coefficient_lower_bounds.as_ref() {
4702 if lb_smooth.len() != design.smooth.total_smooth_cols() {
4703 return Err(SmoothError::dimension_mismatch(format!(
4704 "smooth lower-bound width mismatch: bounds={}, smooth_cols={}",
4705 lb_smooth.len(),
4706 design.smooth.total_smooth_cols()
4707 ))
4708 .into());
4709 }
4710 coefficient_lower_bounds
4711 .slice_mut(s![
4712 smooth_start..(smooth_start + design.smooth.total_smooth_cols())
4713 ])
4714 .assign(lb_smooth);
4715 any_bounds = true;
4716 }
4717 if let Some(lin_smooth) = design.smooth.linear_constraints.as_ref() {
4718 if lin_smooth.a.ncols() != design.smooth.total_smooth_cols() {
4719 return Err(SmoothError::dimension_mismatch(format!(
4720 "smooth linear-constraint width mismatch: cols={}, smooth_cols={}",
4721 lin_smooth.a.ncols(),
4722 design.smooth.total_smooth_cols()
4723 ))
4724 .into());
4725 }
4726 let mut a_global = Array2::<f64>::zeros((lin_smooth.a.nrows(), p_total));
4727 a_global
4728 .slice_mut(s![
4729 ..,
4730 smooth_start..(smooth_start + design.smooth.total_smooth_cols())
4731 ])
4732 .assign(&lin_smooth.a);
4733 for r in 0..a_global.nrows() {
4734 linear_constraintrows.push(a_global.row(r).to_owned());
4735 linear_constraint_b.push(lin_smooth.b[r]);
4736 }
4737 }
4738
4739 let lower_bound_constraints = if any_bounds {
4740 linear_constraints_from_lower_bounds_global(&coefficient_lower_bounds)
4741 } else {
4742 None
4743 };
4744 let explicit_linear_constraints = if linear_constraintrows.is_empty() {
4745 None
4746 } else {
4747 let mut a = Array2::<f64>::zeros((linear_constraintrows.len(), p_total));
4748 for (i, row) in linear_constraintrows.iter().enumerate() {
4749 a.row_mut(i).assign(row);
4750 }
4751 Some(LinearInequalityConstraints {
4752 a,
4753 b: Array1::from_vec(linear_constraint_b),
4754 })
4755 };
4756
4757 design.coefficient_lower_bounds = if any_bounds {
4758 Some(coefficient_lower_bounds)
4759 } else {
4760 None
4761 };
4762 design.linear_constraints =
4763 merge_linear_constraints_global(explicit_linear_constraints, lower_bound_constraints)
4764 .map_err(|error| error.to_string())?;
4765 design.dropped_penaltyinfo = design.smooth.dropped_penaltyinfo.clone();
4766 Ok(())
4767}
4768
4769fn theta_values_match(left: &Array1<f64>, right: &Array1<f64>) -> bool {
4770 left.len() == right.len()
4771 && left
4772 .iter()
4773 .zip(right.iter())
4774 .all(|(&l, &r)| l.to_bits() == r.to_bits())
4775}
4776
4777fn latent_values_match(left: &Array1<f64>, right: &Array1<f64>) -> bool {
4778 theta_values_match(left, right)
4779}
4780
4781fn spatial_aniso_matches(left: Option<&[f64]>, right: Option<&[f64]>) -> bool {
4782 match (left, right) {
4783 (None, None) => true,
4784 (Some(a), Some(b)) => {
4785 a.len() == b.len()
4786 && a.iter()
4787 .zip(b.iter())
4788 .all(|(&x, &y)| x.to_bits() == y.to_bits())
4789 }
4790 _ => false,
4791 }
4792}
4793
4794fn spatial_length_scale_matches(left: Option<f64>, right: Option<f64>) -> bool {
4795 match (left, right) {
4796 (None, None) => true,
4797 (Some(a), Some(b)) => a.to_bits() == b.to_bits(),
4798 _ => false,
4799 }
4800}
4801
4802struct FrozenTermCollectionIncrementalRealizer<'d> {
4803 data: ArrayView2<'d, f64>,
4804 spec: TermCollectionSpec,
4805 design: TermCollectionDesign,
4806 fixed_blocks: Vec<DesignBlock>,
4807 dropped_penaltyinfo_by_term: Vec<Vec<DroppedPenaltyBlockInfo>>,
4808 smooth_penalty_ranges: Vec<Range<usize>>,
4809 full_penalty_ranges: Vec<Range<usize>>,
4810 basisworkspace: gam_terms::basis::BasisWorkspace,
4814 spatial_realization_geometry: Vec<Option<SmoothTermSpec>>,
4827 design_revision: u64,
4833}
4834
4835impl<'d> std::fmt::Debug for FrozenTermCollectionIncrementalRealizer<'d> {
4836 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
4837 f.debug_struct("FrozenTermCollectionIncrementalRealizer")
4838 .field("data_shape", &(self.data.nrows(), self.data.ncols()))
4839 .field("fixed_blocks", &self.fixed_blocks.len())
4840 .finish_non_exhaustive()
4841 }
4842}
4843
4844fn emitted_smooth_penalty_ranges(
4853 design: &TermCollectionDesign,
4854) -> Result<(Vec<Range<usize>>, Vec<Range<usize>>), String> {
4855 let leading = design.leading_penalty_blocks_before_smooth();
4856 let mut smooth_penalty_ranges = Vec::with_capacity(design.smooth.terms.len());
4857 let mut full_penalty_ranges = Vec::with_capacity(design.smooth.terms.len());
4858 let mut smooth_cursor = 0usize;
4859 for term_idx in 0..design.smooth.terms.len() {
4860 let full_range = design.smooth_term_penalty_range(term_idx)?;
4861 match full_range {
4862 Some(full_range) => {
4863 let local_start = full_range.start.checked_sub(leading).ok_or_else(|| {
4864 "incremental realizer smooth penalty range precedes the emitted smooth prefix"
4865 .to_string()
4866 })?;
4867 let local_end = full_range.end.checked_sub(leading).ok_or_else(|| {
4868 "incremental realizer smooth penalty range precedes the emitted smooth prefix"
4869 .to_string()
4870 })?;
4871 if local_start != smooth_cursor {
4872 return Err(format!(
4873 "incremental realizer non-contiguous emitted smooth layout at term {term_idx}: expected local start {smooth_cursor}, got {local_start}"
4874 ));
4875 }
4876 smooth_cursor = local_end;
4877 smooth_penalty_ranges.push(local_start..local_end);
4878 full_penalty_ranges.push(full_range);
4879 }
4880 None => {
4881 smooth_penalty_ranges.push(smooth_cursor..smooth_cursor);
4882 let global_cursor = leading.checked_add(smooth_cursor).ok_or_else(|| {
4883 "incremental realizer empty smooth penalty range overflow".to_string()
4884 })?;
4885 full_penalty_ranges.push(global_cursor..global_cursor);
4886 }
4887 }
4888 }
4889 if smooth_cursor != design.smooth.penalties.len() {
4890 return Err(format!(
4891 "incremental realizer smooth penalty mismatch: ranged={}, actual={}",
4892 smooth_cursor,
4893 design.smooth.penalties.len()
4894 ));
4895 }
4896 Ok((smooth_penalty_ranges, full_penalty_ranges))
4897}
4898
4899impl<'d> FrozenTermCollectionIncrementalRealizer<'d> {
4900 fn new(
4901 data: ArrayView2<'d, f64>,
4902 spec: TermCollectionSpec,
4903 design: TermCollectionDesign,
4904 ) -> Result<Self, String> {
4905 let policy = gam_runtime::resource::ResourcePolicy::default_library();
4906 Self::new_with_policy(data, spec, design, &policy)
4907 }
4908
4909 fn new_with_policy(
4910 data: ArrayView2<'d, f64>,
4911 spec: TermCollectionSpec,
4912 design: TermCollectionDesign,
4913 policy: &gam_runtime::resource::ResourcePolicy,
4914 ) -> Result<Self, String> {
4915 if spec.smooth_terms.len() != design.smooth.terms.len() {
4916 return Err(SmoothError::dimension_mismatch(format!(
4917 "incremental realizer smooth term mismatch: spec_terms={}, design_terms={}",
4918 spec.smooth_terms.len(),
4919 design.smooth.terms.len()
4920 ))
4921 .into());
4922 }
4923
4924 let (smooth_penalty_ranges, full_penalty_ranges) = emitted_smooth_penalty_ranges(&design)?;
4929 let fixed_blocks = build_term_collection_fixed_blocks(data, &spec)
4930 .map_err(|e| format!("failed to cache fixed term-collection blocks: {e}"))?;
4931
4932 let mut dropped_penaltyinfo_by_term = Vec::with_capacity(spec.smooth_terms.len());
4933 for (term_idx, termspec) in spec.smooth_terms.iter().enumerate() {
4934 let realization = build_single_smooth_term_realization_with_policy(
4935 data, termspec, policy,
4936 )
4937 .map_err(|e| {
4938 format!(
4939 "failed to build cached realization for smooth term '{}' (index {}): {e}",
4940 termspec.name, term_idx
4941 )
4942 })?;
4943 let expected_cols = design.smooth.terms[term_idx].coeff_range.len();
4944 if realization.design_local.ncols() != expected_cols {
4945 return Err(SmoothError::dimension_mismatch(format!(
4946 "cached realization width mismatch for term '{}': cached_cols={}, design_cols={}",
4947 termspec.name,
4948 realization.design_local.ncols(),
4949 expected_cols
4950 ))
4951 .into());
4952 }
4953 if realization.active_penalty_count()
4954 != design.smooth.terms[term_idx].active_penalties.len()
4955 {
4956 return Err(SmoothError::dimension_mismatch(format!(
4957 "cached realization penalty mismatch for term '{}': cached_penalties={}, design_penalties={}",
4958 termspec.name,
4959 realization.active_penalty_count(),
4960 design.smooth.terms[term_idx].active_penalties.len()
4961 ))
4962 .into());
4963 }
4964 dropped_penaltyinfo_by_term.push(realization.dropped_penaltyinfo);
4965 }
4966
4967 let geometry_slots = spec.smooth_terms.len();
4968 Ok(Self {
4969 data,
4970 spec,
4971 design,
4972 fixed_blocks,
4973 dropped_penaltyinfo_by_term,
4974 smooth_penalty_ranges,
4975 full_penalty_ranges,
4976 basisworkspace: gam_terms::basis::BasisWorkspace::with_policy(policy.clone()),
4977 spatial_realization_geometry: vec![None; geometry_slots],
4978 design_revision: 0,
4979 })
4980 }
4981
4982 fn design_revision(&self) -> u64 {
4983 self.design_revision
4984 }
4985
4986 fn spec(&self) -> &TermCollectionSpec {
4987 &self.spec
4988 }
4989
4990 fn design(&self) -> &TermCollectionDesign {
4991 &self.design
4992 }
4993
4994 fn supports_nfree_penalty_rekey(&self, spatial_terms: &[usize]) -> bool {
5035 if spatial_terms.len() != 1 {
5036 return false;
5037 }
5038 let term_idx = spatial_terms[0];
5039 matches!(
5040 self.design.smooth.terms.get(term_idx).map(|t| &t.metadata),
5041 Some(BasisMetadata::Duchon { .. } | BasisMetadata::ThinPlate { .. })
5042 )
5043 }
5044
5045 fn supports_nfree_gradient_only_routing(&self, spatial_terms: &[usize]) -> bool {
5054 if spatial_terms.len() != 1 {
5055 return false;
5056 }
5057 let term_idx = spatial_terms[0];
5058 matches!(
5059 self.design.smooth.terms.get(term_idx).map(|t| &t.metadata),
5060 Some(BasisMetadata::Duchon { .. } | BasisMetadata::ThinPlate { .. })
5061 )
5062 }
5063
5064 fn canonical_penalties_at_psi(
5077 &mut self,
5078 spatial_terms: &[usize],
5079 psi: &[f64],
5080 ) -> Result<(Vec<gam_terms::construction::CanonicalPenalty>, Vec<usize>), String> {
5081 if spatial_terms.len() != 1 {
5082 return Err(format!(
5083 "n-free penalty re-key requires exactly one spatial term, found {}",
5084 spatial_terms.len()
5085 ));
5086 }
5087 let term_idx = spatial_terms[0];
5088 let (ls_opt, aniso_from_psi) = spatial_term_psi_to_length_scale_and_aniso(psi);
5094 let termspec =
5097 self.spec.smooth_terms.get(term_idx).ok_or_else(|| {
5098 format!("spatial term {term_idx} out of range for n-free penalty")
5099 })?;
5100 let term = self
5101 .design
5102 .smooth
5103 .terms
5104 .get(term_idx)
5105 .ok_or_else(|| format!("realized smooth term {term_idx} out of range"))?;
5106 let p_total = self.design.design.ncols();
5109 let (locals, nullspace_dims): (Vec<Array2<f64>>, Vec<usize>) = match &term.metadata {
5110 BasisMetadata::Duchon {
5111 centers,
5112 identifiability_transform,
5113 operator_collocation_points,
5114 power,
5115 nullspace_order,
5116 aniso_log_scales,
5117 input_scale,
5118 radial_reparam,
5119 ..
5120 } => {
5121 let operator_penalties = match &termspec.basis {
5122 SmoothBasisSpec::Duchon { spec, .. } => spec.operator_penalties.clone(),
5123 _ => gam_terms::basis::DuchonOperatorPenaltySpec::default(),
5124 };
5125 let effective_ls =
5132 ls_opt.map(|length| input_scale.to_standardized_units(length));
5133 gam_terms::basis::duchon_penalties_at_length_scale(
5134 centers.view(),
5135 identifiability_transform.as_ref(),
5136 operator_collocation_points.as_ref().map(|p| p.view()),
5137 &operator_penalties,
5138 *power,
5139 *nullspace_order,
5140 aniso_log_scales.as_deref(),
5141 radial_reparam.as_ref(),
5142 effective_ls,
5143 &mut self.basisworkspace,
5144 )
5145 .map_err(|e| e.to_string())?
5146 }
5147 BasisMetadata::Matern {
5148 centers,
5149 periodic,
5150 nu,
5151 include_intercept,
5152 identifiability_transform,
5153 aniso_log_scales,
5154 input_scale,
5155 ..
5156 } => {
5157 let ls = ls_opt.ok_or_else(|| {
5164 "Matérn n-free penalty re-key requires a finite length-scale".to_string()
5165 })?;
5166 let effective_ls = input_scale.to_standardized_units(ls);
5167 let aniso_for_penalty = aniso_from_psi.as_deref().or(aniso_log_scales.as_deref());
5168 let filtered = matern_operator_penalty_triplet_at_length_scale(
5179 centers.view(),
5180 periodic.as_deref(),
5181 identifiability_transform.as_ref(),
5182 *nu,
5183 *include_intercept,
5184 aniso_for_penalty,
5185 effective_ls,
5186 )
5187 .map_err(|e| e.to_string())?;
5188 let locals = filtered
5189 .active
5190 .iter()
5191 .map(|penalty| penalty.matrix.clone())
5192 .collect();
5193 let nullspace_dims = filtered
5194 .active
5195 .iter()
5196 .map(|penalty| penalty.nullity)
5197 .collect();
5198 (locals, nullspace_dims)
5199 }
5200 BasisMetadata::ThinPlate {
5201 centers,
5202 identifiability_transform,
5203 radial_reparam,
5204 ..
5205 } => {
5206 let ls = ls_opt.ok_or_else(|| {
5207 "thin-plate n-free penalty re-key requires a finite length-scale".to_string()
5208 })?;
5209 let double_penalty = match &termspec.basis {
5210 SmoothBasisSpec::ThinPlate { spec, .. } => spec.double_penalty,
5211 _ => false,
5212 };
5213 gam_terms::basis::thin_plate_penalties_at_length_scale(
5214 centers.view(),
5215 identifiability_transform.as_ref(),
5216 radial_reparam.as_ref(),
5217 ls,
5218 double_penalty,
5219 &mut self.basisworkspace,
5220 )
5221 .map_err(|e| e.to_string())?
5222 }
5223 other => {
5224 return Err(format!(
5225 "n-free penalty re-key unsupported for basis metadata {:?}",
5226 std::mem::discriminant(other)
5227 ));
5228 }
5229 };
5230 let templates = &self.design.penalties;
5235 if templates.len() != locals.len() {
5236 return Err(format!(
5237 "n-free penalty re-key produced {} blocks but the frozen design carries {} \
5238 — penalty topology is not ψ-stable",
5239 locals.len(),
5240 templates.len()
5241 ));
5242 }
5243 let specs: Vec<gam_solve::estimate::PenaltySpec> = templates
5244 .iter()
5245 .zip(locals.into_iter())
5246 .map(|(tmpl, local)| gam_solve::estimate::PenaltySpec::Block {
5247 local,
5248 col_range: tmpl.col_range.clone(),
5249 prior_mean: tmpl.prior_mean.clone(),
5250 structure_hint: tmpl.structure_hint.clone(),
5251 op: tmpl.op.clone(),
5252 })
5253 .collect();
5254 gam_terms::construction::canonicalize_penalty_specs(
5255 &specs,
5256 &nullspace_dims,
5257 p_total,
5258 "nfree-psi-penalty",
5259 )
5260 .map_err(|e| e.to_string())
5261 }
5262
5263 fn canonical_penalty_derivatives_at_psi(
5264 &mut self,
5265 spatial_terms: &[usize],
5266 psi: &[f64],
5267 ) -> Result<(Range<usize>, usize, Vec<Array2<f64>>), String> {
5268 if spatial_terms.len() != 1 {
5269 return Err(format!(
5270 "n-free penalty derivative re-key requires exactly one spatial term, found {}",
5271 spatial_terms.len()
5272 ));
5273 }
5274 let term_idx = spatial_terms[0];
5275 let (ls_opt, aniso_from_psi) = spatial_term_psi_to_length_scale_and_aniso(psi);
5276 let termspec = self.spec.smooth_terms.get(term_idx).ok_or_else(|| {
5277 format!("spatial term {term_idx} out of range for n-free penalty derivative")
5278 })?;
5279 let term = self
5280 .design
5281 .smooth
5282 .terms
5283 .get(term_idx)
5284 .ok_or_else(|| format!("realized smooth term {term_idx} out of range"))?;
5285 let p_total = self.design.design.ncols();
5286 let smooth_start = p_total.saturating_sub(self.design.smooth.total_smooth_cols());
5287 let global_range =
5288 (smooth_start + term.coeff_range.start)..(smooth_start + term.coeff_range.end);
5289
5290 let locals = match &term.metadata {
5291 BasisMetadata::Duchon {
5292 centers,
5293 identifiability_transform,
5294 operator_collocation_points,
5295 power,
5296 nullspace_order,
5297 aniso_log_scales,
5298 input_scale,
5299 radial_reparam,
5300 ..
5301 } => {
5302 let mut spec = match &termspec.basis {
5303 SmoothBasisSpec::Duchon { spec, .. } => spec.clone(),
5304 _ => {
5305 return Err(
5306 "Duchon n-free penalty derivative requires a Duchon term spec"
5307 .to_string(),
5308 );
5309 }
5310 };
5311 let effective_ls =
5312 ls_opt.map(|length| input_scale.to_standardized_units(length));
5313 spec.length_scale = effective_ls;
5314 spec.power = *power;
5315 spec.nullspace_order = *nullspace_order;
5316 spec.aniso_log_scales = aniso_log_scales.clone();
5317 spec.radial_reparam = radial_reparam.clone();
5320 if spec.length_scale.is_none() {
5321 return Err(
5322 "Duchon n-free penalty derivative requires a hybrid length-scale"
5323 .to_string(),
5324 );
5325 }
5326 let collocation = operator_collocation_points
5327 .as_ref()
5328 .map(|points| points.view())
5329 .unwrap_or_else(|| centers.view());
5330 let (_native_sources, mut first, _native_second) =
5331 gam_terms::basis::build_duchon_native_penalty_psi_derivatives(
5332 centers.view(),
5333 &spec,
5334 identifiability_transform.as_ref(),
5335 &mut self.basisworkspace,
5336 )
5337 .map_err(|e| e.to_string())?;
5338 let (_operator_sources, operator_first, _operator_second) =
5339 gam_terms::basis::build_duchon_operator_penalty_psi_derivatives(
5340 collocation,
5341 centers.view(),
5342 &spec,
5343 identifiability_transform.as_ref(),
5344 &mut self.basisworkspace,
5345 )
5346 .map_err(|e| e.to_string())?;
5347 first.extend(operator_first);
5348 first
5349 }
5350 BasisMetadata::Matern {
5351 centers,
5352 periodic,
5353 nu,
5354 include_intercept,
5355 identifiability_transform,
5356 aniso_log_scales,
5357 input_scale,
5358 ..
5359 } => {
5360 let ls = ls_opt.ok_or_else(|| {
5361 "Matérn n-free penalty derivative requires a finite length-scale".to_string()
5362 })?;
5363 let effective_ls = input_scale.to_standardized_units(ls);
5364 let penalty_centers = gam_terms::basis::expand_periodic_centers(
5365 ¢ers.to_owned(),
5366 periodic.as_deref(),
5367 )
5368 .map_err(|e| e.to_string())?;
5369 let aniso_for_penalty = aniso_from_psi.as_deref().or(aniso_log_scales.as_deref());
5370 let (first, _second) =
5371 gam_terms::basis::build_matern_operator_penalty_psi_derivatives(
5372 penalty_centers.view(),
5373 effective_ls,
5374 *nu,
5375 *include_intercept,
5376 identifiability_transform.as_ref(),
5377 aniso_for_penalty,
5378 )
5379 .map_err(|e| e.to_string())?;
5380 first
5381 }
5382 BasisMetadata::ThinPlate {
5383 centers,
5384 identifiability_transform,
5385 radial_reparam,
5386 ..
5387 } => {
5388 let ls = ls_opt.ok_or_else(|| {
5389 "thin-plate n-free penalty derivative requires a finite length-scale"
5390 .to_string()
5391 })?;
5392 let mut spec = match &termspec.basis {
5393 SmoothBasisSpec::ThinPlate { spec, .. } => spec.clone(),
5394 _ => {
5395 return Err(
5396 "thin-plate n-free penalty derivative requires a ThinPlate term spec"
5397 .to_string(),
5398 );
5399 }
5400 };
5401 spec.length_scale = ls;
5402 if spec.radial_reparam.is_none() {
5403 spec.radial_reparam = radial_reparam.clone();
5404 }
5405 let (primary, _primary_second, nullspace, _nullspace_second) =
5406 gam_terms::basis::build_thin_plate_penalty_psi_derivativeswithworkspace(
5407 centers.view(),
5408 &spec,
5409 identifiability_transform.as_ref(),
5410 &mut self.basisworkspace,
5411 )
5412 .map_err(|e| e.to_string())?;
5413 if self.design.penalties.len() > 1 {
5414 vec![primary, nullspace]
5415 } else {
5416 vec![primary]
5417 }
5418 }
5419 other => {
5420 return Err(format!(
5421 "n-free penalty derivative re-key unsupported for basis metadata {:?}",
5422 std::mem::discriminant(other)
5423 ));
5424 }
5425 };
5426 if locals.len() != self.design.penalties.len() {
5427 return Err(format!(
5428 "n-free penalty derivative re-key produced {} blocks but the frozen design carries {} \
5429 — penalty topology is not ψ-stable",
5430 locals.len(),
5431 self.design.penalties.len()
5432 ));
5433 }
5434 Ok((global_range, p_total, locals))
5435 }
5436
5437 fn apply_log_kappa(
5438 &mut self,
5439 log_kappa: &SpatialLogKappaCoords,
5440 term_indices: &[usize],
5441 ) -> Result<(), String> {
5442 if term_indices.len() != log_kappa.dims_per_term().len() {
5443 return Err(SmoothError::dimension_mismatch(format!(
5444 "incremental realizer log-kappa term mismatch: term_indices={}, dims_per_term={}",
5445 term_indices.len(),
5446 log_kappa.dims_per_term().len()
5447 ))
5448 .into());
5449 }
5450
5451 let mut any_changed = false;
5452 for (slot, &term_idx) in term_indices.iter().enumerate() {
5453 any_changed |= self.apply_log_kappa_to_term(term_idx, log_kappa.term_slice(slot))?;
5454 }
5455
5456 if any_changed {
5457 self.refresh_full_design_operator()?;
5458 rebuild_smooth_auxiliary_state(
5459 &mut self.design.smooth,
5460 &self.dropped_penaltyinfo_by_term,
5461 )?;
5462 rebuild_term_collection_auxiliary_state(&self.spec, &mut self.design)?;
5463 self.design_revision = self.design_revision.wrapping_add(1);
5464 }
5465 Ok(())
5466 }
5467
5468 fn apply_log_kappa_to_term(&mut self, term_idx: usize, psi: &[f64]) -> Result<bool, String> {
5469 if !spatial_term_supports_hyper_optimization(&self.spec, term_idx) {
5470 return Err(SmoothError::invalid_config(format!(
5471 "incremental realizer term {term_idx} does not expose spatial hyperparameters"
5472 ))
5473 .into());
5474 }
5475 let measure_jet_term = measure_jet_term_spec(&self.spec, term_idx).is_some();
5479 let constant_curvature_term = constant_curvature_term_spec(&self.spec, term_idx).is_some();
5483 let mut next_length_scale = None;
5484 let mut next_aniso: Option<Vec<f64>> = None;
5485 if measure_jet_term {
5486 if !set_measure_jet_psi_dials(&mut self.spec, term_idx, psi)
5487 .map_err(|e| e.to_string())?
5488 {
5489 return Ok(false);
5490 }
5491 } else if constant_curvature_term {
5492 if !set_constant_curvature_kappa(&mut self.spec, term_idx, psi)
5493 .map_err(|e| e.to_string())?
5494 {
5495 return Ok(false);
5496 }
5497 } else {
5498 let current_length_scale = get_spatial_length_scale(&self.spec, term_idx);
5499 let current_aniso = get_spatial_aniso_log_scales(&self.spec, term_idx);
5500 let (ls, eta) = spatial_term_psi_to_length_scale_and_aniso(psi);
5501 next_length_scale = ls;
5502 next_aniso = eta;
5503 let same_length = spatial_length_scale_matches(current_length_scale, next_length_scale);
5504 let same_aniso = spatial_aniso_matches(current_aniso.as_deref(), next_aniso.as_deref());
5505 if same_length && same_aniso {
5506 return Ok(false);
5507 }
5508 if let Some(length_scale) = next_length_scale {
5509 set_spatial_length_scale(&mut self.spec, term_idx, length_scale)
5510 .map_err(|e| e.to_string())?;
5511 }
5512 if let Some(eta) = next_aniso.clone() {
5513 set_spatial_aniso_log_scales(&mut self.spec, term_idx, eta)
5514 .map_err(|e| e.to_string())?;
5515 }
5516 }
5517
5518 let geometry_slot = self
5529 .spatial_realization_geometry
5530 .get(term_idx)
5531 .ok_or_else(|| format!("incremental realizer geometry slot {term_idx} out of range"))?;
5532 let mut build_spec = match geometry_slot {
5533 Some(cached) => cached.clone(),
5534 None => self
5535 .spec
5536 .smooth_terms
5537 .get(term_idx)
5538 .ok_or_else(|| format!("incremental realizer smooth term {term_idx} out of range"))?
5539 .clone(),
5540 };
5541 if measure_jet_term {
5542 set_single_term_measure_jet_psi_dials(&mut build_spec, psi)
5546 .map_err(|e| e.to_string())?;
5547 } else if constant_curvature_term {
5548 set_single_term_constant_curvature_kappa(&mut build_spec, psi)
5553 .map_err(|e| e.to_string())?;
5554 } else {
5555 if let Some(length_scale) = next_length_scale {
5556 set_single_term_spatial_length_scale(&mut build_spec, length_scale)
5557 .map_err(|e| e.to_string())?;
5558 }
5559 if let Some(eta) = next_aniso {
5560 set_single_term_spatial_aniso_log_scales(&mut build_spec, eta)
5561 .map_err(|e| e.to_string())?;
5562 }
5563 }
5564
5565 let termname = build_spec.name.clone();
5566 let local = build_single_local_smooth_term(
5567 self.data,
5568 &build_spec,
5569 &mut self.basisworkspace,
5570 )
5571 .map_err(|e| {
5572 format!(
5573 "failed to rebuild smooth term '{termname}' during incremental κ realization: {e}"
5574 )
5575 })?;
5576
5577 if self.spatial_realization_geometry[term_idx].is_none()
5582 && let Some(frozen) = freeze_geometry_from_metadata(&build_spec, &local.metadata)
5583 {
5584 if let (
5596 SmoothBasisSpec::Matern {
5597 spec: frozen_spec, ..
5598 },
5599 Some(SmoothBasisSpec::Matern {
5600 spec: live_spec, ..
5601 }),
5602 ) = (
5603 &frozen.basis,
5604 self.spec
5605 .smooth_terms
5606 .get_mut(term_idx)
5607 .map(|t| &mut t.basis),
5608 ) {
5609 live_spec.identifiability = frozen_spec.identifiability.clone();
5610 live_spec.center_strategy = frozen_spec.center_strategy.clone();
5611 }
5612 self.spatial_realization_geometry[term_idx] = Some(frozen);
5613 }
5614
5615 let realization = wrap_local_build_as_realization(local, &build_spec)?;
5616 self.replace_term_realization(term_idx, realization)?;
5617 Ok(true)
5618 }
5619
5620 fn replace_term_realization(
5621 &mut self,
5622 term_idx: usize,
5623 realization: SingleSmoothTermRealization,
5624 ) -> Result<(), String> {
5625 let t_replace = std::time::Instant::now();
5626 let SingleSmoothTermRealization {
5627 design_local,
5628 term,
5629 dropped_penaltyinfo,
5630 } = realization;
5631 let SmoothTerm {
5632 name,
5633 active_penalties,
5634 dropped_penalties,
5635 metadata,
5636 lower_bounds_local,
5637 linear_constraints_local,
5638 joint_null_rotation,
5639 ..
5640 } = term;
5641 let coeff_range = self
5642 .design
5643 .smooth
5644 .terms
5645 .get(term_idx)
5646 .ok_or_else(|| format!("incremental realizer smooth term {term_idx} out of range"))?
5647 .coeff_range
5648 .clone();
5649 if design_local.ncols() != coeff_range.len() {
5650 return Err(SmoothError::dimension_mismatch(format!(
5651 "incremental realizer width mismatch for term {}: rebuilt_cols={}, cached_cols={}",
5652 term_idx,
5653 design_local.ncols(),
5654 coeff_range.len()
5655 ))
5656 .into());
5657 }
5658 if design_local.nrows() != self.design.design.nrows() {
5659 return Err(SmoothError::dimension_mismatch(format!(
5660 "incremental realizer row mismatch for term {}: rebuilt_rows={}, design_rows={}",
5661 term_idx,
5662 design_local.nrows(),
5663 self.design.design.nrows()
5664 ))
5665 .into());
5666 }
5667
5668 let smooth_penalty_range = self
5669 .smooth_penalty_ranges
5670 .get(term_idx)
5671 .ok_or_else(|| {
5672 format!("incremental realizer missing smooth penalty range for term {term_idx}")
5673 })?
5674 .clone();
5675 let full_penalty_range = self
5676 .full_penalty_ranges
5677 .get(term_idx)
5678 .ok_or_else(|| {
5679 format!("incremental realizer missing full penalty range for term {term_idx}")
5680 })?
5681 .clone();
5682 if active_penalties.len() != smooth_penalty_range.len() {
5683 return Err(SmoothError::dimension_mismatch(format!(
5684 "incremental realizer topology changed for term '{}': active_penalties={}, cached_penalties={}",
5685 name,
5686 active_penalties.len(),
5687 smooth_penalty_range.len()
5688 ))
5689 .into());
5690 }
5691
5692 self.design.smooth.term_designs[term_idx] = design_local;
5693
5694 for (offset, active_penalty) in active_penalties.iter().enumerate() {
5695 let smooth_penalty_idx = smooth_penalty_range.start + offset;
5696 let full_penalty_idx = full_penalty_range.start + offset;
5697 let penalty_local = &active_penalty.matrix;
5698
5699 if penalty_local.nrows() != coeff_range.len()
5700 || penalty_local.ncols() != coeff_range.len()
5701 {
5702 return Err(SmoothError::dimension_mismatch(format!(
5703 "incremental realizer penalty shape mismatch for term '{}' penalty {}: \
5704 penalty is {}x{} but coeff_range has {} columns",
5705 name,
5706 offset,
5707 penalty_local.nrows(),
5708 penalty_local.ncols(),
5709 coeff_range.len()
5710 ))
5711 .into());
5712 }
5713
5714 let smooth_penalty = self
5715 .design
5716 .smooth
5717 .penalties
5718 .get_mut(smooth_penalty_idx)
5719 .ok_or_else(|| {
5720 format!(
5721 "incremental realizer smooth penalty {} out of range for term {}",
5722 smooth_penalty_idx, term_idx
5723 )
5724 })?;
5725 smooth_penalty.local.assign(penalty_local);
5728 smooth_penalty.op = active_penalty.op.clone();
5729
5730 let full_bp = self
5731 .design
5732 .penalties
5733 .get_mut(full_penalty_idx)
5734 .ok_or_else(|| {
5735 format!(
5736 "incremental realizer full penalty {} out of range for term {}",
5737 full_penalty_idx, term_idx
5738 )
5739 })?;
5740 full_bp.local.assign(penalty_local);
5743 full_bp.op = active_penalty.op.clone();
5744
5745 self.design.smooth.nullspace_dims[smooth_penalty_idx] = active_penalty.nullity;
5746 self.design.nullspace_dims[full_penalty_idx] = active_penalty.nullity;
5747
5748 self.design.smooth.penaltyinfo[smooth_penalty_idx].global_index = smooth_penalty_idx;
5749 self.design.smooth.penaltyinfo[smooth_penalty_idx].termname = Some(name.clone());
5750 self.design.smooth.penaltyinfo[smooth_penalty_idx].penalty =
5751 active_penalty.info.clone();
5752
5753 self.design.penaltyinfo[full_penalty_idx].global_index = full_penalty_idx;
5754 self.design.penaltyinfo[full_penalty_idx].termname = Some(name.clone());
5755 self.design.penaltyinfo[full_penalty_idx].penalty = active_penalty.info.clone();
5756 }
5757
5758 let target_term = self.design.smooth.terms.get_mut(term_idx).ok_or_else(|| {
5759 format!("incremental realizer smooth term {term_idx} disappeared during replacement")
5760 })?;
5761 target_term.active_penalties = active_penalties;
5762 target_term.dropped_penalties = dropped_penalties;
5763 target_term.metadata = metadata;
5764 target_term.lower_bounds_local = lower_bounds_local;
5765 target_term.linear_constraints_local = linear_constraints_local;
5766 target_term.joint_null_rotation = joint_null_rotation;
5767 self.dropped_penaltyinfo_by_term[term_idx] = dropped_penaltyinfo;
5768 log::info!(
5769 "[STAGE] smooth basis rebuild (term {}, '{}', cols={}): {:.3}s",
5770 term_idx,
5771 target_term.name,
5772 coeff_range.len(),
5773 t_replace.elapsed().as_secs_f64(),
5774 );
5775 Ok(())
5776 }
5777
5778 fn refresh_full_design_operator(&mut self) -> Result<(), String> {
5779 let mut blocks = Vec::<DesignBlock>::with_capacity(
5780 self.fixed_blocks.len() + self.design.smooth.term_designs.len(),
5781 );
5782 blocks.extend(self.fixed_blocks.iter().cloned());
5783 for term_design in &self.design.smooth.term_designs {
5784 blocks.push(DesignBlock::from(term_design));
5785 }
5786 self.design.design = assemble_term_collection_design_matrix(blocks)
5787 .map_err(|e| format!("failed to refresh term-collection design: {e}"))?;
5788 Ok(())
5789 }
5790}
5791
5792fn build_term_collection_fixed_blocks(
5793 data: ArrayView2<'_, f64>,
5794 spec: &TermCollectionSpec,
5795) -> Result<Vec<DesignBlock>, BasisError> {
5796 let mut blocks = Vec::<DesignBlock>::new();
5797 if !term_collection_has_anchored_bspline(spec) {
5798 blocks.push(DesignBlock::Intercept(data.nrows()));
5799 }
5800
5801 if !spec.linear_terms.is_empty() {
5802 let mut linear_block = Array2::<f64>::zeros((data.nrows(), spec.linear_terms.len()));
5803 for (j, linear) in spec.linear_terms.iter().enumerate() {
5804 let column = linear
5808 .realized_design_column(data)
5809 .map_err(BasisError::InvalidInput)?;
5810 linear_block.column_mut(j).assign(&column);
5811 }
5812 blocks.push(DesignBlock::Dense(
5813 gam_linalg::matrix::DenseDesignMatrix::from(linear_block),
5814 ));
5815 }
5816
5817 for term in &spec.random_effect_terms {
5818 let block = build_random_effect_block(data, term)?;
5819 let re_op = RandomEffectOperator::new(block.group_ids, block.num_groups);
5820 blocks.push(DesignBlock::RandomEffect(Arc::new(re_op)));
5821 }
5822
5823 Ok(blocks)
5824}
5825
5826pub struct SpatialLengthScaleOptimizationResult<FitOut> {
5831 pub resolved_specs: Vec<TermCollectionSpec>,
5832 pub designs: Vec<TermCollectionDesign>,
5833 pub fit: FitOut,
5834 pub certified_outer: Option<gam_solve::rho_optimizer::CertifiedOuterResult>,
5835 pub timing: Option<SpatialLengthScaleOptimizationTiming>,
5836}
5837
5838pub struct ExactJointEvaluation<M> {
5846 pub objective: f64,
5847 pub gradient: Array1<f64>,
5848 pub hessian: gam_problem::HessianValue,
5849 pub mode: M,
5850}
5851
5852pub struct ExactJointEfsEvaluation<M> {
5855 pub evaluation: gam_problem::EfsEval,
5856 pub mode: M,
5857}
5858
5859pub enum SpatialFitProvenance<'a, M> {
5860 NoOuterOptimization,
5861 Certified {
5862 outer: &'a gam_solve::rho_optimizer::CertifiedOuterResult,
5863 mode: M,
5864 },
5865}
5866
5867#[derive(Debug, Clone)]
5869pub struct ExactJointHyperSetup {
5870 rho0: Array1<f64>,
5871 rho_lower: Array1<f64>,
5872 rho_upper: Array1<f64>,
5873 log_kappa0: SpatialLogKappaCoords,
5874 log_kappa_lower: SpatialLogKappaCoords,
5875 log_kappa_upper: SpatialLogKappaCoords,
5876 auxiliary0: Array1<f64>,
5877 auxiliary_lower: Array1<f64>,
5878 auxiliary_upper: Array1<f64>,
5879}
5880
5881impl ExactJointHyperSetup {
5882 fn sanitize_rho_seed(
5883 rho0: Array1<f64>,
5884 rho_lower: &Array1<f64>,
5885 rho_upper: &Array1<f64>,
5886 ) -> Array1<f64> {
5887 Array1::from_iter(rho0.iter().enumerate().map(|(idx, &value)| {
5888 let lo = rho_lower[idx];
5889 let hi = rho_upper[idx];
5890 let fallback = 0.0_f64.clamp(lo, hi);
5891 if value.is_finite() {
5892 value.clamp(lo, hi)
5893 } else {
5894 fallback
5895 }
5896 }))
5897 }
5898
5899 pub(crate) fn new(
5900 rho0: Array1<f64>,
5901 rho_lower: Array1<f64>,
5902 rho_upper: Array1<f64>,
5903 log_kappa0: SpatialLogKappaCoords,
5904 log_kappa_lower: SpatialLogKappaCoords,
5905 log_kappa_upper: SpatialLogKappaCoords,
5906 ) -> Self {
5907 let rho0 = Self::sanitize_rho_seed(rho0, &rho_lower, &rho_upper);
5908 Self {
5909 rho0,
5910 rho_lower,
5911 rho_upper,
5912 log_kappa0,
5913 log_kappa_lower,
5914 log_kappa_upper,
5915 auxiliary0: Array1::zeros(0),
5916 auxiliary_lower: Array1::zeros(0),
5917 auxiliary_upper: Array1::zeros(0),
5918 }
5919 }
5920
5921 pub(crate) fn with_auxiliary(
5922 mut self,
5923 auxiliary0: Array1<f64>,
5924 auxiliary_lower: Array1<f64>,
5925 auxiliary_upper: Array1<f64>,
5926 ) -> Self {
5927 assert_eq!(
5928 auxiliary0.len(),
5929 auxiliary_lower.len(),
5930 "auxiliary lower bound length mismatch"
5931 );
5932 assert_eq!(
5933 auxiliary0.len(),
5934 auxiliary_upper.len(),
5935 "auxiliary upper bound length mismatch"
5936 );
5937 self.auxiliary0 = Self::sanitize_rho_seed(auxiliary0, &auxiliary_lower, &auxiliary_upper);
5938 self.auxiliary_lower = auxiliary_lower;
5939 self.auxiliary_upper = auxiliary_upper;
5940 self
5941 }
5942
5943 pub(crate) fn rho_dim(&self) -> usize {
5944 self.rho0.len()
5945 }
5946
5947 pub(crate) fn log_kappa_dim(&self) -> usize {
5948 self.log_kappa0.len()
5949 }
5950
5951 pub(crate) fn auxiliary_dim(&self) -> usize {
5952 self.auxiliary0.len()
5953 }
5954
5955 pub(crate) fn theta0(&self) -> Array1<f64> {
5956 let mut out =
5957 Array1::<f64>::zeros(self.rho_dim() + self.log_kappa_dim() + self.auxiliary_dim());
5958 out.slice_mut(s![..self.rho_dim()]).assign(&self.rho0);
5959 out.slice_mut(s![self.rho_dim()..self.rho_dim() + self.log_kappa_dim()])
5960 .assign(self.log_kappa0.as_array());
5961 out.slice_mut(s![self.rho_dim() + self.log_kappa_dim()..])
5962 .assign(&self.auxiliary0);
5963 out
5964 }
5965
5966 pub(crate) fn lower(&self) -> Array1<f64> {
5967 let mut out =
5968 Array1::<f64>::zeros(self.rho_dim() + self.log_kappa_dim() + self.auxiliary_dim());
5969 out.slice_mut(s![..self.rho_dim()]).assign(&self.rho_lower);
5970 out.slice_mut(s![self.rho_dim()..self.rho_dim() + self.log_kappa_dim()])
5971 .assign(self.log_kappa_lower.as_array());
5972 out.slice_mut(s![self.rho_dim() + self.log_kappa_dim()..])
5973 .assign(&self.auxiliary_lower);
5974 out
5975 }
5976
5977 pub(crate) fn upper(&self) -> Array1<f64> {
5978 let mut out =
5979 Array1::<f64>::zeros(self.rho_dim() + self.log_kappa_dim() + self.auxiliary_dim());
5980 out.slice_mut(s![..self.rho_dim()]).assign(&self.rho_upper);
5981 out.slice_mut(s![self.rho_dim()..self.rho_dim() + self.log_kappa_dim()])
5982 .assign(self.log_kappa_upper.as_array());
5983 out.slice_mut(s![self.rho_dim() + self.log_kappa_dim()..])
5984 .assign(&self.auxiliary_upper);
5985 out
5986 }
5987
5988 pub(crate) fn log_kappa_dims_per_term(&self) -> Vec<usize> {
5990 self.log_kappa0.dims_per_term().to_vec()
5991 }
5992}
5993
5994struct ExactJointDesignCache<'d> {
6000 realizers: Vec<FrozenTermCollectionIncrementalRealizer<'d>>,
6001 block_term_indices: Vec<Vec<usize>>,
6002 current_theta: Option<Array1<f64>>,
6003 last_cost: Option<f64>,
6004 last_eval: Option<(f64, Array1<f64>, gam_problem::HessianValue)>,
6005 rho_dim: usize,
6006 all_dims: Vec<usize>,
6007 log_kappa_dim: usize,
6008 block_term_counts: Vec<usize>,
6009}
6010
6011impl<'d> ExactJointDesignCache<'d> {
6012 fn new(
6013 data: ArrayView2<'d, f64>,
6014 blocks: Vec<(TermCollectionSpec, TermCollectionDesign, Vec<usize>)>,
6015 rho_dim: usize,
6016 all_dims: Vec<usize>,
6017 ) -> Result<Self, String> {
6018 let n_blocks = blocks.len();
6019 let mut realizers = Vec::with_capacity(n_blocks);
6020 let mut block_term_indices = Vec::with_capacity(n_blocks);
6021 let mut block_term_counts = Vec::with_capacity(n_blocks);
6022
6023 for (spec, design, terms) in blocks {
6024 block_term_counts.push(terms.len());
6025 block_term_indices.push(terms);
6026 realizers.push(FrozenTermCollectionIncrementalRealizer::new(
6027 data, spec, design,
6028 )?);
6029 }
6030
6031 Ok(Self {
6032 realizers,
6033 block_term_indices,
6034 current_theta: None,
6035 last_cost: None,
6036 last_eval: None,
6037 rho_dim,
6038 log_kappa_dim: all_dims.iter().sum(),
6039 all_dims,
6040 block_term_counts,
6041 })
6042 }
6043
6044 fn ensure_theta(&mut self, theta: &Array1<f64>) -> Result<(), String> {
6045 if self
6046 .current_theta
6047 .as_ref()
6048 .is_some_and(|cached| theta_values_match(cached, theta))
6049 {
6050 return Ok(());
6051 }
6052
6053 let t_ensure = std::time::Instant::now();
6054 let kappa_theta_len = self.rho_dim + self.log_kappa_dim;
6055 if theta.len() < kappa_theta_len {
6056 return Err(SmoothError::dimension_mismatch(format!(
6057 "exact-joint theta length mismatch: got {}, expected at least {} (rho_dim={}, log_kappa_dim={})",
6058 theta.len(),
6059 kappa_theta_len,
6060 self.rho_dim,
6061 self.log_kappa_dim
6062 ))
6063 .into());
6064 }
6065 let theta_kappa = theta.slice(s![..kappa_theta_len]).to_owned();
6066 let full_log_kappa = SpatialLogKappaCoords::from_theta_tail_with_dims(
6067 &theta_kappa,
6068 self.rho_dim,
6069 self.all_dims.clone(),
6070 );
6071
6072 let n = self.realizers.len();
6076 let mut remaining = full_log_kappa;
6077 for block_idx in 0..n {
6078 let count = self.block_term_counts[block_idx];
6079 if block_idx < n - 1 {
6080 let (block_lk, rest) = remaining.split_at(count);
6081 self.realizers[block_idx]
6082 .apply_log_kappa(&block_lk, &self.block_term_indices[block_idx])?;
6083 remaining = rest;
6084 } else {
6085 self.realizers[block_idx]
6087 .apply_log_kappa(&remaining, &self.block_term_indices[block_idx])?;
6088 }
6089 }
6090
6091 log::info!(
6092 "[STAGE] ensure_theta (n-block, {} blocks, {} realizers): {:.3}s",
6093 n,
6094 self.realizers.len(),
6095 t_ensure.elapsed().as_secs_f64(),
6096 );
6097 self.current_theta = Some(theta.clone());
6098 self.last_cost = None;
6099 self.last_eval = None;
6100 Ok(())
6101 }
6102
6103 impl_exact_joint_theta_memo!();
6104
6105 fn store_cost_only(&mut self, theta: &Array1<f64>, cost: f64) {
6111 if self
6112 .current_theta
6113 .as_ref()
6114 .is_some_and(|cached| theta_values_match(cached, theta))
6115 {
6116 self.last_cost = Some(cost);
6117 }
6118 }
6119
6120 fn invalidate_objective_memo(&mut self) {
6123 self.last_cost = None;
6124 self.last_eval = None;
6125 }
6126
6127 fn specs(&self) -> Vec<&TermCollectionSpec> {
6128 self.realizers.iter().map(|r| r.spec()).collect()
6129 }
6130
6131 fn designs(&self) -> Vec<&TermCollectionDesign> {
6132 self.realizers.iter().map(|r| r.design()).collect()
6133 }
6134
6135 fn design_revision(&self) -> u64 {
6145 self.realizers
6146 .iter()
6147 .fold(0u64, |acc, r| acc.wrapping_add(r.design_revision()))
6148 }
6149}
6150
6151pub(crate) fn seed_risk_profile_for_likelihood_family(
6152 family: &LikelihoodSpec,
6153) -> gam_problem::SeedRiskProfile {
6154 match &family.response {
6155 ResponseFamily::Gaussian => gam_problem::SeedRiskProfile::Gaussian,
6156 ResponseFamily::RoystonParmar => gam_problem::SeedRiskProfile::Survival,
6157 ResponseFamily::Binomial
6158 | ResponseFamily::Poisson
6159 | ResponseFamily::Tweedie { .. }
6160 | ResponseFamily::NegativeBinomial { .. }
6161 | ResponseFamily::Beta { .. }
6162 | ResponseFamily::Gamma => gam_problem::SeedRiskProfile::GeneralizedLinear,
6163 }
6164}
6165
6166const EXACT_JOINT_SECOND_ORDER_THETA_CAP: usize = 8;
6174
6175fn exact_joint_seed_config(
6176 risk_profile: gam_problem::SeedRiskProfile,
6177 auxiliary_dim: usize,
6178 initial_seed_only: bool,
6179) -> gam_problem::SeedConfig {
6180 let mut config = gam_problem::SeedConfig {
6181 risk_profile,
6182 num_auxiliary_trailing: auxiliary_dim,
6183 ..Default::default()
6184 };
6185 match risk_profile {
6186 gam_problem::SeedRiskProfile::Gaussian
6187 | gam_problem::SeedRiskProfile::GaussianLocationScale => {
6188 config.max_seeds = 4;
6189 config.seed_budget = 2;
6190 }
6191 gam_problem::SeedRiskProfile::GeneralizedLinear => {
6192 config.max_seeds = 1;
6197 config.seed_budget = 1;
6198 config.screen_max_inner_iterations = 8;
6199 }
6200 gam_problem::SeedRiskProfile::Survival => {
6201 config.max_seeds = 8;
6207 config.seed_budget = 4;
6208 config.screen_max_inner_iterations = 8;
6209 }
6210 }
6211 if initial_seed_only {
6212 config.max_seeds = 1;
6219 config.seed_budget = 1;
6220 config.over_smoothing_probe_rho = None;
6221 }
6222 config
6223}
6224
6225#[cfg(test)]
6226mod exact_joint_seed_config_tests {
6227 use super::*;
6228
6229 #[test]
6230 fn exact_joint_marginal_slope_profiles_get_deeper_startup_validation() {
6231 let bms =
6232 exact_joint_seed_config(gam_problem::SeedRiskProfile::GeneralizedLinear, 2, false);
6233 assert_eq!(bms.max_seeds, 1);
6234 assert_eq!(bms.seed_budget, 1);
6235 assert_eq!(bms.screen_max_inner_iterations, 8);
6236 assert_eq!(bms.num_auxiliary_trailing, 2);
6237
6238 let survival = exact_joint_seed_config(gam_problem::SeedRiskProfile::Survival, 3, false);
6239 assert_eq!(survival.max_seeds, 8);
6240 assert_eq!(survival.seed_budget, 4);
6241 assert_eq!(survival.screen_max_inner_iterations, 8);
6242 assert_eq!(survival.num_auxiliary_trailing, 3);
6243 }
6244
6245 #[test]
6246 fn exact_joint_gaussian_keeps_tight_historical_multistart_budget() {
6247 let gaussian = exact_joint_seed_config(gam_problem::SeedRiskProfile::Gaussian, 1, false);
6248 assert_eq!(gaussian.max_seeds, 4);
6249 assert_eq!(gaussian.seed_budget, 2);
6250 assert_eq!(
6251 gaussian.screen_max_inner_iterations,
6252 gam_problem::SeedConfig::default().screen_max_inner_iterations
6253 );
6254 assert_eq!(gaussian.num_auxiliary_trailing, 1);
6255 }
6256
6257 #[test]
6258 fn certified_matern_basin_owns_the_only_joint_start() {
6259 let gaussian = exact_joint_seed_config(gam_problem::SeedRiskProfile::Gaussian, 1, true);
6260 assert_eq!(gaussian.max_seeds, 1);
6261 assert_eq!(gaussian.seed_budget, 1);
6262 assert_eq!(gaussian.over_smoothing_probe_rho, None);
6263 assert_eq!(gaussian.num_auxiliary_trailing, 1);
6264 }
6265}
6266
6267#[cfg(test)]
6268mod wood_reference_df_tests {
6269 use super::*;
6270
6271 #[test]
6277 fn edf1_equals_two_trace_minus_trace_of_square() {
6278 let f = ndarray::array![[0.9_f64, 0.0], [0.0, 0.4]];
6282 let got = wood_reference_df(Some(&f), &(0..2)).unwrap();
6283 assert!(
6284 (got - 1.63).abs() < 1e-12,
6285 "edf1 should be 2*tr - tr(F^2) = 1.63, got {got}"
6286 );
6287 let edf = 1.3;
6290 assert!(got >= edf - 1e-12, "edf1 {got} must be >= edf {edf}");
6291 }
6292
6293 #[test]
6294 fn edf1_never_collapses_below_edf_when_offdiagonals_blow_up() {
6295 let f = ndarray::array![[0.5_f64, 40.0], [40.0, 0.5]];
6302 let tr = 1.0_f64;
6303 let got = wood_reference_df(Some(&f), &(0..2)).unwrap();
6304 assert!(
6305 got >= tr - 1e-12,
6306 "edf1 must be floored at edf (=tr={tr}) even when tr(F^2) explodes, got {got}"
6307 );
6308 assert!(
6309 got.is_finite() && got > 0.0,
6310 "edf1 must stay finite/positive"
6311 );
6312 }
6313
6314 #[test]
6315 fn returns_none_on_nonpositive_or_missing_trace() {
6316 assert!(wood_reference_df(None, &(0..2)).is_none());
6319 let zero = ndarray::array![[0.0_f64, 0.0], [0.0, 0.0]];
6321 assert!(wood_reference_df(Some(&zero), &(0..2)).is_none());
6322 let f = ndarray::array![[0.5_f64, 0.0], [0.0, 0.5]];
6324 assert!(wood_reference_df(Some(&f), &(0..5)).is_none());
6325 }
6326}
6327
6328pub(crate) fn exact_joint_multistart_outer_problem(
6329 theta0: &Array1<f64>,
6330 lower: &Array1<f64>,
6331 upper: &Array1<f64>,
6332 rho_dim: usize,
6333 auxiliary_dim: usize,
6334 n_params: usize,
6335 gradient: gam_problem::Derivative,
6336 hessian: gam_problem::DeclaredHessianForm,
6337 prefer_gradient_only: bool,
6338 disable_fixed_point: bool,
6339 risk_profile: gam_problem::SeedRiskProfile,
6340 tolerance: f64,
6341 max_iter: usize,
6342 bfgs_step_cap: Option<f64>,
6351 bfgs_step_cap_psi: Option<f64>,
6352 screening_cap: Option<Arc<AtomicUsize>>,
6353 profiled_objective_size: Option<(usize, usize)>,
6374 has_constant_curvature: bool,
6383 initial_seed_only: bool,
6388) -> Result<gam_solve::rho_optimizer::OuterProblem, EstimationError> {
6389 if rho_dim > theta0.len() {
6390 crate::bail_invalid_estim!(
6391 "exact joint outer problem declares {rho_dim} smoothing coordinates for theta length {}",
6392 theta0.len(),
6393 );
6394 }
6395 let mut seed_heuristic = theta0.to_vec();
6396 let initial_lambdas = gam_problem::checked_exp_log_strengths(
6397 theta0.iter().take(rho_dim).copied(),
6398 )
6399 .map_err(|error| {
6400 EstimationError::InvalidInput(format!(
6401 "exact joint initial smoothing coordinate is outside the canonical log-strength domain: {error}"
6402 ))
6403 })?;
6404 for (value, lambda) in seed_heuristic[..rho_dim].iter_mut().zip(initial_lambdas) {
6405 *value = lambda;
6406 }
6407 let rho_ceiling = if has_constant_curvature {
6412 gam_solve::estimate::RHO_BOUND
6413 } else {
6414 12.0
6415 };
6416 let mut problem = gam_solve::rho_optimizer::OuterProblem::new(n_params)
6417 .with_gradient(gradient)
6418 .with_hessian(hessian)
6419 .with_prefer_gradient_only(prefer_gradient_only)
6420 .with_disable_fixed_point(disable_fixed_point)
6421 .with_fallback_policy(gam_solve::rho_optimizer::FallbackPolicy::Automatic)
6431 .with_psi_dim(auxiliary_dim)
6432 .with_tolerance(tolerance)
6433 .with_max_iter(max_iter)
6434 .with_bounds(lower.clone(), upper.clone())
6435 .with_initial_rho(theta0.clone())
6436 .with_bfgs_step_cap(bfgs_step_cap)
6437 .with_bfgs_step_cap_psi(bfgs_step_cap_psi)
6438 .with_seed_config({
6439 let mut sc = exact_joint_seed_config(risk_profile, auxiliary_dim, initial_seed_only);
6440 if has_constant_curvature {
6441 sc.bounds = (sc.bounds.0, rho_ceiling);
6445 }
6460 sc
6461 })
6462 .with_rho_bound(rho_ceiling)
6463 .with_heuristic_lambdas(seed_heuristic);
6464 if let Some((n_obs, p_cols)) = profiled_objective_size {
6465 problem = problem
6473 .with_objective_scale(Some(n_obs as f64))
6474 .with_problem_size(n_obs, p_cols)
6475 .with_arc_initial_regularization(Some(0.25))
6476 .with_operator_initial_trust_radius(Some(4.0));
6477 }
6478 if let Some(screening_cap) = screening_cap {
6479 problem = problem
6480 .with_screening_cap(screening_cap)
6481 .with_screen_initial_rho(true);
6482 }
6483 Ok(problem)
6484}
6485
6486pub fn optimize_spatial_length_scale_exact_joint<FitOut, Mode, FitFn, ExactFn, ExactEfsFn, SeedFn>(
6487 data: ArrayView2<'_, f64>,
6488 block_specs: &[TermCollectionSpec],
6489 block_term_indices: &[Vec<usize>],
6490 kappa_options: &SpatialLengthScaleOptimizationOptions,
6491 joint_setup: &ExactJointHyperSetup,
6492 seed_risk_profile: gam_problem::SeedRiskProfile,
6493 analytic_joint_gradient_available: bool,
6494 analytic_joint_hessian_available: bool,
6495 disable_fixed_point: bool,
6496 screening_cap: Option<Arc<AtomicUsize>>,
6497 outer_derivative_policy: gam_model_api::families::custom_family::OuterDerivativePolicy,
6498 mut fit_fn: FitFn,
6499 mut exact_fn: ExactFn,
6500 mut exact_efs_fn: ExactEfsFn,
6501 mut seed_inner_beta_fn: SeedFn,
6502) -> Result<SpatialLengthScaleOptimizationResult<FitOut>, String>
6503where
6504 FitFn: FnMut(
6505 &Array1<f64>,
6506 &[TermCollectionSpec],
6507 &[TermCollectionDesign],
6508 SpatialFitProvenance<'_, Mode>,
6509 ) -> Result<FitOut, String>,
6510 ExactFn: FnMut(
6511 &Array1<f64>,
6512 &[TermCollectionSpec],
6513 &[TermCollectionDesign],
6514 gam_solve::estimate::reml::reml_outer_engine::EvalMode,
6515 &gam_problem::outer_subsample::RowSet,
6516 ) -> Result<ExactJointEvaluation<Mode>, String>,
6517 ExactEfsFn: FnMut(
6518 &Array1<f64>,
6519 &[TermCollectionSpec],
6520 &[TermCollectionDesign],
6521 &gam_problem::outer_subsample::RowSet,
6522 ) -> Result<ExactJointEfsEvaluation<Mode>, String>,
6523 SeedFn: FnMut(&Array1<f64>) -> Result<gam_solve::rho_optimizer::SeedOutcome, EstimationError>,
6524{
6525 let n_blocks = block_specs.len();
6526 if block_term_indices.len() != n_blocks {
6527 return Err(SmoothError::dimension_mismatch(format!(
6528 "block_specs ({}) and block_term_indices ({}) length mismatch",
6529 n_blocks,
6530 block_term_indices.len()
6531 ))
6532 .into());
6533 }
6534
6535 let log_kappa_dim = joint_setup.log_kappa_dim();
6536
6537 log::trace!(
6538 "[spatial-exact-joint] driver entry: aux_dim={} log_kappa_dim={} kappa_enabled={} rho_dim={} theta0_len={}",
6539 joint_setup.auxiliary_dim(),
6540 log_kappa_dim,
6541 kappa_options.enabled,
6542 joint_setup.rho_dim(),
6543 joint_setup.theta0().len()
6544 );
6545
6546 if joint_setup.auxiliary_dim() == 0 && (!kappa_options.enabled || log_kappa_dim == 0) {
6550 log::trace!(
6551 "[spatial-exact-joint] taking fast path (no outer theta optimization in this driver)"
6552 );
6553 let (designs, resolved_specs) = build_term_collection_designs_and_freeze_joint(
6554 data, block_specs,
6555 )
6556 .map_err(|e| {
6557 format!("failed to build and freeze joint block designs during exact joint kappa optimization: {e}")
6558 })?;
6559 let theta0 = joint_setup.theta0();
6560
6561 let spec_refs: Vec<TermCollectionSpec> = resolved_specs.clone();
6563 let design_refs: Vec<TermCollectionDesign> = designs.clone();
6564 let fit = fit_fn(
6565 &theta0,
6566 &spec_refs,
6567 &design_refs,
6568 SpatialFitProvenance::NoOuterOptimization,
6569 )?;
6570 return Ok(SpatialLengthScaleOptimizationResult {
6571 resolved_specs,
6572 designs,
6573 fit,
6574 certified_outer: None,
6575 timing: None,
6576 });
6577 }
6578
6579 let theta0 = joint_setup.theta0();
6583 let lower = joint_setup.lower();
6584 let upper = joint_setup.upper();
6585 if theta0.len() < log_kappa_dim || lower.len() != theta0.len() || upper.len() != theta0.len() {
6586 return Err(SmoothError::dimension_mismatch(format!(
6587 "invalid exact joint theta setup: theta0={}, lower={}, upper={}, required_log_kappa_dim={}",
6588 theta0.len(),
6589 lower.len(),
6590 upper.len(),
6591 log_kappa_dim
6592 ))
6593 .into());
6594 }
6595 let rho_dim = joint_setup.rho_dim();
6596 let all_dims = joint_setup.log_kappa_dims_per_term();
6597
6598 let (boot_designs, best_specs) = build_term_collection_designs_and_freeze_joint(
6600 data,
6601 block_specs,
6602 )
6603 .map_err(|e| {
6604 format!(
6605 "failed to build and freeze joint block designs during exact joint kappa bootstrap: {e}"
6606 )
6607 })?;
6608 let policy_hessian_form = outer_derivative_policy.declared_hessian_form();
6618 let analytic_outer_hessian_available = analytic_joint_hessian_available
6619 && matches!(
6620 policy_hessian_form,
6621 gam_problem::DeclaredHessianForm::Either
6622 | gam_problem::DeclaredHessianForm::Dense
6623 | gam_problem::DeclaredHessianForm::Operator { .. }
6624 );
6625 let prefer_gradient_only = !analytic_outer_hessian_available;
6626
6627 let theta_dim = theta0.len();
6628 let psi_dim = theta_dim - rho_dim;
6629
6630 let cache_blocks: Vec<(TermCollectionSpec, TermCollectionDesign, Vec<usize>)> = best_specs
6632 .iter()
6633 .zip(boot_designs.iter())
6634 .zip(block_term_indices.iter())
6635 .map(|((spec, design), terms)| (spec.clone(), design.clone(), terms.clone()))
6636 .collect();
6637
6638 struct NBlockExactJointState<'d, M> {
6639 cache: ExactJointDesignCache<'d>,
6640 row_set: gam_problem::outer_subsample::RowSet,
6641 staged_pilot_active: bool,
6642 terminal_mode: Option<(Array1<f64>, f64, M)>,
6643 }
6644
6645 impl<M> NBlockExactJointState<'_, M> {
6646 fn ensure_theta(&mut self, theta: &Array1<f64>) -> Result<(), String> {
6647 let theta_changed = !self
6648 .cache
6649 .current_theta
6650 .as_ref()
6651 .is_some_and(|current| theta_values_match(current, theta));
6652 if theta_changed {
6653 self.terminal_mode = None;
6654 }
6655 self.cache.ensure_theta(theta)
6656 }
6657
6658 fn install_terminal_mode(&mut self, theta: &Array1<f64>, objective: f64, mode: M) {
6659 self.terminal_mode = Some((theta.clone(), objective, mode));
6660 }
6661
6662 fn terminal_mode_matches(&self, theta: &Array1<f64>, objective: f64) -> bool {
6663 self.terminal_mode
6664 .as_ref()
6665 .is_some_and(|(mode_theta, mode_objective, _)| {
6666 theta_values_match(mode_theta, theta)
6667 && mode_objective.to_bits() == objective.to_bits()
6668 })
6669 }
6670 }
6671
6672 let mut state = NBlockExactJointState {
6673 cache: ExactJointDesignCache::new(data, cache_blocks, rho_dim, all_dims.clone())?,
6674 row_set: gam_problem::outer_subsample::RowSet::All,
6675 staged_pilot_active: false,
6676 terminal_mode: None,
6677 };
6678
6679 const KAPPA_PILOT_K: usize = 5_000;
6707
6708 let n_total = data.nrows();
6709 let use_staged_kappa = outer_derivative_policy.should_use_staged_kappa(n_total);
6710 if use_staged_kappa {
6711 log::info!(
6712 "[KAPPA-STAGED] auto-engaging pilot+exact schedule: n={} pilot_k={}",
6713 n_total,
6714 KAPPA_PILOT_K,
6715 );
6716 }
6717
6718 fn build_uniform_pilot_subsample(
6735 n_total: usize,
6736 k_target: usize,
6737 seed: u64,
6738 ) -> gam_problem::outer_subsample::OuterScoreSubsample {
6739 use gam_problem::outer_subsample::OuterScoreSubsample;
6740 let k = k_target.min(n_total);
6741 if k == 0 || n_total == 0 {
6742 return OuterScoreSubsample::from_uniform_inclusion_mask(Vec::new(), n_total, seed);
6743 }
6744 let mut mask: Vec<usize> = Vec::with_capacity(k);
6748 let mut state = seed.wrapping_add(0x9E3779B97F4A7C15);
6750 let splitmix = |s: &mut u64| -> u64 { gam_linalg::utils::splitmix64(s) };
6751 let mut taken = std::collections::HashSet::with_capacity(k);
6752 for j in (n_total - k)..n_total {
6753 let r = (splitmix(&mut state) % (j as u64 + 1)) as usize;
6754 if !taken.insert(r) {
6755 taken.insert(j);
6756 mask.push(j);
6757 } else {
6758 mask.push(r);
6759 }
6760 }
6761 mask.sort_unstable();
6762 mask.dedup();
6763 OuterScoreSubsample::from_uniform_inclusion_mask(mask, n_total, seed)
6764 }
6765
6766 if use_staged_kappa {
6767 let pilot = build_uniform_pilot_subsample(n_total, KAPPA_PILOT_K, n_total as u64);
6768 state.row_set = gam_problem::outer_subsample::RowSet::Subsample {
6769 rows: std::sync::Arc::clone(&pilot.rows),
6770 n_full: n_total,
6771 };
6772 state.staged_pilot_active = true;
6773 }
6774
6775 let exact_fn_cell = std::cell::RefCell::new(&mut exact_fn);
6776 let exact_efs_fn_cell = std::cell::RefCell::new(&mut exact_efs_fn);
6777
6778 use std::cell::Cell;
6793 let kphase_cost_calls: Cell<usize> = Cell::new(0);
6794 let kphase_cost_total_s: Cell<f64> = Cell::new(0.0);
6795 let kphase_eval_calls: Cell<usize> = Cell::new(0);
6796 let kphase_eval_total_s: Cell<f64> = Cell::new(0.0);
6797 let kphase_efs_calls: Cell<usize> = Cell::new(0);
6798 let kphase_efs_total_s: Cell<f64> = Cell::new(0.0);
6799 let kphase_optim_start = std::time::Instant::now();
6800 let kphase_log_kappa_dim = log_kappa_dim;
6801 let kphase_log_norms = |theta: &Array1<f64>| -> (f64, f64) {
6802 let theta_norm = theta.iter().map(|v| v * v).sum::<f64>().sqrt();
6803 let log_kappa_norm = if kphase_log_kappa_dim > 0 && theta.len() >= kphase_log_kappa_dim {
6804 let start = theta.len() - kphase_log_kappa_dim;
6805 theta.iter().skip(start).map(|v| v * v).sum::<f64>().sqrt()
6806 } else {
6807 0.0
6808 };
6809 (theta_norm, log_kappa_norm)
6810 };
6811
6812 use gam_problem::{DeclaredHessianForm, Derivative, OuterEval};
6813 use gam_solve::rho_optimizer::OuterEvalOrder;
6814
6815 let joint_p_cols: usize = boot_designs
6819 .iter()
6820 .map(|d| d.design.ncols())
6821 .sum::<usize>()
6822 .max(1);
6823
6824 let problem = exact_joint_multistart_outer_problem(
6825 &theta0,
6826 &lower,
6827 &upper,
6828 rho_dim,
6829 psi_dim,
6830 theta_dim,
6831 if analytic_joint_gradient_available {
6832 Derivative::Analytic
6833 } else {
6834 Derivative::Unavailable
6835 },
6836 if analytic_outer_hessian_available {
6837 DeclaredHessianForm::Either
6838 } else {
6839 DeclaredHessianForm::Unavailable
6840 },
6841 prefer_gradient_only,
6842 disable_fixed_point,
6843 seed_risk_profile,
6844 kappa_options.rel_tol.max(1e-6),
6845 kappa_options.max_outer_iter.max(1),
6846 Some(5.0),
6848 Some(kappa_options.log_step.clamp(0.25, 1.0)),
6850 screening_cap.clone(),
6851 Some((n_total, joint_p_cols)),
6854 block_specs
6857 .iter()
6858 .any(|s| !constant_curvature_term_indices(s).is_empty()),
6859 false,
6862 )
6863 .map_err(|e| e.to_string())?;
6864
6865 fn collect_specs(cache: &ExactJointDesignCache<'_>) -> Vec<TermCollectionSpec> {
6867 cache.specs().into_iter().cloned().collect()
6868 }
6869 fn collect_designs(cache: &ExactJointDesignCache<'_>) -> Vec<TermCollectionDesign> {
6870 cache.designs().into_iter().cloned().collect()
6871 }
6872
6873 let result = {
6874 let eval_outer = |ctx: &mut &mut NBlockExactJointState<'_, Mode>,
6875 theta: &Array1<f64>,
6876 order: OuterEvalOrder|
6877 -> Result<OuterEval, EstimationError> {
6878 if let Some((cost, grad, hess)) = ctx.cache.memoized_eval(theta)
6879 && ctx.terminal_mode_matches(theta, cost)
6880 {
6881 let cached_satisfies_order = match order {
6882 OuterEvalOrder::Value => true,
6883 OuterEvalOrder::ValueAndGradient => grad.len() == theta.len(),
6884 OuterEvalOrder::ValueGradientHessian => {
6885 grad.len() == theta.len() && hess.is_analytic()
6886 }
6887 };
6888 if cached_satisfies_order {
6889 if !cost.is_finite() {
6890 return Ok(OuterEval::infeasible(theta.len()));
6891 }
6892 if grad.iter().any(|v| !v.is_finite()) {
6905 return Ok(OuterEval::infeasible(theta.len()));
6906 }
6907 return Ok(OuterEval {
6908 cost,
6909 gradient: grad,
6910 hessian: hess,
6911 inner_beta_hint: None,
6912 });
6913 }
6914 }
6915 ctx.ensure_theta(theta).map_err(|err| {
6916 EstimationError::InvalidInput(format!(
6917 "n-block exact-joint spatial design realization failed: {err}"
6918 ))
6919 })?;
6920 let design_revision = Some(ctx.cache.design_revision());
6921 let specs = collect_specs(&ctx.cache);
6922 let designs = collect_designs(&ctx.cache);
6923 let clamped = outer_derivative_policy.order_for_evaluation(order);
6931 let value_only = matches!(clamped, OuterEvalOrder::Value);
6932 let need_hessian = matches!(clamped, OuterEvalOrder::ValueGradientHessian)
6933 && analytic_outer_hessian_available;
6934 let eval_mode = if value_only {
6935 gam_solve::estimate::reml::reml_outer_engine::EvalMode::ValueOnly
6936 } else if need_hessian {
6937 gam_solve::estimate::reml::reml_outer_engine::EvalMode::ValueGradientHessian
6938 } else {
6939 gam_solve::estimate::reml::reml_outer_engine::EvalMode::ValueAndGradient
6940 };
6941 let t0 = std::time::Instant::now();
6942 let result =
6943 (*exact_fn_cell.borrow_mut())(theta, &specs, &designs, eval_mode, &ctx.row_set);
6944 let elapsed_s = t0.elapsed().as_secs_f64();
6945 kphase_eval_calls.set(kphase_eval_calls.get() + 1);
6946 kphase_eval_total_s.set(kphase_eval_total_s.get() + elapsed_s);
6947 let (theta_norm, log_kappa_norm) = kphase_log_norms(theta);
6948 log::info!(
6949 "[KAPPA-PHASE] phase=eval_outer call={} order={:?} design_revision={:?} theta_norm={:.4e} log_kappa_norm={:.4e} elapsed_s={:.4}",
6950 kphase_eval_calls.get(),
6951 order,
6952 design_revision,
6953 theta_norm,
6954 log_kappa_norm,
6955 elapsed_s,
6956 );
6957 match result {
6958 Ok(ExactJointEvaluation {
6959 objective: cost,
6960 gradient: grad,
6961 hessian: hess,
6962 mode,
6963 }) => {
6964 ctx.install_terminal_mode(theta, cost, mode);
6965 if value_only {
6966 ctx.cache.store_cost_only(theta, cost);
6967 } else {
6968 ctx.cache.store_eval((cost, grad.clone(), hess.clone()));
6969 }
6970 if !cost.is_finite() {
6971 return Ok(OuterEval::infeasible(theta.len()));
6972 }
6973 if grad.iter().any(|v| !v.is_finite()) {
6986 return Ok(OuterEval::infeasible(theta.len()));
6987 }
6988 Ok(OuterEval {
6989 cost,
6990 gradient: grad,
6991 hessian: hess,
6992 inner_beta_hint: None,
6993 })
6994 }
6995 Err(err) => Err(EstimationError::RemlOptimizationFailed(format!(
6996 "n-block exact-joint spatial evaluation failed: {err}"
6997 ))),
6998 }
6999 };
7000
7001 let obj = problem.build_objective_with_eval_order(
7002 &mut state,
7003 |ctx: &mut &mut NBlockExactJointState<'_, Mode>, theta: &Array1<f64>| {
7004 if let Some(cost) = ctx.cache.memoized_cost(theta)
7005 && ctx.terminal_mode_matches(theta, cost)
7006 {
7007 return Ok(cost);
7008 }
7009 ctx.ensure_theta(theta).map_err(|err| {
7010 EstimationError::InvalidInput(format!(
7011 "n-block exact-joint spatial design realization failed: {err}"
7012 ))
7013 })?;
7014 let design_revision = Some(ctx.cache.design_revision());
7015 let specs = collect_specs(&ctx.cache);
7016 let designs = collect_designs(&ctx.cache);
7017 let t0 = std::time::Instant::now();
7024 let result = (*exact_fn_cell.borrow_mut())(
7025 theta,
7026 &specs,
7027 &designs,
7028 gam_solve::estimate::reml::reml_outer_engine::EvalMode::ValueOnly,
7029 &ctx.row_set,
7030 );
7031 let elapsed_s = t0.elapsed().as_secs_f64();
7032 kphase_cost_calls.set(kphase_cost_calls.get() + 1);
7033 kphase_cost_total_s.set(kphase_cost_total_s.get() + elapsed_s);
7034 let (theta_norm, log_kappa_norm) = kphase_log_norms(theta);
7035 log::info!(
7036 "[KAPPA-PHASE] phase=cost call={} design_revision={:?} theta_norm={:.4e} log_kappa_norm={:.4e} elapsed_s={:.4}",
7037 kphase_cost_calls.get(),
7038 design_revision,
7039 theta_norm,
7040 log_kappa_norm,
7041 elapsed_s,
7042 );
7043 match result {
7044 Ok(ExactJointEvaluation {
7045 objective: cost,
7046 mode,
7047 ..
7048 }) => {
7049 ctx.install_terminal_mode(theta, cost, mode);
7050 ctx.cache.store_cost_only(theta, cost);
7056 Ok(cost)
7057 }
7058 Err(err) => Err(EstimationError::RemlOptimizationFailed(format!(
7059 "n-block exact-joint spatial cost evaluation failed: {err}"
7060 ))),
7061 }
7062 },
7063 |ctx: &mut &mut NBlockExactJointState<'_, Mode>, theta: &Array1<f64>| {
7064 eval_outer(
7065 ctx,
7066 theta,
7067 if analytic_outer_hessian_available {
7068 OuterEvalOrder::ValueGradientHessian
7069 } else {
7070 OuterEvalOrder::ValueAndGradient
7071 },
7072 )
7073 },
7074 |ctx: &mut &mut NBlockExactJointState<'_, Mode>,
7075 theta: &Array1<f64>,
7076 order: OuterEvalOrder| { eval_outer(ctx, theta, order) },
7077 None::<fn(&mut &mut NBlockExactJointState<'_, Mode>)>,
7078 Some(
7079 |ctx: &mut &mut NBlockExactJointState<'_, Mode>, theta: &Array1<f64>| {
7080 ctx
7081 .ensure_theta(theta)
7082 .map_err(EstimationError::InvalidInput)?;
7083 let design_revision = Some(ctx.cache.design_revision());
7084 let specs = collect_specs(&ctx.cache);
7085 let designs = collect_designs(&ctx.cache);
7086 let t0 = std::time::Instant::now();
7087 let eval_result = (*exact_efs_fn_cell.borrow_mut())(
7088 theta,
7089 &specs,
7090 &designs,
7091 &ctx.row_set,
7092 );
7093 let elapsed_s = t0.elapsed().as_secs_f64();
7094 kphase_efs_calls.set(kphase_efs_calls.get() + 1);
7095 kphase_efs_total_s.set(kphase_efs_total_s.get() + elapsed_s);
7096 let (theta_norm, log_kappa_norm) = kphase_log_norms(theta);
7097 log::info!(
7098 "[KAPPA-PHASE] phase=efs call={} design_revision={:?} theta_norm={:.4e} log_kappa_norm={:.4e} elapsed_s={:.4}",
7099 kphase_efs_calls.get(),
7100 design_revision,
7101 theta_norm,
7102 log_kappa_norm,
7103 elapsed_s,
7104 );
7105 let ExactJointEfsEvaluation { evaluation, mode } =
7106 eval_result.map_err(EstimationError::RemlOptimizationFailed)?;
7107 ctx.cache.invalidate_objective_memo();
7114 ctx.cache.store_cost_only(theta, evaluation.cost);
7115 ctx.install_terminal_mode(theta, evaluation.cost, mode);
7116 Ok(evaluation)
7117 },
7118 ),
7119 );
7120 let mut obj = obj
7121 .with_seed_inner_state(
7122 move |_ctx: &mut &mut NBlockExactJointState<'_, Mode>, beta: &Array1<f64>| {
7123 (seed_inner_beta_fn)(beta)
7124 },
7125 )
7126 .with_exact_polish(|ctx: &mut &mut NBlockExactJointState<'_, Mode>| {
7127 if !ctx.staged_pilot_active {
7128 return false;
7129 }
7130 ctx.cache.invalidate_objective_memo();
7135 ctx.terminal_mode = None;
7136 ctx.row_set = gam_problem::outer_subsample::RowSet::All;
7137 ctx.staged_pilot_active = false;
7138 true
7139 });
7140
7141 problem
7142 .run_certified(&mut obj, "n-block exact-joint spatial")
7143 .map_err(|error| error.to_string())?
7144 }; let kphase_total_s = kphase_optim_start.elapsed().as_secs_f64();
7154 log::info!(
7155 "[KAPPA-PHASE-SUMMARY] log_kappa_dim={} n_cost={} cost_total_s={:.4} n_eval={} eval_total_s={:.4} n_efs={} efs_total_s={:.4} optim_total_s={:.4}",
7156 kphase_log_kappa_dim,
7157 kphase_cost_calls.get(),
7158 kphase_cost_total_s.get(),
7159 kphase_eval_calls.get(),
7160 kphase_eval_total_s.get(),
7161 kphase_efs_calls.get(),
7162 kphase_efs_total_s.get(),
7163 kphase_total_s,
7164 );
7165 let timing = SpatialLengthScaleOptimizationTiming {
7166 log_kappa_dim: kphase_log_kappa_dim,
7167 cost_calls: kphase_cost_calls.get(),
7168 cost_total_s: kphase_cost_total_s.get(),
7169 eval_calls: kphase_eval_calls.get(),
7170 eval_total_s: kphase_eval_total_s.get(),
7171 efs_calls: kphase_efs_calls.get(),
7172 efs_total_s: kphase_efs_total_s.get(),
7173 slow_path_resets: 0,
7174 design_revision_delta: 0,
7175 nfree_skip_row_touches: 0,
7176 nfree_miss_shape: 0,
7177 nfree_miss_value: 0,
7178 nfree_miss_gradient: 0,
7179 nfree_miss_penalty: 0,
7180 nfree_miss_revision: 0,
7181 nfree_miss_second_order: 0,
7182 nfree_miss_other: 0,
7183 optim_total_s: kphase_total_s,
7184 };
7185
7186 if !matches!(state.row_set, gam_problem::outer_subsample::RowSet::All) {
7187 return Err(
7188 "n-block exact-joint spatial optimization returned before its exact full-data transition"
7189 .to_string(),
7190 );
7191 }
7192 let certified_outer = result;
7193 let theta_star = certified_outer.rho().clone();
7194
7195 state.ensure_theta(&theta_star)?;
7200 let (mode_theta, mode_objective, mode) = state.terminal_mode.take().ok_or_else(|| {
7201 "n-block exact-joint spatial optimization produced a certificate without retaining the owned terminal coefficient mode"
7202 .to_string()
7203 })?;
7204 if !theta_values_match(&mode_theta, &theta_star) {
7205 return Err(
7206 "n-block exact-joint spatial terminal coefficient mode does not bitwise match the certified hyperparameter vector"
7207 .to_string(),
7208 );
7209 }
7210 if mode_objective.to_bits() != certified_outer.final_value().to_bits() {
7211 return Err(format!(
7212 "n-block exact-joint spatial terminal coefficient mode objective does not bitwise match the certified objective: mode={mode_objective:.17e}, certified={:.17e}",
7213 certified_outer.final_value(),
7214 ));
7215 }
7216
7217 let resolved_specs: Vec<TermCollectionSpec> = collect_specs(&state.cache);
7218 let designs: Vec<TermCollectionDesign> = collect_designs(&state.cache);
7219
7220 let fit = fit_fn(
7221 &theta_star,
7222 &resolved_specs,
7223 &designs,
7224 SpatialFitProvenance::Certified {
7225 outer: &certified_outer,
7226 mode,
7227 },
7228 )?;
7229
7230 for spec in &resolved_specs {
7231 log_spatial_aniso_scales(spec);
7232 }
7233
7234 Ok(SpatialLengthScaleOptimizationResult {
7235 resolved_specs,
7236 designs,
7237 fit,
7238 certified_outer: Some(certified_outer),
7239 timing: Some(timing),
7240 })
7241}
7242
7243fn try_exact_joint_latent_coord_optimization(
7244 data: ArrayView2<'_, f64>,
7245 y: ArrayView1<'_, f64>,
7246 weights: ArrayView1<'_, f64>,
7247 offset: ArrayView1<'_, f64>,
7248 resolvedspec: &TermCollectionSpec,
7249 best: &FittedTermCollection,
7250 family: LikelihoodSpec,
7251 options: &FitOptions,
7252 latent: &StandardLatentCoordConfig,
7253) -> Result<FittedTermCollectionWithSpec, EstimationError> {
7254 use gam_problem::{DeclaredHessianForm, Derivative, OuterEval};
7255 use gam_solve::rho_optimizer::OuterEvalOrder;
7256
7257 let rho_dim = best.fit.lambdas.len();
7258 let latent_flat_dim = latent.values.len();
7259 if latent_flat_dim == 0 {
7260 crate::bail_invalid_estim!(
7261 "latent-coordinate optimization requires a non-empty latent block"
7262 );
7263 }
7264 let direct_hypers =
7265 latent_coord_initial_direct_hypers(latent.values.id_mode(), latent.values.latent_dim())?;
7266 let analytic_rho_count = latent
7267 .analytic_penalties
7268 .as_ref()
7269 .map_or(0, |registry| registry.total_rho_count());
7270 let latent_coord_ext_dim = latent_flat_dim + analytic_rho_count + direct_hypers.len();
7271
7272 let mut theta0 = Array1::<f64>::zeros(rho_dim + latent_coord_ext_dim);
7273 theta0
7274 .slice_mut(s![..rho_dim])
7275 .assign(&best.fit.lambdas.mapv(f64::ln));
7276 theta0
7277 .slice_mut(s![rho_dim..rho_dim + latent_flat_dim])
7278 .assign(latent.values.as_flat());
7279 if !direct_hypers.is_empty() {
7280 let direct_start = rho_dim + latent_flat_dim + analytic_rho_count;
7281 theta0
7282 .slice_mut(s![direct_start..direct_start + direct_hypers.len()])
7283 .assign(&direct_hypers);
7284 }
7285
7286 let mut lower = Array1::<f64>::from_elem(theta0.len(), -12.0);
7287 let mut upper = Array1::<f64>::from_elem(theta0.len(), 12.0);
7288 let latent_bound = latent
7289 .values
7290 .as_flat()
7291 .iter()
7292 .fold(1.0_f64, |acc, &v| acc.max(v.abs()))
7293 + 10.0;
7294 for axis in rho_dim..rho_dim + latent_flat_dim {
7295 lower[axis] = -latent_bound;
7296 upper[axis] = latent_bound;
7297 }
7298 if let Some(registry) = latent.analytic_penalties.as_ref() {
7299 let (domain_lower, domain_upper) = registry
7300 .rho_domain_bounds()
7301 .map_err(EstimationError::InvalidInput)?;
7302 let start = rho_dim + latent_flat_dim;
7303 for local in 0..analytic_rho_count {
7304 lower[start + local] = lower[start + local].max(domain_lower[local]);
7305 upper[start + local] = upper[start + local].min(domain_upper[local]);
7306 if lower[start + local] >= upper[start + local] {
7307 return Err(EstimationError::InvalidInput(format!(
7308 "analytic-penalty rho domain has no searchable interval at coordinate {local}: lower={}, upper={}",
7309 lower[start + local],
7310 upper[start + local]
7311 )));
7312 }
7313 }
7314 }
7315
7316 struct LatentJointContext<'d> {
7317 rho_dim: usize,
7318 cache: SingleBlockLatentCoordDesignCache,
7319 evaluator: gam_solve::estimate::ExternalJointHyperEvaluator<'d>,
7320 }
7321
7322 impl<'d> LatentJointContext<'d> {
7323 fn eval_full(
7324 &mut self,
7325 theta: &Array1<f64>,
7326 order: OuterEvalOrder,
7327 ) -> Result<(f64, Array1<f64>, gam_problem::HessianValue), EstimationError> {
7328 if let Some(eval) = self.cache.memoized_eval(theta) {
7329 return Ok(eval);
7330 }
7331 self.cache
7332 .ensure_theta(theta)
7333 .map_err(EstimationError::InvalidInput)?;
7334 let hyper_dirs = self
7335 .cache
7336 .hyper_dirs()
7337 .map_err(EstimationError::InvalidInput)?;
7338 let design_revision = Some(self.cache.design_revision());
7339 let registry_for_key = self.cache.analytic_penalties();
7340 self.evaluator
7341 .set_analytic_penalty_registry(registry_for_key.as_deref());
7342 let mut eval = evaluate_joint_reml_outer_eval_at_theta(
7343 &mut self.evaluator,
7344 self.cache.design(),
7345 theta,
7346 self.rho_dim,
7347 hyper_dirs,
7348 None,
7349 order,
7350 design_revision,
7351 )?;
7352 let latent = self.cache.latent().map_err(EstimationError::InvalidInput)?;
7353 if let Some(registry) = registry_for_key {
7354 let mut registry = registry.as_ref().clone();
7355 registry.apply_weight_schedules(
7356 gam_solve::estimate::reml::outer_eval::current_outer_iter() as usize,
7357 );
7358 add_analytic_penalty_objective_to_eval(
7359 theta,
7360 self.rho_dim,
7361 latent.as_ref(),
7362 ®istry,
7363 &mut eval,
7364 )?;
7365 }
7366 add_latent_id_objective_to_eval(
7367 theta,
7368 self.rho_dim,
7369 self.cache.analytic_penalty_rho_count(),
7370 latent.as_ref(),
7371 &mut eval,
7372 )?;
7373 self.cache.store_eval(eval.clone());
7374 Ok(eval)
7375 }
7376
7377 fn eval_efs(
7378 &mut self,
7379 theta: &Array1<f64>,
7380 ) -> Result<gam_problem::EfsEval, EstimationError> {
7381 self.cache
7382 .ensure_theta(theta)
7383 .map_err(EstimationError::InvalidInput)?;
7384 let hyper_dirs = self
7385 .cache
7386 .hyper_dirs()
7387 .map_err(EstimationError::InvalidInput)?;
7388 let registry_for_key = self.cache.analytic_penalties();
7389 self.evaluator
7390 .set_analytic_penalty_registry(registry_for_key.as_deref());
7391 let mut efs = evaluate_joint_reml_efs_at_theta(
7392 &mut self.evaluator,
7393 self.cache.design(),
7394 theta,
7395 self.rho_dim,
7396 hyper_dirs,
7397 None,
7398 Some(self.cache.design_revision()),
7399 )?;
7400 if let Some(registry) = registry_for_key {
7401 let mut registry = registry.as_ref().clone();
7402 registry.apply_weight_schedules(
7403 gam_solve::estimate::reml::outer_eval::current_outer_iter() as usize,
7404 );
7405 let latent = self.cache.latent().map_err(EstimationError::InvalidInput)?;
7406 let contribution = analytic_penalty_objective_contribution(
7407 theta,
7408 self.rho_dim,
7409 latent.as_ref(),
7410 ®istry,
7411 )?;
7412 efs.cost += contribution.cost;
7413 if let (Some(psi_gradient), Some(psi_indices)) =
7414 (efs.psi_gradient.as_mut(), efs.psi_indices.as_ref())
7415 {
7416 if psi_gradient.len() != psi_indices.len() {
7417 crate::bail_invalid_estim!(
7418 "latent-coordinate analytic penalty EFS psi gradient length mismatch: gradient={}, indices={}",
7419 psi_gradient.len(),
7420 psi_indices.len()
7421 );
7422 }
7423 for (local_idx, &theta_idx) in psi_indices.iter().enumerate() {
7424 psi_gradient[local_idx] += contribution.gradient[theta_idx];
7425 }
7426 }
7427 }
7428 Ok(efs)
7429 }
7430
7431 fn eval_cost(&mut self, theta: &Array1<f64>) -> f64 {
7432 if let Some(cost) = self.cache.memoized_cost(theta) {
7433 return cost;
7434 }
7435 if self.cache.ensure_theta(theta).is_err() {
7436 return f64::INFINITY;
7437 }
7438 let design_revision = Some(self.cache.design_revision());
7439 let registry_for_key = self.cache.analytic_penalties();
7440 self.evaluator
7441 .set_analytic_penalty_registry(registry_for_key.as_deref());
7442 let result = {
7443 let design = self.cache.design();
7444 self.evaluator.evaluate_cost_only(
7445 &design.design,
7446 &design.penalties,
7447 &design.nullspace_dims,
7448 design.linear_constraints.clone(),
7449 theta,
7450 self.rho_dim,
7451 None,
7452 "latent-coordinate-joint cost-only",
7453 design_revision,
7454 )
7455 };
7456 match result {
7457 Ok(cost) => {
7458 let latent = match self.cache.latent() {
7459 Ok(latent) => latent,
7460 Err(_) => return f64::INFINITY,
7461 };
7462 let contribution = match latent_id_objective_contribution(
7463 theta,
7464 self.rho_dim,
7465 self.cache.analytic_penalty_rho_count(),
7466 latent.as_ref(),
7467 ) {
7468 Ok(contribution) => contribution,
7469 Err(_) => return f64::INFINITY,
7470 };
7471 let cost = cost + contribution.cost;
7472 let cost = if let Some(registry) = registry_for_key {
7473 let mut registry = registry.as_ref().clone();
7474 registry.apply_weight_schedules(
7475 gam_solve::estimate::reml::outer_eval::current_outer_iter() as usize,
7476 );
7477 match analytic_penalty_objective_contribution(
7478 theta,
7479 self.rho_dim,
7480 latent.as_ref(),
7481 ®istry,
7482 ) {
7483 Ok(contribution) => cost + contribution.cost,
7484 Err(_) => return f64::INFINITY,
7485 }
7486 } else {
7487 cost
7488 };
7489 self.cache.store_cost(cost);
7490 cost
7491 }
7492 Err(_) => f64::INFINITY,
7493 }
7494 }
7495 }
7496
7497 let effective_offset = best
7498 .design
7499 .compose_offset(offset, "latent-coordinate joint fit")
7500 .map_err(EstimationError::BasisError)?;
7501 let mut ctx = LatentJointContext {
7502 rho_dim,
7503 cache: SingleBlockLatentCoordDesignCache::new(
7504 data.to_owned(),
7505 resolvedspec.clone(),
7506 best.design.clone(),
7507 latent,
7508 rho_dim,
7509 )
7510 .map_err(EstimationError::InvalidInput)?,
7511 evaluator: gam_solve::estimate::ExternalJointHyperEvaluator::new(
7512 y,
7513 weights,
7514 &best.design.design,
7515 effective_offset.view(),
7516 &best.design.penalties,
7517 &external_opts_for_design(&family, &best.design, options),
7518 "latent-coordinate-joint",
7519 )?,
7520 };
7521 let registry_for_key = ctx.cache.analytic_penalties();
7522 ctx.evaluator
7523 .set_analytic_penalty_registry(registry_for_key.as_deref());
7524 ctx.evaluator
7525 .set_persistent_latent_values_fingerprint(latent.values.id_mode());
7526 if let Some(cached_t) = ctx
7527 .evaluator
7528 .load_persistent_latent_values(latent.values.n_obs(), latent.values.latent_dim())
7529 {
7530 let cached_t: Array2<f64> = cached_t;
7531 for (dst, src) in theta0
7532 .slice_mut(s![rho_dim..rho_dim + latent_flat_dim])
7533 .iter_mut()
7534 .zip(cached_t.iter())
7535 {
7536 *dst = *src;
7537 }
7538 }
7539
7540 let problem = exact_joint_multistart_outer_problem(
7541 &theta0,
7542 &lower,
7543 &upper,
7544 rho_dim,
7545 latent_coord_ext_dim,
7546 theta0.len(),
7547 Derivative::Analytic,
7548 DeclaredHessianForm::Unavailable,
7549 false,
7550 false,
7551 seed_risk_profile_for_likelihood_family(&family),
7552 options.tol,
7553 options.max_iter.max(1),
7554 Some(5.0),
7555 Some(0.5),
7556 None,
7557 Some((data.nrows(), best.design.design.ncols().max(1))),
7560 !constant_curvature_term_indices(resolvedspec).is_empty(),
7563 false,
7565 )?;
7566
7567 let eval_outer = |ctx: &mut &mut LatentJointContext<'_>,
7568 theta: &Array1<f64>,
7569 order: OuterEvalOrder|
7570 -> Result<OuterEval, EstimationError> {
7571 let (cost, gradient, hessian) = ctx.eval_full(theta, order)?;
7572 Ok(OuterEval {
7573 cost,
7574 gradient,
7575 hessian,
7576 inner_beta_hint: None,
7577 })
7578 };
7579
7580 let result = {
7581 let mut obj = problem.build_objective_with_eval_order(
7582 &mut ctx,
7583 |ctx: &mut &mut LatentJointContext<'_>, theta: &Array1<f64>| Ok(ctx.eval_cost(theta)),
7584 |ctx: &mut &mut LatentJointContext<'_>, theta: &Array1<f64>| {
7585 eval_outer(ctx, theta, OuterEvalOrder::ValueAndGradient)
7586 },
7587 |ctx: &mut &mut LatentJointContext<'_>, theta: &Array1<f64>, order: OuterEvalOrder| {
7588 eval_outer(ctx, theta, order)
7589 },
7590 Some(|ctx: &mut &mut LatentJointContext<'_>| {
7591 ctx.cache.reset();
7592 }),
7593 Some(|ctx: &mut &mut LatentJointContext<'_>, theta: &Array1<f64>| ctx.eval_efs(theta)),
7594 );
7595
7596 problem
7597 .run(&mut obj, "latent-coordinate joint REML")
7598 .map_err(|e| {
7599 EstimationError::InvalidInput(format!(
7600 "latent-coordinate joint optimization failed after exhausting strategy fallbacks: {e}"
7601 ))
7602 })?
7603 };
7604 if !result.converged {
7605 crate::bail_invalid_estim!(
7606 "latent-coordinate joint optimization did not converge after {} iterations (final_objective={:.6e}, final_grad_norm={})",
7607 result.iterations,
7608 result.final_value,
7609 result.final_grad_norm_report(),
7610 );
7611 }
7612
7613 let theta_star = result.rho;
7614 let selected_lambdas = Array1::from_vec(
7615 gam_problem::checked_exp_log_strengths(
7616 theta_star.slice(s![..rho_dim]).iter().copied(),
7617 )
7618 .map_err(|error| {
7619 EstimationError::InvalidInput(format!(
7620 "selected latent-coordinate smoothing coordinate is outside the canonical log-strength domain: {error}"
7621 ))
7622 })?,
7623 );
7624 let mut final_data = data.to_owned();
7625 let flat_t = theta_star
7626 .slice(s![rho_dim..rho_dim + latent_flat_dim])
7627 .to_owned();
7628 let mut fitted_latent_values =
7629 Array2::<f64>::zeros((latent.values.n_obs(), latent.values.latent_dim()));
7630 for n in 0..latent.values.n_obs() {
7631 for axis in 0..latent.values.latent_dim() {
7632 let value = flat_t[n * latent.values.latent_dim() + axis];
7633 fitted_latent_values[[n, axis]] = value;
7634 final_data[[n, latent.feature_cols[axis]]] = value;
7635 }
7636 }
7637 let optimized = fit_term_collection_forspecwith_heuristic_lambdas(
7638 final_data.view(),
7639 y,
7640 weights,
7641 offset,
7642 resolvedspec,
7643 selected_lambdas.as_slice(),
7644 family,
7645 options,
7646 )?;
7647 ctx.evaluator
7648 .store_persistent_latent_values(&fitted_latent_values);
7649 let mut fit = optimized.fit;
7650 fit.reml_score = result.final_value;
7651 fit.penalized_objective = result.final_value;
7652 Ok(FittedTermCollectionWithSpec {
7653 fit,
7654 design: optimized.design,
7655 resolvedspec: resolvedspec.clone(),
7656 adaptive_diagnostics: optimized.adaptive_diagnostics,
7657 kappa_timing: None,
7658 })
7659}
7660
7661pub fn fit_term_collectionwith_latent_coord_optimization(
7662 data: ArrayView2<'_, f64>,
7663 y: Array1<f64>,
7664 weights: Array1<f64>,
7665 offset: Array1<f64>,
7666 spec: &TermCollectionSpec,
7667 latent: &StandardLatentCoordConfig,
7668 family: LikelihoodSpec,
7669 options: &FitOptions,
7670) -> Result<FittedTermCollectionWithSpec, EstimationError> {
7671 let n = data.nrows();
7672 if !(y.len() == n && weights.len() == n && offset.len() == n) {
7673 crate::bail_invalid_estim!(
7674 "fit_term_collectionwith_latent_coord_optimization row mismatch: n={}, y={}, weights={}, offset={}",
7675 n,
7676 y.len(),
7677 weights.len(),
7678 offset.len()
7679 );
7680 }
7681 let best = fit_term_collection_forspec(
7682 data,
7683 y.view(),
7684 weights.view(),
7685 offset.view(),
7686 spec,
7687 family.clone(),
7688 options,
7689 )?;
7690 let resolvedspec = freeze_term_collection_from_design(spec, &best.design)?;
7691 try_exact_joint_latent_coord_optimization(
7692 data,
7693 y.view(),
7694 weights.view(),
7695 offset.view(),
7696 &resolvedspec,
7697 &best,
7698 family,
7699 options,
7700 latent,
7701 )
7702}
7703
7704fn select_isotropic_matern_range_basin(
7721 data: ArrayView2<'_, f64>,
7722 y: ArrayView1<'_, f64>,
7723 weights: ArrayView1<'_, f64>,
7724 offset: ArrayView1<'_, f64>,
7725 mut resolvedspec: TermCollectionSpec,
7726 mut best: FittedTermCollection,
7727 family: &LikelihoodSpec,
7728 options: &FitOptions,
7729 kappa_options: &SpatialLengthScaleOptimizationOptions,
7730 spatial_terms: &[usize],
7731) -> Result<(TermCollectionSpec, FittedTermCollection), EstimationError> {
7732 if has_aniso_terms(&resolvedspec, spatial_terms)
7736 || !constant_curvature_term_indices(&resolvedspec).is_empty()
7737 {
7738 return Ok((resolvedspec, best));
7739 }
7740
7741 let mut best_score = fit_score(&best.fit);
7742 if !best_score.is_finite() {
7743 crate::bail_invalid_estim!(
7744 "isotropic Matérn basin selection received a non-finite incumbent profile"
7745 );
7746 }
7747
7748 for &term_idx in spatial_terms {
7749 let Some(SmoothBasisSpec::Matern {
7750 feature_cols,
7751 spec: matern,
7752 ..
7753 }) = resolvedspec
7754 .smooth_terms
7755 .get(term_idx)
7756 .map(|term| &term.basis)
7757 else {
7758 continue;
7759 };
7760 let num_centers = gam_terms::basis::center_strategy_num_centers(&matern.center_strategy)
7761 .ok_or_else(|| {
7762 EstimationError::InvalidInput(format!(
7763 "resolved isotropic Matérn term {term_idx} has no finite center count"
7764 ))
7765 })?;
7766 let companion_length_scale = matern_low_rank_center_resolution_length_scale(
7767 data,
7768 feature_cols,
7769 num_centers,
7770 )
7771 .ok_or_else(|| {
7772 EstimationError::InvalidInput(format!(
7773 "resolved isotropic Matérn term {term_idx} has no finite center-resolution range"
7774 ))
7775 })?;
7776 let (psi_long_bound, psi_short_bound) =
7777 spatial_term_psi_bounds(data, &resolvedspec, term_idx, kappa_options)
7778 .map_err(EstimationError::BasisError)?;
7779 let psi_long = (-companion_length_scale.ln()).clamp(psi_long_bound, psi_short_bound);
7780 let long_length_scale = (-psi_long).exp();
7781 if !(long_length_scale.is_finite() && long_length_scale > 0.0) {
7782 crate::bail_invalid_estim!(
7783 "isotropic Matérn term {term_idx} produced an invalid long-range endpoint from psi={psi_long}"
7784 );
7785 }
7786 if get_spatial_length_scale(&resolvedspec, term_idx)
7787 .is_some_and(|current| current == long_length_scale)
7788 {
7789 continue;
7790 }
7791
7792 let mut endpoint_spec = resolvedspec.clone();
7793 set_spatial_length_scale(&mut endpoint_spec, term_idx, long_length_scale)?;
7794 let endpoint = fit_term_collection_forspecwith_heuristic_lambdas(
7804 data,
7805 y,
7806 weights,
7807 offset,
7808 &endpoint_spec,
7809 best.fit.lambdas.as_slice(),
7810 family.clone(),
7811 options,
7812 )?;
7813 let endpoint_score = fit_score(&endpoint.fit);
7814 if !endpoint_score.is_finite() {
7815 crate::bail_invalid_estim!(
7816 "isotropic Matérn term {term_idx} long-range endpoint returned a non-finite profiled REML score"
7817 );
7818 }
7819
7820 if endpoint_score < best_score {
7821 log::info!(
7822 "[spatial-kappa] term {term_idx} selected certified long-range basin: \
7823 length_scale={long_length_scale:.6}, profiled REML {endpoint_score:.6} \
7824 < short-basin {best_score:.6}"
7825 );
7826 resolvedspec = freeze_term_collection_from_design(&endpoint_spec, &endpoint.design)?;
7827 best = endpoint;
7828 best_score = endpoint_score;
7829 } else {
7830 log::info!(
7831 "[spatial-kappa] term {term_idx} retained certified short-range basin: \
7832 profiled REML {best_score:.6} <= long-endpoint {endpoint_score:.6} \
7833 at length_scale={long_length_scale:.6}"
7834 );
7835 }
7836 }
7837
7838 Ok((resolvedspec, best))
7839}
7840
7841pub fn fit_term_collectionwith_spatial_length_scale_optimization(
7842 data: ArrayView2<'_, f64>,
7843 y: Array1<f64>,
7844 weights: Array1<f64>,
7845 offset: Array1<f64>,
7846 spec: &TermCollectionSpec,
7847 family: LikelihoodSpec,
7848 options: &FitOptions,
7849 kappa_options: &SpatialLengthScaleOptimizationOptions,
7850) -> Result<FittedTermCollectionWithSpec, EstimationError> {
7851 let mut resolvedspec = spec.clone();
7867 let spatial_terms = spatial_length_scale_term_indices(&resolvedspec);
7868 let n = data.nrows();
7869 if !(y.len() == n && weights.len() == n && offset.len() == n) {
7870 crate::bail_invalid_estim!(
7871 "fit_term_collectionwith_spatial_length_scale_optimization row mismatch: n={}, y={}, weights={}, offset={}",
7872 n,
7873 y.len(),
7874 weights.len(),
7875 offset.len()
7876 );
7877 }
7878 if !kappa_options.enabled || spatial_terms.is_empty() {
7879 let out = fit_term_collection_forspec(
7880 data,
7881 y.view(),
7882 weights.view(),
7883 offset.view(),
7884 &resolvedspec,
7885 family,
7886 options,
7887 )?;
7888 let resolvedspec = freeze_term_collection_from_design(&resolvedspec, &out.design)?;
7889 return Ok(FittedTermCollectionWithSpec {
7890 fit: out.fit,
7891 design: out.design,
7892 resolvedspec,
7893 adaptive_diagnostics: out.adaptive_diagnostics,
7894 kappa_timing: None,
7895 });
7896 }
7897 if kappa_options.max_outer_iter == 0 {
7898 crate::bail_invalid_estim!("spatial kappa optimization requires max_outer_iter >= 1");
7899 }
7900 if !(kappa_options.log_step.is_finite() && kappa_options.log_step > 0.0) {
7901 crate::bail_invalid_estim!("spatial kappa optimization requires log_step > 0");
7902 }
7903 if !(kappa_options.min_length_scale.is_finite()
7904 && kappa_options.max_length_scale.is_finite()
7905 && kappa_options.min_length_scale > 0.0
7906 && kappa_options.max_length_scale >= kappa_options.min_length_scale)
7907 {
7908 crate::bail_invalid_estim!(
7909 "spatial kappa optimization requires valid positive length_scale bounds"
7910 );
7911 }
7912
7913 let pilot_threshold = kappa_options.pilot_subsample_threshold;
7914 if pilot_threshold > 0 && n > pilot_threshold * 2 {
7915 log::info!(
7916 "[spatial-kappa] n={n} exceeds pilot threshold {}; using pilot geometry only for deterministic anisotropy initialization",
7917 pilot_threshold * 2,
7918 );
7919 apply_spatial_anisotropy_pilot_initializer(
7920 data,
7921 &mut resolvedspec,
7922 &spatial_terms,
7923 pilot_threshold,
7924 kappa_options,
7925 )?;
7926 }
7927
7928 apply_response_aware_anisotropy_seed(data, y.view(), &mut resolvedspec, &spatial_terms);
7937
7938 let free_curvature_terms: Vec<usize> = constant_curvature_term_indices(&resolvedspec)
7942 .into_iter()
7943 .filter(|&term_idx| !constant_curvature_kappa_is_fixed(&resolvedspec, term_idx))
7944 .collect();
7945 if !free_curvature_terms.is_empty() {
7946 validate_constant_curvature_fair_profile_inputs(weights.view(), offset.view(), &family)?;
7947 }
7948 for term_idx in free_curvature_terms {
7949 let kappa_hat = constant_curvature_kappa_fair_optimum(
7950 data,
7951 y.view(),
7952 &resolvedspec,
7953 term_idx,
7954 options,
7955 )?;
7956 if let Some(SmoothBasisSpec::ConstantCurvature { spec: cc, .. }) = resolvedspec
7957 .smooth_terms
7958 .get_mut(term_idx)
7959 .map(|term| &mut term.basis)
7960 {
7961 cc.kappa = kappa_hat;
7962 }
7963 }
7964
7965 let baseline_options = superseded_fit_options(options);
7966 let best = fit_term_collection_forspec(
7967 data,
7968 y.view(),
7969 weights.view(),
7970 offset.view(),
7971 &resolvedspec,
7972 family.clone(),
7973 &baseline_options,
7974 )?;
7975 resolvedspec = freeze_term_collection_from_design(&resolvedspec, &best.design)?;
7976 let spatial_terms = spatial_length_scale_term_indices(&resolvedspec);
7986 let (next_spec, best) = select_isotropic_matern_range_basin(
7987 data,
7988 y.view(),
7989 weights.view(),
7990 offset.view(),
7991 resolvedspec,
7992 best,
7993 &family,
7994 &baseline_options,
7995 kappa_options,
7996 &spatial_terms,
7997 )?;
7998 resolvedspec = next_spec;
7999 sync_aniso_contrasts_from_metadata(&mut resolvedspec, &best.design.smooth);
8003 if spatial_terms.is_empty() {
8004 let fitted = fit_term_collection_forspecwith_heuristic_lambdas(
8005 data,
8006 y.view(),
8007 weights.view(),
8008 offset.view(),
8009 &resolvedspec,
8010 best.fit.lambdas.as_slice(),
8011 family,
8012 options,
8013 )?;
8014 return Ok(FittedTermCollectionWithSpec {
8015 fit: fitted.fit,
8016 design: fitted.design,
8017 resolvedspec,
8018 adaptive_diagnostics: fitted.adaptive_diagnostics,
8019 kappa_timing: None,
8020 });
8021 }
8022 let initial_score = fit_score(&best.fit);
8023 if !initial_score.is_finite() {
8024 crate::bail_invalid_estim!(
8025 "spatial kappa optimization received a non-finite initial profiled score"
8026 );
8027 }
8028 let exact_joint = try_exact_joint_spatial_length_scale_optimization(
8029 data,
8030 y.view(),
8031 weights.view(),
8032 offset.view(),
8033 &resolvedspec,
8034 &best,
8035 family.clone(),
8036 options,
8037 kappa_options,
8038 &spatial_terms,
8039 )?
8040 .ok_or_else(|| {
8041 EstimationError::RemlOptimizationFailed(
8042 "spatial kappa optimization is unavailable for one or more eligible spatial terms"
8043 .to_string(),
8044 )
8045 })?;
8046 let exact_score = fit_score(&exact_joint.fit);
8047 let exact_joint = require_successful_spatial_optimization_result(
8048 initial_score,
8049 Ok(Some((exact_joint, exact_score))),
8050 )?;
8051
8052 log_spatial_aniso_scales(&exact_joint.resolvedspec);
8053 Ok(exact_joint)
8054}
8055
8056#[derive(Clone, Debug)]
8062pub struct CurvatureInference {
8063 pub term_idx: usize,
8065 pub kappa_hat: f64,
8068 pub ci: gam_geometry::curvature_estimand::KappaProfileCi,
8070 pub flatness: gam_geometry::curvature_estimand::FlatnessTest,
8074}
8075
8076fn curvature_profile_lr_endpoint<F>(
8088 profile: &mut F,
8089 kappa_hat: f64,
8090 value_hat: f64,
8091 bound: f64,
8092 half_threshold: f64,
8093 x_tolerance: f64,
8094 score_tolerance: f64,
8095) -> Result<(f64, bool), String>
8096where
8097 F: FnMut(f64) -> Result<(f64, f64), String>,
8098{
8099 let direction = (bound - kappa_hat).signum();
8100 let span = (bound - kappa_hat).abs();
8101 if direction == 0.0 || span <= x_tolerance {
8102 return Ok((bound, true));
8103 }
8104
8105 let (bound_value, bound_score) = profile(bound)?;
8106 let outward_score = direction * bound_score;
8107 if outward_score < -score_tolerance {
8108 return Err(format!(
8109 "curvature profile is not outward-monotone at chart bound {bound}: \
8110 outward score {outward_score:.6e} is below tolerance {score_tolerance:.6e}"
8111 ));
8112 }
8113 let value_tolerance = score_tolerance * span;
8114 if bound_value < value_hat - value_tolerance {
8115 return Err(format!(
8116 "fitted curvature is not the minimum of its inference profile: \
8117 V(bound={bound})={bound_value:.6e} < V(kappa_hat)={value_hat:.6e}"
8118 ));
8119 }
8120 let bound_residual = bound_value - value_hat - half_threshold;
8121 if bound_residual < 0.0 {
8122 return Ok((bound, true));
8123 }
8124 if bound_residual == 0.0 {
8125 return Ok((bound, false));
8126 }
8127
8128 let mut inside_x = kappa_hat;
8133 let mut outside_x = bound;
8134 let mut outside_residual = bound_residual;
8135 let mut outside_score = bound_score;
8136 while (outside_x - inside_x).abs() > x_tolerance {
8137 let lo = inside_x.min(outside_x);
8138 let hi = inside_x.max(outside_x);
8139 let width = hi - lo;
8140 let central_lo = lo + 0.25 * width;
8141 let central_hi = hi - 0.25 * width;
8142 let newton = outside_x - outside_residual / outside_score;
8143 let probe = if newton.is_finite() && newton > central_lo && newton < central_hi {
8144 newton
8145 } else {
8146 lo + 0.5 * width
8147 };
8148 if !(probe > lo && probe < hi) {
8149 break;
8150 }
8151 let (value, score) = profile(probe)?;
8152 let outward_score = direction * score;
8153 if outward_score < -score_tolerance {
8154 return Err(format!(
8155 "curvature profile changed direction before its likelihood crossing at \
8156 kappa={probe}: outward score {outward_score:.6e} is below tolerance \
8157 {score_tolerance:.6e}"
8158 ));
8159 }
8160 let residual = value - value_hat - half_threshold;
8161 if residual >= 0.0 {
8162 outside_x = probe;
8163 outside_residual = residual;
8164 outside_score = score;
8165 } else {
8166 inside_x = probe;
8167 }
8168 }
8169 Ok((inside_x + 0.5 * (outside_x - inside_x), false))
8170}
8171
8172fn curvature_profile_ci_from_analytic_score<F>(
8173 profile: &mut F,
8174 kappa_hat: f64,
8175 kappa_min: f64,
8176 kappa_max: f64,
8177 level: f64,
8178 relative_tolerance: f64,
8179) -> Result<gam_geometry::curvature_estimand::KappaProfileCi, String>
8180where
8181 F: FnMut(f64) -> Result<(f64, f64), String>,
8182{
8183 if !(kappa_min < kappa_max && kappa_hat >= kappa_min && kappa_hat <= kappa_max) {
8184 return Err("curvature profile requires kappa_hat inside valid chart bounds".to_string());
8185 }
8186 if !(level > 0.0 && level < 1.0) {
8187 return Err("curvature profile level must lie in (0, 1)".to_string());
8188 }
8189 let z = gam_geometry::curvature_estimand::wald_half_width(1.0, level)
8190 .ok_or_else(|| "curvature profile threshold is not finite".to_string())?;
8191 let half_threshold = 0.5 * z * z;
8192 let (value_hat, score_hat) = profile(kappa_hat)?;
8193 let relative_tolerance = relative_tolerance.max(f64::EPSILON.sqrt());
8194 let x_tolerance = relative_tolerance * (1.0 + kappa_min.abs().max(kappa_max.abs()));
8195 let score_tolerance = relative_tolerance * (1.0 + value_hat.abs());
8196 let at_lower = (kappa_hat - kappa_min).abs() <= x_tolerance;
8197 let at_upper = (kappa_hat - kappa_max).abs() <= x_tolerance;
8198 let stationary = if at_lower {
8199 score_hat >= -score_tolerance
8200 } else if at_upper {
8201 score_hat <= score_tolerance
8202 } else {
8203 score_hat.abs() <= score_tolerance
8204 };
8205 if !stationary {
8206 return Err(format!(
8207 "curvature inference rejected a non-stationary point estimate: \
8208 kappa_hat={kappa_hat}, score={score_hat:.6e}, \
8209 stationarity_bound={score_tolerance:.6e}"
8210 ));
8211 }
8212
8213 let (ci_lo, lo_at_bound) = curvature_profile_lr_endpoint(
8214 profile,
8215 kappa_hat,
8216 value_hat,
8217 kappa_min,
8218 half_threshold,
8219 x_tolerance,
8220 score_tolerance,
8221 )?;
8222 let (ci_hi, hi_at_bound) = curvature_profile_lr_endpoint(
8223 profile,
8224 kappa_hat,
8225 value_hat,
8226 kappa_max,
8227 half_threshold,
8228 x_tolerance,
8229 score_tolerance,
8230 )?;
8231 let verdict = if ci_lo > 0.0 {
8232 gam_geometry::curvature_estimand::CurvatureVerdict::Spherical
8233 } else if ci_hi < 0.0 {
8234 gam_geometry::curvature_estimand::CurvatureVerdict::Hyperbolic
8235 } else {
8236 gam_geometry::curvature_estimand::CurvatureVerdict::Flat
8237 };
8238 Ok(gam_geometry::curvature_estimand::KappaProfileCi {
8239 kappa_hat,
8240 ci_lo,
8241 ci_hi,
8242 lo_at_bound,
8243 hi_at_bound,
8244 verdict,
8245 })
8246}
8247
8248pub fn curvature_inference_forspec(
8249 data: ArrayView2<'_, f64>,
8250 y: ArrayView1<'_, f64>,
8251 weights: ArrayView1<'_, f64>,
8252 offset: ArrayView1<'_, f64>,
8253 resolvedspec: &TermCollectionSpec,
8254 term_idx: usize,
8255 family: LikelihoodSpec,
8256 options: &FitOptions,
8257 level: f64,
8258) -> Result<CurvatureInference, EstimationError> {
8259 let kappa_hat = get_constant_curvature_kappa(resolvedspec, term_idx).ok_or_else(|| {
8260 EstimationError::InvalidInput(format!(
8261 "curvature_inference_forspec: term {term_idx} is not a constant-curvature smooth"
8262 ))
8263 })?;
8264 if constant_curvature_kappa_is_fixed(resolvedspec, term_idx) {
8265 crate::bail_invalid_estim!(
8266 "curvature inference requires an estimated curvature; term {term_idx} has user-pinned kappa={kappa_hat}"
8267 );
8268 }
8269 if y.len() != data.nrows() || weights.len() != data.nrows() || offset.len() != data.nrows() {
8270 crate::bail_invalid_estim!(
8271 "curvature inference row mismatch: data={}, y={}, weights={}, offset={}",
8272 data.nrows(),
8273 y.len(),
8274 weights.len(),
8275 offset.len(),
8276 );
8277 }
8278 validate_constant_curvature_fair_profile_inputs(weights, offset, &family)?;
8279 let (kappa_min, kappa_max) = constant_curvature_kappa_bounds(data, resolvedspec, term_idx);
8280 let (feature_cols, base_spec) = match resolvedspec
8281 .smooth_terms
8282 .get(term_idx)
8283 .map(|term| &term.basis)
8284 {
8285 Some(SmoothBasisSpec::ConstantCurvature {
8286 feature_cols, spec, ..
8287 }) => (feature_cols, spec.clone()),
8288 _ => {
8289 return Err(EstimationError::InvalidInput(format!(
8290 "constant-curvature κ profile: smooth term {term_idx} is not a \
8291 constant-curvature basis"
8292 )));
8293 }
8294 };
8295 let x_term = select_columns(data, feature_cols).map_err(EstimationError::from)?;
8296 let radial_reference = constant_curvature_radial_reference(x_term.view(), y)?;
8297 let fair_profile = ConstantCurvatureFairProfile {
8298 data: x_term.view(),
8299 response: y,
8300 radial_reference,
8301 spec: base_spec,
8302 cache: std::cell::RefCell::new(std::collections::HashMap::new()),
8303 };
8304
8305 let mut v_p = |kappa: f64| -> Result<(f64, f64), String> {
8308 if !kappa.is_finite() {
8309 return Err(format!("V_p probed a non-finite κ = {kappa}"));
8310 }
8311 let sample = fair_profile.evaluate(kappa).map_err(|error| {
8312 format!("analytic curvature profile at kappa={kappa} failed: {error}")
8313 })?;
8314 Ok(sample)
8315 };
8316 let ci = curvature_profile_ci_from_analytic_score(
8317 &mut v_p,
8318 kappa_hat,
8319 kappa_min,
8320 kappa_max,
8321 level,
8322 options.tol,
8323 )
8324 .map_err(EstimationError::RemlOptimizationFailed)?;
8325 let flatness = gam_geometry::curvature_estimand::flatness_lr_test(
8326 |kappa| v_p(kappa).map(|(value, _)| value),
8327 kappa_hat,
8328 )
8329 .map_err(EstimationError::RemlOptimizationFailed)?;
8330
8331 Ok(CurvatureInference {
8332 term_idx,
8333 kappa_hat,
8334 ci,
8335 flatness,
8336 })
8337}
8338
8339#[cfg(test)]
8340mod curvature_profile_score_tests {
8341 use super::*;
8342
8343 #[test]
8344 fn analytic_profile_score_finds_exact_quadratic_lr_crossings() {
8345 let kappa_hat = -0.37;
8346 let curvature = 16.0;
8347 let level = 0.95;
8348 let mut profile = |kappa: f64| -> Result<(f64, f64), String> {
8349 let displacement = kappa - kappa_hat;
8350 Ok((
8351 7.0 + 0.5 * curvature * displacement * displacement,
8352 curvature * displacement,
8353 ))
8354 };
8355 let ci = curvature_profile_ci_from_analytic_score(
8356 &mut profile,
8357 kappa_hat,
8358 -3.0,
8359 3.0,
8360 level,
8361 1.0e-10,
8362 )
8363 .expect("analytic quadratic profile CI");
8364 let z = gam_geometry::curvature_estimand::wald_half_width(1.0, level)
8365 .expect("valid normal quantile");
8366 let expected_half_width = z / curvature.sqrt();
8367 assert!((ci.ci_lo - (kappa_hat - expected_half_width)).abs() <= 1.0e-8);
8368 assert!((ci.ci_hi - (kappa_hat + expected_half_width)).abs() <= 1.0e-8);
8369 assert!(!ci.lo_at_bound && !ci.hi_at_bound);
8370 }
8371
8372 #[test]
8373 fn analytic_profile_marks_chart_bound_when_wilks_set_never_crosses() {
8374 let mut profile =
8375 |kappa: f64| -> Result<(f64, f64), String> { Ok((0.5 * kappa * kappa, kappa)) };
8376 let ci =
8377 curvature_profile_ci_from_analytic_score(&mut profile, 0.0, -0.1, 0.1, 0.95, 1.0e-10)
8378 .expect("open bounded profile CI");
8379 assert_eq!(ci.ci_lo, -0.1);
8380 assert_eq!(ci.ci_hi, 0.1);
8381 assert!(ci.lo_at_bound && ci.hi_at_bound);
8382 }
8383}
8384
8385#[derive(Clone, Copy, Debug, PartialEq, Eq)]
8388pub enum SmoothLrCorrection {
8389 LawleyLrEstimatedLambda,
8393 LawleyLrFixedLambda,
8398 None,
8402}
8403
8404impl SmoothLrCorrection {
8405 pub fn label(self) -> &'static str {
8407 match self {
8408 SmoothLrCorrection::LawleyLrEstimatedLambda => "lawley_lr_estimated_lambda",
8409 SmoothLrCorrection::LawleyLrFixedLambda => "lawley_lr_fixed_lambda",
8410 SmoothLrCorrection::None => "none",
8411 }
8412 }
8413}
8414
8415#[derive(Clone, Debug)]
8421pub struct SmoothTermLrInference {
8422 pub name: String,
8424 pub term_idx: usize,
8426 pub statistic_lr: f64,
8429 pub ref_df: f64,
8432 pub bartlett_factor: f64,
8435 pub bartlett_factor_conditional: Option<f64>,
8439 pub rho_variation_shift: Option<f64>,
8442 pub statistic_corrected: f64,
8444 pub p_value_uncorrected: f64,
8446 pub p_value_corrected: f64,
8449 pub material: bool,
8457 pub correction: SmoothLrCorrection,
8459}
8460
8461pub const SMOOTH_LR_MATERIAL_THRESHOLD: f64 = 0.10;
8465
8466fn fitted_rho_penalty_components(
8472 penalties: &[BlockwisePenalty],
8473 lambdas: &[f64],
8474 p_total: usize,
8475) -> Result<Vec<gam_terms::inference::lawley::RhoPenaltyComponent>, EstimationError> {
8476 if penalties.len() != lambdas.len() {
8477 return Err(EstimationError::InvalidInput(format!(
8478 "smooth_term_lr_inference: penalty/lambda count mismatch ({} penalties, {} lambdas)",
8479 penalties.len(),
8480 lambdas.len()
8481 )));
8482 }
8483 let mut components = Vec::with_capacity(penalties.len());
8484 for (idx, (penalty, &lambda)) in penalties.iter().zip(lambdas.iter()).enumerate() {
8485 if !(lambda.is_finite() && lambda >= 0.0) {
8486 return Err(EstimationError::InvalidInput(format!(
8487 "smooth_term_lr_inference: lambda[{idx}] is invalid: {lambda}"
8488 )));
8489 }
8490 let r = &penalty.col_range;
8491 if r.end > p_total {
8492 return Err(EstimationError::InvalidInput(format!(
8493 "smooth_term_lr_inference: penalty[{idx}] range {:?} exceeds coefficient dimension {p_total}",
8494 r
8495 )));
8496 }
8497 let mut s_component = Array2::<f64>::zeros((p_total, p_total));
8498 s_component
8499 .slice_mut(s![r.start..r.end, r.start..r.end])
8500 .scaled_add(lambda, &penalty.local);
8501 components.push(gam_terms::inference::lawley::RhoPenaltyComponent { s_component });
8502 }
8503 Ok(components)
8504}
8505
8506pub fn smooth_term_lr_inference_forspec(
8551 data: ArrayView2<'_, f64>,
8552 y: ArrayView1<'_, f64>,
8553 weights: ArrayView1<'_, f64>,
8554 offset: ArrayView1<'_, f64>,
8555 resolvedspec: &TermCollectionSpec,
8556 family: LikelihoodSpec,
8557 options: &FitOptions,
8558) -> Result<Vec<SmoothTermLrInference>, EstimationError> {
8559 use gam_terms::inference::lawley::{
8560 LAWLEY_PAIR_MATRIX_MAX_ROWS, known_scale_expected_jets_with_dispersion,
8561 lawley_lr_bartlett_factor, lawley_lr_mean_shift_with_rho_variation,
8562 };
8563
8564 let n = data.nrows();
8565 let full = fit_term_collection_forspec(
8568 data,
8569 y,
8570 weights,
8571 offset,
8572 resolvedspec,
8573 family.clone(),
8574 options,
8575 )?;
8576 let ll_full = full.fit.log_likelihood;
8577 let p_total = full.design.design.ncols();
8578 let lambdas = full.fit.lambdas.as_slice().ok_or_else(|| {
8579 EstimationError::InvalidInput(
8580 "smooth_term_lr_inference: non-contiguous lambda vector".to_string(),
8581 )
8582 })?;
8583 let s_lambda = weighted_blockwise_penalty_sum(&full.design.penalties, lambdas, p_total);
8584 let rho_penalty_components =
8585 fitted_rho_penalty_components(&full.design.penalties, lambdas, p_total)?;
8586 let rho_covariance = full.fit.artifacts.rho_covariance.as_ref().filter(|cov| {
8587 cov.nrows() == rho_penalty_components.len() && cov.ncols() == rho_penalty_components.len()
8588 });
8589 let full_design_dense = full.design.design.to_dense();
8591 let influence = full.fit.coefficient_influence();
8592 let fitted_likelihood = resolved_likelihood_for_fit(&full.fit)?;
8593 let family_disp = lawley_dispersion_for_family(&fitted_likelihood, &full.fit)?;
8594 let coefficient_covariance_scale = fitted_likelihood
8595 .coefficient_covariance_scale(family_disp)
8596 .map_err(|error| EstimationError::InvalidInput(error.to_string()))?;
8597
8598 let mut out = Vec::<SmoothTermLrInference>::new();
8599 for (term_idx, design_term) in full.design.smooth.terms.iter().enumerate() {
8600 let penalty_range = full
8601 .design
8602 .smooth_term_penalty_range(term_idx)
8603 .map_err(EstimationError::InvalidInput)?;
8604 let (block_start, k) = penalty_range
8605 .map(|range| (range.start, range.len()))
8606 .unwrap_or((0, 0));
8607 if design_term.shape != ShapeConstraint::None {
8610 continue;
8611 }
8612 let coeff_range = design_term.coeff_range.clone();
8613 if coeff_range.start >= coeff_range.end || coeff_range.end > p_total {
8614 continue;
8615 }
8616 let edf = full.fit.per_term_edf(coeff_range.clone(), block_start, k);
8628 let null_dim = design_term.wald_unpenalized_dim();
8648 let rho_uncertainty_df = match wps_block_uncertainty_df(
8669 full.fit.weighted_gram(),
8670 full.fit.smoothing_correction(),
8671 &coeff_range,
8672 coefficient_covariance_scale,
8673 )? {
8674 Some(extra_df) => extra_df,
8677 None => 0.0,
8678 };
8679 let ref_df = (wood_reference_df(influence, &coeff_range)
8680 .unwrap_or(0.0)
8681 .max(edf)
8682 + rho_uncertainty_df)
8683 .max(null_dim as f64)
8684 .max(1.0);
8685 if !(ref_df.is_finite() && ref_df > 0.0) {
8686 continue;
8687 }
8688
8689 let mut null_spec = resolvedspec.clone();
8692 let Some(spec_pos) = null_spec
8693 .smooth_terms
8694 .iter()
8695 .position(|t| t.name == design_term.name)
8696 else {
8697 continue;
8698 };
8699 null_spec.smooth_terms.remove(spec_pos);
8700 let null_fit = fit_term_collection_forspec(
8701 data,
8702 y,
8703 weights,
8704 offset,
8705 &null_spec,
8706 family.clone(),
8707 options,
8708 );
8709 let (statistic_lr, eta_null) = match null_fit {
8710 Ok(null) if null.fit.log_likelihood.is_finite() => {
8711 let w = (2.0 * (ll_full - null.fit.log_likelihood)).max(0.0);
8712 let null_offset = null
8718 .design
8719 .compose_offset(offset, "smooth likelihood-ratio null model")
8720 .map_err(|error| EstimationError::InvalidInput(error.to_string()))?;
8721 let mut eta = null.design.design.dot(&null.fit.beta);
8722 eta += &null_offset;
8723 (w, Some(eta))
8724 }
8725 _ => (f64::NAN, None),
8726 };
8727
8728 let chi2 = statrs::distribution::ChiSquared::new(ref_df).ok();
8729 let p_uncorrected = match (chi2.as_ref(), statistic_lr.is_finite()) {
8730 (Some(dist), true) => {
8731 use statrs::distribution::ContinuousCDF;
8732 (1.0 - dist.cdf(statistic_lr)).clamp(0.0, 1.0)
8733 }
8734 _ => f64::NAN,
8735 };
8736
8737 let mut bartlett_factor = 1.0;
8741 let mut bartlett_factor_conditional = None;
8742 let mut rho_variation_shift = None;
8743 let mut statistic_corrected = statistic_lr;
8744 let mut p_corrected = p_uncorrected;
8745 let mut correction = SmoothLrCorrection::None;
8746 if let (Some(eta), true, true) = (
8747 eta_null.as_ref(),
8748 statistic_lr.is_finite(),
8749 n <= LAWLEY_PAIR_MATRIX_MAX_ROWS,
8750 ) {
8751 let kappas: Option<Vec<_>> = (0..n)
8752 .map(|i| {
8753 known_scale_expected_jets_with_dispersion(
8754 &fitted_likelihood.spec,
8755 eta[i],
8756 family_disp,
8757 )
8758 .and_then(|jets| jets.kappas().ok())
8759 })
8760 .collect();
8761 if let (Some(kappas), Some(dist)) = (kappas, chi2.as_ref()) {
8762 let fixed_factor = lawley_lr_bartlett_factor(
8763 full_design_dense.view(),
8764 &kappas,
8765 Some(s_lambda.view()),
8766 coeff_range.clone(),
8767 ref_df,
8768 );
8769 if let Ok(c_cond) = fixed_factor
8770 && c_cond.is_finite()
8771 && c_cond > 0.0
8772 {
8773 let mut c_applied = c_cond;
8774 correction = SmoothLrCorrection::LawleyLrFixedLambda;
8775 if let Some(cov) = rho_covariance
8776 && let Ok(total_shift) = lawley_lr_mean_shift_with_rho_variation(
8777 full_design_dense.view(),
8778 &kappas,
8779 s_lambda.view(),
8780 coeff_range.clone(),
8781 &rho_penalty_components,
8782 cov.view(),
8783 )
8784 {
8785 let mean_w = ref_df + total_shift;
8786 if let Some(c_est) =
8787 gam_terms::inference::higher_order::bartlett_factor_from_mean(
8788 mean_w, ref_df,
8789 )
8790 && c_est.is_finite()
8791 && c_est > 0.0
8792 {
8793 let conditional_shift = (c_cond - 1.0) * ref_df;
8794 c_applied = c_est;
8795 bartlett_factor_conditional = Some(c_cond);
8796 rho_variation_shift = Some(total_shift - conditional_shift);
8797 correction = SmoothLrCorrection::LawleyLrEstimatedLambda;
8798 }
8799 }
8800 use statrs::distribution::ContinuousCDF;
8801 bartlett_factor = c_applied;
8802 statistic_corrected = statistic_lr / c_applied;
8803 p_corrected = (1.0 - dist.cdf(statistic_corrected)).clamp(0.0, 1.0);
8804 }
8805 }
8806 }
8807
8808 let material = match correction {
8814 SmoothLrCorrection::LawleyLrEstimatedLambda
8815 | SmoothLrCorrection::LawleyLrFixedLambda => {
8816 let factor_move = (bartlett_factor - 1.0).abs();
8817 let p_denom = p_uncorrected.max(p_corrected).max(f64::MIN_POSITIVE);
8818 let p_move = if p_uncorrected.is_finite() && p_corrected.is_finite() {
8819 (p_corrected - p_uncorrected).abs() / p_denom
8820 } else {
8821 0.0
8822 };
8823 factor_move > SMOOTH_LR_MATERIAL_THRESHOLD || p_move > SMOOTH_LR_MATERIAL_THRESHOLD
8824 }
8825 SmoothLrCorrection::None => false,
8826 };
8827
8828 out.push(SmoothTermLrInference {
8829 name: design_term.name.clone(),
8830 term_idx,
8831 statistic_lr,
8832 ref_df,
8833 bartlett_factor,
8834 bartlett_factor_conditional,
8835 rho_variation_shift,
8836 statistic_corrected,
8837 p_value_uncorrected: p_uncorrected,
8838 p_value_corrected: p_corrected,
8839 material,
8840 correction,
8841 });
8842 }
8843 Ok(out)
8844}
8845
8846fn resolved_likelihood_for_fit(
8847 fit: &UnifiedFitResult,
8848) -> Result<gam_spec::GlmLikelihoodSpec, EstimationError> {
8849 let spec = fit.likelihood_family.as_ref().ok_or_else(|| {
8850 EstimationError::InvalidInput(
8851 "smooth-term LR inference requires an engine-level GLM likelihood".to_string(),
8852 )
8853 })?;
8854 gam_spec::GlmLikelihoodSpec::try_new(spec.clone(), fit.likelihood_scale.clone())
8855 .map_err(|error| EstimationError::InvalidInput(error.to_string()))
8856}
8857
8858fn lawley_dispersion_for_family(
8863 likelihood: &gam_spec::GlmLikelihoodSpec,
8864 fit: &UnifiedFitResult,
8865) -> Result<f64, EstimationError> {
8866 let profiled_standard_deviation = matches!(
8867 likelihood
8868 .resolved_scale()
8869 .map_err(|error| EstimationError::InvalidInput(error.to_string()))?,
8870 gam_spec::ResolvedLikelihoodScale::ProfiledGaussian
8871 )
8872 .then_some(fit.standard_deviation);
8873 gam_solve::estimate::dispersion_from_likelihood(likelihood, profiled_standard_deviation)
8874 .map(|dispersion| dispersion.phi())
8875}
8876
8877fn wps_block_uncertainty_df(
8878 weighted_gram: Option<&Array2<f64>>,
8879 smoothing_correction: Option<&Array2<f64>>,
8880 coeff_range: &Range<usize>,
8881 coefficient_covariance_scale: f64,
8882) -> Result<Option<f64>, EstimationError> {
8883 let (Some(xwx), Some(corr)) = (weighted_gram, smoothing_correction) else {
8884 return Ok(None);
8885 };
8886 let (start, end) = (coeff_range.start, coeff_range.end);
8887 if start >= end {
8888 return Err(EstimationError::InvalidInput(format!(
8889 "WPS coefficient block must be non-empty, got {coeff_range:?}"
8890 )));
8891 }
8892 if xwx.nrows() != xwx.ncols() || corr.nrows() != corr.ncols() {
8893 return Err(EstimationError::InvalidInput(format!(
8894 "WPS matrices must be square, got X'WX={}x{} and correction={}x{}",
8895 xwx.nrows(),
8896 xwx.ncols(),
8897 corr.nrows(),
8898 corr.ncols()
8899 )));
8900 }
8901 if xwx.dim() != corr.dim() || end > xwx.nrows() {
8902 return Err(EstimationError::InvalidInput(format!(
8903 "WPS block {coeff_range:?} is incompatible with X'WX={:?} and correction={:?}",
8904 xwx.dim(),
8905 corr.dim()
8906 )));
8907 }
8908 if !(coefficient_covariance_scale.is_finite() && coefficient_covariance_scale > 0.0) {
8909 return Err(EstimationError::InvalidInput(format!(
8910 "WPS coefficient-covariance scale must be finite and strictly positive, got {coefficient_covariance_scale:?}"
8911 )));
8912 }
8913
8914 let mut trace = gam_linalg::utils::KahanSum::default();
8915 for i in start..end {
8916 for j in start..end {
8917 let gram_value = xwx[[i, j]];
8918 let correction_value = corr[[j, i]];
8919 if !gram_value.is_finite() || !correction_value.is_finite() {
8920 return Err(EstimationError::InvalidInput(format!(
8921 "WPS trace has non-finite matrix entry at ({i}, {j}): X'WX={gram_value:?}, correction-transpose={correction_value:?}"
8922 )));
8923 }
8924 let product = gram_value * correction_value;
8925 if !product.is_finite() {
8926 return Err(EstimationError::InvalidInput(format!(
8927 "WPS trace product is not representable at ({i}, {j}): {gram_value:?} * {correction_value:?}"
8928 )));
8929 }
8930 trace.add(product);
8931 }
8932 }
8933 let trace = trace.sum() / coefficient_covariance_scale;
8934 if !trace.is_finite() {
8935 return Err(EstimationError::InvalidInput(format!(
8936 "WPS corrected-EDF trace is not representable after coefficient scale {coefficient_covariance_scale:?}: {trace:?}"
8937 )));
8938 }
8939 if trace < 0.0 {
8940 return Err(EstimationError::InvalidInput(format!(
8941 "WPS corrected-EDF trace must be non-negative, got {trace:?}"
8942 )));
8943 }
8944 Ok(Some(trace))
8945}
8946
8947fn wood_reference_df(influence: Option<&Array2<f64>>, coeff_range: &Range<usize>) -> Option<f64> {
8971 let f = influence?;
8972 let (start, end) = (coeff_range.start, coeff_range.end);
8973 if start >= end || end > f.nrows() || end > f.ncols() {
8974 return None;
8975 }
8976 let block = f.slice(s![start..end, start..end]);
8977 let tr = (0..block.nrows()).map(|i| block[[i, i]]).sum::<f64>();
8978 let tr2 = block.dot(&block).diag().sum();
8979 (tr.is_finite() && tr2.is_finite() && tr > 0.0).then(|| (2.0 * tr - tr2).max(tr).max(1e-12))
8980}
8981
8982#[cfg(test)]
8983mod likelihood_scale_wps_tests {
8984 use super::wps_block_uncertainty_df;
8985 use ndarray::array;
8986
8987 #[test]
8988 fn wps_trace_uses_coefficient_covariance_scale() {
8989 let xwx = array![[1.0, 0.0], [0.0, 1.0]];
8990 let correction = array![[1.0, 0.0], [0.0, 1.0]];
8991 let extra_df = wps_block_uncertainty_df(Some(&xwx), Some(&correction), &(0..2), 4.0)
8992 .expect("valid WPS geometry")
8993 .expect("correction artifacts are present");
8994 assert_eq!(extra_df, 0.5);
8995 }
8996
8997 #[test]
8998 fn wps_absence_is_distinct_from_invalid_geometry() {
8999 let xwx = array![[1.0]];
9000 assert_eq!(
9001 wps_block_uncertainty_df(Some(&xwx), None, &(0..1), 1.0)
9002 .expect("missing optional artifact is not malformed geometry"),
9003 None
9004 );
9005
9006 let negative_correction = array![[-1.0]];
9007 let error = wps_block_uncertainty_df(Some(&xwx), Some(&negative_correction), &(0..1), 1.0)
9008 .expect_err("negative corrected EDF must not be silently zeroed");
9009 assert!(error.to_string().contains("must be non-negative"));
9010 }
9011}
9012
9013#[cfg(test)]
9014mod nfree_gate_tests {
9015 use super::nfree_skip_gate_status_from_parts;
9016
9017 #[test]
9018 fn value_only_nfree_gate_does_not_require_basis_skip_witness() {
9019 let gate = nfree_skip_gate_status_from_parts(
9020 true, true, false, false, true, true, false, false, );
9029 assert!(
9030 gate.would_skip(false),
9031 "value-only κ cost probes must stay n-free when the Gram value is certified; \
9032 the reduced-basis skip witness is required only for beta/gradient probes"
9033 );
9034 }
9035
9036 #[test]
9037 fn gradient_nfree_gate_still_requires_basis_skip_witness() {
9038 let gate =
9039 nfree_skip_gate_status_from_parts(true, true, false, true, true, true, false, true);
9040 assert!(
9041 !gate.would_skip(true),
9042 "gradient probes return beta/gradient objects in a reduced basis and must not \
9043 skip the row lane without the reduced-basis witness"
9044 );
9045 }
9046}