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}
1566
1567impl SingleBlockLatentCoordDesignCache {
1568 fn new(
1569 data: Array2<f64>,
1570 spec: TermCollectionSpec,
1571 design: TermCollectionDesign,
1572 latent: &StandardLatentCoordConfig,
1573 rho_dim: usize,
1574 ) -> Result<Self, String> {
1575 if latent.term_index.get() >= spec.smooth_terms.len() {
1576 return Err(SmoothError::dimension_mismatch(format!(
1577 "latent-coordinate term index {} out of bounds for {} smooth terms",
1578 latent.term_index,
1579 spec.smooth_terms.len()
1580 ))
1581 .into());
1582 }
1583 if latent.feature_cols.len() != latent.values.latent_dim() {
1584 return Err(SmoothError::dimension_mismatch(format!(
1585 "latent-coordinate feature width mismatch: feature_cols={}, latent_dim={}",
1586 latent.feature_cols.len(),
1587 latent.values.latent_dim()
1588 ))
1589 .into());
1590 }
1591 if latent.values.n_obs() != data.nrows() {
1592 return Err(SmoothError::dimension_mismatch(format!(
1593 "latent-coordinate row mismatch: latent n={}, data n={}",
1594 latent.values.n_obs(),
1595 data.nrows()
1596 ))
1597 .into());
1598 }
1599 let analytic_rho_count = latent
1600 .analytic_penalties
1601 .as_ref()
1602 .map_or(0, |registry| registry.total_rho_count());
1603 Ok(Self {
1604 data,
1605 spec,
1606 design,
1607 current_theta: None,
1608 current_latent: None,
1609 current_hyper_dirs: None,
1610 current_design_cache_id: None,
1611 latent_design_cache: gam_solve::latent_cache::LatentDesignCache::default(),
1612 last_cost: None,
1613 last_eval: None,
1614 term_index: latent.term_index,
1615 feature_cols: latent.feature_cols.clone(),
1616 rho_dim,
1617 n_obs: latent.values.n_obs(),
1618 latent_dim: latent.values.latent_dim(),
1619 id_mode: latent.values.id_mode().clone(),
1620 manifold: latent.values.manifold().clone(),
1621 retraction_registry: latent.values.retraction_registry().clone(),
1622 latent_id: latent.values.latent_id(),
1623 analytic_penalties: latent.analytic_penalties.clone(),
1624 analytic_rho_count,
1625 design_revision: 0,
1626 })
1627 }
1628
1629 fn design_revision(&self) -> u64 {
1630 self.design_revision
1631 }
1632
1633 fn design(&self) -> &TermCollectionDesign {
1634 &self.design
1635 }
1636
1637 fn latent(&self) -> Result<std::sync::Arc<gam_terms::latent::LatentCoordValues>, String> {
1638 self.current_latent
1639 .as_ref()
1640 .cloned()
1641 .ok_or_else(|| "latent-coordinate cache has not been realized".to_string())
1642 }
1643
1644 fn analytic_penalties(&self) -> Option<std::sync::Arc<gam_terms::AnalyticPenaltyRegistry>> {
1645 self.analytic_penalties.clone()
1646 }
1647
1648 fn analytic_penalty_rho_count(&self) -> usize {
1649 self.analytic_rho_count
1650 }
1651
1652 fn hyper_dirs(&self) -> Result<Vec<gam_solve::estimate::reml::DirectionalHyperParam>, String> {
1653 self.current_hyper_dirs
1654 .as_ref()
1655 .cloned()
1656 .ok_or_else(|| "latent-coordinate hyper_dirs cache has not been realized".to_string())
1657 }
1658
1659 fn latent_basis_kind(&self) -> Result<gam_solve::latent_cache::LatentBasisKind, String> {
1660 let smooth_term = self
1661 .design
1662 .smooth
1663 .terms
1664 .get(self.term_index.get())
1665 .ok_or_else(|| {
1666 SmoothError::dimension_mismatch(format!(
1667 "LatentCoord term index {} out of bounds for realized smooth design",
1668 self.term_index
1669 ))
1670 })?;
1671 let termspec = self
1672 .spec
1673 .smooth_terms
1674 .get(self.term_index.get())
1675 .ok_or_else(|| {
1676 SmoothError::dimension_mismatch(format!(
1677 "LatentCoord term index {} out of bounds for resolved smooth spec",
1678 self.term_index
1679 ))
1680 })?;
1681 match (&termspec.basis, &smooth_term.metadata) {
1682 (
1683 SmoothBasisSpec::Matern { .. },
1684 BasisMetadata::Matern {
1685 centers,
1686 length_scale,
1687 nu,
1688 aniso_log_scales,
1689 ..
1690 },
1691 ) => Ok(gam_solve::latent_cache::LatentBasisKind::Matern {
1692 centers: centers.clone(),
1693 length_scale: *length_scale,
1694 nu: *nu,
1695 aniso_log_scales: aniso_log_scales
1696 .clone()
1697 .unwrap_or_else(|| vec![0.0; centers.ncols()]),
1698 chunk_size: gam_terms::basis::auto_streaming_chunk_size_for_dense(
1699 self.n_obs,
1700 centers.nrows(),
1701 ),
1702 }),
1703 (
1704 SmoothBasisSpec::Duchon { .. },
1705 BasisMetadata::Duchon {
1706 centers,
1707 length_scale,
1708 power,
1709 nullspace_order,
1710 aniso_log_scales,
1711 ..
1712 },
1713 ) => Ok(gam_solve::latent_cache::LatentBasisKind::Duchon {
1714 centers: centers.clone(),
1715 length_scale: *length_scale,
1716 power: *power,
1717 nullspace_order: *nullspace_order,
1718 aniso_log_scales: aniso_log_scales
1719 .clone()
1720 .unwrap_or_else(|| vec![0.0; centers.ncols()]),
1721 }),
1722 (
1723 SmoothBasisSpec::Sphere { .. },
1724 BasisMetadata::Sphere {
1725 centers,
1726 penalty_order,
1727 method,
1728 ..
1729 },
1730 ) if matches!(*method, gam_terms::basis::SphereMethod::Wahba) => {
1731 Ok(gam_solve::latent_cache::LatentBasisKind::Sphere {
1732 centers: centers.clone(),
1733 penalty_order: *penalty_order,
1734 chunk_size: gam_terms::basis::auto_streaming_chunk_size_for_dense(
1735 self.n_obs,
1736 centers.nrows(),
1737 ),
1738 })
1739 }
1740 (
1741 SmoothBasisSpec::BSpline1D { spec, .. },
1742 BasisMetadata::BSpline1D {
1743 knots,
1744 periodic,
1745 degree: meta_degree,
1746 ..
1747 },
1748 ) => {
1749 let effective_degree = meta_degree.unwrap_or(spec.degree);
1753 if let Some((domain_start, period, num_basis)) = periodic {
1754 Ok(gam_solve::latent_cache::LatentBasisKind::PeriodicBspline {
1755 domain_start: *domain_start,
1756 period: *period,
1757 degree: effective_degree,
1758 num_basis: *num_basis,
1759 chunk_size: gam_terms::basis::auto_streaming_chunk_size_for_dense(
1760 self.n_obs, *num_basis,
1761 ),
1762 })
1763 } else {
1764 let num_basis_est = knots.len().saturating_sub(effective_degree + 1);
1765 Ok(gam_solve::latent_cache::LatentBasisKind::TensorBspline {
1766 knots: vec![knots.clone()],
1767 degrees: vec![effective_degree],
1768 chunk_size: gam_terms::basis::auto_streaming_chunk_size_for_dense(
1769 self.n_obs,
1770 num_basis_est,
1771 ),
1772 })
1773 }
1774 }
1775 (
1776 SmoothBasisSpec::TensorBSpline { .. },
1777 BasisMetadata::TensorBSpline { knots, degrees, .. },
1778 ) => Ok(gam_solve::latent_cache::LatentBasisKind::TensorBspline {
1779 knots: knots.clone(),
1780 degrees: degrees.clone(),
1781 chunk_size: None,
1782 }),
1783 (
1784 SmoothBasisSpec::Pca { .. },
1785 BasisMetadata::Pca {
1786 basis_matrix,
1787 centered,
1788 smooth_penalty,
1789 center_mean,
1790 pca_basis_path,
1791 chunk_size,
1792 ..
1793 },
1794 ) => {
1795 let center_mean_fingerprint = if *centered && pca_basis_path.is_none() {
1796 let mean = center_mean.as_ref().ok_or_else(|| {
1797 SmoothError::invalid_config(
1798 "latent-coordinate Pca cache key requires center_mean when centered",
1799 )
1800 })?;
1801 Some(gam_solve::latent_cache::pca_center_mean_fingerprint(mean))
1802 } else {
1803 None
1804 };
1805 Ok(gam_solve::latent_cache::LatentBasisKind::Pca {
1806 basis_matrix: basis_matrix.clone(),
1807 centered: *centered,
1808 center_mean_fingerprint,
1809 smooth_penalty: *smooth_penalty,
1810 pca_basis_path: pca_basis_path.clone(),
1811 chunk_size: *chunk_size,
1812 })
1813 }
1814 _ => Err(SmoothError::invalid_config(
1815 "latent-coordinate design cache could not key the realized latent smooth basis"
1816 .to_string(),
1817 )
1818 .into()),
1819 }
1820 }
1821
1822 fn ensure_theta(&mut self, theta: &Array1<f64>) -> Result<(), String> {
1823 if self
1824 .current_theta
1825 .as_ref()
1826 .is_some_and(|cached| theta_values_match(cached, theta))
1827 {
1828 return Ok(());
1829 }
1830 let latent_flat_len = self.n_obs * self.latent_dim;
1831 let direct_hyper_count = latent_coord_direct_hyper_count(&self.id_mode, self.latent_dim);
1832 let expected =
1833 self.rho_dim + latent_flat_len + self.analytic_rho_count + direct_hyper_count;
1834 if theta.len() != expected {
1835 return Err(SmoothError::dimension_mismatch(format!(
1836 "latent-coordinate theta length mismatch: got {}, expected {} (rho_dim={}, n={}, d={}, analytic_rhos={}, direct_hypers={})",
1837 theta.len(),
1838 expected,
1839 self.rho_dim,
1840 self.n_obs,
1841 self.latent_dim,
1842 self.analytic_rho_count,
1843 direct_hyper_count
1844 ))
1845 .into());
1846 }
1847 let flat = theta
1848 .slice(s![self.rho_dim..self.rho_dim + latent_flat_len])
1849 .to_owned();
1850 let latent = std::sync::Arc::new(
1851 gam_terms::latent::LatentCoordValues::from_flat_with_manifold_and_retraction_and_id(
1852 flat,
1853 self.n_obs,
1854 self.latent_dim,
1855 self.id_mode.clone(),
1856 self.manifold.clone(),
1857 self.retraction_registry.clone(),
1858 self.latent_id,
1859 ),
1860 );
1861 let latent_values_changed = self
1862 .current_latent
1863 .as_ref()
1864 .map(|cached| !latent_values_match(cached.as_flat(), latent.as_flat()))
1865 .unwrap_or(true);
1866 if latent_values_changed {
1867 self.latent_design_cache.invalidate_all();
1868 self.current_design_cache_id = None;
1869 self.design_revision = self.design_revision.wrapping_add(1);
1870 }
1871 for n in 0..self.n_obs {
1872 for axis in 0..self.latent_dim {
1873 let col = self.feature_cols[axis];
1874 self.data[[n, col]] = latent.as_flat()[n * self.latent_dim + axis];
1875 }
1876 }
1877
1878 let basis_kind = self.latent_basis_kind()?;
1879 let rebuilt_width = self.design.design.ncols();
1880 let spec = self.spec.clone();
1881 let term_index = self.term_index;
1882 let analytic_rho_count = self.analytic_rho_count;
1883 let data = self.data.view();
1884 let design_context_digest = gam_solve::latent_cache::latent_design_context_cache_digest(
1885 data,
1886 &spec,
1887 term_index,
1888 analytic_rho_count,
1889 &self.feature_cols,
1890 )
1891 .map_err(|e| e.to_string())?;
1892 let lookup = self
1893 .latent_design_cache
1894 .lookup_or_compute(latent.clone(), basis_kind, design_context_digest, || {
1895 let rebuilt = build_term_collection_design(data, &spec).map_err(|e| {
1896 EstimationError::InvalidInput(format!(
1897 "failed to rebuild latent-coordinate design: {e}"
1898 ))
1899 })?;
1900 if rebuilt.design.ncols() != rebuilt_width {
1901 crate::bail_invalid_estim!(
1902 "latent-coordinate design topology changed: rebuilt p={}, cached p={}",
1903 rebuilt.design.ncols(),
1904 rebuilt_width
1905 );
1906 }
1907 let hyper_dirs = try_build_latent_coord_hyper_dirs(
1908 latent.clone(),
1909 &spec,
1910 &rebuilt,
1911 &[term_index],
1912 analytic_rho_count,
1913 )?
1914 .ok_or_else(|| {
1915 EstimationError::InvalidInput(
1916 "failed to build latent-coordinate hyper_dirs".to_string(),
1917 )
1918 })?;
1919 Ok(gam_solve::latent_cache::ComputedLatentDesign {
1920 design: rebuilt,
1921 hyper_dirs,
1922 })
1923 })
1924 .map_err(|e| e.to_string())?;
1925 if lookup.cached.design.design.ncols() != self.design.design.ncols() {
1926 return Err(SmoothError::dimension_mismatch(format!(
1927 "latent-coordinate design topology changed: rebuilt p={}, cached p={}",
1928 lookup.cached.design.design.ncols(),
1929 self.design.design.ncols()
1930 ))
1931 .into());
1932 }
1933 self.design = lookup.cached.design.clone();
1934 self.current_hyper_dirs = Some(lookup.cached.hyper_dirs.clone());
1935 self.current_latent = Some(latent);
1936 self.current_theta = Some(theta.clone());
1937 self.last_cost = None;
1938 self.last_eval = None;
1939 if !latent_values_changed && self.current_design_cache_id != Some(lookup.entry_id) {
1940 self.design_revision = self.design_revision.wrapping_add(1);
1941 }
1942 self.current_design_cache_id = Some(lookup.entry_id);
1943 Ok(())
1944 }
1945
1946 fn memoized_cost(&self, theta: &Array1<f64>) -> Option<f64> {
1947 if self
1948 .current_theta
1949 .as_ref()
1950 .is_some_and(|cached| theta_values_match(cached, theta))
1951 {
1952 self.last_eval
1953 .as_ref()
1954 .map(|cached| cached.0)
1955 .or(self.last_cost)
1956 } else {
1957 None
1958 }
1959 }
1960
1961 fn memoized_eval(
1962 &self,
1963 theta: &Array1<f64>,
1964 ) -> Option<(f64, Array1<f64>, gam_problem::HessianValue)> {
1965 if self
1966 .current_theta
1967 .as_ref()
1968 .is_some_and(|cached| theta_values_match(cached, theta))
1969 {
1970 self.last_eval.clone()
1971 } else {
1972 None
1973 }
1974 }
1975
1976 fn store_eval(&mut self, eval: (f64, Array1<f64>, gam_problem::HessianValue)) {
1977 self.last_cost = Some(eval.0);
1978 self.last_eval = Some(eval);
1979 }
1980
1981 fn store_cost(&mut self, cost: f64) {
1982 self.last_cost = Some(cost);
1983 }
1984
1985 fn reset(&mut self) {
1986 self.current_theta = None;
1987 self.current_latent = None;
1988 self.current_hyper_dirs = None;
1989 self.current_design_cache_id = None;
1990 self.latent_design_cache.invalidate();
1991 self.last_cost = None;
1992 self.last_eval = None;
1993 }
1994}
1995
1996pub fn fixed_kappa_profiled_reml_score(
2012 data: ArrayView2<'_, f64>,
2013 y: ArrayView1<'_, f64>,
2014 weights: ArrayView1<'_, f64>,
2015 offset: ArrayView1<'_, f64>,
2016 resolvedspec: &TermCollectionSpec,
2017 term_idx: usize,
2018 kappa: f64,
2019 family: LikelihoodSpec,
2020 options: &FitOptions,
2021) -> Result<f64, EstimationError> {
2022 if !kappa.is_finite() {
2023 crate::bail_invalid_estim!("fixed-κ profiled score probed a non-finite κ = {kappa}");
2024 }
2025 let (feature_cols, mut probe_basis) =
2028 match resolvedspec.smooth_terms.get(term_idx).map(|t| &t.basis) {
2029 Some(SmoothBasisSpec::ConstantCurvature {
2030 feature_cols, spec, ..
2031 }) => (feature_cols.clone(), spec.clone()),
2032 _ => {
2033 crate::bail_invalid_estim!(
2034 "fixed-κ profiled score: term {term_idx} is not a constant-curvature smooth"
2035 )
2036 }
2037 };
2038 probe_basis.kappa = kappa;
2039
2040 let is_unweighted = weights.iter().all(|&w| (w - 1.0).abs() <= 1e-12);
2060 let is_zero_offset = offset.iter().all(|&o| o.abs() <= 1e-12);
2061 if family == LikelihoodSpec::gaussian_identity() && is_unweighted && is_zero_offset {
2062 let x_term = select_columns(data, &feature_cols).map_err(EstimationError::from)?;
2063 let score = gam_terms::basis::constant_curvature_honest_profiled_reml_score(
2064 x_term.view(),
2065 y,
2066 &probe_basis,
2067 )
2068 .map_err(|e| {
2069 EstimationError::InvalidInput(format!(
2070 "fixed-κ honest profiled-REML score at κ={kappa} failed: {e}"
2071 ))
2072 })?;
2073 if !score.is_finite() {
2074 crate::bail_invalid_estim!(
2075 "fixed-κ honest profiled-REML score at κ={kappa} is non-finite"
2076 );
2077 }
2078 return Ok(score);
2079 }
2080
2081 let mut probe_spec = resolvedspec.clone();
2083 match probe_spec
2084 .smooth_terms
2085 .get_mut(term_idx)
2086 .map(|t| &mut t.basis)
2087 {
2088 Some(SmoothBasisSpec::ConstantCurvature { spec, .. }) => spec.kappa = kappa,
2089 _ => {
2090 crate::bail_invalid_estim!(
2091 "fixed-κ profiled score: term {term_idx} is not a constant-curvature smooth"
2092 )
2093 }
2094 }
2095 let fixed_kappa_options = SpatialLengthScaleOptimizationOptions {
2096 enabled: false,
2097 ..SpatialLengthScaleOptimizationOptions::default()
2098 };
2099 let fit = fit_term_collectionwith_spatial_length_scale_optimization(
2100 data,
2101 y.to_owned(),
2102 weights.to_owned(),
2103 offset.to_owned(),
2104 &probe_spec,
2105 family,
2106 options,
2107 &fixed_kappa_options,
2108 )?;
2109 let score = fit_score(&fit.fit);
2110 if !score.is_finite() {
2111 crate::bail_invalid_estim!("fixed-κ profiled fit at κ={kappa} returned a non-finite score");
2112 }
2113 Ok(score)
2114}
2115
2116fn profiled_gaussian_reml_value_kappa_gradient(
2121 design: &Array2<f64>,
2122 design_kappa: &Array2<f64>,
2123 penalty: &Array2<f64>,
2124 penalty_kappa: &Array2<f64>,
2125 response: ArrayView1<'_, f64>,
2126) -> Result<(f64, f64), EstimationError> {
2127 if design.dim() != design_kappa.dim()
2128 || penalty.dim() != penalty_kappa.dim()
2129 || penalty.dim() != (design.ncols(), design.ncols())
2130 || response.len() != design.nrows()
2131 {
2132 crate::bail_invalid_estim!("constant-curvature profile value/gradient shape mismatch");
2133 }
2134
2135 let response_2d = response.insert_axis(ndarray::Axis(1));
2136 let fit = gam_solve::gaussian_reml::gaussian_reml_multi_closed_form(
2137 design.view(),
2138 response_2d.view(),
2139 penalty.view(),
2140 None,
2141 None,
2142 )?;
2143 let backward = gam_solve::gaussian_reml::gaussian_reml_multi_closed_form_backward_from_fit(
2144 design.view(),
2145 response_2d.view(),
2146 penalty.view(),
2147 None,
2148 &fit,
2149 0.0,
2150 None,
2151 None,
2152 1.0,
2153 0.0,
2154 )?;
2155 let derivative = backward
2156 .grad_x
2157 .iter()
2158 .zip(design_kappa.iter())
2159 .map(|(&adjoint, &direction)| adjoint * direction)
2160 .sum::<f64>()
2161 + backward
2162 .grad_penalty
2163 .iter()
2164 .zip(penalty_kappa.iter())
2165 .map(|(&adjoint, &direction)| adjoint * direction)
2166 .sum::<f64>();
2167 if !(fit.reml_score.is_finite() && derivative.is_finite()) {
2168 crate::bail_invalid_estim!(
2169 "constant-curvature analytic profile returned a non-finite value or derivative"
2170 );
2171 }
2172 Ok((fit.reml_score, derivative))
2173}
2174
2175fn constant_curvature_radial_reference(
2179 data: ArrayView2<'_, f64>,
2180 y: ArrayView1<'_, f64>,
2181) -> Result<Array1<f64>, EstimationError> {
2182 if y.len() != data.nrows() || y.is_empty() {
2183 crate::bail_invalid_estim!(
2184 "constant-curvature radial reference needs one non-empty response per row"
2185 );
2186 }
2187 let radii: Array1<f64> = data.outer_iter().map(|row| row.dot(&row).sqrt()).collect();
2188 let r_max = radii.iter().copied().fold(0.0_f64, f64::max);
2189 if r_max <= f64::MIN_POSITIVE {
2190 let mean = y.sum() / y.len() as f64;
2191 return Ok(Array1::from_elem(y.len(), mean));
2192 }
2193
2194 let bin_count = (data.nrows() as f64).log2().ceil() as usize + 1;
2195 let bin_of = |radius: f64| -> usize {
2196 ((radius / r_max * bin_count as f64) as usize).min(bin_count - 1)
2197 };
2198 let mut sums = vec![0.0; bin_count];
2199 let mut counts = vec![0usize; bin_count];
2200 for (row, &radius) in radii.iter().enumerate() {
2201 let bin = bin_of(radius);
2202 sums[bin] += y[row];
2203 counts[bin] += 1;
2204 }
2205 let means: Vec<f64> = sums
2206 .into_iter()
2207 .zip(counts)
2208 .map(
2209 |(sum, count)| {
2210 if count == 0 { 0.0 } else { sum / count as f64 }
2211 },
2212 )
2213 .collect();
2214 Ok(radii.mapv(|radius| means[bin_of(radius)]))
2215}
2216
2217fn constant_curvature_kappa_fair_profile_value_gradient(
2222 data: ArrayView2<'_, f64>,
2223 y: ArrayView1<'_, f64>,
2224 y_ref: ArrayView1<'_, f64>,
2225 spec: &gam_terms::basis::ConstantCurvatureBasisSpec,
2226) -> Result<(f64, f64), EstimationError> {
2227 if y.len() != data.nrows() || y_ref.len() != data.nrows() {
2228 crate::bail_invalid_estim!(
2229 "constant-curvature fair profile row mismatch: data={}, response={}, reference={}",
2230 data.nrows(),
2231 y.len(),
2232 y_ref.len(),
2233 );
2234 }
2235
2236 let mut profile_spec = spec.clone();
2237 profile_spec.double_penalty = false;
2238 let basis = gam_terms::basis::build_constant_curvature_basis(data, &profile_spec)
2239 .map_err(EstimationError::from)?;
2240 let derivatives =
2241 gam_terms::basis::build_constant_curvature_basis_kappa_derivatives(data, &profile_spec)
2242 .map_err(EstimationError::from)?;
2243 if basis.active_penalties.len() != 1 || derivatives.first.penalties_derivative.len() != 1 {
2244 crate::bail_invalid_estim!(
2245 "constant-curvature fair profile expected one primary penalty; value blocks={}, derivative blocks={}",
2246 basis.active_penalties.len(),
2247 derivatives.first.penalties_derivative.len(),
2248 );
2249 }
2250
2251 let smooth_design = basis.design.to_dense();
2252 let smooth_design_kappa = &derivatives.first.design_derivative;
2253 let smooth_penalty = &basis.active_penalties[0].matrix;
2254 let smooth_penalty_kappa = &derivatives.first.penalties_derivative[0];
2255 let n = smooth_design.nrows();
2256 let p = smooth_design.ncols();
2257 if smooth_design_kappa.dim() != (n, p)
2258 || smooth_penalty.dim() != (p, p)
2259 || smooth_penalty_kappa.dim() != (p, p)
2260 {
2261 crate::bail_invalid_estim!(
2262 "constant-curvature kappa derivative bundle does not match its value basis"
2263 );
2264 }
2265
2266 let mut design = Array2::<f64>::ones((n, p + 1));
2267 design.slice_mut(s![.., 1..]).assign(&smooth_design);
2268 let mut design_kappa = Array2::<f64>::zeros((n, p + 1));
2269 design_kappa
2270 .slice_mut(s![.., 1..])
2271 .assign(smooth_design_kappa);
2272 let mut penalty = Array2::<f64>::zeros((p + 1, p + 1));
2273 penalty.slice_mut(s![1.., 1..]).assign(smooth_penalty);
2274 let mut penalty_kappa = Array2::<f64>::zeros((p + 1, p + 1));
2275 penalty_kappa
2276 .slice_mut(s![1.., 1..])
2277 .assign(smooth_penalty_kappa);
2278
2279 let (value_y, derivative_y) = profiled_gaussian_reml_value_kappa_gradient(
2280 &design,
2281 &design_kappa,
2282 &penalty,
2283 &penalty_kappa,
2284 y,
2285 )?;
2286 let (value_ref, derivative_ref) = profiled_gaussian_reml_value_kappa_gradient(
2287 &design,
2288 &design_kappa,
2289 &penalty,
2290 &penalty_kappa,
2291 y_ref,
2292 )?;
2293 Ok((value_y - value_ref, derivative_y - derivative_ref))
2294}
2295
2296struct ConstantCurvatureFairProfile<'a> {
2297 data: ArrayView2<'a, f64>,
2298 response: ArrayView1<'a, f64>,
2299 radial_reference: Array1<f64>,
2300 spec: gam_terms::basis::ConstantCurvatureBasisSpec,
2301 cache: std::cell::RefCell<std::collections::HashMap<u64, (f64, f64)>>,
2302}
2303
2304impl ConstantCurvatureFairProfile<'_> {
2305 fn evaluate(&self, kappa: f64) -> Result<(f64, f64), EstimationError> {
2306 if !kappa.is_finite() {
2307 crate::bail_invalid_estim!("constant-curvature fair profile probed a non-finite kappa");
2308 }
2309 let key = kappa.to_bits();
2310 if let Some(&cached) = self.cache.borrow().get(&key) {
2311 return Ok(cached);
2312 }
2313 let mut probe_spec = self.spec.clone();
2314 probe_spec.kappa = kappa;
2315 let sample = constant_curvature_kappa_fair_profile_value_gradient(
2316 self.data,
2317 self.response,
2318 self.radial_reference.view(),
2319 &probe_spec,
2320 )?;
2321 self.cache.borrow_mut().insert(key, sample);
2322 Ok(sample)
2323 }
2324}
2325
2326fn validate_constant_curvature_fair_profile_inputs(
2327 weights: ArrayView1<'_, f64>,
2328 offset: ArrayView1<'_, f64>,
2329 family: &LikelihoodSpec,
2330) -> Result<(), EstimationError> {
2331 if *family != LikelihoodSpec::gaussian_identity() {
2332 crate::bail_invalid_estim!(
2333 "curvature-as-an-estimand profile currently requires Gaussian identity likelihood"
2334 );
2335 }
2336 let input_tolerance = f64::EPSILON.sqrt();
2337 if weights
2338 .iter()
2339 .any(|&weight| (weight - 1.0).abs() > input_tolerance)
2340 || offset.iter().any(|&value| value.abs() > input_tolerance)
2341 {
2342 crate::bail_invalid_estim!(
2343 "curvature-as-an-estimand profile requires unit weights and zero offset"
2344 );
2345 }
2346 Ok(())
2347}
2348
2349fn constant_curvature_kappa_fair_optimum(
2356 data: ArrayView2<'_, f64>,
2357 y: ArrayView1<'_, f64>,
2358 resolvedspec: &TermCollectionSpec,
2359 term_idx: usize,
2360 options: &FitOptions,
2361) -> Result<f64, EstimationError> {
2362 let (kappa_min, kappa_max) = constant_curvature_kappa_bounds(data, resolvedspec, term_idx);
2363 if !(kappa_min.is_finite() && kappa_max.is_finite() && kappa_max > kappa_min) {
2364 crate::bail_invalid_estim!(
2365 "constant-curvature term {term_idx} has invalid kappa bounds [{kappa_min}, {kappa_max}]"
2366 );
2367 }
2368 let (feature_cols, base_spec) = match resolvedspec
2369 .smooth_terms
2370 .get(term_idx)
2371 .map(|term| &term.basis)
2372 {
2373 Some(SmoothBasisSpec::ConstantCurvature {
2374 feature_cols, spec, ..
2375 }) => (feature_cols, spec.clone()),
2376 _ => {
2377 crate::bail_invalid_estim!(
2378 "constant-curvature optimum requested for non-curvature term {term_idx}"
2379 )
2380 }
2381 };
2382 let x_term = select_columns(data, feature_cols).map_err(EstimationError::from)?;
2383 let y_ref = constant_curvature_radial_reference(x_term.view(), y)?;
2384 let profile = ConstantCurvatureFairProfile {
2385 data: x_term.view(),
2386 response: y,
2387 radial_reference: y_ref,
2388 spec: base_spec,
2389 cache: std::cell::RefCell::new(std::collections::HashMap::new()),
2390 };
2391 let mut seed_config = gam_problem::SeedConfig::default();
2392 seed_config.bounds = (kappa_min, kappa_max);
2393 seed_config.max_seeds = 1;
2394 seed_config.seed_budget = 1;
2395 seed_config.risk_profile = gam_problem::SeedRiskProfile::Gaussian;
2396 seed_config.num_auxiliary_trailing = 1;
2397 seed_config.over_smoothing_probe_rho = None;
2398 let initial_kappa = profile.spec.kappa.clamp(kappa_min, kappa_max);
2399 let problem = gam_solve::rho_optimizer::OuterProblem::new(1)
2400 .with_gradient(gam_problem::Derivative::Analytic)
2401 .with_hessian(gam_problem::DeclaredHessianForm::Unavailable)
2402 .with_prefer_gradient_only(true)
2403 .with_disable_fixed_point(true)
2404 .with_fallback_policy(gam_solve::rho_optimizer::FallbackPolicy::Disabled)
2405 .with_psi_dim(1)
2406 .with_tolerance(options.tol.max(f64::EPSILON.sqrt()))
2407 .with_max_iter(options.max_iter.max(1))
2408 .with_bounds(
2409 Array1::from_vec(vec![kappa_min]),
2410 Array1::from_vec(vec![kappa_max]),
2411 )
2412 .with_initial_rho(Array1::from_vec(vec![initial_kappa]))
2413 .with_seed_config(seed_config);
2414 let mut objective = problem.build_objective(
2415 profile,
2416 |profile: &mut ConstantCurvatureFairProfile<'_>, theta: &Array1<f64>| {
2417 profile.evaluate(theta[0]).map(|(value, _)| value)
2418 },
2419 |profile: &mut ConstantCurvatureFairProfile<'_>, theta: &Array1<f64>| {
2420 let (cost, derivative) = profile.evaluate(theta[0])?;
2421 Ok(gam_problem::OuterEval {
2422 cost,
2423 gradient: Array1::from_vec(vec![derivative]),
2424 hessian: gam_problem::HessianValue::Unavailable,
2425 inner_beta_hint: None,
2426 })
2427 },
2428 None::<fn(&mut ConstantCurvatureFairProfile<'_>)>,
2429 None::<
2430 fn(
2431 &mut ConstantCurvatureFairProfile<'_>,
2432 &Array1<f64>,
2433 ) -> Result<gam_problem::EfsEval, EstimationError>,
2434 >,
2435 );
2436 let result = problem.run(
2437 &mut objective,
2438 &format!("constant-curvature fair profile term {term_idx}"),
2439 )?;
2440 if !result.converged {
2441 crate::bail_invalid_estim!(
2442 "constant-curvature fair-profile κ optimization did not converge for term {} after {} iterations (negative_log_evidence={:.6e}, final_grad_norm={})",
2443 term_idx,
2444 result.iterations,
2445 result.final_value,
2446 result.final_grad_norm_report(),
2447 );
2448 }
2449 let kappa_hat = result.rho[0];
2450 log::info!(
2451 "[spatial-kappa] continuous fair-profile optimum kappa_hat={:.6} \
2452 (negative_log_evidence={:.6e}, projected_gradient={}) for term {term_idx}",
2453 kappa_hat,
2454 result.final_value,
2455 result.final_grad_norm_report(),
2456 );
2457 Ok(kappa_hat)
2458}
2459
2460fn try_exact_joint_spatial_length_scale_optimization(
2461 data: ArrayView2<'_, f64>,
2462 y: ArrayView1<'_, f64>,
2463 weights: ArrayView1<'_, f64>,
2464 offset: ArrayView1<'_, f64>,
2465 resolvedspec: &TermCollectionSpec,
2466 best: &FittedTermCollection,
2467 family: LikelihoodSpec,
2468 options: &FitOptions,
2469 kappa_options: &SpatialLengthScaleOptimizationOptions,
2470 spatial_terms: &[usize],
2471) -> Result<Option<FittedTermCollectionWithSpec>, EstimationError> {
2472 if spatial_terms.is_empty() {
2473 return Ok(None);
2474 }
2475 kappa_options
2480 .validate()
2481 .map_err(EstimationError::InvalidInput)?;
2482
2483 if try_build_spatial_log_kappa_hyper_dirs(data, resolvedspec, &best.design, spatial_terms)?
2484 .is_none()
2485 {
2486 if !constant_curvature_term_indices(resolvedspec).is_empty() {
2487 log::info!(
2488 "[#1464-trace] try_exact_joint RETURNED None (hyper_dirs unavailable); \
2489 κ̂ comes from a NON-joint path"
2490 );
2491 }
2492 return Ok(None);
2493 }
2494 if !constant_curvature_term_indices(resolvedspec).is_empty() {
2495 log::info!(
2496 "[#1464-trace] try_exact_joint ENTERED for {} spatial term(s); CC present",
2497 spatial_terms.len()
2498 );
2499 }
2500
2501 const JOINT_RHO_BOUND: f64 = 12.0;
2502 let rho_dim = best.fit.lambdas.len();
2503
2504 let has_constant_curvature_term = !constant_curvature_term_indices(resolvedspec).is_empty();
2518 let rho_upper_bound = if has_constant_curvature_term {
2519 gam_solve::estimate::RHO_BOUND
2520 } else {
2521 JOINT_RHO_BOUND
2522 };
2523
2524 let dims_per_term = spatial_dims_per_term(resolvedspec, spatial_terms);
2526 let use_aniso = has_aniso_terms(resolvedspec, spatial_terms);
2527
2528 let log_kappa0 = if use_aniso {
2533 SpatialLogKappaCoords::from_length_scales_aniso(resolvedspec, spatial_terms, kappa_options)
2534 } else {
2535 SpatialLogKappaCoords::from_length_scales(resolvedspec, spatial_terms, kappa_options)
2536 };
2537 let mut log_kappa0 = log_kappa0
2540 .reseed_from_data(data, resolvedspec, spatial_terms, kappa_options)
2541 .map_err(EstimationError::BasisError)?;
2542 let mut cc_profiled_values: Vec<(usize, f64)> = Vec::new();
2547 if has_constant_curvature_term {
2548 for (slot, &term_idx) in spatial_terms.iter().enumerate() {
2549 if constant_curvature_term_spec(resolvedspec, term_idx).is_none() {
2550 continue;
2551 }
2552 let kappa = get_constant_curvature_kappa(resolvedspec, term_idx)
2553 .expect("constant-curvature term exposes its kappa");
2554 log_kappa0.set_scalar_slot(slot, kappa);
2555 cc_profiled_values.push((slot, kappa));
2556 }
2557 }
2558 let log_kappa_lower = if use_aniso {
2559 SpatialLogKappaCoords::lower_bounds_aniso_from_data(
2560 data,
2561 resolvedspec,
2562 spatial_terms,
2563 &dims_per_term,
2564 kappa_options,
2565 )
2566 } else {
2567 SpatialLogKappaCoords::lower_bounds_from_data(
2568 data,
2569 resolvedspec,
2570 spatial_terms,
2571 kappa_options,
2572 )
2573 }
2574 .map_err(EstimationError::BasisError)?;
2575 let log_kappa_upper = if use_aniso {
2576 SpatialLogKappaCoords::upper_bounds_aniso_from_data(
2577 data,
2578 resolvedspec,
2579 spatial_terms,
2580 &dims_per_term,
2581 kappa_options,
2582 )
2583 } else {
2584 SpatialLogKappaCoords::upper_bounds_from_data(
2585 data,
2586 resolvedspec,
2587 spatial_terms,
2588 kappa_options,
2589 )
2590 }
2591 .map_err(EstimationError::BasisError)?;
2592 let mut log_kappa_lower = log_kappa_lower;
2593 let mut log_kappa_upper = log_kappa_upper;
2594 for &(slot, kappa) in &cc_profiled_values {
2595 log_kappa_lower.set_scalar_slot(slot, kappa);
2596 log_kappa_upper.set_scalar_slot(slot, kappa);
2597 log::info!("[spatial-kappa] slot {slot}: profiling rho at certified kappa={kappa}");
2598 }
2599 let log_kappa0 = log_kappa0.clamp_to_bounds(&log_kappa_lower, &log_kappa_upper);
2602 let setup = ExactJointHyperSetup::new(
2603 best.fit.lambdas.mapv(f64::ln),
2604 Array1::<f64>::from_elem(rho_dim, -JOINT_RHO_BOUND),
2605 Array1::<f64>::from_elem(rho_dim, rho_upper_bound),
2606 log_kappa0,
2607 log_kappa_lower,
2608 log_kappa_upper,
2609 );
2610
2611 let theta0 = setup.theta0();
2612 let lower = setup.lower();
2613 let upper = setup.upper();
2614
2615 let kind = if use_aniso {
2627 SpatialHyperKind::Anisotropic
2628 } else {
2629 SpatialHyperKind::Isotropic
2630 };
2631 let (theta_star, joint_final_value, kappa_timing) = run_exact_joint_spatial_optimization(
2632 kind,
2633 data,
2634 y,
2635 weights,
2636 offset,
2637 resolvedspec,
2638 &best.design,
2639 family.clone(),
2640 options,
2641 spatial_terms,
2642 &dims_per_term,
2643 &theta0,
2644 &lower,
2645 &upper,
2646 rho_dim,
2647 kappa_options,
2648 )?;
2649
2650 let baseline_score = fit_score(&best.fit);
2651
2652 let accept_tol = options.tol.max(1e-8 * baseline_score.abs()).max(1e-12);
2657 if joint_final_value > baseline_score + accept_tol {
2658 return Err(EstimationError::RemlOptimizationFailed(format!(
2659 "exact joint spatial optimization failed its objective-monotonicity certificate: \
2660 initial={baseline_score:.6e}, final={joint_final_value:.6e}, \
2661 acceptance_tolerance={accept_tol:.3e}, theta_checkpoint={:?}",
2662 theta_star.to_vec(),
2663 )));
2664 }
2665
2666 let selected_lambdas = Array1::from_vec(
2667 gam_problem::checked_exp_log_strengths(
2668 theta_star.slice(s![..rho_dim]).iter().copied(),
2669 )
2670 .map_err(|error| {
2671 EstimationError::InvalidInput(format!(
2672 "selected joint spatial smoothing coordinate is outside the canonical log-strength domain: {error}"
2673 ))
2674 })?,
2675 );
2676 let log_kappa_star =
2677 SpatialLogKappaCoords::from_theta_tail_with_dims(&theta_star, rho_dim, dims_per_term);
2678 if has_constant_curvature_term {
2684 let star = log_kappa_star.as_array();
2685 let dims = log_kappa_star.dims_per_term();
2686 for (slot, &term_idx) in spatial_terms.iter().enumerate() {
2687 if constant_curvature_term_spec(resolvedspec, term_idx).is_some() {
2688 let off: usize = dims[..slot].iter().sum();
2689 log::info!(
2690 "[#1464-trace] term {term_idx}: joint solver CONVERGED ψ-tail κ = {} \
2691 (this is the optimised candidate; joint_final_value={joint_final_value})",
2692 star[off]
2693 );
2694 }
2695 }
2696 }
2697 let optimized_spec = log_kappa_star.apply_tospec(resolvedspec, spatial_terms)?;
2698 let optimized = fit_term_collection_forspecwith_heuristic_lambdas(
2699 data,
2700 y,
2701 weights,
2702 offset,
2703 &optimized_spec,
2704 selected_lambdas.as_slice(),
2705 family.clone(),
2706 options,
2707 )?;
2708
2709 let mut fit = optimized.fit;
2713 fit.reml_score = joint_final_value;
2714 let optimized_result = FittedTermCollectionWithSpec {
2715 fit,
2716 design: optimized.design,
2717 resolvedspec: optimized_spec,
2718 adaptive_diagnostics: optimized.adaptive_diagnostics,
2719 kappa_timing: Some(kappa_timing),
2720 };
2721
2722 Ok(Some(optimized_result))
2723}
2724
2725#[derive(Clone, Copy, PartialEq, Eq, Debug)]
2737enum SpatialHyperKind {
2738 Anisotropic,
2739 Isotropic,
2740}
2741
2742impl SpatialHyperKind {
2743 fn label(self) -> &'static str {
2746 match self {
2747 SpatialHyperKind::Anisotropic => "spatial-aniso-joint",
2748 SpatialHyperKind::Isotropic => "spatial-iso-joint",
2749 }
2750 }
2751
2752 fn adjective(self) -> &'static str {
2754 match self {
2755 SpatialHyperKind::Anisotropic => "anisotropic",
2756 SpatialHyperKind::Isotropic => "isotropic",
2757 }
2758 }
2759
2760 fn coord_name(self) -> &'static str {
2763 match self {
2764 SpatialHyperKind::Anisotropic => "psi",
2765 SpatialHyperKind::Isotropic => "kappa",
2766 }
2767 }
2768}
2769
2770struct SpatialFrozenGlmInputs {
2776 y: Array1<f64>,
2777 weights: Array1<f64>,
2778 offset: Array1<f64>,
2779 family: LikelihoodSpec,
2780}
2781
2782fn frozen_glm_tensor_eligible_family(family: &LikelihoodSpec) -> bool {
2799 !family.is_gaussian_identity()
2800 && matches!(
2801 &family.response,
2802 ResponseFamily::Binomial
2803 | ResponseFamily::Poisson
2804 | ResponseFamily::Gamma
2805 | ResponseFamily::NegativeBinomial { .. }
2806 )
2807}
2808
2809struct SpatialJointContext<'d> {
2810 data: ArrayView2<'d, f64>,
2811 rho_dim: usize,
2812 kind: SpatialHyperKind,
2813 cache: SingleBlockExactJointDesignCache<'d>,
2814 evaluator: gam_solve::estimate::ExternalJointHyperEvaluator<'d>,
2815 frozen_glm_inputs: Option<SpatialFrozenGlmInputs>,
2816 frozen_glm_psi_bounds: Option<(f64, f64)>,
2817 frozen_glm_tensor: Option<gam_solve::glm_sufficient_lane::FrozenWeightGramTensor>,
2818 frozen_glm_tensor_attempted: bool,
2819 frozen_glm_weight_memo: Option<(Array1<f64>, Array1<f64>)>,
2831 value_realization_failures: usize,
2841 value_evaluation_failures: usize,
2842}
2843
2844#[derive(Clone, Copy, Debug, Default)]
2845struct NfreeSkipGateStatus {
2846 shape: bool,
2847 value: bool,
2848 gradient: bool,
2849 penalty: bool,
2850 revision: bool,
2851 second_order: bool,
2852}
2853
2854impl NfreeSkipGateStatus {
2855 fn would_skip(self, require_gradient: bool) -> bool {
2856 self.shape
2857 && self.value
2858 && (!require_gradient || self.gradient)
2859 && self.penalty
2860 && self.revision
2861 && !self.second_order
2862 }
2863}
2864
2865fn nfree_skip_gate_status_from_parts(
2866 shape: bool,
2867 covers_value: bool,
2868 covers_skip: bool,
2869 covers_gradient: bool,
2870 penalty: bool,
2871 revision: bool,
2872 allow_second_order: bool,
2873 require_gradient: bool,
2874) -> NfreeSkipGateStatus {
2875 NfreeSkipGateStatus {
2876 shape,
2877 value: shape && covers_value && (!require_gradient || covers_skip),
2885 gradient: shape && (!require_gradient || covers_gradient),
2886 penalty,
2887 revision,
2888 second_order: allow_second_order,
2889 }
2890}
2891
2892impl<'d> SpatialJointContext<'d> {
2893 fn nfree_skip_gate_status(
2894 &self,
2895 theta: &Array1<f64>,
2896 allow_second_order: bool,
2897 require_gradient: bool,
2898 ) -> NfreeSkipGateStatus {
2899 let shape = theta.len() == self.rho_dim + 1;
2900 let (covers_value, covers_skip, covers_gradient) = if shape {
2901 let psi = theta[self.rho_dim];
2902 (
2903 self.evaluator.psi_gram_tensor_covers(psi),
2904 self.evaluator.psi_gram_tensor_covers_skip(psi),
2905 self.evaluator.psi_gram_tensor_covers_gradient(psi),
2906 )
2907 } else {
2908 (false, false, false)
2909 };
2910 nfree_skip_gate_status_from_parts(
2911 shape,
2912 covers_value,
2913 covers_skip,
2914 covers_gradient,
2915 self.evaluator.supports_nfree_penalty_rekey(),
2916 self.evaluator.nfree_fast_path_revision().is_some(),
2917 allow_second_order,
2918 require_gradient,
2919 )
2920 }
2921
2922 fn frozen_glm_working_state(
2923 &self,
2924 beta: &Array1<f64>,
2925 ) -> Result<Option<(Array1<f64>, Array1<f64>)>, EstimationError> {
2926 let Some(inputs) = self.frozen_glm_inputs.as_ref() else {
2927 return Ok(None);
2928 };
2929 if beta.len() != self.cache.design().design.ncols() {
2930 return Ok(None);
2931 }
2932 let mut eta = self.cache.design().design.matrixvectormultiply(beta);
2933 if eta.len() != inputs.offset.len() {
2934 crate::bail_invalid_estim!(
2935 "frozen GLM tensor warm-state row mismatch: eta={}, offset={}",
2936 eta.len(),
2937 inputs.offset.len()
2938 );
2939 }
2940 eta += &inputs.offset;
2941 let obs = evaluate_standard_familyobservations(
2942 inputs.family.clone(),
2943 None,
2944 None,
2945 None,
2946 &inputs.y,
2947 &inputs.weights,
2948 &eta,
2949 )?;
2950 let mut working_response = obs.eta.clone();
2951 for i in 0..working_response.len() {
2952 let wi = obs.fisherweight[i].max(1e-12);
2953 working_response[i] += obs.score[i] / wi;
2954 }
2955 Ok(Some((obs.fisherweight, working_response)))
2956 }
2957
2958 fn frozen_glm_trial_weights(
2967 &mut self,
2968 beta: &Array1<f64>,
2969 ) -> Result<Option<Array1<f64>>, EstimationError> {
2970 if let Some((memo_beta, memo_w)) = self.frozen_glm_weight_memo.as_ref()
2971 && memo_beta.len() == beta.len()
2972 && memo_beta
2973 .iter()
2974 .zip(beta.iter())
2975 .all(|(a, b)| a.to_bits() == b.to_bits())
2976 {
2977 return Ok(Some(memo_w.clone()));
2978 }
2979 match self.frozen_glm_working_state(beta)? {
2980 Some((current_w, _)) => {
2981 self.frozen_glm_weight_memo = Some((beta.clone(), current_w.clone()));
2982 Ok(Some(current_w))
2983 }
2984 None => Ok(None),
2985 }
2986 }
2987
2988 fn ensure_frozen_glm_tensor(
2989 &mut self,
2990 theta: &Array1<f64>,
2991 warm_beta: Option<&Array1<f64>>,
2992 ) -> Result<(), EstimationError> {
2993 if self.frozen_glm_tensor.is_some() || self.frozen_glm_tensor_attempted {
2994 return Ok(());
2995 }
2996 let Some((psi_lo, psi_hi)) = self.frozen_glm_psi_bounds else {
2997 return Ok(());
2998 };
2999 if theta.len() != self.rho_dim + 1 {
3000 self.frozen_glm_tensor_attempted = true;
3001 return Ok(());
3002 }
3003 let Some(beta) = warm_beta else {
3004 return Ok(());
3005 };
3006 let Some((frozen_w, working_z)) = self.frozen_glm_working_state(beta)? else {
3007 self.frozen_glm_tensor_attempted = true;
3008 return Ok(());
3009 };
3010 let theta_probe_base = theta.clone();
3011 let rho_dim = self.rho_dim;
3012 let Self {
3019 cache, evaluator, ..
3020 } = self;
3021 let tensor = evaluator.build_frozen_glm_gram_tensor(
3022 |psi| {
3023 let mut theta_probe = theta_probe_base.clone();
3024 theta_probe[rho_dim] = psi;
3025 cache.ensure_theta(&theta_probe)?;
3026 Ok(cache.design().design.clone())
3027 },
3028 frozen_w.view(),
3029 working_z.view(),
3030 psi_lo,
3031 psi_hi,
3032 );
3033 self.cache
3034 .ensure_theta(theta)
3035 .map_err(EstimationError::InvalidInput)?;
3036 self.frozen_glm_tensor_attempted = true;
3037 if let Some(tensor) = tensor {
3038 self.frozen_glm_tensor = Some(tensor);
3039 log::info!(
3040 "[STAGE] {} certified frozen-W GLM ψ tensor over [{psi_lo:.3}, {psi_hi:.3}]",
3041 self.kind.label(),
3042 );
3043 } else {
3044 log::info!(
3045 "[STAGE] {} frozen-W GLM ψ tensor did not certify over [{psi_lo:.3}, {psi_hi:.3}]",
3046 self.kind.label(),
3047 );
3048 }
3049 Ok(())
3050 }
3051
3052 fn stage_frozen_glm_trial_statistics(
3053 &mut self,
3054 theta: &Array1<f64>,
3055 warm_beta: Option<&Array1<f64>>,
3056 allow_gradient: bool,
3057 ) -> Result<(), EstimationError> {
3058 let kind = self.kind;
3059 let mut staged_gram: Option<Array2<f64>> = None;
3060 let mut staged_deriv: Option<(Array2<f64>, Array1<f64>)> = None;
3061 if theta.len() == self.rho_dim + 1 {
3062 let psi = theta[self.rho_dim];
3063 let tensor_covers = self
3070 .frozen_glm_tensor
3071 .as_ref()
3072 .is_some_and(|t| t.contains(psi));
3073 let current_w = if tensor_covers {
3074 match warm_beta {
3075 Some(beta) => self.frozen_glm_trial_weights(beta)?,
3076 None => None,
3077 }
3078 } else {
3079 None
3080 };
3081 if let (Some(tensor), Some(current_w)) =
3082 (self.frozen_glm_tensor.as_ref(), current_w.as_ref())
3083 {
3084 const FROZEN_GLM_WEIGHT_DRIFT_RTOL: f64 = 1e-3;
3085 if tensor.weight_drift_within(current_w.view(), FROZEN_GLM_WEIGHT_DRIFT_RTOL) {
3086 staged_gram = Some(tensor.gram_at(psi));
3087 log::debug!(
3088 "[STAGE] {} trial at psi={psi:.6}: serving frozen-W GLM \
3089 first-Fisher-step XᵀWX n-free (weight drift within tol)",
3090 kind.label(),
3091 );
3092 }
3093 if allow_gradient
3094 && tensor.contains_for_gradient(psi)
3095 && let Some((dgram_dpsi, drhs_dpsi)) =
3096 tensor.gradient_pair_if_sound(psi, current_w.view())
3097 {
3098 staged_deriv = Some((dgram_dpsi, drhs_dpsi));
3099 log::debug!(
3100 "[STAGE] {} trial at psi={psi:.6}: serving frozen-W GLM \
3101 ψ-gradient (∂G/∂ψ, ∂b/∂ψ) n-free (gradient weight drift within \
3102 tight tol); B_j stays exact",
3103 kind.label(),
3104 );
3105 }
3106 }
3107 }
3108 self.evaluator.stage_glm_first_step_gram(staged_gram);
3109 self.evaluator.stage_glm_psi_gram_deriv(staged_deriv);
3110 Ok(())
3111 }
3112
3113 fn eval_full(
3115 &mut self,
3116 theta: &Array1<f64>,
3117 order: gam_solve::rho_optimizer::OuterEvalOrder,
3118 analytic_outer_hessian_available: bool,
3119 ) -> Result<(f64, Array1<f64>, gam_problem::HessianValue), EstimationError> {
3120 use gam_solve::rho_optimizer::OuterEvalOrder;
3121 let allow_second_order = matches!(order, OuterEvalOrder::ValueGradientHessian)
3122 && analytic_outer_hessian_available;
3123 if let Some(eval) = self.cache.memoized_eval(theta) {
3124 let cached_satisfies_order = !allow_second_order || eval.2.is_analytic();
3125 if cached_satisfies_order {
3126 return Ok(eval);
3127 }
3128 }
3129 let kind = self.kind;
3130 let nfree_fast_path_revision = self.evaluator.nfree_fast_path_revision();
3166 let skip_design_realization = !allow_second_order && theta.len() == self.rho_dim + 1 && {
3167 let psi = theta[self.rho_dim];
3168 self.evaluator.psi_gram_tensor_covers(psi)
3169 && self.evaluator.psi_gram_tensor_covers_gradient(psi)
3176 && self.evaluator.psi_gram_tensor_covers_skip(psi)
3193 && self.evaluator.supports_nfree_penalty_rekey()
3198 && nfree_fast_path_revision.is_some()
3199 };
3200 if skip_design_realization {
3212 log::debug!(
3213 "[STAGE] {} eval_full at psi={:.6}: skipping n×k design re-realization \
3214 + reconditioning — criterion/gradient/inner-solve served n-free from \
3215 the certified ψ-gram tensor (GaussianFixedCache + k-space ψ-derivatives)",
3216 kind.label(),
3217 theta[self.rho_dim],
3218 );
3219 } else {
3220 self.cache
3221 .ensure_theta(theta)
3222 .map_err(EstimationError::InvalidInput)?;
3223 }
3224 let warm_beta = self.evaluator.current_beta();
3225 self.ensure_frozen_glm_tensor(theta, warm_beta.as_ref())?;
3226 self.stage_frozen_glm_trial_statistics(theta, warm_beta.as_ref(), !allow_second_order)?;
3234 let hyper_dirs = if skip_design_realization {
3241 self.cache.nfree_tensor_gradient_hyper_dirs(theta)?
3242 } else {
3243 self.cache.hyper_dirs_for_current_design(self.data, kind)?
3244 };
3245
3246 let design_revision = if skip_design_realization {
3247 nfree_fast_path_revision
3248 } else {
3249 Some(self.cache.design_revision())
3250 };
3251 if self.evaluator.supports_nfree_penalty_rekey() {
3265 match self.cache.canonical_penalties_at(theta) {
3266 Ok(penalty) => self.evaluator.stage_fast_path_penalty(Some(penalty)),
3267 Err(e) => {
3268 log::warn!(
3269 "[STAGE] {} eval_full at psi={:.6}: exact n-free S(ψ) rebuild failed \
3270 ({e}); clearing stage (eval falls to slow path)",
3271 kind.label(),
3272 theta[self.rho_dim],
3273 );
3274 self.evaluator.stage_fast_path_penalty(None);
3275 }
3276 }
3277 }
3278 let eval = evaluate_joint_reml_outer_eval_at_theta(
3285 &mut self.evaluator,
3286 self.cache.design(),
3287 theta,
3288 self.rho_dim,
3289 hyper_dirs,
3290 warm_beta.as_ref().map(|b: &Array1<f64>| b.view()),
3291 if allow_second_order {
3292 order
3293 } else {
3294 OuterEvalOrder::ValueAndGradient
3295 },
3296 design_revision,
3297 );
3298 if let Ok(ref value) = eval {
3299 self.cache.store_eval_at(theta, value.clone());
3300 }
3301 eval
3302 }
3303
3304 fn eval_efs(&mut self, theta: &Array1<f64>) -> Result<gam_problem::EfsEval, EstimationError> {
3305 self.cache
3306 .ensure_theta(theta)
3307 .map_err(EstimationError::InvalidInput)?;
3308 let kind = self.kind;
3309 let hyper_dirs = try_build_spatial_log_kappa_hyper_dirs(
3310 self.data,
3311 self.cache.spec(),
3312 self.cache.design(),
3313 &self.cache.spatial_terms,
3314 )?
3315 .ok_or_else(|| {
3316 EstimationError::InvalidInput(format!(
3317 "failed to build {} hyper_dirs for exact-joint EFS",
3318 kind.adjective(),
3319 ))
3320 })?;
3321 let design_revision = Some(self.cache.design_revision());
3322 let warm_beta = self.evaluator.current_beta();
3323 evaluate_joint_reml_efs_at_theta(
3324 &mut self.evaluator,
3325 self.cache.design(),
3326 theta,
3327 self.rho_dim,
3328 hyper_dirs,
3329 warm_beta.as_ref().map(|b: &Array1<f64>| b.view()),
3330 design_revision,
3331 )
3332 }
3333
3334 fn eval_cost(&mut self, theta: &Array1<f64>) -> f64 {
3340 if let Some(cost) = self.cache.memoized_cost(theta) {
3341 return cost;
3342 }
3343 let probe_start = std::time::Instant::now();
3358 let psi_distance = self
3359 .cache
3360 .current_theta
3361 .as_ref()
3362 .filter(|reference| reference.len() == theta.len())
3363 .map(|reference| {
3364 reference
3365 .iter()
3366 .zip(theta.iter())
3367 .map(|(a, b)| (a - b) * (a - b))
3368 .sum::<f64>()
3369 .sqrt()
3370 })
3371 .unwrap_or(f64::NAN);
3372 let nfree_fast_path_revision = self.evaluator.nfree_fast_path_revision();
3386 let skip_value_realization = theta.len() == self.rho_dim + 1 && {
3387 let psi = theta[self.rho_dim];
3388 self.evaluator.psi_gram_tensor_covers(psi)
3389 && self.evaluator.supports_nfree_penalty_rekey()
3423 && nfree_fast_path_revision.is_some()
3424 };
3425 if theta.len() == self.rho_dim + 1
3426 && self.evaluator.has_psi_gram_tensor()
3427 && !self.evaluator.psi_gram_tensor_covers(theta[self.rho_dim])
3428 {
3429 self.cache.store_cost_at(theta, f64::INFINITY);
3430 return f64::INFINITY;
3431 }
3432 if !skip_value_realization && let Err(err) = self.cache.ensure_theta(theta) {
3441 self.value_realization_failures += 1;
3442 let (theta_norm, log_kappa_norm) = kphase_log_norms(theta, self.rho_dim);
3443 if self.value_realization_failures == 1 {
3444 log::warn!(
3445 "[STAGE] {} value-probe: design realization FAILED at theta_norm={:.4e} \
3446 log_kappa_norm={:.4e} ({err}); reporting +inf to the line search, which \
3447 cannot distinguish this from genuine infeasibility (#2481). Further \
3448 occurrences are counted, not logged.",
3449 self.kind.label(),
3450 theta_norm,
3451 log_kappa_norm,
3452 );
3453 } else {
3454 log::debug!(
3455 "[STAGE] {} value-probe: design realization FAILED (occurrence {}) at \
3456 theta_norm={:.4e} log_kappa_norm={:.4e} ({err})",
3457 self.kind.label(),
3458 self.value_realization_failures,
3459 theta_norm,
3460 log_kappa_norm,
3461 );
3462 }
3463 return f64::INFINITY;
3464 }
3465 if self.evaluator.supports_nfree_penalty_rekey() {
3471 match self.cache.canonical_penalties_at(theta) {
3472 Ok(penalty) => self.evaluator.stage_fast_path_penalty(Some(penalty)),
3473 Err(_) => self.evaluator.stage_fast_path_penalty(None),
3474 }
3475 }
3476 let warm_beta = self.evaluator.current_beta();
3477 if let Err(err) = self.ensure_frozen_glm_tensor(theta, warm_beta.as_ref()) {
3478 log::warn!(
3479 "[STAGE] {} value-probe at psi={:.6}: frozen-W GLM tensor setup failed ({err}); \
3480 falling back to exact streamed Gram",
3481 self.kind.label(),
3482 if theta.len() > self.rho_dim {
3483 theta[self.rho_dim]
3484 } else {
3485 f64::NAN
3486 },
3487 );
3488 self.evaluator.stage_glm_first_step_gram(None);
3489 self.evaluator.stage_glm_psi_gram_deriv(None);
3490 } else if let Err(err) =
3491 self.stage_frozen_glm_trial_statistics(theta, warm_beta.as_ref(), false)
3492 {
3493 log::warn!(
3494 "[STAGE] {} value-probe at psi={:.6}: frozen-W GLM staging failed ({err}); \
3495 falling back to exact streamed Gram",
3496 self.kind.label(),
3497 if theta.len() > self.rho_dim {
3498 theta[self.rho_dim]
3499 } else {
3500 f64::NAN
3501 },
3502 );
3503 self.evaluator.stage_glm_first_step_gram(None);
3504 self.evaluator.stage_glm_psi_gram_deriv(None);
3505 }
3506 let design_revision = if skip_value_realization {
3507 nfree_fast_path_revision
3508 } else {
3509 Some(self.cache.design_revision())
3510 };
3511 let cost_label = self.kind.label();
3512 let result = {
3513 let design = self.cache.design();
3514 self.evaluator.evaluate_cost_only(
3515 &design.design,
3516 &design.penalties,
3517 &design.nullspace_dims,
3518 design.linear_constraints.clone(),
3519 theta,
3520 self.rho_dim,
3521 warm_beta.as_ref().map(|b: &Array1<f64>| b.view()),
3522 cost_label,
3523 design_revision,
3524 )
3525 };
3526 match result {
3527 Ok(cost) => {
3528 log::debug!(
3529 "[STAGE] {cost_label} value-probe (order=Value): elapsed={:.3}s \
3530 cost={cost:.6e} trial_theta_distance={psi_distance:.3e}",
3531 probe_start.elapsed().as_secs_f64(),
3532 );
3533 self.cache.store_cost_at(theta, cost);
3534 cost
3535 }
3536 Err(err) => {
3541 self.value_evaluation_failures += 1;
3542 let (theta_norm, log_kappa_norm) = kphase_log_norms(theta, self.rho_dim);
3543 if self.value_evaluation_failures == 1 {
3544 log::warn!(
3545 "[STAGE] {cost_label} value-probe: cost evaluation FAILED at \
3546 theta_norm={theta_norm:.4e} log_kappa_norm={log_kappa_norm:.4e} \
3547 ({err}); reporting +inf to the line search, which cannot distinguish \
3548 this from genuine infeasibility (#2481). Further occurrences are \
3549 counted, not logged.",
3550 );
3551 } else {
3552 log::debug!(
3553 "[STAGE] {cost_label} value-probe: cost evaluation FAILED (occurrence \
3554 {}) at theta_norm={theta_norm:.4e} log_kappa_norm={log_kappa_norm:.4e} \
3555 ({err})",
3556 self.value_evaluation_failures,
3557 );
3558 }
3559 f64::INFINITY
3560 }
3561 }
3562 }
3563
3564 fn reset(&mut self) {
3565 self.cache.current_theta = None;
3566 self.cache.last_eval_theta = None;
3567 self.cache.last_cost = None;
3568 self.cache.last_eval = None;
3569 }
3570}
3571
3572fn kphase_log_norms(theta: &Array1<f64>, rho_dim: usize) -> (f64, f64) {
3596 let theta_norm = theta.iter().map(|v| v * v).sum::<f64>().sqrt();
3597 let log_kappa_norm = theta
3598 .iter()
3599 .skip(rho_dim)
3600 .map(|v| v * v)
3601 .sum::<f64>()
3602 .sqrt();
3603 (theta_norm, log_kappa_norm)
3604}
3605
3606fn run_exact_joint_spatial_optimization(
3607 kind: SpatialHyperKind,
3608 data: ArrayView2<'_, f64>,
3609 y: ArrayView1<'_, f64>,
3610 weights: ArrayView1<'_, f64>,
3611 offset: ArrayView1<'_, f64>,
3612 resolvedspec: &TermCollectionSpec,
3613 baseline_design: &TermCollectionDesign,
3614 family: LikelihoodSpec,
3615 options: &FitOptions,
3616 spatial_terms: &[usize],
3617 dims_per_term: &[usize],
3618 theta0: &Array1<f64>,
3619 lower: &Array1<f64>,
3620 upper: &Array1<f64>,
3621 rho_dim: usize,
3622 kappa_options: &SpatialLengthScaleOptimizationOptions,
3623) -> Result<(Array1<f64>, f64, SpatialLengthScaleOptimizationTiming), EstimationError> {
3624 let label = kind.label();
3625 let effective_offset = baseline_design
3626 .compose_offset(offset, "spatial joint fit")
3627 .map_err(EstimationError::BasisError)?;
3628 let offset = effective_offset.view();
3629 assert!(
3631 lower.len() == theta0.len() && upper.len() == theta0.len(),
3632 "spatial hyperparameter bounds must match theta length: lower_len={}, upper_len={}, theta_len={}",
3633 lower.len(),
3634 upper.len(),
3635 theta0.len()
3636 );
3637 assert!(
3638 baseline_design.smooth.terms.len() >= spatial_terms.len(),
3639 "baseline design must have at least one smooth term per spatial term: baseline_terms={}, spatial_terms={}",
3640 baseline_design.smooth.terms.len(),
3641 spatial_terms.len()
3642 );
3643 use gam_problem::{DeclaredHessianForm, Derivative, OuterEval};
3644 use gam_solve::rho_optimizer::OuterEvalOrder;
3645
3646 let theta_dim = theta0.len();
3647 let coord_dim = theta_dim - rho_dim;
3650 let analytic_outer_hessian_available =
3655 exact_joint_spatial_outer_hessian_available(&family, baseline_design);
3656 if !analytic_outer_hessian_available {
3657 log::info!(
3658 "[{label}] analytic outer Hessian unavailable for family/design; routing without second-order geometry (coord_dim={coord_dim})"
3659 );
3660 }
3661 let mut suppress_outer_hessian_for_nfree = false;
3670
3671 log::trace!(
3672 "[{}] starting analytic optimization: rho_dim={}, coord_dim={}, dims_per_term={:?}",
3673 label,
3674 rho_dim,
3675 coord_dim,
3676 dims_per_term,
3677 );
3678
3679 let mut ctx = SpatialJointContext {
3680 data,
3681 rho_dim,
3682 kind,
3683 value_realization_failures: 0,
3684 value_evaluation_failures: 0,
3685 cache: SingleBlockExactJointDesignCache::new_with_policy(
3686 data,
3687 resolvedspec.clone(),
3688 baseline_design.clone(),
3689 spatial_terms.to_vec(),
3690 rho_dim,
3691 dims_per_term.to_vec(),
3692 &options.resource_policy,
3693 )
3694 .map_err(EstimationError::InvalidInput)?,
3695 evaluator: gam_solve::estimate::ExternalJointHyperEvaluator::new(
3696 y,
3697 weights,
3698 &baseline_design.design,
3699 offset,
3700 &baseline_design.penalties,
3701 &external_opts_for_design(&family, baseline_design, options),
3702 label,
3703 )?,
3704 frozen_glm_inputs: if coord_dim == 1 && frozen_glm_tensor_eligible_family(&family) {
3705 Some(SpatialFrozenGlmInputs {
3706 y: y.to_owned(),
3707 weights: weights.to_owned(),
3708 offset: offset.to_owned(),
3709 family: family.clone(),
3710 })
3711 } else {
3712 None
3713 },
3714 frozen_glm_psi_bounds: if coord_dim == 1 && frozen_glm_tensor_eligible_family(&family) {
3715 Some((lower[rho_dim], upper[rho_dim]))
3716 } else {
3717 None
3718 },
3719 frozen_glm_tensor: None,
3720 frozen_glm_tensor_attempted: false,
3721 frozen_glm_weight_memo: None,
3722 };
3723
3724 let mut psi_rank_stable_floor: Option<f64> = None;
3747 let mut psi_rank_stable_ceiling: Option<f64> = None;
3756 let nfree_penalty_capable =
3757 coord_dim == 1 && family.is_gaussian_identity() && ctx.cache.supports_nfree_penalty_rekey();
3758 if nfree_penalty_capable {
3759 let psi_lo = lower[rho_dim];
3760 let psi_hi = upper[rho_dim];
3761 let z = Array1::from_iter(y.iter().zip(offset.iter()).map(|(yi, oi)| yi - oi));
3762 let theta_probe_base = theta0.clone();
3763 let SpatialJointContext {
3766 cache, evaluator, ..
3767 } = &mut ctx;
3768 let attached = evaluator.build_and_set_psi_gram_tensor(
3769 |psi| {
3770 let mut theta_probe = theta_probe_base.clone();
3771 theta_probe[rho_dim] = psi;
3772 cache.ensure_theta(&theta_probe)?;
3773 Ok(cache.design().design.clone())
3774 },
3775 weights,
3776 z.view(),
3777 psi_lo,
3778 psi_hi,
3779 );
3780 if attached {
3781 log::info!(
3782 "[{label}] certified ψ-gram tensor over [{psi_lo:.3}, {psi_hi:.3}]: \
3783 in-window trials assemble Gaussian sufficient statistics n-free"
3784 );
3785 let psi_anchor = theta0[rho_dim];
3790 let psi_projector_bar = evaluator.psi_gram_projector_error_bar(psi_anchor);
3796 let psi_rank_stable_floor_raw = evaluator.psi_gram_rank_stable_floor(psi_anchor);
3799 psi_rank_stable_floor = psi_rank_stable_floor_raw
3800 .filter(|&f| f.is_finite() && f > psi_lo && f < psi_anchor);
3801 log::info!(
3802 "[KAPPA-PHASE-FLOOR] n_rows={} psi_lo={psi_lo:.6} psi_anchor={psi_anchor:.6} \
3803 rank_stable_floor={psi_rank_stable_floor_raw:?} lifted={} \
3804 projector_error_bar={psi_projector_bar:?}",
3805 data.nrows(),
3806 psi_rank_stable_floor.is_some(),
3807 );
3808 if let Some(floor) = psi_rank_stable_floor {
3809 log::info!(
3810 "[{label}] rank-stable κ-floor ψ_floor={floor:.6} > window floor \
3811 ψ_lo={psi_lo:.6}: lifting the optimizer lower bound to keep every \
3812 in-window trial on the n-free design-realization skip (#1033). The \
3813 conditioned Gram is rank-deficient below ψ_floor (longest-length-scale \
3814 radial mode collapses into the nullspace), where the skip is soundly \
3815 refused. The SEARCH is n-free — O(iters·k³) off the k-space tensor, \
3816 zero row access — but the EDGE IS NOT AN n-INVARIANT CONSTANT of the \
3817 design (#2408): the tensor is built from n rows, so its Gram is an \
3818 O(1/n) relative perturbation of the continuum Gram, which moves the \
3819 rank margin additively and displaces this root by \
3820 sup|δ margin| / inf|d margin/dψ|. A steep cliff pins it to machine \
3821 precision; a grazing crossing does not. Treat it as a clamp carrying \
3822 that transport bound, not as the n-independent answer."
3823 );
3824 }
3825 let psi_rank_stable_ceiling_raw = evaluator.psi_gram_rank_stable_ceiling(psi_anchor);
3834 psi_rank_stable_ceiling = psi_rank_stable_ceiling_raw
3835 .filter(|&c| c.is_finite() && c < psi_hi && c > psi_anchor);
3836 log::info!(
3837 "[KAPPA-PHASE-CEIL] n_rows={} psi_hi={psi_hi:.6} psi_anchor={psi_anchor:.6} \
3838 rank_stable_ceiling={psi_rank_stable_ceiling_raw:?} clamped={} \
3839 projector_error_bar={psi_projector_bar:?}",
3840 data.nrows(),
3841 psi_rank_stable_ceiling.is_some(),
3842 );
3843 if let Some(ceiling) = psi_rank_stable_ceiling {
3844 log::info!(
3845 "[{label}] rank-stable κ-ceiling ψ_ceil={ceiling:.6} < window ceiling \
3846 ψ_hi={psi_hi:.6}: clamping the optimizer upper bound to keep every \
3847 in-window trial on the n-free design-realization skip (#1033). The \
3848 conditioned Gram is rank-deficient above ψ_ceil (longest-frequency \
3849 radial mode goes collinear), where the skip is soundly refused; a \
3850 line-search overshoot there trips the O(n) reset_surface lane (and the \
3851 deficient pinning ψ it records resets the next in-band trial too)."
3852 );
3853 }
3854 if let Some(bar) = psi_projector_bar
3863 && bar > gam_solve::psi_gram_tensor::PSI_GRAM_SKIP_PROJ_ATOL
3864 {
3865 log::warn!(
3866 "[{label}] ψ-gram range projector at the anchor ψ={psi_anchor:.6} is \
3867 UNRESOLVED: Davis–Kahan bar {bar:.3e} exceeds the {:.3e} subspace \
3868 tolerance the design-revision skip gates on (#2448). The conditioned \
3869 Gram has no kept/dropped eigen-gap wide enough to decide subspace \
3870 identity at double precision here — its spectrum decays smoothly \
3871 through the rank cutoff instead of cliffing — so the skip witness \
3872 soundly refuses every trial and the n-free fast path will not fire \
3873 at all. Results are unaffected (the exact O(n) path runs); the cost \
3874 is the fast path. The lever is the geometry (basis size / centers) \
3875 or the rank cutoff, not this clamp.",
3876 gam_solve::psi_gram_tensor::PSI_GRAM_SKIP_PROJ_ATOL
3877 );
3878 }
3879 let gradient_covers_full_window = evaluator.psi_gram_tensor_covers_gradient(psi_lo)
3880 && evaluator.psi_gram_tensor_covers_gradient(psi_hi);
3881 if gradient_covers_full_window {
3882 log::info!(
3883 "[{label}] certified ψ-gram tensor gradient lane covers the full \
3884 optimizer window [{psi_lo:.3}, {psi_hi:.3}]"
3885 );
3886 } else {
3887 log::info!(
3888 "[{label}] ψ-gram tensor value lane certified, but the gradient lane \
3889 does not cover the full optimizer window [{psi_lo:.3}, {psi_hi:.3}]; \
3890 keeping exact streamed kappa routing"
3891 );
3892 }
3893 evaluator.set_supports_nfree_penalty_rekey(true);
3913 log::info!(
3914 "[{label}] exact n-free ψ-penalty re-key enabled over [{psi_lo:.3}, \
3915 {psi_hi:.3}]: in-window fast-path trials rebuild S(ψ) n-free from frozen \
3916 geometry (no reset_surface)"
3917 );
3918 } else {
3919 log::info!(
3920 "[{label}] ψ-gram tensor did not certify over [{psi_lo:.3}, {psi_hi:.3}]; \
3921 keeping the exact per-trial path"
3922 );
3923 }
3924 if attached
3945 && evaluator.psi_gram_tensor_covers_gradient(psi_lo)
3946 && evaluator.psi_gram_tensor_covers_gradient(psi_hi)
3947 && evaluator.supports_nfree_penalty_rekey()
3948 && cache.supports_nfree_gradient_only_routing()
3949 {
3950 suppress_outer_hessian_for_nfree = true;
3951 log::info!(
3952 "[{label}] n-free Gaussian ψ-lane armed; suppressing the analytic outer \
3953 Hessian and routing gradient-only (BFGS) so the κ outer loop never realizes \
3954 the O(n) second-order slab — n-independent outer loop (#1033)"
3955 );
3956 }
3957 } else if coord_dim == 1 && family.is_gaussian_identity() {
3958 log::info!(
3959 "[{label}] exact n-free ψ-penalty re-key unavailable; skipping ψ-gram tensor \
3960 attachment so value, gradient, and Hessian remain on the same exact streamed \
3961 objective"
3962 );
3963 }
3964
3965 let kphase_prime_order = OuterEvalOrder::ValueAndGradient;
3968 let kphase_prime_start = std::time::Instant::now();
3969 drop(ctx.eval_full(theta0, kphase_prime_order, analytic_outer_hessian_available)?);
3970 log::info!(
3971 "[KAPPA-PHASE-PRIME] n_rows={} order={:?} elapsed_s={:.4} slow_path_resets_total={} design_revision={}",
3972 data.nrows(),
3973 kphase_prime_order,
3974 kphase_prime_start.elapsed().as_secs_f64(),
3975 ctx.evaluator.slow_path_reset_count(),
3976 ctx.cache.design_revision(),
3977 );
3978
3979 let kphase_cost_calls = std::cell::Cell::new(0usize);
3980 let kphase_eval_calls = std::cell::Cell::new(0usize);
3981 let kphase_efs_calls = std::cell::Cell::new(0usize);
3982 let kphase_cost_total_s = std::cell::Cell::new(0.0);
3983 let kphase_eval_total_s = std::cell::Cell::new(0.0);
3984 let kphase_efs_total_s = std::cell::Cell::new(0.0);
3985 let kphase_nfree_miss_shape = std::cell::Cell::new(0u64);
3986 let kphase_nfree_miss_value = std::cell::Cell::new(0u64);
3987 let kphase_nfree_miss_gradient = std::cell::Cell::new(0u64);
3988 let kphase_nfree_miss_penalty = std::cell::Cell::new(0u64);
3989 let kphase_nfree_miss_revision = std::cell::Cell::new(0u64);
3990 let kphase_nfree_miss_second_order = std::cell::Cell::new(0u64);
3991 let kphase_nfree_miss_other = std::cell::Cell::new(0u64);
3992 let kphase_optim_start = std::time::Instant::now();
3993 let kphase_log_kappa_dim = coord_dim;
3994 let kphase_slow_resets_start = ctx.evaluator.slow_path_reset_count();
3995 let kphase_design_revision_start = ctx.cache.design_revision();
3996 let kphase_nfree_skip_touches_start = gam_solve::pirls::nfree_skip_row_element_touches();
4000
4001 let lower_effective: std::borrow::Cow<'_, Array1<f64>> = match psi_rank_stable_floor {
4008 Some(floor) if coord_dim == 1 && floor > lower[rho_dim] => {
4009 let mut lifted = lower.clone();
4010 lifted[rho_dim] = floor;
4011 std::borrow::Cow::Owned(lifted)
4012 }
4013 _ => std::borrow::Cow::Borrowed(lower),
4014 };
4015 let lower = lower_effective.as_ref();
4016
4017 let upper_effective: std::borrow::Cow<'_, Array1<f64>> = match psi_rank_stable_ceiling {
4025 Some(ceiling) if coord_dim == 1 && ceiling < upper[rho_dim] => {
4026 let mut clamped = upper.clone();
4027 clamped[rho_dim] = ceiling;
4028 std::borrow::Cow::Owned(clamped)
4029 }
4030 _ => std::borrow::Cow::Borrowed(upper),
4031 };
4032 let upper = upper_effective.as_ref();
4033
4034 let problem = exact_joint_multistart_outer_problem(
4035 theta0,
4036 lower,
4037 upper,
4038 rho_dim,
4039 coord_dim,
4040 theta_dim,
4041 Derivative::Analytic,
4042 if analytic_outer_hessian_available && !suppress_outer_hessian_for_nfree {
4043 DeclaredHessianForm::Either
4044 } else {
4045 DeclaredHessianForm::Unavailable
4050 },
4051 true,
4054 suppress_outer_hessian_for_nfree,
4065 seed_risk_profile_for_likelihood_family(&family),
4066 kappa_options.rel_tol.max(1e-6),
4067 kappa_options.max_outer_iter.max(1),
4068 Some(5.0),
4071 Some(kappa_options.log_step.clamp(0.25, 1.0)),
4073 None,
4074 Some((data.nrows(), baseline_design.design.ncols())),
4079 !constant_curvature_term_indices(resolvedspec).is_empty(),
4083 kind == SpatialHyperKind::Isotropic
4088 && constant_curvature_term_indices(resolvedspec).is_empty()
4089 && spatial_terms.iter().any(|&term_idx| {
4090 matches!(
4091 resolvedspec
4092 .smooth_terms
4093 .get(term_idx)
4094 .map(|term| &term.basis),
4095 Some(SmoothBasisSpec::Matern { .. })
4096 )
4097 }),
4098 )?;
4099
4100 let eval_outer = |ctx: &mut &mut SpatialJointContext<'_>,
4101 theta: &Array1<f64>,
4102 order: OuterEvalOrder|
4103 -> Result<OuterEval, EstimationError> {
4104 let t0 = std::time::Instant::now();
4105 let allow_second_order_for_call = matches!(order, OuterEvalOrder::ValueGradientHessian)
4106 && analytic_outer_hessian_available;
4107 let gate = ctx.nfree_skip_gate_status(theta, allow_second_order_for_call, true);
4108 let resets_before = ctx.evaluator.slow_path_reset_count();
4109 let raw = ctx.eval_full(theta, order, analytic_outer_hessian_available);
4110 let reset_delta = ctx
4111 .evaluator
4112 .slow_path_reset_count()
4113 .saturating_sub(resets_before);
4114 if reset_delta > 0 {
4115 if !gate.shape {
4116 kphase_nfree_miss_shape.set(kphase_nfree_miss_shape.get() + reset_delta);
4117 }
4118 if gate.shape && !gate.value {
4119 kphase_nfree_miss_value.set(kphase_nfree_miss_value.get() + reset_delta);
4120 }
4121 if gate.shape && gate.value && !gate.gradient {
4122 kphase_nfree_miss_gradient.set(kphase_nfree_miss_gradient.get() + reset_delta);
4123 }
4124 if gate.shape && gate.value && gate.gradient && !gate.penalty {
4125 kphase_nfree_miss_penalty.set(kphase_nfree_miss_penalty.get() + reset_delta);
4126 }
4127 if gate.shape && gate.value && gate.gradient && gate.penalty && !gate.revision {
4128 kphase_nfree_miss_revision.set(kphase_nfree_miss_revision.get() + reset_delta);
4129 }
4130 if gate.shape
4131 && gate.value
4132 && gate.gradient
4133 && gate.penalty
4134 && gate.revision
4135 && gate.second_order
4136 {
4137 kphase_nfree_miss_second_order
4138 .set(kphase_nfree_miss_second_order.get() + reset_delta);
4139 }
4140 if gate.would_skip(true) {
4141 kphase_nfree_miss_other.set(kphase_nfree_miss_other.get() + reset_delta);
4142 }
4143 }
4144 let elapsed_s = t0.elapsed().as_secs_f64();
4145 kphase_eval_calls.set(kphase_eval_calls.get() + 1);
4146 kphase_eval_total_s.set(kphase_eval_total_s.get() + elapsed_s);
4147 let (theta_norm, log_kappa_norm) = kphase_log_norms(theta, rho_dim);
4148 log::info!(
4149 "[KAPPA-PHASE] phase=eval_outer call={} order={:?} design_revision={:?} theta_norm={:.4e} log_kappa_norm={:.4e} elapsed_s={:.4}",
4150 kphase_eval_calls.get(),
4151 order,
4152 Some(ctx.cache.design_revision()),
4153 theta_norm,
4154 log_kappa_norm,
4155 elapsed_s,
4156 );
4157 match raw {
4158 Ok((cost, grad, hess)) => Ok(OuterEval {
4159 cost,
4160 gradient: grad,
4161 hessian: hess,
4162 inner_beta_hint: None,
4163 }),
4164 Err(err) if is_recoverable_trial_point_error(&err) => {
4172 log::debug!(
4173 "[{label}] trial point infeasible (kernel design \
4174 not constructible at theta={theta:?}): {err}; retreating",
4175 );
4176 Ok(OuterEval::infeasible(theta_dim))
4177 }
4178 Err(err) => Err(err),
4179 }
4180 };
4181
4182 let mut obj = problem.build_objective_with_eval_order(
4183 &mut ctx,
4184 |ctx: &mut &mut SpatialJointContext<'_>, theta: &Array1<f64>| {
4185 let t0 = std::time::Instant::now();
4186 let gate = ctx.nfree_skip_gate_status(theta, false, false);
4187 let resets_before = ctx.evaluator.slow_path_reset_count();
4188 let cost = ctx.eval_cost(theta);
4189 let reset_delta = ctx
4190 .evaluator
4191 .slow_path_reset_count()
4192 .saturating_sub(resets_before);
4193 if reset_delta > 0 {
4194 if !gate.shape {
4195 kphase_nfree_miss_shape.set(kphase_nfree_miss_shape.get() + reset_delta);
4196 }
4197 if gate.shape && !gate.value {
4198 kphase_nfree_miss_value.set(kphase_nfree_miss_value.get() + reset_delta);
4199 }
4200 if gate.shape && gate.value && !gate.penalty {
4201 kphase_nfree_miss_penalty.set(kphase_nfree_miss_penalty.get() + reset_delta);
4202 }
4203 if gate.shape && gate.value && gate.penalty && !gate.revision {
4204 kphase_nfree_miss_revision.set(kphase_nfree_miss_revision.get() + reset_delta);
4205 }
4206 if gate.would_skip(false) {
4207 kphase_nfree_miss_other.set(kphase_nfree_miss_other.get() + reset_delta);
4208 }
4209 }
4210 let elapsed_s = t0.elapsed().as_secs_f64();
4211 kphase_cost_calls.set(kphase_cost_calls.get() + 1);
4212 kphase_cost_total_s.set(kphase_cost_total_s.get() + elapsed_s);
4213 let (theta_norm, log_kappa_norm) = kphase_log_norms(theta, rho_dim);
4214 log::info!(
4215 "[KAPPA-PHASE] phase=cost call={} design_revision={:?} theta_norm={:.4e} log_kappa_norm={:.4e} elapsed_s={:.4}",
4216 kphase_cost_calls.get(),
4217 Some(ctx.cache.design_revision()),
4218 theta_norm,
4219 log_kappa_norm,
4220 elapsed_s,
4221 );
4222 Ok(cost)
4223 },
4224 |ctx: &mut &mut SpatialJointContext<'_>, theta: &Array1<f64>| {
4225 eval_outer(
4226 ctx,
4227 theta,
4228 OuterEvalOrder::ValueAndGradient,
4232 )
4233 },
4234 |ctx: &mut &mut SpatialJointContext<'_>, theta: &Array1<f64>, order: OuterEvalOrder| {
4235 eval_outer(ctx, theta, order)
4236 },
4237 Some(|ctx: &mut &mut SpatialJointContext<'_>| {
4238 ctx.reset();
4239 }),
4240 Some(|ctx: &mut &mut SpatialJointContext<'_>, theta: &Array1<f64>| {
4241 let t0 = std::time::Instant::now();
4242 let eval = ctx.eval_efs(theta);
4243 let elapsed_s = t0.elapsed().as_secs_f64();
4244 kphase_efs_calls.set(kphase_efs_calls.get() + 1);
4245 kphase_efs_total_s.set(kphase_efs_total_s.get() + elapsed_s);
4246 let (theta_norm, log_kappa_norm) = kphase_log_norms(theta, rho_dim);
4247 log::info!(
4248 "[KAPPA-PHASE] phase=efs call={} design_revision={:?} theta_norm={:.4e} log_kappa_norm={:.4e} elapsed_s={:.4}",
4249 kphase_efs_calls.get(),
4250 Some(ctx.cache.design_revision()),
4251 theta_norm,
4252 log_kappa_norm,
4253 elapsed_s,
4254 );
4255 eval
4256 }),
4257 );
4258
4259 let run_label = match kind {
4260 SpatialHyperKind::Anisotropic => "aniso-psi joint REML",
4261 SpatialHyperKind::Isotropic => "iso-kappa joint REML",
4262 };
4263 let result = problem.run(&mut obj, run_label)?;
4264 if !result.converged {
4265 crate::bail_invalid_estim!(
4266 "{} did not converge after {} iterations (final_objective={:.6e}, final_grad_norm={})",
4267 run_label,
4268 result.iterations,
4269 result.final_value,
4270 result.final_grad_norm_report(),
4271 );
4272 }
4273 drop(obj);
4274 let kphase_total_s = kphase_optim_start.elapsed().as_secs_f64();
4275 let kphase_slow_resets = ctx
4276 .evaluator
4277 .slow_path_reset_count()
4278 .saturating_sub(kphase_slow_resets_start);
4279 let kphase_design_revision_delta = ctx
4280 .cache
4281 .design_revision()
4282 .saturating_sub(kphase_design_revision_start);
4283 let kphase_nfree_skip_touches = gam_solve::pirls::nfree_skip_row_element_touches()
4284 .saturating_sub(kphase_nfree_skip_touches_start);
4285 log::info!(
4286 "[KAPPA-PHASE-SUMMARY] n_rows={} log_kappa_dim={} n_cost={} cost_total_s={:.4} n_eval={} eval_total_s={:.4} n_efs={} efs_total_s={:.4} value_realization_failures={} value_evaluation_failures={} slow_path_resets={} design_revision_delta={} nfree_skip_row_touches={} nfree_miss_shape={} nfree_miss_value={} nfree_miss_gradient={} nfree_miss_penalty={} nfree_miss_revision={} nfree_miss_second_order={} nfree_miss_other={} optim_total_s={:.4}",
4287 data.nrows(),
4288 kphase_log_kappa_dim,
4289 kphase_cost_calls.get(),
4290 kphase_cost_total_s.get(),
4291 kphase_eval_calls.get(),
4292 kphase_eval_total_s.get(),
4293 kphase_efs_calls.get(),
4294 kphase_efs_total_s.get(),
4295 ctx.value_realization_failures,
4296 ctx.value_evaluation_failures,
4297 kphase_slow_resets,
4298 kphase_design_revision_delta,
4299 kphase_nfree_skip_touches,
4300 kphase_nfree_miss_shape.get(),
4301 kphase_nfree_miss_value.get(),
4302 kphase_nfree_miss_gradient.get(),
4303 kphase_nfree_miss_penalty.get(),
4304 kphase_nfree_miss_revision.get(),
4305 kphase_nfree_miss_second_order.get(),
4306 kphase_nfree_miss_other.get(),
4307 kphase_total_s,
4308 );
4309 let timing = SpatialLengthScaleOptimizationTiming {
4310 log_kappa_dim: kphase_log_kappa_dim,
4311 cost_calls: kphase_cost_calls.get(),
4312 cost_total_s: kphase_cost_total_s.get(),
4313 eval_calls: kphase_eval_calls.get(),
4314 eval_total_s: kphase_eval_total_s.get(),
4315 efs_calls: kphase_efs_calls.get(),
4316 efs_total_s: kphase_efs_total_s.get(),
4317 slow_path_resets: kphase_slow_resets,
4318 design_revision_delta: kphase_design_revision_delta,
4319 nfree_skip_row_touches: kphase_nfree_skip_touches,
4320 nfree_miss_shape: kphase_nfree_miss_shape.get(),
4321 nfree_miss_value: kphase_nfree_miss_value.get(),
4322 nfree_miss_gradient: kphase_nfree_miss_gradient.get(),
4323 nfree_miss_penalty: kphase_nfree_miss_penalty.get(),
4324 nfree_miss_revision: kphase_nfree_miss_revision.get(),
4325 nfree_miss_second_order: kphase_nfree_miss_second_order.get(),
4326 nfree_miss_other: kphase_nfree_miss_other.get(),
4327 optim_total_s: kphase_total_s,
4328 };
4329 log::trace!(
4330 "[{}] converged in {} iterations, final_value={:.6e}, grad_norm={}",
4331 label,
4332 result.iterations,
4333 result.final_value,
4334 result.final_grad_norm_report(),
4335 );
4336 let theta_star = result.rho;
4340 Ok((theta_star, result.final_value, timing))
4341}
4342
4343fn set_single_term_spatial_length_scale(
4347 term: &mut SmoothTermSpec,
4348 length_scale: f64,
4349) -> Result<(), EstimationError> {
4350 match &mut term.basis {
4351 SmoothBasisSpec::ThinPlate { spec, .. } => {
4352 spec.length_scale = length_scale;
4353 Ok(())
4354 }
4355 SmoothBasisSpec::Matern { spec, .. } => {
4356 spec.length_scale.set_resolved(length_scale);
4357 Ok(())
4358 }
4359 SmoothBasisSpec::Duchon { spec, .. } => {
4360 spec.length_scale = Some(length_scale);
4361 Ok(())
4362 }
4363 _ => Err(EstimationError::InvalidInput(format!(
4364 "term '{}' does not expose a spatial length scale",
4365 term.name
4366 ))),
4367 }
4368}
4369
4370fn set_single_term_spatial_aniso_log_scales(
4374 term: &mut SmoothTermSpec,
4375 eta: Vec<f64>,
4376) -> Result<(), EstimationError> {
4377 let eta = center_aniso_log_scales(&eta);
4378 match &mut term.basis {
4379 SmoothBasisSpec::Matern { spec, .. } => {
4380 spec.aniso_log_scales = Some(eta);
4381 Ok(())
4382 }
4383 SmoothBasisSpec::Duchon { spec, .. } => {
4384 spec.aniso_log_scales = Some(eta);
4385 Ok(())
4386 }
4387 _ => Err(EstimationError::InvalidInput(format!(
4388 "term '{}' does not support aniso_log_scales",
4389 term.name
4390 ))),
4391 }
4392}
4393
4394pub fn get_constant_curvature_kappa(spec: &TermCollectionSpec, term_idx: usize) -> Option<f64> {
4413 constant_curvature_term_spec(spec, term_idx).map(|cc| cc.kappa)
4414}
4415
4416pub fn constant_curvature_kappa_is_fixed(spec: &TermCollectionSpec, term_idx: usize) -> bool {
4423 constant_curvature_term_spec(spec, term_idx).is_some_and(|cc| cc.kappa_fixed)
4424}
4425
4426pub fn constant_curvature_term_indices(spec: &TermCollectionSpec) -> Vec<usize> {
4428 (0..spec.smooth_terms.len())
4429 .filter(|&idx| constant_curvature_term_spec(spec, idx).is_some())
4430 .collect()
4431}
4432
4433#[derive(Debug, Clone)]
4434struct SingleSmoothTermRealization {
4435 design_local: DesignMatrix,
4436 term: SmoothTerm,
4437 dropped_penaltyinfo: Vec<DroppedPenaltyBlockInfo>,
4438}
4439
4440fn wrap_local_build_as_realization(
4447 mut local: LocalSmoothTermBuild,
4448 termspec: &SmoothTermSpec,
4449) -> Result<SingleSmoothTermRealization, String> {
4450 let p_local = local.dim;
4451 let lb_local = if local.box_reparam {
4452 shape_lower_bounds_local(termspec.shape, p_local)
4453 } else {
4454 None
4455 };
4456
4457 let dropped_penaltyinfo = local
4458 .dropped_penalties
4459 .iter()
4460 .map(|info| DroppedPenaltyBlockInfo {
4461 termname: Some(termspec.name.clone()),
4462 penalty: info.clone(),
4463 })
4464 .collect();
4465
4466 let applied_rotation: Option<gam_terms::basis::JointNullRotation> = match (
4470 local.joint_null_rotation.take(),
4471 lb_local.is_some(),
4472 local.linear_constraints.is_some(),
4473 ) {
4474 (Some(rot), false, false) => {
4475 let q = &rot.rotation;
4476 local.design =
4477 apply_smooth_transform_to_design(local.design.clone(), q, &termspec.name).map_err(
4478 |e| {
4479 format!(
4480 "joint-null absorption rotation failed for term '{}': {}",
4481 termspec.name, e
4482 )
4483 },
4484 )?;
4485 for penalty in &mut local.active_penalties {
4486 let qt_s = gam_linalg::faer_ndarray::fast_atb(q, &penalty.matrix);
4487 penalty.matrix = gam_linalg::faer_ndarray::fast_ab(&qt_s, q);
4488 penalty.null_eigenvectors = penalty
4489 .null_eigenvectors
4490 .as_ref()
4491 .map(|basis| gam_linalg::faer_ndarray::fast_atb(q, basis));
4492 penalty.op = None;
4493 penalty.info.kronecker_factors = None;
4494 }
4495 local.kronecker_factored = None;
4496 Some(rot)
4497 }
4498 (Some(_), _, _) => None,
4499 (None, _, _) => None,
4500 };
4501
4502 let smooth_term = SmoothTerm {
4503 name: termspec.name.clone(),
4504 coeff_range: 0..p_local,
4505 shape: termspec.shape,
4506 active_penalties: local.active_penalties.clone(),
4507 dropped_penalties: local.dropped_penalties.clone(),
4508 metadata: local.metadata.clone(),
4509 lower_bounds_local: lb_local,
4510 linear_constraints_local: local.linear_constraints.clone(),
4511 kronecker_factored: local.kronecker_factored.take(),
4512 joint_null_rotation: applied_rotation,
4513 unabsorbed_global_orthogonality: None,
4516 };
4517
4518 Ok(SingleSmoothTermRealization {
4519 design_local: local.design,
4520 term: smooth_term,
4521 dropped_penaltyinfo,
4522 })
4523}
4524
4525fn freeze_geometry_from_metadata(
4536 termspec: &SmoothTermSpec,
4537 metadata: &BasisMetadata,
4538) -> Option<SmoothTermSpec> {
4539 let mut frozen = termspec.clone();
4540 match (&mut frozen.basis, metadata) {
4541 (
4542 SmoothBasisSpec::Matern {
4543 spec,
4544 input_scale: spec_scale,
4545 ..
4546 },
4547 BasisMetadata::Matern {
4548 centers,
4549 input_scale: metadata_scale,
4550 identifiability_transform,
4551 ..
4552 },
4553 ) => {
4554 spec.center_strategy = CenterStrategy::UserProvided(centers.clone());
4555 *spec_scale = Some(*metadata_scale);
4556 if let Some(transform) = identifiability_transform.clone() {
4560 spec.identifiability = MaternIdentifiability::FrozenTransform { transform };
4561 }
4562 Some(frozen)
4563 }
4564 (
4565 SmoothBasisSpec::Duchon {
4566 spec,
4567 input_scale: spec_scale,
4568 ..
4569 },
4570 BasisMetadata::Duchon {
4571 centers,
4572 input_scale: metadata_scale,
4573 ..
4574 },
4575 ) => {
4576 spec.center_strategy = CenterStrategy::UserProvided(centers.clone());
4577 *spec_scale = Some(*metadata_scale);
4578 Some(frozen)
4579 }
4580 (
4581 SmoothBasisSpec::ThinPlate {
4582 spec,
4583 input_scale: spec_scale,
4584 ..
4585 },
4586 BasisMetadata::ThinPlate {
4587 centers,
4588 input_scale: metadata_scale,
4589 ..
4590 },
4591 ) => {
4592 spec.center_strategy = CenterStrategy::UserProvided(centers.clone());
4593 *spec_scale = Some(*metadata_scale);
4594 Some(frozen)
4595 }
4596 _ => None,
4599 }
4600}
4601
4602fn rebuild_smooth_auxiliary_state(
4603 smooth: &mut SmoothDesign,
4604 dropped_penaltyinfo_by_term: &[Vec<DroppedPenaltyBlockInfo>],
4605) -> Result<(), String> {
4606 if dropped_penaltyinfo_by_term.len() != smooth.terms.len() {
4607 return Err(SmoothError::dimension_mismatch(format!(
4608 "smooth dropped-penalty cache mismatch: terms={}, dropped_sets={}",
4609 smooth.terms.len(),
4610 dropped_penaltyinfo_by_term.len()
4611 ))
4612 .into());
4613 }
4614
4615 let total_p = smooth.total_smooth_cols();
4616 let mut coefficient_lower_bounds = Array1::<f64>::from_elem(total_p, f64::NEG_INFINITY);
4617 let mut any_bounds = false;
4618 let mut linear_constraintrows: Vec<Array1<f64>> = Vec::new();
4619 let mut linear_constraint_b: Vec<f64> = Vec::new();
4620
4621 for term in &smooth.terms {
4622 let range = term.coeff_range.clone();
4623 if let Some(lb_local) = term.lower_bounds_local.as_ref() {
4624 if lb_local.len() != range.len() {
4625 return Err(SmoothError::dimension_mismatch(format!(
4626 "smooth lower-bound cache mismatch for term '{}': bounds={}, coeffs={}",
4627 term.name,
4628 lb_local.len(),
4629 range.len()
4630 ))
4631 .into());
4632 }
4633 coefficient_lower_bounds
4634 .slice_mut(s![range.clone()])
4635 .assign(lb_local);
4636 any_bounds = true;
4637 }
4638 if let Some(lin_local) = term.linear_constraints_local.as_ref() {
4639 if lin_local.a.ncols() != range.len() {
4640 return Err(SmoothError::dimension_mismatch(format!(
4641 "smooth linear-constraint cache mismatch for term '{}': cols={}, coeffs={}",
4642 term.name,
4643 lin_local.a.ncols(),
4644 range.len()
4645 ))
4646 .into());
4647 }
4648 for r in 0..lin_local.a.nrows() {
4649 let mut row = Array1::<f64>::zeros(total_p);
4650 row.slice_mut(s![range.clone()]).assign(&lin_local.a.row(r));
4651 linear_constraintrows.push(row);
4652 linear_constraint_b.push(lin_local.b[r]);
4653 }
4654 }
4655 }
4656
4657 smooth.coefficient_lower_bounds = if any_bounds {
4658 Some(coefficient_lower_bounds)
4659 } else {
4660 None
4661 };
4662 smooth.linear_constraints = if linear_constraintrows.is_empty() {
4663 None
4664 } else {
4665 let mut a = Array2::<f64>::zeros((linear_constraintrows.len(), total_p));
4666 for (i, row) in linear_constraintrows.iter().enumerate() {
4667 a.row_mut(i).assign(row);
4668 }
4669 Some(LinearInequalityConstraints {
4670 a,
4671 b: Array1::from_vec(linear_constraint_b),
4672 })
4673 };
4674 smooth.dropped_penaltyinfo = dropped_penaltyinfo_by_term
4675 .iter()
4676 .flat_map(|infos| infos.iter().cloned())
4677 .collect();
4678 Ok(())
4679}
4680
4681fn rebuild_term_collection_auxiliary_state(
4682 spec: &TermCollectionSpec,
4683 design: &mut TermCollectionDesign,
4684) -> Result<(), String> {
4685 if spec.linear_terms.len() != design.linear_ranges.len() {
4686 return Err(SmoothError::dimension_mismatch(format!(
4687 "term-collection linear bookkeeping mismatch: spec_terms={}, design_ranges={}",
4688 spec.linear_terms.len(),
4689 design.linear_ranges.len()
4690 ))
4691 .into());
4692 }
4693
4694 let p_total = design.design.ncols();
4695 let smooth_start = p_total.saturating_sub(design.smooth.total_smooth_cols());
4696 let mut coefficient_lower_bounds = Array1::<f64>::from_elem(p_total, f64::NEG_INFINITY);
4697 let mut any_bounds = false;
4698 let mut linear_constraintrows: Vec<Array1<f64>> = Vec::new();
4699 let mut linear_constraint_b: Vec<f64> = Vec::new();
4700
4701 for (linear, (_, range)) in spec.linear_terms.iter().zip(design.linear_ranges.iter()) {
4702 if range.len() != 1 {
4703 return Err(SmoothError::dimension_mismatch(format!(
4704 "linear term '{}' expected one coefficient column, found {}",
4705 linear.name,
4706 range.len()
4707 ))
4708 .into());
4709 }
4710 let col = range.start;
4711 if let Some(lb) = linear.coefficient_min {
4712 let mut row = Array1::<f64>::zeros(p_total);
4713 row[col] = 1.0;
4714 linear_constraintrows.push(row);
4715 linear_constraint_b.push(lb);
4716 }
4717 if let Some(ub) = linear.coefficient_max {
4718 let mut row = Array1::<f64>::zeros(p_total);
4719 row[col] = -1.0;
4720 linear_constraintrows.push(row);
4721 linear_constraint_b.push(-ub);
4722 }
4723 }
4724
4725 if let Some(lb_smooth) = design.smooth.coefficient_lower_bounds.as_ref() {
4726 if lb_smooth.len() != design.smooth.total_smooth_cols() {
4727 return Err(SmoothError::dimension_mismatch(format!(
4728 "smooth lower-bound width mismatch: bounds={}, smooth_cols={}",
4729 lb_smooth.len(),
4730 design.smooth.total_smooth_cols()
4731 ))
4732 .into());
4733 }
4734 coefficient_lower_bounds
4735 .slice_mut(s![
4736 smooth_start..(smooth_start + design.smooth.total_smooth_cols())
4737 ])
4738 .assign(lb_smooth);
4739 any_bounds = true;
4740 }
4741 if let Some(lin_smooth) = design.smooth.linear_constraints.as_ref() {
4742 if lin_smooth.a.ncols() != design.smooth.total_smooth_cols() {
4743 return Err(SmoothError::dimension_mismatch(format!(
4744 "smooth linear-constraint width mismatch: cols={}, smooth_cols={}",
4745 lin_smooth.a.ncols(),
4746 design.smooth.total_smooth_cols()
4747 ))
4748 .into());
4749 }
4750 let mut a_global = Array2::<f64>::zeros((lin_smooth.a.nrows(), p_total));
4751 a_global
4752 .slice_mut(s![
4753 ..,
4754 smooth_start..(smooth_start + design.smooth.total_smooth_cols())
4755 ])
4756 .assign(&lin_smooth.a);
4757 for r in 0..a_global.nrows() {
4758 linear_constraintrows.push(a_global.row(r).to_owned());
4759 linear_constraint_b.push(lin_smooth.b[r]);
4760 }
4761 }
4762
4763 let lower_bound_constraints = if any_bounds {
4764 linear_constraints_from_lower_bounds_global(&coefficient_lower_bounds)
4765 } else {
4766 None
4767 };
4768 let explicit_linear_constraints = if linear_constraintrows.is_empty() {
4769 None
4770 } else {
4771 let mut a = Array2::<f64>::zeros((linear_constraintrows.len(), p_total));
4772 for (i, row) in linear_constraintrows.iter().enumerate() {
4773 a.row_mut(i).assign(row);
4774 }
4775 Some(LinearInequalityConstraints {
4776 a,
4777 b: Array1::from_vec(linear_constraint_b),
4778 })
4779 };
4780
4781 design.coefficient_lower_bounds = if any_bounds {
4782 Some(coefficient_lower_bounds)
4783 } else {
4784 None
4785 };
4786 design.linear_constraints =
4787 merge_linear_constraints_global(explicit_linear_constraints, lower_bound_constraints)
4788 .map_err(|error| error.to_string())?;
4789 design.dropped_penaltyinfo = design.smooth.dropped_penaltyinfo.clone();
4790 Ok(())
4791}
4792
4793fn theta_values_match(left: &Array1<f64>, right: &Array1<f64>) -> bool {
4794 left.len() == right.len()
4795 && left
4796 .iter()
4797 .zip(right.iter())
4798 .all(|(&l, &r)| l.to_bits() == r.to_bits())
4799}
4800
4801fn latent_values_match(left: &Array1<f64>, right: &Array1<f64>) -> bool {
4802 theta_values_match(left, right)
4803}
4804
4805fn spatial_aniso_matches(left: Option<&[f64]>, right: Option<&[f64]>) -> bool {
4806 match (left, right) {
4807 (None, None) => true,
4808 (Some(a), Some(b)) => {
4809 a.len() == b.len()
4810 && a.iter()
4811 .zip(b.iter())
4812 .all(|(&x, &y)| x.to_bits() == y.to_bits())
4813 }
4814 _ => false,
4815 }
4816}
4817
4818fn spatial_length_scale_matches(left: Option<f64>, right: Option<f64>) -> bool {
4819 match (left, right) {
4820 (None, None) => true,
4821 (Some(a), Some(b)) => a.to_bits() == b.to_bits(),
4822 _ => false,
4823 }
4824}
4825
4826struct FrozenTermCollectionIncrementalRealizer<'d> {
4827 data: ArrayView2<'d, f64>,
4828 spec: TermCollectionSpec,
4829 design: TermCollectionDesign,
4830 fixed_blocks: Vec<DesignBlock>,
4831 dropped_penaltyinfo_by_term: Vec<Vec<DroppedPenaltyBlockInfo>>,
4832 smooth_penalty_ranges: Vec<Range<usize>>,
4833 full_penalty_ranges: Vec<Range<usize>>,
4834 basisworkspace: gam_terms::basis::BasisWorkspace,
4838 spatial_realization_geometry: Vec<Option<SmoothTermSpec>>,
4851 design_revision: u64,
4857}
4858
4859impl<'d> std::fmt::Debug for FrozenTermCollectionIncrementalRealizer<'d> {
4860 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
4861 f.debug_struct("FrozenTermCollectionIncrementalRealizer")
4862 .field("data_shape", &(self.data.nrows(), self.data.ncols()))
4863 .field("fixed_blocks", &self.fixed_blocks.len())
4864 .finish_non_exhaustive()
4865 }
4866}
4867
4868fn emitted_smooth_penalty_ranges(
4877 design: &TermCollectionDesign,
4878) -> Result<(Vec<Range<usize>>, Vec<Range<usize>>), String> {
4879 let leading = design.leading_penalty_blocks_before_smooth();
4880 let mut smooth_penalty_ranges = Vec::with_capacity(design.smooth.terms.len());
4881 let mut full_penalty_ranges = Vec::with_capacity(design.smooth.terms.len());
4882 let mut smooth_cursor = 0usize;
4883 for term_idx in 0..design.smooth.terms.len() {
4884 let full_range = design.smooth_term_penalty_range(term_idx)?;
4885 match full_range {
4886 Some(full_range) => {
4887 let local_start = full_range.start.checked_sub(leading).ok_or_else(|| {
4888 "incremental realizer smooth penalty range precedes the emitted smooth prefix"
4889 .to_string()
4890 })?;
4891 let local_end = full_range.end.checked_sub(leading).ok_or_else(|| {
4892 "incremental realizer smooth penalty range precedes the emitted smooth prefix"
4893 .to_string()
4894 })?;
4895 if local_start != smooth_cursor {
4896 return Err(format!(
4897 "incremental realizer non-contiguous emitted smooth layout at term {term_idx}: expected local start {smooth_cursor}, got {local_start}"
4898 ));
4899 }
4900 smooth_cursor = local_end;
4901 smooth_penalty_ranges.push(local_start..local_end);
4902 full_penalty_ranges.push(full_range);
4903 }
4904 None => {
4905 smooth_penalty_ranges.push(smooth_cursor..smooth_cursor);
4906 let global_cursor = leading.checked_add(smooth_cursor).ok_or_else(|| {
4907 "incremental realizer empty smooth penalty range overflow".to_string()
4908 })?;
4909 full_penalty_ranges.push(global_cursor..global_cursor);
4910 }
4911 }
4912 }
4913 if smooth_cursor != design.smooth.penalties.len() {
4914 return Err(format!(
4915 "incremental realizer smooth penalty mismatch: ranged={}, actual={}",
4916 smooth_cursor,
4917 design.smooth.penalties.len()
4918 ));
4919 }
4920 Ok((smooth_penalty_ranges, full_penalty_ranges))
4921}
4922
4923impl<'d> FrozenTermCollectionIncrementalRealizer<'d> {
4924 fn new(
4925 data: ArrayView2<'d, f64>,
4926 spec: TermCollectionSpec,
4927 design: TermCollectionDesign,
4928 ) -> Result<Self, String> {
4929 let policy = gam_runtime::resource::ResourcePolicy::default_library();
4930 Self::new_with_policy(data, spec, design, &policy)
4931 }
4932
4933 fn new_with_policy(
4934 data: ArrayView2<'d, f64>,
4935 spec: TermCollectionSpec,
4936 design: TermCollectionDesign,
4937 policy: &gam_runtime::resource::ResourcePolicy,
4938 ) -> Result<Self, String> {
4939 if spec.smooth_terms.len() != design.smooth.terms.len() {
4940 return Err(SmoothError::dimension_mismatch(format!(
4941 "incremental realizer smooth term mismatch: spec_terms={}, design_terms={}",
4942 spec.smooth_terms.len(),
4943 design.smooth.terms.len()
4944 ))
4945 .into());
4946 }
4947
4948 let (smooth_penalty_ranges, full_penalty_ranges) = emitted_smooth_penalty_ranges(&design)?;
4953 let spec = freeze_term_collection_from_design(&spec, &design)
4964 .map_err(|e| format!("failed to freeze incremental replay specification: {e}"))?;
4965 let fixed_blocks = build_term_collection_fixed_blocks(data, &spec)
4966 .map_err(|e| format!("failed to cache fixed term-collection blocks: {e}"))?;
4967
4968 let dropped_penaltyinfo_by_term: Vec<Vec<DroppedPenaltyBlockInfo>> = design
4982 .smooth
4983 .terms
4984 .iter()
4985 .map(|term| {
4986 term.dropped_penalties
4987 .iter()
4988 .cloned()
4989 .map(|penalty| DroppedPenaltyBlockInfo {
4990 termname: Some(term.name.clone()),
4991 penalty,
4992 })
4993 .collect()
4994 })
4995 .collect();
4996
4997 let geometry_slots = spec.smooth_terms.len();
4998 Ok(Self {
4999 data,
5000 spec,
5001 design,
5002 fixed_blocks,
5003 dropped_penaltyinfo_by_term,
5004 smooth_penalty_ranges,
5005 full_penalty_ranges,
5006 basisworkspace: gam_terms::basis::BasisWorkspace::with_policy(policy.clone()),
5007 spatial_realization_geometry: vec![None; geometry_slots],
5008 design_revision: 0,
5009 })
5010 }
5011
5012 fn design_revision(&self) -> u64 {
5013 self.design_revision
5014 }
5015
5016 fn spec(&self) -> &TermCollectionSpec {
5017 &self.spec
5018 }
5019
5020 fn design(&self) -> &TermCollectionDesign {
5021 &self.design
5022 }
5023
5024 fn supports_nfree_penalty_rekey(&self, spatial_terms: &[usize]) -> bool {
5065 if spatial_terms.len() != 1 {
5066 return false;
5067 }
5068 let term_idx = spatial_terms[0];
5069 matches!(
5070 self.design.smooth.terms.get(term_idx).map(|t| &t.metadata),
5071 Some(BasisMetadata::Duchon { .. } | BasisMetadata::ThinPlate { .. })
5072 )
5073 }
5074
5075 fn supports_nfree_gradient_only_routing(&self, spatial_terms: &[usize]) -> bool {
5084 if spatial_terms.len() != 1 {
5085 return false;
5086 }
5087 let term_idx = spatial_terms[0];
5088 matches!(
5089 self.design.smooth.terms.get(term_idx).map(|t| &t.metadata),
5090 Some(BasisMetadata::Duchon { .. } | BasisMetadata::ThinPlate { .. })
5091 )
5092 }
5093
5094 fn canonical_penalties_at_psi(
5107 &mut self,
5108 spatial_terms: &[usize],
5109 psi: &[f64],
5110 ) -> Result<(Vec<gam_terms::construction::CanonicalPenalty>, Vec<usize>), String> {
5111 if spatial_terms.len() != 1 {
5112 return Err(format!(
5113 "n-free penalty re-key requires exactly one spatial term, found {}",
5114 spatial_terms.len()
5115 ));
5116 }
5117 let term_idx = spatial_terms[0];
5118 let (ls_opt, aniso_from_psi) = spatial_term_psi_to_length_scale_and_aniso(psi);
5124 let termspec =
5127 self.spec.smooth_terms.get(term_idx).ok_or_else(|| {
5128 format!("spatial term {term_idx} out of range for n-free penalty")
5129 })?;
5130 let term = self
5131 .design
5132 .smooth
5133 .terms
5134 .get(term_idx)
5135 .ok_or_else(|| format!("realized smooth term {term_idx} out of range"))?;
5136 let p_total = self.design.design.ncols();
5139 let (locals, nullspace_dims): (Vec<Array2<f64>>, Vec<usize>) = match &term.metadata {
5140 BasisMetadata::Duchon {
5141 centers,
5142 identifiability_transform,
5143 operator_collocation_points,
5144 power,
5145 nullspace_order,
5146 aniso_log_scales,
5147 input_scale,
5148 radial_reparam,
5149 ..
5150 } => {
5151 let operator_penalties = match &termspec.basis {
5152 SmoothBasisSpec::Duchon { spec, .. } => spec.operator_penalties.clone(),
5153 _ => gam_terms::basis::DuchonOperatorPenaltySpec::default(),
5154 };
5155 let effective_ls =
5162 ls_opt.map(|length| input_scale.to_standardized_units(length));
5163 gam_terms::basis::duchon_penalties_at_length_scale(
5164 centers.view(),
5165 identifiability_transform.as_ref(),
5166 operator_collocation_points.as_ref().map(|p| p.view()),
5167 &operator_penalties,
5168 *power,
5169 *nullspace_order,
5170 aniso_log_scales.as_deref(),
5171 radial_reparam.as_ref(),
5172 effective_ls,
5173 &mut self.basisworkspace,
5174 )
5175 .map_err(|e| e.to_string())?
5176 }
5177 BasisMetadata::Matern {
5178 centers,
5179 periodic,
5180 nu,
5181 include_intercept,
5182 identifiability_transform,
5183 aniso_log_scales,
5184 input_scale,
5185 ..
5186 } => {
5187 let ls = ls_opt.ok_or_else(|| {
5194 "Matérn n-free penalty re-key requires a finite length-scale".to_string()
5195 })?;
5196 let effective_ls = input_scale.to_standardized_units(ls);
5197 let aniso_for_penalty = aniso_from_psi.as_deref().or(aniso_log_scales.as_deref());
5198 let filtered = matern_operator_penalty_triplet_at_length_scale(
5209 centers.view(),
5210 periodic.as_deref(),
5211 identifiability_transform.as_ref(),
5212 *nu,
5213 *include_intercept,
5214 aniso_for_penalty,
5215 effective_ls,
5216 )
5217 .map_err(|e| e.to_string())?;
5218 let locals = filtered
5219 .active
5220 .iter()
5221 .map(|penalty| penalty.matrix.clone())
5222 .collect();
5223 let nullspace_dims = filtered
5224 .active
5225 .iter()
5226 .map(|penalty| penalty.nullity)
5227 .collect();
5228 (locals, nullspace_dims)
5229 }
5230 BasisMetadata::ThinPlate {
5231 centers,
5232 identifiability_transform,
5233 radial_reparam,
5234 ..
5235 } => {
5236 let ls = ls_opt.ok_or_else(|| {
5237 "thin-plate n-free penalty re-key requires a finite length-scale".to_string()
5238 })?;
5239 let double_penalty = match &termspec.basis {
5240 SmoothBasisSpec::ThinPlate { spec, .. } => spec.double_penalty,
5241 _ => false,
5242 };
5243 gam_terms::basis::thin_plate_penalties_at_length_scale(
5244 centers.view(),
5245 identifiability_transform.as_ref(),
5246 radial_reparam.as_ref(),
5247 ls,
5248 double_penalty,
5249 &mut self.basisworkspace,
5250 )
5251 .map_err(|e| e.to_string())?
5252 }
5253 other => {
5254 return Err(format!(
5255 "n-free penalty re-key unsupported for basis metadata {:?}",
5256 std::mem::discriminant(other)
5257 ));
5258 }
5259 };
5260 let templates = &self.design.penalties;
5265 if templates.len() != locals.len() {
5266 return Err(format!(
5267 "n-free penalty re-key produced {} blocks but the frozen design carries {} \
5268 — penalty topology is not ψ-stable",
5269 locals.len(),
5270 templates.len()
5271 ));
5272 }
5273 let specs: Vec<gam_solve::estimate::PenaltySpec> = templates
5274 .iter()
5275 .zip(locals.into_iter())
5276 .map(|(tmpl, local)| gam_solve::estimate::PenaltySpec::Block {
5277 local,
5278 col_range: tmpl.col_range.clone(),
5279 prior_mean: tmpl.prior_mean.clone(),
5280 structure_hint: tmpl.structure_hint.clone(),
5281 op: tmpl.op.clone(),
5282 })
5283 .collect();
5284 gam_terms::construction::canonicalize_penalty_specs(
5285 &specs,
5286 &nullspace_dims,
5287 p_total,
5288 "nfree-psi-penalty",
5289 )
5290 .map_err(|e| e.to_string())
5291 }
5292
5293 fn canonical_penalty_derivatives_at_psi(
5294 &mut self,
5295 spatial_terms: &[usize],
5296 psi: &[f64],
5297 ) -> Result<(Range<usize>, usize, Vec<Array2<f64>>), String> {
5298 if spatial_terms.len() != 1 {
5299 return Err(format!(
5300 "n-free penalty derivative re-key requires exactly one spatial term, found {}",
5301 spatial_terms.len()
5302 ));
5303 }
5304 let term_idx = spatial_terms[0];
5305 let (ls_opt, aniso_from_psi) = spatial_term_psi_to_length_scale_and_aniso(psi);
5306 let termspec = self.spec.smooth_terms.get(term_idx).ok_or_else(|| {
5307 format!("spatial term {term_idx} out of range for n-free penalty derivative")
5308 })?;
5309 let term = self
5310 .design
5311 .smooth
5312 .terms
5313 .get(term_idx)
5314 .ok_or_else(|| format!("realized smooth term {term_idx} out of range"))?;
5315 let p_total = self.design.design.ncols();
5316 let smooth_start = p_total.saturating_sub(self.design.smooth.total_smooth_cols());
5317 let global_range =
5318 (smooth_start + term.coeff_range.start)..(smooth_start + term.coeff_range.end);
5319
5320 let locals = match &term.metadata {
5321 BasisMetadata::Duchon {
5322 centers,
5323 identifiability_transform,
5324 operator_collocation_points,
5325 power,
5326 nullspace_order,
5327 aniso_log_scales,
5328 input_scale,
5329 radial_reparam,
5330 ..
5331 } => {
5332 let mut spec = match &termspec.basis {
5333 SmoothBasisSpec::Duchon { spec, .. } => spec.clone(),
5334 _ => {
5335 return Err(
5336 "Duchon n-free penalty derivative requires a Duchon term spec"
5337 .to_string(),
5338 );
5339 }
5340 };
5341 let effective_ls =
5342 ls_opt.map(|length| input_scale.to_standardized_units(length));
5343 spec.length_scale = effective_ls;
5344 spec.power = *power;
5345 spec.nullspace_order = *nullspace_order;
5346 spec.aniso_log_scales = aniso_log_scales.clone();
5347 spec.radial_reparam = radial_reparam.clone();
5350 if spec.length_scale.is_none() {
5351 return Err(
5352 "Duchon n-free penalty derivative requires a hybrid length-scale"
5353 .to_string(),
5354 );
5355 }
5356 let collocation = operator_collocation_points
5357 .as_ref()
5358 .map(|points| points.view())
5359 .unwrap_or_else(|| centers.view());
5360 let (_native_sources, mut first, _native_second) =
5361 gam_terms::basis::build_duchon_native_penalty_psi_derivatives(
5362 centers.view(),
5363 &spec,
5364 identifiability_transform.as_ref(),
5365 &mut self.basisworkspace,
5366 )
5367 .map_err(|e| e.to_string())?;
5368 let (_operator_sources, operator_first, _operator_second) =
5369 gam_terms::basis::build_duchon_operator_penalty_psi_derivatives(
5370 collocation,
5371 centers.view(),
5372 &spec,
5373 identifiability_transform.as_ref(),
5374 &mut self.basisworkspace,
5375 )
5376 .map_err(|e| e.to_string())?;
5377 first.extend(operator_first);
5378 first
5379 }
5380 BasisMetadata::Matern {
5381 centers,
5382 periodic,
5383 nu,
5384 include_intercept,
5385 identifiability_transform,
5386 aniso_log_scales,
5387 input_scale,
5388 ..
5389 } => {
5390 let ls = ls_opt.ok_or_else(|| {
5391 "Matérn n-free penalty derivative requires a finite length-scale".to_string()
5392 })?;
5393 let effective_ls = input_scale.to_standardized_units(ls);
5394 let penalty_centers = gam_terms::basis::expand_periodic_centers(
5395 ¢ers.to_owned(),
5396 periodic.as_deref(),
5397 )
5398 .map_err(|e| e.to_string())?;
5399 let aniso_for_penalty = aniso_from_psi.as_deref().or(aniso_log_scales.as_deref());
5400 let (first, _second) =
5401 gam_terms::basis::build_matern_operator_penalty_psi_derivatives(
5402 penalty_centers.view(),
5403 effective_ls,
5404 *nu,
5405 *include_intercept,
5406 identifiability_transform.as_ref(),
5407 aniso_for_penalty,
5408 )
5409 .map_err(|e| e.to_string())?;
5410 first
5411 }
5412 BasisMetadata::ThinPlate {
5413 centers,
5414 identifiability_transform,
5415 radial_reparam,
5416 ..
5417 } => {
5418 let ls = ls_opt.ok_or_else(|| {
5419 "thin-plate n-free penalty derivative requires a finite length-scale"
5420 .to_string()
5421 })?;
5422 let mut spec = match &termspec.basis {
5423 SmoothBasisSpec::ThinPlate { spec, .. } => spec.clone(),
5424 _ => {
5425 return Err(
5426 "thin-plate n-free penalty derivative requires a ThinPlate term spec"
5427 .to_string(),
5428 );
5429 }
5430 };
5431 spec.length_scale = ls;
5432 if spec.radial_reparam.is_none() {
5433 spec.radial_reparam = radial_reparam.clone();
5434 }
5435 let (primary, _primary_second, nullspace, _nullspace_second) =
5436 gam_terms::basis::build_thin_plate_penalty_psi_derivativeswithworkspace(
5437 centers.view(),
5438 &spec,
5439 identifiability_transform.as_ref(),
5440 &mut self.basisworkspace,
5441 )
5442 .map_err(|e| e.to_string())?;
5443 if self.design.penalties.len() > 1 {
5444 vec![primary, nullspace]
5445 } else {
5446 vec![primary]
5447 }
5448 }
5449 other => {
5450 return Err(format!(
5451 "n-free penalty derivative re-key unsupported for basis metadata {:?}",
5452 std::mem::discriminant(other)
5453 ));
5454 }
5455 };
5456 if locals.len() != self.design.penalties.len() {
5457 return Err(format!(
5458 "n-free penalty derivative re-key produced {} blocks but the frozen design carries {} \
5459 — penalty topology is not ψ-stable",
5460 locals.len(),
5461 self.design.penalties.len()
5462 ));
5463 }
5464 Ok((global_range, p_total, locals))
5465 }
5466
5467 fn apply_log_kappa(
5468 &mut self,
5469 log_kappa: &SpatialLogKappaCoords,
5470 term_indices: &[usize],
5471 ) -> Result<(), String> {
5472 if term_indices.len() != log_kappa.dims_per_term().len() {
5473 return Err(SmoothError::dimension_mismatch(format!(
5474 "incremental realizer log-kappa term mismatch: term_indices={}, dims_per_term={}",
5475 term_indices.len(),
5476 log_kappa.dims_per_term().len()
5477 ))
5478 .into());
5479 }
5480
5481 let mut any_changed = false;
5482 for (slot, &term_idx) in term_indices.iter().enumerate() {
5483 any_changed |= self.apply_log_kappa_to_term(term_idx, log_kappa.term_slice(slot))?;
5484 }
5485
5486 if any_changed {
5487 self.refresh_full_design_operator()?;
5488 rebuild_smooth_auxiliary_state(
5489 &mut self.design.smooth,
5490 &self.dropped_penaltyinfo_by_term,
5491 )?;
5492 rebuild_term_collection_auxiliary_state(&self.spec, &mut self.design)?;
5493 self.design_revision = self.design_revision.wrapping_add(1);
5494 }
5495 Ok(())
5496 }
5497
5498 fn apply_log_kappa_to_term(&mut self, term_idx: usize, psi: &[f64]) -> Result<bool, String> {
5499 if !spatial_term_supports_hyper_optimization(&self.spec, term_idx) {
5500 return Err(SmoothError::invalid_config(format!(
5501 "incremental realizer term {term_idx} does not expose spatial hyperparameters"
5502 ))
5503 .into());
5504 }
5505 let measure_jet_term = measure_jet_term_spec(&self.spec, term_idx).is_some();
5509 let constant_curvature_term = constant_curvature_term_spec(&self.spec, term_idx).is_some();
5513 let mut next_length_scale = None;
5514 let mut next_aniso: Option<Vec<f64>> = None;
5515 if measure_jet_term {
5516 if !set_measure_jet_psi_dials(&mut self.spec, term_idx, psi)
5517 .map_err(|e| e.to_string())?
5518 {
5519 return Ok(false);
5520 }
5521 } else if constant_curvature_term {
5522 if !set_constant_curvature_kappa(&mut self.spec, term_idx, psi)
5523 .map_err(|e| e.to_string())?
5524 {
5525 return Ok(false);
5526 }
5527 } else {
5528 let current_length_scale = get_spatial_length_scale(&self.spec, term_idx);
5529 let current_aniso = get_spatial_aniso_log_scales(&self.spec, term_idx);
5530 let (ls, eta) = spatial_term_psi_to_length_scale_and_aniso(psi);
5531 next_length_scale = ls;
5532 next_aniso = eta;
5533 let same_length = spatial_length_scale_matches(current_length_scale, next_length_scale);
5534 let same_aniso = spatial_aniso_matches(current_aniso.as_deref(), next_aniso.as_deref());
5535 if same_length && same_aniso {
5536 return Ok(false);
5537 }
5538 if let Some(length_scale) = next_length_scale {
5539 set_spatial_length_scale(&mut self.spec, term_idx, length_scale)
5540 .map_err(|e| e.to_string())?;
5541 }
5542 if let Some(eta) = next_aniso.clone() {
5543 set_spatial_aniso_log_scales(&mut self.spec, term_idx, eta)
5544 .map_err(|e| e.to_string())?;
5545 }
5546 }
5547
5548 let geometry_slot = self
5559 .spatial_realization_geometry
5560 .get(term_idx)
5561 .ok_or_else(|| format!("incremental realizer geometry slot {term_idx} out of range"))?;
5562 let mut build_spec = match geometry_slot {
5563 Some(cached) => cached.clone(),
5564 None => self
5565 .spec
5566 .smooth_terms
5567 .get(term_idx)
5568 .ok_or_else(|| format!("incremental realizer smooth term {term_idx} out of range"))?
5569 .clone(),
5570 };
5571 if measure_jet_term {
5572 set_single_term_measure_jet_psi_dials(&mut build_spec, psi)
5576 .map_err(|e| e.to_string())?;
5577 } else if constant_curvature_term {
5578 set_single_term_constant_curvature_kappa(&mut build_spec, psi)
5583 .map_err(|e| e.to_string())?;
5584 } else {
5585 if let Some(length_scale) = next_length_scale {
5586 set_single_term_spatial_length_scale(&mut build_spec, length_scale)
5587 .map_err(|e| e.to_string())?;
5588 }
5589 if let Some(eta) = next_aniso {
5590 set_single_term_spatial_aniso_log_scales(&mut build_spec, eta)
5591 .map_err(|e| e.to_string())?;
5592 }
5593 }
5594
5595 let termname = build_spec.name.clone();
5596 let local = build_single_local_smooth_term(
5597 self.data,
5598 &build_spec,
5599 &mut self.basisworkspace,
5600 )
5601 .map_err(|e| {
5602 format!(
5603 "failed to rebuild smooth term '{termname}' during incremental κ realization: {e}"
5604 )
5605 })?;
5606
5607 if self.spatial_realization_geometry[term_idx].is_none()
5612 && let Some(frozen) = freeze_geometry_from_metadata(&build_spec, &local.metadata)
5613 {
5614 if let (
5626 SmoothBasisSpec::Matern {
5627 spec: frozen_spec, ..
5628 },
5629 Some(SmoothBasisSpec::Matern {
5630 spec: live_spec, ..
5631 }),
5632 ) = (
5633 &frozen.basis,
5634 self.spec
5635 .smooth_terms
5636 .get_mut(term_idx)
5637 .map(|t| &mut t.basis),
5638 ) {
5639 live_spec.identifiability = frozen_spec.identifiability.clone();
5640 live_spec.center_strategy = frozen_spec.center_strategy.clone();
5641 }
5642 self.spatial_realization_geometry[term_idx] = Some(frozen);
5643 }
5644
5645 let realization = wrap_local_build_as_realization(local, &build_spec)?;
5646 self.replace_term_realization(term_idx, realization)?;
5647 Ok(true)
5648 }
5649
5650 fn replace_term_realization(
5651 &mut self,
5652 term_idx: usize,
5653 realization: SingleSmoothTermRealization,
5654 ) -> Result<(), String> {
5655 let t_replace = std::time::Instant::now();
5656 let SingleSmoothTermRealization {
5657 design_local,
5658 term,
5659 dropped_penaltyinfo,
5660 } = realization;
5661 let SmoothTerm {
5662 name,
5663 active_penalties,
5664 dropped_penalties,
5665 metadata,
5666 lower_bounds_local,
5667 linear_constraints_local,
5668 joint_null_rotation,
5669 ..
5670 } = term;
5671 let coeff_range = self
5672 .design
5673 .smooth
5674 .terms
5675 .get(term_idx)
5676 .ok_or_else(|| format!("incremental realizer smooth term {term_idx} out of range"))?
5677 .coeff_range
5678 .clone();
5679 if design_local.ncols() != coeff_range.len() {
5680 return Err(SmoothError::dimension_mismatch(format!(
5681 "incremental realizer width mismatch for term {}: rebuilt_cols={}, cached_cols={}",
5682 term_idx,
5683 design_local.ncols(),
5684 coeff_range.len()
5685 ))
5686 .into());
5687 }
5688 if design_local.nrows() != self.design.design.nrows() {
5689 return Err(SmoothError::dimension_mismatch(format!(
5690 "incremental realizer row mismatch for term {}: rebuilt_rows={}, design_rows={}",
5691 term_idx,
5692 design_local.nrows(),
5693 self.design.design.nrows()
5694 ))
5695 .into());
5696 }
5697
5698 let smooth_penalty_range = self
5699 .smooth_penalty_ranges
5700 .get(term_idx)
5701 .ok_or_else(|| {
5702 format!("incremental realizer missing smooth penalty range for term {term_idx}")
5703 })?
5704 .clone();
5705 let full_penalty_range = self
5706 .full_penalty_ranges
5707 .get(term_idx)
5708 .ok_or_else(|| {
5709 format!("incremental realizer missing full penalty range for term {term_idx}")
5710 })?
5711 .clone();
5712 if active_penalties.len() != smooth_penalty_range.len() {
5713 return Err(SmoothError::dimension_mismatch(format!(
5714 "incremental realizer topology changed for term '{}': active_penalties={}, cached_penalties={}",
5715 name,
5716 active_penalties.len(),
5717 smooth_penalty_range.len()
5718 ))
5719 .into());
5720 }
5721
5722 self.design.smooth.term_designs[term_idx] = design_local;
5723
5724 for (offset, active_penalty) in active_penalties.iter().enumerate() {
5725 let smooth_penalty_idx = smooth_penalty_range.start + offset;
5726 let full_penalty_idx = full_penalty_range.start + offset;
5727 let penalty_local = &active_penalty.matrix;
5728
5729 if penalty_local.nrows() != coeff_range.len()
5730 || penalty_local.ncols() != coeff_range.len()
5731 {
5732 return Err(SmoothError::dimension_mismatch(format!(
5733 "incremental realizer penalty shape mismatch for term '{}' penalty {}: \
5734 penalty is {}x{} but coeff_range has {} columns",
5735 name,
5736 offset,
5737 penalty_local.nrows(),
5738 penalty_local.ncols(),
5739 coeff_range.len()
5740 ))
5741 .into());
5742 }
5743
5744 let smooth_penalty = self
5745 .design
5746 .smooth
5747 .penalties
5748 .get_mut(smooth_penalty_idx)
5749 .ok_or_else(|| {
5750 format!(
5751 "incremental realizer smooth penalty {} out of range for term {}",
5752 smooth_penalty_idx, term_idx
5753 )
5754 })?;
5755 smooth_penalty.local.assign(penalty_local);
5758 smooth_penalty.op = active_penalty.op.clone();
5759
5760 let full_bp = self
5761 .design
5762 .penalties
5763 .get_mut(full_penalty_idx)
5764 .ok_or_else(|| {
5765 format!(
5766 "incremental realizer full penalty {} out of range for term {}",
5767 full_penalty_idx, term_idx
5768 )
5769 })?;
5770 full_bp.local.assign(penalty_local);
5773 full_bp.op = active_penalty.op.clone();
5774
5775 self.design.smooth.nullspace_dims[smooth_penalty_idx] = active_penalty.nullity;
5776 self.design.nullspace_dims[full_penalty_idx] = active_penalty.nullity;
5777
5778 self.design.smooth.penaltyinfo[smooth_penalty_idx].global_index = smooth_penalty_idx;
5779 self.design.smooth.penaltyinfo[smooth_penalty_idx].termname = Some(name.clone());
5780 self.design.smooth.penaltyinfo[smooth_penalty_idx].penalty =
5781 active_penalty.info.clone();
5782
5783 self.design.penaltyinfo[full_penalty_idx].global_index = full_penalty_idx;
5784 self.design.penaltyinfo[full_penalty_idx].termname = Some(name.clone());
5785 self.design.penaltyinfo[full_penalty_idx].penalty = active_penalty.info.clone();
5786 }
5787
5788 let target_term = self.design.smooth.terms.get_mut(term_idx).ok_or_else(|| {
5789 format!("incremental realizer smooth term {term_idx} disappeared during replacement")
5790 })?;
5791 target_term.active_penalties = active_penalties;
5792 target_term.dropped_penalties = dropped_penalties;
5793 target_term.metadata = metadata;
5794 target_term.lower_bounds_local = lower_bounds_local;
5795 target_term.linear_constraints_local = linear_constraints_local;
5796 target_term.joint_null_rotation = joint_null_rotation;
5797 self.dropped_penaltyinfo_by_term[term_idx] = dropped_penaltyinfo;
5798 log::info!(
5799 "[STAGE] smooth basis rebuild (term {}, '{}', cols={}): {:.3}s",
5800 term_idx,
5801 target_term.name,
5802 coeff_range.len(),
5803 t_replace.elapsed().as_secs_f64(),
5804 );
5805 Ok(())
5806 }
5807
5808 fn refresh_full_design_operator(&mut self) -> Result<(), String> {
5809 let mut blocks = Vec::<DesignBlock>::with_capacity(
5810 self.fixed_blocks.len() + self.design.smooth.term_designs.len(),
5811 );
5812 blocks.extend(self.fixed_blocks.iter().cloned());
5813 for term_design in &self.design.smooth.term_designs {
5814 blocks.push(DesignBlock::from(term_design));
5815 }
5816 self.design.design = assemble_term_collection_design_matrix(blocks)
5817 .map_err(|e| format!("failed to refresh term-collection design: {e}"))?;
5818 Ok(())
5819 }
5820}
5821
5822fn build_term_collection_fixed_blocks(
5823 data: ArrayView2<'_, f64>,
5824 spec: &TermCollectionSpec,
5825) -> Result<Vec<DesignBlock>, BasisError> {
5826 let mut blocks = Vec::<DesignBlock>::new();
5827 if !term_collection_has_anchored_bspline(spec) {
5828 blocks.push(DesignBlock::Intercept(data.nrows()));
5829 }
5830
5831 if !spec.linear_terms.is_empty() {
5832 let mut linear_block = Array2::<f64>::zeros((data.nrows(), spec.linear_terms.len()));
5833 for (j, linear) in spec.linear_terms.iter().enumerate() {
5834 let column = linear
5838 .realized_design_column(data)
5839 .map_err(BasisError::InvalidInput)?;
5840 linear_block.column_mut(j).assign(&column);
5841 }
5842 blocks.push(DesignBlock::Dense(
5843 gam_linalg::matrix::DenseDesignMatrix::from(linear_block),
5844 ));
5845 }
5846
5847 for term in &spec.random_effect_terms {
5848 let block = build_random_effect_block(data, term)?;
5849 let re_op = RandomEffectOperator::new(block.group_ids, block.num_groups);
5850 blocks.push(DesignBlock::RandomEffect(Arc::new(re_op)));
5851 }
5852
5853 Ok(blocks)
5854}
5855
5856pub struct SpatialLengthScaleOptimizationResult<FitOut> {
5861 pub resolved_specs: Vec<TermCollectionSpec>,
5862 pub designs: Vec<TermCollectionDesign>,
5863 pub fit: FitOut,
5864 pub certified_outer: Option<gam_solve::rho_optimizer::CertifiedOuterResult>,
5865 pub timing: Option<SpatialLengthScaleOptimizationTiming>,
5866}
5867
5868pub struct ExactJointEvaluation<M> {
5876 pub objective: f64,
5877 pub gradient: Array1<f64>,
5878 pub hessian: gam_problem::HessianValue,
5879 pub mode: M,
5880}
5881
5882pub struct ExactJointEfsEvaluation<M> {
5885 pub evaluation: gam_problem::EfsEval,
5886 pub mode: M,
5887}
5888
5889pub enum SpatialFitProvenance<'a, M> {
5890 NoOuterOptimization,
5891 Certified {
5892 outer: &'a gam_solve::rho_optimizer::CertifiedOuterResult,
5893 mode: M,
5894 },
5895}
5896
5897#[derive(Debug, Clone)]
5899pub struct ExactJointHyperSetup {
5900 rho0: Array1<f64>,
5901 rho_lower: Array1<f64>,
5902 rho_upper: Array1<f64>,
5903 log_kappa0: SpatialLogKappaCoords,
5904 log_kappa_lower: SpatialLogKappaCoords,
5905 log_kappa_upper: SpatialLogKappaCoords,
5906 auxiliary0: Array1<f64>,
5907 auxiliary_lower: Array1<f64>,
5908 auxiliary_upper: Array1<f64>,
5909}
5910
5911impl ExactJointHyperSetup {
5912 fn sanitize_rho_seed(
5913 rho0: Array1<f64>,
5914 rho_lower: &Array1<f64>,
5915 rho_upper: &Array1<f64>,
5916 ) -> Array1<f64> {
5917 Array1::from_iter(rho0.iter().enumerate().map(|(idx, &value)| {
5918 let lo = rho_lower[idx];
5919 let hi = rho_upper[idx];
5920 let fallback = 0.0_f64.clamp(lo, hi);
5921 if value.is_finite() {
5922 value.clamp(lo, hi)
5923 } else {
5924 fallback
5925 }
5926 }))
5927 }
5928
5929 pub(crate) fn new(
5930 rho0: Array1<f64>,
5931 rho_lower: Array1<f64>,
5932 rho_upper: Array1<f64>,
5933 log_kappa0: SpatialLogKappaCoords,
5934 log_kappa_lower: SpatialLogKappaCoords,
5935 log_kappa_upper: SpatialLogKappaCoords,
5936 ) -> Self {
5937 let rho0 = Self::sanitize_rho_seed(rho0, &rho_lower, &rho_upper);
5938 Self {
5939 rho0,
5940 rho_lower,
5941 rho_upper,
5942 log_kappa0,
5943 log_kappa_lower,
5944 log_kappa_upper,
5945 auxiliary0: Array1::zeros(0),
5946 auxiliary_lower: Array1::zeros(0),
5947 auxiliary_upper: Array1::zeros(0),
5948 }
5949 }
5950
5951 pub(crate) fn with_auxiliary(
5952 mut self,
5953 auxiliary0: Array1<f64>,
5954 auxiliary_lower: Array1<f64>,
5955 auxiliary_upper: Array1<f64>,
5956 ) -> Self {
5957 assert_eq!(
5958 auxiliary0.len(),
5959 auxiliary_lower.len(),
5960 "auxiliary lower bound length mismatch"
5961 );
5962 assert_eq!(
5963 auxiliary0.len(),
5964 auxiliary_upper.len(),
5965 "auxiliary upper bound length mismatch"
5966 );
5967 self.auxiliary0 = Self::sanitize_rho_seed(auxiliary0, &auxiliary_lower, &auxiliary_upper);
5968 self.auxiliary_lower = auxiliary_lower;
5969 self.auxiliary_upper = auxiliary_upper;
5970 self
5971 }
5972
5973 pub(crate) fn rho_dim(&self) -> usize {
5974 self.rho0.len()
5975 }
5976
5977 pub(crate) fn log_kappa_dim(&self) -> usize {
5978 self.log_kappa0.len()
5979 }
5980
5981 pub(crate) fn auxiliary_dim(&self) -> usize {
5982 self.auxiliary0.len()
5983 }
5984
5985 pub(crate) fn theta0(&self) -> Array1<f64> {
5986 let mut out =
5987 Array1::<f64>::zeros(self.rho_dim() + self.log_kappa_dim() + self.auxiliary_dim());
5988 out.slice_mut(s![..self.rho_dim()]).assign(&self.rho0);
5989 out.slice_mut(s![self.rho_dim()..self.rho_dim() + self.log_kappa_dim()])
5990 .assign(self.log_kappa0.as_array());
5991 out.slice_mut(s![self.rho_dim() + self.log_kappa_dim()..])
5992 .assign(&self.auxiliary0);
5993 out
5994 }
5995
5996 pub(crate) fn lower(&self) -> Array1<f64> {
5997 let mut out =
5998 Array1::<f64>::zeros(self.rho_dim() + self.log_kappa_dim() + self.auxiliary_dim());
5999 out.slice_mut(s![..self.rho_dim()]).assign(&self.rho_lower);
6000 out.slice_mut(s![self.rho_dim()..self.rho_dim() + self.log_kappa_dim()])
6001 .assign(self.log_kappa_lower.as_array());
6002 out.slice_mut(s![self.rho_dim() + self.log_kappa_dim()..])
6003 .assign(&self.auxiliary_lower);
6004 out
6005 }
6006
6007 pub(crate) fn upper(&self) -> Array1<f64> {
6008 let mut out =
6009 Array1::<f64>::zeros(self.rho_dim() + self.log_kappa_dim() + self.auxiliary_dim());
6010 out.slice_mut(s![..self.rho_dim()]).assign(&self.rho_upper);
6011 out.slice_mut(s![self.rho_dim()..self.rho_dim() + self.log_kappa_dim()])
6012 .assign(self.log_kappa_upper.as_array());
6013 out.slice_mut(s![self.rho_dim() + self.log_kappa_dim()..])
6014 .assign(&self.auxiliary_upper);
6015 out
6016 }
6017
6018 pub(crate) fn log_kappa_dims_per_term(&self) -> Vec<usize> {
6020 self.log_kappa0.dims_per_term().to_vec()
6021 }
6022}
6023
6024struct ExactJointDesignCache<'d> {
6030 realizers: Vec<FrozenTermCollectionIncrementalRealizer<'d>>,
6031 block_term_indices: Vec<Vec<usize>>,
6032 current_theta: Option<Array1<f64>>,
6033 last_cost: Option<f64>,
6034 last_eval: Option<(f64, Array1<f64>, gam_problem::HessianValue)>,
6035 rho_dim: usize,
6036 all_dims: Vec<usize>,
6037 log_kappa_dim: usize,
6038 block_term_counts: Vec<usize>,
6039}
6040
6041impl<'d> ExactJointDesignCache<'d> {
6042 fn new(
6043 data: ArrayView2<'d, f64>,
6044 blocks: Vec<(TermCollectionSpec, TermCollectionDesign, Vec<usize>)>,
6045 rho_dim: usize,
6046 all_dims: Vec<usize>,
6047 ) -> Result<Self, String> {
6048 let n_blocks = blocks.len();
6049 let mut realizers = Vec::with_capacity(n_blocks);
6050 let mut block_term_indices = Vec::with_capacity(n_blocks);
6051 let mut block_term_counts = Vec::with_capacity(n_blocks);
6052
6053 for (spec, design, terms) in blocks {
6054 block_term_counts.push(terms.len());
6055 block_term_indices.push(terms);
6056 realizers.push(FrozenTermCollectionIncrementalRealizer::new(
6057 data, spec, design,
6058 )?);
6059 }
6060
6061 Ok(Self {
6062 realizers,
6063 block_term_indices,
6064 current_theta: None,
6065 last_cost: None,
6066 last_eval: None,
6067 rho_dim,
6068 log_kappa_dim: all_dims.iter().sum(),
6069 all_dims,
6070 block_term_counts,
6071 })
6072 }
6073
6074 fn ensure_theta(&mut self, theta: &Array1<f64>) -> Result<(), String> {
6075 if self
6076 .current_theta
6077 .as_ref()
6078 .is_some_and(|cached| theta_values_match(cached, theta))
6079 {
6080 return Ok(());
6081 }
6082
6083 let t_ensure = std::time::Instant::now();
6084 let kappa_theta_len = self.rho_dim + self.log_kappa_dim;
6085 if theta.len() < kappa_theta_len {
6086 return Err(SmoothError::dimension_mismatch(format!(
6087 "exact-joint theta length mismatch: got {}, expected at least {} (rho_dim={}, log_kappa_dim={})",
6088 theta.len(),
6089 kappa_theta_len,
6090 self.rho_dim,
6091 self.log_kappa_dim
6092 ))
6093 .into());
6094 }
6095 let theta_kappa = theta.slice(s![..kappa_theta_len]).to_owned();
6096 let full_log_kappa = SpatialLogKappaCoords::from_theta_tail_with_dims(
6097 &theta_kappa,
6098 self.rho_dim,
6099 self.all_dims.clone(),
6100 );
6101
6102 let n = self.realizers.len();
6106 let mut remaining = full_log_kappa;
6107 for block_idx in 0..n {
6108 let count = self.block_term_counts[block_idx];
6109 if block_idx < n - 1 {
6110 let (block_lk, rest) = remaining.split_at(count);
6111 self.realizers[block_idx]
6112 .apply_log_kappa(&block_lk, &self.block_term_indices[block_idx])?;
6113 remaining = rest;
6114 } else {
6115 self.realizers[block_idx]
6117 .apply_log_kappa(&remaining, &self.block_term_indices[block_idx])?;
6118 }
6119 }
6120
6121 log::info!(
6122 "[STAGE] ensure_theta (n-block, {} blocks, {} realizers): {:.3}s",
6123 n,
6124 self.realizers.len(),
6125 t_ensure.elapsed().as_secs_f64(),
6126 );
6127 self.current_theta = Some(theta.clone());
6128 self.last_cost = None;
6129 self.last_eval = None;
6130 Ok(())
6131 }
6132
6133 impl_exact_joint_theta_memo!();
6134
6135 fn store_cost_only(&mut self, theta: &Array1<f64>, cost: f64) {
6141 if self
6142 .current_theta
6143 .as_ref()
6144 .is_some_and(|cached| theta_values_match(cached, theta))
6145 {
6146 self.last_cost = Some(cost);
6147 }
6148 }
6149
6150 fn invalidate_objective_memo(&mut self) {
6153 self.last_cost = None;
6154 self.last_eval = None;
6155 }
6156
6157 fn specs(&self) -> Vec<&TermCollectionSpec> {
6158 self.realizers.iter().map(|r| r.spec()).collect()
6159 }
6160
6161 fn designs(&self) -> Vec<&TermCollectionDesign> {
6162 self.realizers.iter().map(|r| r.design()).collect()
6163 }
6164
6165 fn design_revision(&self) -> u64 {
6175 self.realizers
6176 .iter()
6177 .fold(0u64, |acc, r| acc.wrapping_add(r.design_revision()))
6178 }
6179}
6180
6181pub(crate) fn seed_risk_profile_for_likelihood_family(
6182 family: &LikelihoodSpec,
6183) -> gam_problem::SeedRiskProfile {
6184 match &family.response {
6185 ResponseFamily::Gaussian => gam_problem::SeedRiskProfile::Gaussian,
6186 ResponseFamily::RoystonParmar => gam_problem::SeedRiskProfile::Survival,
6187 ResponseFamily::Binomial
6188 | ResponseFamily::Poisson
6189 | ResponseFamily::Tweedie { .. }
6190 | ResponseFamily::NegativeBinomial { .. }
6191 | ResponseFamily::Beta { .. }
6192 | ResponseFamily::Gamma => gam_problem::SeedRiskProfile::GeneralizedLinear,
6193 }
6194}
6195
6196fn exact_joint_seed_config(
6197 risk_profile: gam_problem::SeedRiskProfile,
6198 auxiliary_dim: usize,
6199 initial_seed_only: bool,
6200) -> gam_problem::SeedConfig {
6201 let mut config = gam_problem::SeedConfig {
6202 risk_profile,
6203 num_auxiliary_trailing: auxiliary_dim,
6204 ..Default::default()
6205 };
6206 match risk_profile {
6207 gam_problem::SeedRiskProfile::Gaussian
6208 | gam_problem::SeedRiskProfile::GaussianLocationScale => {
6209 config.max_seeds = 4;
6210 config.seed_budget = 2;
6211 }
6212 gam_problem::SeedRiskProfile::GeneralizedLinear => {
6213 config.max_seeds = 1;
6218 config.seed_budget = 1;
6219 config.screen_max_inner_iterations = 8;
6220 }
6221 gam_problem::SeedRiskProfile::Survival => {
6222 config.max_seeds = 8;
6228 config.seed_budget = 4;
6229 config.screen_max_inner_iterations = 8;
6230 }
6231 }
6232 if initial_seed_only {
6233 config.max_seeds = 1;
6240 config.seed_budget = 1;
6241 config.over_smoothing_probe_rho = None;
6242 }
6243 config
6244}
6245
6246#[cfg(test)]
6247mod exact_joint_seed_config_tests {
6248 use super::*;
6249
6250 #[test]
6251 fn exact_joint_marginal_slope_profiles_get_deeper_startup_validation() {
6252 let bms =
6253 exact_joint_seed_config(gam_problem::SeedRiskProfile::GeneralizedLinear, 2, false);
6254 assert_eq!(bms.max_seeds, 1);
6255 assert_eq!(bms.seed_budget, 1);
6256 assert_eq!(bms.screen_max_inner_iterations, 8);
6257 assert_eq!(bms.num_auxiliary_trailing, 2);
6258
6259 let survival = exact_joint_seed_config(gam_problem::SeedRiskProfile::Survival, 3, false);
6260 assert_eq!(survival.max_seeds, 8);
6261 assert_eq!(survival.seed_budget, 4);
6262 assert_eq!(survival.screen_max_inner_iterations, 8);
6263 assert_eq!(survival.num_auxiliary_trailing, 3);
6264 }
6265
6266 #[test]
6267 fn exact_joint_gaussian_keeps_tight_historical_multistart_budget() {
6268 let gaussian = exact_joint_seed_config(gam_problem::SeedRiskProfile::Gaussian, 1, false);
6269 assert_eq!(gaussian.max_seeds, 4);
6270 assert_eq!(gaussian.seed_budget, 2);
6271 assert_eq!(
6272 gaussian.screen_max_inner_iterations,
6273 gam_problem::SeedConfig::default().screen_max_inner_iterations
6274 );
6275 assert_eq!(gaussian.num_auxiliary_trailing, 1);
6276 }
6277
6278 #[test]
6279 fn certified_matern_basin_owns_the_only_joint_start() {
6280 let gaussian = exact_joint_seed_config(gam_problem::SeedRiskProfile::Gaussian, 1, true);
6281 assert_eq!(gaussian.max_seeds, 1);
6282 assert_eq!(gaussian.seed_budget, 1);
6283 assert_eq!(gaussian.over_smoothing_probe_rho, None);
6284 assert_eq!(gaussian.num_auxiliary_trailing, 1);
6285 }
6286}
6287
6288#[cfg(test)]
6289mod wood_reference_df_tests {
6290 use super::*;
6291
6292 #[test]
6298 fn edf1_equals_two_trace_minus_trace_of_square() {
6299 let f = ndarray::array![[0.9_f64, 0.0], [0.0, 0.4]];
6303 let got = wood_reference_df(Some(&f), &(0..2)).unwrap();
6304 assert!(
6305 (got - 1.63).abs() < 1e-12,
6306 "edf1 should be 2*tr - tr(F^2) = 1.63, got {got}"
6307 );
6308 let edf = 1.3;
6311 assert!(got >= edf - 1e-12, "edf1 {got} must be >= edf {edf}");
6312 }
6313
6314 #[test]
6315 fn edf1_never_collapses_below_edf_when_offdiagonals_blow_up() {
6316 let f = ndarray::array![[0.5_f64, 40.0], [40.0, 0.5]];
6323 let tr = 1.0_f64;
6324 let got = wood_reference_df(Some(&f), &(0..2)).unwrap();
6325 assert!(
6326 got >= tr - 1e-12,
6327 "edf1 must be floored at edf (=tr={tr}) even when tr(F^2) explodes, got {got}"
6328 );
6329 assert!(
6330 got.is_finite() && got > 0.0,
6331 "edf1 must stay finite/positive"
6332 );
6333 }
6334
6335 #[test]
6336 fn returns_none_on_nonpositive_or_missing_trace() {
6337 assert!(wood_reference_df(None, &(0..2)).is_none());
6340 let zero = ndarray::array![[0.0_f64, 0.0], [0.0, 0.0]];
6342 assert!(wood_reference_df(Some(&zero), &(0..2)).is_none());
6343 let f = ndarray::array![[0.5_f64, 0.0], [0.0, 0.5]];
6345 assert!(wood_reference_df(Some(&f), &(0..5)).is_none());
6346 }
6347}
6348
6349pub(crate) fn exact_joint_multistart_outer_problem(
6350 theta0: &Array1<f64>,
6351 lower: &Array1<f64>,
6352 upper: &Array1<f64>,
6353 rho_dim: usize,
6354 auxiliary_dim: usize,
6355 n_params: usize,
6356 gradient: gam_problem::Derivative,
6357 hessian: gam_problem::DeclaredHessianForm,
6358 reserve_analytic_hessian_for_certificate: bool,
6363 disable_fixed_point: bool,
6364 risk_profile: gam_problem::SeedRiskProfile,
6365 tolerance: f64,
6366 max_iter: usize,
6367 bfgs_step_cap: Option<f64>,
6376 bfgs_step_cap_psi: Option<f64>,
6377 screening_cap: Option<Arc<AtomicUsize>>,
6378 profiled_objective_size: Option<(usize, usize)>,
6399 has_constant_curvature: bool,
6408 initial_seed_only: bool,
6413) -> Result<gam_solve::rho_optimizer::OuterProblem, EstimationError> {
6414 if rho_dim > theta0.len() {
6415 crate::bail_invalid_estim!(
6416 "exact joint outer problem declares {rho_dim} smoothing coordinates for theta length {}",
6417 theta0.len(),
6418 );
6419 }
6420 let mut seed_heuristic = theta0.to_vec();
6421 let initial_lambdas = gam_problem::checked_exp_log_strengths(
6422 theta0.iter().take(rho_dim).copied(),
6423 )
6424 .map_err(|error| {
6425 EstimationError::InvalidInput(format!(
6426 "exact joint initial smoothing coordinate is outside the canonical log-strength domain: {error}"
6427 ))
6428 })?;
6429 for (value, lambda) in seed_heuristic[..rho_dim].iter_mut().zip(initial_lambdas) {
6430 *value = lambda;
6431 }
6432 let rho_ceiling = if has_constant_curvature {
6437 gam_solve::estimate::RHO_BOUND
6438 } else {
6439 12.0
6440 };
6441 let mut problem = gam_solve::rho_optimizer::OuterProblem::new(n_params)
6442 .with_gradient(gradient)
6443 .with_hessian(hessian)
6444 .with_prefer_gradient_only(reserve_analytic_hessian_for_certificate)
6445 .with_disable_fixed_point(disable_fixed_point)
6446 .with_fallback_policy(gam_solve::rho_optimizer::FallbackPolicy::Automatic)
6456 .with_psi_dim(auxiliary_dim)
6457 .with_tolerance(tolerance)
6458 .with_max_iter(max_iter)
6459 .with_bounds(lower.clone(), upper.clone())
6460 .with_initial_rho(theta0.clone())
6461 .with_bfgs_step_cap(bfgs_step_cap)
6462 .with_bfgs_step_cap_psi(bfgs_step_cap_psi)
6463 .with_seed_config({
6464 let mut sc = exact_joint_seed_config(risk_profile, auxiliary_dim, initial_seed_only);
6465 if has_constant_curvature {
6466 sc.bounds = (sc.bounds.0, rho_ceiling);
6470 }
6485 sc
6486 })
6487 .with_rho_bound(rho_ceiling)
6488 .with_heuristic_lambdas(seed_heuristic);
6489 if let Some((n_obs, p_cols)) = profiled_objective_size {
6490 problem = problem
6495 .with_objective_scale(Some(n_obs as f64))
6496 .with_problem_size(n_obs, p_cols);
6497 }
6498 if let Some(screening_cap) = screening_cap {
6499 problem = problem
6500 .with_screening_cap(screening_cap)
6501 .with_screen_initial_rho(true);
6502 }
6503 Ok(problem)
6504}
6505
6506pub fn optimize_spatial_length_scale_exact_joint<FitOut, Mode, FitFn, ExactFn, ExactEfsFn, SeedFn>(
6507 data: ArrayView2<'_, f64>,
6508 block_specs: &[TermCollectionSpec],
6509 block_term_indices: &[Vec<usize>],
6510 kappa_options: &SpatialLengthScaleOptimizationOptions,
6511 joint_setup: &ExactJointHyperSetup,
6512 seed_risk_profile: gam_problem::SeedRiskProfile,
6513 analytic_joint_gradient_available: bool,
6514 analytic_joint_hessian_available: bool,
6515 disable_fixed_point: bool,
6516 screening_cap: Option<Arc<AtomicUsize>>,
6517 outer_derivative_policy: gam_model_api::families::custom_family::OuterDerivativePolicy,
6518 mut fit_fn: FitFn,
6519 mut exact_fn: ExactFn,
6520 mut exact_efs_fn: ExactEfsFn,
6521 mut seed_inner_beta_fn: SeedFn,
6522) -> Result<SpatialLengthScaleOptimizationResult<FitOut>, String>
6523where
6524 FitFn: FnMut(
6525 &Array1<f64>,
6526 &[TermCollectionSpec],
6527 &[TermCollectionDesign],
6528 SpatialFitProvenance<'_, Mode>,
6529 ) -> Result<FitOut, String>,
6530 ExactFn: FnMut(
6531 &Array1<f64>,
6532 &[TermCollectionSpec],
6533 &[TermCollectionDesign],
6534 gam_solve::estimate::reml::reml_outer_engine::EvalMode,
6535 &gam_problem::outer_subsample::RowSet,
6536 ) -> Result<ExactJointEvaluation<Mode>, String>,
6537 ExactEfsFn: FnMut(
6538 &Array1<f64>,
6539 &[TermCollectionSpec],
6540 &[TermCollectionDesign],
6541 &gam_problem::outer_subsample::RowSet,
6542 ) -> Result<ExactJointEfsEvaluation<Mode>, String>,
6543 SeedFn: FnMut(&Array1<f64>) -> Result<gam_solve::rho_optimizer::SeedOutcome, EstimationError>,
6544{
6545 let n_blocks = block_specs.len();
6546 if block_term_indices.len() != n_blocks {
6547 return Err(SmoothError::dimension_mismatch(format!(
6548 "block_specs ({}) and block_term_indices ({}) length mismatch",
6549 n_blocks,
6550 block_term_indices.len()
6551 ))
6552 .into());
6553 }
6554
6555 let log_kappa_dim = joint_setup.log_kappa_dim();
6556
6557 log::trace!(
6558 "[spatial-exact-joint] driver entry: aux_dim={} log_kappa_dim={} kappa_enabled={} rho_dim={} theta0_len={}",
6559 joint_setup.auxiliary_dim(),
6560 log_kappa_dim,
6561 kappa_options.enabled,
6562 joint_setup.rho_dim(),
6563 joint_setup.theta0().len()
6564 );
6565
6566 if joint_setup.auxiliary_dim() == 0 && (!kappa_options.enabled || log_kappa_dim == 0) {
6570 log::trace!(
6571 "[spatial-exact-joint] taking fast path (no outer theta optimization in this driver)"
6572 );
6573 let (designs, resolved_specs) = build_term_collection_designs_and_freeze_joint(
6574 data, block_specs,
6575 )
6576 .map_err(|e| {
6577 format!("failed to build and freeze joint block designs during exact joint kappa optimization: {e}")
6578 })?;
6579 let theta0 = joint_setup.theta0();
6580
6581 let spec_refs: Vec<TermCollectionSpec> = resolved_specs.clone();
6583 let design_refs: Vec<TermCollectionDesign> = designs.clone();
6584 let fit = fit_fn(
6585 &theta0,
6586 &spec_refs,
6587 &design_refs,
6588 SpatialFitProvenance::NoOuterOptimization,
6589 )?;
6590 return Ok(SpatialLengthScaleOptimizationResult {
6591 resolved_specs,
6592 designs,
6593 fit,
6594 certified_outer: None,
6595 timing: None,
6596 });
6597 }
6598
6599 let theta0 = joint_setup.theta0();
6603 let lower = joint_setup.lower();
6604 let upper = joint_setup.upper();
6605 if theta0.len() < log_kappa_dim || lower.len() != theta0.len() || upper.len() != theta0.len() {
6606 return Err(SmoothError::dimension_mismatch(format!(
6607 "invalid exact joint theta setup: theta0={}, lower={}, upper={}, required_log_kappa_dim={}",
6608 theta0.len(),
6609 lower.len(),
6610 upper.len(),
6611 log_kappa_dim
6612 ))
6613 .into());
6614 }
6615 let rho_dim = joint_setup.rho_dim();
6616 let all_dims = joint_setup.log_kappa_dims_per_term();
6617
6618 let (boot_designs, best_specs) = build_term_collection_designs_and_freeze_joint(
6620 data,
6621 block_specs,
6622 )
6623 .map_err(|e| {
6624 format!(
6625 "failed to build and freeze joint block designs during exact joint kappa bootstrap: {e}"
6626 )
6627 })?;
6628 let policy_hessian_form = outer_derivative_policy.declared_hessian_form();
6638 let analytic_outer_hessian_available = analytic_joint_hessian_available
6639 && matches!(
6640 policy_hessian_form,
6641 gam_problem::DeclaredHessianForm::Either
6642 | gam_problem::DeclaredHessianForm::Dense
6643 | gam_problem::DeclaredHessianForm::Operator { .. }
6644 );
6645 let theta_dim = theta0.len();
6646 let psi_dim = theta_dim - rho_dim;
6647
6648 let cache_blocks: Vec<(TermCollectionSpec, TermCollectionDesign, Vec<usize>)> = best_specs
6650 .iter()
6651 .zip(boot_designs.iter())
6652 .zip(block_term_indices.iter())
6653 .map(|((spec, design), terms)| (spec.clone(), design.clone(), terms.clone()))
6654 .collect();
6655
6656 struct NBlockExactJointState<'d, M> {
6657 cache: ExactJointDesignCache<'d>,
6658 row_set: gam_problem::outer_subsample::RowSet,
6659 staged_pilot_active: bool,
6660 terminal_mode: Option<(Array1<f64>, f64, M)>,
6661 }
6662
6663 impl<M> NBlockExactJointState<'_, M> {
6664 fn ensure_theta(&mut self, theta: &Array1<f64>) -> Result<(), String> {
6665 let theta_changed = !self
6666 .cache
6667 .current_theta
6668 .as_ref()
6669 .is_some_and(|current| theta_values_match(current, theta));
6670 if theta_changed {
6671 self.terminal_mode = None;
6672 }
6673 self.cache.ensure_theta(theta)
6674 }
6675
6676 fn install_terminal_mode(&mut self, theta: &Array1<f64>, objective: f64, mode: M) {
6677 self.terminal_mode = Some((theta.clone(), objective, mode));
6678 }
6679
6680 fn terminal_mode_matches(&self, theta: &Array1<f64>, objective: f64) -> bool {
6681 self.terminal_mode
6682 .as_ref()
6683 .is_some_and(|(mode_theta, mode_objective, _)| {
6684 theta_values_match(mode_theta, theta)
6685 && mode_objective.to_bits() == objective.to_bits()
6686 })
6687 }
6688 }
6689
6690 let mut state = NBlockExactJointState {
6691 cache: ExactJointDesignCache::new(data, cache_blocks, rho_dim, all_dims.clone())?,
6692 row_set: gam_problem::outer_subsample::RowSet::All,
6693 staged_pilot_active: false,
6694 terminal_mode: None,
6695 };
6696
6697 const KAPPA_PILOT_K: usize = 5_000;
6725
6726 let n_total = data.nrows();
6727 let use_staged_kappa = outer_derivative_policy.should_use_staged_kappa(n_total);
6728 if use_staged_kappa {
6729 log::info!(
6730 "[KAPPA-STAGED] auto-engaging pilot+exact schedule: n={} pilot_k={}",
6731 n_total,
6732 KAPPA_PILOT_K,
6733 );
6734 }
6735
6736 fn build_uniform_pilot_subsample(
6753 n_total: usize,
6754 k_target: usize,
6755 seed: u64,
6756 ) -> gam_problem::outer_subsample::OuterScoreSubsample {
6757 use gam_problem::outer_subsample::OuterScoreSubsample;
6758 let k = k_target.min(n_total);
6759 if k == 0 || n_total == 0 {
6760 return OuterScoreSubsample::from_uniform_inclusion_mask(Vec::new(), n_total, seed);
6761 }
6762 let mut mask: Vec<usize> = Vec::with_capacity(k);
6766 let mut state = seed.wrapping_add(0x9E3779B97F4A7C15);
6768 let splitmix = |s: &mut u64| -> u64 { gam_linalg::utils::splitmix64(s) };
6769 let mut taken = std::collections::HashSet::with_capacity(k);
6770 for j in (n_total - k)..n_total {
6771 let r = (splitmix(&mut state) % (j as u64 + 1)) as usize;
6772 if !taken.insert(r) {
6773 taken.insert(j);
6774 mask.push(j);
6775 } else {
6776 mask.push(r);
6777 }
6778 }
6779 mask.sort_unstable();
6780 mask.dedup();
6781 OuterScoreSubsample::from_uniform_inclusion_mask(mask, n_total, seed)
6782 }
6783
6784 if use_staged_kappa {
6785 let pilot = build_uniform_pilot_subsample(n_total, KAPPA_PILOT_K, n_total as u64);
6786 state.row_set = gam_problem::outer_subsample::RowSet::Subsample {
6787 rows: std::sync::Arc::clone(&pilot.rows),
6788 n_full: n_total,
6789 };
6790 state.staged_pilot_active = true;
6791 }
6792
6793 let exact_fn_cell = std::cell::RefCell::new(&mut exact_fn);
6794 let exact_efs_fn_cell = std::cell::RefCell::new(&mut exact_efs_fn);
6795
6796 use std::cell::Cell;
6811 let kphase_cost_calls: Cell<usize> = Cell::new(0);
6812 let kphase_cost_total_s: Cell<f64> = Cell::new(0.0);
6813 let kphase_eval_calls: Cell<usize> = Cell::new(0);
6814 let kphase_eval_total_s: Cell<f64> = Cell::new(0.0);
6815 let kphase_efs_calls: Cell<usize> = Cell::new(0);
6816 let kphase_efs_total_s: Cell<f64> = Cell::new(0.0);
6817 let kphase_optim_start = std::time::Instant::now();
6818 let kphase_log_kappa_dim = log_kappa_dim;
6819 let kphase_log_norms = |theta: &Array1<f64>| -> (f64, f64) {
6820 let theta_norm = theta.iter().map(|v| v * v).sum::<f64>().sqrt();
6821 let log_kappa_norm = if kphase_log_kappa_dim > 0 && theta.len() >= kphase_log_kappa_dim {
6822 let start = theta.len() - kphase_log_kappa_dim;
6823 theta.iter().skip(start).map(|v| v * v).sum::<f64>().sqrt()
6824 } else {
6825 0.0
6826 };
6827 (theta_norm, log_kappa_norm)
6828 };
6829
6830 use gam_problem::{DeclaredHessianForm, Derivative, OuterEval};
6831 use gam_solve::rho_optimizer::OuterEvalOrder;
6832
6833 let joint_p_cols: usize = boot_designs
6837 .iter()
6838 .map(|d| d.design.ncols())
6839 .sum::<usize>()
6840 .max(1);
6841
6842 let problem = exact_joint_multistart_outer_problem(
6843 &theta0,
6844 &lower,
6845 &upper,
6846 rho_dim,
6847 psi_dim,
6848 theta_dim,
6849 if analytic_joint_gradient_available {
6850 Derivative::Analytic
6851 } else {
6852 Derivative::Unavailable
6853 },
6854 if analytic_outer_hessian_available {
6855 DeclaredHessianForm::Either
6856 } else {
6857 DeclaredHessianForm::Unavailable
6858 },
6859 false,
6863 disable_fixed_point,
6864 seed_risk_profile,
6865 kappa_options.rel_tol.max(1e-6),
6866 kappa_options.max_outer_iter.max(1),
6867 Some(5.0),
6869 Some(kappa_options.log_step.clamp(0.25, 1.0)),
6871 screening_cap.clone(),
6872 Some((n_total, joint_p_cols)),
6875 block_specs
6878 .iter()
6879 .any(|s| !constant_curvature_term_indices(s).is_empty()),
6880 false,
6883 )
6884 .map_err(|e| e.to_string())?;
6885
6886 fn collect_specs(cache: &ExactJointDesignCache<'_>) -> Vec<TermCollectionSpec> {
6888 cache.specs().into_iter().cloned().collect()
6889 }
6890 fn collect_designs(cache: &ExactJointDesignCache<'_>) -> Vec<TermCollectionDesign> {
6891 cache.designs().into_iter().cloned().collect()
6892 }
6893
6894 let result = {
6895 let eval_outer = |ctx: &mut &mut NBlockExactJointState<'_, Mode>,
6896 theta: &Array1<f64>,
6897 order: OuterEvalOrder|
6898 -> Result<OuterEval, EstimationError> {
6899 if let Some((cost, grad, hess)) = ctx.cache.memoized_eval(theta)
6900 && ctx.terminal_mode_matches(theta, cost)
6901 {
6902 let cached_satisfies_order = match order {
6903 OuterEvalOrder::Value => true,
6904 OuterEvalOrder::ValueAndGradient => grad.len() == theta.len(),
6905 OuterEvalOrder::ValueGradientHessian => {
6906 grad.len() == theta.len() && hess.is_analytic()
6907 }
6908 };
6909 if cached_satisfies_order {
6910 if !cost.is_finite() {
6911 return Ok(OuterEval::infeasible(theta.len()));
6912 }
6913 if grad.iter().any(|v| !v.is_finite()) {
6926 return Ok(OuterEval::infeasible(theta.len()));
6927 }
6928 return Ok(OuterEval {
6929 cost,
6930 gradient: grad,
6931 hessian: hess,
6932 inner_beta_hint: None,
6933 });
6934 }
6935 }
6936 ctx.ensure_theta(theta).map_err(|err| {
6937 EstimationError::InvalidInput(format!(
6938 "n-block exact-joint spatial design realization failed: {err}"
6939 ))
6940 })?;
6941 let design_revision = Some(ctx.cache.design_revision());
6942 let specs = collect_specs(&ctx.cache);
6943 let designs = collect_designs(&ctx.cache);
6944 let clamped = outer_derivative_policy.order_for_evaluation(order);
6952 let value_only = matches!(clamped, OuterEvalOrder::Value);
6953 let need_hessian = matches!(clamped, OuterEvalOrder::ValueGradientHessian)
6954 && analytic_outer_hessian_available;
6955 let eval_mode = if value_only {
6956 gam_solve::estimate::reml::reml_outer_engine::EvalMode::ValueOnly
6957 } else if need_hessian {
6958 gam_solve::estimate::reml::reml_outer_engine::EvalMode::ValueGradientHessian
6959 } else {
6960 gam_solve::estimate::reml::reml_outer_engine::EvalMode::ValueAndGradient
6961 };
6962 let t0 = std::time::Instant::now();
6963 let result =
6964 (*exact_fn_cell.borrow_mut())(theta, &specs, &designs, eval_mode, &ctx.row_set);
6965 let elapsed_s = t0.elapsed().as_secs_f64();
6966 kphase_eval_calls.set(kphase_eval_calls.get() + 1);
6967 kphase_eval_total_s.set(kphase_eval_total_s.get() + elapsed_s);
6968 let (theta_norm, log_kappa_norm) = kphase_log_norms(theta);
6969 log::info!(
6970 "[KAPPA-PHASE] phase=eval_outer call={} order={:?} design_revision={:?} theta_norm={:.4e} log_kappa_norm={:.4e} elapsed_s={:.4}",
6971 kphase_eval_calls.get(),
6972 order,
6973 design_revision,
6974 theta_norm,
6975 log_kappa_norm,
6976 elapsed_s,
6977 );
6978 match result {
6979 Ok(ExactJointEvaluation {
6980 objective: cost,
6981 gradient: grad,
6982 hessian: hess,
6983 mode,
6984 }) => {
6985 ctx.install_terminal_mode(theta, cost, mode);
6986 if value_only {
6987 ctx.cache.store_cost_only(theta, cost);
6988 } else {
6989 ctx.cache.store_eval((cost, grad.clone(), hess.clone()));
6990 }
6991 if !cost.is_finite() {
6992 return Ok(OuterEval::infeasible(theta.len()));
6993 }
6994 if grad.iter().any(|v| !v.is_finite()) {
7007 return Ok(OuterEval::infeasible(theta.len()));
7008 }
7009 Ok(OuterEval {
7010 cost,
7011 gradient: grad,
7012 hessian: hess,
7013 inner_beta_hint: None,
7014 })
7015 }
7016 Err(err) => Err(EstimationError::RemlOptimizationFailed(format!(
7017 "n-block exact-joint spatial evaluation failed: {err}"
7018 ))),
7019 }
7020 };
7021
7022 let obj = problem.build_objective_with_eval_order(
7023 &mut state,
7024 |ctx: &mut &mut NBlockExactJointState<'_, Mode>, theta: &Array1<f64>| {
7025 if let Some(cost) = ctx.cache.memoized_cost(theta)
7026 && ctx.terminal_mode_matches(theta, cost)
7027 {
7028 return Ok(cost);
7029 }
7030 ctx.ensure_theta(theta).map_err(|err| {
7031 EstimationError::InvalidInput(format!(
7032 "n-block exact-joint spatial design realization failed: {err}"
7033 ))
7034 })?;
7035 let design_revision = Some(ctx.cache.design_revision());
7036 let specs = collect_specs(&ctx.cache);
7037 let designs = collect_designs(&ctx.cache);
7038 let t0 = std::time::Instant::now();
7045 let result = (*exact_fn_cell.borrow_mut())(
7046 theta,
7047 &specs,
7048 &designs,
7049 gam_solve::estimate::reml::reml_outer_engine::EvalMode::ValueOnly,
7050 &ctx.row_set,
7051 );
7052 let elapsed_s = t0.elapsed().as_secs_f64();
7053 kphase_cost_calls.set(kphase_cost_calls.get() + 1);
7054 kphase_cost_total_s.set(kphase_cost_total_s.get() + elapsed_s);
7055 let (theta_norm, log_kappa_norm) = kphase_log_norms(theta);
7056 log::info!(
7057 "[KAPPA-PHASE] phase=cost call={} design_revision={:?} theta_norm={:.4e} log_kappa_norm={:.4e} elapsed_s={:.4}",
7058 kphase_cost_calls.get(),
7059 design_revision,
7060 theta_norm,
7061 log_kappa_norm,
7062 elapsed_s,
7063 );
7064 match result {
7065 Ok(ExactJointEvaluation {
7066 objective: cost,
7067 mode,
7068 ..
7069 }) => {
7070 ctx.install_terminal_mode(theta, cost, mode);
7071 ctx.cache.store_cost_only(theta, cost);
7077 Ok(cost)
7078 }
7079 Err(err) => Err(EstimationError::RemlOptimizationFailed(format!(
7080 "n-block exact-joint spatial cost evaluation failed: {err}"
7081 ))),
7082 }
7083 },
7084 |ctx: &mut &mut NBlockExactJointState<'_, Mode>, theta: &Array1<f64>| {
7085 eval_outer(ctx, theta, OuterEvalOrder::ValueAndGradient)
7088 },
7089 |ctx: &mut &mut NBlockExactJointState<'_, Mode>,
7090 theta: &Array1<f64>,
7091 order: OuterEvalOrder| { eval_outer(ctx, theta, order) },
7092 None::<fn(&mut &mut NBlockExactJointState<'_, Mode>)>,
7093 Some(
7094 |ctx: &mut &mut NBlockExactJointState<'_, Mode>, theta: &Array1<f64>| {
7095 ctx
7096 .ensure_theta(theta)
7097 .map_err(EstimationError::InvalidInput)?;
7098 let design_revision = Some(ctx.cache.design_revision());
7099 let specs = collect_specs(&ctx.cache);
7100 let designs = collect_designs(&ctx.cache);
7101 let t0 = std::time::Instant::now();
7102 let eval_result = (*exact_efs_fn_cell.borrow_mut())(
7103 theta,
7104 &specs,
7105 &designs,
7106 &ctx.row_set,
7107 );
7108 let elapsed_s = t0.elapsed().as_secs_f64();
7109 kphase_efs_calls.set(kphase_efs_calls.get() + 1);
7110 kphase_efs_total_s.set(kphase_efs_total_s.get() + elapsed_s);
7111 let (theta_norm, log_kappa_norm) = kphase_log_norms(theta);
7112 log::info!(
7113 "[KAPPA-PHASE] phase=efs call={} design_revision={:?} theta_norm={:.4e} log_kappa_norm={:.4e} elapsed_s={:.4}",
7114 kphase_efs_calls.get(),
7115 design_revision,
7116 theta_norm,
7117 log_kappa_norm,
7118 elapsed_s,
7119 );
7120 let ExactJointEfsEvaluation { evaluation, mode } =
7121 eval_result.map_err(EstimationError::RemlOptimizationFailed)?;
7122 ctx.cache.invalidate_objective_memo();
7129 ctx.cache.store_cost_only(theta, evaluation.cost);
7130 ctx.install_terminal_mode(theta, evaluation.cost, mode);
7131 Ok(evaluation)
7132 },
7133 ),
7134 );
7135 let mut obj = obj
7136 .with_seed_inner_state(
7137 move |_ctx: &mut &mut NBlockExactJointState<'_, Mode>, beta: &Array1<f64>| {
7138 (seed_inner_beta_fn)(beta)
7139 },
7140 )
7141 .with_exact_polish(|ctx: &mut &mut NBlockExactJointState<'_, Mode>| {
7142 if !ctx.staged_pilot_active {
7143 return false;
7144 }
7145 ctx.cache.invalidate_objective_memo();
7150 ctx.terminal_mode = None;
7151 ctx.row_set = gam_problem::outer_subsample::RowSet::All;
7152 ctx.staged_pilot_active = false;
7153 true
7154 });
7155
7156 problem
7157 .run_certified(&mut obj, "n-block exact-joint spatial")
7158 .map_err(|error| error.to_string())?
7159 }; let kphase_total_s = kphase_optim_start.elapsed().as_secs_f64();
7169 log::info!(
7170 "[KAPPA-PHASE-SUMMARY] log_kappa_dim={} n_cost={} cost_total_s={:.4} n_eval={} eval_total_s={:.4} n_efs={} efs_total_s={:.4} optim_total_s={:.4}",
7171 kphase_log_kappa_dim,
7172 kphase_cost_calls.get(),
7173 kphase_cost_total_s.get(),
7174 kphase_eval_calls.get(),
7175 kphase_eval_total_s.get(),
7176 kphase_efs_calls.get(),
7177 kphase_efs_total_s.get(),
7178 kphase_total_s,
7179 );
7180 let timing = SpatialLengthScaleOptimizationTiming {
7181 log_kappa_dim: kphase_log_kappa_dim,
7182 cost_calls: kphase_cost_calls.get(),
7183 cost_total_s: kphase_cost_total_s.get(),
7184 eval_calls: kphase_eval_calls.get(),
7185 eval_total_s: kphase_eval_total_s.get(),
7186 efs_calls: kphase_efs_calls.get(),
7187 efs_total_s: kphase_efs_total_s.get(),
7188 slow_path_resets: 0,
7189 design_revision_delta: 0,
7190 nfree_skip_row_touches: 0,
7191 nfree_miss_shape: 0,
7192 nfree_miss_value: 0,
7193 nfree_miss_gradient: 0,
7194 nfree_miss_penalty: 0,
7195 nfree_miss_revision: 0,
7196 nfree_miss_second_order: 0,
7197 nfree_miss_other: 0,
7198 optim_total_s: kphase_total_s,
7199 };
7200
7201 if !matches!(state.row_set, gam_problem::outer_subsample::RowSet::All) {
7202 return Err(
7203 "n-block exact-joint spatial optimization returned before its exact full-data transition"
7204 .to_string(),
7205 );
7206 }
7207 let certified_outer = result;
7208 let theta_star = certified_outer.rho().clone();
7209
7210 state.ensure_theta(&theta_star)?;
7215 let (mode_theta, mode_objective, mode) = state.terminal_mode.take().ok_or_else(|| {
7216 "n-block exact-joint spatial optimization produced a certificate without retaining the owned terminal coefficient mode"
7217 .to_string()
7218 })?;
7219 if !theta_values_match(&mode_theta, &theta_star) {
7220 return Err(
7221 "n-block exact-joint spatial terminal coefficient mode does not bitwise match the certified hyperparameter vector"
7222 .to_string(),
7223 );
7224 }
7225 if mode_objective.to_bits() != certified_outer.final_value().to_bits() {
7226 return Err(format!(
7227 "n-block exact-joint spatial terminal coefficient mode objective does not bitwise match the certified objective: mode={mode_objective:.17e}, certified={:.17e}",
7228 certified_outer.final_value(),
7229 ));
7230 }
7231
7232 let resolved_specs: Vec<TermCollectionSpec> = collect_specs(&state.cache);
7233 let designs: Vec<TermCollectionDesign> = collect_designs(&state.cache);
7234
7235 let fit = fit_fn(
7236 &theta_star,
7237 &resolved_specs,
7238 &designs,
7239 SpatialFitProvenance::Certified {
7240 outer: &certified_outer,
7241 mode,
7242 },
7243 )?;
7244
7245 for spec in &resolved_specs {
7246 log_spatial_aniso_scales(spec);
7247 }
7248
7249 Ok(SpatialLengthScaleOptimizationResult {
7250 resolved_specs,
7251 designs,
7252 fit,
7253 certified_outer: Some(certified_outer),
7254 timing: Some(timing),
7255 })
7256}
7257
7258fn try_exact_joint_latent_coord_optimization(
7259 data: ArrayView2<'_, f64>,
7260 y: ArrayView1<'_, f64>,
7261 weights: ArrayView1<'_, f64>,
7262 offset: ArrayView1<'_, f64>,
7263 resolvedspec: &TermCollectionSpec,
7264 best: &FittedTermCollection,
7265 family: LikelihoodSpec,
7266 options: &FitOptions,
7267 latent: &StandardLatentCoordConfig,
7268) -> Result<FittedTermCollectionWithSpec, EstimationError> {
7269 use gam_problem::{DeclaredHessianForm, Derivative, OuterEval};
7270 use gam_solve::rho_optimizer::OuterEvalOrder;
7271
7272 let rho_dim = best.fit.lambdas.len();
7273 let latent_flat_dim = latent.values.len();
7274 if latent_flat_dim == 0 {
7275 crate::bail_invalid_estim!(
7276 "latent-coordinate optimization requires a non-empty latent block"
7277 );
7278 }
7279 let direct_hypers =
7280 latent_coord_initial_direct_hypers(latent.values.id_mode(), latent.values.latent_dim())?;
7281 let analytic_rho_count = latent
7282 .analytic_penalties
7283 .as_ref()
7284 .map_or(0, |registry| registry.total_rho_count());
7285 let latent_coord_ext_dim = latent_flat_dim + analytic_rho_count + direct_hypers.len();
7286
7287 let mut theta0 = Array1::<f64>::zeros(rho_dim + latent_coord_ext_dim);
7288 theta0
7289 .slice_mut(s![..rho_dim])
7290 .assign(&best.fit.lambdas.mapv(f64::ln));
7291 theta0
7292 .slice_mut(s![rho_dim..rho_dim + latent_flat_dim])
7293 .assign(latent.values.as_flat());
7294 if !direct_hypers.is_empty() {
7295 let direct_start = rho_dim + latent_flat_dim + analytic_rho_count;
7296 theta0
7297 .slice_mut(s![direct_start..direct_start + direct_hypers.len()])
7298 .assign(&direct_hypers);
7299 }
7300
7301 let mut lower = Array1::<f64>::from_elem(theta0.len(), -12.0);
7302 let mut upper = Array1::<f64>::from_elem(theta0.len(), 12.0);
7303 let latent_bound = latent
7304 .values
7305 .as_flat()
7306 .iter()
7307 .fold(1.0_f64, |acc, &v| acc.max(v.abs()))
7308 + 10.0;
7309 for axis in rho_dim..rho_dim + latent_flat_dim {
7310 lower[axis] = -latent_bound;
7311 upper[axis] = latent_bound;
7312 }
7313 if let Some(registry) = latent.analytic_penalties.as_ref() {
7314 let (domain_lower, domain_upper) = registry
7315 .rho_domain_bounds()
7316 .map_err(EstimationError::InvalidInput)?;
7317 let start = rho_dim + latent_flat_dim;
7318 for local in 0..analytic_rho_count {
7319 lower[start + local] = lower[start + local].max(domain_lower[local]);
7320 upper[start + local] = upper[start + local].min(domain_upper[local]);
7321 if lower[start + local] >= upper[start + local] {
7322 return Err(EstimationError::InvalidInput(format!(
7323 "analytic-penalty rho domain has no searchable interval at coordinate {local}: lower={}, upper={}",
7324 lower[start + local],
7325 upper[start + local]
7326 )));
7327 }
7328 }
7329 }
7330
7331 struct LatentJointContext<'d> {
7332 rho_dim: usize,
7333 cache: SingleBlockLatentCoordDesignCache,
7334 evaluator: gam_solve::estimate::ExternalJointHyperEvaluator<'d>,
7335 }
7336
7337 impl<'d> LatentJointContext<'d> {
7338 fn eval_full(
7339 &mut self,
7340 theta: &Array1<f64>,
7341 order: OuterEvalOrder,
7342 ) -> Result<(f64, Array1<f64>, gam_problem::HessianValue), EstimationError> {
7343 if let Some(eval) = self.cache.memoized_eval(theta) {
7344 return Ok(eval);
7345 }
7346 self.cache
7347 .ensure_theta(theta)
7348 .map_err(EstimationError::InvalidInput)?;
7349 let hyper_dirs = self
7350 .cache
7351 .hyper_dirs()
7352 .map_err(EstimationError::InvalidInput)?;
7353 let design_revision = Some(self.cache.design_revision());
7354 let registry_for_key = self.cache.analytic_penalties();
7355 self.evaluator
7356 .set_analytic_penalty_registry(registry_for_key.as_deref());
7357 let mut eval = evaluate_joint_reml_outer_eval_at_theta(
7358 &mut self.evaluator,
7359 self.cache.design(),
7360 theta,
7361 self.rho_dim,
7362 hyper_dirs,
7363 None,
7364 order,
7365 design_revision,
7366 )?;
7367 let latent = self.cache.latent().map_err(EstimationError::InvalidInput)?;
7368 if let Some(registry) = registry_for_key {
7369 add_analytic_penalty_objective_to_eval(
7370 theta,
7371 self.rho_dim,
7372 latent.as_ref(),
7373 registry.as_ref(),
7374 &mut eval,
7375 )?;
7376 }
7377 add_latent_id_objective_to_eval(
7378 theta,
7379 self.rho_dim,
7380 self.cache.analytic_penalty_rho_count(),
7381 latent.as_ref(),
7382 &mut eval,
7383 )?;
7384 self.cache.store_eval(eval.clone());
7385 Ok(eval)
7386 }
7387
7388 fn eval_efs(
7389 &mut self,
7390 theta: &Array1<f64>,
7391 ) -> Result<gam_problem::EfsEval, EstimationError> {
7392 self.cache
7393 .ensure_theta(theta)
7394 .map_err(EstimationError::InvalidInput)?;
7395 let hyper_dirs = self
7396 .cache
7397 .hyper_dirs()
7398 .map_err(EstimationError::InvalidInput)?;
7399 let registry_for_key = self.cache.analytic_penalties();
7400 self.evaluator
7401 .set_analytic_penalty_registry(registry_for_key.as_deref());
7402 let mut efs = evaluate_joint_reml_efs_at_theta(
7403 &mut self.evaluator,
7404 self.cache.design(),
7405 theta,
7406 self.rho_dim,
7407 hyper_dirs,
7408 None,
7409 Some(self.cache.design_revision()),
7410 )?;
7411 if let Some(registry) = registry_for_key {
7412 let latent = self.cache.latent().map_err(EstimationError::InvalidInput)?;
7413 let contribution = analytic_penalty_objective_contribution(
7414 theta,
7415 self.rho_dim,
7416 latent.as_ref(),
7417 registry.as_ref(),
7418 )?;
7419 efs.cost += contribution.cost;
7420 if let (Some(psi_gradient), Some(psi_indices)) =
7421 (efs.psi_gradient.as_mut(), efs.psi_indices.as_ref())
7422 {
7423 if psi_gradient.len() != psi_indices.len() {
7424 crate::bail_invalid_estim!(
7425 "latent-coordinate analytic penalty EFS psi gradient length mismatch: gradient={}, indices={}",
7426 psi_gradient.len(),
7427 psi_indices.len()
7428 );
7429 }
7430 for (local_idx, &theta_idx) in psi_indices.iter().enumerate() {
7431 psi_gradient[local_idx] += contribution.gradient[theta_idx];
7432 }
7433 }
7434 }
7435 Ok(efs)
7436 }
7437
7438 fn eval_cost(&mut self, theta: &Array1<f64>) -> f64 {
7439 if let Some(cost) = self.cache.memoized_cost(theta) {
7440 return cost;
7441 }
7442 if self.cache.ensure_theta(theta).is_err() {
7443 return f64::INFINITY;
7444 }
7445 let design_revision = Some(self.cache.design_revision());
7446 let registry_for_key = self.cache.analytic_penalties();
7447 self.evaluator
7448 .set_analytic_penalty_registry(registry_for_key.as_deref());
7449 let result = {
7450 let design = self.cache.design();
7451 self.evaluator.evaluate_cost_only(
7452 &design.design,
7453 &design.penalties,
7454 &design.nullspace_dims,
7455 design.linear_constraints.clone(),
7456 theta,
7457 self.rho_dim,
7458 None,
7459 "latent-coordinate-joint cost-only",
7460 design_revision,
7461 )
7462 };
7463 match result {
7464 Ok(cost) => {
7465 let latent = match self.cache.latent() {
7466 Ok(latent) => latent,
7467 Err(_) => return f64::INFINITY,
7468 };
7469 let contribution = match latent_id_objective_contribution(
7470 theta,
7471 self.rho_dim,
7472 self.cache.analytic_penalty_rho_count(),
7473 latent.as_ref(),
7474 ) {
7475 Ok(contribution) => contribution,
7476 Err(_) => return f64::INFINITY,
7477 };
7478 let cost = cost + contribution.cost;
7479 let cost = if let Some(registry) = registry_for_key {
7480 match analytic_penalty_objective_contribution(
7481 theta,
7482 self.rho_dim,
7483 latent.as_ref(),
7484 registry.as_ref(),
7485 ) {
7486 Ok(contribution) => cost + contribution.cost,
7487 Err(_) => return f64::INFINITY,
7488 }
7489 } else {
7490 cost
7491 };
7492 self.cache.store_cost(cost);
7493 cost
7494 }
7495 Err(_) => f64::INFINITY,
7496 }
7497 }
7498 }
7499
7500 let effective_offset = best
7501 .design
7502 .compose_offset(offset, "latent-coordinate joint fit")
7503 .map_err(EstimationError::BasisError)?;
7504 let mut ctx = LatentJointContext {
7505 rho_dim,
7506 cache: SingleBlockLatentCoordDesignCache::new(
7507 data.to_owned(),
7508 resolvedspec.clone(),
7509 best.design.clone(),
7510 latent,
7511 rho_dim,
7512 )
7513 .map_err(EstimationError::InvalidInput)?,
7514 evaluator: gam_solve::estimate::ExternalJointHyperEvaluator::new(
7515 y,
7516 weights,
7517 &best.design.design,
7518 effective_offset.view(),
7519 &best.design.penalties,
7520 &external_opts_for_design(&family, &best.design, options),
7521 "latent-coordinate-joint",
7522 )?,
7523 };
7524 let registry_for_key = ctx.cache.analytic_penalties();
7525 ctx.evaluator
7526 .set_analytic_penalty_registry(registry_for_key.as_deref());
7527 ctx.evaluator
7528 .set_persistent_latent_values_fingerprint(latent.values.id_mode());
7529 if let Some(cached_t) = ctx
7530 .evaluator
7531 .load_persistent_latent_values(latent.values.n_obs(), latent.values.latent_dim())
7532 {
7533 let cached_t: Array2<f64> = cached_t;
7534 for (dst, src) in theta0
7535 .slice_mut(s![rho_dim..rho_dim + latent_flat_dim])
7536 .iter_mut()
7537 .zip(cached_t.iter())
7538 {
7539 *dst = *src;
7540 }
7541 }
7542
7543 let problem = exact_joint_multistart_outer_problem(
7544 &theta0,
7545 &lower,
7546 &upper,
7547 rho_dim,
7548 latent_coord_ext_dim,
7549 theta0.len(),
7550 Derivative::Analytic,
7551 DeclaredHessianForm::Unavailable,
7552 true,
7555 false,
7556 seed_risk_profile_for_likelihood_family(&family),
7557 options.tol,
7558 options.max_iter.max(1),
7559 Some(5.0),
7560 Some(0.5),
7561 None,
7562 Some((data.nrows(), best.design.design.ncols().max(1))),
7565 !constant_curvature_term_indices(resolvedspec).is_empty(),
7568 false,
7570 )?;
7571
7572 let eval_outer = |ctx: &mut &mut LatentJointContext<'_>,
7573 theta: &Array1<f64>,
7574 order: OuterEvalOrder|
7575 -> Result<OuterEval, EstimationError> {
7576 let (cost, gradient, hessian) = ctx.eval_full(theta, order)?;
7577 Ok(OuterEval {
7578 cost,
7579 gradient,
7580 hessian,
7581 inner_beta_hint: None,
7582 })
7583 };
7584
7585 let result = {
7586 let mut obj = problem.build_objective_with_eval_order(
7587 &mut ctx,
7588 |ctx: &mut &mut LatentJointContext<'_>, theta: &Array1<f64>| Ok(ctx.eval_cost(theta)),
7589 |ctx: &mut &mut LatentJointContext<'_>, theta: &Array1<f64>| {
7590 eval_outer(ctx, theta, OuterEvalOrder::ValueAndGradient)
7591 },
7592 |ctx: &mut &mut LatentJointContext<'_>, theta: &Array1<f64>, order: OuterEvalOrder| {
7593 eval_outer(ctx, theta, order)
7594 },
7595 Some(|ctx: &mut &mut LatentJointContext<'_>| {
7596 ctx.cache.reset();
7597 }),
7598 Some(|ctx: &mut &mut LatentJointContext<'_>, theta: &Array1<f64>| ctx.eval_efs(theta)),
7599 );
7600
7601 problem
7602 .run(&mut obj, "latent-coordinate joint REML")
7603 .map_err(|e| {
7604 EstimationError::InvalidInput(format!(
7605 "latent-coordinate joint optimization failed after exhausting strategy fallbacks: {e}"
7606 ))
7607 })?
7608 };
7609 if !result.converged {
7610 crate::bail_invalid_estim!(
7611 "latent-coordinate joint optimization did not converge after {} iterations (final_objective={:.6e}, final_grad_norm={})",
7612 result.iterations,
7613 result.final_value,
7614 result.final_grad_norm_report(),
7615 );
7616 }
7617
7618 let theta_star = result.rho;
7619 let selected_lambdas = Array1::from_vec(
7620 gam_problem::checked_exp_log_strengths(
7621 theta_star.slice(s![..rho_dim]).iter().copied(),
7622 )
7623 .map_err(|error| {
7624 EstimationError::InvalidInput(format!(
7625 "selected latent-coordinate smoothing coordinate is outside the canonical log-strength domain: {error}"
7626 ))
7627 })?,
7628 );
7629 let mut final_data = data.to_owned();
7630 let flat_t = theta_star
7631 .slice(s![rho_dim..rho_dim + latent_flat_dim])
7632 .to_owned();
7633 let mut fitted_latent_values =
7634 Array2::<f64>::zeros((latent.values.n_obs(), latent.values.latent_dim()));
7635 for n in 0..latent.values.n_obs() {
7636 for axis in 0..latent.values.latent_dim() {
7637 let value = flat_t[n * latent.values.latent_dim() + axis];
7638 fitted_latent_values[[n, axis]] = value;
7639 final_data[[n, latent.feature_cols[axis]]] = value;
7640 }
7641 }
7642 let optimized = fit_term_collection_forspecwith_heuristic_lambdas(
7643 final_data.view(),
7644 y,
7645 weights,
7646 offset,
7647 resolvedspec,
7648 selected_lambdas.as_slice(),
7649 family,
7650 options,
7651 )?;
7652 ctx.evaluator
7653 .store_persistent_latent_values(&fitted_latent_values);
7654 let mut fit = optimized.fit;
7655 fit.reml_score = result.final_value;
7656 fit.penalized_objective = result.final_value;
7657 Ok(FittedTermCollectionWithSpec {
7658 fit,
7659 design: optimized.design,
7660 resolvedspec: resolvedspec.clone(),
7661 adaptive_diagnostics: optimized.adaptive_diagnostics,
7662 kappa_timing: None,
7663 })
7664}
7665
7666pub fn fit_term_collectionwith_latent_coord_optimization(
7667 data: ArrayView2<'_, f64>,
7668 y: Array1<f64>,
7669 weights: Array1<f64>,
7670 offset: Array1<f64>,
7671 spec: &TermCollectionSpec,
7672 latent: &StandardLatentCoordConfig,
7673 family: LikelihoodSpec,
7674 options: &FitOptions,
7675) -> Result<FittedTermCollectionWithSpec, EstimationError> {
7676 let n = data.nrows();
7677 if !(y.len() == n && weights.len() == n && offset.len() == n) {
7678 crate::bail_invalid_estim!(
7679 "fit_term_collectionwith_latent_coord_optimization row mismatch: n={}, y={}, weights={}, offset={}",
7680 n,
7681 y.len(),
7682 weights.len(),
7683 offset.len()
7684 );
7685 }
7686 let best = fit_term_collection_forspec(
7687 data,
7688 y.view(),
7689 weights.view(),
7690 offset.view(),
7691 spec,
7692 family.clone(),
7693 options,
7694 )?;
7695 let resolvedspec = freeze_term_collection_from_design(spec, &best.design)?;
7696 try_exact_joint_latent_coord_optimization(
7697 data,
7698 y.view(),
7699 weights.view(),
7700 offset.view(),
7701 &resolvedspec,
7702 &best,
7703 family,
7704 options,
7705 latent,
7706 )
7707}
7708
7709fn select_isotropic_matern_range_basin(
7726 data: ArrayView2<'_, f64>,
7727 y: ArrayView1<'_, f64>,
7728 weights: ArrayView1<'_, f64>,
7729 offset: ArrayView1<'_, f64>,
7730 mut resolvedspec: TermCollectionSpec,
7731 mut best: FittedTermCollection,
7732 family: &LikelihoodSpec,
7733 options: &FitOptions,
7734 kappa_options: &SpatialLengthScaleOptimizationOptions,
7735 spatial_terms: &[usize],
7736) -> Result<(TermCollectionSpec, FittedTermCollection), EstimationError> {
7737 if has_aniso_terms(&resolvedspec, spatial_terms)
7741 || !constant_curvature_term_indices(&resolvedspec).is_empty()
7742 {
7743 return Ok((resolvedspec, best));
7744 }
7745
7746 let mut best_score = fit_score(&best.fit);
7747 if !best_score.is_finite() {
7748 crate::bail_invalid_estim!(
7749 "isotropic Matérn basin selection received a non-finite incumbent profile"
7750 );
7751 }
7752
7753 for &term_idx in spatial_terms {
7754 let Some(SmoothBasisSpec::Matern {
7755 feature_cols,
7756 spec: matern,
7757 ..
7758 }) = resolvedspec
7759 .smooth_terms
7760 .get(term_idx)
7761 .map(|term| &term.basis)
7762 else {
7763 continue;
7764 };
7765 let num_centers = gam_terms::basis::center_strategy_num_centers(&matern.center_strategy)
7766 .ok_or_else(|| {
7767 EstimationError::InvalidInput(format!(
7768 "resolved isotropic Matérn term {term_idx} has no finite center count"
7769 ))
7770 })?;
7771 let companion_length_scale = matern_low_rank_center_resolution_length_scale(
7772 data,
7773 feature_cols,
7774 num_centers,
7775 )
7776 .ok_or_else(|| {
7777 EstimationError::InvalidInput(format!(
7778 "resolved isotropic Matérn term {term_idx} has no finite center-resolution range"
7779 ))
7780 })?;
7781 let (psi_long_bound, psi_short_bound) =
7782 spatial_term_psi_bounds(data, &resolvedspec, term_idx, kappa_options)
7783 .map_err(EstimationError::BasisError)?;
7784 let psi_long = (-companion_length_scale.ln()).clamp(psi_long_bound, psi_short_bound);
7785 let long_length_scale = (-psi_long).exp();
7786 if !(long_length_scale.is_finite() && long_length_scale > 0.0) {
7787 crate::bail_invalid_estim!(
7788 "isotropic Matérn term {term_idx} produced an invalid long-range endpoint from psi={psi_long}"
7789 );
7790 }
7791 if get_spatial_length_scale(&resolvedspec, term_idx)
7792 .is_some_and(|current| current == long_length_scale)
7793 {
7794 continue;
7795 }
7796
7797 let mut endpoint_spec = resolvedspec.clone();
7798 set_spatial_length_scale(&mut endpoint_spec, term_idx, long_length_scale)?;
7799 let endpoint = fit_term_collection_forspecwith_heuristic_lambdas(
7809 data,
7810 y,
7811 weights,
7812 offset,
7813 &endpoint_spec,
7814 best.fit.lambdas.as_slice(),
7815 family.clone(),
7816 options,
7817 )?;
7818 let endpoint_score = fit_score(&endpoint.fit);
7819 if !endpoint_score.is_finite() {
7820 crate::bail_invalid_estim!(
7821 "isotropic Matérn term {term_idx} long-range endpoint returned a non-finite profiled REML score"
7822 );
7823 }
7824
7825 if endpoint_score < best_score {
7826 log::info!(
7827 "[spatial-kappa] term {term_idx} selected certified long-range basin: \
7828 length_scale={long_length_scale:.6}, profiled REML {endpoint_score:.6} \
7829 < short-basin {best_score:.6}"
7830 );
7831 resolvedspec = freeze_term_collection_from_design(&endpoint_spec, &endpoint.design)?;
7832 best = endpoint;
7833 best_score = endpoint_score;
7834 } else {
7835 log::info!(
7836 "[spatial-kappa] term {term_idx} retained certified short-range basin: \
7837 profiled REML {best_score:.6} <= long-endpoint {endpoint_score:.6} \
7838 at length_scale={long_length_scale:.6}"
7839 );
7840 }
7841 }
7842
7843 Ok((resolvedspec, best))
7844}
7845
7846pub fn fit_term_collectionwith_spatial_length_scale_optimization(
7847 data: ArrayView2<'_, f64>,
7848 y: Array1<f64>,
7849 weights: Array1<f64>,
7850 offset: Array1<f64>,
7851 spec: &TermCollectionSpec,
7852 family: LikelihoodSpec,
7853 options: &FitOptions,
7854 kappa_options: &SpatialLengthScaleOptimizationOptions,
7855) -> Result<FittedTermCollectionWithSpec, EstimationError> {
7856 let mut resolvedspec = spec.clone();
7872 let spatial_terms = spatial_length_scale_term_indices(&resolvedspec);
7873 let n = data.nrows();
7874 if !(y.len() == n && weights.len() == n && offset.len() == n) {
7875 crate::bail_invalid_estim!(
7876 "fit_term_collectionwith_spatial_length_scale_optimization row mismatch: n={}, y={}, weights={}, offset={}",
7877 n,
7878 y.len(),
7879 weights.len(),
7880 offset.len()
7881 );
7882 }
7883 if !kappa_options.enabled || spatial_terms.is_empty() {
7884 let out = fit_term_collection_forspec(
7885 data,
7886 y.view(),
7887 weights.view(),
7888 offset.view(),
7889 &resolvedspec,
7890 family,
7891 options,
7892 )?;
7893 let resolvedspec = freeze_term_collection_from_design(&resolvedspec, &out.design)?;
7894 return Ok(FittedTermCollectionWithSpec {
7895 fit: out.fit,
7896 design: out.design,
7897 resolvedspec,
7898 adaptive_diagnostics: out.adaptive_diagnostics,
7899 kappa_timing: None,
7900 });
7901 }
7902 if kappa_options.max_outer_iter == 0 {
7903 crate::bail_invalid_estim!("spatial kappa optimization requires max_outer_iter >= 1");
7904 }
7905 if !(kappa_options.log_step.is_finite() && kappa_options.log_step > 0.0) {
7906 crate::bail_invalid_estim!("spatial kappa optimization requires log_step > 0");
7907 }
7908 if !(kappa_options.min_length_scale.is_finite()
7909 && kappa_options.max_length_scale.is_finite()
7910 && kappa_options.min_length_scale > 0.0
7911 && kappa_options.max_length_scale >= kappa_options.min_length_scale)
7912 {
7913 crate::bail_invalid_estim!(
7914 "spatial kappa optimization requires valid positive length_scale bounds"
7915 );
7916 }
7917
7918 let pilot_threshold = kappa_options.pilot_subsample_threshold;
7919 if pilot_threshold > 0 && n > pilot_threshold * 2 {
7920 log::info!(
7921 "[spatial-kappa] n={n} exceeds pilot threshold {}; using pilot geometry only for deterministic anisotropy initialization",
7922 pilot_threshold * 2,
7923 );
7924 apply_spatial_anisotropy_pilot_initializer(
7925 data,
7926 &mut resolvedspec,
7927 &spatial_terms,
7928 pilot_threshold,
7929 kappa_options,
7930 )?;
7931 }
7932
7933 apply_response_aware_anisotropy_seed(data, y.view(), &mut resolvedspec, &spatial_terms);
7942
7943 let free_curvature_terms: Vec<usize> = constant_curvature_term_indices(&resolvedspec)
7947 .into_iter()
7948 .filter(|&term_idx| !constant_curvature_kappa_is_fixed(&resolvedspec, term_idx))
7949 .collect();
7950 if !free_curvature_terms.is_empty() {
7951 validate_constant_curvature_fair_profile_inputs(weights.view(), offset.view(), &family)?;
7952 }
7953 for term_idx in free_curvature_terms {
7954 let kappa_hat = constant_curvature_kappa_fair_optimum(
7955 data,
7956 y.view(),
7957 &resolvedspec,
7958 term_idx,
7959 options,
7960 )?;
7961 if let Some(SmoothBasisSpec::ConstantCurvature { spec: cc, .. }) = resolvedspec
7962 .smooth_terms
7963 .get_mut(term_idx)
7964 .map(|term| &mut term.basis)
7965 {
7966 cc.kappa = kappa_hat;
7967 }
7968 }
7969
7970 let baseline_options = superseded_fit_options(options);
7971 let best = fit_term_collection_forspec(
7972 data,
7973 y.view(),
7974 weights.view(),
7975 offset.view(),
7976 &resolvedspec,
7977 family.clone(),
7978 &baseline_options,
7979 )?;
7980 resolvedspec = freeze_term_collection_from_design(&resolvedspec, &best.design)?;
7981 let spatial_terms = spatial_length_scale_term_indices(&resolvedspec);
7991 let (next_spec, best) = select_isotropic_matern_range_basin(
7992 data,
7993 y.view(),
7994 weights.view(),
7995 offset.view(),
7996 resolvedspec,
7997 best,
7998 &family,
7999 &baseline_options,
8000 kappa_options,
8001 &spatial_terms,
8002 )?;
8003 resolvedspec = next_spec;
8004 sync_aniso_contrasts_from_metadata(&mut resolvedspec, &best.design.smooth);
8008 if spatial_terms.is_empty() {
8009 let fitted = fit_term_collection_forspecwith_heuristic_lambdas(
8010 data,
8011 y.view(),
8012 weights.view(),
8013 offset.view(),
8014 &resolvedspec,
8015 best.fit.lambdas.as_slice(),
8016 family,
8017 options,
8018 )?;
8019 return Ok(FittedTermCollectionWithSpec {
8020 fit: fitted.fit,
8021 design: fitted.design,
8022 resolvedspec,
8023 adaptive_diagnostics: fitted.adaptive_diagnostics,
8024 kappa_timing: None,
8025 });
8026 }
8027 let initial_score = fit_score(&best.fit);
8028 if !initial_score.is_finite() {
8029 crate::bail_invalid_estim!(
8030 "spatial kappa optimization received a non-finite initial profiled score"
8031 );
8032 }
8033 let exact_joint = try_exact_joint_spatial_length_scale_optimization(
8034 data,
8035 y.view(),
8036 weights.view(),
8037 offset.view(),
8038 &resolvedspec,
8039 &best,
8040 family.clone(),
8041 options,
8042 kappa_options,
8043 &spatial_terms,
8044 )?
8045 .ok_or_else(|| {
8046 EstimationError::RemlOptimizationFailed(
8047 "spatial kappa optimization is unavailable for one or more eligible spatial terms"
8048 .to_string(),
8049 )
8050 })?;
8051 let exact_score = fit_score(&exact_joint.fit);
8052 let exact_joint = require_successful_spatial_optimization_result(
8053 initial_score,
8054 Ok(Some((exact_joint, exact_score))),
8055 )?;
8056
8057 log_spatial_aniso_scales(&exact_joint.resolvedspec);
8058 Ok(exact_joint)
8059}
8060
8061#[derive(Clone, Debug)]
8067pub struct CurvatureInference {
8068 pub term_idx: usize,
8070 pub kappa_hat: f64,
8073 pub ci: gam_geometry::curvature_estimand::KappaProfileCi,
8075 pub flatness: gam_geometry::curvature_estimand::FlatnessTest,
8079}
8080
8081fn curvature_profile_lr_endpoint<F>(
8093 profile: &mut F,
8094 kappa_hat: f64,
8095 value_hat: f64,
8096 bound: f64,
8097 half_threshold: f64,
8098 x_tolerance: f64,
8099 score_tolerance: f64,
8100) -> Result<(f64, bool), String>
8101where
8102 F: FnMut(f64) -> Result<(f64, f64), String>,
8103{
8104 let direction = (bound - kappa_hat).signum();
8105 let span = (bound - kappa_hat).abs();
8106 if direction == 0.0 || span <= x_tolerance {
8107 return Ok((bound, true));
8108 }
8109
8110 let (bound_value, bound_score) = profile(bound)?;
8111 let outward_score = direction * bound_score;
8112 if outward_score < -score_tolerance {
8113 return Err(format!(
8114 "curvature profile is not outward-monotone at chart bound {bound}: \
8115 outward score {outward_score:.6e} is below tolerance {score_tolerance:.6e}"
8116 ));
8117 }
8118 let value_tolerance = score_tolerance * span;
8119 if bound_value < value_hat - value_tolerance {
8120 return Err(format!(
8121 "fitted curvature is not the minimum of its inference profile: \
8122 V(bound={bound})={bound_value:.6e} < V(kappa_hat)={value_hat:.6e}"
8123 ));
8124 }
8125 let bound_residual = bound_value - value_hat - half_threshold;
8126 if bound_residual < 0.0 {
8127 return Ok((bound, true));
8128 }
8129 if bound_residual == 0.0 {
8130 return Ok((bound, false));
8131 }
8132
8133 let mut inside_x = kappa_hat;
8138 let mut outside_x = bound;
8139 let mut outside_residual = bound_residual;
8140 let mut outside_score = bound_score;
8141 while (outside_x - inside_x).abs() > x_tolerance {
8142 let lo = inside_x.min(outside_x);
8143 let hi = inside_x.max(outside_x);
8144 let width = hi - lo;
8145 let central_lo = lo + 0.25 * width;
8146 let central_hi = hi - 0.25 * width;
8147 let newton = outside_x - outside_residual / outside_score;
8148 let probe = if newton.is_finite() && newton > central_lo && newton < central_hi {
8149 newton
8150 } else {
8151 lo + 0.5 * width
8152 };
8153 if !(probe > lo && probe < hi) {
8154 break;
8155 }
8156 let (value, score) = profile(probe)?;
8157 let outward_score = direction * score;
8158 if outward_score < -score_tolerance {
8159 return Err(format!(
8160 "curvature profile changed direction before its likelihood crossing at \
8161 kappa={probe}: outward score {outward_score:.6e} is below tolerance \
8162 {score_tolerance:.6e}"
8163 ));
8164 }
8165 let residual = value - value_hat - half_threshold;
8166 if residual >= 0.0 {
8167 outside_x = probe;
8168 outside_residual = residual;
8169 outside_score = score;
8170 } else {
8171 inside_x = probe;
8172 }
8173 }
8174 Ok((inside_x + 0.5 * (outside_x - inside_x), false))
8175}
8176
8177fn curvature_profile_ci_from_analytic_score<F>(
8178 profile: &mut F,
8179 kappa_hat: f64,
8180 kappa_min: f64,
8181 kappa_max: f64,
8182 level: f64,
8183 relative_tolerance: f64,
8184) -> Result<gam_geometry::curvature_estimand::KappaProfileCi, String>
8185where
8186 F: FnMut(f64) -> Result<(f64, f64), String>,
8187{
8188 if !(kappa_min < kappa_max && kappa_hat >= kappa_min && kappa_hat <= kappa_max) {
8189 return Err("curvature profile requires kappa_hat inside valid chart bounds".to_string());
8190 }
8191 if !(level > 0.0 && level < 1.0) {
8192 return Err("curvature profile level must lie in (0, 1)".to_string());
8193 }
8194 let z = gam_geometry::curvature_estimand::wald_half_width(1.0, level)
8195 .ok_or_else(|| "curvature profile threshold is not finite".to_string())?;
8196 let half_threshold = 0.5 * z * z;
8197 let (value_hat, score_hat) = profile(kappa_hat)?;
8198 let relative_tolerance = relative_tolerance.max(f64::EPSILON.sqrt());
8199 let x_tolerance = relative_tolerance * (1.0 + kappa_min.abs().max(kappa_max.abs()));
8200 let score_tolerance = relative_tolerance * (1.0 + value_hat.abs());
8201 let at_lower = (kappa_hat - kappa_min).abs() <= x_tolerance;
8202 let at_upper = (kappa_hat - kappa_max).abs() <= x_tolerance;
8203 let stationary = if at_lower {
8204 score_hat >= -score_tolerance
8205 } else if at_upper {
8206 score_hat <= score_tolerance
8207 } else {
8208 score_hat.abs() <= score_tolerance
8209 };
8210 if !stationary {
8211 return Err(format!(
8212 "curvature inference rejected a non-stationary point estimate: \
8213 kappa_hat={kappa_hat}, score={score_hat:.6e}, \
8214 stationarity_bound={score_tolerance:.6e}"
8215 ));
8216 }
8217
8218 let (ci_lo, lo_at_bound) = curvature_profile_lr_endpoint(
8219 profile,
8220 kappa_hat,
8221 value_hat,
8222 kappa_min,
8223 half_threshold,
8224 x_tolerance,
8225 score_tolerance,
8226 )?;
8227 let (ci_hi, hi_at_bound) = curvature_profile_lr_endpoint(
8228 profile,
8229 kappa_hat,
8230 value_hat,
8231 kappa_max,
8232 half_threshold,
8233 x_tolerance,
8234 score_tolerance,
8235 )?;
8236 let verdict = if ci_lo > 0.0 {
8237 gam_geometry::curvature_estimand::CurvatureVerdict::Spherical
8238 } else if ci_hi < 0.0 {
8239 gam_geometry::curvature_estimand::CurvatureVerdict::Hyperbolic
8240 } else {
8241 gam_geometry::curvature_estimand::CurvatureVerdict::Flat
8242 };
8243 Ok(gam_geometry::curvature_estimand::KappaProfileCi {
8244 kappa_hat,
8245 ci_lo,
8246 ci_hi,
8247 lo_at_bound,
8248 hi_at_bound,
8249 verdict,
8250 })
8251}
8252
8253pub fn curvature_inference_forspec(
8254 data: ArrayView2<'_, f64>,
8255 y: ArrayView1<'_, f64>,
8256 weights: ArrayView1<'_, f64>,
8257 offset: ArrayView1<'_, f64>,
8258 resolvedspec: &TermCollectionSpec,
8259 term_idx: usize,
8260 family: LikelihoodSpec,
8261 options: &FitOptions,
8262 level: f64,
8263) -> Result<CurvatureInference, EstimationError> {
8264 let kappa_hat = get_constant_curvature_kappa(resolvedspec, term_idx).ok_or_else(|| {
8265 EstimationError::InvalidInput(format!(
8266 "curvature_inference_forspec: term {term_idx} is not a constant-curvature smooth"
8267 ))
8268 })?;
8269 if constant_curvature_kappa_is_fixed(resolvedspec, term_idx) {
8270 crate::bail_invalid_estim!(
8271 "curvature inference requires an estimated curvature; term {term_idx} has user-pinned kappa={kappa_hat}"
8272 );
8273 }
8274 if y.len() != data.nrows() || weights.len() != data.nrows() || offset.len() != data.nrows() {
8275 crate::bail_invalid_estim!(
8276 "curvature inference row mismatch: data={}, y={}, weights={}, offset={}",
8277 data.nrows(),
8278 y.len(),
8279 weights.len(),
8280 offset.len(),
8281 );
8282 }
8283 validate_constant_curvature_fair_profile_inputs(weights, offset, &family)?;
8284 let (kappa_min, kappa_max) = constant_curvature_kappa_bounds(data, resolvedspec, term_idx);
8285 let (feature_cols, base_spec) = match resolvedspec
8286 .smooth_terms
8287 .get(term_idx)
8288 .map(|term| &term.basis)
8289 {
8290 Some(SmoothBasisSpec::ConstantCurvature {
8291 feature_cols, spec, ..
8292 }) => (feature_cols, spec.clone()),
8293 _ => {
8294 return Err(EstimationError::InvalidInput(format!(
8295 "constant-curvature κ profile: smooth term {term_idx} is not a \
8296 constant-curvature basis"
8297 )));
8298 }
8299 };
8300 let x_term = select_columns(data, feature_cols).map_err(EstimationError::from)?;
8301 let radial_reference = constant_curvature_radial_reference(x_term.view(), y)?;
8302 let fair_profile = ConstantCurvatureFairProfile {
8303 data: x_term.view(),
8304 response: y,
8305 radial_reference,
8306 spec: base_spec,
8307 cache: std::cell::RefCell::new(std::collections::HashMap::new()),
8308 };
8309
8310 let mut v_p = |kappa: f64| -> Result<(f64, f64), String> {
8313 if !kappa.is_finite() {
8314 return Err(format!("V_p probed a non-finite κ = {kappa}"));
8315 }
8316 let sample = fair_profile.evaluate(kappa).map_err(|error| {
8317 format!("analytic curvature profile at kappa={kappa} failed: {error}")
8318 })?;
8319 Ok(sample)
8320 };
8321 let ci = curvature_profile_ci_from_analytic_score(
8322 &mut v_p,
8323 kappa_hat,
8324 kappa_min,
8325 kappa_max,
8326 level,
8327 options.tol,
8328 )
8329 .map_err(EstimationError::RemlOptimizationFailed)?;
8330 let flatness = gam_geometry::curvature_estimand::flatness_lr_test(
8331 |kappa| v_p(kappa).map(|(value, _)| value),
8332 kappa_hat,
8333 )
8334 .map_err(EstimationError::RemlOptimizationFailed)?;
8335
8336 Ok(CurvatureInference {
8337 term_idx,
8338 kappa_hat,
8339 ci,
8340 flatness,
8341 })
8342}
8343
8344#[cfg(test)]
8345mod curvature_profile_score_tests {
8346 use super::*;
8347
8348 #[test]
8349 fn analytic_profile_score_finds_exact_quadratic_lr_crossings() {
8350 let kappa_hat = -0.37;
8351 let curvature = 16.0;
8352 let level = 0.95;
8353 let mut profile = |kappa: f64| -> Result<(f64, f64), String> {
8354 let displacement = kappa - kappa_hat;
8355 Ok((
8356 7.0 + 0.5 * curvature * displacement * displacement,
8357 curvature * displacement,
8358 ))
8359 };
8360 let ci = curvature_profile_ci_from_analytic_score(
8361 &mut profile,
8362 kappa_hat,
8363 -3.0,
8364 3.0,
8365 level,
8366 1.0e-10,
8367 )
8368 .expect("analytic quadratic profile CI");
8369 let z = gam_geometry::curvature_estimand::wald_half_width(1.0, level)
8370 .expect("valid normal quantile");
8371 let expected_half_width = z / curvature.sqrt();
8372 assert!((ci.ci_lo - (kappa_hat - expected_half_width)).abs() <= 1.0e-8);
8373 assert!((ci.ci_hi - (kappa_hat + expected_half_width)).abs() <= 1.0e-8);
8374 assert!(!ci.lo_at_bound && !ci.hi_at_bound);
8375 }
8376
8377 #[test]
8378 fn analytic_profile_marks_chart_bound_when_wilks_set_never_crosses() {
8379 let mut profile =
8380 |kappa: f64| -> Result<(f64, f64), String> { Ok((0.5 * kappa * kappa, kappa)) };
8381 let ci =
8382 curvature_profile_ci_from_analytic_score(&mut profile, 0.0, -0.1, 0.1, 0.95, 1.0e-10)
8383 .expect("open bounded profile CI");
8384 assert_eq!(ci.ci_lo, -0.1);
8385 assert_eq!(ci.ci_hi, 0.1);
8386 assert!(ci.lo_at_bound && ci.hi_at_bound);
8387 }
8388}
8389
8390#[derive(Clone, Copy, Debug, PartialEq, Eq)]
8393pub enum SmoothLrCorrection {
8394 LawleyLrEstimatedLambda,
8398 LawleyLrFixedLambda,
8403 None,
8407}
8408
8409impl SmoothLrCorrection {
8410 pub fn label(self) -> &'static str {
8412 match self {
8413 SmoothLrCorrection::LawleyLrEstimatedLambda => "lawley_lr_estimated_lambda",
8414 SmoothLrCorrection::LawleyLrFixedLambda => "lawley_lr_fixed_lambda",
8415 SmoothLrCorrection::None => "none",
8416 }
8417 }
8418}
8419
8420#[derive(Clone, Debug)]
8426pub struct SmoothTermLrInference {
8427 pub name: String,
8429 pub term_idx: usize,
8431 pub statistic_lr: f64,
8434 pub ref_df: f64,
8437 pub bartlett_factor: f64,
8440 pub bartlett_factor_conditional: Option<f64>,
8444 pub rho_variation_shift: Option<f64>,
8447 pub statistic_corrected: f64,
8449 pub p_value_uncorrected: f64,
8451 pub p_value_corrected: f64,
8454 pub material: bool,
8462 pub correction: SmoothLrCorrection,
8464}
8465
8466pub const SMOOTH_LR_MATERIAL_THRESHOLD: f64 = 0.10;
8470
8471fn fitted_rho_penalty_components(
8477 penalties: &[BlockwisePenalty],
8478 lambdas: &[f64],
8479 p_total: usize,
8480) -> Result<Vec<gam_terms::inference::lawley::RhoPenaltyComponent>, EstimationError> {
8481 if penalties.len() != lambdas.len() {
8482 return Err(EstimationError::InvalidInput(format!(
8483 "smooth_term_lr_inference: penalty/lambda count mismatch ({} penalties, {} lambdas)",
8484 penalties.len(),
8485 lambdas.len()
8486 )));
8487 }
8488 let mut components = Vec::with_capacity(penalties.len());
8489 for (idx, (penalty, &lambda)) in penalties.iter().zip(lambdas.iter()).enumerate() {
8490 if !(lambda.is_finite() && lambda >= 0.0) {
8491 return Err(EstimationError::InvalidInput(format!(
8492 "smooth_term_lr_inference: lambda[{idx}] is invalid: {lambda}"
8493 )));
8494 }
8495 let r = &penalty.col_range;
8496 if r.end > p_total {
8497 return Err(EstimationError::InvalidInput(format!(
8498 "smooth_term_lr_inference: penalty[{idx}] range {:?} exceeds coefficient dimension {p_total}",
8499 r
8500 )));
8501 }
8502 let mut s_component = Array2::<f64>::zeros((p_total, p_total));
8503 s_component
8504 .slice_mut(s![r.start..r.end, r.start..r.end])
8505 .scaled_add(lambda, &penalty.local);
8506 components.push(gam_terms::inference::lawley::RhoPenaltyComponent { s_component });
8507 }
8508 Ok(components)
8509}
8510
8511pub fn smooth_term_lr_inference_forspec(
8556 data: ArrayView2<'_, f64>,
8557 y: ArrayView1<'_, f64>,
8558 weights: ArrayView1<'_, f64>,
8559 offset: ArrayView1<'_, f64>,
8560 resolvedspec: &TermCollectionSpec,
8561 family: LikelihoodSpec,
8562 options: &FitOptions,
8563) -> Result<Vec<SmoothTermLrInference>, EstimationError> {
8564 use gam_terms::inference::lawley::{
8565 LAWLEY_PAIR_MATRIX_MAX_ROWS, known_scale_expected_jets_with_dispersion,
8566 lawley_lr_bartlett_factor, lawley_lr_mean_shift_with_rho_variation,
8567 };
8568
8569 let n = data.nrows();
8570 let full = fit_term_collection_forspec(
8573 data,
8574 y,
8575 weights,
8576 offset,
8577 resolvedspec,
8578 family.clone(),
8579 options,
8580 )?;
8581 let ll_full = full.fit.log_likelihood;
8582 let p_total = full.design.design.ncols();
8583 let lambdas = full.fit.lambdas.as_slice().ok_or_else(|| {
8584 EstimationError::InvalidInput(
8585 "smooth_term_lr_inference: non-contiguous lambda vector".to_string(),
8586 )
8587 })?;
8588 let s_lambda = weighted_blockwise_penalty_sum(&full.design.penalties, lambdas, p_total);
8589 let rho_penalty_components =
8590 fitted_rho_penalty_components(&full.design.penalties, lambdas, p_total)?;
8591 let rho_covariance = full.fit.artifacts.rho_covariance.as_ref().filter(|cov| {
8592 cov.nrows() == rho_penalty_components.len() && cov.ncols() == rho_penalty_components.len()
8593 });
8594 let full_design_dense = full.design.design.to_dense();
8596 let influence = full.fit.coefficient_influence();
8597 let fitted_likelihood = resolved_likelihood_for_fit(&full.fit)?;
8598 let family_disp = lawley_dispersion_for_family(&fitted_likelihood, &full.fit)?;
8599 let coefficient_covariance_scale = fitted_likelihood
8600 .coefficient_covariance_scale(family_disp)
8601 .map_err(|error| EstimationError::InvalidInput(error.to_string()))?;
8602
8603 let mut out = Vec::<SmoothTermLrInference>::new();
8604 for (term_idx, design_term) in full.design.smooth.terms.iter().enumerate() {
8605 let penalty_range = full
8606 .design
8607 .smooth_term_penalty_range(term_idx)
8608 .map_err(EstimationError::InvalidInput)?;
8609 let (block_start, k) = penalty_range
8610 .map(|range| (range.start, range.len()))
8611 .unwrap_or((0, 0));
8612 if design_term.shape != ShapeConstraint::None {
8615 continue;
8616 }
8617 let coeff_range = design_term.coeff_range.clone();
8618 if coeff_range.start >= coeff_range.end || coeff_range.end > p_total {
8619 continue;
8620 }
8621 let edf = full.fit.per_term_edf(coeff_range.clone(), block_start, k);
8633 let null_dim = design_term.wald_unpenalized_dim();
8653 let rho_uncertainty_df = match wps_block_uncertainty_df(
8674 full.fit.weighted_gram(),
8675 full.fit.smoothing_correction(),
8676 &coeff_range,
8677 coefficient_covariance_scale,
8678 )? {
8679 Some(extra_df) => extra_df,
8682 None => 0.0,
8683 };
8684 let ref_df = (wood_reference_df(influence, &coeff_range)
8685 .unwrap_or(0.0)
8686 .max(edf)
8687 + rho_uncertainty_df)
8688 .max(null_dim as f64)
8689 .max(1.0);
8690 if !(ref_df.is_finite() && ref_df > 0.0) {
8691 continue;
8692 }
8693
8694 let mut null_spec = resolvedspec.clone();
8697 let Some(spec_pos) = null_spec
8698 .smooth_terms
8699 .iter()
8700 .position(|t| t.name == design_term.name)
8701 else {
8702 continue;
8703 };
8704 null_spec.smooth_terms.remove(spec_pos);
8705 let null_fit = fit_term_collection_forspec(
8706 data,
8707 y,
8708 weights,
8709 offset,
8710 &null_spec,
8711 family.clone(),
8712 options,
8713 );
8714 let (statistic_lr, eta_null) = match null_fit {
8715 Ok(null) if null.fit.log_likelihood.is_finite() => {
8716 let w = (2.0 * (ll_full - null.fit.log_likelihood)).max(0.0);
8717 let null_offset = null
8723 .design
8724 .compose_offset(offset, "smooth likelihood-ratio null model")
8725 .map_err(|error| EstimationError::InvalidInput(error.to_string()))?;
8726 let mut eta = null.design.design.dot(&null.fit.beta);
8727 eta += &null_offset;
8728 (w, Some(eta))
8729 }
8730 _ => (f64::NAN, None),
8731 };
8732
8733 let chi2 = statrs::distribution::ChiSquared::new(ref_df).ok();
8734 let p_uncorrected = match (chi2.as_ref(), statistic_lr.is_finite()) {
8735 (Some(dist), true) => {
8736 use statrs::distribution::ContinuousCDF;
8737 (1.0 - dist.cdf(statistic_lr)).clamp(0.0, 1.0)
8738 }
8739 _ => f64::NAN,
8740 };
8741
8742 let mut bartlett_factor = 1.0;
8746 let mut bartlett_factor_conditional = None;
8747 let mut rho_variation_shift = None;
8748 let mut statistic_corrected = statistic_lr;
8749 let mut p_corrected = p_uncorrected;
8750 let mut correction = SmoothLrCorrection::None;
8751 if let (Some(eta), true, true) = (
8752 eta_null.as_ref(),
8753 statistic_lr.is_finite(),
8754 n <= LAWLEY_PAIR_MATRIX_MAX_ROWS,
8755 ) {
8756 let kappas: Option<Vec<_>> = (0..n)
8757 .map(|i| {
8758 known_scale_expected_jets_with_dispersion(
8759 &fitted_likelihood.spec,
8760 eta[i],
8761 family_disp,
8762 )
8763 .and_then(|jets| jets.kappas().ok())
8764 })
8765 .collect();
8766 if let (Some(kappas), Some(dist)) = (kappas, chi2.as_ref()) {
8767 let fixed_factor = lawley_lr_bartlett_factor(
8768 full_design_dense.view(),
8769 &kappas,
8770 Some(s_lambda.view()),
8771 coeff_range.clone(),
8772 ref_df,
8773 );
8774 if let Ok(c_cond) = fixed_factor
8775 && c_cond.is_finite()
8776 && c_cond > 0.0
8777 {
8778 let mut c_applied = c_cond;
8779 correction = SmoothLrCorrection::LawleyLrFixedLambda;
8780 if let Some(cov) = rho_covariance
8781 && let Ok(total_shift) = lawley_lr_mean_shift_with_rho_variation(
8782 full_design_dense.view(),
8783 &kappas,
8784 s_lambda.view(),
8785 coeff_range.clone(),
8786 &rho_penalty_components,
8787 cov.view(),
8788 )
8789 {
8790 let mean_w = ref_df + total_shift;
8791 if let Some(c_est) =
8792 gam_terms::inference::higher_order::bartlett_factor_from_mean(
8793 mean_w, ref_df,
8794 )
8795 && c_est.is_finite()
8796 && c_est > 0.0
8797 {
8798 let conditional_shift = (c_cond - 1.0) * ref_df;
8799 c_applied = c_est;
8800 bartlett_factor_conditional = Some(c_cond);
8801 rho_variation_shift = Some(total_shift - conditional_shift);
8802 correction = SmoothLrCorrection::LawleyLrEstimatedLambda;
8803 }
8804 }
8805 use statrs::distribution::ContinuousCDF;
8806 bartlett_factor = c_applied;
8807 statistic_corrected = statistic_lr / c_applied;
8808 p_corrected = (1.0 - dist.cdf(statistic_corrected)).clamp(0.0, 1.0);
8809 }
8810 }
8811 }
8812
8813 let material = match correction {
8819 SmoothLrCorrection::LawleyLrEstimatedLambda
8820 | SmoothLrCorrection::LawleyLrFixedLambda => {
8821 let factor_move = (bartlett_factor - 1.0).abs();
8822 let p_denom = p_uncorrected.max(p_corrected).max(f64::MIN_POSITIVE);
8823 let p_move = if p_uncorrected.is_finite() && p_corrected.is_finite() {
8824 (p_corrected - p_uncorrected).abs() / p_denom
8825 } else {
8826 0.0
8827 };
8828 factor_move > SMOOTH_LR_MATERIAL_THRESHOLD || p_move > SMOOTH_LR_MATERIAL_THRESHOLD
8829 }
8830 SmoothLrCorrection::None => false,
8831 };
8832
8833 out.push(SmoothTermLrInference {
8834 name: design_term.name.clone(),
8835 term_idx,
8836 statistic_lr,
8837 ref_df,
8838 bartlett_factor,
8839 bartlett_factor_conditional,
8840 rho_variation_shift,
8841 statistic_corrected,
8842 p_value_uncorrected: p_uncorrected,
8843 p_value_corrected: p_corrected,
8844 material,
8845 correction,
8846 });
8847 }
8848 Ok(out)
8849}
8850
8851fn resolved_likelihood_for_fit(
8852 fit: &UnifiedFitResult,
8853) -> Result<gam_spec::GlmLikelihoodSpec, EstimationError> {
8854 let spec = fit.likelihood_family.as_ref().ok_or_else(|| {
8855 EstimationError::InvalidInput(
8856 "smooth-term LR inference requires an engine-level GLM likelihood".to_string(),
8857 )
8858 })?;
8859 gam_spec::GlmLikelihoodSpec::try_new(spec.clone(), fit.likelihood_scale.clone())
8860 .map_err(|error| EstimationError::InvalidInput(error.to_string()))
8861}
8862
8863fn lawley_dispersion_for_family(
8868 likelihood: &gam_spec::GlmLikelihoodSpec,
8869 fit: &UnifiedFitResult,
8870) -> Result<f64, EstimationError> {
8871 let profiled_standard_deviation = matches!(
8872 likelihood
8873 .resolved_scale()
8874 .map_err(|error| EstimationError::InvalidInput(error.to_string()))?,
8875 gam_spec::ResolvedLikelihoodScale::ProfiledGaussian
8876 )
8877 .then_some(fit.standard_deviation);
8878 gam_solve::estimate::dispersion_from_likelihood(likelihood, profiled_standard_deviation)
8879 .map(|dispersion| dispersion.phi())
8880}
8881
8882fn wps_block_uncertainty_df(
8883 weighted_gram: Option<&Array2<f64>>,
8884 smoothing_correction: Option<&Array2<f64>>,
8885 coeff_range: &Range<usize>,
8886 coefficient_covariance_scale: f64,
8887) -> Result<Option<f64>, EstimationError> {
8888 let (Some(xwx), Some(corr)) = (weighted_gram, smoothing_correction) else {
8889 return Ok(None);
8890 };
8891 let (start, end) = (coeff_range.start, coeff_range.end);
8892 if start >= end {
8893 return Err(EstimationError::InvalidInput(format!(
8894 "WPS coefficient block must be non-empty, got {coeff_range:?}"
8895 )));
8896 }
8897 if xwx.nrows() != xwx.ncols() || corr.nrows() != corr.ncols() {
8898 return Err(EstimationError::InvalidInput(format!(
8899 "WPS matrices must be square, got X'WX={}x{} and correction={}x{}",
8900 xwx.nrows(),
8901 xwx.ncols(),
8902 corr.nrows(),
8903 corr.ncols()
8904 )));
8905 }
8906 if xwx.dim() != corr.dim() || end > xwx.nrows() {
8907 return Err(EstimationError::InvalidInput(format!(
8908 "WPS block {coeff_range:?} is incompatible with X'WX={:?} and correction={:?}",
8909 xwx.dim(),
8910 corr.dim()
8911 )));
8912 }
8913 if !(coefficient_covariance_scale.is_finite() && coefficient_covariance_scale > 0.0) {
8914 return Err(EstimationError::InvalidInput(format!(
8915 "WPS coefficient-covariance scale must be finite and strictly positive, got {coefficient_covariance_scale:?}"
8916 )));
8917 }
8918
8919 let mut trace = gam_linalg::utils::KahanSum::default();
8920 for i in start..end {
8921 for j in start..end {
8922 let gram_value = xwx[[i, j]];
8923 let correction_value = corr[[j, i]];
8924 if !gram_value.is_finite() || !correction_value.is_finite() {
8925 return Err(EstimationError::InvalidInput(format!(
8926 "WPS trace has non-finite matrix entry at ({i}, {j}): X'WX={gram_value:?}, correction-transpose={correction_value:?}"
8927 )));
8928 }
8929 let product = gram_value * correction_value;
8930 if !product.is_finite() {
8931 return Err(EstimationError::InvalidInput(format!(
8932 "WPS trace product is not representable at ({i}, {j}): {gram_value:?} * {correction_value:?}"
8933 )));
8934 }
8935 trace.add(product);
8936 }
8937 }
8938 let trace = trace.sum() / coefficient_covariance_scale;
8939 if !trace.is_finite() {
8940 return Err(EstimationError::InvalidInput(format!(
8941 "WPS corrected-EDF trace is not representable after coefficient scale {coefficient_covariance_scale:?}: {trace:?}"
8942 )));
8943 }
8944 if trace < 0.0 {
8945 return Err(EstimationError::InvalidInput(format!(
8946 "WPS corrected-EDF trace must be non-negative, got {trace:?}"
8947 )));
8948 }
8949 Ok(Some(trace))
8950}
8951
8952fn wood_reference_df(influence: Option<&Array2<f64>>, coeff_range: &Range<usize>) -> Option<f64> {
8976 let f = influence?;
8977 let (start, end) = (coeff_range.start, coeff_range.end);
8978 if start >= end || end > f.nrows() || end > f.ncols() {
8979 return None;
8980 }
8981 let block = f.slice(s![start..end, start..end]);
8982 let tr = (0..block.nrows()).map(|i| block[[i, i]]).sum::<f64>();
8983 let tr2 = block.dot(&block).diag().sum();
8984 (tr.is_finite() && tr2.is_finite() && tr > 0.0).then(|| (2.0 * tr - tr2).max(tr).max(1e-12))
8985}
8986
8987#[cfg(test)]
8988mod likelihood_scale_wps_tests {
8989 use super::wps_block_uncertainty_df;
8990 use ndarray::array;
8991
8992 #[test]
8993 fn wps_trace_uses_coefficient_covariance_scale() {
8994 let xwx = array![[1.0, 0.0], [0.0, 1.0]];
8995 let correction = array![[1.0, 0.0], [0.0, 1.0]];
8996 let extra_df = wps_block_uncertainty_df(Some(&xwx), Some(&correction), &(0..2), 4.0)
8997 .expect("valid WPS geometry")
8998 .expect("correction artifacts are present");
8999 assert_eq!(extra_df, 0.5);
9000 }
9001
9002 #[test]
9003 fn wps_absence_is_distinct_from_invalid_geometry() {
9004 let xwx = array![[1.0]];
9005 assert_eq!(
9006 wps_block_uncertainty_df(Some(&xwx), None, &(0..1), 1.0)
9007 .expect("missing optional artifact is not malformed geometry"),
9008 None
9009 );
9010
9011 let negative_correction = array![[-1.0]];
9012 let error = wps_block_uncertainty_df(Some(&xwx), Some(&negative_correction), &(0..1), 1.0)
9013 .expect_err("negative corrected EDF must not be silently zeroed");
9014 assert!(error.to_string().contains("must be non-negative"));
9015 }
9016}
9017
9018#[cfg(test)]
9019mod nfree_gate_tests {
9020 use super::nfree_skip_gate_status_from_parts;
9021
9022 #[test]
9023 fn value_only_nfree_gate_does_not_require_basis_skip_witness() {
9024 let gate = nfree_skip_gate_status_from_parts(
9025 true, true, false, false, true, true, false, false, );
9034 assert!(
9035 gate.would_skip(false),
9036 "value-only κ cost probes must stay n-free when the Gram value is certified; \
9037 the reduced-basis skip witness is required only for beta/gradient probes"
9038 );
9039 }
9040
9041 #[test]
9042 fn gradient_nfree_gate_still_requires_basis_skip_witness() {
9043 let gate =
9044 nfree_skip_gate_status_from_parts(true, true, false, true, true, true, false, true);
9045 assert!(
9046 !gate.would_skip(true),
9047 "gradient probes return beta/gradient objects in a reduced basis and must not \
9048 skip the row lane without the reduced-basis witness"
9049 );
9050 }
9051}