1fn try_build_spatial_term_log_kappa_derivative(
2 data: ArrayView2<'_, f64>,
3 resolvedspec: &TermCollectionSpec,
4 design: &TermCollectionDesign,
5 term_idx: usize,
6) -> Result<
7 Option<(
8 Range<usize>,
9 usize,
10 Array2<f64>,
11 Array2<f64>,
12 Array2<f64>,
13 Array2<f64>,
14 Vec<Array2<f64>>,
15 Vec<Array2<f64>>,
16 Option<std::sync::Arc<gam_terms::basis::ImplicitDesignPsiDerivative>>,
17 )>,
18 EstimationError,
19> {
20 let Some(smooth_term) = design.smooth.terms.get(term_idx) else {
21 return Ok(None);
22 };
23 let Some(termspec) = resolvedspec.smooth_terms.get(term_idx) else {
24 return Ok(None);
25 };
26
27 let derivative_bundle = match &termspec.basis {
28 SmoothBasisSpec::ThinPlate {
29 feature_cols,
30 spec,
31 input_scales,
32 } => {
33 let mut x = select_columns(data, feature_cols).map_err(EstimationError::from)?;
34 let mut spec_local = spec.clone();
35 if let Some(s) = input_scales {
36 apply_input_standardization(&mut x, s);
37 spec_local.length_scale =
38 compensate_length_scale_for_standardization(spec.length_scale, s);
39 }
40 build_thin_plate_basis_log_kappa_derivatives(x.view(), &spec_local)
41 .map_err(EstimationError::from)?
42 }
43 SmoothBasisSpec::Sphere { .. } => return Ok(None),
44 SmoothBasisSpec::ConstantCurvature { feature_cols, spec } => {
53 let x = select_columns(data, feature_cols).map_err(EstimationError::from)?;
54 build_constant_curvature_basis_kappa_derivatives(x.view(), spec)
55 .map_err(EstimationError::from)?
56 }
57 SmoothBasisSpec::MeasureJet { .. } => return Ok(None),
63 SmoothBasisSpec::Matern {
64 feature_cols,
65 spec,
66 input_scales,
67 } => {
68 let mut x = select_columns(data, feature_cols).map_err(EstimationError::from)?;
69 let mut spec_local = spec.clone();
70 if let Some(s) = input_scales {
71 apply_input_standardization(&mut x, s);
72 spec_local.length_scale =
73 compensate_length_scale_for_standardization(spec.length_scale, s);
74 }
75 spec_local.double_penalty = false;
90 build_matern_basis_log_kappa_derivatives(x.view(), &spec_local)
91 .map_err(EstimationError::from)?
92 }
93 SmoothBasisSpec::Duchon {
94 feature_cols,
95 spec,
96 input_scales,
97 } => {
98 let mut x = select_columns(data, feature_cols).map_err(EstimationError::from)?;
99 let mut spec_local = spec.clone();
100 if let Some(s) = input_scales {
101 apply_input_standardization(&mut x, s);
102 spec_local.length_scale =
103 compensate_optional_length_scale_for_standardization(spec.length_scale, s);
104 }
105 let BasisMetadata::Duchon {
106 centers,
107 identifiability_transform,
108 operator_collocation_points,
109 radial_reparam,
110 ..
111 } = &smooth_term.metadata
112 else {
113 return Ok(None);
114 };
115 if spec_local.radial_reparam.is_none() {
118 spec_local.radial_reparam = radial_reparam.clone();
119 }
120 gam_terms::basis::build_duchon_basis_log_kappa_derivativeswith_collocationwithworkspace(
121 x.view(),
122 &spec_local,
123 centers.view(),
124 identifiability_transform.as_ref(),
125 operator_collocation_points
126 .as_ref()
127 .map(|points| points.view()),
128 &mut BasisWorkspace::default(),
129 )
130 .map_err(EstimationError::from)?
131 }
132 SmoothBasisSpec::BSpline1D { .. }
133 | SmoothBasisSpec::TensorBSpline { .. }
134 | SmoothBasisSpec::ByVariable { .. }
135 | SmoothBasisSpec::FactorSumToZero { .. }
136 | SmoothBasisSpec::BySmooth { .. }
137 | SmoothBasisSpec::FactorSmooth { .. }
138 | SmoothBasisSpec::Pca { .. } => {
139 return Ok(None);
140 }
141 };
142 let mut implicit_operator = derivative_bundle.implicit_operator;
143 let BasisPsiDerivativeResult {
144 design_derivative: mut local_x_psi,
145 penalties_derivative: mut local_s_psi,
146 implicit_operator: local_implicit_first_unused,
147 } = derivative_bundle.first;
148 let BasisPsiSecondDerivativeResult {
149 designsecond_derivative: mut local_x_psi_psi,
150 penaltiessecond_derivative: mut local_s_psi_psi,
151 implicit_operator: local_implicit_second_unused,
152 } = derivative_bundle.second;
153 assert!(local_implicit_first_unused.is_none());
154 assert!(local_implicit_second_unused.is_none());
155
156 if let Some(rotation) = smooth_term.joint_null_rotation.as_ref() {
157 let q = &rotation.rotation;
158 if let Some(op) = implicit_operator.take() {
159 implicit_operator = Some(op.append_full_transform(q).map_err(EstimationError::from)?);
160 } else {
161 if local_x_psi.ncols() != q.nrows() || local_x_psi_psi.ncols() != q.nrows() {
162 return Ok(None);
163 }
164 local_x_psi = fast_ab(&local_x_psi, q);
165 local_x_psi_psi = fast_ab(&local_x_psi_psi, q);
166 }
167 let rotate_penalty = |s_local: Array2<f64>| -> Option<Array2<f64>> {
168 if s_local.nrows() != q.nrows() || s_local.ncols() != q.nrows() {
169 return None;
170 }
171 let qt_s = gam_linalg::faer_ndarray::fast_atb(q, &s_local);
172 Some(gam_linalg::faer_ndarray::fast_ab(&qt_s, q))
173 };
174 let Some(rotated_s_psi) = local_s_psi
175 .into_iter()
176 .map(|s| rotate_penalty(s))
177 .collect::<Option<Vec<_>>>()
178 else {
179 return Ok(None);
180 };
181 local_s_psi = rotated_s_psi;
182 let Some(rotated_s_psi_psi) = local_s_psi_psi
183 .into_iter()
184 .map(|s| rotate_penalty(s))
185 .collect::<Option<Vec<_>>>()
186 else {
187 return Ok(None);
188 };
189 local_s_psi_psi = rotated_s_psi_psi;
190 }
191 let implicit_operator = implicit_operator.map(std::sync::Arc::new);
192
193 if let Some(ref op) = implicit_operator {
194 if op.p_out() != smooth_term.coeff_range.len() {
195 return Ok(None);
196 }
197 } else {
198 if local_x_psi.ncols() != smooth_term.coeff_range.len() {
199 return Ok(None);
200 }
201 if local_x_psi_psi.ncols() != smooth_term.coeff_range.len() {
202 return Ok(None);
203 }
204 }
205 if local_s_psi.is_empty() || local_s_psi.len() != local_s_psi_psi.len() {
206 return Ok(None);
207 }
208 if local_s_psi.iter().any(|s| {
209 s.nrows() != smooth_term.coeff_range.len() || s.ncols() != smooth_term.coeff_range.len()
210 }) {
211 return Ok(None);
212 }
213 if local_s_psi_psi.iter().any(|s| {
214 s.nrows() != smooth_term.coeff_range.len() || s.ncols() != smooth_term.coeff_range.len()
215 }) {
216 return Ok(None);
217 }
218
219 let p_total = design.design.ncols();
220 let smooth_start = p_total.saturating_sub(design.smooth.total_smooth_cols());
221 let global_range = (smooth_start + smooth_term.coeff_range.start)
222 ..(smooth_start + smooth_term.coeff_range.end);
223
224 Ok(Some((
225 global_range,
226 p_total,
227 local_x_psi,
228 local_s_psi.iter().fold(
229 Array2::<f64>::zeros((smooth_term.coeff_range.len(), smooth_term.coeff_range.len())),
230 |acc, m| acc + m,
231 ),
232 local_x_psi_psi,
233 local_s_psi_psi.iter().fold(
234 Array2::<f64>::zeros((smooth_term.coeff_range.len(), smooth_term.coeff_range.len())),
235 |acc, m| acc + m,
236 ),
237 local_s_psi,
238 local_s_psi_psi,
239 implicit_operator,
240 )))
241}
242
243fn try_build_spatial_log_kappa_hyper_dirs(
244 data: ArrayView2<'_, f64>,
245 resolvedspec: &TermCollectionSpec,
246 design: &TermCollectionDesign,
247 spatial_terms: &[usize],
248) -> Result<Option<Vec<DirectionalHyperParam>>, EstimationError> {
249 let Some(info_list) =
256 try_build_spatial_log_kappa_derivativeinfo_list(data, resolvedspec, design, spatial_terms)?
257 else {
258 return Ok(None);
259 };
260 Ok(Some(spatial_log_kappa_hyper_dirs_frominfo_list(info_list)?))
261}
262
263pub(crate) fn try_build_latent_coord_hyper_dirs(
264 latent: std::sync::Arc<gam_terms::latent::LatentCoordValues>,
265 resolvedspec: &TermCollectionSpec,
266 design: &TermCollectionDesign,
267 latent_terms: &[gam_problem::types::SmoothTermIdx],
268 analytic_rho_count: usize,
269) -> Result<Option<Vec<DirectionalHyperParam>>, EstimationError> {
270 if latent_terms.is_empty() || latent.is_empty() {
271 return Ok(None);
272 }
273 if latent_terms.len() != 1 {
274 crate::bail_invalid_estim!(
275 "LatentCoord standard-fit hyper_dirs currently require exactly one latent smooth term"
276 .to_string(),
277 );
278 }
279 let term_idx = latent_terms[0];
280 let smooth_term = design.smooth.terms.get(term_idx.get()).ok_or_else(|| {
281 EstimationError::InvalidInput(format!(
282 "LatentCoord term index {term_idx} out of bounds for realized smooth design"
283 ))
284 })?;
285 let termspec = resolvedspec
286 .smooth_terms
287 .get(term_idx.get())
288 .ok_or_else(|| {
289 EstimationError::InvalidInput(format!(
290 "LatentCoord term index {term_idx} out of bounds for resolved smooth spec"
291 ))
292 })?;
293 let p_total = design.design.ncols();
294 let smooth_start = p_total.saturating_sub(design.smooth.total_smooth_cols());
295 let global_range = (smooth_start + smooth_term.coeff_range.start)
296 ..(smooth_start + smooth_term.coeff_range.end);
297
298 let operator = match (&termspec.basis, &smooth_term.metadata) {
303 (
304 SmoothBasisSpec::Matern { .. },
305 BasisMetadata::Matern {
306 centers,
307 length_scale,
308 nu,
309 include_intercept,
310 identifiability_transform,
311 ..
312 },
313 ) => gam_terms::basis::LatentCoordDesignDerivative::new_matern(
314 latent.clone(),
315 std::sync::Arc::new(centers.clone()),
316 *length_scale,
317 *nu,
318 *include_intercept,
319 identifiability_transform.clone(),
320 )
321 .map_err(EstimationError::from)?,
322 (
323 SmoothBasisSpec::Duchon { .. },
324 BasisMetadata::Duchon {
325 centers,
326 length_scale,
327 power,
328 nullspace_order,
329 identifiability_transform,
330 ..
331 },
332 ) => gam_terms::basis::LatentCoordDesignDerivative::new_duchon(
333 latent.clone(),
334 std::sync::Arc::new(centers.clone()),
335 *length_scale,
336 *power,
337 *nullspace_order,
338 identifiability_transform.clone(),
339 )
340 .map_err(EstimationError::from)?,
341 (
342 SmoothBasisSpec::Sphere { .. },
343 BasisMetadata::Sphere {
344 centers,
345 penalty_order,
346 method,
347 constraint_transform,
348 ..
349 },
350 ) if matches!(*method, gam_terms::basis::SphereMethod::Wahba) => {
351 gam_terms::basis::LatentCoordDesignDerivative::new_sphere(
352 latent.clone(),
353 std::sync::Arc::new(centers.clone()),
354 *penalty_order,
355 constraint_transform.clone(),
356 )
357 .map_err(EstimationError::from)?
358 }
359 (
360 SmoothBasisSpec::BSpline1D { spec, .. },
361 BasisMetadata::BSpline1D {
362 knots,
363 identifiability_transform,
364 periodic,
365 degree: meta_degree,
366 ..
367 },
368 ) => {
369 let effective_degree = meta_degree.unwrap_or(spec.degree);
373 if let Some((domain_start, period, num_basis)) = periodic {
374 gam_terms::basis::LatentCoordDesignDerivative::new_periodic_bspline(
375 latent.clone(),
376 (*domain_start, *domain_start + *period),
377 effective_degree,
378 *num_basis,
379 identifiability_transform.clone(),
380 )
381 .map_err(EstimationError::from)?
382 } else {
383 gam_terms::basis::LatentCoordDesignDerivative::new_tensor_bspline(
384 latent.clone(),
385 vec![knots.clone()],
386 vec![effective_degree],
387 identifiability_transform.clone(),
388 )
389 .map_err(EstimationError::from)?
390 }
391 }
392 (
393 SmoothBasisSpec::TensorBSpline { .. },
394 BasisMetadata::TensorBSpline {
395 knots,
396 degrees,
397 identifiability_transform,
398 ..
399 },
400 ) => gam_terms::basis::LatentCoordDesignDerivative::new_tensor_bspline(
401 latent.clone(),
402 knots.clone(),
403 degrees.clone(),
404 identifiability_transform.clone(),
405 )
406 .map_err(EstimationError::from)?,
407 (SmoothBasisSpec::Pca { .. }, BasisMetadata::Pca { basis_matrix, .. }) => {
408 gam_terms::basis::LatentCoordDesignDerivative::new_pca(
409 latent.clone(),
410 std::sync::Arc::new(basis_matrix.clone()),
411 )
412 .map_err(EstimationError::from)?
413 }
414 _ => return Ok(None),
415 };
416 if operator.p_out() != global_range.len() {
417 crate::bail_invalid_estim!(
418 "LatentCoord derivative width mismatch for term '{}': operator p={}, coeff range={}",
419 smooth_term.name,
420 operator.p_out(),
421 global_range.len()
422 );
423 }
424 let operator = std::sync::Arc::new(operator);
425 let mut hyper_dirs = Vec::with_capacity(operator.n_axes());
426 for flat_axis in 0..operator.n_axes() {
427 let dir = DirectionalHyperParam::new_compact(
428 gam_solve::estimate::reml::HyperDesignDerivative::from_latent_coord(
429 operator.clone(),
430 flat_axis,
431 global_range.clone(),
432 p_total,
433 ),
434 Vec::new(),
435 None,
436 None,
437 )?
438 .not_penalty_like();
439 hyper_dirs.push(dir);
440 }
441 let direct_dim = latent_coord_direct_hyper_count(latent.id_mode(), latent.latent_dim());
442 if analytic_rho_count + direct_dim > 0 {
443 let zero_x = gam_solve::estimate::reml::HyperDesignDerivative::from(Array2::<f64>::zeros(
444 (design.design.nrows(), p_total),
445 ));
446 for _ in 0..analytic_rho_count {
447 hyper_dirs.push(
448 DirectionalHyperParam::new_compact(zero_x.clone(), Vec::new(), None, None)?
449 .not_penalty_like(),
450 );
451 }
452 for _ in 0..direct_dim {
453 hyper_dirs.push(
454 DirectionalHyperParam::new_compact(zero_x.clone(), Vec::new(), None, None)?
455 .not_penalty_like(),
456 );
457 }
458 }
459 Ok(Some(hyper_dirs))
460}
461
462fn latent_coord_direct_hyper_count(
463 id_mode: &gam_terms::latent::LatentIdMode,
464 latent_dim: usize,
465) -> usize {
466 use gam_terms::latent::{AuxPriorStrength, LatentIdMode};
467 match id_mode {
468 LatentIdMode::AuxPrior { strength, .. } => match strength {
469 AuxPriorStrength::Auto => 1,
470 AuxPriorStrength::Fixed(_) => 0,
471 },
472 LatentIdMode::AuxPriorDimSelection { strength, .. } => {
473 latent_dim
474 + match strength {
475 AuxPriorStrength::Auto => 1,
476 AuxPriorStrength::Fixed(_) => 0,
477 }
478 }
479 LatentIdMode::DimSelection { .. } => latent_dim,
480 LatentIdMode::IsometryToReference { strength, .. } => match strength {
483 AuxPriorStrength::Auto => 1,
484 AuxPriorStrength::Fixed(_) => 0,
485 },
486 LatentIdMode::AuxOutcome { head, .. } => head.n_coeffs(latent_dim) + latent_dim,
489 LatentIdMode::None => 0,
490 }
491}
492
493fn latent_coord_initial_direct_hypers(
494 id_mode: &gam_terms::latent::LatentIdMode,
495 latent_dim: usize,
496) -> Result<Array1<f64>, EstimationError> {
497 use gam_terms::latent::{AuxPriorStrength, LatentIdMode};
498 let mut values = Vec::with_capacity(latent_coord_direct_hyper_count(id_mode, latent_dim));
499 match id_mode {
500 LatentIdMode::AuxPrior { strength, .. } => {
501 if matches!(strength, AuxPriorStrength::Auto) {
502 values.push(0.0);
503 }
504 }
505 LatentIdMode::AuxPriorDimSelection {
506 strength,
507 init_log_precision,
508 ..
509 } => {
510 if matches!(strength, AuxPriorStrength::Auto) {
511 values.push(0.0);
512 }
513 append_latent_ard_seed(&mut values, init_log_precision.as_ref(), latent_dim)?;
514 }
515 LatentIdMode::DimSelection { init_log_precision } => {
516 append_latent_ard_seed(&mut values, init_log_precision.as_ref(), latent_dim)?;
517 }
518 LatentIdMode::IsometryToReference { strength, .. } => {
519 if matches!(strength, AuxPriorStrength::Auto) {
520 values.push(0.0);
521 }
522 }
523 LatentIdMode::AuxOutcome {
524 head,
525 init_log_precision,
526 } => {
527 values.extend(std::iter::repeat_n(0.0, head.n_coeffs(latent_dim)));
531 append_latent_ard_seed(&mut values, init_log_precision.as_ref(), latent_dim)?;
532 }
533 LatentIdMode::None => {}
534 }
535 Ok(Array1::from_vec(values))
536}
537
538fn append_latent_ard_seed(
539 values: &mut Vec<f64>,
540 init: Option<&Array1<f64>>,
541 latent_dim: usize,
542) -> Result<(), EstimationError> {
543 if let Some(init) = init {
544 if init.len() != latent_dim {
545 crate::bail_invalid_estim!(
546 "latent dim_selection init_log_precision length mismatch: got {}, expected {}",
547 init.len(),
548 latent_dim
549 );
550 }
551 values.extend(init.iter().copied());
552 } else {
553 values.extend(std::iter::repeat_n(0.0, latent_dim));
554 }
555 Ok(())
556}
557
558struct LatentIdObjectiveContribution {
559 cost: f64,
560 gradient: Array1<f64>,
561}
562
563fn latent_id_objective_contribution(
564 theta: &Array1<f64>,
565 rho_dim: usize,
566 analytic_rho_count: usize,
567 latent: &gam_terms::latent::LatentCoordValues,
568) -> Result<LatentIdObjectiveContribution, EstimationError> {
569 use gam_terms::latent::{AuxPriorStrength, LatentIdMode, aux_prior_targets};
570 let n_obs = latent.n_obs();
571 let latent_dim = latent.latent_dim();
572 let flat_len = latent.len();
573 let mut gradient = Array1::<f64>::zeros(theta.len());
574 let t_start = rho_dim;
575 let direct_start = t_start + flat_len + analytic_rho_count;
576 if theta.len() < direct_start {
577 crate::bail_invalid_estim!(
578 "latent-coordinate theta too short for id objective: got {}, need at least {}",
579 theta.len(),
580 direct_start
581 );
582 }
583 let t = latent.as_matrix();
584 let mut cost = 0.0;
585 let mut cursor = direct_start;
586
587 match latent.id_mode() {
588 LatentIdMode::AuxPrior {
589 u,
590 family,
591 strength,
592 }
593 | LatentIdMode::AuxPriorDimSelection {
594 u,
595 family,
596 strength,
597 ..
598 } => {
599 let (log_mu, mu) = match strength {
600 AuxPriorStrength::Fixed(mu) => (mu.ln(), *mu),
601 AuxPriorStrength::Auto => {
602 let log_mu = theta[cursor];
603 cursor += 1;
604 (log_mu, log_mu.exp())
605 }
606 };
607 let targets = aux_prior_targets(t.view(), u.view(), *family)
608 .map_err(EstimationError::InvalidInput)?;
609 let residual = &t - &targets;
610 let q = residual.iter().map(|v| v * v).sum::<f64>();
611 let k = (n_obs * latent_dim) as f64;
618 cost += 0.5 * mu * q - 0.5 * k * log_mu;
619
620 let projected_residual = aux_prior_targets(residual.view(), u.view(), *family)
621 .map_err(EstimationError::InvalidInput)?;
622 let grad_base = residual - projected_residual;
623 for n in 0..n_obs {
624 for axis in 0..latent_dim {
625 gradient[t_start + n * latent_dim + axis] += mu * grad_base[[n, axis]];
626 }
627 }
628 if matches!(strength, AuxPriorStrength::Auto) {
629 gradient[direct_start] += 0.5 * mu * q - 0.5 * k;
630 }
631 }
632 LatentIdMode::IsometryToReference {
633 reference,
634 strength,
635 } => {
636 if reference.dim() != (n_obs, latent_dim) {
643 crate::bail_invalid_estim!(
644 "IsometryToReference reference shape {:?} must equal (n_obs, latent_dim) = ({}, {})",
645 reference.dim(),
646 n_obs,
647 latent_dim
648 );
649 }
650 let mu_slot = cursor;
651 let (log_mu, mu) = match strength {
652 AuxPriorStrength::Fixed(mu) => (mu.ln(), *mu),
653 AuxPriorStrength::Auto => {
654 let log_mu = theta[cursor];
655 cursor += 1;
656 (log_mu, log_mu.exp())
657 }
658 };
659 let residual = &t - reference;
660 let q = residual.iter().map(|v| v * v).sum::<f64>();
661 let k = (n_obs * latent_dim) as f64;
665 cost += 0.5 * mu * q - 0.5 * k * log_mu;
666 for n in 0..n_obs {
667 for axis in 0..latent_dim {
668 gradient[t_start + n * latent_dim + axis] += mu * residual[[n, axis]];
669 }
670 }
671 if matches!(strength, AuxPriorStrength::Auto) {
672 gradient[mu_slot] += 0.5 * mu * q - 0.5 * k;
673 }
674 }
675 LatentIdMode::AuxOutcome { head, .. } => {
676 let n_coeffs = head.n_coeffs(latent_dim);
684 let coeffs = theta
685 .slice(ndarray::s![cursor..cursor + n_coeffs])
686 .to_owned();
687 let (head_nll, grad_coeffs, grad_t) = head
688 .neg_loglik_and_grad(t.view(), coeffs.view())
689 .map_err(EstimationError::InvalidInput)?;
690 cost += head_nll;
691 for (offset, &g) in grad_coeffs.iter().enumerate() {
692 gradient[cursor + offset] += g;
693 }
694 for n in 0..n_obs {
695 for axis in 0..latent_dim {
696 gradient[t_start + n * latent_dim + axis] += grad_t[[n, axis]];
697 }
698 }
699 cursor += n_coeffs;
700 }
701 LatentIdMode::DimSelection { .. } | LatentIdMode::None => {}
702 }
703
704 match latent.id_mode() {
705 LatentIdMode::AuxPriorDimSelection { .. }
706 | LatentIdMode::DimSelection { .. }
707 | LatentIdMode::AuxOutcome { .. } => {
708 for axis in 0..latent_dim {
709 let log_alpha = theta[cursor + axis];
710 let alpha = log_alpha.exp();
711 let mut q_axis = 0.0;
712 for n in 0..n_obs {
713 let flat_idx = n * latent_dim + axis;
714 let value = latent.as_flat()[flat_idx];
715 q_axis += value * value;
716 gradient[t_start + flat_idx] += alpha * value;
717 }
718 cost += 0.5 * alpha * q_axis - 0.5 * n_obs as f64 * log_alpha;
719 gradient[cursor + axis] += 0.5 * alpha * q_axis - 0.5 * n_obs as f64;
720 }
721 cursor += latent_dim;
722 }
723 LatentIdMode::AuxPrior { .. }
724 | LatentIdMode::IsometryToReference { .. }
725 | LatentIdMode::None => {}
726 }
727
728 if cursor != theta.len() {
729 crate::bail_invalid_estim!(
730 "latent-coordinate direct hyperparameter length mismatch: consumed {}, theta len {}",
731 cursor,
732 theta.len()
733 );
734 }
735 Ok(LatentIdObjectiveContribution { cost, gradient })
736}
737
738fn add_latent_id_objective_to_eval(
739 theta: &Array1<f64>,
740 rho_dim: usize,
741 analytic_rho_count: usize,
742 latent: &gam_terms::latent::LatentCoordValues,
743 eval: &mut (f64, Array1<f64>, gam_problem::HessianValue),
744) -> Result<(), EstimationError> {
745 let contribution =
746 latent_id_objective_contribution(theta, rho_dim, analytic_rho_count, latent)?;
747 eval.0 += contribution.cost;
748 if eval.1.len() != contribution.gradient.len() {
749 crate::bail_invalid_estim!(
750 "latent-coordinate REML gradient length mismatch: base={}, id={}",
751 eval.1.len(),
752 contribution.gradient.len()
753 );
754 }
755 eval.1 += &contribution.gradient;
756 if eval.2.is_analytic() {
757 eval.2 = gam_problem::HessianValue::Unavailable;
758 }
759 Ok(())
760}
761
762fn analytic_penalty_objective_contribution(
763 theta: &Array1<f64>,
764 rho_dim: usize,
765 latent: &gam_terms::latent::LatentCoordValues,
766 registry: &gam_terms::AnalyticPenaltyRegistry,
767) -> Result<LatentIdObjectiveContribution, EstimationError> {
768 let flat_len = latent.len();
769 let t_start = rho_dim;
770 let t_end = t_start + flat_len;
771 let rho_start = t_end;
772 let rho_end = rho_start + registry.total_rho_count();
773 if theta.len() < rho_end {
774 crate::bail_invalid_estim!(
775 "latent-coordinate theta too short for analytic penalties: got {}, need at least {}",
776 theta.len(),
777 rho_end
778 );
779 }
780 let target_t = theta.slice(s![t_start..t_end]);
781 let rho = theta.slice(s![rho_start..rho_end]);
782 let mut cost = 0.0_f64;
783 let mut gradient = Array1::<f64>::zeros(theta.len());
784 for (penalty, (rho_slice, tier, name)) in registry.penalties.iter().zip(registry.rho_layout()) {
785 let rho_local = rho.slice(s![rho_slice.clone()]);
786 match tier {
787 gam_terms::PenaltyTier::Psi => {
788 cost += penalty.value(target_t.view(), rho_local);
789 let grad = penalty.grad_target(target_t.view(), rho_local);
790 if grad.len() != flat_len {
791 crate::bail_invalid_estim!(
792 "analytic penalty {name:?} gradient length mismatch: got {}, expected {}",
793 grad.len(),
794 flat_len
795 );
796 }
797 for i in 0..flat_len {
798 gradient[t_start + i] += grad[i];
799 }
800 let grad_rho_local = penalty.grad_rho(target_t.view(), rho_local);
801 if grad_rho_local.len() != rho_slice.len() {
802 crate::bail_invalid_estim!(
803 "analytic penalty {name:?} rho-gradient length mismatch: got {}, expected {}",
804 grad_rho_local.len(),
805 rho_slice.len()
806 );
807 }
808 for local_idx in 0..grad_rho_local.len() {
809 gradient[rho_start + rho_slice.start + local_idx] += grad_rho_local[local_idx];
810 }
811 }
812 gam_terms::PenaltyTier::Beta => {}
813 gam_terms::PenaltyTier::Rho => {}
814 }
815 }
816 Ok(LatentIdObjectiveContribution { cost, gradient })
817}
818
819fn add_analytic_penalty_hessian_to_eval(
820 theta: &Array1<f64>,
821 rho_dim: usize,
822 latent: &gam_terms::latent::LatentCoordValues,
823 registry: &gam_terms::AnalyticPenaltyRegistry,
824 eval: &mut (f64, Array1<f64>, gam_problem::HessianValue),
825) -> Result<(), EstimationError> {
826 let flat_len = latent.len();
827 let t_start = rho_dim;
828 let t_end = t_start + flat_len;
829 let rho_start = t_end;
830 let rho_end = rho_start + registry.total_rho_count();
831 if theta.len() < rho_end {
832 crate::bail_invalid_estim!(
833 "latent-coordinate theta too short for analytic penalty Hessian: got {}, need at least {}",
834 theta.len(),
835 rho_end
836 );
837 }
838 let gam_problem::HessianValue::Dense(hessian) = &mut eval.2 else {
839 if eval.2.is_analytic() {
840 eval.2 = gam_problem::HessianValue::Unavailable;
841 }
842 return Ok(());
843 };
844 if hessian.dim() != (theta.len(), theta.len()) {
845 crate::bail_invalid_estim!(
846 "analytic penalty Hessian target shape mismatch: got {}x{}, expected {}x{}",
847 hessian.nrows(),
848 hessian.ncols(),
849 theta.len(),
850 theta.len()
851 );
852 }
853 let target_t = theta.slice(s![t_start..t_end]);
854 let rho = theta.slice(s![rho_start..rho_end]);
855 for (penalty, (rho_slice, tier, _name)) in registry.penalties.iter().zip(registry.rho_layout())
856 {
857 let rho_local = rho.slice(s![rho_slice]);
858 if !matches!(tier, gam_terms::PenaltyTier::Psi) {
859 continue;
860 }
861 if let Some(diag) = penalty.hessian_diag(target_t.view(), rho_local) {
862 if diag.len() != flat_len {
863 crate::bail_invalid_estim!(
864 "analytic penalty Hessian diagonal length mismatch: got {}, expected {}",
865 diag.len(),
866 flat_len
867 );
868 }
869 for i in 0..flat_len {
870 hessian[[t_start + i, t_start + i]] += diag[i];
871 }
872 continue;
873 }
874 let mut probe = Array1::<f64>::zeros(flat_len);
875 for col in 0..flat_len {
876 probe[col] = 1.0;
877 let hv = penalty.hvp(target_t.view(), rho_local, probe.view());
878 if hv.len() != flat_len {
879 crate::bail_invalid_estim!(
880 "analytic penalty Hessian-vector length mismatch: got {}, expected {}",
881 hv.len(),
882 flat_len
883 );
884 }
885 for row in 0..flat_len {
886 hessian[[t_start + row, t_start + col]] += hv[row];
887 }
888 probe[col] = 0.0;
889 }
890 }
891 Ok(())
892}
893
894fn add_analytic_penalty_objective_to_eval(
895 theta: &Array1<f64>,
896 rho_dim: usize,
897 latent: &gam_terms::latent::LatentCoordValues,
898 registry: &gam_terms::AnalyticPenaltyRegistry,
899 eval: &mut (f64, Array1<f64>, gam_problem::HessianValue),
900) -> Result<(), EstimationError> {
901 let contribution = analytic_penalty_objective_contribution(theta, rho_dim, latent, registry)?;
902 eval.0 += contribution.cost;
903 if eval.1.len() != contribution.gradient.len() {
904 crate::bail_invalid_estim!(
905 "latent-coordinate REML gradient length mismatch: base={}, analytic_penalty={}",
906 eval.1.len(),
907 contribution.gradient.len()
908 );
909 }
910 eval.1 += &contribution.gradient;
911 add_analytic_penalty_hessian_to_eval(theta, rho_dim, latent, registry, eval)?;
912 Ok(())
913}
914
915fn spatial_log_kappa_hyper_dirs_frominfo_list(
916 info_list: Vec<SpatialPsiDerivative>,
917) -> Result<Vec<DirectionalHyperParam>, EstimationError> {
918 use gam_solve::estimate::reml::ImplicitDerivLevel;
919 use std::collections::HashMap;
920
921 let log_kappa_dim = info_list.len();
922 let group_ids: Vec<Option<usize>> = info_list.iter().map(|e| e.aniso_group_id).collect();
928 let mut group_indices_map: HashMap<usize, Vec<usize>> = HashMap::new();
929 for (idx, gid) in group_ids.iter().enumerate() {
930 if let Some(g) = gid {
931 group_indices_map.entry(*g).or_default().push(idx);
932 }
933 }
934
935 let mut hyper_dirs = Vec::with_capacity(log_kappa_dim);
936 for (i, info) in info_list.into_iter().enumerate() {
937 let SpatialPsiDerivative {
938 penalty_index: _,
939 penalty_indices,
940 global_range,
941 total_p,
942 x_psi_local,
943 s_psi_components_local,
944 x_psi_psi_local,
945 s_psi_psi_components_local,
946 aniso_group_id,
947 aniso_cross_designs,
948 aniso_cross_penalty_provider,
949 implicit_operator,
950 implicit_axis,
951 } = info;
952
953 let mut xsecond = vec![None; log_kappa_dim];
954 xsecond[i] = Some(if let Some(ref op) = implicit_operator {
956 gam_solve::estimate::reml::HyperDesignDerivative::from_implicit(
957 op.clone(),
958 ImplicitDerivLevel::SecondDiag(implicit_axis),
959 global_range.clone(),
960 total_p,
961 )
962 } else {
963 gam_solve::estimate::reml::HyperDesignDerivative::from_embedded(
964 x_psi_psi_local,
965 global_range.clone(),
966 total_p,
967 )
968 });
969 if let Some(cross_designs) = aniso_cross_designs {
971 if let Some(gid) = aniso_group_id {
975 let base = group_indices_map
976 .get(&gid)
977 .and_then(|v| v.first().copied())
978 .unwrap_or(i);
979 for (b_axis, cross_mat) in cross_designs.into_iter() {
980 let j = base + b_axis;
981 if j < log_kappa_dim {
982 xsecond[j] = Some(if let Some(ref op) = implicit_operator {
983 gam_solve::estimate::reml::HyperDesignDerivative::from_implicit(
984 op.clone(),
985 ImplicitDerivLevel::SecondCross(implicit_axis, b_axis),
986 global_range.clone(),
987 total_p,
988 )
989 } else {
990 gam_solve::estimate::reml::HyperDesignDerivative::from_embedded(
991 cross_mat,
992 global_range.clone(),
993 total_p,
994 )
995 });
996 }
997 }
998 }
999 }
1000 let s_components = penalty_indices
1001 .iter()
1002 .copied()
1003 .zip(s_psi_components_local.into_iter().map(|local| {
1004 gam_solve::estimate::reml::HyperPenaltyDerivative::from_embedded(
1005 local,
1006 global_range.clone(),
1007 total_p,
1008 )
1009 }))
1010 .collect::<Vec<_>>();
1011 let s2_components = penalty_indices
1012 .iter()
1013 .copied()
1014 .zip(s_psi_psi_components_local.into_iter().map(|local| {
1015 gam_solve::estimate::reml::HyperPenaltyDerivative::from_embedded(
1016 local,
1017 global_range.clone(),
1018 total_p,
1019 )
1020 }))
1021 .collect::<Vec<_>>();
1022 let mut ssecond_components = vec![None; log_kappa_dim];
1023 ssecond_components[i] = Some(s2_components);
1024 let mut penaltysecond_partner_indices: Option<Vec<usize>> = None;
1025 let penaltysecond_component_provider = if let (Some(provider), Some(gid)) =
1026 (aniso_cross_penalty_provider, aniso_group_id)
1027 {
1028 let group_indices = group_indices_map.get(&gid).cloned().unwrap_or_default();
1029 let axis_in_group =
1030 group_indices
1031 .iter()
1032 .position(|&idx| idx == i)
1033 .ok_or_else(|| {
1034 EstimationError::InvalidInput(format!(
1035 "missing spatial hyper axis {} in anisotropy group {}",
1036 i, gid
1037 ))
1038 })?;
1039 penaltysecond_partner_indices = Some(
1040 group_indices
1041 .iter()
1042 .copied()
1043 .filter(|&idx| idx != i)
1044 .collect(),
1045 );
1046 let penalty_indices_inner = penalty_indices.clone();
1047 let global_range_inner = global_range.clone();
1048 let total_p_inner = total_p;
1049 let group_indices_inner = group_indices;
1050 Some(std::sync::Arc::new(
1051 move |j: usize| -> Result<
1052 Option<Vec<gam_solve::estimate::reml::PenaltyDerivativeComponent>>,
1053 EstimationError,
1054 > {
1055 let Some(other_axis_in_group) =
1056 group_indices_inner.iter().position(|&idx| idx == j)
1057 else {
1058 return Ok(None);
1059 };
1060 if other_axis_in_group == axis_in_group {
1061 return Ok(None);
1062 }
1063 let cross_pens = provider(other_axis_in_group)?;
1064 if cross_pens.is_empty() {
1065 return Ok(None);
1066 }
1067 Ok(Some(
1068 penalty_indices_inner
1069 .iter()
1070 .copied()
1071 .zip(cross_pens.into_iter().map(|local| {
1072 gam_solve::estimate::reml::HyperPenaltyDerivative::from_embedded(
1073 local,
1074 global_range_inner.clone(),
1075 total_p_inner,
1076 )
1077 }))
1078 .map(|(penalty_index, matrix)| {
1079 gam_solve::estimate::reml::PenaltyDerivativeComponent {
1080 penalty_index,
1081 matrix,
1082 }
1083 })
1084 .collect(),
1085 ))
1086 },
1087 )
1088 as std::sync::Arc<
1089 dyn Fn(
1090 usize,
1091 ) -> Result<
1092 Option<Vec<gam_solve::estimate::reml::PenaltyDerivativeComponent>>,
1093 EstimationError,
1094 > + Send
1095 + Sync
1096 + 'static,
1097 >)
1098 } else {
1099 None
1100 };
1101 let x_first_hyper = if let Some(ref op) = implicit_operator {
1104 gam_solve::estimate::reml::HyperDesignDerivative::from_implicit(
1105 op.clone(),
1106 ImplicitDerivLevel::First(implicit_axis),
1107 global_range.clone(),
1108 total_p,
1109 )
1110 } else {
1111 gam_solve::estimate::reml::HyperDesignDerivative::from_embedded(
1112 x_psi_local,
1113 global_range.clone(),
1114 total_p,
1115 )
1116 };
1117 let mut dir = DirectionalHyperParam::new_compact(
1118 x_first_hyper,
1119 s_components,
1120 Some(xsecond),
1121 Some(ssecond_components),
1122 )?
1123 .not_penalty_like();
1124 if let Some(provider) = penaltysecond_component_provider {
1125 dir = dir.with_penaltysecond_component_provider(provider);
1126 }
1127 if let Some(partner_indices) = penaltysecond_partner_indices {
1128 dir = dir.with_penaltysecond_partner_indices(partner_indices);
1129 }
1130 hyper_dirs.push(dir);
1131 }
1132 Ok(hyper_dirs)
1133}
1134
1135pub(crate) fn spatial_dims_per_term(
1141 resolvedspec: &TermCollectionSpec,
1142 spatial_terms: &[usize],
1143) -> Vec<usize> {
1144 spatial_terms
1145 .iter()
1146 .map(|&term_idx| {
1147 if let Some(mj) = measure_jet_term_spec(resolvedspec, term_idx) {
1148 measure_jet_psi_dim(mj)
1151 } else if spatial_term_uses_per_axis_psi(resolvedspec, term_idx) {
1152 get_spatial_feature_dim(resolvedspec, term_idx).unwrap_or(1)
1153 } else {
1154 1
1155 }
1156 })
1157 .collect()
1158}
1159
1160fn has_aniso_terms(resolvedspec: &TermCollectionSpec, spatial_terms: &[usize]) -> bool {
1164 spatial_terms
1165 .iter()
1166 .any(|&term_idx| spatial_term_uses_per_axis_psi(resolvedspec, term_idx))
1167}
1168
1169macro_rules! impl_exact_joint_theta_memo {
1175 () => {
1176 fn memoized_cost(&self, theta: &Array1<f64>) -> Option<f64> {
1177 if self
1178 .current_theta
1179 .as_ref()
1180 .is_some_and(|cached| theta_values_match(cached, theta))
1181 {
1182 self.last_eval
1183 .as_ref()
1184 .map(|cached| cached.0)
1185 .or(self.last_cost)
1186 } else {
1187 None
1188 }
1189 }
1190
1191 fn memoized_eval(
1192 &self,
1193 theta: &Array1<f64>,
1194 ) -> Option<(f64, Array1<f64>, gam_problem::HessianValue)> {
1195 if self
1196 .current_theta
1197 .as_ref()
1198 .is_some_and(|cached| theta_values_match(cached, theta))
1199 {
1200 self.last_eval.clone()
1201 } else {
1202 None
1203 }
1204 }
1205
1206 fn store_eval(&mut self, eval: (f64, Array1<f64>, gam_problem::HessianValue)) {
1207 self.last_cost = Some(eval.0);
1208 self.last_eval = Some(eval);
1209 }
1210 };
1211}
1212
1213struct SingleBlockExactJointDesignCache<'d> {
1214 realizer: FrozenTermCollectionIncrementalRealizer<'d>,
1215 current_theta: Option<Array1<f64>>,
1216 last_eval_theta: Option<Array1<f64>>,
1223 last_cost: Option<f64>,
1224 last_eval: Option<(f64, Array1<f64>, gam_problem::HessianValue)>,
1225 cached_hyper_dirs: Option<(u64, Vec<DirectionalHyperParam>)>,
1237 spatial_terms: Vec<usize>,
1238 rho_dim: usize,
1239 dims_per_term: Vec<usize>,
1240}
1241
1242impl<'d> SingleBlockExactJointDesignCache<'d> {
1243 fn new_with_policy(
1244 data: ArrayView2<'d, f64>,
1245 spec: TermCollectionSpec,
1246 design: TermCollectionDesign,
1247 spatial_terms: Vec<usize>,
1248 rho_dim: usize,
1249 dims_per_term: Vec<usize>,
1250 policy: &gam_runtime::resource::ResourcePolicy,
1251 ) -> Result<Self, String> {
1252 Ok(Self {
1253 realizer: FrozenTermCollectionIncrementalRealizer::new_with_policy(
1254 data, spec, design, policy,
1255 )?,
1256 current_theta: None,
1257 last_eval_theta: None,
1258 last_cost: None,
1259 last_eval: None,
1260 cached_hyper_dirs: None,
1261 spatial_terms,
1262 rho_dim,
1263 dims_per_term,
1264 })
1265 }
1266
1267 fn design_revision(&self) -> u64 {
1268 self.realizer.design_revision()
1269 }
1270
1271 fn hyper_dirs_for_current_design(
1281 &mut self,
1282 data: ArrayView2<'_, f64>,
1283 kind: SpatialHyperKind,
1284 ) -> Result<Vec<DirectionalHyperParam>, EstimationError> {
1285 let revision = self.realizer.design_revision();
1286 if let Some((cached_rev, dirs)) = self.cached_hyper_dirs.as_ref()
1287 && *cached_rev == revision
1288 {
1289 return Ok(dirs.clone());
1290 }
1291 let dirs = try_build_spatial_log_kappa_hyper_dirs(
1292 data,
1293 self.realizer.spec(),
1294 self.realizer.design(),
1295 &self.spatial_terms,
1296 )?
1297 .ok_or_else(|| {
1298 EstimationError::InvalidInput(format!(
1299 "failed to build {} hyper_dirs at current {}",
1300 kind.adjective(),
1301 kind.coord_name(),
1302 ))
1303 })?;
1304 self.cached_hyper_dirs = Some((revision, dirs.clone()));
1305 Ok(dirs)
1306 }
1307
1308 fn nfree_tensor_gradient_hyper_dirs(
1309 &mut self,
1310 theta: &Array1<f64>,
1311 ) -> Result<Vec<DirectionalHyperParam>, EstimationError> {
1312 let psi = &theta.as_slice().ok_or_else(|| {
1313 EstimationError::InvalidInput(
1314 "nfree_tensor_gradient_hyper_dirs: theta is not contiguous".to_string(),
1315 )
1316 })?[self.rho_dim..];
1317 let (global_range, p_total, s_psi_components) = self
1318 .realizer
1319 .canonical_penalty_derivatives_at_psi(&self.spatial_terms, psi)
1320 .map_err(EstimationError::InvalidInput)?;
1321 let zero_x = gam_solve::estimate::reml::HyperDesignDerivative::zero(
1322 self.realizer.design().design.nrows(),
1323 p_total,
1324 );
1325 let components = s_psi_components
1326 .into_iter()
1327 .enumerate()
1328 .map(|(penalty_index, local)| {
1329 (
1330 penalty_index,
1331 gam_solve::estimate::reml::HyperPenaltyDerivative::from_embedded(
1332 local,
1333 global_range.clone(),
1334 p_total,
1335 ),
1336 )
1337 })
1338 .collect::<Vec<_>>();
1339 Ok(DirectionalHyperParam::new_compact(zero_x, components, None, None)?.not_penalty_like())
1340 .map(|dir| vec![dir])
1341 }
1342
1343 fn ensure_theta(&mut self, theta: &Array1<f64>) -> Result<(), String> {
1344 if self
1345 .current_theta
1346 .as_ref()
1347 .is_some_and(|cached| theta_values_match(cached, theta))
1348 {
1349 return Ok(());
1350 }
1351 let t_ensure = std::time::Instant::now();
1352 let log_kappa = SpatialLogKappaCoords::from_theta_tail_with_dims(
1353 theta,
1354 self.rho_dim,
1355 self.dims_per_term.clone(),
1356 );
1357 self.realizer
1358 .apply_log_kappa(&log_kappa, &self.spatial_terms)?;
1359 log::info!(
1360 "[STAGE] ensure_theta (apply_log_kappa, {} terms): {:.3}s",
1361 self.spatial_terms.len(),
1362 t_ensure.elapsed().as_secs_f64(),
1363 );
1364 self.current_theta = Some(theta.clone());
1365 self.last_eval_theta = None;
1366 self.last_cost = None;
1367 self.last_eval = None;
1368 Ok(())
1369 }
1370
1371 fn memoized_cost(&self, theta: &Array1<f64>) -> Option<f64> {
1378 if self
1379 .last_eval_theta
1380 .as_ref()
1381 .is_some_and(|cached| theta_values_match(cached, theta))
1382 {
1383 self.last_eval
1384 .as_ref()
1385 .map(|cached| cached.0)
1386 .or(self.last_cost)
1387 } else {
1388 None
1389 }
1390 }
1391
1392 fn memoized_eval(
1393 &self,
1394 theta: &Array1<f64>,
1395 ) -> Option<(f64, Array1<f64>, gam_problem::HessianValue)> {
1396 if self
1397 .last_eval_theta
1398 .as_ref()
1399 .is_some_and(|cached| theta_values_match(cached, theta))
1400 {
1401 self.last_eval.clone()
1402 } else {
1403 None
1404 }
1405 }
1406
1407 fn store_eval_at(
1411 &mut self,
1412 theta: &Array1<f64>,
1413 eval: (f64, Array1<f64>, gam_problem::HessianValue),
1414 ) {
1415 self.last_eval_theta = Some(theta.clone());
1416 self.last_cost = Some(eval.0);
1417 self.last_eval = Some(eval);
1418 }
1419
1420 fn store_cost_at(&mut self, theta: &Array1<f64>, cost: f64) {
1423 self.last_eval_theta = Some(theta.clone());
1424 self.last_cost = Some(cost);
1425 self.last_eval = None;
1429 }
1430
1431 fn spec(&self) -> &TermCollectionSpec {
1432 self.realizer.spec()
1433 }
1434
1435 fn design(&self) -> &TermCollectionDesign {
1436 self.realizer.design()
1437 }
1438
1439 fn supports_nfree_penalty_rekey(&self) -> bool {
1445 self.realizer
1446 .supports_nfree_penalty_rekey(&self.spatial_terms)
1447 }
1448
1449 fn supports_nfree_gradient_only_routing(&self) -> bool {
1450 self.realizer
1451 .supports_nfree_gradient_only_routing(&self.spatial_terms)
1452 }
1453
1454 fn canonical_penalties_at(
1464 &mut self,
1465 theta: &Array1<f64>,
1466 ) -> Result<(Vec<gam_terms::construction::CanonicalPenalty>, Vec<usize>), String> {
1467 let psi = &theta
1468 .as_slice()
1469 .ok_or_else(|| "canonical_penalties_at: theta is not contiguous".to_string())?
1470 [self.rho_dim..];
1471 self.realizer
1472 .canonical_penalties_at_psi(&self.spatial_terms, psi)
1473 }
1474}
1475
1476struct SingleBlockLatentCoordDesignCache {
1477 data: Array2<f64>,
1478 spec: TermCollectionSpec,
1479 design: TermCollectionDesign,
1480 current_theta: Option<Array1<f64>>,
1481 current_latent: Option<std::sync::Arc<gam_terms::latent::LatentCoordValues>>,
1482 current_hyper_dirs: Option<Vec<gam_solve::estimate::reml::DirectionalHyperParam>>,
1483 current_design_cache_id: Option<u64>,
1484 latent_design_cache: gam_solve::latent_cache::LatentDesignCache,
1485 last_cost: Option<f64>,
1486 last_eval: Option<(f64, Array1<f64>, gam_problem::HessianValue)>,
1487 term_index: gam_problem::types::SmoothTermIdx,
1488 feature_cols: Vec<usize>,
1489 rho_dim: usize,
1490 n_obs: usize,
1491 latent_dim: usize,
1492 id_mode: gam_terms::latent::LatentIdMode,
1493 manifold: gam_terms::latent::LatentManifold,
1494 retraction_registry: gam_solve::latent_cache::LatentRetractionRegistry,
1495 latent_id: u64,
1496 analytic_penalties: Option<std::sync::Arc<gam_terms::AnalyticPenaltyRegistry>>,
1497 analytic_rho_count: usize,
1498 design_revision: u64,
1499 last_outer_iter: Option<u64>,
1503}
1504
1505impl SingleBlockLatentCoordDesignCache {
1506 fn new(
1507 data: Array2<f64>,
1508 spec: TermCollectionSpec,
1509 design: TermCollectionDesign,
1510 latent: &StandardLatentCoordConfig,
1511 rho_dim: usize,
1512 ) -> Result<Self, String> {
1513 if latent.term_index.get() >= spec.smooth_terms.len() {
1514 return Err(SmoothError::dimension_mismatch(format!(
1515 "latent-coordinate term index {} out of bounds for {} smooth terms",
1516 latent.term_index,
1517 spec.smooth_terms.len()
1518 ))
1519 .into());
1520 }
1521 if latent.feature_cols.len() != latent.values.latent_dim() {
1522 return Err(SmoothError::dimension_mismatch(format!(
1523 "latent-coordinate feature width mismatch: feature_cols={}, latent_dim={}",
1524 latent.feature_cols.len(),
1525 latent.values.latent_dim()
1526 ))
1527 .into());
1528 }
1529 if latent.values.n_obs() != data.nrows() {
1530 return Err(SmoothError::dimension_mismatch(format!(
1531 "latent-coordinate row mismatch: latent n={}, data n={}",
1532 latent.values.n_obs(),
1533 data.nrows()
1534 ))
1535 .into());
1536 }
1537 let analytic_rho_count = latent
1538 .analytic_penalties
1539 .as_ref()
1540 .map_or(0, |registry| registry.total_rho_count());
1541 Ok(Self {
1542 data,
1543 spec,
1544 design,
1545 current_theta: None,
1546 current_latent: None,
1547 current_hyper_dirs: None,
1548 current_design_cache_id: None,
1549 latent_design_cache: gam_solve::latent_cache::LatentDesignCache::default(),
1550 last_cost: None,
1551 last_eval: None,
1552 term_index: latent.term_index,
1553 feature_cols: latent.feature_cols.clone(),
1554 rho_dim,
1555 n_obs: latent.values.n_obs(),
1556 latent_dim: latent.values.latent_dim(),
1557 id_mode: latent.values.id_mode().clone(),
1558 manifold: latent.values.manifold().clone(),
1559 retraction_registry: latent.values.retraction_registry().clone(),
1560 latent_id: latent.values.latent_id(),
1561 analytic_penalties: latent.analytic_penalties.clone(),
1562 analytic_rho_count,
1563 design_revision: 0,
1564 last_outer_iter: None,
1565 })
1566 }
1567
1568 fn design_revision(&self) -> u64 {
1569 self.design_revision
1570 }
1571
1572 fn design(&self) -> &TermCollectionDesign {
1573 &self.design
1574 }
1575
1576 fn latent(&self) -> Result<std::sync::Arc<gam_terms::latent::LatentCoordValues>, String> {
1577 self.current_latent
1578 .as_ref()
1579 .cloned()
1580 .ok_or_else(|| "latent-coordinate cache has not been realized".to_string())
1581 }
1582
1583 fn analytic_penalties(&self) -> Option<std::sync::Arc<gam_terms::AnalyticPenaltyRegistry>> {
1584 self.analytic_penalties.clone()
1585 }
1586
1587 fn analytic_penalty_rho_count(&self) -> usize {
1588 self.analytic_rho_count
1589 }
1590
1591 fn hyper_dirs(&self) -> Result<Vec<gam_solve::estimate::reml::DirectionalHyperParam>, String> {
1592 self.current_hyper_dirs
1593 .as_ref()
1594 .cloned()
1595 .ok_or_else(|| "latent-coordinate hyper_dirs cache has not been realized".to_string())
1596 }
1597
1598 fn latent_basis_kind(&self) -> Result<gam_solve::latent_cache::LatentBasisKind, String> {
1599 let smooth_term = self
1600 .design
1601 .smooth
1602 .terms
1603 .get(self.term_index.get())
1604 .ok_or_else(|| {
1605 SmoothError::dimension_mismatch(format!(
1606 "LatentCoord term index {} out of bounds for realized smooth design",
1607 self.term_index
1608 ))
1609 })?;
1610 let termspec = self
1611 .spec
1612 .smooth_terms
1613 .get(self.term_index.get())
1614 .ok_or_else(|| {
1615 SmoothError::dimension_mismatch(format!(
1616 "LatentCoord term index {} out of bounds for resolved smooth spec",
1617 self.term_index
1618 ))
1619 })?;
1620 match (&termspec.basis, &smooth_term.metadata) {
1621 (
1622 SmoothBasisSpec::Matern { .. },
1623 BasisMetadata::Matern {
1624 centers,
1625 length_scale,
1626 nu,
1627 aniso_log_scales,
1628 ..
1629 },
1630 ) => Ok(gam_solve::latent_cache::LatentBasisKind::Matern {
1631 centers: centers.clone(),
1632 length_scale: *length_scale,
1633 nu: *nu,
1634 aniso_log_scales: aniso_log_scales
1635 .clone()
1636 .unwrap_or_else(|| vec![0.0; centers.ncols()]),
1637 chunk_size: gam_terms::basis::auto_streaming_chunk_size_for_dense(
1638 self.n_obs,
1639 centers.nrows(),
1640 ),
1641 }),
1642 (
1643 SmoothBasisSpec::Duchon { .. },
1644 BasisMetadata::Duchon {
1645 centers,
1646 length_scale,
1647 power,
1648 nullspace_order,
1649 aniso_log_scales,
1650 ..
1651 },
1652 ) => Ok(gam_solve::latent_cache::LatentBasisKind::Duchon {
1653 centers: centers.clone(),
1654 length_scale: *length_scale,
1655 power: *power,
1656 nullspace_order: *nullspace_order,
1657 aniso_log_scales: aniso_log_scales
1658 .clone()
1659 .unwrap_or_else(|| vec![0.0; centers.ncols()]),
1660 }),
1661 (
1662 SmoothBasisSpec::Sphere { .. },
1663 BasisMetadata::Sphere {
1664 centers,
1665 penalty_order,
1666 method,
1667 ..
1668 },
1669 ) if matches!(*method, gam_terms::basis::SphereMethod::Wahba) => {
1670 Ok(gam_solve::latent_cache::LatentBasisKind::Sphere {
1671 centers: centers.clone(),
1672 penalty_order: *penalty_order,
1673 chunk_size: gam_terms::basis::auto_streaming_chunk_size_for_dense(
1674 self.n_obs,
1675 centers.nrows(),
1676 ),
1677 })
1678 }
1679 (
1680 SmoothBasisSpec::BSpline1D { spec, .. },
1681 BasisMetadata::BSpline1D {
1682 knots,
1683 periodic,
1684 degree: meta_degree,
1685 ..
1686 },
1687 ) => {
1688 let effective_degree = meta_degree.unwrap_or(spec.degree);
1692 if let Some((domain_start, period, num_basis)) = periodic {
1693 Ok(gam_solve::latent_cache::LatentBasisKind::PeriodicBspline {
1694 domain_start: *domain_start,
1695 period: *period,
1696 degree: effective_degree,
1697 num_basis: *num_basis,
1698 chunk_size: gam_terms::basis::auto_streaming_chunk_size_for_dense(
1699 self.n_obs, *num_basis,
1700 ),
1701 })
1702 } else {
1703 let num_basis_est = knots.len().saturating_sub(effective_degree + 1);
1704 Ok(gam_solve::latent_cache::LatentBasisKind::TensorBspline {
1705 knots: vec![knots.clone()],
1706 degrees: vec![effective_degree],
1707 chunk_size: gam_terms::basis::auto_streaming_chunk_size_for_dense(
1708 self.n_obs,
1709 num_basis_est,
1710 ),
1711 })
1712 }
1713 }
1714 (
1715 SmoothBasisSpec::TensorBSpline { .. },
1716 BasisMetadata::TensorBSpline { knots, degrees, .. },
1717 ) => Ok(gam_solve::latent_cache::LatentBasisKind::TensorBspline {
1718 knots: knots.clone(),
1719 degrees: degrees.clone(),
1720 chunk_size: None,
1721 }),
1722 (
1723 SmoothBasisSpec::Pca { .. },
1724 BasisMetadata::Pca {
1725 basis_matrix,
1726 centered,
1727 smooth_penalty,
1728 center_mean,
1729 pca_basis_path,
1730 chunk_size,
1731 ..
1732 },
1733 ) => {
1734 let center_mean_fingerprint = if *centered && pca_basis_path.is_none() {
1735 let mean = center_mean.as_ref().ok_or_else(|| {
1736 SmoothError::invalid_config(
1737 "latent-coordinate Pca cache key requires center_mean when centered",
1738 )
1739 })?;
1740 Some(gam_solve::latent_cache::pca_center_mean_fingerprint(mean))
1741 } else {
1742 None
1743 };
1744 Ok(gam_solve::latent_cache::LatentBasisKind::Pca {
1745 basis_matrix: basis_matrix.clone(),
1746 centered: *centered,
1747 center_mean_fingerprint,
1748 smooth_penalty: *smooth_penalty,
1749 pca_basis_path: pca_basis_path.clone(),
1750 chunk_size: *chunk_size,
1751 })
1752 }
1753 _ => Err(SmoothError::invalid_config(
1754 "latent-coordinate design cache could not key the realized latent smooth basis"
1755 .to_string(),
1756 )
1757 .into()),
1758 }
1759 }
1760
1761 fn ensure_theta(&mut self, theta: &Array1<f64>) -> Result<(), String> {
1762 if self
1763 .current_theta
1764 .as_ref()
1765 .is_some_and(|cached| theta_values_match(cached, theta))
1766 {
1767 return Ok(());
1768 }
1769 let latent_flat_len = self.n_obs * self.latent_dim;
1770 let direct_hyper_count = latent_coord_direct_hyper_count(&self.id_mode, self.latent_dim);
1771 let expected =
1772 self.rho_dim + latent_flat_len + self.analytic_rho_count + direct_hyper_count;
1773 if theta.len() != expected {
1774 return Err(SmoothError::dimension_mismatch(format!(
1775 "latent-coordinate theta length mismatch: got {}, expected {} (rho_dim={}, n={}, d={}, analytic_rhos={}, direct_hypers={})",
1776 theta.len(),
1777 expected,
1778 self.rho_dim,
1779 self.n_obs,
1780 self.latent_dim,
1781 self.analytic_rho_count,
1782 direct_hyper_count
1783 ))
1784 .into());
1785 }
1786 let flat = theta
1787 .slice(s![self.rho_dim..self.rho_dim + latent_flat_len])
1788 .to_owned();
1789 let latent = std::sync::Arc::new(
1790 gam_terms::latent::LatentCoordValues::from_flat_with_manifold_and_retraction_and_id(
1791 flat,
1792 self.n_obs,
1793 self.latent_dim,
1794 self.id_mode.clone(),
1795 self.manifold.clone(),
1796 self.retraction_registry.clone(),
1797 self.latent_id,
1798 ),
1799 );
1800 let latent_values_changed = self
1801 .current_latent
1802 .as_ref()
1803 .map(|cached| !latent_values_match(cached.as_flat(), latent.as_flat()))
1804 .unwrap_or(true);
1805 if latent_values_changed {
1806 self.latent_design_cache.invalidate_all();
1807 self.current_design_cache_id = None;
1808 self.design_revision = self.design_revision.wrapping_add(1);
1809 }
1810 for n in 0..self.n_obs {
1811 for axis in 0..self.latent_dim {
1812 let col = self.feature_cols[axis];
1813 self.data[[n, col]] = latent.as_flat()[n * self.latent_dim + axis];
1814 }
1815 }
1816
1817 let basis_kind = self.latent_basis_kind()?;
1818 let rebuilt_width = self.design.design.ncols();
1819 let spec = self.spec.clone();
1820 let term_index = self.term_index;
1821 let analytic_rho_count = self.analytic_rho_count;
1822 let data = self.data.view();
1823 let design_context_digest = gam_solve::latent_cache::latent_design_context_cache_digest(
1824 data,
1825 &spec,
1826 term_index,
1827 analytic_rho_count,
1828 &self.feature_cols,
1829 )
1830 .map_err(|e| e.to_string())?;
1831 let lookup = self
1832 .latent_design_cache
1833 .lookup_or_compute(latent.clone(), basis_kind, design_context_digest, || {
1834 let rebuilt = build_term_collection_design(data, &spec).map_err(|e| {
1835 EstimationError::InvalidInput(format!(
1836 "failed to rebuild latent-coordinate design: {e}"
1837 ))
1838 })?;
1839 if rebuilt.design.ncols() != rebuilt_width {
1840 crate::bail_invalid_estim!(
1841 "latent-coordinate design topology changed: rebuilt p={}, cached p={}",
1842 rebuilt.design.ncols(),
1843 rebuilt_width
1844 );
1845 }
1846 let hyper_dirs = try_build_latent_coord_hyper_dirs(
1847 latent.clone(),
1848 &spec,
1849 &rebuilt,
1850 &[term_index],
1851 analytic_rho_count,
1852 )?
1853 .ok_or_else(|| {
1854 EstimationError::InvalidInput(
1855 "failed to build latent-coordinate hyper_dirs".to_string(),
1856 )
1857 })?;
1858 Ok(gam_solve::latent_cache::ComputedLatentDesign {
1859 design: rebuilt,
1860 hyper_dirs,
1861 })
1862 })
1863 .map_err(|e| e.to_string())?;
1864 if lookup.cached.design.design.ncols() != self.design.design.ncols() {
1865 return Err(SmoothError::dimension_mismatch(format!(
1866 "latent-coordinate design topology changed: rebuilt p={}, cached p={}",
1867 lookup.cached.design.design.ncols(),
1868 self.design.design.ncols()
1869 ))
1870 .into());
1871 }
1872 self.design = lookup.cached.design.clone();
1873 self.current_hyper_dirs = Some(lookup.cached.hyper_dirs.clone());
1874 self.current_latent = Some(latent);
1875 self.current_theta = Some(theta.clone());
1876 self.last_cost = None;
1877 self.last_eval = None;
1878 self.last_outer_iter = None;
1879 if !latent_values_changed && self.current_design_cache_id != Some(lookup.entry_id) {
1880 self.design_revision = self.design_revision.wrapping_add(1);
1881 }
1882 self.current_design_cache_id = Some(lookup.entry_id);
1883 Ok(())
1884 }
1885
1886 fn memoized_cost(&self, theta: &Array1<f64>) -> Option<f64> {
1887 if self
1888 .current_theta
1889 .as_ref()
1890 .is_some_and(|cached| theta_values_match(cached, theta))
1891 && self.last_outer_iter
1892 == Some(gam_solve::estimate::reml::outer_eval::current_outer_iter())
1893 {
1894 self.last_eval
1895 .as_ref()
1896 .map(|cached| cached.0)
1897 .or(self.last_cost)
1898 } else {
1899 None
1900 }
1901 }
1902
1903 fn memoized_eval(
1904 &self,
1905 theta: &Array1<f64>,
1906 ) -> Option<(f64, Array1<f64>, gam_problem::HessianValue)> {
1907 if self
1908 .current_theta
1909 .as_ref()
1910 .is_some_and(|cached| theta_values_match(cached, theta))
1911 && self.last_outer_iter
1912 == Some(gam_solve::estimate::reml::outer_eval::current_outer_iter())
1913 {
1914 self.last_eval.clone()
1915 } else {
1916 None
1917 }
1918 }
1919
1920 fn store_eval(&mut self, eval: (f64, Array1<f64>, gam_problem::HessianValue)) {
1921 self.last_cost = Some(eval.0);
1922 self.last_eval = Some(eval);
1923 self.last_outer_iter = Some(gam_solve::estimate::reml::outer_eval::current_outer_iter());
1924 }
1925
1926 fn store_cost(&mut self, cost: f64) {
1927 self.last_cost = Some(cost);
1928 self.last_outer_iter = Some(gam_solve::estimate::reml::outer_eval::current_outer_iter());
1929 }
1930
1931 fn reset(&mut self) {
1932 self.current_theta = None;
1933 self.current_latent = None;
1934 self.current_hyper_dirs = None;
1935 self.current_design_cache_id = None;
1936 self.latent_design_cache.invalidate();
1937 self.last_cost = None;
1938 self.last_eval = None;
1939 self.last_outer_iter = None;
1940 }
1941}
1942
1943pub fn fixed_kappa_profiled_reml_score(
1959 data: ArrayView2<'_, f64>,
1960 y: ArrayView1<'_, f64>,
1961 weights: ArrayView1<'_, f64>,
1962 offset: ArrayView1<'_, f64>,
1963 resolvedspec: &TermCollectionSpec,
1964 term_idx: usize,
1965 kappa: f64,
1966 family: LikelihoodSpec,
1967 options: &FitOptions,
1968) -> Result<f64, EstimationError> {
1969 if !kappa.is_finite() {
1970 crate::bail_invalid_estim!("fixed-κ profiled score probed a non-finite κ = {kappa}");
1971 }
1972 let (feature_cols, mut probe_basis) =
1975 match resolvedspec.smooth_terms.get(term_idx).map(|t| &t.basis) {
1976 Some(SmoothBasisSpec::ConstantCurvature {
1977 feature_cols, spec, ..
1978 }) => (feature_cols.clone(), spec.clone()),
1979 _ => {
1980 crate::bail_invalid_estim!(
1981 "fixed-κ profiled score: term {term_idx} is not a constant-curvature smooth"
1982 )
1983 }
1984 };
1985 probe_basis.kappa = kappa;
1986
1987 let is_unweighted = weights.iter().all(|&w| (w - 1.0).abs() <= 1e-12);
2007 let is_zero_offset = offset.iter().all(|&o| o.abs() <= 1e-12);
2008 if family == LikelihoodSpec::gaussian_identity() && is_unweighted && is_zero_offset {
2009 let x_term = select_columns(data, &feature_cols).map_err(EstimationError::from)?;
2010 let score = gam_terms::basis::constant_curvature_honest_profiled_reml_score(
2011 x_term.view(),
2012 y,
2013 &probe_basis,
2014 )
2015 .map_err(|e| {
2016 EstimationError::InvalidInput(format!(
2017 "fixed-κ honest profiled-REML score at κ={kappa} failed: {e}"
2018 ))
2019 })?;
2020 if !score.is_finite() {
2021 crate::bail_invalid_estim!(
2022 "fixed-κ honest profiled-REML score at κ={kappa} is non-finite"
2023 );
2024 }
2025 return Ok(score);
2026 }
2027
2028 let mut probe_spec = resolvedspec.clone();
2030 match probe_spec
2031 .smooth_terms
2032 .get_mut(term_idx)
2033 .map(|t| &mut t.basis)
2034 {
2035 Some(SmoothBasisSpec::ConstantCurvature { spec, .. }) => spec.kappa = kappa,
2036 _ => {
2037 crate::bail_invalid_estim!(
2038 "fixed-κ profiled score: term {term_idx} is not a constant-curvature smooth"
2039 )
2040 }
2041 }
2042 let fixed_kappa_options = SpatialLengthScaleOptimizationOptions {
2043 enabled: false,
2044 ..SpatialLengthScaleOptimizationOptions::default()
2045 };
2046 let fit = fit_term_collectionwith_spatial_length_scale_optimization(
2047 data,
2048 y.to_owned(),
2049 weights.to_owned(),
2050 offset.to_owned(),
2051 &probe_spec,
2052 family,
2053 options,
2054 &fixed_kappa_options,
2055 )?;
2056 let score = fit_score(&fit.fit);
2057 if !score.is_finite() {
2058 crate::bail_invalid_estim!("fixed-κ profiled fit at κ={kappa} returned a non-finite score");
2059 }
2060 Ok(score)
2061}
2062
2063fn profiled_gaussian_reml_value_kappa_gradient(
2068 design: &Array2<f64>,
2069 design_kappa: &Array2<f64>,
2070 penalty: &Array2<f64>,
2071 penalty_kappa: &Array2<f64>,
2072 response: ArrayView1<'_, f64>,
2073) -> Result<(f64, f64), EstimationError> {
2074 if design.dim() != design_kappa.dim()
2075 || penalty.dim() != penalty_kappa.dim()
2076 || penalty.dim() != (design.ncols(), design.ncols())
2077 || response.len() != design.nrows()
2078 {
2079 crate::bail_invalid_estim!("constant-curvature profile value/gradient shape mismatch");
2080 }
2081
2082 let response_2d = response.insert_axis(ndarray::Axis(1));
2083 let fit = gam_solve::gaussian_reml::gaussian_reml_multi_closed_form(
2084 design.view(),
2085 response_2d.view(),
2086 penalty.view(),
2087 None,
2088 None,
2089 )?;
2090 let backward = gam_solve::gaussian_reml::gaussian_reml_multi_closed_form_backward_from_fit(
2091 design.view(),
2092 response_2d.view(),
2093 penalty.view(),
2094 None,
2095 &fit,
2096 0.0,
2097 None,
2098 None,
2099 1.0,
2100 0.0,
2101 )?;
2102 let derivative = backward
2103 .grad_x
2104 .iter()
2105 .zip(design_kappa.iter())
2106 .map(|(&adjoint, &direction)| adjoint * direction)
2107 .sum::<f64>()
2108 + backward
2109 .grad_penalty
2110 .iter()
2111 .zip(penalty_kappa.iter())
2112 .map(|(&adjoint, &direction)| adjoint * direction)
2113 .sum::<f64>();
2114 if !(fit.reml_score.is_finite() && derivative.is_finite()) {
2115 crate::bail_invalid_estim!(
2116 "constant-curvature analytic profile returned a non-finite value or derivative"
2117 );
2118 }
2119 Ok((fit.reml_score, derivative))
2120}
2121
2122fn constant_curvature_radial_reference(
2126 data: ArrayView2<'_, f64>,
2127 y: ArrayView1<'_, f64>,
2128) -> Result<Array1<f64>, EstimationError> {
2129 if y.len() != data.nrows() || y.is_empty() {
2130 crate::bail_invalid_estim!(
2131 "constant-curvature radial reference needs one non-empty response per row"
2132 );
2133 }
2134 let radii: Array1<f64> = data.outer_iter().map(|row| row.dot(&row).sqrt()).collect();
2135 let r_max = radii.iter().copied().fold(0.0_f64, f64::max);
2136 if r_max <= f64::MIN_POSITIVE {
2137 let mean = y.sum() / y.len() as f64;
2138 return Ok(Array1::from_elem(y.len(), mean));
2139 }
2140
2141 let bin_count = (data.nrows() as f64).log2().ceil() as usize + 1;
2142 let bin_of = |radius: f64| -> usize {
2143 ((radius / r_max * bin_count as f64) as usize).min(bin_count - 1)
2144 };
2145 let mut sums = vec![0.0; bin_count];
2146 let mut counts = vec![0usize; bin_count];
2147 for (row, &radius) in radii.iter().enumerate() {
2148 let bin = bin_of(radius);
2149 sums[bin] += y[row];
2150 counts[bin] += 1;
2151 }
2152 let means: Vec<f64> = sums
2153 .into_iter()
2154 .zip(counts)
2155 .map(
2156 |(sum, count)| {
2157 if count == 0 { 0.0 } else { sum / count as f64 }
2158 },
2159 )
2160 .collect();
2161 Ok(radii.mapv(|radius| means[bin_of(radius)]))
2162}
2163
2164fn constant_curvature_kappa_fair_profile_value_gradient(
2169 data: ArrayView2<'_, f64>,
2170 y: ArrayView1<'_, f64>,
2171 y_ref: ArrayView1<'_, f64>,
2172 spec: &gam_terms::basis::ConstantCurvatureBasisSpec,
2173) -> Result<(f64, f64), EstimationError> {
2174 if y.len() != data.nrows() || y_ref.len() != data.nrows() {
2175 crate::bail_invalid_estim!(
2176 "constant-curvature fair profile row mismatch: data={}, response={}, reference={}",
2177 data.nrows(),
2178 y.len(),
2179 y_ref.len(),
2180 );
2181 }
2182
2183 let mut profile_spec = spec.clone();
2184 profile_spec.double_penalty = false;
2185 let basis = gam_terms::basis::build_constant_curvature_basis(data, &profile_spec)
2186 .map_err(EstimationError::from)?;
2187 let derivatives =
2188 gam_terms::basis::build_constant_curvature_basis_kappa_derivatives(data, &profile_spec)
2189 .map_err(EstimationError::from)?;
2190 if basis.penalties.len() != 1 || derivatives.first.penalties_derivative.len() != 1 {
2191 crate::bail_invalid_estim!(
2192 "constant-curvature fair profile expected one primary penalty; value blocks={}, derivative blocks={}",
2193 basis.penalties.len(),
2194 derivatives.first.penalties_derivative.len(),
2195 );
2196 }
2197
2198 let smooth_design = basis.design.to_dense();
2199 let smooth_design_kappa = &derivatives.first.design_derivative;
2200 let smooth_penalty = &basis.penalties[0];
2201 let smooth_penalty_kappa = &derivatives.first.penalties_derivative[0];
2202 let n = smooth_design.nrows();
2203 let p = smooth_design.ncols();
2204 if smooth_design_kappa.dim() != (n, p)
2205 || smooth_penalty.dim() != (p, p)
2206 || smooth_penalty_kappa.dim() != (p, p)
2207 {
2208 crate::bail_invalid_estim!(
2209 "constant-curvature kappa derivative bundle does not match its value basis"
2210 );
2211 }
2212
2213 let mut design = Array2::<f64>::ones((n, p + 1));
2214 design.slice_mut(s![.., 1..]).assign(&smooth_design);
2215 let mut design_kappa = Array2::<f64>::zeros((n, p + 1));
2216 design_kappa
2217 .slice_mut(s![.., 1..])
2218 .assign(smooth_design_kappa);
2219 let mut penalty = Array2::<f64>::zeros((p + 1, p + 1));
2220 penalty.slice_mut(s![1.., 1..]).assign(smooth_penalty);
2221 let mut penalty_kappa = Array2::<f64>::zeros((p + 1, p + 1));
2222 penalty_kappa
2223 .slice_mut(s![1.., 1..])
2224 .assign(smooth_penalty_kappa);
2225
2226 let (value_y, derivative_y) = profiled_gaussian_reml_value_kappa_gradient(
2227 &design,
2228 &design_kappa,
2229 &penalty,
2230 &penalty_kappa,
2231 y,
2232 )?;
2233 let (value_ref, derivative_ref) = profiled_gaussian_reml_value_kappa_gradient(
2234 &design,
2235 &design_kappa,
2236 &penalty,
2237 &penalty_kappa,
2238 y_ref,
2239 )?;
2240 Ok((value_y - value_ref, derivative_y - derivative_ref))
2241}
2242
2243struct ConstantCurvatureFairProfile<'a> {
2244 data: ArrayView2<'a, f64>,
2245 response: ArrayView1<'a, f64>,
2246 radial_reference: Array1<f64>,
2247 spec: gam_terms::basis::ConstantCurvatureBasisSpec,
2248 cache: std::cell::RefCell<std::collections::HashMap<u64, (f64, f64)>>,
2249}
2250
2251impl ConstantCurvatureFairProfile<'_> {
2252 fn evaluate(&self, kappa: f64) -> Result<(f64, f64), EstimationError> {
2253 if !kappa.is_finite() {
2254 crate::bail_invalid_estim!("constant-curvature fair profile probed a non-finite kappa");
2255 }
2256 let key = kappa.to_bits();
2257 if let Some(&cached) = self.cache.borrow().get(&key) {
2258 return Ok(cached);
2259 }
2260 let mut probe_spec = self.spec.clone();
2261 probe_spec.kappa = kappa;
2262 let sample = constant_curvature_kappa_fair_profile_value_gradient(
2263 self.data,
2264 self.response,
2265 self.radial_reference.view(),
2266 &probe_spec,
2267 )?;
2268 self.cache.borrow_mut().insert(key, sample);
2269 Ok(sample)
2270 }
2271}
2272
2273fn validate_constant_curvature_fair_profile_inputs(
2274 weights: ArrayView1<'_, f64>,
2275 offset: ArrayView1<'_, f64>,
2276 family: &LikelihoodSpec,
2277) -> Result<(), EstimationError> {
2278 if *family != LikelihoodSpec::gaussian_identity() {
2279 crate::bail_invalid_estim!(
2280 "curvature-as-an-estimand profile currently requires Gaussian identity likelihood"
2281 );
2282 }
2283 let input_tolerance = f64::EPSILON.sqrt();
2284 if weights
2285 .iter()
2286 .any(|&weight| (weight - 1.0).abs() > input_tolerance)
2287 || offset.iter().any(|&value| value.abs() > input_tolerance)
2288 {
2289 crate::bail_invalid_estim!(
2290 "curvature-as-an-estimand profile requires unit weights and zero offset"
2291 );
2292 }
2293 Ok(())
2294}
2295
2296fn constant_curvature_kappa_fair_optimum(
2303 data: ArrayView2<'_, f64>,
2304 y: ArrayView1<'_, f64>,
2305 resolvedspec: &TermCollectionSpec,
2306 term_idx: usize,
2307 options: &FitOptions,
2308) -> Result<f64, EstimationError> {
2309 let (kappa_min, kappa_max) = constant_curvature_kappa_bounds(data, resolvedspec, term_idx);
2310 if !(kappa_min.is_finite() && kappa_max.is_finite() && kappa_max > kappa_min) {
2311 crate::bail_invalid_estim!(
2312 "constant-curvature term {term_idx} has invalid kappa bounds [{kappa_min}, {kappa_max}]"
2313 );
2314 }
2315 let (feature_cols, base_spec) = match resolvedspec
2316 .smooth_terms
2317 .get(term_idx)
2318 .map(|term| &term.basis)
2319 {
2320 Some(SmoothBasisSpec::ConstantCurvature {
2321 feature_cols, spec, ..
2322 }) => (feature_cols, spec.clone()),
2323 _ => {
2324 crate::bail_invalid_estim!(
2325 "constant-curvature optimum requested for non-curvature term {term_idx}"
2326 )
2327 }
2328 };
2329 let x_term = select_columns(data, feature_cols).map_err(EstimationError::from)?;
2330 let y_ref = constant_curvature_radial_reference(x_term.view(), y)?;
2331 let profile = ConstantCurvatureFairProfile {
2332 data: x_term.view(),
2333 response: y,
2334 radial_reference: y_ref,
2335 spec: base_spec,
2336 cache: std::cell::RefCell::new(std::collections::HashMap::new()),
2337 };
2338 let mut seed_config = gam_problem::SeedConfig::default();
2339 seed_config.bounds = (kappa_min, kappa_max);
2340 seed_config.max_seeds = 1;
2341 seed_config.seed_budget = 1;
2342 seed_config.risk_profile = gam_problem::SeedRiskProfile::Gaussian;
2343 seed_config.num_auxiliary_trailing = 1;
2344 seed_config.over_smoothing_probe_rho = None;
2345 let initial_kappa = profile.spec.kappa.clamp(kappa_min, kappa_max);
2346 let problem = gam_solve::rho_optimizer::OuterProblem::new(1)
2347 .with_gradient(gam_problem::Derivative::Analytic)
2348 .with_hessian(gam_problem::DeclaredHessianForm::Unavailable)
2349 .with_prefer_gradient_only(true)
2350 .with_disable_fixed_point(true)
2351 .with_fallback_policy(gam_solve::rho_optimizer::FallbackPolicy::Disabled)
2352 .with_continuation_prewarm(false)
2353 .with_psi_dim(1)
2354 .with_tolerance(options.tol.max(f64::EPSILON.sqrt()))
2355 .with_max_iter(options.max_iter.max(1))
2356 .with_bounds(
2357 Array1::from_vec(vec![kappa_min]),
2358 Array1::from_vec(vec![kappa_max]),
2359 )
2360 .with_initial_rho(Array1::from_vec(vec![initial_kappa]))
2361 .with_seed_config(seed_config);
2362 let mut objective = problem.build_objective(
2363 profile,
2364 |profile: &mut ConstantCurvatureFairProfile<'_>, theta: &Array1<f64>| {
2365 profile.evaluate(theta[0]).map(|(value, _)| value)
2366 },
2367 |profile: &mut ConstantCurvatureFairProfile<'_>, theta: &Array1<f64>| {
2368 let (cost, derivative) = profile.evaluate(theta[0])?;
2369 Ok(gam_problem::OuterEval {
2370 cost,
2371 gradient: Array1::from_vec(vec![derivative]),
2372 hessian: gam_problem::HessianValue::Unavailable,
2373 inner_beta_hint: None,
2374 })
2375 },
2376 None::<fn(&mut ConstantCurvatureFairProfile<'_>)>,
2377 None::<
2378 fn(
2379 &mut ConstantCurvatureFairProfile<'_>,
2380 &Array1<f64>,
2381 ) -> Result<gam_problem::EfsEval, EstimationError>,
2382 >,
2383 );
2384 let result = problem.run(
2385 &mut objective,
2386 &format!("constant-curvature fair profile term {term_idx}"),
2387 )?;
2388 if !result.converged {
2389 crate::bail_invalid_estim!(
2390 "constant-curvature fair-profile κ optimization did not converge for term {} after {} iterations (negative_log_evidence={:.6e}, final_grad_norm={})",
2391 term_idx,
2392 result.iterations,
2393 result.final_value,
2394 result.final_grad_norm_report(),
2395 );
2396 }
2397 let kappa_hat = result.rho[0];
2398 log::info!(
2399 "[spatial-kappa] continuous fair-profile optimum kappa_hat={:.6} \
2400 (negative_log_evidence={:.6e}, projected_gradient={}) for term {term_idx}",
2401 kappa_hat,
2402 result.final_value,
2403 result.final_grad_norm_report(),
2404 );
2405 Ok(kappa_hat)
2406}
2407
2408fn try_exact_joint_spatial_length_scale_optimization(
2409 data: ArrayView2<'_, f64>,
2410 y: ArrayView1<'_, f64>,
2411 weights: ArrayView1<'_, f64>,
2412 offset: ArrayView1<'_, f64>,
2413 resolvedspec: &TermCollectionSpec,
2414 best: &FittedTermCollection,
2415 family: LikelihoodSpec,
2416 options: &FitOptions,
2417 kappa_options: &SpatialLengthScaleOptimizationOptions,
2418 spatial_terms: &[usize],
2419) -> Result<Option<FittedTermCollectionWithSpec>, EstimationError> {
2420 if spatial_terms.is_empty() {
2421 return Ok(None);
2422 }
2423 kappa_options
2428 .validate()
2429 .map_err(EstimationError::InvalidInput)?;
2430
2431 if try_build_spatial_log_kappa_hyper_dirs(data, resolvedspec, &best.design, spatial_terms)?
2432 .is_none()
2433 {
2434 if !constant_curvature_term_indices(resolvedspec).is_empty() {
2435 log::info!(
2436 "[#1464-trace] try_exact_joint RETURNED None (hyper_dirs unavailable); \
2437 κ̂ comes from a NON-joint path"
2438 );
2439 }
2440 return Ok(None);
2441 }
2442 if !constant_curvature_term_indices(resolvedspec).is_empty() {
2443 log::info!(
2444 "[#1464-trace] try_exact_joint ENTERED for {} spatial term(s); CC present",
2445 spatial_terms.len()
2446 );
2447 }
2448
2449 const JOINT_RHO_BOUND: f64 = 12.0;
2450 let rho_dim = best.fit.lambdas.len();
2451
2452 let has_constant_curvature_term = !constant_curvature_term_indices(resolvedspec).is_empty();
2466 let rho_upper_bound = if has_constant_curvature_term {
2467 gam_solve::estimate::RHO_BOUND
2468 } else {
2469 JOINT_RHO_BOUND
2470 };
2471
2472 let dims_per_term = spatial_dims_per_term(resolvedspec, spatial_terms);
2474 let use_aniso = has_aniso_terms(resolvedspec, spatial_terms);
2475
2476 let log_kappa0 = if use_aniso {
2481 SpatialLogKappaCoords::from_length_scales_aniso(resolvedspec, spatial_terms, kappa_options)
2482 } else {
2483 SpatialLogKappaCoords::from_length_scales(resolvedspec, spatial_terms, kappa_options)
2484 };
2485 let mut log_kappa0 =
2488 log_kappa0.reseed_from_data(data, resolvedspec, spatial_terms, kappa_options);
2489 let mut cc_profiled_values: Vec<(usize, f64)> = Vec::new();
2494 if has_constant_curvature_term {
2495 for (slot, &term_idx) in spatial_terms.iter().enumerate() {
2496 if constant_curvature_term_spec(resolvedspec, term_idx).is_none() {
2497 continue;
2498 }
2499 let kappa = get_constant_curvature_kappa(resolvedspec, term_idx)
2500 .expect("constant-curvature term exposes its kappa");
2501 log_kappa0.set_scalar_slot(slot, kappa);
2502 cc_profiled_values.push((slot, kappa));
2503 }
2504 }
2505 let log_kappa_lower = if use_aniso {
2506 SpatialLogKappaCoords::lower_bounds_aniso_from_data(
2507 data,
2508 resolvedspec,
2509 spatial_terms,
2510 &dims_per_term,
2511 kappa_options,
2512 )
2513 } else {
2514 SpatialLogKappaCoords::lower_bounds_from_data(
2515 data,
2516 resolvedspec,
2517 spatial_terms,
2518 kappa_options,
2519 )
2520 };
2521 let log_kappa_upper = if use_aniso {
2522 SpatialLogKappaCoords::upper_bounds_aniso_from_data(
2523 data,
2524 resolvedspec,
2525 spatial_terms,
2526 &dims_per_term,
2527 kappa_options,
2528 )
2529 } else {
2530 SpatialLogKappaCoords::upper_bounds_from_data(
2531 data,
2532 resolvedspec,
2533 spatial_terms,
2534 kappa_options,
2535 )
2536 };
2537 let mut log_kappa_lower = log_kappa_lower;
2538 let mut log_kappa_upper = log_kappa_upper;
2539 for &(slot, kappa) in &cc_profiled_values {
2540 log_kappa_lower.set_scalar_slot(slot, kappa);
2541 log_kappa_upper.set_scalar_slot(slot, kappa);
2542 log::info!("[spatial-kappa] slot {slot}: profiling rho at certified kappa={kappa}");
2543 }
2544 let log_kappa0 = log_kappa0.clamp_to_bounds(&log_kappa_lower, &log_kappa_upper);
2547 let setup = ExactJointHyperSetup::new(
2548 best.fit.lambdas.mapv(f64::ln),
2549 Array1::<f64>::from_elem(rho_dim, -JOINT_RHO_BOUND),
2550 Array1::<f64>::from_elem(rho_dim, rho_upper_bound),
2551 log_kappa0,
2552 log_kappa_lower,
2553 log_kappa_upper,
2554 );
2555
2556 let theta0 = setup.theta0();
2557 let lower = setup.lower();
2558 let upper = setup.upper();
2559
2560 let kind = if use_aniso {
2572 SpatialHyperKind::Anisotropic
2573 } else {
2574 SpatialHyperKind::Isotropic
2575 };
2576 let (theta_star, joint_final_value, kappa_timing) = run_exact_joint_spatial_optimization(
2577 kind,
2578 data,
2579 y,
2580 weights,
2581 offset,
2582 resolvedspec,
2583 &best.design,
2584 family.clone(),
2585 options,
2586 spatial_terms,
2587 &dims_per_term,
2588 &theta0,
2589 &lower,
2590 &upper,
2591 rho_dim,
2592 kappa_options,
2593 )?;
2594
2595 let baseline_score = fit_score(&best.fit);
2596
2597 let accept_tol = options.tol.max(1e-8 * baseline_score.abs()).max(1e-12);
2602 if joint_final_value > baseline_score + accept_tol {
2603 return Err(EstimationError::RemlOptimizationFailed(format!(
2604 "exact joint spatial optimization failed its objective-monotonicity certificate: \
2605 initial={baseline_score:.6e}, final={joint_final_value:.6e}, \
2606 acceptance_tolerance={accept_tol:.3e}, theta_checkpoint={:?}",
2607 theta_star.to_vec(),
2608 )));
2609 }
2610
2611 let rho_star = theta_star.slice(s![..rho_dim]).mapv(f64::exp);
2612 let log_kappa_star =
2613 SpatialLogKappaCoords::from_theta_tail_with_dims(&theta_star, rho_dim, dims_per_term);
2614 if has_constant_curvature_term {
2620 let star = log_kappa_star.as_array();
2621 let dims = log_kappa_star.dims_per_term();
2622 for (slot, &term_idx) in spatial_terms.iter().enumerate() {
2623 if constant_curvature_term_spec(resolvedspec, term_idx).is_some() {
2624 let off: usize = dims[..slot].iter().sum();
2625 log::info!(
2626 "[#1464-trace] term {term_idx}: joint solver CONVERGED ψ-tail κ = {} \
2627 (this is the optimised candidate; joint_final_value={joint_final_value})",
2628 star[off]
2629 );
2630 }
2631 }
2632 }
2633 let optimized_spec = log_kappa_star.apply_tospec(resolvedspec, spatial_terms)?;
2634 let optimized = fit_term_collection_forspecwith_heuristic_lambdas(
2635 data,
2636 y,
2637 weights,
2638 offset,
2639 &optimized_spec,
2640 rho_star.as_slice(),
2641 family.clone(),
2642 options,
2643 )?;
2644
2645 let mut fit = optimized.fit;
2649 fit.reml_score = joint_final_value;
2650 let optimized_result = FittedTermCollectionWithSpec {
2651 fit,
2652 design: optimized.design,
2653 resolvedspec: optimized_spec,
2654 adaptive_diagnostics: optimized.adaptive_diagnostics,
2655 kappa_timing: Some(kappa_timing),
2656 };
2657
2658 Ok(Some(optimized_result))
2659}
2660
2661#[derive(Clone, Copy, PartialEq, Eq, Debug)]
2673enum SpatialHyperKind {
2674 Anisotropic,
2675 Isotropic,
2676}
2677
2678impl SpatialHyperKind {
2679 fn label(self) -> &'static str {
2682 match self {
2683 SpatialHyperKind::Anisotropic => "spatial-aniso-joint",
2684 SpatialHyperKind::Isotropic => "spatial-iso-joint",
2685 }
2686 }
2687
2688 fn adjective(self) -> &'static str {
2690 match self {
2691 SpatialHyperKind::Anisotropic => "anisotropic",
2692 SpatialHyperKind::Isotropic => "isotropic",
2693 }
2694 }
2695
2696 fn coord_name(self) -> &'static str {
2699 match self {
2700 SpatialHyperKind::Anisotropic => "psi",
2701 SpatialHyperKind::Isotropic => "kappa",
2702 }
2703 }
2704}
2705
2706struct SpatialFrozenGlmInputs {
2712 y: Array1<f64>,
2713 weights: Array1<f64>,
2714 offset: Array1<f64>,
2715 family: LikelihoodSpec,
2716}
2717
2718fn frozen_glm_tensor_eligible_family(family: &LikelihoodSpec) -> bool {
2735 !family.is_gaussian_identity()
2736 && matches!(
2737 &family.response,
2738 ResponseFamily::Binomial
2739 | ResponseFamily::Poisson
2740 | ResponseFamily::Gamma
2741 | ResponseFamily::NegativeBinomial { .. }
2742 )
2743}
2744
2745struct SpatialJointContext<'d> {
2746 data: ArrayView2<'d, f64>,
2747 rho_dim: usize,
2748 kind: SpatialHyperKind,
2749 cache: SingleBlockExactJointDesignCache<'d>,
2750 evaluator: gam_solve::estimate::ExternalJointHyperEvaluator<'d>,
2751 frozen_glm_inputs: Option<SpatialFrozenGlmInputs>,
2752 frozen_glm_psi_bounds: Option<(f64, f64)>,
2753 frozen_glm_tensor: Option<gam_solve::glm_sufficient_lane::FrozenWeightGramTensor>,
2754 frozen_glm_tensor_attempted: bool,
2755 frozen_glm_weight_memo: Option<(Array1<f64>, Array1<f64>)>,
2767}
2768
2769#[derive(Clone, Copy, Debug, Default)]
2770struct NfreeSkipGateStatus {
2771 shape: bool,
2772 value: bool,
2773 gradient: bool,
2774 penalty: bool,
2775 revision: bool,
2776 second_order: bool,
2777}
2778
2779impl NfreeSkipGateStatus {
2780 fn would_skip(self, require_gradient: bool) -> bool {
2781 self.shape
2782 && self.value
2783 && (!require_gradient || self.gradient)
2784 && self.penalty
2785 && self.revision
2786 && !self.second_order
2787 }
2788}
2789
2790fn nfree_skip_gate_status_from_parts(
2791 shape: bool,
2792 covers_value: bool,
2793 covers_skip: bool,
2794 covers_gradient: bool,
2795 penalty: bool,
2796 revision: bool,
2797 allow_second_order: bool,
2798 require_gradient: bool,
2799) -> NfreeSkipGateStatus {
2800 NfreeSkipGateStatus {
2801 shape,
2802 value: shape && covers_value && (!require_gradient || covers_skip),
2810 gradient: shape && (!require_gradient || covers_gradient),
2811 penalty,
2812 revision,
2813 second_order: allow_second_order,
2814 }
2815}
2816
2817impl<'d> SpatialJointContext<'d> {
2818 fn nfree_skip_gate_status(
2819 &self,
2820 theta: &Array1<f64>,
2821 allow_second_order: bool,
2822 require_gradient: bool,
2823 ) -> NfreeSkipGateStatus {
2824 let shape = theta.len() == self.rho_dim + 1;
2825 let (covers_value, covers_skip, covers_gradient) = if shape {
2826 let psi = theta[self.rho_dim];
2827 (
2828 self.evaluator.psi_gram_tensor_covers(psi),
2829 self.evaluator.psi_gram_tensor_covers_skip(psi),
2830 self.evaluator.psi_gram_tensor_covers_gradient(psi),
2831 )
2832 } else {
2833 (false, false, false)
2834 };
2835 nfree_skip_gate_status_from_parts(
2836 shape,
2837 covers_value,
2838 covers_skip,
2839 covers_gradient,
2840 self.evaluator.supports_nfree_penalty_rekey(),
2841 self.evaluator.nfree_fast_path_revision().is_some(),
2842 allow_second_order,
2843 require_gradient,
2844 )
2845 }
2846
2847 fn frozen_glm_working_state(
2848 &self,
2849 beta: &Array1<f64>,
2850 ) -> Result<Option<(Array1<f64>, Array1<f64>)>, EstimationError> {
2851 let Some(inputs) = self.frozen_glm_inputs.as_ref() else {
2852 return Ok(None);
2853 };
2854 if beta.len() != self.cache.design().design.ncols() {
2855 return Ok(None);
2856 }
2857 let mut eta = self.cache.design().design.matrixvectormultiply(beta);
2858 if eta.len() != inputs.offset.len() {
2859 crate::bail_invalid_estim!(
2860 "frozen GLM tensor warm-state row mismatch: eta={}, offset={}",
2861 eta.len(),
2862 inputs.offset.len()
2863 );
2864 }
2865 eta += &inputs.offset;
2866 let obs = evaluate_standard_familyobservations(
2867 inputs.family.clone(),
2868 None,
2869 None,
2870 None,
2871 &inputs.y,
2872 &inputs.weights,
2873 &eta,
2874 )?;
2875 let mut working_response = obs.eta.clone();
2876 for i in 0..working_response.len() {
2877 let wi = obs.fisherweight[i].max(1e-12);
2878 working_response[i] += obs.score[i] / wi;
2879 }
2880 Ok(Some((obs.fisherweight, working_response)))
2881 }
2882
2883 fn frozen_glm_trial_weights(
2892 &mut self,
2893 beta: &Array1<f64>,
2894 ) -> Result<Option<Array1<f64>>, EstimationError> {
2895 if let Some((memo_beta, memo_w)) = self.frozen_glm_weight_memo.as_ref()
2896 && memo_beta.len() == beta.len()
2897 && memo_beta
2898 .iter()
2899 .zip(beta.iter())
2900 .all(|(a, b)| a.to_bits() == b.to_bits())
2901 {
2902 return Ok(Some(memo_w.clone()));
2903 }
2904 match self.frozen_glm_working_state(beta)? {
2905 Some((current_w, _)) => {
2906 self.frozen_glm_weight_memo = Some((beta.clone(), current_w.clone()));
2907 Ok(Some(current_w))
2908 }
2909 None => Ok(None),
2910 }
2911 }
2912
2913 fn ensure_frozen_glm_tensor(
2914 &mut self,
2915 theta: &Array1<f64>,
2916 warm_beta: Option<&Array1<f64>>,
2917 ) -> Result<(), EstimationError> {
2918 if self.frozen_glm_tensor.is_some() || self.frozen_glm_tensor_attempted {
2919 return Ok(());
2920 }
2921 let Some((psi_lo, psi_hi)) = self.frozen_glm_psi_bounds else {
2922 return Ok(());
2923 };
2924 if theta.len() != self.rho_dim + 1 {
2925 self.frozen_glm_tensor_attempted = true;
2926 return Ok(());
2927 }
2928 let Some(beta) = warm_beta else {
2929 return Ok(());
2930 };
2931 let Some((frozen_w, working_z)) = self.frozen_glm_working_state(beta)? else {
2932 self.frozen_glm_tensor_attempted = true;
2933 return Ok(());
2934 };
2935 let theta_probe_base = theta.clone();
2936 let rho_dim = self.rho_dim;
2937 let Self {
2944 cache, evaluator, ..
2945 } = self;
2946 let tensor = evaluator.build_frozen_glm_gram_tensor(
2947 |psi| {
2948 let mut theta_probe = theta_probe_base.clone();
2949 theta_probe[rho_dim] = psi;
2950 cache.ensure_theta(&theta_probe)?;
2951 Ok(cache.design().design.clone())
2952 },
2953 frozen_w.view(),
2954 working_z.view(),
2955 psi_lo,
2956 psi_hi,
2957 );
2958 self.cache
2959 .ensure_theta(theta)
2960 .map_err(EstimationError::InvalidInput)?;
2961 self.frozen_glm_tensor_attempted = true;
2962 if let Some(tensor) = tensor {
2963 self.frozen_glm_tensor = Some(tensor);
2964 log::info!(
2965 "[STAGE] {} certified frozen-W GLM ψ tensor over [{psi_lo:.3}, {psi_hi:.3}]",
2966 self.kind.label(),
2967 );
2968 } else {
2969 log::info!(
2970 "[STAGE] {} frozen-W GLM ψ tensor did not certify over [{psi_lo:.3}, {psi_hi:.3}]",
2971 self.kind.label(),
2972 );
2973 }
2974 Ok(())
2975 }
2976
2977 fn stage_frozen_glm_trial_statistics(
2978 &mut self,
2979 theta: &Array1<f64>,
2980 warm_beta: Option<&Array1<f64>>,
2981 allow_gradient: bool,
2982 ) -> Result<(), EstimationError> {
2983 let kind = self.kind;
2984 let mut staged_gram: Option<Array2<f64>> = None;
2985 let mut staged_deriv: Option<(Array2<f64>, Array1<f64>)> = None;
2986 if theta.len() == self.rho_dim + 1 {
2987 let psi = theta[self.rho_dim];
2988 let tensor_covers = self
2995 .frozen_glm_tensor
2996 .as_ref()
2997 .is_some_and(|t| t.contains(psi));
2998 let current_w = if tensor_covers {
2999 match warm_beta {
3000 Some(beta) => self.frozen_glm_trial_weights(beta)?,
3001 None => None,
3002 }
3003 } else {
3004 None
3005 };
3006 if let (Some(tensor), Some(current_w)) =
3007 (self.frozen_glm_tensor.as_ref(), current_w.as_ref())
3008 {
3009 const FROZEN_GLM_WEIGHT_DRIFT_RTOL: f64 = 1e-3;
3010 if tensor.weight_drift_within(current_w.view(), FROZEN_GLM_WEIGHT_DRIFT_RTOL) {
3011 staged_gram = Some(tensor.gram_at(psi));
3012 log::debug!(
3013 "[STAGE] {} trial at psi={psi:.6}: serving frozen-W GLM \
3014 first-Fisher-step XᵀWX n-free (weight drift within tol)",
3015 kind.label(),
3016 );
3017 }
3018 if allow_gradient
3019 && tensor.contains_for_gradient(psi)
3020 && let Some((dgram_dpsi, drhs_dpsi)) =
3021 tensor.gradient_pair_if_sound(psi, current_w.view())
3022 {
3023 staged_deriv = Some((dgram_dpsi, drhs_dpsi));
3024 log::debug!(
3025 "[STAGE] {} trial at psi={psi:.6}: serving frozen-W GLM \
3026 ψ-gradient (∂G/∂ψ, ∂b/∂ψ) n-free (gradient weight drift within \
3027 tight tol); B_j stays exact",
3028 kind.label(),
3029 );
3030 }
3031 }
3032 }
3033 self.evaluator.stage_glm_first_step_gram(staged_gram);
3034 self.evaluator.stage_glm_psi_gram_deriv(staged_deriv);
3035 Ok(())
3036 }
3037
3038 fn eval_full(
3040 &mut self,
3041 theta: &Array1<f64>,
3042 order: gam_solve::rho_optimizer::OuterEvalOrder,
3043 analytic_outer_hessian_available: bool,
3044 ) -> Result<(f64, Array1<f64>, gam_problem::HessianValue), EstimationError> {
3045 use gam_solve::rho_optimizer::OuterEvalOrder;
3046 let allow_second_order = matches!(order, OuterEvalOrder::ValueGradientHessian)
3047 && analytic_outer_hessian_available;
3048 if let Some(eval) = self.cache.memoized_eval(theta) {
3049 let cached_satisfies_order = !allow_second_order || eval.2.is_analytic();
3050 if cached_satisfies_order {
3051 return Ok(eval);
3052 }
3053 }
3054 let kind = self.kind;
3055 let nfree_fast_path_revision = self.evaluator.nfree_fast_path_revision();
3091 let skip_design_realization = !allow_second_order && theta.len() == self.rho_dim + 1 && {
3092 let psi = theta[self.rho_dim];
3093 self.evaluator.psi_gram_tensor_covers(psi)
3094 && self.evaluator.psi_gram_tensor_covers_gradient(psi)
3101 && self.evaluator.psi_gram_tensor_covers_skip(psi)
3118 && self.evaluator.supports_nfree_penalty_rekey()
3123 && nfree_fast_path_revision.is_some()
3124 };
3125 if skip_design_realization {
3137 log::debug!(
3138 "[STAGE] {} eval_full at psi={:.6}: skipping n×k design re-realization \
3139 + reconditioning — criterion/gradient/inner-solve served n-free from \
3140 the certified ψ-gram tensor (GaussianFixedCache + k-space ψ-derivatives)",
3141 kind.label(),
3142 theta[self.rho_dim],
3143 );
3144 } else {
3145 self.cache
3146 .ensure_theta(theta)
3147 .map_err(EstimationError::InvalidInput)?;
3148 }
3149 let warm_beta = self.evaluator.current_beta();
3150 self.ensure_frozen_glm_tensor(theta, warm_beta.as_ref())?;
3151 self.stage_frozen_glm_trial_statistics(theta, warm_beta.as_ref(), !allow_second_order)?;
3159 let hyper_dirs = if skip_design_realization {
3166 self.cache.nfree_tensor_gradient_hyper_dirs(theta)?
3167 } else {
3168 self.cache.hyper_dirs_for_current_design(self.data, kind)?
3169 };
3170
3171 let design_revision = if skip_design_realization {
3172 nfree_fast_path_revision
3173 } else {
3174 Some(self.cache.design_revision())
3175 };
3176 if self.evaluator.supports_nfree_penalty_rekey() {
3190 match self.cache.canonical_penalties_at(theta) {
3191 Ok(penalty) => self.evaluator.stage_fast_path_penalty(Some(penalty)),
3192 Err(e) => {
3193 log::warn!(
3194 "[STAGE] {} eval_full at psi={:.6}: exact n-free S(ψ) rebuild failed \
3195 ({e}); clearing stage (eval falls to slow path)",
3196 kind.label(),
3197 theta[self.rho_dim],
3198 );
3199 self.evaluator.stage_fast_path_penalty(None);
3200 }
3201 }
3202 }
3203 let eval = evaluate_joint_reml_outer_eval_at_theta(
3210 &mut self.evaluator,
3211 self.cache.design(),
3212 theta,
3213 self.rho_dim,
3214 hyper_dirs,
3215 warm_beta.as_ref().map(|b: &Array1<f64>| b.view()),
3216 if allow_second_order {
3217 order
3218 } else {
3219 OuterEvalOrder::ValueAndGradient
3220 },
3221 design_revision,
3222 );
3223 if let Ok(ref value) = eval {
3224 self.cache.store_eval_at(theta, value.clone());
3225 }
3226 eval
3227 }
3228
3229 fn eval_efs(&mut self, theta: &Array1<f64>) -> Result<gam_problem::EfsEval, EstimationError> {
3230 self.cache
3231 .ensure_theta(theta)
3232 .map_err(EstimationError::InvalidInput)?;
3233 let kind = self.kind;
3234 let hyper_dirs = try_build_spatial_log_kappa_hyper_dirs(
3235 self.data,
3236 self.cache.spec(),
3237 self.cache.design(),
3238 &self.cache.spatial_terms,
3239 )?
3240 .ok_or_else(|| {
3241 EstimationError::InvalidInput(format!(
3242 "failed to build {} hyper_dirs for exact-joint EFS",
3243 kind.adjective(),
3244 ))
3245 })?;
3246 let design_revision = Some(self.cache.design_revision());
3247 let warm_beta = self.evaluator.current_beta();
3248 evaluate_joint_reml_efs_at_theta(
3249 &mut self.evaluator,
3250 self.cache.design(),
3251 theta,
3252 self.rho_dim,
3253 hyper_dirs,
3254 warm_beta.as_ref().map(|b: &Array1<f64>| b.view()),
3255 design_revision,
3256 )
3257 }
3258
3259 fn eval_cost(&mut self, theta: &Array1<f64>) -> f64 {
3265 if let Some(cost) = self.cache.memoized_cost(theta) {
3266 return cost;
3267 }
3268 let probe_start = std::time::Instant::now();
3283 let psi_distance = self
3284 .cache
3285 .current_theta
3286 .as_ref()
3287 .filter(|reference| reference.len() == theta.len())
3288 .map(|reference| {
3289 reference
3290 .iter()
3291 .zip(theta.iter())
3292 .map(|(a, b)| (a - b) * (a - b))
3293 .sum::<f64>()
3294 .sqrt()
3295 })
3296 .unwrap_or(f64::NAN);
3297 let nfree_fast_path_revision = self.evaluator.nfree_fast_path_revision();
3311 let skip_value_realization = theta.len() == self.rho_dim + 1 && {
3312 let psi = theta[self.rho_dim];
3313 self.evaluator.psi_gram_tensor_covers(psi)
3314 && self.evaluator.supports_nfree_penalty_rekey()
3348 && nfree_fast_path_revision.is_some()
3349 };
3350 if theta.len() == self.rho_dim + 1
3351 && self.evaluator.has_psi_gram_tensor()
3352 && !self.evaluator.psi_gram_tensor_covers(theta[self.rho_dim])
3353 {
3354 self.cache.store_cost_at(theta, f64::INFINITY);
3355 return f64::INFINITY;
3356 }
3357 if !skip_value_realization && self.cache.ensure_theta(theta).is_err() {
3358 return f64::INFINITY;
3359 }
3360 if self.evaluator.supports_nfree_penalty_rekey() {
3366 match self.cache.canonical_penalties_at(theta) {
3367 Ok(penalty) => self.evaluator.stage_fast_path_penalty(Some(penalty)),
3368 Err(_) => self.evaluator.stage_fast_path_penalty(None),
3369 }
3370 }
3371 let warm_beta = self.evaluator.current_beta();
3372 if let Err(err) = self.ensure_frozen_glm_tensor(theta, warm_beta.as_ref()) {
3373 log::warn!(
3374 "[STAGE] {} value-probe at psi={:.6}: frozen-W GLM tensor setup failed ({err}); \
3375 falling back to exact streamed Gram",
3376 self.kind.label(),
3377 if theta.len() > self.rho_dim {
3378 theta[self.rho_dim]
3379 } else {
3380 f64::NAN
3381 },
3382 );
3383 self.evaluator.stage_glm_first_step_gram(None);
3384 self.evaluator.stage_glm_psi_gram_deriv(None);
3385 } else if let Err(err) =
3386 self.stage_frozen_glm_trial_statistics(theta, warm_beta.as_ref(), false)
3387 {
3388 log::warn!(
3389 "[STAGE] {} value-probe at psi={:.6}: frozen-W GLM staging failed ({err}); \
3390 falling back to exact streamed Gram",
3391 self.kind.label(),
3392 if theta.len() > self.rho_dim {
3393 theta[self.rho_dim]
3394 } else {
3395 f64::NAN
3396 },
3397 );
3398 self.evaluator.stage_glm_first_step_gram(None);
3399 self.evaluator.stage_glm_psi_gram_deriv(None);
3400 }
3401 let design_revision = if skip_value_realization {
3402 nfree_fast_path_revision
3403 } else {
3404 Some(self.cache.design_revision())
3405 };
3406 let cost_label = self.kind.label();
3407 let result = {
3408 let design = self.cache.design();
3409 self.evaluator.evaluate_cost_only(
3410 &design.design,
3411 &design.penalties,
3412 &design.nullspace_dims,
3413 design.linear_constraints.clone(),
3414 theta,
3415 self.rho_dim,
3416 warm_beta.as_ref().map(|b: &Array1<f64>| b.view()),
3417 cost_label,
3418 design_revision,
3419 )
3420 };
3421 match result {
3422 Ok(cost) => {
3423 log::debug!(
3424 "[STAGE] {cost_label} value-probe (order=Value): elapsed={:.3}s \
3425 cost={cost:.6e} trial_theta_distance={psi_distance:.3e}",
3426 probe_start.elapsed().as_secs_f64(),
3427 );
3428 self.cache.store_cost_at(theta, cost);
3429 cost
3430 }
3431 Err(_) => f64::INFINITY,
3432 }
3433 }
3434
3435 fn reset(&mut self) {
3436 self.cache.current_theta = None;
3437 self.cache.last_eval_theta = None;
3438 self.cache.last_cost = None;
3439 self.cache.last_eval = None;
3440 }
3441}
3442
3443fn kphase_log_norms(theta: &Array1<f64>, rho_dim: usize) -> (f64, f64) {
3467 let theta_norm = theta.iter().map(|v| v * v).sum::<f64>().sqrt();
3468 let log_kappa_norm = theta
3469 .iter()
3470 .skip(rho_dim)
3471 .map(|v| v * v)
3472 .sum::<f64>()
3473 .sqrt();
3474 (theta_norm, log_kappa_norm)
3475}
3476
3477fn run_exact_joint_spatial_optimization(
3478 kind: SpatialHyperKind,
3479 data: ArrayView2<'_, f64>,
3480 y: ArrayView1<'_, f64>,
3481 weights: ArrayView1<'_, f64>,
3482 offset: ArrayView1<'_, f64>,
3483 resolvedspec: &TermCollectionSpec,
3484 baseline_design: &TermCollectionDesign,
3485 family: LikelihoodSpec,
3486 options: &FitOptions,
3487 spatial_terms: &[usize],
3488 dims_per_term: &[usize],
3489 theta0: &Array1<f64>,
3490 lower: &Array1<f64>,
3491 upper: &Array1<f64>,
3492 rho_dim: usize,
3493 kappa_options: &SpatialLengthScaleOptimizationOptions,
3494) -> Result<(Array1<f64>, f64, SpatialLengthScaleOptimizationTiming), EstimationError> {
3495 let label = kind.label();
3496 assert!(
3498 lower.len() == theta0.len() && upper.len() == theta0.len(),
3499 "spatial hyperparameter bounds must match theta length: lower_len={}, upper_len={}, theta_len={}",
3500 lower.len(),
3501 upper.len(),
3502 theta0.len()
3503 );
3504 assert!(
3505 baseline_design.smooth.terms.len() >= spatial_terms.len(),
3506 "baseline design must have at least one smooth term per spatial term: baseline_terms={}, spatial_terms={}",
3507 baseline_design.smooth.terms.len(),
3508 spatial_terms.len()
3509 );
3510 use gam_problem::{DeclaredHessianForm, Derivative, OuterEval};
3511 use gam_solve::rho_optimizer::OuterEvalOrder;
3512
3513 let theta_dim = theta0.len();
3514 let coord_dim = theta_dim - rho_dim;
3517 let analytic_outer_hessian_available =
3527 exact_joint_spatial_outer_hessian_available(&family, baseline_design);
3528 if !analytic_outer_hessian_available {
3529 log::info!(
3530 "[{label}] analytic outer Hessian unavailable for family/design; routing without second-order geometry (coord_dim={coord_dim})"
3531 );
3532 }
3533 let mut prefer_gradient_only = theta_dim > EXACT_JOINT_SECOND_ORDER_THETA_CAP;
3539 if prefer_gradient_only {
3540 log::info!(
3541 "[{label}] joint θ-dim {theta_dim} exceeds the exact pair-Hessian budget \
3542 ({EXACT_JOINT_SECOND_ORDER_THETA_CAP}); routing gradient-only quasi-Newton"
3543 );
3544 }
3545 let mut suppress_outer_hessian_for_nfree = false;
3555
3556 log::trace!(
3557 "[{}] starting analytic optimization: rho_dim={}, coord_dim={}, dims_per_term={:?}",
3558 label,
3559 rho_dim,
3560 coord_dim,
3561 dims_per_term,
3562 );
3563
3564 let mut ctx = SpatialJointContext {
3565 data,
3566 rho_dim,
3567 kind,
3568 cache: SingleBlockExactJointDesignCache::new_with_policy(
3569 data,
3570 resolvedspec.clone(),
3571 baseline_design.clone(),
3572 spatial_terms.to_vec(),
3573 rho_dim,
3574 dims_per_term.to_vec(),
3575 &options.resource_policy,
3576 )
3577 .map_err(EstimationError::InvalidInput)?,
3578 evaluator: gam_solve::estimate::ExternalJointHyperEvaluator::new(
3579 y,
3580 weights,
3581 &baseline_design.design,
3582 offset,
3583 &baseline_design.penalties,
3584 &external_opts_for_design(&family, baseline_design, options),
3585 label,
3586 )?,
3587 frozen_glm_inputs: if coord_dim == 1 && frozen_glm_tensor_eligible_family(&family) {
3588 Some(SpatialFrozenGlmInputs {
3589 y: y.to_owned(),
3590 weights: weights.to_owned(),
3591 offset: offset.to_owned(),
3592 family: family.clone(),
3593 })
3594 } else {
3595 None
3596 },
3597 frozen_glm_psi_bounds: if coord_dim == 1 && frozen_glm_tensor_eligible_family(&family) {
3598 Some((lower[rho_dim], upper[rho_dim]))
3599 } else {
3600 None
3601 },
3602 frozen_glm_tensor: None,
3603 frozen_glm_tensor_attempted: false,
3604 frozen_glm_weight_memo: None,
3605 };
3606
3607 let mut psi_rank_stable_floor: Option<f64> = None;
3630 let mut psi_rank_stable_ceiling: Option<f64> = None;
3639 let nfree_penalty_capable =
3640 coord_dim == 1 && family.is_gaussian_identity() && ctx.cache.supports_nfree_penalty_rekey();
3641 if nfree_penalty_capable {
3642 let psi_lo = lower[rho_dim];
3643 let psi_hi = upper[rho_dim];
3644 let z = Array1::from_iter(y.iter().zip(offset.iter()).map(|(yi, oi)| yi - oi));
3645 let theta_probe_base = theta0.clone();
3646 let SpatialJointContext {
3649 cache, evaluator, ..
3650 } = &mut ctx;
3651 let attached = evaluator.build_and_set_psi_gram_tensor(
3652 |psi| {
3653 let mut theta_probe = theta_probe_base.clone();
3654 theta_probe[rho_dim] = psi;
3655 cache.ensure_theta(&theta_probe)?;
3656 Ok(cache.design().design.clone())
3657 },
3658 weights,
3659 z.view(),
3660 psi_lo,
3661 psi_hi,
3662 );
3663 if attached {
3664 log::info!(
3665 "[{label}] certified ψ-gram tensor over [{psi_lo:.3}, {psi_hi:.3}]: \
3666 in-window trials assemble Gaussian sufficient statistics n-free"
3667 );
3668 let psi_anchor = theta0[rho_dim];
3673 psi_rank_stable_floor = evaluator
3674 .psi_gram_rank_stable_floor(psi_anchor)
3675 .filter(|&f| f.is_finite() && f > psi_lo && f < psi_anchor);
3676 log::info!(
3677 "[KAPPA-PHASE-FLOOR] n_rows={} psi_lo={psi_lo:.6} psi_anchor={psi_anchor:.6} \
3678 rank_stable_floor={:?} lifted={}",
3679 data.nrows(),
3680 evaluator.psi_gram_rank_stable_floor(psi_anchor),
3681 psi_rank_stable_floor.is_some(),
3682 );
3683 if let Some(floor) = psi_rank_stable_floor {
3684 log::info!(
3685 "[{label}] rank-stable κ-floor ψ_floor={floor:.6} > window floor \
3686 ψ_lo={psi_lo:.6}: lifting the optimizer lower bound to keep every \
3687 in-window trial on the n-free design-realization skip (#1033). The \
3688 conditioned Gram is rank-deficient below ψ_floor (longest-length-scale \
3689 radial mode collapses into the nullspace), where the skip is soundly \
3690 refused; that band drifts with n via the sample-std standardization, \
3691 so this n-free k-space floor is the n-independent fix."
3692 );
3693 }
3694 psi_rank_stable_ceiling = evaluator
3703 .psi_gram_rank_stable_ceiling(psi_anchor)
3704 .filter(|&c| c.is_finite() && c < psi_hi && c > psi_anchor);
3705 log::info!(
3706 "[KAPPA-PHASE-CEIL] n_rows={} psi_hi={psi_hi:.6} psi_anchor={psi_anchor:.6} \
3707 rank_stable_ceiling={:?} clamped={}",
3708 data.nrows(),
3709 evaluator.psi_gram_rank_stable_ceiling(psi_anchor),
3710 psi_rank_stable_ceiling.is_some(),
3711 );
3712 if let Some(ceiling) = psi_rank_stable_ceiling {
3713 log::info!(
3714 "[{label}] rank-stable κ-ceiling ψ_ceil={ceiling:.6} < window ceiling \
3715 ψ_hi={psi_hi:.6}: clamping the optimizer upper bound to keep every \
3716 in-window trial on the n-free design-realization skip (#1033). The \
3717 conditioned Gram is rank-deficient above ψ_ceil (longest-frequency \
3718 radial mode goes collinear), where the skip is soundly refused; a \
3719 line-search overshoot there trips the O(n) reset_surface lane (and the \
3720 deficient pinning ψ it records resets the next in-band trial too)."
3721 );
3722 }
3723 let gradient_covers_full_window = evaluator.psi_gram_tensor_covers_gradient(psi_lo)
3724 && evaluator.psi_gram_tensor_covers_gradient(psi_hi);
3725 if gradient_covers_full_window {
3726 log::info!(
3727 "[{label}] certified ψ-gram tensor gradient lane covers the full \
3728 optimizer window [{psi_lo:.3}, {psi_hi:.3}]"
3729 );
3730 } else {
3731 log::info!(
3732 "[{label}] ψ-gram tensor value lane certified, but the gradient lane \
3733 does not cover the full optimizer window [{psi_lo:.3}, {psi_hi:.3}]; \
3734 keeping exact streamed kappa routing"
3735 );
3736 }
3737 evaluator.set_supports_nfree_penalty_rekey(true);
3757 log::info!(
3758 "[{label}] exact n-free ψ-penalty re-key enabled over [{psi_lo:.3}, \
3759 {psi_hi:.3}]: in-window fast-path trials rebuild S(ψ) n-free from frozen \
3760 geometry (no reset_surface)"
3761 );
3762 } else {
3763 log::info!(
3764 "[{label}] ψ-gram tensor did not certify over [{psi_lo:.3}, {psi_hi:.3}]; \
3765 keeping the exact per-trial path"
3766 );
3767 }
3768 if attached
3789 && evaluator.psi_gram_tensor_covers_gradient(psi_lo)
3790 && evaluator.psi_gram_tensor_covers_gradient(psi_hi)
3791 && evaluator.supports_nfree_penalty_rekey()
3792 && cache.supports_nfree_gradient_only_routing()
3793 {
3794 suppress_outer_hessian_for_nfree = true;
3795 prefer_gradient_only = true;
3796 log::info!(
3797 "[{label}] n-free Gaussian ψ-lane armed; suppressing the analytic outer \
3798 Hessian and routing gradient-only (BFGS) so the κ outer loop never realizes \
3799 the O(n) second-order slab — n-independent outer loop (#1033)"
3800 );
3801 }
3802 } else if coord_dim == 1 && family.is_gaussian_identity() {
3803 log::info!(
3804 "[{label}] exact n-free ψ-penalty re-key unavailable; skipping ψ-gram tensor \
3805 attachment so value, gradient, and Hessian remain on the same exact streamed \
3806 objective"
3807 );
3808 }
3809
3810 let kphase_prime_order =
3811 if analytic_outer_hessian_available && !suppress_outer_hessian_for_nfree {
3812 OuterEvalOrder::ValueGradientHessian
3813 } else {
3814 OuterEvalOrder::ValueAndGradient
3815 };
3816 let kphase_prime_start = std::time::Instant::now();
3817 drop(ctx.eval_full(theta0, kphase_prime_order, analytic_outer_hessian_available)?);
3818 log::info!(
3819 "[KAPPA-PHASE-PRIME] n_rows={} order={:?} elapsed_s={:.4} slow_path_resets_total={} design_revision={}",
3820 data.nrows(),
3821 kphase_prime_order,
3822 kphase_prime_start.elapsed().as_secs_f64(),
3823 ctx.evaluator.slow_path_reset_count(),
3824 ctx.cache.design_revision(),
3825 );
3826
3827 let kphase_cost_calls = std::cell::Cell::new(0usize);
3828 let kphase_eval_calls = std::cell::Cell::new(0usize);
3829 let kphase_efs_calls = std::cell::Cell::new(0usize);
3830 let kphase_cost_total_s = std::cell::Cell::new(0.0);
3831 let kphase_eval_total_s = std::cell::Cell::new(0.0);
3832 let kphase_efs_total_s = std::cell::Cell::new(0.0);
3833 let kphase_nfree_miss_shape = std::cell::Cell::new(0u64);
3834 let kphase_nfree_miss_value = std::cell::Cell::new(0u64);
3835 let kphase_nfree_miss_gradient = std::cell::Cell::new(0u64);
3836 let kphase_nfree_miss_penalty = std::cell::Cell::new(0u64);
3837 let kphase_nfree_miss_revision = std::cell::Cell::new(0u64);
3838 let kphase_nfree_miss_second_order = std::cell::Cell::new(0u64);
3839 let kphase_nfree_miss_other = std::cell::Cell::new(0u64);
3840 let kphase_optim_start = std::time::Instant::now();
3841 let kphase_log_kappa_dim = coord_dim;
3842 let kphase_slow_resets_start = ctx.evaluator.slow_path_reset_count();
3843 let kphase_design_revision_start = ctx.cache.design_revision();
3844 let kphase_nfree_skip_touches_start = gam_solve::pirls::nfree_skip_row_element_touches();
3848
3849 let lower_effective: std::borrow::Cow<'_, Array1<f64>> = match psi_rank_stable_floor {
3856 Some(floor) if coord_dim == 1 && floor > lower[rho_dim] => {
3857 let mut lifted = lower.clone();
3858 lifted[rho_dim] = floor;
3859 std::borrow::Cow::Owned(lifted)
3860 }
3861 _ => std::borrow::Cow::Borrowed(lower),
3862 };
3863 let lower = lower_effective.as_ref();
3864
3865 let upper_effective: std::borrow::Cow<'_, Array1<f64>> = match psi_rank_stable_ceiling {
3873 Some(ceiling) if coord_dim == 1 && ceiling < upper[rho_dim] => {
3874 let mut clamped = upper.clone();
3875 clamped[rho_dim] = ceiling;
3876 std::borrow::Cow::Owned(clamped)
3877 }
3878 _ => std::borrow::Cow::Borrowed(upper),
3879 };
3880 let upper = upper_effective.as_ref();
3881
3882 let problem = exact_joint_multistart_outer_problem(
3883 theta0,
3884 lower,
3885 upper,
3886 rho_dim,
3887 coord_dim,
3888 theta_dim,
3889 Derivative::Analytic,
3890 if analytic_outer_hessian_available && !suppress_outer_hessian_for_nfree {
3891 DeclaredHessianForm::Either
3892 } else {
3893 DeclaredHessianForm::Unavailable
3898 },
3899 prefer_gradient_only,
3900 suppress_outer_hessian_for_nfree,
3911 seed_risk_profile_for_likelihood_family(&family),
3912 kappa_options.rel_tol.max(1e-6),
3913 kappa_options.max_outer_iter.max(1),
3914 Some(5.0),
3918 Some(kappa_options.log_step.clamp(0.25, 1.0)),
3920 None,
3921 Some((data.nrows(), baseline_design.design.ncols())),
3926 !constant_curvature_term_indices(resolvedspec).is_empty(),
3930 kind == SpatialHyperKind::Isotropic
3935 && constant_curvature_term_indices(resolvedspec).is_empty()
3936 && spatial_terms.iter().any(|&term_idx| {
3937 matches!(
3938 resolvedspec
3939 .smooth_terms
3940 .get(term_idx)
3941 .map(|term| &term.basis),
3942 Some(SmoothBasisSpec::Matern { .. })
3943 )
3944 }),
3945 );
3946
3947 let eval_outer = |ctx: &mut &mut SpatialJointContext<'_>,
3948 theta: &Array1<f64>,
3949 order: OuterEvalOrder|
3950 -> Result<OuterEval, EstimationError> {
3951 let t0 = std::time::Instant::now();
3952 let allow_second_order_for_call = matches!(order, OuterEvalOrder::ValueGradientHessian)
3953 && analytic_outer_hessian_available;
3954 let gate = ctx.nfree_skip_gate_status(theta, allow_second_order_for_call, true);
3955 let resets_before = ctx.evaluator.slow_path_reset_count();
3956 let raw = ctx.eval_full(theta, order, analytic_outer_hessian_available);
3957 let reset_delta = ctx
3958 .evaluator
3959 .slow_path_reset_count()
3960 .saturating_sub(resets_before);
3961 if reset_delta > 0 {
3962 if !gate.shape {
3963 kphase_nfree_miss_shape.set(kphase_nfree_miss_shape.get() + reset_delta);
3964 }
3965 if gate.shape && !gate.value {
3966 kphase_nfree_miss_value.set(kphase_nfree_miss_value.get() + reset_delta);
3967 }
3968 if gate.shape && gate.value && !gate.gradient {
3969 kphase_nfree_miss_gradient.set(kphase_nfree_miss_gradient.get() + reset_delta);
3970 }
3971 if gate.shape && gate.value && gate.gradient && !gate.penalty {
3972 kphase_nfree_miss_penalty.set(kphase_nfree_miss_penalty.get() + reset_delta);
3973 }
3974 if gate.shape && gate.value && gate.gradient && gate.penalty && !gate.revision {
3975 kphase_nfree_miss_revision.set(kphase_nfree_miss_revision.get() + reset_delta);
3976 }
3977 if gate.shape
3978 && gate.value
3979 && gate.gradient
3980 && gate.penalty
3981 && gate.revision
3982 && gate.second_order
3983 {
3984 kphase_nfree_miss_second_order
3985 .set(kphase_nfree_miss_second_order.get() + reset_delta);
3986 }
3987 if gate.would_skip(true) {
3988 kphase_nfree_miss_other.set(kphase_nfree_miss_other.get() + reset_delta);
3989 }
3990 }
3991 let elapsed_s = t0.elapsed().as_secs_f64();
3992 kphase_eval_calls.set(kphase_eval_calls.get() + 1);
3993 kphase_eval_total_s.set(kphase_eval_total_s.get() + elapsed_s);
3994 let (theta_norm, log_kappa_norm) = kphase_log_norms(theta, rho_dim);
3995 log::info!(
3996 "[KAPPA-PHASE] phase=eval_outer call={} order={:?} design_revision={:?} theta_norm={:.4e} log_kappa_norm={:.4e} elapsed_s={:.4}",
3997 kphase_eval_calls.get(),
3998 order,
3999 Some(ctx.cache.design_revision()),
4000 theta_norm,
4001 log_kappa_norm,
4002 elapsed_s,
4003 );
4004 match raw {
4005 Ok((cost, grad, hess)) => Ok(OuterEval {
4006 cost,
4007 gradient: grad,
4008 hessian: hess,
4009 inner_beta_hint: None,
4010 }),
4011 Err(err) if is_recoverable_trial_point_error(&err) => {
4019 log::debug!(
4020 "[{label}] trial point infeasible (kernel design \
4021 not constructible at theta={theta:?}): {err}; retreating",
4022 );
4023 Ok(OuterEval::infeasible(theta_dim))
4024 }
4025 Err(err) => Err(err),
4026 }
4027 };
4028
4029 let mut obj = problem.build_objective_with_eval_order(
4030 &mut ctx,
4031 |ctx: &mut &mut SpatialJointContext<'_>, theta: &Array1<f64>| {
4032 let t0 = std::time::Instant::now();
4033 let gate = ctx.nfree_skip_gate_status(theta, false, false);
4034 let resets_before = ctx.evaluator.slow_path_reset_count();
4035 let cost = ctx.eval_cost(theta);
4036 let reset_delta = ctx
4037 .evaluator
4038 .slow_path_reset_count()
4039 .saturating_sub(resets_before);
4040 if reset_delta > 0 {
4041 if !gate.shape {
4042 kphase_nfree_miss_shape.set(kphase_nfree_miss_shape.get() + reset_delta);
4043 }
4044 if gate.shape && !gate.value {
4045 kphase_nfree_miss_value.set(kphase_nfree_miss_value.get() + reset_delta);
4046 }
4047 if gate.shape && gate.value && !gate.penalty {
4048 kphase_nfree_miss_penalty.set(kphase_nfree_miss_penalty.get() + reset_delta);
4049 }
4050 if gate.shape && gate.value && gate.penalty && !gate.revision {
4051 kphase_nfree_miss_revision.set(kphase_nfree_miss_revision.get() + reset_delta);
4052 }
4053 if gate.would_skip(false) {
4054 kphase_nfree_miss_other.set(kphase_nfree_miss_other.get() + reset_delta);
4055 }
4056 }
4057 let elapsed_s = t0.elapsed().as_secs_f64();
4058 kphase_cost_calls.set(kphase_cost_calls.get() + 1);
4059 kphase_cost_total_s.set(kphase_cost_total_s.get() + elapsed_s);
4060 let (theta_norm, log_kappa_norm) = kphase_log_norms(theta, rho_dim);
4061 log::info!(
4062 "[KAPPA-PHASE] phase=cost call={} design_revision={:?} theta_norm={:.4e} log_kappa_norm={:.4e} elapsed_s={:.4}",
4063 kphase_cost_calls.get(),
4064 Some(ctx.cache.design_revision()),
4065 theta_norm,
4066 log_kappa_norm,
4067 elapsed_s,
4068 );
4069 Ok(cost)
4070 },
4071 |ctx: &mut &mut SpatialJointContext<'_>, theta: &Array1<f64>| {
4072 eval_outer(
4073 ctx,
4074 theta,
4075 if analytic_outer_hessian_available && !suppress_outer_hessian_for_nfree {
4085 OuterEvalOrder::ValueGradientHessian
4086 } else {
4087 OuterEvalOrder::ValueAndGradient
4088 },
4089 )
4090 },
4091 |ctx: &mut &mut SpatialJointContext<'_>, theta: &Array1<f64>, order: OuterEvalOrder| {
4092 eval_outer(ctx, theta, order)
4093 },
4094 Some(|ctx: &mut &mut SpatialJointContext<'_>| {
4095 ctx.reset();
4096 }),
4097 Some(|ctx: &mut &mut SpatialJointContext<'_>, theta: &Array1<f64>| {
4098 let t0 = std::time::Instant::now();
4099 let eval = ctx.eval_efs(theta);
4100 let elapsed_s = t0.elapsed().as_secs_f64();
4101 kphase_efs_calls.set(kphase_efs_calls.get() + 1);
4102 kphase_efs_total_s.set(kphase_efs_total_s.get() + elapsed_s);
4103 let (theta_norm, log_kappa_norm) = kphase_log_norms(theta, rho_dim);
4104 log::info!(
4105 "[KAPPA-PHASE] phase=efs call={} design_revision={:?} theta_norm={:.4e} log_kappa_norm={:.4e} elapsed_s={:.4}",
4106 kphase_efs_calls.get(),
4107 Some(ctx.cache.design_revision()),
4108 theta_norm,
4109 log_kappa_norm,
4110 elapsed_s,
4111 );
4112 eval
4113 }),
4114 );
4115
4116 let run_label = match kind {
4117 SpatialHyperKind::Anisotropic => "aniso-psi joint REML",
4118 SpatialHyperKind::Isotropic => "iso-kappa joint REML",
4119 };
4120 let result = problem.run(&mut obj, run_label)?;
4121 if !result.converged {
4122 crate::bail_invalid_estim!(
4123 "{} did not converge after {} iterations (final_objective={:.6e}, final_grad_norm={})",
4124 run_label,
4125 result.iterations,
4126 result.final_value,
4127 result.final_grad_norm_report(),
4128 );
4129 }
4130 drop(obj);
4131 let kphase_total_s = kphase_optim_start.elapsed().as_secs_f64();
4132 let kphase_slow_resets = ctx
4133 .evaluator
4134 .slow_path_reset_count()
4135 .saturating_sub(kphase_slow_resets_start);
4136 let kphase_design_revision_delta = ctx
4137 .cache
4138 .design_revision()
4139 .saturating_sub(kphase_design_revision_start);
4140 let kphase_nfree_skip_touches = gam_solve::pirls::nfree_skip_row_element_touches()
4141 .saturating_sub(kphase_nfree_skip_touches_start);
4142 log::info!(
4143 "[KAPPA-PHASE-SUMMARY] n_rows={} log_kappa_dim={} n_cost={} cost_total_s={:.4} n_eval={} eval_total_s={:.4} n_efs={} efs_total_s={:.4} slow_path_resets={} design_revision_delta={} nfree_skip_row_touches={} nfree_miss_shape={} nfree_miss_value={} nfree_miss_gradient={} nfree_miss_penalty={} nfree_miss_revision={} nfree_miss_second_order={} nfree_miss_other={} optim_total_s={:.4}",
4144 data.nrows(),
4145 kphase_log_kappa_dim,
4146 kphase_cost_calls.get(),
4147 kphase_cost_total_s.get(),
4148 kphase_eval_calls.get(),
4149 kphase_eval_total_s.get(),
4150 kphase_efs_calls.get(),
4151 kphase_efs_total_s.get(),
4152 kphase_slow_resets,
4153 kphase_design_revision_delta,
4154 kphase_nfree_skip_touches,
4155 kphase_nfree_miss_shape.get(),
4156 kphase_nfree_miss_value.get(),
4157 kphase_nfree_miss_gradient.get(),
4158 kphase_nfree_miss_penalty.get(),
4159 kphase_nfree_miss_revision.get(),
4160 kphase_nfree_miss_second_order.get(),
4161 kphase_nfree_miss_other.get(),
4162 kphase_total_s,
4163 );
4164 let timing = SpatialLengthScaleOptimizationTiming {
4165 log_kappa_dim: kphase_log_kappa_dim,
4166 cost_calls: kphase_cost_calls.get(),
4167 cost_total_s: kphase_cost_total_s.get(),
4168 eval_calls: kphase_eval_calls.get(),
4169 eval_total_s: kphase_eval_total_s.get(),
4170 efs_calls: kphase_efs_calls.get(),
4171 efs_total_s: kphase_efs_total_s.get(),
4172 slow_path_resets: kphase_slow_resets,
4173 design_revision_delta: kphase_design_revision_delta,
4174 nfree_skip_row_touches: kphase_nfree_skip_touches,
4175 nfree_miss_shape: kphase_nfree_miss_shape.get(),
4176 nfree_miss_value: kphase_nfree_miss_value.get(),
4177 nfree_miss_gradient: kphase_nfree_miss_gradient.get(),
4178 nfree_miss_penalty: kphase_nfree_miss_penalty.get(),
4179 nfree_miss_revision: kphase_nfree_miss_revision.get(),
4180 nfree_miss_second_order: kphase_nfree_miss_second_order.get(),
4181 nfree_miss_other: kphase_nfree_miss_other.get(),
4182 optim_total_s: kphase_total_s,
4183 };
4184 log::trace!(
4185 "[{}] converged in {} iterations, final_value={:.6e}, grad_norm={}",
4186 label,
4187 result.iterations,
4188 result.final_value,
4189 result.final_grad_norm_report(),
4190 );
4191 let theta_star = result.rho;
4195 Ok((theta_star, result.final_value, timing))
4196}
4197
4198fn set_single_term_spatial_length_scale(
4202 term: &mut SmoothTermSpec,
4203 length_scale: f64,
4204) -> Result<(), EstimationError> {
4205 match &mut term.basis {
4206 SmoothBasisSpec::ThinPlate { spec, .. } => {
4207 spec.length_scale = length_scale;
4208 Ok(())
4209 }
4210 SmoothBasisSpec::Matern { spec, .. } => {
4211 spec.length_scale = length_scale;
4212 Ok(())
4213 }
4214 SmoothBasisSpec::Duchon { spec, .. } => {
4215 spec.length_scale = Some(length_scale);
4216 Ok(())
4217 }
4218 _ => Err(EstimationError::InvalidInput(format!(
4219 "term '{}' does not expose a spatial length scale",
4220 term.name
4221 ))),
4222 }
4223}
4224
4225fn set_single_term_spatial_aniso_log_scales(
4229 term: &mut SmoothTermSpec,
4230 eta: Vec<f64>,
4231) -> Result<(), EstimationError> {
4232 let eta = center_aniso_log_scales(&eta);
4233 match &mut term.basis {
4234 SmoothBasisSpec::Matern { spec, .. } => {
4235 spec.aniso_log_scales = Some(eta);
4236 Ok(())
4237 }
4238 SmoothBasisSpec::Duchon { spec, .. } => {
4239 spec.aniso_log_scales = Some(eta);
4240 Ok(())
4241 }
4242 _ => Err(EstimationError::InvalidInput(format!(
4243 "term '{}' does not support aniso_log_scales",
4244 term.name
4245 ))),
4246 }
4247}
4248
4249pub fn get_constant_curvature_kappa(spec: &TermCollectionSpec, term_idx: usize) -> Option<f64> {
4268 constant_curvature_term_spec(spec, term_idx).map(|cc| cc.kappa)
4269}
4270
4271pub fn constant_curvature_kappa_is_fixed(spec: &TermCollectionSpec, term_idx: usize) -> bool {
4278 constant_curvature_term_spec(spec, term_idx).is_some_and(|cc| cc.kappa_fixed)
4279}
4280
4281pub fn constant_curvature_term_indices(spec: &TermCollectionSpec) -> Vec<usize> {
4283 (0..spec.smooth_terms.len())
4284 .filter(|&idx| constant_curvature_term_spec(spec, idx).is_some())
4285 .collect()
4286}
4287
4288#[derive(Debug, Clone)]
4289struct SingleSmoothTermRealization {
4290 design_local: DesignMatrix,
4291 term: SmoothTerm,
4292 dropped_penaltyinfo: Vec<DroppedPenaltyBlockInfo>,
4293}
4294
4295impl SingleSmoothTermRealization {
4296 fn active_penaltyinfo(&self) -> Vec<PenaltyInfo> {
4297 self.term
4298 .penaltyinfo_local
4299 .iter()
4300 .filter(|info| info.active)
4301 .cloned()
4302 .collect()
4303 }
4304}
4305
4306fn build_single_smooth_term_realization_with_policy(
4307 data: ArrayView2<'_, f64>,
4308 termspec: &SmoothTermSpec,
4309 policy: &gam_runtime::resource::ResourcePolicy,
4310) -> Result<SingleSmoothTermRealization, BasisError> {
4311 let mut workspace = gam_terms::basis::BasisWorkspace::with_policy(policy.clone());
4312 let raw =
4313 build_smooth_design_withworkspace(data, std::slice::from_ref(termspec), &mut workspace)?;
4314 finish_single_smooth_term_realization(raw)
4315}
4316
4317fn finish_single_smooth_term_realization(
4318 raw: RawSmoothDesign,
4319) -> Result<SingleSmoothTermRealization, BasisError> {
4320 let RawSmoothDesign {
4321 term_designs,
4322 dropped_penaltyinfo,
4323 terms,
4324 ..
4325 } = raw;
4326 let term = terms.into_iter().next().ok_or_else(|| {
4327 BasisError::InvalidInput("single-term smooth build returned no term".to_string())
4328 })?;
4329 let design = term_designs.into_iter().next().ok_or_else(|| {
4330 BasisError::InvalidInput("single-term smooth build returned no term design".to_string())
4331 })?;
4332
4333 Ok(SingleSmoothTermRealization {
4334 design_local: design,
4335 term,
4336 dropped_penaltyinfo,
4337 })
4338}
4339
4340fn wrap_local_build_as_realization(
4347 mut local: LocalSmoothTermBuild,
4348 termspec: &SmoothTermSpec,
4349) -> Result<SingleSmoothTermRealization, String> {
4350 let p_local = local.dim;
4351 let lb_local = if local.box_reparam {
4352 shape_lower_bounds_local(termspec.shape, p_local)
4353 } else {
4354 None
4355 };
4356
4357 let active_count = local.penaltyinfo.iter().filter(|info| info.active).count();
4358 if active_count != local.penalties.len() {
4359 return Err(format!(
4360 "internal penalty info mismatch for term '{}': active_infos={}, penalties={}",
4361 termspec.name,
4362 active_count,
4363 local.penalties.len()
4364 ));
4365 }
4366
4367 let mut dropped_penaltyinfo = Vec::<DroppedPenaltyBlockInfo>::new();
4368 for info in local.penaltyinfo.iter().filter(|info| !info.active) {
4369 dropped_penaltyinfo.push(DroppedPenaltyBlockInfo {
4370 termname: Some(termspec.name.clone()),
4371 penalty: info.clone(),
4372 });
4373 }
4374 for info in &local.pre_dropped_penaltyinfo {
4375 dropped_penaltyinfo.push(DroppedPenaltyBlockInfo {
4376 termname: Some(termspec.name.clone()),
4377 penalty: info.clone(),
4378 });
4379 }
4380
4381 let applied_rotation: Option<gam_terms::basis::JointNullRotation> = match (
4385 local.joint_null_rotation.take(),
4386 lb_local.is_some(),
4387 local.linear_constraints.is_some(),
4388 ) {
4389 (Some(rot), false, false) => {
4390 let q = &rot.rotation;
4391 local.design =
4392 apply_smooth_transform_to_design(local.design.clone(), q, &termspec.name).map_err(
4393 |e| {
4394 format!(
4395 "joint-null absorption rotation failed for term '{}': {}",
4396 termspec.name, e
4397 )
4398 },
4399 )?;
4400 local.penalties = local
4401 .penalties
4402 .into_iter()
4403 .map(|s_local| {
4404 let qt_s = gam_linalg::faer_ndarray::fast_atb(q, &s_local);
4405 gam_linalg::faer_ndarray::fast_ab(&qt_s, q)
4406 })
4407 .collect();
4408 local.ops = vec![None; local.penalties.len()];
4409 local.kronecker_factored = None;
4410 Some(rot)
4411 }
4412 (Some(_), _, _) => None,
4413 (None, _, _) => None,
4414 };
4415
4416 let smooth_term = SmoothTerm {
4417 name: termspec.name.clone(),
4418 coeff_range: 0..p_local,
4419 shape: termspec.shape,
4420 penalties_local: local.penalties.clone(),
4421 nullspace_dims: local.nullspaces.clone(),
4422 penaltyinfo_local: local.penaltyinfo.clone(),
4423 metadata: local.metadata.clone(),
4424 lower_bounds_local: lb_local,
4425 linear_constraints_local: local.linear_constraints.clone(),
4426 kronecker_factored: local.kronecker_factored.take(),
4427 joint_null_rotation: applied_rotation,
4428 unabsorbed_global_orthogonality: None,
4431 };
4432
4433 Ok(SingleSmoothTermRealization {
4434 design_local: local.design,
4435 term: smooth_term,
4436 dropped_penaltyinfo,
4437 })
4438}
4439
4440fn freeze_geometry_from_metadata(
4451 termspec: &SmoothTermSpec,
4452 metadata: &BasisMetadata,
4453) -> Option<SmoothTermSpec> {
4454 let mut frozen = termspec.clone();
4455 match (&mut frozen.basis, metadata) {
4456 (
4457 SmoothBasisSpec::Matern {
4458 spec,
4459 input_scales: spec_scales,
4460 ..
4461 },
4462 BasisMetadata::Matern {
4463 centers,
4464 input_scales: meta_scales,
4465 identifiability_transform,
4466 nullspace_shrinkage_survived,
4467 ..
4468 },
4469 ) => {
4470 spec.center_strategy = CenterStrategy::UserProvided(centers.clone());
4471 if spec_scales.is_none()
4472 && let Some(s) = meta_scales.clone()
4473 {
4474 *spec_scales = Some(s);
4475 }
4476 if let Some(transform) = identifiability_transform.clone() {
4494 spec.identifiability = MaternIdentifiability::FrozenTransform {
4495 transform,
4496 nullspace_shrinkage_survived: Some(*nullspace_shrinkage_survived),
4497 };
4498 }
4499 Some(frozen)
4500 }
4501 (
4502 SmoothBasisSpec::Duchon {
4503 spec,
4504 input_scales: spec_scales,
4505 ..
4506 },
4507 BasisMetadata::Duchon {
4508 centers,
4509 input_scales: meta_scales,
4510 ..
4511 },
4512 ) => {
4513 spec.center_strategy = CenterStrategy::UserProvided(centers.clone());
4514 if spec_scales.is_none()
4515 && let Some(s) = meta_scales.clone()
4516 {
4517 *spec_scales = Some(s);
4518 }
4519 Some(frozen)
4520 }
4521 (
4522 SmoothBasisSpec::ThinPlate {
4523 spec,
4524 input_scales: spec_scales,
4525 ..
4526 },
4527 BasisMetadata::ThinPlate {
4528 centers,
4529 input_scales: meta_scales,
4530 ..
4531 },
4532 ) => {
4533 spec.center_strategy = CenterStrategy::UserProvided(centers.clone());
4534 if spec_scales.is_none()
4535 && let Some(s) = meta_scales.clone()
4536 {
4537 *spec_scales = Some(s);
4538 }
4539 Some(frozen)
4540 }
4541 _ => None,
4544 }
4545}
4546
4547fn rebuild_smooth_auxiliary_state(
4548 smooth: &mut SmoothDesign,
4549 dropped_penaltyinfo_by_term: &[Vec<DroppedPenaltyBlockInfo>],
4550) -> Result<(), String> {
4551 if dropped_penaltyinfo_by_term.len() != smooth.terms.len() {
4552 return Err(SmoothError::dimension_mismatch(format!(
4553 "smooth dropped-penalty cache mismatch: terms={}, dropped_sets={}",
4554 smooth.terms.len(),
4555 dropped_penaltyinfo_by_term.len()
4556 ))
4557 .into());
4558 }
4559
4560 let total_p = smooth.total_smooth_cols();
4561 let mut coefficient_lower_bounds = Array1::<f64>::from_elem(total_p, f64::NEG_INFINITY);
4562 let mut any_bounds = false;
4563 let mut linear_constraintrows: Vec<Array1<f64>> = Vec::new();
4564 let mut linear_constraint_b: Vec<f64> = Vec::new();
4565
4566 for term in &smooth.terms {
4567 let range = term.coeff_range.clone();
4568 if let Some(lb_local) = term.lower_bounds_local.as_ref() {
4569 if lb_local.len() != range.len() {
4570 return Err(SmoothError::dimension_mismatch(format!(
4571 "smooth lower-bound cache mismatch for term '{}': bounds={}, coeffs={}",
4572 term.name,
4573 lb_local.len(),
4574 range.len()
4575 ))
4576 .into());
4577 }
4578 coefficient_lower_bounds
4579 .slice_mut(s![range.clone()])
4580 .assign(lb_local);
4581 any_bounds = true;
4582 }
4583 if let Some(lin_local) = term.linear_constraints_local.as_ref() {
4584 if lin_local.a.ncols() != range.len() {
4585 return Err(SmoothError::dimension_mismatch(format!(
4586 "smooth linear-constraint cache mismatch for term '{}': cols={}, coeffs={}",
4587 term.name,
4588 lin_local.a.ncols(),
4589 range.len()
4590 ))
4591 .into());
4592 }
4593 for r in 0..lin_local.a.nrows() {
4594 let mut row = Array1::<f64>::zeros(total_p);
4595 row.slice_mut(s![range.clone()]).assign(&lin_local.a.row(r));
4596 linear_constraintrows.push(row);
4597 linear_constraint_b.push(lin_local.b[r]);
4598 }
4599 }
4600 }
4601
4602 smooth.coefficient_lower_bounds = if any_bounds {
4603 Some(coefficient_lower_bounds)
4604 } else {
4605 None
4606 };
4607 smooth.linear_constraints = if linear_constraintrows.is_empty() {
4608 None
4609 } else {
4610 let mut a = Array2::<f64>::zeros((linear_constraintrows.len(), total_p));
4611 for (i, row) in linear_constraintrows.iter().enumerate() {
4612 a.row_mut(i).assign(row);
4613 }
4614 Some(LinearInequalityConstraints {
4615 a,
4616 b: Array1::from_vec(linear_constraint_b),
4617 })
4618 };
4619 smooth.dropped_penaltyinfo = dropped_penaltyinfo_by_term
4620 .iter()
4621 .flat_map(|infos| infos.iter().cloned())
4622 .collect();
4623 Ok(())
4624}
4625
4626fn rebuild_term_collection_auxiliary_state(
4627 spec: &TermCollectionSpec,
4628 design: &mut TermCollectionDesign,
4629) -> Result<(), String> {
4630 if spec.linear_terms.len() != design.linear_ranges.len() {
4631 return Err(SmoothError::dimension_mismatch(format!(
4632 "term-collection linear bookkeeping mismatch: spec_terms={}, design_ranges={}",
4633 spec.linear_terms.len(),
4634 design.linear_ranges.len()
4635 ))
4636 .into());
4637 }
4638
4639 let p_total = design.design.ncols();
4640 let smooth_start = p_total.saturating_sub(design.smooth.total_smooth_cols());
4641 let mut coefficient_lower_bounds = Array1::<f64>::from_elem(p_total, f64::NEG_INFINITY);
4642 let mut any_bounds = false;
4643 let mut linear_constraintrows: Vec<Array1<f64>> = Vec::new();
4644 let mut linear_constraint_b: Vec<f64> = Vec::new();
4645
4646 for (linear, (_, range)) in spec.linear_terms.iter().zip(design.linear_ranges.iter()) {
4647 if range.len() != 1 {
4648 return Err(SmoothError::dimension_mismatch(format!(
4649 "linear term '{}' expected one coefficient column, found {}",
4650 linear.name,
4651 range.len()
4652 ))
4653 .into());
4654 }
4655 let col = range.start;
4656 if let Some(lb) = linear.coefficient_min {
4657 let mut row = Array1::<f64>::zeros(p_total);
4658 row[col] = 1.0;
4659 linear_constraintrows.push(row);
4660 linear_constraint_b.push(lb);
4661 }
4662 if let Some(ub) = linear.coefficient_max {
4663 let mut row = Array1::<f64>::zeros(p_total);
4664 row[col] = -1.0;
4665 linear_constraintrows.push(row);
4666 linear_constraint_b.push(-ub);
4667 }
4668 }
4669
4670 if let Some(lb_smooth) = design.smooth.coefficient_lower_bounds.as_ref() {
4671 if lb_smooth.len() != design.smooth.total_smooth_cols() {
4672 return Err(SmoothError::dimension_mismatch(format!(
4673 "smooth lower-bound width mismatch: bounds={}, smooth_cols={}",
4674 lb_smooth.len(),
4675 design.smooth.total_smooth_cols()
4676 ))
4677 .into());
4678 }
4679 coefficient_lower_bounds
4680 .slice_mut(s![
4681 smooth_start..(smooth_start + design.smooth.total_smooth_cols())
4682 ])
4683 .assign(lb_smooth);
4684 any_bounds = true;
4685 }
4686 if let Some(lin_smooth) = design.smooth.linear_constraints.as_ref() {
4687 if lin_smooth.a.ncols() != design.smooth.total_smooth_cols() {
4688 return Err(SmoothError::dimension_mismatch(format!(
4689 "smooth linear-constraint width mismatch: cols={}, smooth_cols={}",
4690 lin_smooth.a.ncols(),
4691 design.smooth.total_smooth_cols()
4692 ))
4693 .into());
4694 }
4695 let mut a_global = Array2::<f64>::zeros((lin_smooth.a.nrows(), p_total));
4696 a_global
4697 .slice_mut(s![
4698 ..,
4699 smooth_start..(smooth_start + design.smooth.total_smooth_cols())
4700 ])
4701 .assign(&lin_smooth.a);
4702 for r in 0..a_global.nrows() {
4703 linear_constraintrows.push(a_global.row(r).to_owned());
4704 linear_constraint_b.push(lin_smooth.b[r]);
4705 }
4706 }
4707
4708 let lower_bound_constraints = if any_bounds {
4709 linear_constraints_from_lower_bounds_global(&coefficient_lower_bounds)
4710 } else {
4711 None
4712 };
4713 let explicit_linear_constraints = if linear_constraintrows.is_empty() {
4714 None
4715 } else {
4716 let mut a = Array2::<f64>::zeros((linear_constraintrows.len(), p_total));
4717 for (i, row) in linear_constraintrows.iter().enumerate() {
4718 a.row_mut(i).assign(row);
4719 }
4720 Some(LinearInequalityConstraints {
4721 a,
4722 b: Array1::from_vec(linear_constraint_b),
4723 })
4724 };
4725
4726 design.coefficient_lower_bounds = if any_bounds {
4727 Some(coefficient_lower_bounds)
4728 } else {
4729 None
4730 };
4731 design.linear_constraints =
4732 merge_linear_constraints_global(explicit_linear_constraints, lower_bound_constraints)
4733 .map_err(|error| error.to_string())?;
4734 design.dropped_penaltyinfo = design.smooth.dropped_penaltyinfo.clone();
4735 Ok(())
4736}
4737
4738fn theta_values_match(left: &Array1<f64>, right: &Array1<f64>) -> bool {
4739 left.len() == right.len()
4740 && left
4741 .iter()
4742 .zip(right.iter())
4743 .all(|(&l, &r)| l.to_bits() == r.to_bits())
4744}
4745
4746fn latent_values_match(left: &Array1<f64>, right: &Array1<f64>) -> bool {
4747 theta_values_match(left, right)
4748}
4749
4750fn spatial_aniso_matches(left: Option<&[f64]>, right: Option<&[f64]>) -> bool {
4751 match (left, right) {
4752 (None, None) => true,
4753 (Some(a), Some(b)) => {
4754 a.len() == b.len()
4755 && a.iter()
4756 .zip(b.iter())
4757 .all(|(&x, &y)| x.to_bits() == y.to_bits())
4758 }
4759 _ => false,
4760 }
4761}
4762
4763fn spatial_length_scale_matches(left: Option<f64>, right: Option<f64>) -> bool {
4764 match (left, right) {
4765 (None, None) => true,
4766 (Some(a), Some(b)) => a.to_bits() == b.to_bits(),
4767 _ => false,
4768 }
4769}
4770
4771struct FrozenTermCollectionIncrementalRealizer<'d> {
4772 data: ArrayView2<'d, f64>,
4773 spec: TermCollectionSpec,
4774 design: TermCollectionDesign,
4775 fixed_blocks: Vec<DesignBlock>,
4776 dropped_penaltyinfo_by_term: Vec<Vec<DroppedPenaltyBlockInfo>>,
4777 smooth_penalty_ranges: Vec<Range<usize>>,
4778 full_penalty_ranges: Vec<Range<usize>>,
4779 basisworkspace: gam_terms::basis::BasisWorkspace,
4783 spatial_realization_geometry: Vec<Option<SmoothTermSpec>>,
4796 design_revision: u64,
4802}
4803
4804impl<'d> std::fmt::Debug for FrozenTermCollectionIncrementalRealizer<'d> {
4805 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
4806 f.debug_struct("FrozenTermCollectionIncrementalRealizer")
4807 .field("data_shape", &(self.data.nrows(), self.data.ncols()))
4808 .field("fixed_blocks", &self.fixed_blocks.len())
4809 .finish_non_exhaustive()
4810 }
4811}
4812
4813impl<'d> FrozenTermCollectionIncrementalRealizer<'d> {
4814 fn new(
4815 data: ArrayView2<'d, f64>,
4816 spec: TermCollectionSpec,
4817 design: TermCollectionDesign,
4818 ) -> Result<Self, String> {
4819 let policy = gam_runtime::resource::ResourcePolicy::default_library();
4820 Self::new_with_policy(data, spec, design, &policy)
4821 }
4822
4823 fn new_with_policy(
4824 data: ArrayView2<'d, f64>,
4825 spec: TermCollectionSpec,
4826 design: TermCollectionDesign,
4827 policy: &gam_runtime::resource::ResourcePolicy,
4828 ) -> Result<Self, String> {
4829 if spec.smooth_terms.len() != design.smooth.terms.len() {
4830 return Err(SmoothError::dimension_mismatch(format!(
4831 "incremental realizer smooth term mismatch: spec_terms={}, design_terms={}",
4832 spec.smooth_terms.len(),
4833 design.smooth.terms.len()
4834 ))
4835 .into());
4836 }
4837
4838 let mut smooth_cursor = 0usize;
4839 let mut smooth_penalty_ranges = Vec::with_capacity(design.smooth.terms.len());
4840 for term in &design.smooth.terms {
4841 let next = smooth_cursor + term.penalties_local.len();
4842 smooth_penalty_ranges.push(smooth_cursor..next);
4843 smooth_cursor = next;
4844 }
4845 if smooth_cursor != design.smooth.penalties.len() {
4846 return Err(SmoothError::dimension_mismatch(format!(
4847 "incremental realizer smooth penalty mismatch: ranged={}, actual={}",
4848 smooth_cursor,
4849 design.smooth.penalties.len()
4850 ))
4851 .into());
4852 }
4853
4854 let fixed_penalty_offset = design
4855 .penalties
4856 .len()
4857 .checked_sub(design.smooth.penalties.len())
4858 .ok_or_else(|| {
4859 "incremental realizer encountered invalid penalty bookkeeping".to_string()
4860 })?;
4861 let full_penalty_ranges = smooth_penalty_ranges
4862 .iter()
4863 .map(|range| (fixed_penalty_offset + range.start)..(fixed_penalty_offset + range.end))
4864 .collect::<Vec<_>>();
4865 let fixed_blocks = build_term_collection_fixed_blocks(data, &spec)
4866 .map_err(|e| format!("failed to cache fixed term-collection blocks: {e}"))?;
4867
4868 let mut dropped_penaltyinfo_by_term = Vec::with_capacity(spec.smooth_terms.len());
4869 for (term_idx, termspec) in spec.smooth_terms.iter().enumerate() {
4870 let realization = build_single_smooth_term_realization_with_policy(
4871 data, termspec, policy,
4872 )
4873 .map_err(|e| {
4874 format!(
4875 "failed to build cached realization for smooth term '{}' (index {}): {e}",
4876 termspec.name, term_idx
4877 )
4878 })?;
4879 let expected_cols = design.smooth.terms[term_idx].coeff_range.len();
4880 if realization.design_local.ncols() != expected_cols {
4881 return Err(SmoothError::dimension_mismatch(format!(
4882 "cached realization width mismatch for term '{}': cached_cols={}, design_cols={}",
4883 termspec.name,
4884 realization.design_local.ncols(),
4885 expected_cols
4886 ))
4887 .into());
4888 }
4889 if realization.active_penaltyinfo().len()
4890 != design.smooth.terms[term_idx].penalties_local.len()
4891 {
4892 return Err(SmoothError::dimension_mismatch(format!(
4893 "cached realization penalty mismatch for term '{}': cached_penalties={}, design_penalties={}",
4894 termspec.name,
4895 realization.active_penaltyinfo().len(),
4896 design.smooth.terms[term_idx].penalties_local.len()
4897 ))
4898 .into());
4899 }
4900 dropped_penaltyinfo_by_term.push(realization.dropped_penaltyinfo);
4901 }
4902
4903 let geometry_slots = spec.smooth_terms.len();
4904 Ok(Self {
4905 data,
4906 spec,
4907 design,
4908 fixed_blocks,
4909 dropped_penaltyinfo_by_term,
4910 smooth_penalty_ranges,
4911 full_penalty_ranges,
4912 basisworkspace: gam_terms::basis::BasisWorkspace::with_policy(policy.clone()),
4913 spatial_realization_geometry: vec![None; geometry_slots],
4914 design_revision: 0,
4915 })
4916 }
4917
4918 fn design_revision(&self) -> u64 {
4919 self.design_revision
4920 }
4921
4922 fn spec(&self) -> &TermCollectionSpec {
4923 &self.spec
4924 }
4925
4926 fn design(&self) -> &TermCollectionDesign {
4927 &self.design
4928 }
4929
4930 fn supports_nfree_penalty_rekey(&self, spatial_terms: &[usize]) -> bool {
4971 if spatial_terms.len() != 1 {
4972 return false;
4973 }
4974 let term_idx = spatial_terms[0];
4975 matches!(
4976 self.design.smooth.terms.get(term_idx).map(|t| &t.metadata),
4977 Some(BasisMetadata::Duchon { .. } | BasisMetadata::ThinPlate { .. })
4978 )
4979 }
4980
4981 fn supports_nfree_gradient_only_routing(&self, spatial_terms: &[usize]) -> bool {
4990 if spatial_terms.len() != 1 {
4991 return false;
4992 }
4993 let term_idx = spatial_terms[0];
4994 matches!(
4995 self.design.smooth.terms.get(term_idx).map(|t| &t.metadata),
4996 Some(BasisMetadata::Duchon { .. } | BasisMetadata::ThinPlate { .. })
4997 )
4998 }
4999
5000 fn canonical_penalties_at_psi(
5013 &mut self,
5014 spatial_terms: &[usize],
5015 psi: &[f64],
5016 ) -> Result<(Vec<gam_terms::construction::CanonicalPenalty>, Vec<usize>), String> {
5017 if spatial_terms.len() != 1 {
5018 return Err(format!(
5019 "n-free penalty re-key requires exactly one spatial term, found {}",
5020 spatial_terms.len()
5021 ));
5022 }
5023 let term_idx = spatial_terms[0];
5024 let (ls_opt, aniso_from_psi) = spatial_term_psi_to_length_scale_and_aniso(psi);
5030 let termspec =
5033 self.spec.smooth_terms.get(term_idx).ok_or_else(|| {
5034 format!("spatial term {term_idx} out of range for n-free penalty")
5035 })?;
5036 let term = self
5037 .design
5038 .smooth
5039 .terms
5040 .get(term_idx)
5041 .ok_or_else(|| format!("realized smooth term {term_idx} out of range"))?;
5042 let p_total = self.design.design.ncols();
5045 let (locals, nullspace_dims): (Vec<Array2<f64>>, Vec<usize>) = match &term.metadata {
5046 BasisMetadata::Duchon {
5047 centers,
5048 identifiability_transform,
5049 operator_collocation_points,
5050 power,
5051 nullspace_order,
5052 aniso_log_scales,
5053 input_scales,
5054 radial_reparam,
5055 ..
5056 } => {
5057 let operator_penalties = match &termspec.basis {
5058 SmoothBasisSpec::Duchon { spec, .. } => spec.operator_penalties.clone(),
5059 _ => gam_terms::basis::DuchonOperatorPenaltySpec::default(),
5060 };
5061 let effective_ls = match input_scales.as_deref() {
5068 Some(scales) => {
5069 compensate_optional_length_scale_for_standardization(ls_opt, scales)
5070 }
5071 None => ls_opt,
5072 };
5073 gam_terms::basis::duchon_penalties_at_length_scale(
5074 centers.view(),
5075 identifiability_transform.as_ref(),
5076 operator_collocation_points.as_ref().map(|p| p.view()),
5077 &operator_penalties,
5078 *power,
5079 *nullspace_order,
5080 aniso_log_scales.as_deref(),
5081 radial_reparam.as_ref(),
5082 effective_ls,
5083 &mut self.basisworkspace,
5084 )
5085 .map_err(|e| e.to_string())?
5086 }
5087 BasisMetadata::Matern {
5088 centers,
5089 periodic,
5090 nu,
5091 include_intercept,
5092 identifiability_transform,
5093 aniso_log_scales,
5094 input_scales,
5095 ..
5096 } => {
5097 let ls = ls_opt.ok_or_else(|| {
5104 "Matérn n-free penalty re-key requires a finite length-scale".to_string()
5105 })?;
5106 let effective_ls = match input_scales.as_deref() {
5107 Some(scales) => compensate_length_scale_for_standardization(ls, scales),
5108 None => ls,
5109 };
5110 let aniso_for_penalty = aniso_from_psi.as_deref().or(aniso_log_scales.as_deref());
5111 let (penalties, nullspace_dims, _info) =
5122 matern_operator_penalty_triplet_at_length_scale(
5123 centers.view(),
5124 periodic.as_deref(),
5125 identifiability_transform.as_ref(),
5126 *nu,
5127 *include_intercept,
5128 aniso_for_penalty,
5129 effective_ls,
5130 )
5131 .map_err(|e| e.to_string())?;
5132 (penalties, nullspace_dims)
5133 }
5134 BasisMetadata::ThinPlate {
5135 centers,
5136 identifiability_transform,
5137 radial_reparam,
5138 ..
5139 } => {
5140 let ls = ls_opt.ok_or_else(|| {
5141 "thin-plate n-free penalty re-key requires a finite length-scale".to_string()
5142 })?;
5143 let double_penalty = match &termspec.basis {
5144 SmoothBasisSpec::ThinPlate { spec, .. } => spec.double_penalty,
5145 _ => false,
5146 };
5147 gam_terms::basis::thin_plate_penalties_at_length_scale(
5148 centers.view(),
5149 identifiability_transform.as_ref(),
5150 radial_reparam.as_ref(),
5151 ls,
5152 double_penalty,
5153 &mut self.basisworkspace,
5154 )
5155 .map_err(|e| e.to_string())?
5156 }
5157 other => {
5158 return Err(format!(
5159 "n-free penalty re-key unsupported for basis metadata {:?}",
5160 std::mem::discriminant(other)
5161 ));
5162 }
5163 };
5164 let templates = &self.design.penalties;
5169 if templates.len() != locals.len() {
5170 return Err(format!(
5171 "n-free penalty re-key produced {} blocks but the frozen design carries {} \
5172 — penalty topology is not ψ-stable",
5173 locals.len(),
5174 templates.len()
5175 ));
5176 }
5177 let specs: Vec<gam_solve::estimate::PenaltySpec> = templates
5178 .iter()
5179 .zip(locals.into_iter())
5180 .map(|(tmpl, local)| gam_solve::estimate::PenaltySpec::Block {
5181 local,
5182 col_range: tmpl.col_range.clone(),
5183 prior_mean: tmpl.prior_mean.clone(),
5184 structure_hint: tmpl.structure_hint.clone(),
5185 op: tmpl.op.clone(),
5186 })
5187 .collect();
5188 gam_terms::construction::canonicalize_penalty_specs(
5189 &specs,
5190 &nullspace_dims,
5191 p_total,
5192 "nfree-psi-penalty",
5193 )
5194 .map_err(|e| e.to_string())
5195 }
5196
5197 fn canonical_penalty_derivatives_at_psi(
5198 &mut self,
5199 spatial_terms: &[usize],
5200 psi: &[f64],
5201 ) -> Result<(Range<usize>, usize, Vec<Array2<f64>>), String> {
5202 if spatial_terms.len() != 1 {
5203 return Err(format!(
5204 "n-free penalty derivative re-key requires exactly one spatial term, found {}",
5205 spatial_terms.len()
5206 ));
5207 }
5208 let term_idx = spatial_terms[0];
5209 let (ls_opt, aniso_from_psi) = spatial_term_psi_to_length_scale_and_aniso(psi);
5210 let termspec = self.spec.smooth_terms.get(term_idx).ok_or_else(|| {
5211 format!("spatial term {term_idx} out of range for n-free penalty derivative")
5212 })?;
5213 let term = self
5214 .design
5215 .smooth
5216 .terms
5217 .get(term_idx)
5218 .ok_or_else(|| format!("realized smooth term {term_idx} out of range"))?;
5219 let p_total = self.design.design.ncols();
5220 let smooth_start = p_total.saturating_sub(self.design.smooth.total_smooth_cols());
5221 let global_range =
5222 (smooth_start + term.coeff_range.start)..(smooth_start + term.coeff_range.end);
5223
5224 let locals = match &term.metadata {
5225 BasisMetadata::Duchon {
5226 centers,
5227 identifiability_transform,
5228 operator_collocation_points,
5229 power,
5230 nullspace_order,
5231 aniso_log_scales,
5232 input_scales,
5233 radial_reparam,
5234 ..
5235 } => {
5236 let mut spec = match &termspec.basis {
5237 SmoothBasisSpec::Duchon { spec, .. } => spec.clone(),
5238 _ => {
5239 return Err(
5240 "Duchon n-free penalty derivative requires a Duchon term spec"
5241 .to_string(),
5242 );
5243 }
5244 };
5245 let effective_ls = match input_scales.as_deref() {
5246 Some(scales) => {
5247 compensate_optional_length_scale_for_standardization(ls_opt, scales)
5248 }
5249 None => ls_opt,
5250 };
5251 spec.length_scale = effective_ls;
5252 spec.power = *power;
5253 spec.nullspace_order = *nullspace_order;
5254 spec.aniso_log_scales = aniso_log_scales.clone();
5255 spec.radial_reparam = radial_reparam.clone();
5258 if spec.length_scale.is_none() {
5259 return Err(
5260 "Duchon n-free penalty derivative requires a hybrid length-scale"
5261 .to_string(),
5262 );
5263 }
5264 let collocation = operator_collocation_points
5265 .as_ref()
5266 .map(|points| points.view())
5267 .unwrap_or_else(|| centers.view());
5268 let (_native_sources, mut first, _native_second) =
5269 gam_terms::basis::build_duchon_native_penalty_psi_derivatives(
5270 centers.view(),
5271 &spec,
5272 identifiability_transform.as_ref(),
5273 &mut self.basisworkspace,
5274 )
5275 .map_err(|e| e.to_string())?;
5276 let (_operator_sources, operator_first, _operator_second) =
5277 gam_terms::basis::build_duchon_operator_penalty_psi_derivatives(
5278 collocation,
5279 centers.view(),
5280 &spec,
5281 identifiability_transform.as_ref(),
5282 &mut self.basisworkspace,
5283 )
5284 .map_err(|e| e.to_string())?;
5285 first.extend(operator_first);
5286 first
5287 }
5288 BasisMetadata::Matern {
5289 centers,
5290 periodic,
5291 nu,
5292 include_intercept,
5293 identifiability_transform,
5294 aniso_log_scales,
5295 input_scales,
5296 ..
5297 } => {
5298 let ls = ls_opt.ok_or_else(|| {
5299 "Matérn n-free penalty derivative requires a finite length-scale".to_string()
5300 })?;
5301 let effective_ls = match input_scales.as_deref() {
5302 Some(scales) => compensate_length_scale_for_standardization(ls, scales),
5303 None => ls,
5304 };
5305 let penalty_centers = gam_terms::basis::expand_periodic_centers(
5306 ¢ers.to_owned(),
5307 periodic.as_deref(),
5308 )
5309 .map_err(|e| e.to_string())?;
5310 let aniso_for_penalty = aniso_from_psi.as_deref().or(aniso_log_scales.as_deref());
5311 let (first, _second) =
5312 gam_terms::basis::build_matern_operator_penalty_psi_derivatives(
5313 penalty_centers.view(),
5314 effective_ls,
5315 *nu,
5316 *include_intercept,
5317 identifiability_transform.as_ref(),
5318 aniso_for_penalty,
5319 )
5320 .map_err(|e| e.to_string())?;
5321 first
5322 }
5323 BasisMetadata::ThinPlate {
5324 centers,
5325 identifiability_transform,
5326 radial_reparam,
5327 ..
5328 } => {
5329 let ls = ls_opt.ok_or_else(|| {
5330 "thin-plate n-free penalty derivative requires a finite length-scale"
5331 .to_string()
5332 })?;
5333 let mut spec = match &termspec.basis {
5334 SmoothBasisSpec::ThinPlate { spec, .. } => spec.clone(),
5335 _ => {
5336 return Err(
5337 "thin-plate n-free penalty derivative requires a ThinPlate term spec"
5338 .to_string(),
5339 );
5340 }
5341 };
5342 spec.length_scale = ls;
5343 if spec.radial_reparam.is_none() {
5344 spec.radial_reparam = radial_reparam.clone();
5345 }
5346 let (primary, _primary_second) =
5347 gam_terms::basis::build_thin_plate_penalty_psi_derivativeswithworkspace(
5348 centers.view(),
5349 &spec,
5350 identifiability_transform.as_ref(),
5351 &mut self.basisworkspace,
5352 )
5353 .map_err(|e| e.to_string())?;
5354 if self.design.penalties.len() > 1 {
5355 vec![primary.clone(), Array2::<f64>::zeros(primary.raw_dim())]
5356 } else {
5357 vec![primary]
5358 }
5359 }
5360 other => {
5361 return Err(format!(
5362 "n-free penalty derivative re-key unsupported for basis metadata {:?}",
5363 std::mem::discriminant(other)
5364 ));
5365 }
5366 };
5367 if locals.len() != self.design.penalties.len() {
5368 return Err(format!(
5369 "n-free penalty derivative re-key produced {} blocks but the frozen design carries {} \
5370 — penalty topology is not ψ-stable",
5371 locals.len(),
5372 self.design.penalties.len()
5373 ));
5374 }
5375 Ok((global_range, p_total, locals))
5376 }
5377
5378 fn apply_log_kappa(
5379 &mut self,
5380 log_kappa: &SpatialLogKappaCoords,
5381 term_indices: &[usize],
5382 ) -> Result<(), String> {
5383 if term_indices.len() != log_kappa.dims_per_term().len() {
5384 return Err(SmoothError::dimension_mismatch(format!(
5385 "incremental realizer log-kappa term mismatch: term_indices={}, dims_per_term={}",
5386 term_indices.len(),
5387 log_kappa.dims_per_term().len()
5388 ))
5389 .into());
5390 }
5391
5392 let mut any_changed = false;
5393 for (slot, &term_idx) in term_indices.iter().enumerate() {
5394 any_changed |= self.apply_log_kappa_to_term(term_idx, log_kappa.term_slice(slot))?;
5395 }
5396
5397 if any_changed {
5398 self.refresh_full_design_operator()?;
5399 rebuild_smooth_auxiliary_state(
5400 &mut self.design.smooth,
5401 &self.dropped_penaltyinfo_by_term,
5402 )?;
5403 rebuild_term_collection_auxiliary_state(&self.spec, &mut self.design)?;
5404 self.design_revision = self.design_revision.wrapping_add(1);
5405 }
5406 Ok(())
5407 }
5408
5409 fn apply_log_kappa_to_term(&mut self, term_idx: usize, psi: &[f64]) -> Result<bool, String> {
5410 if !spatial_term_supports_hyper_optimization(&self.spec, term_idx) {
5411 return Err(SmoothError::invalid_config(format!(
5412 "incremental realizer term {term_idx} does not expose spatial hyperparameters"
5413 ))
5414 .into());
5415 }
5416 let measure_jet_term = measure_jet_term_spec(&self.spec, term_idx).is_some();
5420 let constant_curvature_term = constant_curvature_term_spec(&self.spec, term_idx).is_some();
5424 let mut next_length_scale = None;
5425 let mut next_aniso: Option<Vec<f64>> = None;
5426 if measure_jet_term {
5427 if !set_measure_jet_psi_dials(&mut self.spec, term_idx, psi)
5428 .map_err(|e| e.to_string())?
5429 {
5430 return Ok(false);
5431 }
5432 } else if constant_curvature_term {
5433 if !set_constant_curvature_kappa(&mut self.spec, term_idx, psi)
5434 .map_err(|e| e.to_string())?
5435 {
5436 return Ok(false);
5437 }
5438 } else {
5439 let current_length_scale = get_spatial_length_scale(&self.spec, term_idx);
5440 let current_aniso = get_spatial_aniso_log_scales(&self.spec, term_idx);
5441 let (ls, eta) = spatial_term_psi_to_length_scale_and_aniso(psi);
5442 next_length_scale = ls;
5443 next_aniso = eta;
5444 let same_length = spatial_length_scale_matches(current_length_scale, next_length_scale);
5445 let same_aniso = spatial_aniso_matches(current_aniso.as_deref(), next_aniso.as_deref());
5446 if same_length && same_aniso {
5447 return Ok(false);
5448 }
5449 if let Some(length_scale) = next_length_scale {
5450 set_spatial_length_scale(&mut self.spec, term_idx, length_scale)
5451 .map_err(|e| e.to_string())?;
5452 }
5453 if let Some(eta) = next_aniso.clone() {
5454 set_spatial_aniso_log_scales(&mut self.spec, term_idx, eta)
5455 .map_err(|e| e.to_string())?;
5456 }
5457 }
5458
5459 let geometry_slot = self
5470 .spatial_realization_geometry
5471 .get(term_idx)
5472 .ok_or_else(|| format!("incremental realizer geometry slot {term_idx} out of range"))?;
5473 let mut build_spec = match geometry_slot {
5474 Some(cached) => cached.clone(),
5475 None => self
5476 .spec
5477 .smooth_terms
5478 .get(term_idx)
5479 .ok_or_else(|| format!("incremental realizer smooth term {term_idx} out of range"))?
5480 .clone(),
5481 };
5482 if measure_jet_term {
5483 set_single_term_measure_jet_psi_dials(&mut build_spec, psi)
5487 .map_err(|e| e.to_string())?;
5488 } else if constant_curvature_term {
5489 set_single_term_constant_curvature_kappa(&mut build_spec, psi)
5494 .map_err(|e| e.to_string())?;
5495 } else {
5496 if let Some(length_scale) = next_length_scale {
5497 set_single_term_spatial_length_scale(&mut build_spec, length_scale)
5498 .map_err(|e| e.to_string())?;
5499 }
5500 if let Some(eta) = next_aniso {
5501 set_single_term_spatial_aniso_log_scales(&mut build_spec, eta)
5502 .map_err(|e| e.to_string())?;
5503 }
5504 }
5505
5506 let termname = build_spec.name.clone();
5507 let local = build_single_local_smooth_term(
5508 self.data,
5509 &build_spec,
5510 &mut self.basisworkspace,
5511 )
5512 .map_err(|e| {
5513 format!(
5514 "failed to rebuild smooth term '{termname}' during incremental κ realization: {e}"
5515 )
5516 })?;
5517
5518 if self.spatial_realization_geometry[term_idx].is_none()
5523 && let Some(frozen) = freeze_geometry_from_metadata(&build_spec, &local.metadata)
5524 {
5525 if let (
5537 SmoothBasisSpec::Matern {
5538 spec: frozen_spec, ..
5539 },
5540 Some(SmoothBasisSpec::Matern {
5541 spec: live_spec, ..
5542 }),
5543 ) = (
5544 &frozen.basis,
5545 self.spec
5546 .smooth_terms
5547 .get_mut(term_idx)
5548 .map(|t| &mut t.basis),
5549 ) {
5550 live_spec.identifiability = frozen_spec.identifiability.clone();
5551 live_spec.center_strategy = frozen_spec.center_strategy.clone();
5552 }
5553 self.spatial_realization_geometry[term_idx] = Some(frozen);
5554 }
5555
5556 let realization = wrap_local_build_as_realization(local, &build_spec)?;
5557 self.replace_term_realization(term_idx, realization)?;
5558 Ok(true)
5559 }
5560
5561 fn replace_term_realization(
5562 &mut self,
5563 term_idx: usize,
5564 realization: SingleSmoothTermRealization,
5565 ) -> Result<(), String> {
5566 let t_replace = std::time::Instant::now();
5567 let SingleSmoothTermRealization {
5568 design_local,
5569 term,
5570 dropped_penaltyinfo,
5571 } = realization;
5572 let SmoothTerm {
5573 name,
5574 penalties_local,
5575 nullspace_dims,
5576 penaltyinfo_local,
5577 metadata,
5578 lower_bounds_local,
5579 linear_constraints_local,
5580 joint_null_rotation,
5581 ..
5582 } = term;
5583 let coeff_range = self
5584 .design
5585 .smooth
5586 .terms
5587 .get(term_idx)
5588 .ok_or_else(|| format!("incremental realizer smooth term {term_idx} out of range"))?
5589 .coeff_range
5590 .clone();
5591 if design_local.ncols() != coeff_range.len() {
5592 return Err(SmoothError::dimension_mismatch(format!(
5593 "incremental realizer width mismatch for term {}: rebuilt_cols={}, cached_cols={}",
5594 term_idx,
5595 design_local.ncols(),
5596 coeff_range.len()
5597 ))
5598 .into());
5599 }
5600 if design_local.nrows() != self.design.design.nrows() {
5601 return Err(SmoothError::dimension_mismatch(format!(
5602 "incremental realizer row mismatch for term {}: rebuilt_rows={}, design_rows={}",
5603 term_idx,
5604 design_local.nrows(),
5605 self.design.design.nrows()
5606 ))
5607 .into());
5608 }
5609
5610 let active_penaltyinfo = penaltyinfo_local
5611 .iter()
5612 .filter(|info| info.active)
5613 .cloned()
5614 .collect::<Vec<_>>();
5615 let smooth_penalty_range = self
5616 .smooth_penalty_ranges
5617 .get(term_idx)
5618 .ok_or_else(|| {
5619 format!("incremental realizer missing smooth penalty range for term {term_idx}")
5620 })?
5621 .clone();
5622 let full_penalty_range = self
5623 .full_penalty_ranges
5624 .get(term_idx)
5625 .ok_or_else(|| {
5626 format!("incremental realizer missing full penalty range for term {term_idx}")
5627 })?
5628 .clone();
5629 if active_penaltyinfo.len() != smooth_penalty_range.len()
5630 || penalties_local.len() != smooth_penalty_range.len()
5631 || nullspace_dims.len() != smooth_penalty_range.len()
5632 {
5633 return Err(SmoothError::dimension_mismatch(format!(
5634 "incremental realizer topology changed for term '{}': penalties={}, infos={}, nullspaces={}, cached_penalties={}",
5635 name,
5636 penalties_local.len(),
5637 active_penaltyinfo.len(),
5638 nullspace_dims.len(),
5639 smooth_penalty_range.len()
5640 ))
5641 .into());
5642 }
5643
5644 self.design.smooth.term_designs[term_idx] = design_local;
5645
5646 for (offset, penalty_local) in penalties_local.iter().enumerate() {
5647 let smooth_penalty_idx = smooth_penalty_range.start + offset;
5648 let full_penalty_idx = full_penalty_range.start + offset;
5649 let nullspace_dim = nullspace_dims[offset];
5650 let penalty_info = active_penaltyinfo[offset].clone();
5651
5652 if penalty_local.nrows() != coeff_range.len()
5653 || penalty_local.ncols() != coeff_range.len()
5654 {
5655 return Err(SmoothError::dimension_mismatch(format!(
5656 "incremental realizer penalty shape mismatch for term '{}' penalty {}: \
5657 penalty is {}x{} but coeff_range has {} columns",
5658 name,
5659 offset,
5660 penalty_local.nrows(),
5661 penalty_local.ncols(),
5662 coeff_range.len()
5663 ))
5664 .into());
5665 }
5666
5667 let smooth_penalty = self
5668 .design
5669 .smooth
5670 .penalties
5671 .get_mut(smooth_penalty_idx)
5672 .ok_or_else(|| {
5673 format!(
5674 "incremental realizer smooth penalty {} out of range for term {}",
5675 smooth_penalty_idx, term_idx
5676 )
5677 })?;
5678 smooth_penalty.local.assign(penalty_local);
5681
5682 let full_bp = self
5683 .design
5684 .penalties
5685 .get_mut(full_penalty_idx)
5686 .ok_or_else(|| {
5687 format!(
5688 "incremental realizer full penalty {} out of range for term {}",
5689 full_penalty_idx, term_idx
5690 )
5691 })?;
5692 full_bp.local.assign(penalty_local);
5695
5696 self.design.smooth.nullspace_dims[smooth_penalty_idx] = nullspace_dim;
5697 self.design.nullspace_dims[full_penalty_idx] = nullspace_dim;
5698
5699 self.design.smooth.penaltyinfo[smooth_penalty_idx].global_index = smooth_penalty_idx;
5700 self.design.smooth.penaltyinfo[smooth_penalty_idx].termname = Some(name.clone());
5701 self.design.smooth.penaltyinfo[smooth_penalty_idx].penalty = penalty_info.clone();
5702
5703 self.design.penaltyinfo[full_penalty_idx].global_index = full_penalty_idx;
5704 self.design.penaltyinfo[full_penalty_idx].termname = Some(name.clone());
5705 self.design.penaltyinfo[full_penalty_idx].penalty = penalty_info;
5706 }
5707
5708 let target_term = self.design.smooth.terms.get_mut(term_idx).ok_or_else(|| {
5709 format!("incremental realizer smooth term {term_idx} disappeared during replacement")
5710 })?;
5711 target_term.penalties_local = penalties_local;
5712 target_term.nullspace_dims = nullspace_dims;
5713 target_term.penaltyinfo_local = penaltyinfo_local;
5714 target_term.metadata = metadata;
5715 target_term.lower_bounds_local = lower_bounds_local;
5716 target_term.linear_constraints_local = linear_constraints_local;
5717 target_term.joint_null_rotation = joint_null_rotation;
5718 self.dropped_penaltyinfo_by_term[term_idx] = dropped_penaltyinfo;
5719 log::info!(
5720 "[STAGE] smooth basis rebuild (term {}, '{}', cols={}): {:.3}s",
5721 term_idx,
5722 target_term.name,
5723 coeff_range.len(),
5724 t_replace.elapsed().as_secs_f64(),
5725 );
5726 Ok(())
5727 }
5728
5729 fn refresh_full_design_operator(&mut self) -> Result<(), String> {
5730 let mut blocks = Vec::<DesignBlock>::with_capacity(
5731 self.fixed_blocks.len() + self.design.smooth.term_designs.len(),
5732 );
5733 blocks.extend(self.fixed_blocks.iter().cloned());
5734 for term_design in &self.design.smooth.term_designs {
5735 blocks.push(DesignBlock::from(term_design));
5736 }
5737 self.design.design = assemble_term_collection_design_matrix(blocks)
5738 .map_err(|e| format!("failed to refresh term-collection design: {e}"))?;
5739 Ok(())
5740 }
5741}
5742
5743fn build_term_collection_fixed_blocks(
5744 data: ArrayView2<'_, f64>,
5745 spec: &TermCollectionSpec,
5746) -> Result<Vec<DesignBlock>, BasisError> {
5747 let mut blocks = Vec::<DesignBlock>::new();
5748 if !term_collection_has_one_sided_anchored_bspline(spec) {
5749 blocks.push(DesignBlock::Intercept(data.nrows()));
5750 }
5751
5752 if !spec.linear_terms.is_empty() {
5753 let mut linear_block = Array2::<f64>::zeros((data.nrows(), spec.linear_terms.len()));
5754 for (j, linear) in spec.linear_terms.iter().enumerate() {
5755 let column = linear
5759 .realized_design_column(data)
5760 .map_err(BasisError::InvalidInput)?;
5761 linear_block.column_mut(j).assign(&column);
5762 }
5763 blocks.push(DesignBlock::Dense(
5764 gam_linalg::matrix::DenseDesignMatrix::from(linear_block),
5765 ));
5766 }
5767
5768 for term in &spec.random_effect_terms {
5769 let block = build_random_effect_block(data, term)?;
5770 let re_op = RandomEffectOperator::new(block.group_ids, block.num_groups);
5771 blocks.push(DesignBlock::RandomEffect(Arc::new(re_op)));
5772 }
5773
5774 Ok(blocks)
5775}
5776
5777pub struct SpatialLengthScaleOptimizationResult<FitOut> {
5782 pub resolved_specs: Vec<TermCollectionSpec>,
5783 pub designs: Vec<TermCollectionDesign>,
5784 pub fit: FitOut,
5785 pub timing: Option<SpatialLengthScaleOptimizationTiming>,
5786}
5787
5788#[derive(Debug, Clone)]
5790pub struct ExactJointHyperSetup {
5791 rho0: Array1<f64>,
5792 rho_lower: Array1<f64>,
5793 rho_upper: Array1<f64>,
5794 log_kappa0: SpatialLogKappaCoords,
5795 log_kappa_lower: SpatialLogKappaCoords,
5796 log_kappa_upper: SpatialLogKappaCoords,
5797 auxiliary0: Array1<f64>,
5798 auxiliary_lower: Array1<f64>,
5799 auxiliary_upper: Array1<f64>,
5800}
5801
5802impl ExactJointHyperSetup {
5803 fn sanitize_rho_seed(
5804 rho0: Array1<f64>,
5805 rho_lower: &Array1<f64>,
5806 rho_upper: &Array1<f64>,
5807 ) -> Array1<f64> {
5808 Array1::from_iter(rho0.iter().enumerate().map(|(idx, &value)| {
5809 let lo = rho_lower[idx];
5810 let hi = rho_upper[idx];
5811 let fallback = 0.0_f64.clamp(lo, hi);
5812 if value.is_finite() {
5813 value.clamp(lo, hi)
5814 } else {
5815 fallback
5816 }
5817 }))
5818 }
5819
5820 pub(crate) fn new(
5821 rho0: Array1<f64>,
5822 rho_lower: Array1<f64>,
5823 rho_upper: Array1<f64>,
5824 log_kappa0: SpatialLogKappaCoords,
5825 log_kappa_lower: SpatialLogKappaCoords,
5826 log_kappa_upper: SpatialLogKappaCoords,
5827 ) -> Self {
5828 let rho0 = Self::sanitize_rho_seed(rho0, &rho_lower, &rho_upper);
5829 Self {
5830 rho0,
5831 rho_lower,
5832 rho_upper,
5833 log_kappa0,
5834 log_kappa_lower,
5835 log_kappa_upper,
5836 auxiliary0: Array1::zeros(0),
5837 auxiliary_lower: Array1::zeros(0),
5838 auxiliary_upper: Array1::zeros(0),
5839 }
5840 }
5841
5842 pub(crate) fn with_auxiliary(
5843 mut self,
5844 auxiliary0: Array1<f64>,
5845 auxiliary_lower: Array1<f64>,
5846 auxiliary_upper: Array1<f64>,
5847 ) -> Self {
5848 assert_eq!(
5849 auxiliary0.len(),
5850 auxiliary_lower.len(),
5851 "auxiliary lower bound length mismatch"
5852 );
5853 assert_eq!(
5854 auxiliary0.len(),
5855 auxiliary_upper.len(),
5856 "auxiliary upper bound length mismatch"
5857 );
5858 self.auxiliary0 = Self::sanitize_rho_seed(auxiliary0, &auxiliary_lower, &auxiliary_upper);
5859 self.auxiliary_lower = auxiliary_lower;
5860 self.auxiliary_upper = auxiliary_upper;
5861 self
5862 }
5863
5864 pub(crate) fn rho_dim(&self) -> usize {
5865 self.rho0.len()
5866 }
5867
5868 pub(crate) fn log_kappa_dim(&self) -> usize {
5869 self.log_kappa0.len()
5870 }
5871
5872 pub(crate) fn auxiliary_dim(&self) -> usize {
5873 self.auxiliary0.len()
5874 }
5875
5876 pub(crate) fn theta0(&self) -> Array1<f64> {
5877 let mut out =
5878 Array1::<f64>::zeros(self.rho_dim() + self.log_kappa_dim() + self.auxiliary_dim());
5879 out.slice_mut(s![..self.rho_dim()]).assign(&self.rho0);
5880 out.slice_mut(s![self.rho_dim()..self.rho_dim() + self.log_kappa_dim()])
5881 .assign(self.log_kappa0.as_array());
5882 out.slice_mut(s![self.rho_dim() + self.log_kappa_dim()..])
5883 .assign(&self.auxiliary0);
5884 out
5885 }
5886
5887 pub(crate) fn lower(&self) -> Array1<f64> {
5888 let mut out =
5889 Array1::<f64>::zeros(self.rho_dim() + self.log_kappa_dim() + self.auxiliary_dim());
5890 out.slice_mut(s![..self.rho_dim()]).assign(&self.rho_lower);
5891 out.slice_mut(s![self.rho_dim()..self.rho_dim() + self.log_kappa_dim()])
5892 .assign(self.log_kappa_lower.as_array());
5893 out.slice_mut(s![self.rho_dim() + self.log_kappa_dim()..])
5894 .assign(&self.auxiliary_lower);
5895 out
5896 }
5897
5898 pub(crate) fn upper(&self) -> Array1<f64> {
5899 let mut out =
5900 Array1::<f64>::zeros(self.rho_dim() + self.log_kappa_dim() + self.auxiliary_dim());
5901 out.slice_mut(s![..self.rho_dim()]).assign(&self.rho_upper);
5902 out.slice_mut(s![self.rho_dim()..self.rho_dim() + self.log_kappa_dim()])
5903 .assign(self.log_kappa_upper.as_array());
5904 out.slice_mut(s![self.rho_dim() + self.log_kappa_dim()..])
5905 .assign(&self.auxiliary_upper);
5906 out
5907 }
5908
5909 pub(crate) fn log_kappa_dims_per_term(&self) -> Vec<usize> {
5911 self.log_kappa0.dims_per_term().to_vec()
5912 }
5913}
5914
5915struct ExactJointDesignCache<'d> {
5921 realizers: Vec<FrozenTermCollectionIncrementalRealizer<'d>>,
5922 block_term_indices: Vec<Vec<usize>>,
5923 current_theta: Option<Array1<f64>>,
5924 last_cost: Option<f64>,
5925 last_eval: Option<(f64, Array1<f64>, gam_problem::HessianValue)>,
5926 rho_dim: usize,
5927 all_dims: Vec<usize>,
5928 log_kappa_dim: usize,
5929 block_term_counts: Vec<usize>,
5930}
5931
5932impl<'d> ExactJointDesignCache<'d> {
5933 fn new(
5934 data: ArrayView2<'d, f64>,
5935 blocks: Vec<(TermCollectionSpec, TermCollectionDesign, Vec<usize>)>,
5936 rho_dim: usize,
5937 all_dims: Vec<usize>,
5938 ) -> Result<Self, String> {
5939 let n_blocks = blocks.len();
5940 let mut realizers = Vec::with_capacity(n_blocks);
5941 let mut block_term_indices = Vec::with_capacity(n_blocks);
5942 let mut block_term_counts = Vec::with_capacity(n_blocks);
5943
5944 for (spec, design, terms) in blocks {
5945 block_term_counts.push(terms.len());
5946 block_term_indices.push(terms);
5947 realizers.push(FrozenTermCollectionIncrementalRealizer::new(
5948 data, spec, design,
5949 )?);
5950 }
5951
5952 Ok(Self {
5953 realizers,
5954 block_term_indices,
5955 current_theta: None,
5956 last_cost: None,
5957 last_eval: None,
5958 rho_dim,
5959 log_kappa_dim: all_dims.iter().sum(),
5960 all_dims,
5961 block_term_counts,
5962 })
5963 }
5964
5965 fn ensure_theta(&mut self, theta: &Array1<f64>) -> Result<(), String> {
5966 if self
5967 .current_theta
5968 .as_ref()
5969 .is_some_and(|cached| theta_values_match(cached, theta))
5970 {
5971 return Ok(());
5972 }
5973
5974 let t_ensure = std::time::Instant::now();
5975 let kappa_theta_len = self.rho_dim + self.log_kappa_dim;
5976 if theta.len() < kappa_theta_len {
5977 return Err(SmoothError::dimension_mismatch(format!(
5978 "exact-joint theta length mismatch: got {}, expected at least {} (rho_dim={}, log_kappa_dim={})",
5979 theta.len(),
5980 kappa_theta_len,
5981 self.rho_dim,
5982 self.log_kappa_dim
5983 ))
5984 .into());
5985 }
5986 let theta_kappa = theta.slice(s![..kappa_theta_len]).to_owned();
5987 let full_log_kappa = SpatialLogKappaCoords::from_theta_tail_with_dims(
5988 &theta_kappa,
5989 self.rho_dim,
5990 self.all_dims.clone(),
5991 );
5992
5993 let n = self.realizers.len();
5997 let mut remaining = full_log_kappa;
5998 for block_idx in 0..n {
5999 let count = self.block_term_counts[block_idx];
6000 if block_idx < n - 1 {
6001 let (block_lk, rest) = remaining.split_at(count);
6002 self.realizers[block_idx]
6003 .apply_log_kappa(&block_lk, &self.block_term_indices[block_idx])?;
6004 remaining = rest;
6005 } else {
6006 self.realizers[block_idx]
6008 .apply_log_kappa(&remaining, &self.block_term_indices[block_idx])?;
6009 }
6010 }
6011
6012 log::info!(
6013 "[STAGE] ensure_theta (n-block, {} blocks, {} realizers): {:.3}s",
6014 n,
6015 self.realizers.len(),
6016 t_ensure.elapsed().as_secs_f64(),
6017 );
6018 self.current_theta = Some(theta.clone());
6019 self.last_cost = None;
6020 self.last_eval = None;
6021 Ok(())
6022 }
6023
6024 impl_exact_joint_theta_memo!();
6025
6026 fn store_cost_only(&mut self, theta: &Array1<f64>, cost: f64) {
6032 if self
6033 .current_theta
6034 .as_ref()
6035 .is_some_and(|cached| theta_values_match(cached, theta))
6036 {
6037 self.last_cost = Some(cost);
6038 }
6039 }
6040
6041 fn specs(&self) -> Vec<&TermCollectionSpec> {
6042 self.realizers.iter().map(|r| r.spec()).collect()
6043 }
6044
6045 fn designs(&self) -> Vec<&TermCollectionDesign> {
6046 self.realizers.iter().map(|r| r.design()).collect()
6047 }
6048
6049 fn design_revision(&self) -> u64 {
6059 self.realizers
6060 .iter()
6061 .fold(0u64, |acc, r| acc.wrapping_add(r.design_revision()))
6062 }
6063}
6064
6065pub(crate) fn seed_risk_profile_for_likelihood_family(
6066 family: &LikelihoodSpec,
6067) -> gam_problem::SeedRiskProfile {
6068 match &family.response {
6069 ResponseFamily::Gaussian => gam_problem::SeedRiskProfile::Gaussian,
6070 ResponseFamily::RoystonParmar => gam_problem::SeedRiskProfile::Survival,
6071 ResponseFamily::Binomial
6072 | ResponseFamily::Poisson
6073 | ResponseFamily::Tweedie { .. }
6074 | ResponseFamily::NegativeBinomial { .. }
6075 | ResponseFamily::Beta { .. }
6076 | ResponseFamily::Gamma => gam_problem::SeedRiskProfile::GeneralizedLinear,
6077 }
6078}
6079
6080const EXACT_JOINT_SECOND_ORDER_THETA_CAP: usize = 8;
6088
6089fn exact_joint_seed_config(
6090 risk_profile: gam_problem::SeedRiskProfile,
6091 auxiliary_dim: usize,
6092 initial_seed_only: bool,
6093) -> gam_problem::SeedConfig {
6094 let mut config = gam_problem::SeedConfig {
6095 risk_profile,
6096 num_auxiliary_trailing: auxiliary_dim,
6097 ..Default::default()
6098 };
6099 match risk_profile {
6100 gam_problem::SeedRiskProfile::Gaussian
6101 | gam_problem::SeedRiskProfile::GaussianLocationScale => {
6102 config.max_seeds = 4;
6103 config.seed_budget = 2;
6104 }
6105 gam_problem::SeedRiskProfile::GeneralizedLinear => {
6106 config.max_seeds = 1;
6111 config.seed_budget = 1;
6112 config.screen_max_inner_iterations = 8;
6113 }
6114 gam_problem::SeedRiskProfile::Survival => {
6115 config.max_seeds = 8;
6121 config.seed_budget = 4;
6122 config.screen_max_inner_iterations = 8;
6123 }
6124 }
6125 if initial_seed_only {
6126 config.max_seeds = 1;
6133 config.seed_budget = 1;
6134 config.over_smoothing_probe_rho = None;
6135 }
6136 config
6137}
6138
6139#[cfg(test)]
6140mod exact_joint_seed_config_tests {
6141 use super::*;
6142
6143 #[test]
6144 fn exact_joint_marginal_slope_profiles_get_deeper_startup_validation() {
6145 let bms = exact_joint_seed_config(
6146 gam_problem::SeedRiskProfile::GeneralizedLinear,
6147 2,
6148 false,
6149 );
6150 assert_eq!(bms.max_seeds, 1);
6151 assert_eq!(bms.seed_budget, 1);
6152 assert_eq!(bms.screen_max_inner_iterations, 8);
6153 assert_eq!(bms.num_auxiliary_trailing, 2);
6154
6155 let survival =
6156 exact_joint_seed_config(gam_problem::SeedRiskProfile::Survival, 3, false);
6157 assert_eq!(survival.max_seeds, 8);
6158 assert_eq!(survival.seed_budget, 4);
6159 assert_eq!(survival.screen_max_inner_iterations, 8);
6160 assert_eq!(survival.num_auxiliary_trailing, 3);
6161 }
6162
6163 #[test]
6164 fn exact_joint_gaussian_keeps_tight_historical_multistart_budget() {
6165 let gaussian = exact_joint_seed_config(gam_problem::SeedRiskProfile::Gaussian, 1, false);
6166 assert_eq!(gaussian.max_seeds, 4);
6167 assert_eq!(gaussian.seed_budget, 2);
6168 assert_eq!(
6169 gaussian.screen_max_inner_iterations,
6170 gam_problem::SeedConfig::default().screen_max_inner_iterations
6171 );
6172 assert_eq!(gaussian.num_auxiliary_trailing, 1);
6173 }
6174
6175 #[test]
6176 fn certified_matern_basin_owns_the_only_joint_start() {
6177 let gaussian = exact_joint_seed_config(gam_problem::SeedRiskProfile::Gaussian, 1, true);
6178 assert_eq!(gaussian.max_seeds, 1);
6179 assert_eq!(gaussian.seed_budget, 1);
6180 assert_eq!(gaussian.over_smoothing_probe_rho, None);
6181 assert_eq!(gaussian.num_auxiliary_trailing, 1);
6182 }
6183}
6184
6185#[cfg(test)]
6186mod wood_reference_df_tests {
6187 use super::*;
6188
6189 #[test]
6195 fn edf1_equals_two_trace_minus_trace_of_square() {
6196 let f = ndarray::array![[0.9_f64, 0.0], [0.0, 0.4]];
6200 let got = wood_reference_df(Some(&f), &(0..2)).unwrap();
6201 assert!(
6202 (got - 1.63).abs() < 1e-12,
6203 "edf1 should be 2*tr - tr(F^2) = 1.63, got {got}"
6204 );
6205 let edf = 1.3;
6208 assert!(got >= edf - 1e-12, "edf1 {got} must be >= edf {edf}");
6209 }
6210
6211 #[test]
6212 fn edf1_never_collapses_below_edf_when_offdiagonals_blow_up() {
6213 let f = ndarray::array![[0.5_f64, 40.0], [40.0, 0.5]];
6220 let tr = 1.0_f64;
6221 let got = wood_reference_df(Some(&f), &(0..2)).unwrap();
6222 assert!(
6223 got >= tr - 1e-12,
6224 "edf1 must be floored at edf (=tr={tr}) even when tr(F^2) explodes, got {got}"
6225 );
6226 assert!(
6227 got.is_finite() && got > 0.0,
6228 "edf1 must stay finite/positive"
6229 );
6230 }
6231
6232 #[test]
6233 fn returns_none_on_nonpositive_or_missing_trace() {
6234 assert!(wood_reference_df(None, &(0..2)).is_none());
6237 let zero = ndarray::array![[0.0_f64, 0.0], [0.0, 0.0]];
6239 assert!(wood_reference_df(Some(&zero), &(0..2)).is_none());
6240 let f = ndarray::array![[0.5_f64, 0.0], [0.0, 0.5]];
6242 assert!(wood_reference_df(Some(&f), &(0..5)).is_none());
6243 }
6244}
6245
6246pub(crate) fn exact_joint_multistart_outer_problem(
6247 theta0: &Array1<f64>,
6248 lower: &Array1<f64>,
6249 upper: &Array1<f64>,
6250 rho_dim: usize,
6251 auxiliary_dim: usize,
6252 n_params: usize,
6253 gradient: gam_problem::Derivative,
6254 hessian: gam_problem::DeclaredHessianForm,
6255 prefer_gradient_only: bool,
6256 disable_fixed_point: bool,
6257 risk_profile: gam_problem::SeedRiskProfile,
6258 tolerance: f64,
6259 max_iter: usize,
6260 bfgs_step_cap: Option<f64>,
6269 bfgs_step_cap_psi: Option<f64>,
6270 screening_cap: Option<Arc<AtomicUsize>>,
6271 profiled_objective_size: Option<(usize, usize)>,
6292 has_constant_curvature: bool,
6301 initial_seed_only: bool,
6306) -> gam_solve::rho_optimizer::OuterProblem {
6307 let mut seed_heuristic = theta0.to_vec();
6308 for value in &mut seed_heuristic[..rho_dim] {
6309 *value = value.exp();
6310 }
6311 let rho_ceiling = if has_constant_curvature {
6316 gam_solve::estimate::RHO_BOUND
6317 } else {
6318 12.0
6319 };
6320 let mut problem = gam_solve::rho_optimizer::OuterProblem::new(n_params)
6321 .with_gradient(gradient)
6322 .with_hessian(hessian)
6323 .with_prefer_gradient_only(prefer_gradient_only)
6324 .with_disable_fixed_point(disable_fixed_point)
6325 .with_fallback_policy(gam_solve::rho_optimizer::FallbackPolicy::Automatic)
6335 .with_psi_dim(auxiliary_dim)
6336 .with_tolerance(tolerance)
6337 .with_max_iter(max_iter)
6338 .with_bounds(lower.clone(), upper.clone())
6339 .with_initial_rho(theta0.clone())
6340 .with_bfgs_step_cap(bfgs_step_cap)
6341 .with_bfgs_step_cap_psi(bfgs_step_cap_psi)
6342 .with_seed_config({
6343 let mut sc =
6344 exact_joint_seed_config(risk_profile, auxiliary_dim, initial_seed_only);
6345 if has_constant_curvature {
6346 sc.bounds = (sc.bounds.0, rho_ceiling);
6350 }
6365 sc
6366 })
6367 .with_rho_bound(rho_ceiling)
6368 .with_heuristic_lambdas(seed_heuristic);
6369 if let Some((n_obs, p_cols)) = profiled_objective_size {
6370 problem = problem
6378 .with_objective_scale(Some(n_obs as f64))
6379 .with_problem_size(n_obs, p_cols)
6380 .with_arc_initial_regularization(Some(0.25))
6381 .with_operator_initial_trust_radius(Some(4.0));
6382 }
6383 if let Some(screening_cap) = screening_cap {
6384 problem = problem
6385 .with_screening_cap(screening_cap)
6386 .with_screen_initial_rho(true);
6387 }
6388 problem
6389}
6390
6391pub fn optimize_spatial_length_scale_exact_joint<FitOut, FitFn, ExactFn, ExactEfsFn, SeedFn>(
6392 data: ArrayView2<'_, f64>,
6393 block_specs: &[TermCollectionSpec],
6394 block_term_indices: &[Vec<usize>],
6395 kappa_options: &SpatialLengthScaleOptimizationOptions,
6396 joint_setup: &ExactJointHyperSetup,
6397 seed_risk_profile: gam_problem::SeedRiskProfile,
6398 analytic_joint_gradient_available: bool,
6399 analytic_joint_hessian_available: bool,
6400 disable_fixed_point: bool,
6401 screening_cap: Option<Arc<AtomicUsize>>,
6402 outer_derivative_policy: gam_model_api::families::custom_family::OuterDerivativePolicy,
6403 mut fit_fn: FitFn,
6404 mut exact_fn: ExactFn,
6405 mut exact_efs_fn: ExactEfsFn,
6406 mut seed_inner_beta_fn: SeedFn,
6407) -> Result<SpatialLengthScaleOptimizationResult<FitOut>, String>
6408where
6409 FitOut: Clone,
6410 FitFn: FnMut(
6411 &Array1<f64>,
6412 &[TermCollectionSpec],
6413 &[TermCollectionDesign],
6414 ) -> Result<FitOut, String>,
6415 ExactFn: FnMut(
6416 &Array1<f64>,
6417 &[TermCollectionSpec],
6418 &[TermCollectionDesign],
6419 gam_solve::estimate::reml::reml_outer_engine::EvalMode,
6420 &gam_problem::outer_subsample::RowSet,
6421 ) -> Result<(f64, Array1<f64>, gam_problem::HessianValue), String>,
6422 ExactEfsFn: FnMut(
6423 &Array1<f64>,
6424 &[TermCollectionSpec],
6425 &[TermCollectionDesign],
6426 ) -> Result<gam_problem::EfsEval, String>,
6427 SeedFn: FnMut(&Array1<f64>) -> Result<gam_solve::rho_optimizer::SeedOutcome, EstimationError>,
6428{
6429 let n_blocks = block_specs.len();
6430 if block_term_indices.len() != n_blocks {
6431 return Err(SmoothError::dimension_mismatch(format!(
6432 "block_specs ({}) and block_term_indices ({}) length mismatch",
6433 n_blocks,
6434 block_term_indices.len()
6435 ))
6436 .into());
6437 }
6438
6439 let log_kappa_dim = joint_setup.log_kappa_dim();
6440
6441 log::trace!(
6442 "[spatial-exact-joint] driver entry: aux_dim={} log_kappa_dim={} kappa_enabled={} rho_dim={} theta0_len={}",
6443 joint_setup.auxiliary_dim(),
6444 log_kappa_dim,
6445 kappa_options.enabled,
6446 joint_setup.rho_dim(),
6447 joint_setup.theta0().len()
6448 );
6449
6450 if joint_setup.auxiliary_dim() == 0 && (!kappa_options.enabled || log_kappa_dim == 0) {
6454 log::trace!(
6455 "[spatial-exact-joint] taking fast path (no outer theta optimization in this driver)"
6456 );
6457 let (designs, resolved_specs) = build_term_collection_designs_and_freeze_joint(
6458 data, block_specs,
6459 )
6460 .map_err(|e| {
6461 format!("failed to build and freeze joint block designs during exact joint kappa optimization: {e}")
6462 })?;
6463 let theta0 = joint_setup.theta0();
6464
6465 let spec_refs: Vec<TermCollectionSpec> = resolved_specs.clone();
6467 let design_refs: Vec<TermCollectionDesign> = designs.clone();
6468 let fit = fit_fn(&theta0, &spec_refs, &design_refs)?;
6469 return Ok(SpatialLengthScaleOptimizationResult {
6470 resolved_specs,
6471 designs,
6472 fit,
6473 timing: None,
6474 });
6475 }
6476
6477 let theta0 = joint_setup.theta0();
6481 let lower = joint_setup.lower();
6482 let upper = joint_setup.upper();
6483 if theta0.len() < log_kappa_dim || lower.len() != theta0.len() || upper.len() != theta0.len() {
6484 return Err(SmoothError::dimension_mismatch(format!(
6485 "invalid exact joint theta setup: theta0={}, lower={}, upper={}, required_log_kappa_dim={}",
6486 theta0.len(),
6487 lower.len(),
6488 upper.len(),
6489 log_kappa_dim
6490 ))
6491 .into());
6492 }
6493 let rho_dim = joint_setup.rho_dim();
6494 let all_dims = joint_setup.log_kappa_dims_per_term();
6495
6496 let (boot_designs, best_specs) = build_term_collection_designs_and_freeze_joint(
6498 data,
6499 block_specs,
6500 )
6501 .map_err(|e| {
6502 format!(
6503 "failed to build and freeze joint block designs during exact joint kappa bootstrap: {e}"
6504 )
6505 })?;
6506 let policy_hessian_form = outer_derivative_policy.declared_hessian_form();
6516 let analytic_outer_hessian_available = analytic_joint_hessian_available
6517 && matches!(
6518 policy_hessian_form,
6519 gam_problem::DeclaredHessianForm::Either
6520 | gam_problem::DeclaredHessianForm::Dense
6521 | gam_problem::DeclaredHessianForm::Operator { .. }
6522 );
6523 let prefer_gradient_only = !analytic_outer_hessian_available;
6524
6525 let theta_dim = theta0.len();
6526 let psi_dim = theta_dim - rho_dim;
6527
6528 let cache_blocks: Vec<(TermCollectionSpec, TermCollectionDesign, Vec<usize>)> = best_specs
6530 .iter()
6531 .zip(boot_designs.iter())
6532 .zip(block_term_indices.iter())
6533 .map(|((spec, design), terms)| (spec.clone(), design.clone(), terms.clone()))
6534 .collect();
6535
6536 struct NBlockExactJointState<'d> {
6537 cache: ExactJointDesignCache<'d>,
6538 }
6539
6540 let mut state = NBlockExactJointState {
6541 cache: ExactJointDesignCache::new(data, cache_blocks, rho_dim, all_dims.clone())?,
6542 };
6543
6544 const KAPPA_PILOT_K: usize = 5_000;
6569 const KAPPA_POLISH_K: usize = 25_000;
6570 const KAPPA_POLISH_TRIGGER_N: usize = 100_000;
6571
6572 let n_total = data.nrows();
6573 let use_staged_kappa = outer_derivative_policy.should_use_staged_kappa(n_total);
6574 if use_staged_kappa {
6575 log::info!(
6576 "[KAPPA-STAGED] auto-engaging pilot+polish schedule: n={} pilot_k={} polish_k={}",
6577 n_total,
6578 KAPPA_PILOT_K,
6579 KAPPA_POLISH_K,
6580 );
6581 }
6582
6583 fn build_uniform_pilot_subsample(
6600 n_total: usize,
6601 k_target: usize,
6602 seed: u64,
6603 ) -> gam_problem::outer_subsample::OuterScoreSubsample {
6604 use gam_problem::outer_subsample::OuterScoreSubsample;
6605 let k = k_target.min(n_total);
6606 if k == 0 || n_total == 0 {
6607 return OuterScoreSubsample::from_uniform_inclusion_mask(Vec::new(), n_total, seed);
6608 }
6609 let mut mask: Vec<usize> = Vec::with_capacity(k);
6613 let mut state = seed.wrapping_add(0x9E3779B97F4A7C15);
6615 let splitmix = |s: &mut u64| -> u64 { gam_linalg::utils::splitmix64(s) };
6616 let mut taken = std::collections::HashSet::with_capacity(k);
6617 for j in (n_total - k)..n_total {
6618 let r = (splitmix(&mut state) % (j as u64 + 1)) as usize;
6619 if !taken.insert(r) {
6620 taken.insert(j);
6621 mask.push(j);
6622 } else {
6623 mask.push(r);
6624 }
6625 }
6626 mask.sort_unstable();
6627 mask.dedup();
6628 OuterScoreSubsample::from_uniform_inclusion_mask(mask, n_total, seed)
6629 }
6630
6631 let current_row_set: std::cell::RefCell<gam_problem::outer_subsample::RowSet> =
6632 if use_staged_kappa {
6633 let pilot = build_uniform_pilot_subsample(n_total, KAPPA_PILOT_K, n_total as u64);
6634 std::cell::RefCell::new(gam_problem::outer_subsample::RowSet::Subsample {
6635 rows: std::sync::Arc::clone(&pilot.rows),
6636 n_full: n_total,
6637 })
6638 } else {
6639 std::cell::RefCell::new(gam_problem::outer_subsample::RowSet::All)
6640 };
6641
6642 let exact_fn_cell = std::cell::RefCell::new(&mut exact_fn);
6643 let exact_efs_fn_cell = std::cell::RefCell::new(&mut exact_efs_fn);
6644
6645 use std::cell::Cell;
6660 let kphase_cost_calls: Cell<usize> = Cell::new(0);
6661 let kphase_cost_total_s: Cell<f64> = Cell::new(0.0);
6662 let kphase_eval_calls: Cell<usize> = Cell::new(0);
6663 let kphase_eval_total_s: Cell<f64> = Cell::new(0.0);
6664 let kphase_efs_calls: Cell<usize> = Cell::new(0);
6665 let kphase_efs_total_s: Cell<f64> = Cell::new(0.0);
6666 let kphase_optim_start = std::time::Instant::now();
6667 let kphase_log_kappa_dim = log_kappa_dim;
6668 let kphase_log_norms = |theta: &Array1<f64>| -> (f64, f64) {
6669 let theta_norm = theta.iter().map(|v| v * v).sum::<f64>().sqrt();
6670 let log_kappa_norm = if kphase_log_kappa_dim > 0 && theta.len() >= kphase_log_kappa_dim {
6671 let start = theta.len() - kphase_log_kappa_dim;
6672 theta.iter().skip(start).map(|v| v * v).sum::<f64>().sqrt()
6673 } else {
6674 0.0
6675 };
6676 (theta_norm, log_kappa_norm)
6677 };
6678
6679 use gam_problem::{DeclaredHessianForm, Derivative, OuterEval};
6680 use gam_solve::rho_optimizer::OuterEvalOrder;
6681
6682 let joint_p_cols: usize = boot_designs
6686 .iter()
6687 .map(|d| d.design.ncols())
6688 .sum::<usize>()
6689 .max(1);
6690
6691 let problem = exact_joint_multistart_outer_problem(
6692 &theta0,
6693 &lower,
6694 &upper,
6695 rho_dim,
6696 psi_dim,
6697 theta_dim,
6698 if analytic_joint_gradient_available {
6699 Derivative::Analytic
6700 } else {
6701 Derivative::Unavailable
6702 },
6703 if analytic_outer_hessian_available {
6704 DeclaredHessianForm::Either
6705 } else {
6706 DeclaredHessianForm::Unavailable
6707 },
6708 prefer_gradient_only,
6709 disable_fixed_point,
6710 seed_risk_profile,
6711 kappa_options.rel_tol.max(1e-6),
6712 kappa_options.max_outer_iter.max(1),
6713 Some(5.0),
6715 Some(kappa_options.log_step.clamp(0.25, 1.0)),
6717 screening_cap.clone(),
6718 Some((n_total, joint_p_cols)),
6721 block_specs
6724 .iter()
6725 .any(|s| !constant_curvature_term_indices(s).is_empty()),
6726 false,
6729 );
6730
6731 fn collect_specs(cache: &ExactJointDesignCache<'_>) -> Vec<TermCollectionSpec> {
6733 cache.specs().into_iter().cloned().collect()
6734 }
6735 fn collect_designs(cache: &ExactJointDesignCache<'_>) -> Vec<TermCollectionDesign> {
6736 cache.designs().into_iter().cloned().collect()
6737 }
6738
6739 let result = {
6740 let eval_outer = |ctx: &mut &mut NBlockExactJointState<'_>,
6741 theta: &Array1<f64>,
6742 order: OuterEvalOrder|
6743 -> Result<OuterEval, EstimationError> {
6744 if let Some((cost, grad, hess)) = ctx.cache.memoized_eval(theta) {
6745 let cached_satisfies_order = match order {
6746 OuterEvalOrder::Value => true,
6747 OuterEvalOrder::ValueAndGradient => true,
6748 OuterEvalOrder::ValueGradientHessian => hess.is_analytic(),
6749 };
6750 if cached_satisfies_order {
6751 if !cost.is_finite() {
6752 return Ok(OuterEval::infeasible(theta.len()));
6753 }
6754 if grad.iter().any(|v| !v.is_finite()) {
6767 return Ok(OuterEval::infeasible(theta.len()));
6768 }
6769 return Ok(OuterEval {
6770 cost,
6771 gradient: grad,
6772 hessian: hess,
6773 inner_beta_hint: None,
6774 });
6775 }
6776 }
6777 if let Err(err) = ctx.cache.ensure_theta(theta) {
6778 log::warn!(
6779 "[OUTER] n-block exact-joint spatial: ensure_theta failed during gradient evaluation: {err}"
6780 );
6781 return Ok(OuterEval::infeasible(theta.len()));
6782 }
6783 let design_revision = Some(ctx.cache.design_revision());
6784 let specs = collect_specs(&ctx.cache);
6785 let designs = collect_designs(&ctx.cache);
6786 let clamped = outer_derivative_policy.order_for_evaluation(order);
6794 let need_hessian = matches!(clamped, OuterEvalOrder::ValueGradientHessian)
6795 && analytic_outer_hessian_available;
6796 let eval_mode = if need_hessian {
6797 gam_solve::estimate::reml::reml_outer_engine::EvalMode::ValueGradientHessian
6798 } else {
6799 gam_solve::estimate::reml::reml_outer_engine::EvalMode::ValueAndGradient
6800 };
6801 let t0 = std::time::Instant::now();
6802 let result = {
6803 let row_set_borrow = current_row_set.borrow();
6804 (*exact_fn_cell.borrow_mut())(theta, &specs, &designs, eval_mode, &row_set_borrow)
6805 };
6806 let elapsed_s = t0.elapsed().as_secs_f64();
6807 kphase_eval_calls.set(kphase_eval_calls.get() + 1);
6808 kphase_eval_total_s.set(kphase_eval_total_s.get() + elapsed_s);
6809 let (theta_norm, log_kappa_norm) = kphase_log_norms(theta);
6810 log::info!(
6811 "[KAPPA-PHASE] phase=eval_outer call={} order={:?} design_revision={:?} theta_norm={:.4e} log_kappa_norm={:.4e} elapsed_s={:.4}",
6812 kphase_eval_calls.get(),
6813 order,
6814 design_revision,
6815 theta_norm,
6816 log_kappa_norm,
6817 elapsed_s,
6818 );
6819 match result {
6820 Ok((cost, grad, hess)) => {
6821 ctx.cache.store_eval((cost, grad.clone(), hess.clone()));
6822 if !cost.is_finite() {
6823 return Ok(OuterEval::infeasible(theta.len()));
6824 }
6825 if grad.iter().any(|v| !v.is_finite()) {
6838 return Ok(OuterEval::infeasible(theta.len()));
6839 }
6840 Ok(OuterEval {
6841 cost,
6842 gradient: grad,
6843 hessian: hess,
6844 inner_beta_hint: None,
6845 })
6846 }
6847 Err(err) => {
6848 log::warn!(
6849 "[OUTER] n-block exact-joint spatial: exact evaluation failed: {err}"
6850 );
6851 Ok(OuterEval::infeasible(theta.len()))
6852 }
6853 }
6854 };
6855
6856 let obj = problem.build_objective_with_eval_order(
6857 &mut state,
6858 |ctx: &mut &mut NBlockExactJointState<'_>, theta: &Array1<f64>| {
6859 if let Some(cost) = ctx.cache.memoized_cost(theta) {
6860 return Ok(cost);
6861 }
6862 if let Err(err) = ctx.cache.ensure_theta(theta) {
6863 log::warn!(
6864 "[OUTER] n-block exact-joint spatial: ensure_theta failed during cost evaluation: {err}"
6865 );
6866 return Ok(f64::INFINITY);
6867 }
6868 let design_revision = Some(ctx.cache.design_revision());
6869 let specs = collect_specs(&ctx.cache);
6870 let designs = collect_designs(&ctx.cache);
6871 let t0 = std::time::Instant::now();
6878 let result = {
6879 let row_set_borrow = current_row_set.borrow();
6880 (*exact_fn_cell.borrow_mut())(
6881 theta,
6882 &specs,
6883 &designs,
6884 gam_solve::estimate::reml::reml_outer_engine::EvalMode::ValueOnly,
6885 &row_set_borrow,
6886 )
6887 };
6888 let elapsed_s = t0.elapsed().as_secs_f64();
6889 kphase_cost_calls.set(kphase_cost_calls.get() + 1);
6890 kphase_cost_total_s.set(kphase_cost_total_s.get() + elapsed_s);
6891 let (theta_norm, log_kappa_norm) = kphase_log_norms(theta);
6892 log::info!(
6893 "[KAPPA-PHASE] phase=cost call={} design_revision={:?} theta_norm={:.4e} log_kappa_norm={:.4e} elapsed_s={:.4}",
6894 kphase_cost_calls.get(),
6895 design_revision,
6896 theta_norm,
6897 log_kappa_norm,
6898 elapsed_s,
6899 );
6900 match result {
6901 Ok((cost, _grad, _hess)) => {
6902 ctx.cache.store_cost_only(theta, cost);
6908 Ok(cost)
6909 }
6910 Err(err) => {
6911 log::warn!(
6912 "[OUTER] n-block exact-joint spatial: exact cost evaluation failed: {err}"
6913 );
6914 Ok(f64::INFINITY)
6915 }
6916 }
6917 },
6918 |ctx: &mut &mut NBlockExactJointState<'_>, theta: &Array1<f64>| {
6919 eval_outer(
6920 ctx,
6921 theta,
6922 if analytic_outer_hessian_available {
6923 OuterEvalOrder::ValueGradientHessian
6924 } else {
6925 OuterEvalOrder::ValueAndGradient
6926 },
6927 )
6928 },
6929 |ctx: &mut &mut NBlockExactJointState<'_>,
6930 theta: &Array1<f64>,
6931 order: OuterEvalOrder| { eval_outer(ctx, theta, order) },
6932 None::<fn(&mut &mut NBlockExactJointState<'_>)>,
6933 Some(
6934 |ctx: &mut &mut NBlockExactJointState<'_>, theta: &Array1<f64>| {
6935 ctx.cache
6936 .ensure_theta(theta)
6937 .map_err(EstimationError::InvalidInput)?;
6938 let design_revision = Some(ctx.cache.design_revision());
6939 let specs = collect_specs(&ctx.cache);
6940 let designs = collect_designs(&ctx.cache);
6941 let t0 = std::time::Instant::now();
6942 let eval_result = (*exact_efs_fn_cell.borrow_mut())(
6943 theta,
6944 &specs,
6945 &designs,
6946 );
6947 let elapsed_s = t0.elapsed().as_secs_f64();
6948 kphase_efs_calls.set(kphase_efs_calls.get() + 1);
6949 kphase_efs_total_s.set(kphase_efs_total_s.get() + elapsed_s);
6950 let (theta_norm, log_kappa_norm) = kphase_log_norms(theta);
6951 log::info!(
6952 "[KAPPA-PHASE] phase=efs call={} design_revision={:?} theta_norm={:.4e} log_kappa_norm={:.4e} elapsed_s={:.4}",
6953 kphase_efs_calls.get(),
6954 design_revision,
6955 theta_norm,
6956 log_kappa_norm,
6957 elapsed_s,
6958 );
6959 let eval = eval_result.map_err(EstimationError::RemlOptimizationFailed)?;
6960 Ok(eval)
6961 },
6962 ),
6963 );
6964 let mut obj = obj.with_seed_inner_state(
6965 move |_ctx: &mut &mut NBlockExactJointState<'_>, beta: &Array1<f64>| {
6966 (seed_inner_beta_fn)(beta)
6967 },
6968 );
6969
6970 problem
6971 .run(&mut obj, "n-block exact-joint spatial")
6972 .map_err(|error| error.to_string())?
6973 }; let kphase_total_s = kphase_optim_start.elapsed().as_secs_f64();
6983 log::info!(
6984 "[KAPPA-PHASE-SUMMARY] log_kappa_dim={} n_cost={} cost_total_s={:.4} n_eval={} eval_total_s={:.4} n_efs={} efs_total_s={:.4} optim_total_s={:.4}",
6985 kphase_log_kappa_dim,
6986 kphase_cost_calls.get(),
6987 kphase_cost_total_s.get(),
6988 kphase_eval_calls.get(),
6989 kphase_eval_total_s.get(),
6990 kphase_efs_calls.get(),
6991 kphase_efs_total_s.get(),
6992 kphase_total_s,
6993 );
6994 let timing = SpatialLengthScaleOptimizationTiming {
6995 log_kappa_dim: kphase_log_kappa_dim,
6996 cost_calls: kphase_cost_calls.get(),
6997 cost_total_s: kphase_cost_total_s.get(),
6998 eval_calls: kphase_eval_calls.get(),
6999 eval_total_s: kphase_eval_total_s.get(),
7000 efs_calls: kphase_efs_calls.get(),
7001 efs_total_s: kphase_efs_total_s.get(),
7002 slow_path_resets: 0,
7003 design_revision_delta: 0,
7004 nfree_skip_row_touches: 0,
7005 nfree_miss_shape: 0,
7006 nfree_miss_value: 0,
7007 nfree_miss_gradient: 0,
7008 nfree_miss_penalty: 0,
7009 nfree_miss_revision: 0,
7010 nfree_miss_second_order: 0,
7011 nfree_miss_other: 0,
7012 optim_total_s: kphase_total_s,
7013 };
7014
7015 if !result.converged {
7016 return Err(format!(
7017 "n-block exact-joint spatial κ optimization did not converge after {} iterations (final_objective={:.6e}, final_grad_norm={})",
7018 result.iterations,
7019 result.final_value,
7020 result.final_grad_norm_report(),
7021 ));
7022 }
7023 let theta_star = result.rho;
7024
7025 if use_staged_kappa && n_total >= KAPPA_POLISH_TRIGGER_N {
7042 let polish = build_uniform_pilot_subsample(
7043 n_total,
7044 KAPPA_POLISH_K,
7045 (n_total as u64).wrapping_add(0xA5A5A5A5),
7046 );
7047 *current_row_set.borrow_mut() = gam_problem::outer_subsample::RowSet::Subsample {
7048 rows: std::sync::Arc::clone(&polish.rows),
7049 n_full: n_total,
7050 };
7051 log::info!(
7052 "[KAPPA-STAGED] rotating to polish subsample: k={} at theta_star",
7053 polish.rows.len(),
7054 );
7055 state.cache.ensure_theta(&theta_star)?;
7059 let (polish_cost, polish_grad, _) = {
7060 let specs = collect_specs(&state.cache);
7061 let designs = collect_designs(&state.cache);
7062 let row_set_borrow = current_row_set.borrow();
7063 exact_fn(
7064 &theta_star,
7065 &specs,
7066 &designs,
7067 gam_solve::estimate::reml::reml_outer_engine::EvalMode::ValueAndGradient,
7068 &row_set_borrow,
7069 )?
7070 };
7071 if !polish_cost.is_finite() || polish_grad.iter().any(|value| !value.is_finite()) {
7072 return Err(
7073 "polish subsample exact-joint evaluation produced non-finite objective pieces"
7074 .to_string(),
7075 );
7076 }
7077 }
7078 *current_row_set.borrow_mut() = gam_problem::outer_subsample::RowSet::All;
7079 if use_staged_kappa {
7080 log::info!(
7081 "[KAPPA-STAGED] rotating to full data for final coefficient fit (n={})",
7082 n_total,
7083 );
7084 }
7085
7086 state.cache.ensure_theta(&theta_star)?;
7087
7088 let resolved_specs: Vec<TermCollectionSpec> = collect_specs(&state.cache);
7089 let designs: Vec<TermCollectionDesign> = collect_designs(&state.cache);
7090
7091 let fit = fit_fn(&theta_star, &resolved_specs, &designs)?;
7092
7093 for spec in &resolved_specs {
7094 log_spatial_aniso_scales(spec);
7095 }
7096
7097 Ok(SpatialLengthScaleOptimizationResult {
7098 resolved_specs,
7099 designs,
7100 fit,
7101 timing: Some(timing),
7102 })
7103}
7104
7105fn try_exact_joint_latent_coord_optimization(
7106 data: ArrayView2<'_, f64>,
7107 y: ArrayView1<'_, f64>,
7108 weights: ArrayView1<'_, f64>,
7109 offset: ArrayView1<'_, f64>,
7110 resolvedspec: &TermCollectionSpec,
7111 best: &FittedTermCollection,
7112 family: LikelihoodSpec,
7113 options: &FitOptions,
7114 latent: &StandardLatentCoordConfig,
7115) -> Result<FittedTermCollectionWithSpec, EstimationError> {
7116 use gam_problem::{DeclaredHessianForm, Derivative, OuterEval};
7117 use gam_solve::rho_optimizer::OuterEvalOrder;
7118
7119 let rho_dim = best.fit.lambdas.len();
7120 let latent_flat_dim = latent.values.len();
7121 if latent_flat_dim == 0 {
7122 crate::bail_invalid_estim!(
7123 "latent-coordinate optimization requires a non-empty latent block"
7124 );
7125 }
7126 let direct_hypers =
7127 latent_coord_initial_direct_hypers(latent.values.id_mode(), latent.values.latent_dim())?;
7128 let analytic_rho_count = latent
7129 .analytic_penalties
7130 .as_ref()
7131 .map_or(0, |registry| registry.total_rho_count());
7132 let latent_coord_ext_dim = latent_flat_dim + analytic_rho_count + direct_hypers.len();
7133
7134 let mut theta0 = Array1::<f64>::zeros(rho_dim + latent_coord_ext_dim);
7135 theta0
7136 .slice_mut(s![..rho_dim])
7137 .assign(&best.fit.lambdas.mapv(f64::ln));
7138 theta0
7139 .slice_mut(s![rho_dim..rho_dim + latent_flat_dim])
7140 .assign(latent.values.as_flat());
7141 if !direct_hypers.is_empty() {
7142 let direct_start = rho_dim + latent_flat_dim + analytic_rho_count;
7143 theta0
7144 .slice_mut(s![direct_start..direct_start + direct_hypers.len()])
7145 .assign(&direct_hypers);
7146 }
7147
7148 let mut lower = Array1::<f64>::from_elem(theta0.len(), -12.0);
7149 let mut upper = Array1::<f64>::from_elem(theta0.len(), 12.0);
7150 let latent_bound = latent
7151 .values
7152 .as_flat()
7153 .iter()
7154 .fold(1.0_f64, |acc, &v| acc.max(v.abs()))
7155 + 10.0;
7156 for axis in rho_dim..rho_dim + latent_flat_dim {
7157 lower[axis] = -latent_bound;
7158 upper[axis] = latent_bound;
7159 }
7160
7161 struct LatentJointContext<'d> {
7162 rho_dim: usize,
7163 cache: SingleBlockLatentCoordDesignCache,
7164 evaluator: gam_solve::estimate::ExternalJointHyperEvaluator<'d>,
7165 }
7166
7167 impl<'d> LatentJointContext<'d> {
7168 fn eval_full(
7169 &mut self,
7170 theta: &Array1<f64>,
7171 order: OuterEvalOrder,
7172 ) -> Result<(f64, Array1<f64>, gam_problem::HessianValue), EstimationError> {
7173 if let Some(eval) = self.cache.memoized_eval(theta) {
7174 return Ok(eval);
7175 }
7176 self.cache
7177 .ensure_theta(theta)
7178 .map_err(EstimationError::InvalidInput)?;
7179 let hyper_dirs = self
7180 .cache
7181 .hyper_dirs()
7182 .map_err(EstimationError::InvalidInput)?;
7183 let design_revision = Some(self.cache.design_revision());
7184 let registry_for_key = self.cache.analytic_penalties();
7185 self.evaluator
7186 .set_analytic_penalty_registry(registry_for_key.as_deref());
7187 let mut eval = evaluate_joint_reml_outer_eval_at_theta(
7188 &mut self.evaluator,
7189 self.cache.design(),
7190 theta,
7191 self.rho_dim,
7192 hyper_dirs,
7193 None,
7194 order,
7195 design_revision,
7196 )?;
7197 let latent = self.cache.latent().map_err(EstimationError::InvalidInput)?;
7198 if let Some(registry) = registry_for_key {
7199 let mut registry = registry.as_ref().clone();
7200 registry.apply_weight_schedules(
7201 gam_solve::estimate::reml::outer_eval::current_outer_iter() as usize,
7202 );
7203 add_analytic_penalty_objective_to_eval(
7204 theta,
7205 self.rho_dim,
7206 latent.as_ref(),
7207 ®istry,
7208 &mut eval,
7209 )?;
7210 }
7211 add_latent_id_objective_to_eval(
7212 theta,
7213 self.rho_dim,
7214 self.cache.analytic_penalty_rho_count(),
7215 latent.as_ref(),
7216 &mut eval,
7217 )?;
7218 self.cache.store_eval(eval.clone());
7219 Ok(eval)
7220 }
7221
7222 fn eval_efs(
7223 &mut self,
7224 theta: &Array1<f64>,
7225 ) -> Result<gam_problem::EfsEval, EstimationError> {
7226 self.cache
7227 .ensure_theta(theta)
7228 .map_err(EstimationError::InvalidInput)?;
7229 let hyper_dirs = self
7230 .cache
7231 .hyper_dirs()
7232 .map_err(EstimationError::InvalidInput)?;
7233 let registry_for_key = self.cache.analytic_penalties();
7234 self.evaluator
7235 .set_analytic_penalty_registry(registry_for_key.as_deref());
7236 let mut efs = evaluate_joint_reml_efs_at_theta(
7237 &mut self.evaluator,
7238 self.cache.design(),
7239 theta,
7240 self.rho_dim,
7241 hyper_dirs,
7242 None,
7243 Some(self.cache.design_revision()),
7244 )?;
7245 if let Some(registry) = registry_for_key {
7246 let mut registry = registry.as_ref().clone();
7247 registry.apply_weight_schedules(
7248 gam_solve::estimate::reml::outer_eval::current_outer_iter() as usize,
7249 );
7250 let latent = self.cache.latent().map_err(EstimationError::InvalidInput)?;
7251 let contribution = analytic_penalty_objective_contribution(
7252 theta,
7253 self.rho_dim,
7254 latent.as_ref(),
7255 ®istry,
7256 )?;
7257 efs.cost += contribution.cost;
7258 if let (Some(psi_gradient), Some(psi_indices)) =
7259 (efs.psi_gradient.as_mut(), efs.psi_indices.as_ref())
7260 {
7261 if psi_gradient.len() != psi_indices.len() {
7262 crate::bail_invalid_estim!(
7263 "latent-coordinate analytic penalty EFS psi gradient length mismatch: gradient={}, indices={}",
7264 psi_gradient.len(),
7265 psi_indices.len()
7266 );
7267 }
7268 for (local_idx, &theta_idx) in psi_indices.iter().enumerate() {
7269 psi_gradient[local_idx] += contribution.gradient[theta_idx];
7270 }
7271 }
7272 }
7273 Ok(efs)
7274 }
7275
7276 fn eval_cost(&mut self, theta: &Array1<f64>) -> f64 {
7277 if let Some(cost) = self.cache.memoized_cost(theta) {
7278 return cost;
7279 }
7280 if self.cache.ensure_theta(theta).is_err() {
7281 return f64::INFINITY;
7282 }
7283 let design_revision = Some(self.cache.design_revision());
7284 let registry_for_key = self.cache.analytic_penalties();
7285 self.evaluator
7286 .set_analytic_penalty_registry(registry_for_key.as_deref());
7287 let result = {
7288 let design = self.cache.design();
7289 self.evaluator.evaluate_cost_only(
7290 &design.design,
7291 &design.penalties,
7292 &design.nullspace_dims,
7293 design.linear_constraints.clone(),
7294 theta,
7295 self.rho_dim,
7296 None,
7297 "latent-coordinate-joint cost-only",
7298 design_revision,
7299 )
7300 };
7301 match result {
7302 Ok(cost) => {
7303 let latent = match self.cache.latent() {
7304 Ok(latent) => latent,
7305 Err(_) => return f64::INFINITY,
7306 };
7307 let contribution = match latent_id_objective_contribution(
7308 theta,
7309 self.rho_dim,
7310 self.cache.analytic_penalty_rho_count(),
7311 latent.as_ref(),
7312 ) {
7313 Ok(contribution) => contribution,
7314 Err(_) => return f64::INFINITY,
7315 };
7316 let cost = cost + contribution.cost;
7317 let cost = if let Some(registry) = registry_for_key {
7318 let mut registry = registry.as_ref().clone();
7319 registry.apply_weight_schedules(
7320 gam_solve::estimate::reml::outer_eval::current_outer_iter() as usize,
7321 );
7322 match analytic_penalty_objective_contribution(
7323 theta,
7324 self.rho_dim,
7325 latent.as_ref(),
7326 ®istry,
7327 ) {
7328 Ok(contribution) => cost + contribution.cost,
7329 Err(_) => return f64::INFINITY,
7330 }
7331 } else {
7332 cost
7333 };
7334 self.cache.store_cost(cost);
7335 cost
7336 }
7337 Err(_) => f64::INFINITY,
7338 }
7339 }
7340 }
7341
7342 let mut ctx = LatentJointContext {
7343 rho_dim,
7344 cache: SingleBlockLatentCoordDesignCache::new(
7345 data.to_owned(),
7346 resolvedspec.clone(),
7347 best.design.clone(),
7348 latent,
7349 rho_dim,
7350 )
7351 .map_err(EstimationError::InvalidInput)?,
7352 evaluator: gam_solve::estimate::ExternalJointHyperEvaluator::new(
7353 y,
7354 weights,
7355 &best.design.design,
7356 offset,
7357 &best.design.penalties,
7358 &external_opts_for_design(&family, &best.design, options),
7359 "latent-coordinate-joint",
7360 )?,
7361 };
7362 let registry_for_key = ctx.cache.analytic_penalties();
7363 ctx.evaluator
7364 .set_analytic_penalty_registry(registry_for_key.as_deref());
7365 ctx.evaluator
7366 .set_persistent_latent_values_fingerprint(latent.values.id_mode());
7367 if let Some(cached_t) = ctx
7368 .evaluator
7369 .load_persistent_latent_values(latent.values.n_obs(), latent.values.latent_dim())
7370 {
7371 let cached_t: Array2<f64> = cached_t;
7372 for (dst, src) in theta0
7373 .slice_mut(s![rho_dim..rho_dim + latent_flat_dim])
7374 .iter_mut()
7375 .zip(cached_t.iter())
7376 {
7377 *dst = *src;
7378 }
7379 }
7380
7381 let problem = exact_joint_multistart_outer_problem(
7382 &theta0,
7383 &lower,
7384 &upper,
7385 rho_dim,
7386 latent_coord_ext_dim,
7387 theta0.len(),
7388 Derivative::Analytic,
7389 DeclaredHessianForm::Unavailable,
7390 false,
7391 false,
7392 seed_risk_profile_for_likelihood_family(&family),
7393 options.tol,
7394 options.max_iter.max(1),
7395 Some(5.0),
7396 Some(0.5),
7397 None,
7398 Some((data.nrows(), best.design.design.ncols().max(1))),
7401 !constant_curvature_term_indices(resolvedspec).is_empty(),
7404 false,
7406 );
7407
7408 let eval_outer = |ctx: &mut &mut LatentJointContext<'_>,
7409 theta: &Array1<f64>,
7410 order: OuterEvalOrder|
7411 -> Result<OuterEval, EstimationError> {
7412 let (cost, gradient, hessian) = ctx.eval_full(theta, order)?;
7413 Ok(OuterEval {
7414 cost,
7415 gradient,
7416 hessian,
7417 inner_beta_hint: None,
7418 })
7419 };
7420
7421 let result = {
7422 let mut obj = problem.build_objective_with_eval_order(
7423 &mut ctx,
7424 |ctx: &mut &mut LatentJointContext<'_>, theta: &Array1<f64>| Ok(ctx.eval_cost(theta)),
7425 |ctx: &mut &mut LatentJointContext<'_>, theta: &Array1<f64>| {
7426 eval_outer(ctx, theta, OuterEvalOrder::ValueAndGradient)
7427 },
7428 |ctx: &mut &mut LatentJointContext<'_>, theta: &Array1<f64>, order: OuterEvalOrder| {
7429 eval_outer(ctx, theta, order)
7430 },
7431 Some(|ctx: &mut &mut LatentJointContext<'_>| {
7432 ctx.cache.reset();
7433 }),
7434 Some(|ctx: &mut &mut LatentJointContext<'_>, theta: &Array1<f64>| ctx.eval_efs(theta)),
7435 );
7436
7437 problem
7438 .run(&mut obj, "latent-coordinate joint REML")
7439 .map_err(|e| {
7440 EstimationError::InvalidInput(format!(
7441 "latent-coordinate joint optimization failed after exhausting strategy fallbacks: {e}"
7442 ))
7443 })?
7444 };
7445 if !result.converged {
7446 crate::bail_invalid_estim!(
7447 "latent-coordinate joint optimization did not converge after {} iterations (final_objective={:.6e}, final_grad_norm={})",
7448 result.iterations,
7449 result.final_value,
7450 result.final_grad_norm_report(),
7451 );
7452 }
7453
7454 let theta_star = result.rho;
7455 let rho_star = theta_star.slice(s![..rho_dim]).mapv(f64::exp);
7456 let mut final_data = data.to_owned();
7457 let flat_t = theta_star
7458 .slice(s![rho_dim..rho_dim + latent_flat_dim])
7459 .to_owned();
7460 let mut fitted_latent_values =
7461 Array2::<f64>::zeros((latent.values.n_obs(), latent.values.latent_dim()));
7462 for n in 0..latent.values.n_obs() {
7463 for axis in 0..latent.values.latent_dim() {
7464 let value = flat_t[n * latent.values.latent_dim() + axis];
7465 fitted_latent_values[[n, axis]] = value;
7466 final_data[[n, latent.feature_cols[axis]]] = value;
7467 }
7468 }
7469 let optimized = fit_term_collection_forspecwith_heuristic_lambdas(
7470 final_data.view(),
7471 y,
7472 weights,
7473 offset,
7474 resolvedspec,
7475 rho_star.as_slice(),
7476 family,
7477 options,
7478 )?;
7479 ctx.evaluator
7480 .store_persistent_latent_values(&fitted_latent_values);
7481 let mut fit = optimized.fit;
7482 fit.reml_score = result.final_value;
7483 fit.penalized_objective = result.final_value;
7484 Ok(FittedTermCollectionWithSpec {
7485 fit,
7486 design: optimized.design,
7487 resolvedspec: resolvedspec.clone(),
7488 adaptive_diagnostics: optimized.adaptive_diagnostics,
7489 kappa_timing: None,
7490 })
7491}
7492
7493pub fn fit_term_collectionwith_latent_coord_optimization(
7494 data: ArrayView2<'_, f64>,
7495 y: Array1<f64>,
7496 weights: Array1<f64>,
7497 offset: Array1<f64>,
7498 spec: &TermCollectionSpec,
7499 latent: &StandardLatentCoordConfig,
7500 family: LikelihoodSpec,
7501 options: &FitOptions,
7502) -> Result<FittedTermCollectionWithSpec, EstimationError> {
7503 let n = data.nrows();
7504 if !(y.len() == n && weights.len() == n && offset.len() == n) {
7505 crate::bail_invalid_estim!(
7506 "fit_term_collectionwith_latent_coord_optimization row mismatch: n={}, y={}, weights={}, offset={}",
7507 n,
7508 y.len(),
7509 weights.len(),
7510 offset.len()
7511 );
7512 }
7513 let best = fit_term_collection_forspec(
7514 data,
7515 y.view(),
7516 weights.view(),
7517 offset.view(),
7518 spec,
7519 family.clone(),
7520 options,
7521 )?;
7522 let resolvedspec = freeze_term_collection_from_design(spec, &best.design)?;
7523 try_exact_joint_latent_coord_optimization(
7524 data,
7525 y.view(),
7526 weights.view(),
7527 offset.view(),
7528 &resolvedspec,
7529 &best,
7530 family,
7531 options,
7532 latent,
7533 )
7534}
7535
7536fn select_isotropic_matern_range_basin(
7553 data: ArrayView2<'_, f64>,
7554 y: ArrayView1<'_, f64>,
7555 weights: ArrayView1<'_, f64>,
7556 offset: ArrayView1<'_, f64>,
7557 mut resolvedspec: TermCollectionSpec,
7558 mut best: FittedTermCollection,
7559 family: &LikelihoodSpec,
7560 options: &FitOptions,
7561 kappa_options: &SpatialLengthScaleOptimizationOptions,
7562 spatial_terms: &[usize],
7563) -> Result<(TermCollectionSpec, FittedTermCollection), EstimationError> {
7564 if has_aniso_terms(&resolvedspec, spatial_terms)
7568 || !constant_curvature_term_indices(&resolvedspec).is_empty()
7569 {
7570 return Ok((resolvedspec, best));
7571 }
7572
7573 let mut best_score = fit_score(&best.fit);
7574 if !best_score.is_finite() {
7575 crate::bail_invalid_estim!(
7576 "isotropic Matérn basin selection received a non-finite incumbent profile"
7577 );
7578 }
7579
7580 for &term_idx in spatial_terms {
7581 let Some(SmoothBasisSpec::Matern {
7582 feature_cols,
7583 spec: matern,
7584 ..
7585 }) = resolvedspec
7586 .smooth_terms
7587 .get(term_idx)
7588 .map(|term| &term.basis)
7589 else {
7590 continue;
7591 };
7592 let num_centers =
7593 gam_terms::basis::center_strategy_num_centers(&matern.center_strategy).ok_or_else(
7594 || {
7595 EstimationError::InvalidInput(format!(
7596 "resolved isotropic Matérn term {term_idx} has no finite center count"
7597 ))
7598 },
7599 )?;
7600 let companion_length_scale = matern_low_rank_center_resolution_length_scale(
7601 data,
7602 feature_cols,
7603 num_centers,
7604 )
7605 .ok_or_else(|| {
7606 EstimationError::InvalidInput(format!(
7607 "resolved isotropic Matérn term {term_idx} has no finite center-resolution range"
7608 ))
7609 })?;
7610 let (psi_long_bound, psi_short_bound) =
7611 spatial_term_psi_bounds(data, &resolvedspec, term_idx, kappa_options);
7612 let psi_long = (-companion_length_scale.ln()).clamp(psi_long_bound, psi_short_bound);
7613 let long_length_scale = (-psi_long).exp();
7614 if !(long_length_scale.is_finite() && long_length_scale > 0.0) {
7615 crate::bail_invalid_estim!(
7616 "isotropic Matérn term {term_idx} produced an invalid long-range endpoint from psi={psi_long}"
7617 );
7618 }
7619 if get_spatial_length_scale(&resolvedspec, term_idx)
7620 .is_some_and(|current| current == long_length_scale)
7621 {
7622 continue;
7623 }
7624
7625 let mut endpoint_spec = resolvedspec.clone();
7626 set_spatial_length_scale(&mut endpoint_spec, term_idx, long_length_scale)?;
7627 let endpoint = fit_term_collection_forspecwith_heuristic_lambdas(
7636 data,
7637 y,
7638 weights,
7639 offset,
7640 &endpoint_spec,
7641 best.fit.lambdas.as_slice(),
7642 family.clone(),
7643 options,
7644 )?;
7645 let endpoint_score = fit_score(&endpoint.fit);
7646 if !endpoint_score.is_finite() {
7647 crate::bail_invalid_estim!(
7648 "isotropic Matérn term {term_idx} long-range endpoint returned a non-finite profiled REML score"
7649 );
7650 }
7651
7652 if endpoint_score < best_score {
7653 log::info!(
7654 "[spatial-kappa] term {term_idx} selected certified long-range basin: \
7655 length_scale={long_length_scale:.6}, profiled REML {endpoint_score:.6} \
7656 < short-basin {best_score:.6}"
7657 );
7658 resolvedspec = freeze_term_collection_from_design(&endpoint_spec, &endpoint.design)?;
7659 best = endpoint;
7660 best_score = endpoint_score;
7661 } else {
7662 log::info!(
7663 "[spatial-kappa] term {term_idx} retained certified short-range basin: \
7664 profiled REML {best_score:.6} <= long-endpoint {endpoint_score:.6} \
7665 at length_scale={long_length_scale:.6}"
7666 );
7667 }
7668 }
7669
7670 Ok((resolvedspec, best))
7671}
7672
7673pub fn fit_term_collectionwith_spatial_length_scale_optimization(
7674 data: ArrayView2<'_, f64>,
7675 y: Array1<f64>,
7676 weights: Array1<f64>,
7677 offset: Array1<f64>,
7678 spec: &TermCollectionSpec,
7679 family: LikelihoodSpec,
7680 options: &FitOptions,
7681 kappa_options: &SpatialLengthScaleOptimizationOptions,
7682) -> Result<FittedTermCollectionWithSpec, EstimationError> {
7683 let mut resolvedspec = spec.clone();
7699 let spatial_terms = spatial_length_scale_term_indices(&resolvedspec);
7700 let n = data.nrows();
7701 if !(y.len() == n && weights.len() == n && offset.len() == n) {
7702 crate::bail_invalid_estim!(
7703 "fit_term_collectionwith_spatial_length_scale_optimization row mismatch: n={}, y={}, weights={}, offset={}",
7704 n,
7705 y.len(),
7706 weights.len(),
7707 offset.len()
7708 );
7709 }
7710 if !kappa_options.enabled || spatial_terms.is_empty() {
7711 let out = fit_term_collection_forspec(
7712 data,
7713 y.view(),
7714 weights.view(),
7715 offset.view(),
7716 &resolvedspec,
7717 family,
7718 options,
7719 )?;
7720 let resolvedspec = freeze_term_collection_from_design(&resolvedspec, &out.design)?;
7721 return Ok(FittedTermCollectionWithSpec {
7722 fit: out.fit,
7723 design: out.design,
7724 resolvedspec,
7725 adaptive_diagnostics: out.adaptive_diagnostics,
7726 kappa_timing: None,
7727 });
7728 }
7729 if kappa_options.max_outer_iter == 0 {
7730 crate::bail_invalid_estim!("spatial kappa optimization requires max_outer_iter >= 1");
7731 }
7732 if !(kappa_options.log_step.is_finite() && kappa_options.log_step > 0.0) {
7733 crate::bail_invalid_estim!("spatial kappa optimization requires log_step > 0");
7734 }
7735 if !(kappa_options.min_length_scale.is_finite()
7736 && kappa_options.max_length_scale.is_finite()
7737 && kappa_options.min_length_scale > 0.0
7738 && kappa_options.max_length_scale >= kappa_options.min_length_scale)
7739 {
7740 crate::bail_invalid_estim!(
7741 "spatial kappa optimization requires valid positive length_scale bounds"
7742 );
7743 }
7744
7745 let pilot_threshold = kappa_options.pilot_subsample_threshold;
7746 if pilot_threshold > 0 && n > pilot_threshold * 2 {
7747 log::info!(
7748 "[spatial-kappa] n={n} exceeds pilot threshold {}; using pilot geometry only for deterministic anisotropy initialization",
7749 pilot_threshold * 2,
7750 );
7751 apply_spatial_anisotropy_pilot_initializer(
7752 data,
7753 &mut resolvedspec,
7754 &spatial_terms,
7755 pilot_threshold,
7756 kappa_options,
7757 );
7758 }
7759
7760 apply_response_aware_anisotropy_seed(data, y.view(), &mut resolvedspec, &spatial_terms);
7769
7770 let free_curvature_terms: Vec<usize> = constant_curvature_term_indices(&resolvedspec)
7774 .into_iter()
7775 .filter(|&term_idx| !constant_curvature_kappa_is_fixed(&resolvedspec, term_idx))
7776 .collect();
7777 if !free_curvature_terms.is_empty() {
7778 validate_constant_curvature_fair_profile_inputs(weights.view(), offset.view(), &family)?;
7779 }
7780 for term_idx in free_curvature_terms {
7781 let kappa_hat = constant_curvature_kappa_fair_optimum(
7782 data,
7783 y.view(),
7784 &resolvedspec,
7785 term_idx,
7786 options,
7787 )?;
7788 if let Some(SmoothBasisSpec::ConstantCurvature { spec: cc, .. }) = resolvedspec
7789 .smooth_terms
7790 .get_mut(term_idx)
7791 .map(|term| &mut term.basis)
7792 {
7793 cc.kappa = kappa_hat;
7794 }
7795 }
7796
7797 let baseline_options = superseded_fit_options(options);
7798 let best = fit_term_collection_forspec(
7799 data,
7800 y.view(),
7801 weights.view(),
7802 offset.view(),
7803 &resolvedspec,
7804 family.clone(),
7805 &baseline_options,
7806 )?;
7807 resolvedspec = freeze_term_collection_from_design(&resolvedspec, &best.design)?;
7808 let spatial_terms = spatial_length_scale_term_indices(&resolvedspec);
7818 let (next_spec, best) = select_isotropic_matern_range_basin(
7819 data,
7820 y.view(),
7821 weights.view(),
7822 offset.view(),
7823 resolvedspec,
7824 best,
7825 &family,
7826 &baseline_options,
7827 kappa_options,
7828 &spatial_terms,
7829 )?;
7830 resolvedspec = next_spec;
7831 sync_aniso_contrasts_from_metadata(&mut resolvedspec, &best.design.smooth);
7835 if spatial_terms.is_empty() {
7836 let fitted = fit_term_collection_forspecwith_heuristic_lambdas(
7837 data,
7838 y.view(),
7839 weights.view(),
7840 offset.view(),
7841 &resolvedspec,
7842 best.fit.lambdas.as_slice(),
7843 family,
7844 options,
7845 )?;
7846 return Ok(FittedTermCollectionWithSpec {
7847 fit: fitted.fit,
7848 design: fitted.design,
7849 resolvedspec,
7850 adaptive_diagnostics: fitted.adaptive_diagnostics,
7851 kappa_timing: None,
7852 });
7853 }
7854 let initial_score = fit_score(&best.fit);
7855 if !initial_score.is_finite() {
7856 crate::bail_invalid_estim!(
7857 "spatial kappa optimization received a non-finite initial profiled score"
7858 );
7859 }
7860 let exact_joint = try_exact_joint_spatial_length_scale_optimization(
7861 data,
7862 y.view(),
7863 weights.view(),
7864 offset.view(),
7865 &resolvedspec,
7866 &best,
7867 family.clone(),
7868 options,
7869 kappa_options,
7870 &spatial_terms,
7871 )?
7872 .ok_or_else(|| {
7873 EstimationError::RemlOptimizationFailed(
7874 "spatial kappa optimization is unavailable for one or more eligible spatial terms"
7875 .to_string(),
7876 )
7877 })?;
7878 let exact_score = fit_score(&exact_joint.fit);
7879 let exact_joint = require_successful_spatial_optimization_result(
7880 initial_score,
7881 Ok(Some((exact_joint, exact_score))),
7882 )?;
7883
7884 log_spatial_aniso_scales(&exact_joint.resolvedspec);
7885 Ok(exact_joint)
7886}
7887
7888#[derive(Clone, Debug)]
7894pub struct CurvatureInference {
7895 pub term_idx: usize,
7897 pub kappa_hat: f64,
7900 pub ci: gam_geometry::curvature_estimand::KappaProfileCi,
7902 pub flatness: gam_geometry::curvature_estimand::FlatnessTest,
7906}
7907
7908fn curvature_profile_lr_endpoint<F>(
7920 profile: &mut F,
7921 kappa_hat: f64,
7922 value_hat: f64,
7923 bound: f64,
7924 half_threshold: f64,
7925 x_tolerance: f64,
7926 score_tolerance: f64,
7927) -> Result<(f64, bool), String>
7928where
7929 F: FnMut(f64) -> Result<(f64, f64), String>,
7930{
7931 let direction = (bound - kappa_hat).signum();
7932 let span = (bound - kappa_hat).abs();
7933 if direction == 0.0 || span <= x_tolerance {
7934 return Ok((bound, true));
7935 }
7936
7937 let (bound_value, bound_score) = profile(bound)?;
7938 let outward_score = direction * bound_score;
7939 if outward_score < -score_tolerance {
7940 return Err(format!(
7941 "curvature profile is not outward-monotone at chart bound {bound}: \
7942 outward score {outward_score:.6e} is below tolerance {score_tolerance:.6e}"
7943 ));
7944 }
7945 let value_tolerance = score_tolerance * span;
7946 if bound_value < value_hat - value_tolerance {
7947 return Err(format!(
7948 "fitted curvature is not the minimum of its inference profile: \
7949 V(bound={bound})={bound_value:.6e} < V(kappa_hat)={value_hat:.6e}"
7950 ));
7951 }
7952 let bound_residual = bound_value - value_hat - half_threshold;
7953 if bound_residual < 0.0 {
7954 return Ok((bound, true));
7955 }
7956 if bound_residual == 0.0 {
7957 return Ok((bound, false));
7958 }
7959
7960 let mut inside_x = kappa_hat;
7965 let mut outside_x = bound;
7966 let mut outside_residual = bound_residual;
7967 let mut outside_score = bound_score;
7968 while (outside_x - inside_x).abs() > x_tolerance {
7969 let lo = inside_x.min(outside_x);
7970 let hi = inside_x.max(outside_x);
7971 let width = hi - lo;
7972 let central_lo = lo + 0.25 * width;
7973 let central_hi = hi - 0.25 * width;
7974 let newton = outside_x - outside_residual / outside_score;
7975 let probe = if newton.is_finite() && newton > central_lo && newton < central_hi {
7976 newton
7977 } else {
7978 lo + 0.5 * width
7979 };
7980 if !(probe > lo && probe < hi) {
7981 break;
7982 }
7983 let (value, score) = profile(probe)?;
7984 let outward_score = direction * score;
7985 if outward_score < -score_tolerance {
7986 return Err(format!(
7987 "curvature profile changed direction before its likelihood crossing at \
7988 kappa={probe}: outward score {outward_score:.6e} is below tolerance \
7989 {score_tolerance:.6e}"
7990 ));
7991 }
7992 let residual = value - value_hat - half_threshold;
7993 if residual >= 0.0 {
7994 outside_x = probe;
7995 outside_residual = residual;
7996 outside_score = score;
7997 } else {
7998 inside_x = probe;
7999 }
8000 }
8001 Ok((inside_x + 0.5 * (outside_x - inside_x), false))
8002}
8003
8004fn curvature_profile_ci_from_analytic_score<F>(
8005 profile: &mut F,
8006 kappa_hat: f64,
8007 kappa_min: f64,
8008 kappa_max: f64,
8009 level: f64,
8010 relative_tolerance: f64,
8011) -> Result<gam_geometry::curvature_estimand::KappaProfileCi, String>
8012where
8013 F: FnMut(f64) -> Result<(f64, f64), String>,
8014{
8015 if !(kappa_min < kappa_max && kappa_hat >= kappa_min && kappa_hat <= kappa_max) {
8016 return Err("curvature profile requires kappa_hat inside valid chart bounds".to_string());
8017 }
8018 if !(level > 0.0 && level < 1.0) {
8019 return Err("curvature profile level must lie in (0, 1)".to_string());
8020 }
8021 let z = gam_geometry::curvature_estimand::wald_half_width(1.0, level)
8022 .ok_or_else(|| "curvature profile threshold is not finite".to_string())?;
8023 let half_threshold = 0.5 * z * z;
8024 let (value_hat, score_hat) = profile(kappa_hat)?;
8025 let relative_tolerance = relative_tolerance.max(f64::EPSILON.sqrt());
8026 let x_tolerance = relative_tolerance * (1.0 + kappa_min.abs().max(kappa_max.abs()));
8027 let score_tolerance = relative_tolerance * (1.0 + value_hat.abs());
8028 let at_lower = (kappa_hat - kappa_min).abs() <= x_tolerance;
8029 let at_upper = (kappa_hat - kappa_max).abs() <= x_tolerance;
8030 let stationary = if at_lower {
8031 score_hat >= -score_tolerance
8032 } else if at_upper {
8033 score_hat <= score_tolerance
8034 } else {
8035 score_hat.abs() <= score_tolerance
8036 };
8037 if !stationary {
8038 return Err(format!(
8039 "curvature inference rejected a non-stationary point estimate: \
8040 kappa_hat={kappa_hat}, score={score_hat:.6e}, \
8041 stationarity_bound={score_tolerance:.6e}"
8042 ));
8043 }
8044
8045 let (ci_lo, lo_at_bound) = curvature_profile_lr_endpoint(
8046 profile,
8047 kappa_hat,
8048 value_hat,
8049 kappa_min,
8050 half_threshold,
8051 x_tolerance,
8052 score_tolerance,
8053 )?;
8054 let (ci_hi, hi_at_bound) = curvature_profile_lr_endpoint(
8055 profile,
8056 kappa_hat,
8057 value_hat,
8058 kappa_max,
8059 half_threshold,
8060 x_tolerance,
8061 score_tolerance,
8062 )?;
8063 let verdict = if ci_lo > 0.0 {
8064 gam_geometry::curvature_estimand::CurvatureVerdict::Spherical
8065 } else if ci_hi < 0.0 {
8066 gam_geometry::curvature_estimand::CurvatureVerdict::Hyperbolic
8067 } else {
8068 gam_geometry::curvature_estimand::CurvatureVerdict::Flat
8069 };
8070 Ok(gam_geometry::curvature_estimand::KappaProfileCi {
8071 kappa_hat,
8072 ci_lo,
8073 ci_hi,
8074 lo_at_bound,
8075 hi_at_bound,
8076 verdict,
8077 })
8078}
8079
8080pub fn curvature_inference_forspec(
8081 data: ArrayView2<'_, f64>,
8082 y: ArrayView1<'_, f64>,
8083 weights: ArrayView1<'_, f64>,
8084 offset: ArrayView1<'_, f64>,
8085 resolvedspec: &TermCollectionSpec,
8086 term_idx: usize,
8087 family: LikelihoodSpec,
8088 options: &FitOptions,
8089 level: f64,
8090) -> Result<CurvatureInference, EstimationError> {
8091 let kappa_hat = get_constant_curvature_kappa(resolvedspec, term_idx).ok_or_else(|| {
8092 EstimationError::InvalidInput(format!(
8093 "curvature_inference_forspec: term {term_idx} is not a constant-curvature smooth"
8094 ))
8095 })?;
8096 if constant_curvature_kappa_is_fixed(resolvedspec, term_idx) {
8097 crate::bail_invalid_estim!(
8098 "curvature inference requires an estimated curvature; term {term_idx} has user-pinned kappa={kappa_hat}"
8099 );
8100 }
8101 if y.len() != data.nrows() || weights.len() != data.nrows() || offset.len() != data.nrows() {
8102 crate::bail_invalid_estim!(
8103 "curvature inference row mismatch: data={}, y={}, weights={}, offset={}",
8104 data.nrows(),
8105 y.len(),
8106 weights.len(),
8107 offset.len(),
8108 );
8109 }
8110 validate_constant_curvature_fair_profile_inputs(weights, offset, &family)?;
8111 let (kappa_min, kappa_max) = constant_curvature_kappa_bounds(data, resolvedspec, term_idx);
8112 let (feature_cols, base_spec) = match resolvedspec
8113 .smooth_terms
8114 .get(term_idx)
8115 .map(|term| &term.basis)
8116 {
8117 Some(SmoothBasisSpec::ConstantCurvature {
8118 feature_cols, spec, ..
8119 }) => (feature_cols, spec.clone()),
8120 _ => {
8121 return Err(EstimationError::InvalidInput(format!(
8122 "constant-curvature κ profile: smooth term {term_idx} is not a \
8123 constant-curvature basis"
8124 )));
8125 }
8126 };
8127 let x_term = select_columns(data, feature_cols).map_err(EstimationError::from)?;
8128 let radial_reference = constant_curvature_radial_reference(x_term.view(), y)?;
8129 let fair_profile = ConstantCurvatureFairProfile {
8130 data: x_term.view(),
8131 response: y,
8132 radial_reference,
8133 spec: base_spec,
8134 cache: std::cell::RefCell::new(std::collections::HashMap::new()),
8135 };
8136
8137 let mut v_p = |kappa: f64| -> Result<(f64, f64), String> {
8140 if !kappa.is_finite() {
8141 return Err(format!("V_p probed a non-finite κ = {kappa}"));
8142 }
8143 let sample = fair_profile.evaluate(kappa).map_err(|error| {
8144 format!("analytic curvature profile at kappa={kappa} failed: {error}")
8145 })?;
8146 Ok(sample)
8147 };
8148 let ci = curvature_profile_ci_from_analytic_score(
8149 &mut v_p,
8150 kappa_hat,
8151 kappa_min,
8152 kappa_max,
8153 level,
8154 options.tol,
8155 )
8156 .map_err(EstimationError::RemlOptimizationFailed)?;
8157 let flatness = gam_geometry::curvature_estimand::flatness_lr_test(
8158 |kappa| v_p(kappa).map(|(value, _)| value),
8159 kappa_hat,
8160 )
8161 .map_err(EstimationError::RemlOptimizationFailed)?;
8162
8163 Ok(CurvatureInference {
8164 term_idx,
8165 kappa_hat,
8166 ci,
8167 flatness,
8168 })
8169}
8170
8171#[cfg(test)]
8172mod curvature_profile_score_tests {
8173 use super::*;
8174
8175 #[test]
8176 fn analytic_profile_score_finds_exact_quadratic_lr_crossings() {
8177 let kappa_hat = -0.37;
8178 let curvature = 16.0;
8179 let level = 0.95;
8180 let mut profile = |kappa: f64| -> Result<(f64, f64), String> {
8181 let displacement = kappa - kappa_hat;
8182 Ok((
8183 7.0 + 0.5 * curvature * displacement * displacement,
8184 curvature * displacement,
8185 ))
8186 };
8187 let ci = curvature_profile_ci_from_analytic_score(
8188 &mut profile,
8189 kappa_hat,
8190 -3.0,
8191 3.0,
8192 level,
8193 1.0e-10,
8194 )
8195 .expect("analytic quadratic profile CI");
8196 let z = gam_geometry::curvature_estimand::wald_half_width(1.0, level)
8197 .expect("valid normal quantile");
8198 let expected_half_width = z / curvature.sqrt();
8199 assert!((ci.ci_lo - (kappa_hat - expected_half_width)).abs() <= 1.0e-8);
8200 assert!((ci.ci_hi - (kappa_hat + expected_half_width)).abs() <= 1.0e-8);
8201 assert!(!ci.lo_at_bound && !ci.hi_at_bound);
8202 }
8203
8204 #[test]
8205 fn analytic_profile_marks_chart_bound_when_wilks_set_never_crosses() {
8206 let mut profile =
8207 |kappa: f64| -> Result<(f64, f64), String> { Ok((0.5 * kappa * kappa, kappa)) };
8208 let ci =
8209 curvature_profile_ci_from_analytic_score(&mut profile, 0.0, -0.1, 0.1, 0.95, 1.0e-10)
8210 .expect("open bounded profile CI");
8211 assert_eq!(ci.ci_lo, -0.1);
8212 assert_eq!(ci.ci_hi, 0.1);
8213 assert!(ci.lo_at_bound && ci.hi_at_bound);
8214 }
8215}
8216
8217#[derive(Clone, Copy, Debug, PartialEq, Eq)]
8220pub enum SmoothLrCorrection {
8221 LawleyLrEstimatedLambda,
8225 LawleyLrFixedLambda,
8230 None,
8234}
8235
8236impl SmoothLrCorrection {
8237 pub fn label(self) -> &'static str {
8239 match self {
8240 SmoothLrCorrection::LawleyLrEstimatedLambda => "lawley_lr_estimated_lambda",
8241 SmoothLrCorrection::LawleyLrFixedLambda => "lawley_lr_fixed_lambda",
8242 SmoothLrCorrection::None => "none",
8243 }
8244 }
8245}
8246
8247#[derive(Clone, Debug)]
8253pub struct SmoothTermLrInference {
8254 pub name: String,
8256 pub term_idx: usize,
8258 pub statistic_lr: f64,
8261 pub ref_df: f64,
8264 pub bartlett_factor: f64,
8267 pub bartlett_factor_conditional: Option<f64>,
8271 pub rho_variation_shift: Option<f64>,
8274 pub statistic_corrected: f64,
8276 pub p_value_uncorrected: f64,
8278 pub p_value_corrected: f64,
8281 pub material: bool,
8289 pub correction: SmoothLrCorrection,
8291}
8292
8293pub const SMOOTH_LR_MATERIAL_THRESHOLD: f64 = 0.10;
8297
8298fn fitted_rho_penalty_components(
8304 penalties: &[BlockwisePenalty],
8305 lambdas: &[f64],
8306 p_total: usize,
8307) -> Result<Vec<gam_terms::inference::lawley::RhoPenaltyComponent>, EstimationError> {
8308 if penalties.len() != lambdas.len() {
8309 return Err(EstimationError::InvalidInput(format!(
8310 "smooth_term_lr_inference: penalty/lambda count mismatch ({} penalties, {} lambdas)",
8311 penalties.len(),
8312 lambdas.len()
8313 )));
8314 }
8315 let mut components = Vec::with_capacity(penalties.len());
8316 for (idx, (penalty, &lambda)) in penalties.iter().zip(lambdas.iter()).enumerate() {
8317 if !(lambda.is_finite() && lambda >= 0.0) {
8318 return Err(EstimationError::InvalidInput(format!(
8319 "smooth_term_lr_inference: lambda[{idx}] is invalid: {lambda}"
8320 )));
8321 }
8322 let r = &penalty.col_range;
8323 if r.end > p_total {
8324 return Err(EstimationError::InvalidInput(format!(
8325 "smooth_term_lr_inference: penalty[{idx}] range {:?} exceeds coefficient dimension {p_total}",
8326 r
8327 )));
8328 }
8329 let mut s_component = Array2::<f64>::zeros((p_total, p_total));
8330 s_component
8331 .slice_mut(s![r.start..r.end, r.start..r.end])
8332 .scaled_add(lambda, &penalty.local);
8333 components.push(gam_terms::inference::lawley::RhoPenaltyComponent { s_component });
8334 }
8335 Ok(components)
8336}
8337
8338pub fn smooth_term_lr_inference_forspec(
8383 data: ArrayView2<'_, f64>,
8384 y: ArrayView1<'_, f64>,
8385 weights: ArrayView1<'_, f64>,
8386 offset: ArrayView1<'_, f64>,
8387 resolvedspec: &TermCollectionSpec,
8388 family: LikelihoodSpec,
8389 options: &FitOptions,
8390) -> Result<Vec<SmoothTermLrInference>, EstimationError> {
8391 use gam_terms::inference::lawley::{
8392 LAWLEY_PAIR_MATRIX_MAX_ROWS, known_scale_expected_jets_with_dispersion,
8393 lawley_lr_bartlett_factor, lawley_lr_mean_shift_with_rho_variation,
8394 };
8395
8396 let n = data.nrows();
8397 let full = fit_term_collection_forspec(
8400 data,
8401 y,
8402 weights,
8403 offset,
8404 resolvedspec,
8405 family.clone(),
8406 options,
8407 )?;
8408 let ll_full = full.fit.log_likelihood;
8409 let p_total = full.design.design.ncols();
8410 let lambdas = full.fit.lambdas.as_slice().ok_or_else(|| {
8411 EstimationError::InvalidInput(
8412 "smooth_term_lr_inference: non-contiguous lambda vector".to_string(),
8413 )
8414 })?;
8415 let s_lambda = weighted_blockwise_penalty_sum(&full.design.penalties, lambdas, p_total);
8416 let rho_penalty_components =
8417 fitted_rho_penalty_components(&full.design.penalties, lambdas, p_total)?;
8418 let rho_covariance = full.fit.artifacts.rho_covariance.as_ref().filter(|cov| {
8419 cov.nrows() == rho_penalty_components.len() && cov.ncols() == rho_penalty_components.len()
8420 });
8421 let full_design_dense = full.design.design.to_dense();
8423 let influence = full.fit.coefficient_influence();
8424 let family_disp = lawley_dispersion_for_family(&family, &full.fit);
8425
8426 let mut penalty_cursor = full.design.leading_penalty_blocks_before_smooth();
8430 let mut out = Vec::<SmoothTermLrInference>::new();
8431 for (term_idx, design_term) in full.design.smooth.terms.iter().enumerate() {
8432 let k = design_term.penalties_local.len();
8433 let block_start = penalty_cursor;
8434 penalty_cursor += k;
8435 if design_term.shape != ShapeConstraint::None {
8438 continue;
8439 }
8440 let coeff_range = design_term.coeff_range.clone();
8441 if coeff_range.start >= coeff_range.end || coeff_range.end > p_total {
8442 continue;
8443 }
8444 let edf = full.fit.per_term_edf(coeff_range.clone(), block_start, k);
8456 let null_dim = design_term.wald_unpenalized_dim();
8476 let rho_uncertainty_df = wps_block_uncertainty_df(
8497 full.fit.weighted_gram(),
8498 full.fit.smoothing_correction(),
8499 &coeff_range,
8500 family_disp,
8501 );
8502 let ref_df = (wood_reference_df(influence, &coeff_range)
8503 .unwrap_or(0.0)
8504 .max(edf)
8505 + rho_uncertainty_df)
8506 .max(null_dim as f64)
8507 .max(1.0);
8508 if !(ref_df.is_finite() && ref_df > 0.0) {
8509 continue;
8510 }
8511
8512 let mut null_spec = resolvedspec.clone();
8515 let Some(spec_pos) = null_spec
8516 .smooth_terms
8517 .iter()
8518 .position(|t| t.name == design_term.name)
8519 else {
8520 continue;
8521 };
8522 null_spec.smooth_terms.remove(spec_pos);
8523 let null_fit = fit_term_collection_forspec(
8524 data,
8525 y,
8526 weights,
8527 offset,
8528 &null_spec,
8529 family.clone(),
8530 options,
8531 );
8532 let (statistic_lr, eta_null) = match null_fit {
8533 Ok(null) if null.fit.log_likelihood.is_finite() => {
8534 let w = (2.0 * (ll_full - null.fit.log_likelihood)).max(0.0);
8535 let mut eta = null.design.design.dot(&null.fit.beta);
8539 eta += &offset;
8540 (w, Some(eta))
8541 }
8542 _ => (f64::NAN, None),
8543 };
8544
8545 let chi2 = statrs::distribution::ChiSquared::new(ref_df).ok();
8546 let p_uncorrected = match (chi2.as_ref(), statistic_lr.is_finite()) {
8547 (Some(dist), true) => {
8548 use statrs::distribution::ContinuousCDF;
8549 (1.0 - dist.cdf(statistic_lr)).clamp(0.0, 1.0)
8550 }
8551 _ => f64::NAN,
8552 };
8553
8554 let mut bartlett_factor = 1.0;
8558 let mut bartlett_factor_conditional = None;
8559 let mut rho_variation_shift = None;
8560 let mut statistic_corrected = statistic_lr;
8561 let mut p_corrected = p_uncorrected;
8562 let mut correction = SmoothLrCorrection::None;
8563 if let (Some(eta), true, true) = (
8564 eta_null.as_ref(),
8565 statistic_lr.is_finite(),
8566 n <= LAWLEY_PAIR_MATRIX_MAX_ROWS,
8567 ) {
8568 let kappas: Option<Vec<_>> = (0..n)
8569 .map(|i| {
8570 known_scale_expected_jets_with_dispersion(&family, eta[i], family_disp)
8571 .and_then(|jets| jets.kappas().ok())
8572 })
8573 .collect();
8574 if let (Some(kappas), Some(dist)) = (kappas, chi2.as_ref()) {
8575 let fixed_factor = lawley_lr_bartlett_factor(
8576 full_design_dense.view(),
8577 &kappas,
8578 Some(s_lambda.view()),
8579 coeff_range.clone(),
8580 ref_df,
8581 );
8582 if let Ok(c_cond) = fixed_factor
8583 && c_cond.is_finite()
8584 && c_cond > 0.0
8585 {
8586 let mut c_applied = c_cond;
8587 correction = SmoothLrCorrection::LawleyLrFixedLambda;
8588 if let Some(cov) = rho_covariance
8589 && let Ok(total_shift) = lawley_lr_mean_shift_with_rho_variation(
8590 full_design_dense.view(),
8591 &kappas,
8592 s_lambda.view(),
8593 coeff_range.clone(),
8594 &rho_penalty_components,
8595 cov.view(),
8596 )
8597 {
8598 let mean_w = ref_df + total_shift;
8599 if let Some(c_est) =
8600 gam_terms::inference::higher_order::bartlett_factor_from_mean(
8601 mean_w, ref_df,
8602 )
8603 && c_est.is_finite()
8604 && c_est > 0.0
8605 {
8606 let conditional_shift = (c_cond - 1.0) * ref_df;
8607 c_applied = c_est;
8608 bartlett_factor_conditional = Some(c_cond);
8609 rho_variation_shift = Some(total_shift - conditional_shift);
8610 correction = SmoothLrCorrection::LawleyLrEstimatedLambda;
8611 }
8612 }
8613 use statrs::distribution::ContinuousCDF;
8614 bartlett_factor = c_applied;
8615 statistic_corrected = statistic_lr / c_applied;
8616 p_corrected = (1.0 - dist.cdf(statistic_corrected)).clamp(0.0, 1.0);
8617 }
8618 }
8619 }
8620
8621 let material = match correction {
8627 SmoothLrCorrection::LawleyLrEstimatedLambda
8628 | SmoothLrCorrection::LawleyLrFixedLambda => {
8629 let factor_move = (bartlett_factor - 1.0).abs();
8630 let p_denom = p_uncorrected.max(p_corrected).max(f64::MIN_POSITIVE);
8631 let p_move = if p_uncorrected.is_finite() && p_corrected.is_finite() {
8632 (p_corrected - p_uncorrected).abs() / p_denom
8633 } else {
8634 0.0
8635 };
8636 factor_move > SMOOTH_LR_MATERIAL_THRESHOLD || p_move > SMOOTH_LR_MATERIAL_THRESHOLD
8637 }
8638 SmoothLrCorrection::None => false,
8639 };
8640
8641 out.push(SmoothTermLrInference {
8642 name: design_term.name.clone(),
8643 term_idx,
8644 statistic_lr,
8645 ref_df,
8646 bartlett_factor,
8647 bartlett_factor_conditional,
8648 rho_variation_shift,
8649 statistic_corrected,
8650 p_value_uncorrected: p_uncorrected,
8651 p_value_corrected: p_corrected,
8652 material,
8653 correction,
8654 });
8655 }
8656 Ok(out)
8657}
8658
8659fn lawley_dispersion_for_family(family: &LikelihoodSpec, fit: &UnifiedFitResult) -> f64 {
8662 match family.response {
8663 gam_spec::ResponseFamily::Gaussian => {
8664 let sd = fit.standard_deviation;
8665 (sd * sd).max(f64::MIN_POSITIVE)
8666 }
8667 gam_spec::ResponseFamily::Gamma => {
8668 let shape = fit.standard_deviation;
8669 if shape.is_finite() && shape > 0.0 {
8670 1.0 / shape
8671 } else {
8672 1.0
8673 }
8674 }
8675 _ => 1.0,
8676 }
8677}
8678
8679fn wps_block_uncertainty_df(
8680 weighted_gram: Option<&Array2<f64>>,
8681 smoothing_correction: Option<&Array2<f64>>,
8682 coeff_range: &Range<usize>,
8683 phi: f64,
8684) -> f64 {
8685 let (Some(xwx), Some(corr)) = (weighted_gram, smoothing_correction) else {
8686 return 0.0;
8687 };
8688 let (start, end) = (coeff_range.start, coeff_range.end);
8689 if start >= end
8690 || end > xwx.nrows()
8691 || end > xwx.ncols()
8692 || end > corr.nrows()
8693 || end > corr.ncols()
8694 || !(phi.is_finite() && phi > 0.0)
8695 {
8696 return 0.0;
8697 }
8698
8699 let mut trace = 0.0;
8700 for i in start..end {
8701 for j in start..end {
8702 trace += xwx[[i, j]] * corr[[j, i]];
8703 }
8704 }
8705 trace /= phi;
8706 if trace.is_finite() && trace > 0.0 {
8707 trace
8708 } else {
8709 0.0
8710 }
8711}
8712
8713fn wood_reference_df(influence: Option<&Array2<f64>>, coeff_range: &Range<usize>) -> Option<f64> {
8737 let f = influence?;
8738 let (start, end) = (coeff_range.start, coeff_range.end);
8739 if start >= end || end > f.nrows() || end > f.ncols() {
8740 return None;
8741 }
8742 let block = f.slice(s![start..end, start..end]);
8743 let tr = (0..block.nrows()).map(|i| block[[i, i]]).sum::<f64>();
8744 let tr2 = block.dot(&block).diag().sum();
8745 (tr.is_finite() && tr2.is_finite() && tr > 0.0).then(|| (2.0 * tr - tr2).max(tr).max(1e-12))
8746}
8747
8748#[cfg(test)]
8749mod nfree_gate_tests {
8750 use super::nfree_skip_gate_status_from_parts;
8751
8752 #[test]
8753 fn value_only_nfree_gate_does_not_require_basis_skip_witness() {
8754 let gate = nfree_skip_gate_status_from_parts(
8755 true, true, false, false, true, true, false, false, );
8764 assert!(
8765 gate.would_skip(false),
8766 "value-only κ cost probes must stay n-free when the Gram value is certified; \
8767 the reduced-basis skip witness is required only for beta/gradient probes"
8768 );
8769 }
8770
8771 #[test]
8772 fn gradient_nfree_gate_still_requires_basis_skip_witness() {
8773 let gate =
8774 nfree_skip_gate_status_from_parts(true, true, false, true, true, true, false, true);
8775 assert!(
8776 !gate.would_skip(true),
8777 "gradient probes return beta/gradient objects in a reduced basis and must not \
8778 skip the row lane without the reduced-basis witness"
8779 );
8780 }
8781}