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(),
445 p_total,
446 )));
447 for _ in 0..analytic_rho_count {
448 hyper_dirs.push(
449 DirectionalHyperParam::new_compact(zero_x.clone(), Vec::new(), None, None)?
450 .not_penalty_like(),
451 );
452 }
453 for _ in 0..direct_dim {
454 hyper_dirs.push(
455 DirectionalHyperParam::new_compact(zero_x.clone(), Vec::new(), None, None)?
456 .not_penalty_like(),
457 );
458 }
459 }
460 Ok(Some(hyper_dirs))
461}
462
463fn latent_coord_direct_hyper_count(
464 id_mode: &gam_terms::latent::LatentIdMode,
465 latent_dim: usize,
466) -> usize {
467 use gam_terms::latent::{AuxPriorStrength, LatentIdMode};
468 match id_mode {
469 LatentIdMode::AuxPrior { strength, .. } => match strength {
470 AuxPriorStrength::Auto => 1,
471 AuxPriorStrength::Fixed(_) => 0,
472 },
473 LatentIdMode::AuxPriorDimSelection { strength, .. } => {
474 latent_dim
475 + match strength {
476 AuxPriorStrength::Auto => 1,
477 AuxPriorStrength::Fixed(_) => 0,
478 }
479 }
480 LatentIdMode::DimSelection { .. } => latent_dim,
481 LatentIdMode::IsometryToReference { strength, .. } => match strength {
484 AuxPriorStrength::Auto => 1,
485 AuxPriorStrength::Fixed(_) => 0,
486 },
487 LatentIdMode::AuxOutcome { head, .. } => head.n_coeffs(latent_dim) + latent_dim,
490 LatentIdMode::None => 0,
491 }
492}
493
494fn latent_coord_initial_direct_hypers(
495 id_mode: &gam_terms::latent::LatentIdMode,
496 latent_dim: usize,
497) -> Result<Array1<f64>, EstimationError> {
498 use gam_terms::latent::{AuxPriorStrength, LatentIdMode};
499 let mut values = Vec::with_capacity(latent_coord_direct_hyper_count(id_mode, latent_dim));
500 match id_mode {
501 LatentIdMode::AuxPrior { strength, .. } => {
502 if matches!(strength, AuxPriorStrength::Auto) {
503 values.push(0.0);
504 }
505 }
506 LatentIdMode::AuxPriorDimSelection {
507 strength,
508 init_log_precision,
509 ..
510 } => {
511 if matches!(strength, AuxPriorStrength::Auto) {
512 values.push(0.0);
513 }
514 append_latent_ard_seed(&mut values, init_log_precision.as_ref(), latent_dim)?;
515 }
516 LatentIdMode::DimSelection { init_log_precision } => {
517 append_latent_ard_seed(&mut values, init_log_precision.as_ref(), latent_dim)?;
518 }
519 LatentIdMode::IsometryToReference { strength, .. } => {
520 if matches!(strength, AuxPriorStrength::Auto) {
521 values.push(0.0);
522 }
523 }
524 LatentIdMode::AuxOutcome {
525 head,
526 init_log_precision,
527 } => {
528 values.extend(std::iter::repeat_n(0.0, head.n_coeffs(latent_dim)));
532 append_latent_ard_seed(&mut values, init_log_precision.as_ref(), latent_dim)?;
533 }
534 LatentIdMode::None => {}
535 }
536 Ok(Array1::from_vec(values))
537}
538
539fn append_latent_ard_seed(
540 values: &mut Vec<f64>,
541 init: Option<&Array1<f64>>,
542 latent_dim: usize,
543) -> Result<(), EstimationError> {
544 if let Some(init) = init {
545 if init.len() != latent_dim {
546 crate::bail_invalid_estim!(
547 "latent dim_selection init_log_precision length mismatch: got {}, expected {}",
548 init.len(),
549 latent_dim
550 );
551 }
552 values.extend(init.iter().copied());
553 } else {
554 values.extend(std::iter::repeat_n(0.0, latent_dim));
555 }
556 Ok(())
557}
558
559struct LatentIdObjectiveContribution {
560 cost: f64,
561 gradient: Array1<f64>,
562}
563
564fn latent_id_objective_contribution(
565 theta: &Array1<f64>,
566 rho_dim: usize,
567 analytic_rho_count: usize,
568 latent: &gam_terms::latent::LatentCoordValues,
569) -> Result<LatentIdObjectiveContribution, EstimationError> {
570 use gam_terms::latent::{AuxPriorStrength, LatentIdMode, aux_prior_targets};
571 let n_obs = latent.n_obs();
572 let latent_dim = latent.latent_dim();
573 let flat_len = latent.len();
574 let mut gradient = Array1::<f64>::zeros(theta.len());
575 let t_start = rho_dim;
576 let direct_start = t_start + flat_len + analytic_rho_count;
577 if theta.len() < direct_start {
578 crate::bail_invalid_estim!(
579 "latent-coordinate theta too short for id objective: got {}, need at least {}",
580 theta.len(),
581 direct_start
582 );
583 }
584 let t = latent.as_matrix();
585 let mut cost = 0.0;
586 let mut cursor = direct_start;
587
588 match latent.id_mode() {
589 LatentIdMode::AuxPrior {
590 u,
591 family,
592 strength,
593 }
594 | LatentIdMode::AuxPriorDimSelection {
595 u,
596 family,
597 strength,
598 ..
599 } => {
600 let (log_mu, mu) = match strength {
601 AuxPriorStrength::Fixed(mu) => (mu.ln(), *mu),
602 AuxPriorStrength::Auto => {
603 let log_mu = theta[cursor];
604 cursor += 1;
605 (log_mu, log_mu.exp())
606 }
607 };
608 let targets = aux_prior_targets(t.view(), u.view(), *family)
609 .map_err(EstimationError::InvalidInput)?;
610 let residual = &t - &targets;
611 let q = residual.iter().map(|v| v * v).sum::<f64>();
612 let k = (n_obs * latent_dim) as f64;
619 cost += 0.5 * mu * q - 0.5 * k * log_mu;
620
621 let projected_residual = aux_prior_targets(residual.view(), u.view(), *family)
622 .map_err(EstimationError::InvalidInput)?;
623 let grad_base = residual - projected_residual;
624 for n in 0..n_obs {
625 for axis in 0..latent_dim {
626 gradient[t_start + n * latent_dim + axis] += mu * grad_base[[n, axis]];
627 }
628 }
629 if matches!(strength, AuxPriorStrength::Auto) {
630 gradient[direct_start] += 0.5 * mu * q - 0.5 * k;
631 }
632 }
633 LatentIdMode::IsometryToReference { reference, strength } => {
634 if reference.dim() != (n_obs, latent_dim) {
641 crate::bail_invalid_estim!(
642 "IsometryToReference reference shape {:?} must equal (n_obs, latent_dim) = ({}, {})",
643 reference.dim(),
644 n_obs,
645 latent_dim
646 );
647 }
648 let mu_slot = cursor;
649 let (log_mu, mu) = match strength {
650 AuxPriorStrength::Fixed(mu) => (mu.ln(), *mu),
651 AuxPriorStrength::Auto => {
652 let log_mu = theta[cursor];
653 cursor += 1;
654 (log_mu, log_mu.exp())
655 }
656 };
657 let residual = &t - reference;
658 let q = residual.iter().map(|v| v * v).sum::<f64>();
659 let k = (n_obs * latent_dim) as f64;
663 cost += 0.5 * mu * q - 0.5 * k * log_mu;
664 for n in 0..n_obs {
665 for axis in 0..latent_dim {
666 gradient[t_start + n * latent_dim + axis] += mu * residual[[n, axis]];
667 }
668 }
669 if matches!(strength, AuxPriorStrength::Auto) {
670 gradient[mu_slot] += 0.5 * mu * q - 0.5 * k;
671 }
672 }
673 LatentIdMode::AuxOutcome { head, .. } => {
674 let n_coeffs = head.n_coeffs(latent_dim);
682 let coeffs = theta
683 .slice(ndarray::s![cursor..cursor + n_coeffs])
684 .to_owned();
685 let (head_nll, grad_coeffs, grad_t) = head
686 .neg_loglik_and_grad(t.view(), coeffs.view())
687 .map_err(EstimationError::InvalidInput)?;
688 cost += head_nll;
689 for (offset, &g) in grad_coeffs.iter().enumerate() {
690 gradient[cursor + offset] += g;
691 }
692 for n in 0..n_obs {
693 for axis in 0..latent_dim {
694 gradient[t_start + n * latent_dim + axis] += grad_t[[n, axis]];
695 }
696 }
697 cursor += n_coeffs;
698 }
699 LatentIdMode::DimSelection { .. } | LatentIdMode::None => {}
700 }
701
702 match latent.id_mode() {
703 LatentIdMode::AuxPriorDimSelection { .. }
704 | LatentIdMode::DimSelection { .. }
705 | LatentIdMode::AuxOutcome { .. } => {
706 for axis in 0..latent_dim {
707 let log_alpha = theta[cursor + axis];
708 let alpha = log_alpha.exp();
709 let mut q_axis = 0.0;
710 for n in 0..n_obs {
711 let flat_idx = n * latent_dim + axis;
712 let value = latent.as_flat()[flat_idx];
713 q_axis += value * value;
714 gradient[t_start + flat_idx] += alpha * value;
715 }
716 cost += 0.5 * alpha * q_axis - 0.5 * n_obs as f64 * log_alpha;
717 gradient[cursor + axis] += 0.5 * alpha * q_axis - 0.5 * n_obs as f64;
718 }
719 cursor += latent_dim;
720 }
721 LatentIdMode::AuxPrior { .. }
722 | LatentIdMode::IsometryToReference { .. }
723 | LatentIdMode::None => {}
724 }
725
726 if cursor != theta.len() {
727 crate::bail_invalid_estim!(
728 "latent-coordinate direct hyperparameter length mismatch: consumed {}, theta len {}",
729 cursor,
730 theta.len()
731 );
732 }
733 Ok(LatentIdObjectiveContribution { cost, gradient })
734}
735
736fn add_latent_id_objective_to_eval(
737 theta: &Array1<f64>,
738 rho_dim: usize,
739 analytic_rho_count: usize,
740 latent: &gam_terms::latent::LatentCoordValues,
741 eval: &mut (
742 f64,
743 Array1<f64>,
744 gam_problem::HessianResult,
745 ),
746) -> Result<(), EstimationError> {
747 let contribution =
748 latent_id_objective_contribution(theta, rho_dim, analytic_rho_count, latent)?;
749 eval.0 += contribution.cost;
750 if eval.1.len() != contribution.gradient.len() {
751 crate::bail_invalid_estim!(
752 "latent-coordinate REML gradient length mismatch: base={}, id={}",
753 eval.1.len(),
754 contribution.gradient.len()
755 );
756 }
757 eval.1 += &contribution.gradient;
758 if eval.2.is_analytic() {
759 eval.2 = gam_problem::HessianResult::Unavailable;
760 }
761 Ok(())
762}
763
764fn analytic_penalty_objective_contribution(
765 theta: &Array1<f64>,
766 rho_dim: usize,
767 latent: &gam_terms::latent::LatentCoordValues,
768 registry: &gam_terms::AnalyticPenaltyRegistry,
769) -> Result<LatentIdObjectiveContribution, EstimationError> {
770 let flat_len = latent.len();
771 let t_start = rho_dim;
772 let t_end = t_start + flat_len;
773 let rho_start = t_end;
774 let rho_end = rho_start + registry.total_rho_count();
775 if theta.len() < rho_end {
776 crate::bail_invalid_estim!(
777 "latent-coordinate theta too short for analytic penalties: got {}, need at least {}",
778 theta.len(),
779 rho_end
780 );
781 }
782 let target_t = theta.slice(s![t_start..t_end]);
783 let rho = theta.slice(s![rho_start..rho_end]);
784 let mut cost = 0.0_f64;
785 let mut gradient = Array1::<f64>::zeros(theta.len());
786 for (penalty, (rho_slice, tier, name)) in registry.penalties.iter().zip(registry.rho_layout()) {
787 let rho_local = rho.slice(s![rho_slice.clone()]);
788 match tier {
789 gam_terms::PenaltyTier::Psi => {
790 cost += penalty.value(target_t.view(), rho_local);
791 let grad = penalty.grad_target(target_t.view(), rho_local);
792 if grad.len() != flat_len {
793 crate::bail_invalid_estim!(
794 "analytic penalty {name:?} gradient length mismatch: got {}, expected {}",
795 grad.len(),
796 flat_len
797 );
798 }
799 for i in 0..flat_len {
800 gradient[t_start + i] += grad[i];
801 }
802 let grad_rho_local = penalty.grad_rho(target_t.view(), rho_local);
803 if grad_rho_local.len() != rho_slice.len() {
804 crate::bail_invalid_estim!(
805 "analytic penalty {name:?} rho-gradient length mismatch: got {}, expected {}",
806 grad_rho_local.len(),
807 rho_slice.len()
808 );
809 }
810 for local_idx in 0..grad_rho_local.len() {
811 gradient[rho_start + rho_slice.start + local_idx] += grad_rho_local[local_idx];
812 }
813 }
814 gam_terms::PenaltyTier::Beta => {}
815 gam_terms::PenaltyTier::Rho => {}
816 }
817 }
818 Ok(LatentIdObjectiveContribution { cost, gradient })
819}
820
821fn add_analytic_penalty_hessian_to_eval(
822 theta: &Array1<f64>,
823 rho_dim: usize,
824 latent: &gam_terms::latent::LatentCoordValues,
825 registry: &gam_terms::AnalyticPenaltyRegistry,
826 eval: &mut (
827 f64,
828 Array1<f64>,
829 gam_problem::HessianResult,
830 ),
831) -> Result<(), EstimationError> {
832 let flat_len = latent.len();
833 let t_start = rho_dim;
834 let t_end = t_start + flat_len;
835 let rho_start = t_end;
836 let rho_end = rho_start + registry.total_rho_count();
837 if theta.len() < rho_end {
838 crate::bail_invalid_estim!(
839 "latent-coordinate theta too short for analytic penalty Hessian: got {}, need at least {}",
840 theta.len(),
841 rho_end
842 );
843 }
844 let gam_problem::HessianResult::Analytic(hessian) = &mut eval.2 else {
845 if eval.2.is_analytic() {
846 eval.2 = gam_problem::HessianResult::Unavailable;
847 }
848 return Ok(());
849 };
850 if hessian.dim() != (theta.len(), theta.len()) {
851 crate::bail_invalid_estim!(
852 "analytic penalty Hessian target shape mismatch: got {}x{}, expected {}x{}",
853 hessian.nrows(),
854 hessian.ncols(),
855 theta.len(),
856 theta.len()
857 );
858 }
859 let target_t = theta.slice(s![t_start..t_end]);
860 let rho = theta.slice(s![rho_start..rho_end]);
861 for (penalty, (rho_slice, tier, _name)) in registry.penalties.iter().zip(registry.rho_layout())
862 {
863 let rho_local = rho.slice(s![rho_slice]);
864 if !matches!(tier, gam_terms::PenaltyTier::Psi) {
865 continue;
866 }
867 if let Some(diag) = penalty.hessian_diag(target_t.view(), rho_local) {
868 if diag.len() != flat_len {
869 crate::bail_invalid_estim!(
870 "analytic penalty Hessian diagonal length mismatch: got {}, expected {}",
871 diag.len(),
872 flat_len
873 );
874 }
875 for i in 0..flat_len {
876 hessian[[t_start + i, t_start + i]] += diag[i];
877 }
878 continue;
879 }
880 let mut probe = Array1::<f64>::zeros(flat_len);
881 for col in 0..flat_len {
882 probe[col] = 1.0;
883 let hv = penalty.hvp(target_t.view(), rho_local, probe.view());
884 if hv.len() != flat_len {
885 crate::bail_invalid_estim!(
886 "analytic penalty Hessian-vector length mismatch: got {}, expected {}",
887 hv.len(),
888 flat_len
889 );
890 }
891 for row in 0..flat_len {
892 hessian[[t_start + row, t_start + col]] += hv[row];
893 }
894 probe[col] = 0.0;
895 }
896 }
897 Ok(())
898}
899
900fn add_analytic_penalty_objective_to_eval(
901 theta: &Array1<f64>,
902 rho_dim: usize,
903 latent: &gam_terms::latent::LatentCoordValues,
904 registry: &gam_terms::AnalyticPenaltyRegistry,
905 eval: &mut (
906 f64,
907 Array1<f64>,
908 gam_problem::HessianResult,
909 ),
910) -> Result<(), EstimationError> {
911 let contribution = analytic_penalty_objective_contribution(theta, rho_dim, latent, registry)?;
912 eval.0 += contribution.cost;
913 if eval.1.len() != contribution.gradient.len() {
914 crate::bail_invalid_estim!(
915 "latent-coordinate REML gradient length mismatch: base={}, analytic_penalty={}",
916 eval.1.len(),
917 contribution.gradient.len()
918 );
919 }
920 eval.1 += &contribution.gradient;
921 add_analytic_penalty_hessian_to_eval(theta, rho_dim, latent, registry, eval)?;
922 Ok(())
923}
924
925fn spatial_log_kappa_hyper_dirs_frominfo_list(
926 info_list: Vec<SpatialPsiDerivative>,
927) -> Result<Vec<DirectionalHyperParam>, EstimationError> {
928 use gam_solve::estimate::reml::ImplicitDerivLevel;
929 use std::collections::HashMap;
930
931 let log_kappa_dim = info_list.len();
932 let group_ids: Vec<Option<usize>> = info_list.iter().map(|e| e.aniso_group_id).collect();
938 let mut group_indices_map: HashMap<usize, Vec<usize>> = HashMap::new();
939 for (idx, gid) in group_ids.iter().enumerate() {
940 if let Some(g) = gid {
941 group_indices_map.entry(*g).or_default().push(idx);
942 }
943 }
944
945 let mut hyper_dirs = Vec::with_capacity(log_kappa_dim);
946 for (i, info) in info_list.into_iter().enumerate() {
947 let SpatialPsiDerivative {
948 penalty_index: _,
949 penalty_indices,
950 global_range,
951 total_p,
952 x_psi_local,
953 s_psi_components_local,
954 x_psi_psi_local,
955 s_psi_psi_components_local,
956 aniso_group_id,
957 aniso_cross_designs,
958 aniso_cross_penalty_provider,
959 implicit_operator,
960 implicit_axis,
961 } = info;
962
963 let mut xsecond = vec![None; log_kappa_dim];
964 xsecond[i] = Some(if let Some(ref op) = implicit_operator {
966 gam_solve::estimate::reml::HyperDesignDerivative::from_implicit(
967 op.clone(),
968 ImplicitDerivLevel::SecondDiag(implicit_axis),
969 global_range.clone(),
970 total_p,
971 )
972 } else {
973 gam_solve::estimate::reml::HyperDesignDerivative::from_embedded(
974 x_psi_psi_local,
975 global_range.clone(),
976 total_p,
977 )
978 });
979 if let Some(cross_designs) = aniso_cross_designs {
981 if let Some(gid) = aniso_group_id {
985 let base = group_indices_map
986 .get(&gid)
987 .and_then(|v| v.first().copied())
988 .unwrap_or(i);
989 for (b_axis, cross_mat) in cross_designs.into_iter() {
990 let j = base + b_axis;
991 if j < log_kappa_dim {
992 xsecond[j] = Some(if let Some(ref op) = implicit_operator {
993 gam_solve::estimate::reml::HyperDesignDerivative::from_implicit(
994 op.clone(),
995 ImplicitDerivLevel::SecondCross(implicit_axis, b_axis),
996 global_range.clone(),
997 total_p,
998 )
999 } else {
1000 gam_solve::estimate::reml::HyperDesignDerivative::from_embedded(
1001 cross_mat,
1002 global_range.clone(),
1003 total_p,
1004 )
1005 });
1006 }
1007 }
1008 }
1009 }
1010 let s_components = penalty_indices
1011 .iter()
1012 .copied()
1013 .zip(s_psi_components_local.into_iter().map(|local| {
1014 gam_solve::estimate::reml::HyperPenaltyDerivative::from_embedded(
1015 local,
1016 global_range.clone(),
1017 total_p,
1018 )
1019 }))
1020 .collect::<Vec<_>>();
1021 let s2_components = penalty_indices
1022 .iter()
1023 .copied()
1024 .zip(s_psi_psi_components_local.into_iter().map(|local| {
1025 gam_solve::estimate::reml::HyperPenaltyDerivative::from_embedded(
1026 local,
1027 global_range.clone(),
1028 total_p,
1029 )
1030 }))
1031 .collect::<Vec<_>>();
1032 let mut ssecond_components = vec![None; log_kappa_dim];
1033 ssecond_components[i] = Some(s2_components);
1034 let mut penaltysecond_partner_indices: Option<Vec<usize>> = None;
1035 let penaltysecond_component_provider =
1036 if let (Some(provider), Some(gid)) = (aniso_cross_penalty_provider, aniso_group_id) {
1037 let group_indices = group_indices_map.get(&gid).cloned().unwrap_or_default();
1038 let axis_in_group =
1039 group_indices
1040 .iter()
1041 .position(|&idx| idx == i)
1042 .ok_or_else(|| {
1043 EstimationError::InvalidInput(format!(
1044 "missing spatial hyper axis {} in anisotropy group {}",
1045 i, gid
1046 ))
1047 })?;
1048 penaltysecond_partner_indices = Some(
1049 group_indices
1050 .iter()
1051 .copied()
1052 .filter(|&idx| idx != i)
1053 .collect(),
1054 );
1055 let penalty_indices_inner = penalty_indices.clone();
1056 let global_range_inner = global_range.clone();
1057 let total_p_inner = total_p;
1058 let group_indices_inner = group_indices;
1059 Some(std::sync::Arc::new(
1060 move |j: usize| -> Result<
1061 Option<Vec<gam_solve::estimate::reml::PenaltyDerivativeComponent>>,
1062 EstimationError,
1063 > {
1064 let Some(other_axis_in_group) =
1065 group_indices_inner.iter().position(|&idx| idx == j)
1066 else {
1067 return Ok(None);
1068 };
1069 if other_axis_in_group == axis_in_group {
1070 return Ok(None);
1071 }
1072 let cross_pens = provider(other_axis_in_group)?;
1073 if cross_pens.is_empty() {
1074 return Ok(None);
1075 }
1076 Ok(Some(
1077 penalty_indices_inner
1078 .iter()
1079 .copied()
1080 .zip(cross_pens.into_iter().map(|local| {
1081 gam_solve::estimate::reml::HyperPenaltyDerivative::from_embedded(
1082 local,
1083 global_range_inner.clone(),
1084 total_p_inner,
1085 )
1086 }))
1087 .map(|(penalty_index, matrix)| {
1088 gam_solve::estimate::reml::PenaltyDerivativeComponent {
1089 penalty_index,
1090 matrix,
1091 }
1092 })
1093 .collect(),
1094 ))
1095 },
1096 )
1097 as std::sync::Arc<
1098 dyn Fn(
1099 usize,
1100 ) -> Result<
1101 Option<Vec<gam_solve::estimate::reml::PenaltyDerivativeComponent>>,
1102 EstimationError,
1103 > + Send
1104 + Sync
1105 + 'static,
1106 >)
1107 } else {
1108 None
1109 };
1110 let x_first_hyper = if let Some(ref op) = implicit_operator {
1113 gam_solve::estimate::reml::HyperDesignDerivative::from_implicit(
1114 op.clone(),
1115 ImplicitDerivLevel::First(implicit_axis),
1116 global_range.clone(),
1117 total_p,
1118 )
1119 } else {
1120 gam_solve::estimate::reml::HyperDesignDerivative::from_embedded(
1121 x_psi_local,
1122 global_range.clone(),
1123 total_p,
1124 )
1125 };
1126 let mut dir = DirectionalHyperParam::new_compact(
1127 x_first_hyper,
1128 s_components,
1129 Some(xsecond),
1130 Some(ssecond_components),
1131 )?
1132 .not_penalty_like();
1133 if let Some(provider) = penaltysecond_component_provider {
1134 dir = dir.with_penaltysecond_component_provider(provider);
1135 }
1136 if let Some(partner_indices) = penaltysecond_partner_indices {
1137 dir = dir.with_penaltysecond_partner_indices(partner_indices);
1138 }
1139 hyper_dirs.push(dir);
1140 }
1141 Ok(hyper_dirs)
1142}
1143
1144pub(crate) fn spatial_dims_per_term(
1150 resolvedspec: &TermCollectionSpec,
1151 spatial_terms: &[usize],
1152) -> Vec<usize> {
1153 spatial_terms
1154 .iter()
1155 .map(|&term_idx| {
1156 if let Some(mj) = measure_jet_term_spec(resolvedspec, term_idx) {
1157 measure_jet_psi_dim(mj)
1160 } else if spatial_term_uses_per_axis_psi(resolvedspec, term_idx) {
1161 get_spatial_feature_dim(resolvedspec, term_idx).unwrap_or(1)
1162 } else {
1163 1
1164 }
1165 })
1166 .collect()
1167}
1168
1169fn has_aniso_terms(resolvedspec: &TermCollectionSpec, spatial_terms: &[usize]) -> bool {
1173 spatial_terms
1174 .iter()
1175 .any(|&term_idx| spatial_term_uses_per_axis_psi(resolvedspec, term_idx))
1176}
1177
1178macro_rules! impl_exact_joint_theta_memo {
1184 () => {
1185 fn memoized_cost(&self, theta: &Array1<f64>) -> Option<f64> {
1186 if self
1187 .current_theta
1188 .as_ref()
1189 .is_some_and(|cached| theta_values_match(cached, theta))
1190 {
1191 self.last_eval
1192 .as_ref()
1193 .map(|cached| cached.0)
1194 .or(self.last_cost)
1195 } else {
1196 None
1197 }
1198 }
1199
1200 fn memoized_eval(
1201 &self,
1202 theta: &Array1<f64>,
1203 ) -> Option<(
1204 f64,
1205 Array1<f64>,
1206 gam_problem::HessianResult,
1207 )> {
1208 if self
1209 .current_theta
1210 .as_ref()
1211 .is_some_and(|cached| theta_values_match(cached, theta))
1212 {
1213 self.last_eval.clone()
1214 } else {
1215 None
1216 }
1217 }
1218
1219 fn store_eval(
1220 &mut self,
1221 eval: (
1222 f64,
1223 Array1<f64>,
1224 gam_problem::HessianResult,
1225 ),
1226 ) {
1227 self.last_cost = Some(eval.0);
1228 self.last_eval = Some(eval);
1229 }
1230 };
1231}
1232
1233struct SingleBlockExactJointDesignCache<'d> {
1234 realizer: FrozenTermCollectionIncrementalRealizer<'d>,
1235 current_theta: Option<Array1<f64>>,
1236 last_eval_theta: Option<Array1<f64>>,
1243 last_cost: Option<f64>,
1244 last_eval: Option<(
1245 f64,
1246 Array1<f64>,
1247 gam_problem::HessianResult,
1248 )>,
1249 cached_hyper_dirs: Option<(u64, Vec<DirectionalHyperParam>)>,
1261 spatial_terms: Vec<usize>,
1262 rho_dim: usize,
1263 dims_per_term: Vec<usize>,
1264}
1265
1266impl<'d> SingleBlockExactJointDesignCache<'d> {
1267 fn new(
1268 data: ArrayView2<'d, f64>,
1269 spec: TermCollectionSpec,
1270 design: TermCollectionDesign,
1271 spatial_terms: Vec<usize>,
1272 rho_dim: usize,
1273 dims_per_term: Vec<usize>,
1274 ) -> Result<Self, String> {
1275 Ok(Self {
1276 realizer: FrozenTermCollectionIncrementalRealizer::new(data, spec, design)?,
1277 current_theta: None,
1278 last_eval_theta: None,
1279 last_cost: None,
1280 last_eval: None,
1281 cached_hyper_dirs: None,
1282 spatial_terms,
1283 rho_dim,
1284 dims_per_term,
1285 })
1286 }
1287
1288 fn design_revision(&self) -> u64 {
1289 self.realizer.design_revision()
1290 }
1291
1292 fn hyper_dirs_for_current_design(
1302 &mut self,
1303 data: ArrayView2<'_, f64>,
1304 kind: SpatialHyperKind,
1305 ) -> Result<Vec<DirectionalHyperParam>, EstimationError> {
1306 let revision = self.realizer.design_revision();
1307 if let Some((cached_rev, dirs)) = self.cached_hyper_dirs.as_ref()
1308 && *cached_rev == revision
1309 {
1310 return Ok(dirs.clone());
1311 }
1312 let dirs = try_build_spatial_log_kappa_hyper_dirs(
1313 data,
1314 self.realizer.spec(),
1315 self.realizer.design(),
1316 &self.spatial_terms,
1317 )?
1318 .ok_or_else(|| {
1319 EstimationError::InvalidInput(format!(
1320 "failed to build {} hyper_dirs at current {}",
1321 kind.adjective(),
1322 kind.coord_name(),
1323 ))
1324 })?;
1325 self.cached_hyper_dirs = Some((revision, dirs.clone()));
1326 Ok(dirs)
1327 }
1328
1329 fn nfree_tensor_gradient_hyper_dirs(
1330 &mut self,
1331 theta: &Array1<f64>,
1332 ) -> Result<Vec<DirectionalHyperParam>, EstimationError> {
1333 let psi = &theta.as_slice().ok_or_else(|| {
1334 EstimationError::InvalidInput(
1335 "nfree_tensor_gradient_hyper_dirs: theta is not contiguous".to_string(),
1336 )
1337 })?[self.rho_dim..];
1338 let (global_range, p_total, s_psi_components) = self
1339 .realizer
1340 .canonical_penalty_derivatives_at_psi(&self.spatial_terms, psi)
1341 .map_err(EstimationError::InvalidInput)?;
1342 let zero_x = gam_solve::estimate::reml::HyperDesignDerivative::zero(
1343 self.realizer.design().design.nrows(),
1344 p_total,
1345 );
1346 let components = s_psi_components
1347 .into_iter()
1348 .enumerate()
1349 .map(|(penalty_index, local)| {
1350 (
1351 penalty_index,
1352 gam_solve::estimate::reml::HyperPenaltyDerivative::from_embedded(
1353 local,
1354 global_range.clone(),
1355 p_total,
1356 ),
1357 )
1358 })
1359 .collect::<Vec<_>>();
1360 Ok(DirectionalHyperParam::new_compact(zero_x, components, None, None)?.not_penalty_like())
1361 .map(|dir| vec![dir])
1362 }
1363
1364 fn ensure_theta(&mut self, theta: &Array1<f64>) -> Result<(), String> {
1365 if self
1366 .current_theta
1367 .as_ref()
1368 .is_some_and(|cached| theta_values_match(cached, theta))
1369 {
1370 return Ok(());
1371 }
1372 let t_ensure = std::time::Instant::now();
1373 let log_kappa = SpatialLogKappaCoords::from_theta_tail_with_dims(
1374 theta,
1375 self.rho_dim,
1376 self.dims_per_term.clone(),
1377 );
1378 self.realizer
1379 .apply_log_kappa(&log_kappa, &self.spatial_terms)?;
1380 log::info!(
1381 "[STAGE] ensure_theta (apply_log_kappa, {} terms): {:.3}s",
1382 self.spatial_terms.len(),
1383 t_ensure.elapsed().as_secs_f64(),
1384 );
1385 self.current_theta = Some(theta.clone());
1386 self.last_eval_theta = None;
1387 self.last_cost = None;
1388 self.last_eval = None;
1389 Ok(())
1390 }
1391
1392 fn memoized_cost(&self, theta: &Array1<f64>) -> Option<f64> {
1399 if self
1400 .last_eval_theta
1401 .as_ref()
1402 .is_some_and(|cached| theta_values_match(cached, theta))
1403 {
1404 self.last_eval
1405 .as_ref()
1406 .map(|cached| cached.0)
1407 .or(self.last_cost)
1408 } else {
1409 None
1410 }
1411 }
1412
1413 fn memoized_eval(
1414 &self,
1415 theta: &Array1<f64>,
1416 ) -> Option<(
1417 f64,
1418 Array1<f64>,
1419 gam_problem::HessianResult,
1420 )> {
1421 if self
1422 .last_eval_theta
1423 .as_ref()
1424 .is_some_and(|cached| theta_values_match(cached, theta))
1425 {
1426 self.last_eval.clone()
1427 } else {
1428 None
1429 }
1430 }
1431
1432 fn store_eval_at(
1436 &mut self,
1437 theta: &Array1<f64>,
1438 eval: (
1439 f64,
1440 Array1<f64>,
1441 gam_problem::HessianResult,
1442 ),
1443 ) {
1444 self.last_eval_theta = Some(theta.clone());
1445 self.last_cost = Some(eval.0);
1446 self.last_eval = Some(eval);
1447 }
1448
1449 fn store_cost_at(&mut self, theta: &Array1<f64>, cost: f64) {
1452 self.last_eval_theta = Some(theta.clone());
1453 self.last_cost = Some(cost);
1454 self.last_eval = None;
1458 }
1459
1460 fn spec(&self) -> &TermCollectionSpec {
1461 self.realizer.spec()
1462 }
1463
1464 fn design(&self) -> &TermCollectionDesign {
1465 self.realizer.design()
1466 }
1467
1468 fn supports_nfree_penalty_rekey(&self) -> bool {
1474 self.realizer
1475 .supports_nfree_penalty_rekey(&self.spatial_terms)
1476 }
1477
1478 fn supports_nfree_gradient_only_routing(&self) -> bool {
1479 self.realizer
1480 .supports_nfree_gradient_only_routing(&self.spatial_terms)
1481 }
1482
1483 fn canonical_penalties_at(
1493 &mut self,
1494 theta: &Array1<f64>,
1495 ) -> Result<(Vec<gam_terms::construction::CanonicalPenalty>, Vec<usize>), String> {
1496 let psi = &theta
1497 .as_slice()
1498 .ok_or_else(|| "canonical_penalties_at: theta is not contiguous".to_string())?
1499 [self.rho_dim..];
1500 self.realizer
1501 .canonical_penalties_at_psi(&self.spatial_terms, psi)
1502 }
1503}
1504
1505struct SingleBlockLatentCoordDesignCache {
1506 data: Array2<f64>,
1507 spec: TermCollectionSpec,
1508 design: TermCollectionDesign,
1509 current_theta: Option<Array1<f64>>,
1510 current_latent: Option<std::sync::Arc<gam_terms::latent::LatentCoordValues>>,
1511 current_hyper_dirs: Option<Vec<gam_solve::estimate::reml::DirectionalHyperParam>>,
1512 current_design_cache_id: Option<u64>,
1513 latent_design_cache: gam_solve::latent_cache::LatentDesignCache,
1514 last_cost: Option<f64>,
1515 last_eval: Option<(
1516 f64,
1517 Array1<f64>,
1518 gam_problem::HessianResult,
1519 )>,
1520 term_index: gam_problem::types::SmoothTermIdx,
1521 feature_cols: Vec<usize>,
1522 rho_dim: usize,
1523 n_obs: usize,
1524 latent_dim: usize,
1525 id_mode: gam_terms::latent::LatentIdMode,
1526 manifold: gam_terms::latent::LatentManifold,
1527 retraction_registry: gam_solve::latent_cache::LatentRetractionRegistry,
1528 latent_id: u64,
1529 analytic_penalties: Option<std::sync::Arc<gam_terms::AnalyticPenaltyRegistry>>,
1530 analytic_rho_count: usize,
1531 design_revision: u64,
1532 last_outer_iter: Option<u64>,
1536}
1537
1538impl SingleBlockLatentCoordDesignCache {
1539 fn new(
1540 data: Array2<f64>,
1541 spec: TermCollectionSpec,
1542 design: TermCollectionDesign,
1543 latent: &StandardLatentCoordConfig,
1544 rho_dim: usize,
1545 ) -> Result<Self, String> {
1546 if latent.term_index.get() >= spec.smooth_terms.len() {
1547 return Err(SmoothError::dimension_mismatch(format!(
1548 "latent-coordinate term index {} out of bounds for {} smooth terms",
1549 latent.term_index,
1550 spec.smooth_terms.len()
1551 ))
1552 .into());
1553 }
1554 if latent.feature_cols.len() != latent.values.latent_dim() {
1555 return Err(SmoothError::dimension_mismatch(format!(
1556 "latent-coordinate feature width mismatch: feature_cols={}, latent_dim={}",
1557 latent.feature_cols.len(),
1558 latent.values.latent_dim()
1559 ))
1560 .into());
1561 }
1562 if latent.values.n_obs() != data.nrows() {
1563 return Err(SmoothError::dimension_mismatch(format!(
1564 "latent-coordinate row mismatch: latent n={}, data n={}",
1565 latent.values.n_obs(),
1566 data.nrows()
1567 ))
1568 .into());
1569 }
1570 let analytic_rho_count = latent
1571 .analytic_penalties
1572 .as_ref()
1573 .map_or(0, |registry| registry.total_rho_count());
1574 Ok(Self {
1575 data,
1576 spec,
1577 design,
1578 current_theta: None,
1579 current_latent: None,
1580 current_hyper_dirs: None,
1581 current_design_cache_id: None,
1582 latent_design_cache: gam_solve::latent_cache::LatentDesignCache::default(),
1583 last_cost: None,
1584 last_eval: None,
1585 term_index: latent.term_index,
1586 feature_cols: latent.feature_cols.clone(),
1587 rho_dim,
1588 n_obs: latent.values.n_obs(),
1589 latent_dim: latent.values.latent_dim(),
1590 id_mode: latent.values.id_mode().clone(),
1591 manifold: latent.values.manifold().clone(),
1592 retraction_registry: latent.values.retraction_registry().clone(),
1593 latent_id: latent.values.latent_id(),
1594 analytic_penalties: latent.analytic_penalties.clone(),
1595 analytic_rho_count,
1596 design_revision: 0,
1597 last_outer_iter: None,
1598 })
1599 }
1600
1601 fn design_revision(&self) -> u64 {
1602 self.design_revision
1603 }
1604
1605 fn design(&self) -> &TermCollectionDesign {
1606 &self.design
1607 }
1608
1609 fn latent(&self) -> Result<std::sync::Arc<gam_terms::latent::LatentCoordValues>, String> {
1610 self.current_latent
1611 .as_ref()
1612 .cloned()
1613 .ok_or_else(|| "latent-coordinate cache has not been realized".to_string())
1614 }
1615
1616 fn analytic_penalties(&self) -> Option<std::sync::Arc<gam_terms::AnalyticPenaltyRegistry>> {
1617 self.analytic_penalties.clone()
1618 }
1619
1620 fn analytic_penalty_rho_count(&self) -> usize {
1621 self.analytic_rho_count
1622 }
1623
1624 fn hyper_dirs(&self) -> Result<Vec<gam_solve::estimate::reml::DirectionalHyperParam>, String> {
1625 self.current_hyper_dirs
1626 .as_ref()
1627 .cloned()
1628 .ok_or_else(|| "latent-coordinate hyper_dirs cache has not been realized".to_string())
1629 }
1630
1631 fn latent_basis_kind(&self) -> Result<gam_solve::latent_cache::LatentBasisKind, String> {
1632 let smooth_term = self
1633 .design
1634 .smooth
1635 .terms
1636 .get(self.term_index.get())
1637 .ok_or_else(|| {
1638 SmoothError::dimension_mismatch(format!(
1639 "LatentCoord term index {} out of bounds for realized smooth design",
1640 self.term_index
1641 ))
1642 })?;
1643 let termspec = self
1644 .spec
1645 .smooth_terms
1646 .get(self.term_index.get())
1647 .ok_or_else(|| {
1648 SmoothError::dimension_mismatch(format!(
1649 "LatentCoord term index {} out of bounds for resolved smooth spec",
1650 self.term_index
1651 ))
1652 })?;
1653 match (&termspec.basis, &smooth_term.metadata) {
1654 (
1655 SmoothBasisSpec::Matern { .. },
1656 BasisMetadata::Matern {
1657 centers,
1658 length_scale,
1659 nu,
1660 aniso_log_scales,
1661 ..
1662 },
1663 ) => Ok(gam_solve::latent_cache::LatentBasisKind::Matern {
1664 centers: centers.clone(),
1665 length_scale: *length_scale,
1666 nu: *nu,
1667 aniso_log_scales: aniso_log_scales
1668 .clone()
1669 .unwrap_or_else(|| vec![0.0; centers.ncols()]),
1670 chunk_size: gam_terms::basis::auto_streaming_chunk_size_for_dense(
1671 self.n_obs,
1672 centers.nrows(),
1673 ),
1674 }),
1675 (
1676 SmoothBasisSpec::Duchon { .. },
1677 BasisMetadata::Duchon {
1678 centers,
1679 length_scale,
1680 power,
1681 nullspace_order,
1682 aniso_log_scales,
1683 ..
1684 },
1685 ) => Ok(gam_solve::latent_cache::LatentBasisKind::Duchon {
1686 centers: centers.clone(),
1687 length_scale: *length_scale,
1688 power: *power,
1689 nullspace_order: *nullspace_order,
1690 aniso_log_scales: aniso_log_scales
1691 .clone()
1692 .unwrap_or_else(|| vec![0.0; centers.ncols()]),
1693 }),
1694 (
1695 SmoothBasisSpec::Sphere { .. },
1696 BasisMetadata::Sphere {
1697 centers,
1698 penalty_order,
1699 method,
1700 ..
1701 },
1702 ) if matches!(*method, gam_terms::basis::SphereMethod::Wahba) => {
1703 Ok(gam_solve::latent_cache::LatentBasisKind::Sphere {
1704 centers: centers.clone(),
1705 penalty_order: *penalty_order,
1706 chunk_size: gam_terms::basis::auto_streaming_chunk_size_for_dense(
1707 self.n_obs,
1708 centers.nrows(),
1709 ),
1710 })
1711 }
1712 (
1713 SmoothBasisSpec::BSpline1D { spec, .. },
1714 BasisMetadata::BSpline1D {
1715 knots,
1716 periodic,
1717 degree: meta_degree,
1718 ..
1719 },
1720 ) => {
1721 let effective_degree = meta_degree.unwrap_or(spec.degree);
1725 if let Some((domain_start, period, num_basis)) = periodic {
1726 Ok(
1727 gam_solve::latent_cache::LatentBasisKind::PeriodicBspline {
1728 domain_start: *domain_start,
1729 period: *period,
1730 degree: effective_degree,
1731 num_basis: *num_basis,
1732 chunk_size: gam_terms::basis::auto_streaming_chunk_size_for_dense(
1733 self.n_obs, *num_basis,
1734 ),
1735 },
1736 )
1737 } else {
1738 let num_basis_est = knots.len().saturating_sub(effective_degree + 1);
1739 Ok(
1740 gam_solve::latent_cache::LatentBasisKind::TensorBspline {
1741 knots: vec![knots.clone()],
1742 degrees: vec![effective_degree],
1743 chunk_size: gam_terms::basis::auto_streaming_chunk_size_for_dense(
1744 self.n_obs,
1745 num_basis_est,
1746 ),
1747 },
1748 )
1749 }
1750 }
1751 (
1752 SmoothBasisSpec::TensorBSpline { .. },
1753 BasisMetadata::TensorBSpline { knots, degrees, .. },
1754 ) => Ok(
1755 gam_solve::latent_cache::LatentBasisKind::TensorBspline {
1756 knots: knots.clone(),
1757 degrees: degrees.clone(),
1758 chunk_size: None,
1759 },
1760 ),
1761 (
1762 SmoothBasisSpec::Pca { .. },
1763 BasisMetadata::Pca {
1764 basis_matrix,
1765 centered,
1766 smooth_penalty,
1767 center_mean,
1768 pca_basis_path,
1769 chunk_size,
1770 ..
1771 },
1772 ) => {
1773 let center_mean_fingerprint = if *centered && pca_basis_path.is_none() {
1774 let mean = center_mean.as_ref().ok_or_else(|| {
1775 SmoothError::invalid_config(
1776 "latent-coordinate Pca cache key requires center_mean when centered",
1777 )
1778 })?;
1779 Some(gam_solve::latent_cache::pca_center_mean_fingerprint(
1780 mean,
1781 ))
1782 } else {
1783 None
1784 };
1785 Ok(gam_solve::latent_cache::LatentBasisKind::Pca {
1786 basis_matrix: basis_matrix.clone(),
1787 centered: *centered,
1788 center_mean_fingerprint,
1789 smooth_penalty: *smooth_penalty,
1790 pca_basis_path: pca_basis_path.clone(),
1791 chunk_size: *chunk_size,
1792 })
1793 }
1794 _ => Err(SmoothError::invalid_config(
1795 "latent-coordinate design cache could not key the realized latent smooth basis"
1796 .to_string(),
1797 )
1798 .into()),
1799 }
1800 }
1801
1802 fn ensure_theta(&mut self, theta: &Array1<f64>) -> Result<(), String> {
1803 if self
1804 .current_theta
1805 .as_ref()
1806 .is_some_and(|cached| theta_values_match(cached, theta))
1807 {
1808 return Ok(());
1809 }
1810 let latent_flat_len = self.n_obs * self.latent_dim;
1811 let direct_hyper_count = latent_coord_direct_hyper_count(&self.id_mode, self.latent_dim);
1812 let expected =
1813 self.rho_dim + latent_flat_len + self.analytic_rho_count + direct_hyper_count;
1814 if theta.len() != expected {
1815 return Err(SmoothError::dimension_mismatch(format!(
1816 "latent-coordinate theta length mismatch: got {}, expected {} (rho_dim={}, n={}, d={}, analytic_rhos={}, direct_hypers={})",
1817 theta.len(),
1818 expected,
1819 self.rho_dim,
1820 self.n_obs,
1821 self.latent_dim,
1822 self.analytic_rho_count,
1823 direct_hyper_count
1824 ))
1825 .into());
1826 }
1827 let flat = theta
1828 .slice(s![self.rho_dim..self.rho_dim + latent_flat_len])
1829 .to_owned();
1830 let latent = std::sync::Arc::new(
1831 gam_terms::latent::LatentCoordValues::from_flat_with_manifold_and_retraction_and_id(
1832 flat,
1833 self.n_obs,
1834 self.latent_dim,
1835 self.id_mode.clone(),
1836 self.manifold.clone(),
1837 self.retraction_registry.clone(),
1838 self.latent_id,
1839 ),
1840 );
1841 let latent_values_changed = self
1842 .current_latent
1843 .as_ref()
1844 .map(|cached| !latent_values_match(cached.as_flat(), latent.as_flat()))
1845 .unwrap_or(true);
1846 if latent_values_changed {
1847 self.latent_design_cache.invalidate_all();
1848 self.current_design_cache_id = None;
1849 self.design_revision = self.design_revision.wrapping_add(1);
1850 }
1851 for n in 0..self.n_obs {
1852 for axis in 0..self.latent_dim {
1853 let col = self.feature_cols[axis];
1854 self.data[[n, col]] = latent.as_flat()[n * self.latent_dim + axis];
1855 }
1856 }
1857
1858 let basis_kind = self.latent_basis_kind()?;
1859 let rebuilt_width = self.design.design.ncols();
1860 let spec = self.spec.clone();
1861 let term_index = self.term_index;
1862 let analytic_rho_count = self.analytic_rho_count;
1863 let data = self.data.view();
1864 let design_context_digest =
1865 gam_solve::latent_cache::latent_design_context_cache_digest(
1866 data,
1867 &spec,
1868 term_index,
1869 analytic_rho_count,
1870 &self.feature_cols,
1871 )
1872 .map_err(|e| e.to_string())?;
1873 let lookup = self
1874 .latent_design_cache
1875 .lookup_or_compute(latent.clone(), basis_kind, design_context_digest, || {
1876 let rebuilt = build_term_collection_design(data, &spec).map_err(|e| {
1877 EstimationError::InvalidInput(format!(
1878 "failed to rebuild latent-coordinate design: {e}"
1879 ))
1880 })?;
1881 if rebuilt.design.ncols() != rebuilt_width {
1882 crate::bail_invalid_estim!(
1883 "latent-coordinate design topology changed: rebuilt p={}, cached p={}",
1884 rebuilt.design.ncols(),
1885 rebuilt_width
1886 );
1887 }
1888 let hyper_dirs = try_build_latent_coord_hyper_dirs(
1889 latent.clone(),
1890 &spec,
1891 &rebuilt,
1892 &[term_index],
1893 analytic_rho_count,
1894 )?
1895 .ok_or_else(|| {
1896 EstimationError::InvalidInput(
1897 "failed to build latent-coordinate hyper_dirs".to_string(),
1898 )
1899 })?;
1900 Ok(gam_solve::latent_cache::ComputedLatentDesign {
1901 design: rebuilt,
1902 hyper_dirs,
1903 })
1904 })
1905 .map_err(|e| e.to_string())?;
1906 if lookup.cached.design.design.ncols() != self.design.design.ncols() {
1907 return Err(SmoothError::dimension_mismatch(format!(
1908 "latent-coordinate design topology changed: rebuilt p={}, cached p={}",
1909 lookup.cached.design.design.ncols(),
1910 self.design.design.ncols()
1911 ))
1912 .into());
1913 }
1914 self.design = lookup.cached.design.clone();
1915 self.current_hyper_dirs = Some(lookup.cached.hyper_dirs.clone());
1916 self.current_latent = Some(latent);
1917 self.current_theta = Some(theta.clone());
1918 self.last_cost = None;
1919 self.last_eval = None;
1920 self.last_outer_iter = None;
1921 if !latent_values_changed && self.current_design_cache_id != Some(lookup.entry_id) {
1922 self.design_revision = self.design_revision.wrapping_add(1);
1923 }
1924 self.current_design_cache_id = Some(lookup.entry_id);
1925 Ok(())
1926 }
1927
1928 fn memoized_cost(&self, theta: &Array1<f64>) -> Option<f64> {
1929 if self
1930 .current_theta
1931 .as_ref()
1932 .is_some_and(|cached| theta_values_match(cached, theta))
1933 && self.last_outer_iter
1934 == Some(gam_solve::estimate::reml::outer_eval::current_outer_iter())
1935 {
1936 self.last_eval
1937 .as_ref()
1938 .map(|cached| cached.0)
1939 .or(self.last_cost)
1940 } else {
1941 None
1942 }
1943 }
1944
1945 fn memoized_eval(
1946 &self,
1947 theta: &Array1<f64>,
1948 ) -> Option<(
1949 f64,
1950 Array1<f64>,
1951 gam_problem::HessianResult,
1952 )> {
1953 if self
1954 .current_theta
1955 .as_ref()
1956 .is_some_and(|cached| theta_values_match(cached, theta))
1957 && self.last_outer_iter
1958 == Some(gam_solve::estimate::reml::outer_eval::current_outer_iter())
1959 {
1960 self.last_eval.clone()
1961 } else {
1962 None
1963 }
1964 }
1965
1966 fn store_eval(
1967 &mut self,
1968 eval: (
1969 f64,
1970 Array1<f64>,
1971 gam_problem::HessianResult,
1972 ),
1973 ) {
1974 self.last_cost = Some(eval.0);
1975 self.last_eval = Some(eval);
1976 self.last_outer_iter =
1977 Some(gam_solve::estimate::reml::outer_eval::current_outer_iter());
1978 }
1979
1980 fn store_cost(&mut self, cost: f64) {
1981 self.last_cost = Some(cost);
1982 self.last_outer_iter =
1983 Some(gam_solve::estimate::reml::outer_eval::current_outer_iter());
1984 }
1985
1986 fn reset(&mut self) {
1987 self.current_theta = None;
1988 self.current_latent = None;
1989 self.current_hyper_dirs = None;
1990 self.current_design_cache_id = None;
1991 self.latent_design_cache.invalidate();
1992 self.last_cost = None;
1993 self.last_eval = None;
1994 self.last_outer_iter = None;
1995 }
1996}
1997
1998pub fn fixed_kappa_profiled_reml_score(
2014 data: ArrayView2<'_, f64>,
2015 y: ArrayView1<'_, f64>,
2016 weights: ArrayView1<'_, f64>,
2017 offset: ArrayView1<'_, f64>,
2018 resolvedspec: &TermCollectionSpec,
2019 term_idx: usize,
2020 kappa: f64,
2021 family: LikelihoodSpec,
2022 options: &FitOptions,
2023) -> Result<f64, EstimationError> {
2024 if !kappa.is_finite() {
2025 crate::bail_invalid_estim!("fixed-κ profiled score probed a non-finite κ = {kappa}");
2026 }
2027 let (feature_cols, mut probe_basis) = match resolvedspec
2030 .smooth_terms
2031 .get(term_idx)
2032 .map(|t| &t.basis)
2033 {
2034 Some(SmoothBasisSpec::ConstantCurvature {
2035 feature_cols, spec, ..
2036 }) => (feature_cols.clone(), spec.clone()),
2037 _ => {
2038 crate::bail_invalid_estim!(
2039 "fixed-κ profiled score: term {term_idx} is not a constant-curvature smooth"
2040 )
2041 }
2042 };
2043 probe_basis.kappa = kappa;
2044
2045 let is_unweighted = weights.iter().all(|&w| (w - 1.0).abs() <= 1e-12);
2065 let is_zero_offset = offset.iter().all(|&o| o.abs() <= 1e-12);
2066 if family == LikelihoodSpec::gaussian_identity() && is_unweighted && is_zero_offset {
2067 let x_term = select_columns(data, &feature_cols).map_err(EstimationError::from)?;
2068 let score =
2069 gam_terms::basis::constant_curvature_honest_profiled_reml_score(x_term.view(), y, &probe_basis)
2070 .map_err(|e| {
2071 EstimationError::InvalidInput(format!(
2072 "fixed-κ honest profiled-REML score at κ={kappa} failed: {e}"
2073 ))
2074 })?;
2075 if !score.is_finite() {
2076 crate::bail_invalid_estim!(
2077 "fixed-κ honest profiled-REML score at κ={kappa} is non-finite"
2078 );
2079 }
2080 return Ok(score);
2081 }
2082
2083 let mut probe_spec = resolvedspec.clone();
2085 match probe_spec.smooth_terms.get_mut(term_idx).map(|t| &mut t.basis) {
2086 Some(SmoothBasisSpec::ConstantCurvature { spec, .. }) => spec.kappa = kappa,
2087 _ => {
2088 crate::bail_invalid_estim!(
2089 "fixed-κ profiled score: term {term_idx} is not a constant-curvature smooth"
2090 )
2091 }
2092 }
2093 let fixed_kappa_options = SpatialLengthScaleOptimizationOptions {
2094 enabled: false,
2095 ..SpatialLengthScaleOptimizationOptions::default()
2096 };
2097 let fit = fit_term_collectionwith_spatial_length_scale_optimization(
2098 data,
2099 y.to_owned(),
2100 weights.to_owned(),
2101 offset.to_owned(),
2102 &probe_spec,
2103 family,
2104 options,
2105 &fixed_kappa_options,
2106 )?;
2107 let score = fit_score(&fit.fit);
2108 if !score.is_finite() {
2109 crate::bail_invalid_estim!("fixed-κ profiled fit at κ={kappa} returned a non-finite score");
2110 }
2111 Ok(score)
2112}
2113
2114fn constant_curvature_kappa_fair_argmin(
2139 data: ArrayView2<'_, f64>,
2140 y: ArrayView1<'_, f64>,
2141 resolvedspec: &TermCollectionSpec,
2142 term_idx: usize,
2143) -> Option<f64> {
2144 let (kappa_min, kappa_max) = constant_curvature_kappa_bounds(data, resolvedspec, term_idx);
2145 if !(kappa_min.is_finite() && kappa_max.is_finite() && kappa_max > kappa_min) {
2146 return None;
2147 }
2148 let (feature_cols, base_spec) = match resolvedspec.smooth_terms.get(term_idx).map(|t| &t.basis) {
2149 Some(SmoothBasisSpec::ConstantCurvature {
2150 feature_cols, spec, ..
2151 }) => (feature_cols, spec.clone()),
2152 _ => return None,
2153 };
2154 let x_term = match select_columns(data, feature_cols) {
2155 Ok(x) => x,
2156 Err(e) => {
2157 log::info!("[spatial-kappa] #1464 κ-fair argmin column select failed ({e}); skipping");
2158 return None;
2159 }
2160 };
2161 const GRID_STEPS: usize = 24;
2167 let mut best: Option<(f64, f64)> = None; for i in 0..=GRID_STEPS {
2169 let t = i as f64 / GRID_STEPS as f64;
2170 let kappa = kappa_min + (kappa_max - kappa_min) * t;
2171 let mut probe_spec = base_spec.clone();
2172 probe_spec.kappa = kappa;
2173 match gam_terms::basis::constant_curvature_kappa_fair_sign_score(x_term.view(), y, &probe_spec) {
2174 Ok(score) => {
2175 if best.as_ref().is_none_or(|(b, _)| score < *b) {
2176 best = Some((score, kappa));
2177 }
2178 }
2179 Err(e) => {
2180 log::info!(
2181 "[spatial-kappa] #1464 κ-fair argmin probe at κ={kappa:.4} failed ({e}); skipping"
2182 );
2183 }
2184 }
2185 }
2186 best.map(|(score, kappa)| {
2187 log::info!(
2188 "[spatial-kappa] #1464 κ-fair argmin κ̂={kappa:.4} (κ-fair score={score:.6e}) for term {term_idx}"
2189 );
2190 kappa
2191 })
2192}
2193
2194fn select_constant_curvature_kappa_sign_seed(
2202 data: ArrayView2<'_, f64>,
2203 y: ArrayView1<'_, f64>,
2204 resolvedspec: &TermCollectionSpec,
2205 term_idx: usize,
2206) -> Option<f64> {
2207 let (kappa_min, kappa_max) = constant_curvature_kappa_bounds(data, resolvedspec, term_idx);
2208 if !(kappa_min.is_finite() && kappa_max.is_finite() && kappa_max > kappa_min) {
2209 return None;
2210 }
2211 let (feature_cols, base_spec) = match resolvedspec.smooth_terms.get(term_idx).map(|t| &t.basis) {
2223 Some(SmoothBasisSpec::ConstantCurvature {
2224 feature_cols, spec, ..
2225 }) => (feature_cols, spec.clone()),
2226 _ => return None,
2227 };
2228 let x_term = match select_columns(data, feature_cols) {
2229 Ok(x) => x,
2230 Err(e) => {
2231 log::info!("[spatial-kappa] #1464 sign-basin scan column select failed ({e}); skipping");
2232 return None;
2233 }
2234 };
2235 let probes = [
2239 kappa_min,
2240 0.5 * kappa_min,
2241 0.0,
2242 0.5 * kappa_max,
2243 kappa_max,
2244 ];
2245 let mut best: Option<(f64, f64)> = None; for &kappa in &probes {
2247 let mut probe_spec = base_spec.clone();
2248 probe_spec.kappa = kappa;
2249 match gam_terms::basis::constant_curvature_kappa_fair_sign_score(
2250 x_term.view(),
2251 y,
2252 &probe_spec,
2253 ) {
2254 Ok(score) => {
2255 if best.as_ref().is_none_or(|(b, _)| score < *b) {
2256 best = Some((score, kappa));
2257 }
2258 }
2259 Err(e) => {
2260 log::info!(
2261 "[spatial-kappa] #1464 sign-basin probe at κ={kappa:.4} failed ({e}); skipping"
2262 );
2263 }
2264 }
2265 }
2266 best.map(|(score, kappa)| {
2267 log::info!(
2268 "[spatial-kappa] #1464 κ-fair sign-basin scan selected κ_seed={kappa:.4} \
2269 (κ-fair score={score:.6e}) for term {term_idx}"
2270 );
2271 kappa
2272 })
2273}
2274
2275const SPATIAL_RANGE_PRESCAN_GRID: usize = 7;
2278
2279fn prescan_isotropic_spatial_range_seed(
2311 data: ArrayView2<'_, f64>,
2312 y: ArrayView1<'_, f64>,
2313 weights: ArrayView1<'_, f64>,
2314 offset: ArrayView1<'_, f64>,
2315 resolvedspec: &TermCollectionSpec,
2316 baseline_score: f64,
2317 family: &LikelihoodSpec,
2318 options: &FitOptions,
2319 kappa_options: &SpatialLengthScaleOptimizationOptions,
2320 spatial_terms: &[usize],
2321) -> Result<Vec<(usize, f64)>, EstimationError> {
2322 if has_aniso_terms(resolvedspec, spatial_terms)
2324 || !constant_curvature_term_indices(resolvedspec).is_empty()
2325 {
2326 return Ok(Vec::new());
2327 }
2328 let dims = spatial_dims_per_term(resolvedspec, spatial_terms);
2329 let mut working = resolvedspec.clone();
2333 let mut best_score = if baseline_score.is_finite() {
2334 baseline_score
2335 } else {
2336 f64::INFINITY
2337 };
2338 let mut overrides: Vec<(usize, f64)> = Vec::new();
2339 for (slot, &term_idx) in spatial_terms.iter().enumerate() {
2340 if dims.get(slot).copied().unwrap_or(1) != 1 {
2343 continue;
2344 }
2345 if get_spatial_length_scale(&working, term_idx).is_none() {
2348 continue;
2349 }
2350 let (psi_lo, psi_hi) = spatial_term_psi_bounds(data, &working, term_idx, kappa_options);
2351 if !(psi_lo.is_finite() && psi_hi.is_finite()) || psi_hi <= psi_lo {
2352 continue;
2353 }
2354 let mut term_best: Option<f64> = None;
2355 for g in 0..SPATIAL_RANGE_PRESCAN_GRID {
2356 let frac = g as f64 / (SPATIAL_RANGE_PRESCAN_GRID - 1) as f64;
2357 let psi = psi_lo + (psi_hi - psi_lo) * frac;
2358 let ls = (-psi).exp();
2362 if !ls.is_finite() || ls <= 0.0 {
2363 continue;
2364 }
2365 let mut probe = working.clone();
2366 if set_spatial_length_scale(&mut probe, term_idx, ls).is_err() {
2367 continue;
2368 }
2369 let fit = match fit_term_collection_forspec(
2378 data,
2379 y,
2380 weights,
2381 offset,
2382 &probe,
2383 family.clone(),
2384 options,
2385 ) {
2386 Ok(fit) => fit,
2387 Err(_) => continue,
2390 };
2391 let score = fit_score(&fit.fit);
2392 if score.is_finite() && score < best_score - 1e-7 * best_score.abs().max(1.0) {
2395 best_score = score;
2396 term_best = Some(ls);
2397 }
2398 }
2399 if let Some(ls) = term_best {
2400 set_spatial_length_scale(&mut working, term_idx, ls)?;
2401 overrides.push((term_idx, ls));
2402 log::info!(
2403 "[spatial-kappa] #1074 range pre-scan: term {term_idx} re-seeded at \
2404 length_scale={ls:.5} (profiled REML {best_score:.5}, was {baseline_score:.5})"
2405 );
2406 }
2407 }
2408 Ok(overrides)
2409}
2410
2411const JOINT_RESTART_WINDOW_FRACTIONS: [f64; 5] = [0.0, 0.2, 0.45, 0.7, 1.0];
2420
2421fn joint_solve_from_window_fraction(
2437 data: ArrayView2<'_, f64>,
2438 y: ArrayView1<'_, f64>,
2439 weights: ArrayView1<'_, f64>,
2440 offset: ArrayView1<'_, f64>,
2441 base_spec: &TermCollectionSpec,
2442 spatial_terms: &[usize],
2443 fraction: f64,
2444 family: &LikelihoodSpec,
2445 options: &FitOptions,
2446 baseline_options: &FitOptions,
2447 kappa_options: &SpatialLengthScaleOptimizationOptions,
2448) -> Result<Option<(FittedTermCollectionWithSpec, f64)>, EstimationError> {
2449 let mut seed_spec = base_spec.clone();
2450 let mut any_set = false;
2451 for &term_idx in spatial_terms {
2452 if get_spatial_length_scale(&seed_spec, term_idx).is_none() {
2453 continue;
2454 }
2455 let (psi_lo, psi_hi) = spatial_term_psi_bounds(data, &seed_spec, term_idx, kappa_options);
2456 if !(psi_lo.is_finite() && psi_hi.is_finite()) || psi_hi <= psi_lo {
2457 continue;
2458 }
2459 let psi = psi_lo + (psi_hi - psi_lo) * fraction;
2460 let ls = (-psi).exp();
2461 if !ls.is_finite() || ls <= 0.0 {
2462 continue;
2463 }
2464 if set_spatial_length_scale(&mut seed_spec, term_idx, ls).is_ok() {
2465 any_set = true;
2466 }
2467 }
2468 if !any_set {
2469 return Ok(None);
2470 }
2471 let seed_best = match fit_term_collection_forspec(
2475 data,
2476 y,
2477 weights,
2478 offset,
2479 &seed_spec,
2480 family.clone(),
2481 baseline_options,
2482 ) {
2483 Ok(fit) => fit,
2484 Err(_) => return Ok(None),
2485 };
2486 let seed_spec = freeze_term_collection_from_design(&seed_spec, &seed_best.design)?;
2487 let seed_terms = spatial_length_scale_term_indices(&seed_spec);
2490 if seed_terms.is_empty() {
2491 let score = fit_score(&seed_best.fit);
2492 return Ok(Some((
2493 FittedTermCollectionWithSpec {
2494 fit: seed_best.fit,
2495 design: seed_best.design,
2496 resolvedspec: seed_spec,
2497 adaptive_diagnostics: seed_best.adaptive_diagnostics,
2498 kappa_timing: None,
2499 },
2500 score,
2501 )));
2502 }
2503 let joint = try_exact_joint_spatial_length_scale_optimization(
2504 data,
2505 y,
2506 weights,
2507 offset,
2508 &seed_spec,
2509 &seed_best,
2510 family.clone(),
2511 options,
2512 kappa_options,
2513 &seed_terms,
2514 )?;
2515 match joint {
2516 Some(fit) => {
2517 let score = fit_score(&fit.fit);
2518 Ok(Some((fit, score)))
2519 }
2520 None => {
2523 let score = fit_score(&seed_best.fit);
2524 Ok(Some((
2525 FittedTermCollectionWithSpec {
2526 fit: seed_best.fit,
2527 design: seed_best.design,
2528 resolvedspec: seed_spec,
2529 adaptive_diagnostics: seed_best.adaptive_diagnostics,
2530 kappa_timing: None,
2531 },
2532 score,
2533 )))
2534 }
2535 }
2536}
2537
2538fn try_exact_joint_spatial_length_scale_optimization(
2539 data: ArrayView2<'_, f64>,
2540 y: ArrayView1<'_, f64>,
2541 weights: ArrayView1<'_, f64>,
2542 offset: ArrayView1<'_, f64>,
2543 resolvedspec: &TermCollectionSpec,
2544 best: &FittedTermCollection,
2545 family: LikelihoodSpec,
2546 options: &FitOptions,
2547 kappa_options: &SpatialLengthScaleOptimizationOptions,
2548 spatial_terms: &[usize],
2549) -> Result<Option<FittedTermCollectionWithSpec>, EstimationError> {
2550 if spatial_terms.is_empty() {
2551 return Ok(None);
2552 }
2553 kappa_options
2558 .validate()
2559 .map_err(EstimationError::InvalidInput)?;
2560
2561 let cc_term_set = constant_curvature_term_indices(resolvedspec);
2581 let all_spatial_are_cc =
2582 !cc_term_set.is_empty() && spatial_terms.iter().all(|t| cc_term_set.contains(t));
2583 if all_spatial_are_cc {
2584 let mut fixed_kappa_spec = resolvedspec.clone();
2585 let mut any_kappa_chosen = false;
2586 for &term_idx in spatial_terms {
2587 if constant_curvature_kappa_is_fixed(resolvedspec, term_idx) {
2591 continue;
2592 }
2593 if let Some(kappa_hat) =
2604 constant_curvature_kappa_fair_argmin(data, y, resolvedspec, term_idx)
2605 .filter(|&k| k < 0.0)
2606 {
2607 if let Some(SmoothBasisSpec::ConstantCurvature { spec: cc, .. }) = fixed_kappa_spec
2608 .smooth_terms
2609 .get_mut(term_idx)
2610 .map(|t| &mut t.basis)
2611 {
2612 cc.kappa = kappa_hat;
2613 any_kappa_chosen = true;
2614 log::info!(
2615 "[spatial-kappa] #1464 term {term_idx}: fixed κ̂ = {kappa_hat:.4} from κ-fair argmin (hyperbolic basin; profiling ρ only)"
2616 );
2617 }
2618 }
2619 }
2620 if any_kappa_chosen {
2621 let baseline_score = fit_score(&best.fit);
2625 let fitted = fit_term_collection_forspec(
2626 data,
2627 y,
2628 weights,
2629 offset,
2630 &fixed_kappa_spec,
2631 family.clone(),
2632 options,
2633 )?;
2634 let frozen_spec =
2635 freeze_term_collection_from_design(&fixed_kappa_spec, &fitted.design)?;
2636 let mut fit = fitted.fit;
2637 fit.reml_score = baseline_score;
2649 return Ok(Some(FittedTermCollectionWithSpec {
2650 fit,
2651 design: fitted.design,
2652 resolvedspec: frozen_spec,
2653 adaptive_diagnostics: fitted.adaptive_diagnostics,
2654 kappa_timing: None,
2655 }));
2656 }
2657 }
2658
2659 if try_build_spatial_log_kappa_hyper_dirs(data, resolvedspec, &best.design, spatial_terms)?
2660 .is_none()
2661 {
2662 if !constant_curvature_term_indices(resolvedspec).is_empty() {
2663 log::info!(
2664 "[#1464-trace] try_exact_joint RETURNED None (hyper_dirs unavailable); \
2665 κ̂ comes from a NON-joint path"
2666 );
2667 }
2668 return Ok(None);
2669 }
2670 if !constant_curvature_term_indices(resolvedspec).is_empty() {
2671 log::info!(
2672 "[#1464-trace] try_exact_joint ENTERED for {} spatial term(s); CC present",
2673 spatial_terms.len()
2674 );
2675 }
2676
2677 const JOINT_RHO_BOUND: f64 = 12.0;
2678 let rho_dim = best.fit.lambdas.len();
2679
2680 let has_constant_curvature_term = !constant_curvature_term_indices(resolvedspec).is_empty();
2694 let rho_upper_bound = if has_constant_curvature_term {
2695 gam_solve::estimate::RHO_BOUND
2696 } else {
2697 JOINT_RHO_BOUND
2698 };
2699
2700 let dims_per_term = spatial_dims_per_term(resolvedspec, spatial_terms);
2702 let use_aniso = has_aniso_terms(resolvedspec, spatial_terms);
2703
2704 let log_kappa0 = if use_aniso {
2709 SpatialLogKappaCoords::from_length_scales_aniso(resolvedspec, spatial_terms, kappa_options)
2710 } else {
2711 SpatialLogKappaCoords::from_length_scales(resolvedspec, spatial_terms, kappa_options)
2712 };
2713 let mut log_kappa0 =
2716 log_kappa0.reseed_from_data(data, resolvedspec, spatial_terms, kappa_options);
2717 let mut cc_sign_seeds: Vec<(usize, f64)> = Vec::new();
2733 let mut cc_fixed_seeds: Vec<(usize, f64)> = Vec::new();
2739 if has_constant_curvature_term {
2740 for (slot, &term_idx) in spatial_terms.iter().enumerate() {
2741 if constant_curvature_term_spec(resolvedspec, term_idx).is_none() {
2742 continue;
2743 }
2744 if constant_curvature_kappa_is_fixed(resolvedspec, term_idx) {
2745 let fixed_kappa = get_constant_curvature_kappa(resolvedspec, term_idx)
2746 .expect("constant-curvature term exposes its κ");
2747 log::info!(
2748 "[#2152] term {term_idx}: κ PINNED at user kappa={fixed_kappa}; \
2749 freezing the joint ψ coordinate (ρ-only refinement, no κ scan)"
2750 );
2751 log_kappa0.set_scalar_slot(slot, fixed_kappa);
2752 cc_fixed_seeds.push((slot, fixed_kappa));
2753 continue;
2754 }
2755 let scan = select_constant_curvature_kappa_sign_seed(
2756 data,
2757 y,
2758 resolvedspec,
2759 term_idx,
2760 );
2761 match scan {
2766 Some(kappa_seed) => {
2767 log::info!(
2768 "[#1464-trace] term {term_idx}: κ-fair sign-basin scan picked κ_seed = {kappa_seed}"
2769 );
2770 log_kappa0.set_scalar_slot(slot, kappa_seed);
2771 cc_sign_seeds.push((slot, kappa_seed));
2772 }
2773 None => {
2774 log::info!(
2775 "[#1464-trace] term {term_idx}: fixed-κ sign-basin scan returned NONE (no seed applied)"
2776 );
2777 }
2778 }
2779 }
2780 }
2781 let log_kappa_lower = if use_aniso {
2782 SpatialLogKappaCoords::lower_bounds_aniso_from_data(
2783 data,
2784 resolvedspec,
2785 spatial_terms,
2786 &dims_per_term,
2787 kappa_options,
2788 )
2789 } else {
2790 SpatialLogKappaCoords::lower_bounds_from_data(
2791 data,
2792 resolvedspec,
2793 spatial_terms,
2794 kappa_options,
2795 )
2796 };
2797 let log_kappa_upper = if use_aniso {
2798 SpatialLogKappaCoords::upper_bounds_aniso_from_data(
2799 data,
2800 resolvedspec,
2801 spatial_terms,
2802 &dims_per_term,
2803 kappa_options,
2804 )
2805 } else {
2806 SpatialLogKappaCoords::upper_bounds_from_data(
2807 data,
2808 resolvedspec,
2809 spatial_terms,
2810 kappa_options,
2811 )
2812 };
2813 let mut log_kappa_lower = log_kappa_lower;
2837 let mut log_kappa_upper = log_kappa_upper;
2838 for &(slot, kappa_seed) in &cc_sign_seeds {
2839 if kappa_seed != 0.0 {
2840 log_kappa_lower.set_scalar_slot(slot, kappa_seed);
2841 log_kappa_upper.set_scalar_slot(slot, kappa_seed);
2842 }
2843 log::info!(
2844 "[#1464-trace] slot {slot}: FROZE joint ψ coordinate at κ_seed={kappa_seed} \
2845 (window [{}, {}]); raw fit_score is sign-blind so the κ-fair scan is authoritative",
2846 log_kappa_lower.as_array()[log_kappa_lower.dims_per_term()[..slot].iter().sum::<usize>()],
2847 log_kappa_upper.as_array()[log_kappa_upper.dims_per_term()[..slot].iter().sum::<usize>()],
2848 );
2849 }
2850 for &(slot, fixed_kappa) in &cc_fixed_seeds {
2857 log_kappa_lower.set_scalar_slot(slot, fixed_kappa);
2858 log_kappa_upper.set_scalar_slot(slot, fixed_kappa);
2859 log::info!(
2860 "[#2152] slot {slot}: FROZE joint ψ coordinate at PINNED κ={fixed_kappa} \
2861 (window [{fixed_kappa}, {fixed_kappa}]); ρ-only refinement at the fixed geometry"
2862 );
2863 }
2864 let log_kappa0 = log_kappa0.clamp_to_bounds(&log_kappa_lower, &log_kappa_upper);
2867 let setup = ExactJointHyperSetup::new(
2868 best.fit.lambdas.mapv(f64::ln),
2869 Array1::<f64>::from_elem(rho_dim, -JOINT_RHO_BOUND),
2870 Array1::<f64>::from_elem(rho_dim, rho_upper_bound),
2871 log_kappa0,
2872 log_kappa_lower,
2873 log_kappa_upper,
2874 );
2875
2876 let theta0 = setup.theta0();
2877 let lower = setup.lower();
2878 let upper = setup.upper();
2879
2880 let kind = if use_aniso {
2892 SpatialHyperKind::Anisotropic
2893 } else {
2894 SpatialHyperKind::Isotropic
2895 };
2896 let (outcome, kappa_timing) = run_exact_joint_spatial_optimization(
2897 kind,
2898 data,
2899 y,
2900 weights,
2901 offset,
2902 resolvedspec,
2903 &best.design,
2904 family.clone(),
2905 options,
2906 spatial_terms,
2907 &dims_per_term,
2908 &theta0,
2909 &lower,
2910 &upper,
2911 rho_dim,
2912 kappa_options,
2913 )?;
2914
2915 let baseline_score = fit_score(&best.fit);
2916
2917 let (theta_star, joint_final_value) = match outcome {
2927 SpatialJointOutcome::Optimized {
2928 theta_star,
2929 final_value,
2930 } => (theta_star, final_value),
2931 SpatialJointOutcome::NonConverged {
2932 iterations,
2933 final_value,
2934 final_grad_norm,
2935 } => {
2936 if has_constant_curvature_term {
2937 log::info!(
2938 "[#1464-trace] joint solve NONCONVERGED (iters={iterations}, \
2939 final_value={final_value}); returning FROZEN BASELINE geometry \
2940 (κ̂ = spec default, NOT the joint candidate)"
2941 );
2942 }
2943 log::info!(
2944 "[spatial-kappa] joint spatial optimization did not converge \
2945 (iterations={}, final_objective={:.6e}, final_grad_norm={}); \
2946 keeping the frozen baseline geometry",
2947 iterations,
2948 final_value,
2949 final_grad_norm.map_or_else(|| "n/a".to_string(), |g| format!("{g:.3e}")),
2950 );
2951 let baseline = fit_frozen_baseline_geometry(
2952 data,
2953 y,
2954 weights,
2955 offset,
2956 resolvedspec,
2957 best,
2958 family.clone(),
2959 options,
2960 baseline_score,
2961 Some(kappa_timing),
2962 )?;
2963 let baseline_edf = baseline.fit.inference.as_ref().map(|inf| inf.edf_total);
2981 if let Some(base_edf) = baseline_edf
2982 && base_edf < SPATIAL_COLLAPSE_EDF_FLOOR
2983 && let Some(adaptive) = fit_data_adaptive_geometry(
2984 data,
2985 y,
2986 weights,
2987 offset,
2988 resolvedspec,
2989 spatial_terms,
2990 &dims_per_term,
2991 &theta0,
2992 &lower,
2993 &upper,
2994 rho_dim,
2995 family,
2996 options,
2997 baseline_score,
2998 Some(kappa_timing),
2999 )?
3000 {
3001 let adaptive_edf = adaptive.fit.inference.as_ref().map(|inf| inf.edf_total);
3002 if let Some(adapt_edf) = adaptive_edf
3003 && adapt_edf >= base_edf + SPATIAL_COLLAPSE_EDF_MARGIN
3004 {
3005 log::info!(
3006 "[spatial-kappa] #2122 stalled joint solve collapsed the frozen \
3007 baseline (edf={base_edf:.3}); data-adaptive geometry recovers \
3008 edf={adapt_edf:.3} — shipping the data-adaptive fit"
3009 );
3010 return Ok(Some(adaptive));
3011 }
3012 }
3013 return Ok(Some(baseline));
3014 }
3015 };
3016
3017 let accept_tol = options.tol.max(1e-8 * baseline_score.abs()).max(1e-12);
3022 if joint_final_value > baseline_score + accept_tol {
3023 if has_constant_curvature_term {
3024 log::info!(
3025 "[#1464-trace] joint candidate WORSENED score (joint={joint_final_value}, \
3026 baseline={baseline_score}); returning FROZEN BASELINE geometry \
3027 (κ̂ = spec default, NOT the joint candidate)"
3028 );
3029 }
3030 log::info!(
3031 "[spatial-kappa] exact joint spatial candidate worsened the profiled score (joint={:.6e}, baseline={:.6e}, tol={:.2e}); keeping the frozen baseline geometry",
3032 joint_final_value,
3033 baseline_score,
3034 accept_tol,
3035 );
3036 return Ok(Some(fit_frozen_baseline_geometry(
3037 data,
3038 y,
3039 weights,
3040 offset,
3041 resolvedspec,
3042 best,
3043 family,
3044 options,
3045 baseline_score,
3046 Some(kappa_timing),
3047 )?));
3048 }
3049
3050 let rho_star = theta_star.slice(s![..rho_dim]).mapv(f64::exp);
3051 let log_kappa_star =
3052 SpatialLogKappaCoords::from_theta_tail_with_dims(&theta_star, rho_dim, dims_per_term);
3053 if has_constant_curvature_term {
3059 let star = log_kappa_star.as_array();
3060 let dims = log_kappa_star.dims_per_term();
3061 for (slot, &term_idx) in spatial_terms.iter().enumerate() {
3062 if constant_curvature_term_spec(resolvedspec, term_idx).is_some() {
3063 let off: usize = dims[..slot].iter().sum();
3064 log::info!(
3065 "[#1464-trace] term {term_idx}: joint solver CONVERGED ψ-tail κ = {} \
3066 (this is the optimised candidate; joint_final_value={joint_final_value})",
3067 star[off]
3068 );
3069 }
3070 }
3071 }
3072 let baseline_spec = resolvedspec;
3076 let optimized_spec = log_kappa_star.apply_tospec(resolvedspec, spatial_terms)?;
3077 let optimized = fit_term_collection_forspecwith_heuristic_lambdas(
3078 data,
3079 y,
3080 weights,
3081 offset,
3082 &optimized_spec,
3083 rho_star.as_slice(),
3084 family.clone(),
3085 options,
3086 )?;
3087
3088 let optimized_edf = optimized.fit.inference.as_ref().map(|inf| inf.edf_total);
3102 if let Some(opt_edf) = optimized_edf
3103 && opt_edf < SPATIAL_COLLAPSE_EDF_FLOOR
3104 {
3105 let baseline = fit_frozen_baseline_geometry(
3106 data,
3107 y,
3108 weights,
3109 offset,
3110 baseline_spec,
3111 best,
3112 family.clone(),
3113 options,
3114 baseline_score,
3115 Some(kappa_timing),
3116 )?;
3117 let baseline_edf = baseline.fit.inference.as_ref().map(|inf| inf.edf_total);
3118 if let Some(base_edf) = baseline_edf
3119 && base_edf >= opt_edf + SPATIAL_COLLAPSE_EDF_MARGIN
3120 {
3121 log::info!(
3122 "[spatial-kappa] joint candidate collapsed to the null (edf={opt_edf:.3}); \
3123 baseline geometry retains edf={base_edf:.3} — keeping the frozen baseline",
3124 );
3125 return Ok(Some(baseline));
3126 }
3127 }
3130
3131 let mut fit = optimized.fit;
3135 fit.reml_score = joint_final_value;
3136 let optimized_result = FittedTermCollectionWithSpec {
3137 fit,
3138 design: optimized.design,
3139 resolvedspec: optimized_spec,
3140 adaptive_diagnostics: optimized.adaptive_diagnostics,
3141 kappa_timing: Some(kappa_timing),
3142 };
3143
3144 Ok(Some(optimized_result))
3145}
3146
3147const SPATIAL_COLLAPSE_EDF_FLOOR: f64 = 2.5;
3151
3152const SPATIAL_COLLAPSE_EDF_MARGIN: f64 = 1.0;
3157
3158fn fit_frozen_baseline_geometry(
3194 data: ArrayView2<'_, f64>,
3195 y: ArrayView1<'_, f64>,
3196 weights: ArrayView1<'_, f64>,
3197 offset: ArrayView1<'_, f64>,
3198 resolvedspec: &TermCollectionSpec,
3199 best: &FittedTermCollection,
3200 family: LikelihoodSpec,
3201 options: &FitOptions,
3202 baseline_score: f64,
3203 kappa_timing: Option<SpatialLengthScaleOptimizationTiming>,
3204) -> Result<FittedTermCollectionWithSpec, EstimationError> {
3205 let baseline = fit_term_collection_forspecwith_heuristic_lambdas(
3206 data,
3207 y,
3208 weights,
3209 offset,
3210 resolvedspec,
3211 best.fit.lambdas.as_slice(),
3212 family.clone(),
3213 options,
3214 )?;
3215 let best_edf = best.fit.inference.as_ref().map(|inf| inf.edf_total);
3220 let baseline_edf = baseline.fit.inference.as_ref().map(|inf| inf.edf_total);
3221 let baseline = match (best_edf, baseline_edf) {
3222 (Some(best_edf), Some(base_edf))
3223 if base_edf < SPATIAL_COLLAPSE_EDF_FLOOR
3224 && best_edf >= base_edf + SPATIAL_COLLAPSE_EDF_MARGIN =>
3225 {
3226 log::info!(
3227 "[spatial-kappa] warm-started frozen baseline collapsed (edf={base_edf:.3}) \
3228 below the certified baseline (edf={best_edf:.3}); refitting from scratch",
3229 );
3230 fit_term_collection_forspec(data, y, weights, offset, resolvedspec, family, options)?
3231 }
3232 _ => baseline,
3233 };
3234 let mut fit = baseline.fit;
3235 fit.reml_score = baseline_score;
3236 Ok(FittedTermCollectionWithSpec {
3237 fit,
3238 design: baseline.design,
3239 resolvedspec: resolvedspec.clone(),
3240 adaptive_diagnostics: baseline.adaptive_diagnostics,
3241 kappa_timing,
3242 })
3243}
3244
3245fn fit_data_adaptive_geometry(
3271 data: ArrayView2<'_, f64>,
3272 y: ArrayView1<'_, f64>,
3273 weights: ArrayView1<'_, f64>,
3274 offset: ArrayView1<'_, f64>,
3275 resolvedspec: &TermCollectionSpec,
3276 spatial_terms: &[usize],
3277 dims_per_term: &[usize],
3278 theta_seed: &Array1<f64>,
3279 theta_lower: &Array1<f64>,
3280 theta_upper: &Array1<f64>,
3281 rho_dim: usize,
3282 family: LikelihoodSpec,
3283 options: &FitOptions,
3284 baseline_score: f64,
3285 kappa_timing: Option<SpatialLengthScaleOptimizationTiming>,
3286) -> Result<Option<FittedTermCollectionWithSpec>, EstimationError> {
3287 let dims = dims_per_term.to_vec();
3288 let seed = SpatialLogKappaCoords::from_theta_tail_with_dims(theta_seed, rho_dim, dims.clone());
3289 let lower = SpatialLogKappaCoords::from_theta_tail_with_dims(theta_lower, rho_dim, dims.clone());
3290 let upper = SpatialLogKappaCoords::from_theta_tail_with_dims(theta_upper, rho_dim, dims.clone());
3291
3292 let mut values = seed.as_array().clone();
3293 let mut cursor = 0usize;
3294 let mut any_override = false;
3295 for (slot, &term_idx) in spatial_terms.iter().enumerate() {
3296 let d = dims[slot];
3297 let is_length_scale_term = constant_curvature_term_spec(resolvedspec, term_idx).is_none()
3301 && measure_jet_term_spec(resolvedspec, term_idx).is_none();
3302 if is_length_scale_term {
3303 for off in 0..d {
3304 let lo = lower.as_array()[cursor + off];
3305 let hi = upper.as_array()[cursor + off];
3306 if lo.is_finite() && hi.is_finite() && lo < hi {
3307 values[cursor + off] = 0.5 * (lo + hi);
3308 any_override = true;
3309 }
3310 }
3311 }
3312 cursor += d;
3313 }
3314 if !any_override {
3315 return Ok(None);
3316 }
3317 let coords = SpatialLogKappaCoords::new_with_dims(values, dims);
3318 let adaptive_spec = coords.apply_tospec(resolvedspec, spatial_terms)?;
3319 let adaptive =
3320 fit_term_collection_forspec(data, y, weights, offset, &adaptive_spec, family, options)?;
3321 let mut fit = adaptive.fit;
3328 fit.reml_score = baseline_score;
3329 Ok(Some(FittedTermCollectionWithSpec {
3330 fit,
3331 design: adaptive.design,
3332 resolvedspec: adaptive_spec,
3333 adaptive_diagnostics: adaptive.adaptive_diagnostics,
3334 kappa_timing,
3335 }))
3336}
3337
3338#[derive(Clone, Copy, PartialEq, Eq, Debug)]
3350enum SpatialHyperKind {
3351 Anisotropic,
3352 Isotropic,
3353}
3354
3355impl SpatialHyperKind {
3356 fn label(self) -> &'static str {
3359 match self {
3360 SpatialHyperKind::Anisotropic => "spatial-aniso-joint",
3361 SpatialHyperKind::Isotropic => "spatial-iso-joint",
3362 }
3363 }
3364
3365 fn adjective(self) -> &'static str {
3367 match self {
3368 SpatialHyperKind::Anisotropic => "anisotropic",
3369 SpatialHyperKind::Isotropic => "isotropic",
3370 }
3371 }
3372
3373 fn coord_name(self) -> &'static str {
3376 match self {
3377 SpatialHyperKind::Anisotropic => "psi",
3378 SpatialHyperKind::Isotropic => "kappa",
3379 }
3380 }
3381}
3382
3383struct SpatialFrozenGlmInputs {
3389 y: Array1<f64>,
3390 weights: Array1<f64>,
3391 offset: Array1<f64>,
3392 family: LikelihoodSpec,
3393}
3394
3395fn frozen_glm_tensor_eligible_family(family: &LikelihoodSpec) -> bool {
3412 !family.is_gaussian_identity()
3413 && matches!(
3414 &family.response,
3415 ResponseFamily::Binomial
3416 | ResponseFamily::Poisson
3417 | ResponseFamily::Gamma
3418 | ResponseFamily::NegativeBinomial { .. }
3419 )
3420}
3421
3422struct SpatialJointContext<'d> {
3423 data: ArrayView2<'d, f64>,
3424 rho_dim: usize,
3425 kind: SpatialHyperKind,
3426 cache: SingleBlockExactJointDesignCache<'d>,
3427 evaluator: gam_solve::estimate::ExternalJointHyperEvaluator<'d>,
3428 frozen_glm_inputs: Option<SpatialFrozenGlmInputs>,
3429 frozen_glm_psi_bounds: Option<(f64, f64)>,
3430 frozen_glm_tensor: Option<gam_solve::glm_sufficient_lane::FrozenWeightGramTensor>,
3431 frozen_glm_tensor_attempted: bool,
3432 frozen_glm_weight_memo: Option<(Array1<f64>, Array1<f64>)>,
3444}
3445
3446#[derive(Clone, Copy, Debug, Default)]
3447struct NfreeSkipGateStatus {
3448 shape: bool,
3449 value: bool,
3450 gradient: bool,
3451 penalty: bool,
3452 revision: bool,
3453 second_order: bool,
3454}
3455
3456impl NfreeSkipGateStatus {
3457 fn would_skip(self, require_gradient: bool) -> bool {
3458 self.shape
3459 && self.value
3460 && (!require_gradient || self.gradient)
3461 && self.penalty
3462 && self.revision
3463 && !self.second_order
3464 }
3465}
3466
3467impl<'d> SpatialJointContext<'d> {
3468 fn nfree_skip_gate_status(
3469 &self,
3470 theta: &Array1<f64>,
3471 allow_second_order: bool,
3472 require_gradient: bool,
3473 ) -> NfreeSkipGateStatus {
3474 let shape = theta.len() == self.rho_dim + 1;
3475 let (value, gradient) = if shape {
3476 let psi = theta[self.rho_dim];
3477 (
3478 self.evaluator.psi_gram_tensor_covers(psi)
3479 && self.evaluator.psi_gram_tensor_covers_skip(psi),
3480 !require_gradient || self.evaluator.psi_gram_tensor_covers_gradient(psi),
3481 )
3482 } else {
3483 (false, false)
3484 };
3485 NfreeSkipGateStatus {
3486 shape,
3487 value,
3488 gradient,
3489 penalty: self.evaluator.supports_nfree_penalty_rekey(),
3490 revision: self.evaluator.nfree_fast_path_revision().is_some(),
3491 second_order: allow_second_order,
3492 }
3493 }
3494
3495 fn frozen_glm_working_state(
3496 &self,
3497 beta: &Array1<f64>,
3498 ) -> Result<Option<(Array1<f64>, Array1<f64>)>, EstimationError> {
3499 let Some(inputs) = self.frozen_glm_inputs.as_ref() else {
3500 return Ok(None);
3501 };
3502 if beta.len() != self.cache.design().design.ncols() {
3503 return Ok(None);
3504 }
3505 let mut eta = self.cache.design().design.matrixvectormultiply(beta);
3506 if eta.len() != inputs.offset.len() {
3507 crate::bail_invalid_estim!(
3508 "frozen GLM tensor warm-state row mismatch: eta={}, offset={}",
3509 eta.len(),
3510 inputs.offset.len()
3511 );
3512 }
3513 eta += &inputs.offset;
3514 let obs = evaluate_standard_familyobservations(
3515 inputs.family.clone(),
3516 None,
3517 None,
3518 None,
3519 &inputs.y,
3520 &inputs.weights,
3521 &eta,
3522 )?;
3523 let mut working_response = obs.eta.clone();
3524 for i in 0..working_response.len() {
3525 let wi = obs.fisherweight[i].max(1e-12);
3526 working_response[i] += obs.score[i] / wi;
3527 }
3528 Ok(Some((obs.fisherweight, working_response)))
3529 }
3530
3531 fn frozen_glm_trial_weights(
3540 &mut self,
3541 beta: &Array1<f64>,
3542 ) -> Result<Option<Array1<f64>>, EstimationError> {
3543 if let Some((memo_beta, memo_w)) = self.frozen_glm_weight_memo.as_ref()
3544 && memo_beta.len() == beta.len()
3545 && memo_beta
3546 .iter()
3547 .zip(beta.iter())
3548 .all(|(a, b)| a.to_bits() == b.to_bits())
3549 {
3550 return Ok(Some(memo_w.clone()));
3551 }
3552 match self.frozen_glm_working_state(beta)? {
3553 Some((current_w, _)) => {
3554 self.frozen_glm_weight_memo = Some((beta.clone(), current_w.clone()));
3555 Ok(Some(current_w))
3556 }
3557 None => Ok(None),
3558 }
3559 }
3560
3561 fn ensure_frozen_glm_tensor(
3562 &mut self,
3563 theta: &Array1<f64>,
3564 warm_beta: Option<&Array1<f64>>,
3565 ) -> Result<(), EstimationError> {
3566 if self.frozen_glm_tensor.is_some() || self.frozen_glm_tensor_attempted {
3567 return Ok(());
3568 }
3569 let Some((psi_lo, psi_hi)) = self.frozen_glm_psi_bounds else {
3570 return Ok(());
3571 };
3572 if theta.len() != self.rho_dim + 1 {
3573 self.frozen_glm_tensor_attempted = true;
3574 return Ok(());
3575 }
3576 let Some(beta) = warm_beta else {
3577 return Ok(());
3578 };
3579 let Some((frozen_w, working_z)) = self.frozen_glm_working_state(beta)? else {
3580 self.frozen_glm_tensor_attempted = true;
3581 return Ok(());
3582 };
3583 let theta_probe_base = theta.clone();
3584 let rho_dim = self.rho_dim;
3585 let Self {
3592 cache, evaluator, ..
3593 } = self;
3594 let tensor = evaluator.build_frozen_glm_gram_tensor(
3595 |psi| {
3596 let mut theta_probe = theta_probe_base.clone();
3597 theta_probe[rho_dim] = psi;
3598 cache.ensure_theta(&theta_probe)?;
3599 Ok(cache.design().design.clone())
3600 },
3601 frozen_w.view(),
3602 working_z.view(),
3603 psi_lo,
3604 psi_hi,
3605 );
3606 self.cache
3607 .ensure_theta(theta)
3608 .map_err(EstimationError::InvalidInput)?;
3609 self.frozen_glm_tensor_attempted = true;
3610 if let Some(tensor) = tensor {
3611 self.frozen_glm_tensor = Some(tensor);
3612 log::info!(
3613 "[STAGE] {} certified frozen-W GLM ψ tensor over [{psi_lo:.3}, {psi_hi:.3}]",
3614 self.kind.label(),
3615 );
3616 } else {
3617 log::info!(
3618 "[STAGE] {} frozen-W GLM ψ tensor did not certify over [{psi_lo:.3}, {psi_hi:.3}]",
3619 self.kind.label(),
3620 );
3621 }
3622 Ok(())
3623 }
3624
3625 fn stage_frozen_glm_trial_statistics(
3626 &mut self,
3627 theta: &Array1<f64>,
3628 warm_beta: Option<&Array1<f64>>,
3629 allow_gradient: bool,
3630 ) -> Result<(), EstimationError> {
3631 let kind = self.kind;
3632 let mut staged_gram: Option<Array2<f64>> = None;
3633 let mut staged_deriv: Option<(Array2<f64>, Array1<f64>)> = None;
3634 if theta.len() == self.rho_dim + 1 {
3635 let psi = theta[self.rho_dim];
3636 let tensor_covers = self
3643 .frozen_glm_tensor
3644 .as_ref()
3645 .is_some_and(|t| t.contains(psi));
3646 let current_w = if tensor_covers {
3647 match warm_beta {
3648 Some(beta) => self.frozen_glm_trial_weights(beta)?,
3649 None => None,
3650 }
3651 } else {
3652 None
3653 };
3654 if let (Some(tensor), Some(current_w)) =
3655 (self.frozen_glm_tensor.as_ref(), current_w.as_ref())
3656 {
3657 const FROZEN_GLM_WEIGHT_DRIFT_RTOL: f64 = 1e-3;
3658 if tensor.weight_drift_within(current_w.view(), FROZEN_GLM_WEIGHT_DRIFT_RTOL) {
3659 staged_gram = Some(tensor.gram_at(psi));
3660 log::debug!(
3661 "[STAGE] {} trial at psi={psi:.6}: serving frozen-W GLM \
3662 first-Fisher-step XᵀWX n-free (weight drift within tol)",
3663 kind.label(),
3664 );
3665 }
3666 if allow_gradient
3667 && tensor.contains_for_gradient(psi)
3668 && let Some((dgram_dpsi, drhs_dpsi)) =
3669 tensor.gradient_pair_if_sound(psi, current_w.view())
3670 {
3671 staged_deriv = Some((dgram_dpsi, drhs_dpsi));
3672 log::debug!(
3673 "[STAGE] {} trial at psi={psi:.6}: serving frozen-W GLM \
3674 ψ-gradient (∂G/∂ψ, ∂b/∂ψ) n-free (gradient weight drift within \
3675 tight tol); B_j stays exact",
3676 kind.label(),
3677 );
3678 }
3679 }
3680 }
3681 self.evaluator.stage_glm_first_step_gram(staged_gram);
3682 self.evaluator.stage_glm_psi_gram_deriv(staged_deriv);
3683 Ok(())
3684 }
3685
3686 fn eval_full(
3688 &mut self,
3689 theta: &Array1<f64>,
3690 order: gam_solve::rho_optimizer::OuterEvalOrder,
3691 analytic_outer_hessian_available: bool,
3692 ) -> Result<
3693 (
3694 f64,
3695 Array1<f64>,
3696 gam_problem::HessianResult,
3697 ),
3698 EstimationError,
3699 > {
3700 use gam_solve::rho_optimizer::OuterEvalOrder;
3701 let allow_second_order = matches!(order, OuterEvalOrder::ValueGradientHessian)
3702 && analytic_outer_hessian_available;
3703 if let Some(eval) = self.cache.memoized_eval(theta) {
3704 let cached_satisfies_order = !allow_second_order || eval.2.is_analytic();
3705 if cached_satisfies_order {
3706 return Ok(eval);
3707 }
3708 }
3709 let kind = self.kind;
3710 let nfree_fast_path_revision = self.evaluator.nfree_fast_path_revision();
3746 let skip_design_realization = !allow_second_order && theta.len() == self.rho_dim + 1 && {
3747 let psi = theta[self.rho_dim];
3748 self.evaluator.psi_gram_tensor_covers(psi)
3749 && self.evaluator.psi_gram_tensor_covers_gradient(psi)
3756 && self.evaluator.psi_gram_tensor_covers_skip(psi)
3773 && self.evaluator.supports_nfree_penalty_rekey()
3778 && nfree_fast_path_revision.is_some()
3779 };
3780 if skip_design_realization {
3792 log::debug!(
3793 "[STAGE] {} eval_full at psi={:.6}: skipping n×k design re-realization \
3794 + reconditioning — criterion/gradient/inner-solve served n-free from \
3795 the certified ψ-gram tensor (GaussianFixedCache + k-space ψ-derivatives)",
3796 kind.label(),
3797 theta[self.rho_dim],
3798 );
3799 } else {
3800 self.cache
3801 .ensure_theta(theta)
3802 .map_err(EstimationError::InvalidInput)?;
3803 }
3804 let warm_beta = self.evaluator.current_beta();
3805 self.ensure_frozen_glm_tensor(theta, warm_beta.as_ref())?;
3806 self.stage_frozen_glm_trial_statistics(theta, warm_beta.as_ref(), !allow_second_order)?;
3814 let hyper_dirs = if skip_design_realization {
3821 self.cache.nfree_tensor_gradient_hyper_dirs(theta)?
3822 } else {
3823 self.cache.hyper_dirs_for_current_design(self.data, kind)?
3824 };
3825
3826 let design_revision = if skip_design_realization {
3827 nfree_fast_path_revision
3828 } else {
3829 Some(self.cache.design_revision())
3830 };
3831 if self.evaluator.supports_nfree_penalty_rekey() {
3845 match self.cache.canonical_penalties_at(theta) {
3846 Ok(penalty) => self.evaluator.stage_fast_path_penalty(Some(penalty)),
3847 Err(e) => {
3848 log::warn!(
3849 "[STAGE] {} eval_full at psi={:.6}: exact n-free S(ψ) rebuild failed \
3850 ({e}); clearing stage (eval falls to slow path)",
3851 kind.label(),
3852 theta[self.rho_dim],
3853 );
3854 self.evaluator.stage_fast_path_penalty(None);
3855 }
3856 }
3857 }
3858 let eval = evaluate_joint_reml_outer_eval_at_theta(
3865 &mut self.evaluator,
3866 self.cache.design(),
3867 theta,
3868 self.rho_dim,
3869 hyper_dirs,
3870 warm_beta.as_ref().map(|b: &Array1<f64>| b.view()),
3871 if allow_second_order {
3872 order
3873 } else {
3874 OuterEvalOrder::ValueAndGradient
3875 },
3876 design_revision,
3877 );
3878 if let Ok(ref value) = eval {
3879 self.cache.store_eval_at(theta, value.clone());
3880 }
3881 eval
3882 }
3883
3884 fn eval_efs(
3885 &mut self,
3886 theta: &Array1<f64>,
3887 ) -> Result<gam_problem::EfsEval, EstimationError> {
3888 self.cache
3889 .ensure_theta(theta)
3890 .map_err(EstimationError::InvalidInput)?;
3891 let kind = self.kind;
3892 let hyper_dirs = try_build_spatial_log_kappa_hyper_dirs(
3893 self.data,
3894 self.cache.spec(),
3895 self.cache.design(),
3896 &self.cache.spatial_terms,
3897 )?
3898 .ok_or_else(|| {
3899 EstimationError::InvalidInput(format!(
3900 "failed to build {} hyper_dirs for exact-joint EFS",
3901 kind.adjective(),
3902 ))
3903 })?;
3904 let design_revision = Some(self.cache.design_revision());
3905 let warm_beta = self.evaluator.current_beta();
3906 evaluate_joint_reml_efs_at_theta(
3907 &mut self.evaluator,
3908 self.cache.design(),
3909 theta,
3910 self.rho_dim,
3911 hyper_dirs,
3912 warm_beta.as_ref().map(|b: &Array1<f64>| b.view()),
3913 design_revision,
3914 )
3915 }
3916
3917 fn eval_cost(&mut self, theta: &Array1<f64>) -> f64 {
3923 if let Some(cost) = self.cache.memoized_cost(theta) {
3924 return cost;
3925 }
3926 let probe_start = std::time::Instant::now();
3941 let psi_distance = self
3942 .cache
3943 .current_theta
3944 .as_ref()
3945 .filter(|reference| reference.len() == theta.len())
3946 .map(|reference| {
3947 reference
3948 .iter()
3949 .zip(theta.iter())
3950 .map(|(a, b)| (a - b) * (a - b))
3951 .sum::<f64>()
3952 .sqrt()
3953 })
3954 .unwrap_or(f64::NAN);
3955 let nfree_fast_path_revision = self.evaluator.nfree_fast_path_revision();
3969 let skip_value_realization = theta.len() == self.rho_dim + 1 && {
3970 let psi = theta[self.rho_dim];
3971 self.evaluator.psi_gram_tensor_covers(psi)
3972 && self.evaluator.supports_nfree_penalty_rekey()
4006 && nfree_fast_path_revision.is_some()
4007 };
4008 if theta.len() == self.rho_dim + 1
4009 && self.evaluator.has_psi_gram_tensor()
4010 && !self.evaluator.psi_gram_tensor_covers(theta[self.rho_dim])
4011 {
4012 self.cache.store_cost_at(theta, f64::INFINITY);
4013 return f64::INFINITY;
4014 }
4015 if !skip_value_realization && self.cache.ensure_theta(theta).is_err() {
4016 return f64::INFINITY;
4017 }
4018 if self.evaluator.supports_nfree_penalty_rekey() {
4024 match self.cache.canonical_penalties_at(theta) {
4025 Ok(penalty) => self.evaluator.stage_fast_path_penalty(Some(penalty)),
4026 Err(_) => self.evaluator.stage_fast_path_penalty(None),
4027 }
4028 }
4029 let warm_beta = self.evaluator.current_beta();
4030 if let Err(err) = self.ensure_frozen_glm_tensor(theta, warm_beta.as_ref()) {
4031 log::warn!(
4032 "[STAGE] {} value-probe at psi={:.6}: frozen-W GLM tensor setup failed ({err}); \
4033 falling back to exact streamed Gram",
4034 self.kind.label(),
4035 if theta.len() > self.rho_dim {
4036 theta[self.rho_dim]
4037 } else {
4038 f64::NAN
4039 },
4040 );
4041 self.evaluator.stage_glm_first_step_gram(None);
4042 self.evaluator.stage_glm_psi_gram_deriv(None);
4043 } else if let Err(err) =
4044 self.stage_frozen_glm_trial_statistics(theta, warm_beta.as_ref(), false)
4045 {
4046 log::warn!(
4047 "[STAGE] {} value-probe at psi={:.6}: frozen-W GLM staging failed ({err}); \
4048 falling back to exact streamed Gram",
4049 self.kind.label(),
4050 if theta.len() > self.rho_dim {
4051 theta[self.rho_dim]
4052 } else {
4053 f64::NAN
4054 },
4055 );
4056 self.evaluator.stage_glm_first_step_gram(None);
4057 self.evaluator.stage_glm_psi_gram_deriv(None);
4058 }
4059 let design_revision = if skip_value_realization {
4060 nfree_fast_path_revision
4061 } else {
4062 Some(self.cache.design_revision())
4063 };
4064 let cost_label = self.kind.label();
4065 let result = {
4066 let design = self.cache.design();
4067 self.evaluator.evaluate_cost_only(
4068 &design.design,
4069 &design.penalties,
4070 &design.nullspace_dims,
4071 design.linear_constraints.clone(),
4072 theta,
4073 self.rho_dim,
4074 warm_beta.as_ref().map(|b: &Array1<f64>| b.view()),
4075 cost_label,
4076 design_revision,
4077 )
4078 };
4079 match result {
4080 Ok(cost) => {
4081 log::debug!(
4082 "[STAGE] {cost_label} value-probe (order=Value): elapsed={:.3}s \
4083 cost={cost:.6e} trial_theta_distance={psi_distance:.3e}",
4084 probe_start.elapsed().as_secs_f64(),
4085 );
4086 self.cache.store_cost_at(theta, cost);
4087 cost
4088 }
4089 Err(_) => f64::INFINITY,
4090 }
4091 }
4092
4093 fn reset(&mut self) {
4094 self.cache.current_theta = None;
4095 self.cache.last_eval_theta = None;
4096 self.cache.last_cost = None;
4097 self.cache.last_eval = None;
4098 }
4099}
4100
4101enum SpatialJointOutcome {
4134 Optimized {
4138 theta_star: Array1<f64>,
4139 final_value: f64,
4140 },
4141 NonConverged {
4145 iterations: usize,
4146 final_value: f64,
4147 final_grad_norm: Option<f64>,
4148 },
4149}
4150
4151fn kphase_log_norms(theta: &Array1<f64>, rho_dim: usize) -> (f64, f64) {
4152 let theta_norm = theta.iter().map(|v| v * v).sum::<f64>().sqrt();
4153 let log_kappa_norm = theta
4154 .iter()
4155 .skip(rho_dim)
4156 .map(|v| v * v)
4157 .sum::<f64>()
4158 .sqrt();
4159 (theta_norm, log_kappa_norm)
4160}
4161
4162fn run_exact_joint_spatial_optimization(
4163 kind: SpatialHyperKind,
4164 data: ArrayView2<'_, f64>,
4165 y: ArrayView1<'_, f64>,
4166 weights: ArrayView1<'_, f64>,
4167 offset: ArrayView1<'_, f64>,
4168 resolvedspec: &TermCollectionSpec,
4169 baseline_design: &TermCollectionDesign,
4170 family: LikelihoodSpec,
4171 options: &FitOptions,
4172 spatial_terms: &[usize],
4173 dims_per_term: &[usize],
4174 theta0: &Array1<f64>,
4175 lower: &Array1<f64>,
4176 upper: &Array1<f64>,
4177 rho_dim: usize,
4178 kappa_options: &SpatialLengthScaleOptimizationOptions,
4179) -> Result<(SpatialJointOutcome, SpatialLengthScaleOptimizationTiming), EstimationError> {
4180 let label = kind.label();
4181 assert!(
4183 lower.len() == theta0.len() && upper.len() == theta0.len(),
4184 "spatial hyperparameter bounds must match theta length: lower_len={}, upper_len={}, theta_len={}",
4185 lower.len(),
4186 upper.len(),
4187 theta0.len()
4188 );
4189 assert!(
4190 baseline_design.smooth.terms.len() >= spatial_terms.len(),
4191 "baseline design must have at least one smooth term per spatial term: baseline_terms={}, spatial_terms={}",
4192 baseline_design.smooth.terms.len(),
4193 spatial_terms.len()
4194 );
4195 use gam_solve::rho_optimizer::OuterEvalOrder;
4196 use gam_problem::{DeclaredHessianForm, Derivative, OuterEval};
4197
4198 let theta_dim = theta0.len();
4199 let coord_dim = theta_dim - rho_dim;
4202 let analytic_outer_hessian_available =
4212 exact_joint_spatial_outer_hessian_available(&family, baseline_design);
4213 if !analytic_outer_hessian_available {
4214 log::info!(
4215 "[{label}] analytic outer Hessian unavailable for family/design; routing without second-order geometry (coord_dim={coord_dim})"
4216 );
4217 }
4218 let mut prefer_gradient_only = theta_dim > EXACT_JOINT_SECOND_ORDER_THETA_CAP;
4224 if prefer_gradient_only {
4225 log::info!(
4226 "[{label}] joint θ-dim {theta_dim} exceeds the exact pair-Hessian budget \
4227 ({EXACT_JOINT_SECOND_ORDER_THETA_CAP}); routing gradient-only quasi-Newton"
4228 );
4229 }
4230 let mut suppress_outer_hessian_for_nfree = false;
4240
4241 log::trace!(
4242 "[{}] starting analytic optimization: rho_dim={}, coord_dim={}, dims_per_term={:?}",
4243 label,
4244 rho_dim,
4245 coord_dim,
4246 dims_per_term,
4247 );
4248
4249 let mut ctx = SpatialJointContext {
4250 data,
4251 rho_dim,
4252 kind,
4253 cache: SingleBlockExactJointDesignCache::new(
4254 data,
4255 resolvedspec.clone(),
4256 baseline_design.clone(),
4257 spatial_terms.to_vec(),
4258 rho_dim,
4259 dims_per_term.to_vec(),
4260 )
4261 .map_err(EstimationError::InvalidInput)?,
4262 evaluator: gam_solve::estimate::ExternalJointHyperEvaluator::new(
4263 y,
4264 weights,
4265 &baseline_design.design,
4266 offset,
4267 &baseline_design.penalties,
4268 &external_opts_for_design(&family, baseline_design, options),
4269 label,
4270 )?,
4271 frozen_glm_inputs: if coord_dim == 1 && frozen_glm_tensor_eligible_family(&family) {
4272 Some(SpatialFrozenGlmInputs {
4273 y: y.to_owned(),
4274 weights: weights.to_owned(),
4275 offset: offset.to_owned(),
4276 family: family.clone(),
4277 })
4278 } else {
4279 None
4280 },
4281 frozen_glm_psi_bounds: if coord_dim == 1 && frozen_glm_tensor_eligible_family(&family) {
4282 Some((lower[rho_dim], upper[rho_dim]))
4283 } else {
4284 None
4285 },
4286 frozen_glm_tensor: None,
4287 frozen_glm_tensor_attempted: false,
4288 frozen_glm_weight_memo: None,
4289 };
4290
4291 let mut psi_rank_stable_floor: Option<f64> = None;
4314 let mut psi_rank_stable_ceiling: Option<f64> = None;
4323 let nfree_penalty_capable = coord_dim == 1
4324 && family.is_gaussian_identity()
4325 && ctx.cache.supports_nfree_penalty_rekey();
4326 if nfree_penalty_capable {
4327 let psi_lo = lower[rho_dim];
4328 let psi_hi = upper[rho_dim];
4329 let z = Array1::from_iter(y.iter().zip(offset.iter()).map(|(yi, oi)| yi - oi));
4330 let theta_probe_base = theta0.clone();
4331 let SpatialJointContext {
4334 cache, evaluator, ..
4335 } = &mut ctx;
4336 let attached = evaluator.build_and_set_psi_gram_tensor(
4337 |psi| {
4338 let mut theta_probe = theta_probe_base.clone();
4339 theta_probe[rho_dim] = psi;
4340 cache.ensure_theta(&theta_probe)?;
4341 Ok(cache.design().design.clone())
4342 },
4343 weights,
4344 z.view(),
4345 psi_lo,
4346 psi_hi,
4347 );
4348 if attached {
4349 log::info!(
4350 "[{label}] certified ψ-gram tensor over [{psi_lo:.3}, {psi_hi:.3}]: \
4351 in-window trials assemble Gaussian sufficient statistics n-free"
4352 );
4353 let psi_anchor = theta0[rho_dim];
4358 psi_rank_stable_floor = evaluator
4359 .psi_gram_rank_stable_floor(psi_anchor)
4360 .filter(|&f| f.is_finite() && f > psi_lo && f < psi_anchor);
4361 log::info!(
4362 "[KAPPA-PHASE-FLOOR] n_rows={} psi_lo={psi_lo:.6} psi_anchor={psi_anchor:.6} \
4363 rank_stable_floor={:?} lifted={}",
4364 data.nrows(),
4365 evaluator.psi_gram_rank_stable_floor(psi_anchor),
4366 psi_rank_stable_floor.is_some(),
4367 );
4368 if let Some(floor) = psi_rank_stable_floor {
4369 log::info!(
4370 "[{label}] rank-stable κ-floor ψ_floor={floor:.6} > window floor \
4371 ψ_lo={psi_lo:.6}: lifting the optimizer lower bound to keep every \
4372 in-window trial on the n-free design-realization skip (#1033). The \
4373 conditioned Gram is rank-deficient below ψ_floor (longest-length-scale \
4374 radial mode collapses into the nullspace), where the skip is soundly \
4375 refused; that band drifts with n via the sample-std standardization, \
4376 so this n-free k-space floor is the n-independent fix."
4377 );
4378 }
4379 psi_rank_stable_ceiling = evaluator
4388 .psi_gram_rank_stable_ceiling(psi_anchor)
4389 .filter(|&c| c.is_finite() && c < psi_hi && c > psi_anchor);
4390 log::info!(
4391 "[KAPPA-PHASE-CEIL] n_rows={} psi_hi={psi_hi:.6} psi_anchor={psi_anchor:.6} \
4392 rank_stable_ceiling={:?} clamped={}",
4393 data.nrows(),
4394 evaluator.psi_gram_rank_stable_ceiling(psi_anchor),
4395 psi_rank_stable_ceiling.is_some(),
4396 );
4397 if let Some(ceiling) = psi_rank_stable_ceiling {
4398 log::info!(
4399 "[{label}] rank-stable κ-ceiling ψ_ceil={ceiling:.6} < window ceiling \
4400 ψ_hi={psi_hi:.6}: clamping the optimizer upper bound to keep every \
4401 in-window trial on the n-free design-realization skip (#1033). The \
4402 conditioned Gram is rank-deficient above ψ_ceil (longest-frequency \
4403 radial mode goes collinear), where the skip is soundly refused; a \
4404 line-search overshoot there trips the O(n) reset_surface lane (and the \
4405 deficient pinning ψ it records resets the next in-band trial too)."
4406 );
4407 }
4408 let gradient_covers_full_window = evaluator.psi_gram_tensor_covers_gradient(psi_lo)
4409 && evaluator.psi_gram_tensor_covers_gradient(psi_hi);
4410 if gradient_covers_full_window {
4411 log::info!(
4412 "[{label}] certified ψ-gram tensor gradient lane covers the full \
4413 optimizer window [{psi_lo:.3}, {psi_hi:.3}]"
4414 );
4415 } else {
4416 log::info!(
4417 "[{label}] ψ-gram tensor value lane certified, but the gradient lane \
4418 does not cover the full optimizer window [{psi_lo:.3}, {psi_hi:.3}]; \
4419 keeping exact streamed kappa routing"
4420 );
4421 }
4422 evaluator.set_supports_nfree_penalty_rekey(true);
4442 log::info!(
4443 "[{label}] exact n-free ψ-penalty re-key enabled over [{psi_lo:.3}, \
4444 {psi_hi:.3}]: in-window fast-path trials rebuild S(ψ) n-free from frozen \
4445 geometry (no reset_surface)"
4446 );
4447 } else {
4448 log::info!(
4449 "[{label}] ψ-gram tensor did not certify over [{psi_lo:.3}, {psi_hi:.3}]; \
4450 keeping the exact per-trial path"
4451 );
4452 }
4453 if attached
4474 && evaluator.psi_gram_tensor_covers_gradient(psi_lo)
4475 && evaluator.psi_gram_tensor_covers_gradient(psi_hi)
4476 && evaluator.supports_nfree_penalty_rekey()
4477 && cache.supports_nfree_gradient_only_routing()
4478 {
4479 suppress_outer_hessian_for_nfree = true;
4480 prefer_gradient_only = true;
4481 log::info!(
4482 "[{label}] n-free Gaussian ψ-lane armed; suppressing the analytic outer \
4483 Hessian and routing gradient-only (BFGS) so the κ outer loop never realizes \
4484 the O(n) second-order slab — n-independent outer loop (#1033)"
4485 );
4486 }
4487 } else if coord_dim == 1 && family.is_gaussian_identity() {
4488 log::info!(
4489 "[{label}] exact n-free ψ-penalty re-key unavailable; skipping ψ-gram tensor \
4490 attachment so value, gradient, and Hessian remain on the same exact streamed \
4491 objective"
4492 );
4493 }
4494
4495 const OUTER_FD_AUDIT_MAX_N: usize = 4_000; const OUTER_FD_AUDIT_MAX_THETA_DIM: usize = 32; let n_total = data.nrows();
4523 let outer_fd_audit_eligible = log::log_enabled!(log::Level::Info) && analytic_outer_hessian_available && n_total <= OUTER_FD_AUDIT_MAX_N && theta_dim <= OUTER_FD_AUDIT_MAX_THETA_DIM; log::info!(
4528 "[OUTER-FD-AUDIT/spatial-exact-joint] gate eligible={outer_fd_audit_eligible} \
4529 analytic_grad={analytic_outer_hessian_available} n_total={n_total} \
4530 theta_dim={theta_dim} rho_dim={rho_dim} psi_dim={coord_dim}"
4531 );
4532 if outer_fd_audit_eligible {
4533 let audit = (|| -> Result<gam_solve::rho_optimizer::OuterGradientFdAudit, String> {
4535 let mut eval_at = |theta: &Array1<f64>,
4536 mode: gam_solve::estimate::reml::reml_outer_engine::EvalMode|
4537 -> Result<
4538 (
4539 f64,
4540 Array1<f64>,
4541 gam_problem::HessianResult,
4542 ),
4543 String,
4544 > {
4545 use gam_solve::estimate::reml::reml_outer_engine::EvalMode;
4546 let order = if matches!(mode, EvalMode::ValueGradientHessian) {
4547 OuterEvalOrder::ValueGradientHessian
4548 } else {
4549 OuterEvalOrder::Value
4550 };
4551 ctx.eval_full(theta, order, analytic_outer_hessian_available)
4552 .map_err(|e| format!("fd-audit eval_full: {e}"))
4553 };
4554 let rho_dim_audit = rho_dim;
4555 let label_fn = move |i: usize| -> String {
4556 if i < rho_dim_audit {
4557 format!("rho[{i}]")
4558 } else {
4559 format!("psi_kappa[{}]", i - rho_dim_audit)
4560 }
4561 };
4562 gam_solve::rho_optimizer::outer_gradient_fd_audit(
4563 theta0,
4565 1e-4,
4566 label_fn,
4567 &mut eval_at,
4568 )
4569 })();
4570 match audit {
4572 Ok(audit) => audit.log_verdict("spatial-exact-joint"),
4573 Err(e) => log::warn!("[OUTER-FD-AUDIT/spatial-exact-joint] skipped: {e}"),
4574 }
4575 }
4576
4577 let kphase_prime_order = if analytic_outer_hessian_available && !suppress_outer_hessian_for_nfree {
4578 OuterEvalOrder::ValueGradientHessian
4579 } else {
4580 OuterEvalOrder::ValueAndGradient
4581 };
4582 let kphase_prime_start = std::time::Instant::now();
4583 drop(ctx.eval_full(theta0, kphase_prime_order, analytic_outer_hessian_available)?);
4584 log::info!(
4585 "[KAPPA-PHASE-PRIME] n_rows={} order={:?} elapsed_s={:.4} slow_path_resets_total={} design_revision={}",
4586 data.nrows(),
4587 kphase_prime_order,
4588 kphase_prime_start.elapsed().as_secs_f64(),
4589 ctx.evaluator.slow_path_reset_count(),
4590 ctx.cache.design_revision(),
4591 );
4592
4593 let kphase_cost_calls = std::cell::Cell::new(0usize);
4594 let kphase_eval_calls = std::cell::Cell::new(0usize);
4595 let kphase_efs_calls = std::cell::Cell::new(0usize);
4596 let kphase_cost_total_s = std::cell::Cell::new(0.0);
4597 let kphase_eval_total_s = std::cell::Cell::new(0.0);
4598 let kphase_efs_total_s = std::cell::Cell::new(0.0);
4599 let kphase_nfree_miss_shape = std::cell::Cell::new(0u64);
4600 let kphase_nfree_miss_value = std::cell::Cell::new(0u64);
4601 let kphase_nfree_miss_gradient = std::cell::Cell::new(0u64);
4602 let kphase_nfree_miss_penalty = std::cell::Cell::new(0u64);
4603 let kphase_nfree_miss_revision = std::cell::Cell::new(0u64);
4604 let kphase_nfree_miss_second_order = std::cell::Cell::new(0u64);
4605 let kphase_nfree_miss_other = std::cell::Cell::new(0u64);
4606 let kphase_optim_start = std::time::Instant::now();
4607 let kphase_log_kappa_dim = coord_dim;
4608 let kphase_slow_resets_start = ctx.evaluator.slow_path_reset_count();
4609 let kphase_design_revision_start = ctx.cache.design_revision();
4610 let kphase_nfree_skip_touches_start = gam_solve::pirls::nfree_skip_row_element_touches();
4614
4615 let lower_effective: std::borrow::Cow<'_, Array1<f64>> = match psi_rank_stable_floor {
4622 Some(floor) if coord_dim == 1 && floor > lower[rho_dim] => {
4623 let mut lifted = lower.clone();
4624 lifted[rho_dim] = floor;
4625 std::borrow::Cow::Owned(lifted)
4626 }
4627 _ => std::borrow::Cow::Borrowed(lower),
4628 };
4629 let lower = lower_effective.as_ref();
4630
4631 let upper_effective: std::borrow::Cow<'_, Array1<f64>> = match psi_rank_stable_ceiling {
4639 Some(ceiling) if coord_dim == 1 && ceiling < upper[rho_dim] => {
4640 let mut clamped = upper.clone();
4641 clamped[rho_dim] = ceiling;
4642 std::borrow::Cow::Owned(clamped)
4643 }
4644 _ => std::borrow::Cow::Borrowed(upper),
4645 };
4646 let upper = upper_effective.as_ref();
4647
4648 let problem = exact_joint_multistart_outer_problem(
4649 theta0,
4650 lower,
4651 upper,
4652 rho_dim,
4653 coord_dim,
4654 theta_dim,
4655 Derivative::Analytic,
4656 if analytic_outer_hessian_available && !suppress_outer_hessian_for_nfree {
4657 DeclaredHessianForm::Either
4658 } else {
4659 DeclaredHessianForm::Unavailable
4664 },
4665 prefer_gradient_only,
4666 suppress_outer_hessian_for_nfree,
4677 seed_risk_profile_for_likelihood_family(&family),
4678 kappa_options.rel_tol.max(1e-6),
4679 kappa_options.max_outer_iter.max(1),
4680 Some(5.0),
4684 Some(kappa_options.log_step.clamp(0.25, 1.0)),
4686 None,
4687 Some((data.nrows(), baseline_design.design.ncols())),
4692 !constant_curvature_term_indices(resolvedspec).is_empty(),
4696 );
4697
4698 let eval_outer = |ctx: &mut &mut SpatialJointContext<'_>,
4699 theta: &Array1<f64>,
4700 order: OuterEvalOrder|
4701 -> Result<OuterEval, EstimationError> {
4702 let t0 = std::time::Instant::now();
4703 let allow_second_order_for_call = matches!(order, OuterEvalOrder::ValueGradientHessian)
4704 && analytic_outer_hessian_available;
4705 let gate = ctx.nfree_skip_gate_status(theta, allow_second_order_for_call, true);
4706 let resets_before = ctx.evaluator.slow_path_reset_count();
4707 let raw = ctx.eval_full(theta, order, analytic_outer_hessian_available);
4708 let reset_delta = ctx
4709 .evaluator
4710 .slow_path_reset_count()
4711 .saturating_sub(resets_before);
4712 if reset_delta > 0 {
4713 if !gate.shape {
4714 kphase_nfree_miss_shape.set(kphase_nfree_miss_shape.get() + reset_delta);
4715 }
4716 if gate.shape && !gate.value {
4717 kphase_nfree_miss_value.set(kphase_nfree_miss_value.get() + reset_delta);
4718 }
4719 if gate.shape && gate.value && !gate.gradient {
4720 kphase_nfree_miss_gradient.set(kphase_nfree_miss_gradient.get() + reset_delta);
4721 }
4722 if gate.shape && gate.value && gate.gradient && !gate.penalty {
4723 kphase_nfree_miss_penalty.set(kphase_nfree_miss_penalty.get() + reset_delta);
4724 }
4725 if gate.shape && gate.value && gate.gradient && gate.penalty && !gate.revision {
4726 kphase_nfree_miss_revision.set(kphase_nfree_miss_revision.get() + reset_delta);
4727 }
4728 if gate.shape
4729 && gate.value
4730 && gate.gradient
4731 && gate.penalty
4732 && gate.revision
4733 && gate.second_order
4734 {
4735 kphase_nfree_miss_second_order
4736 .set(kphase_nfree_miss_second_order.get() + reset_delta);
4737 }
4738 if gate.would_skip(true) {
4739 kphase_nfree_miss_other.set(kphase_nfree_miss_other.get() + reset_delta);
4740 }
4741 }
4742 let elapsed_s = t0.elapsed().as_secs_f64();
4743 kphase_eval_calls.set(kphase_eval_calls.get() + 1);
4744 kphase_eval_total_s.set(kphase_eval_total_s.get() + elapsed_s);
4745 let (theta_norm, log_kappa_norm) = kphase_log_norms(theta, rho_dim);
4746 log::info!(
4747 "[KAPPA-PHASE] phase=eval_outer call={} order={:?} design_revision={:?} theta_norm={:.4e} log_kappa_norm={:.4e} elapsed_s={:.4}",
4748 kphase_eval_calls.get(),
4749 order,
4750 Some(ctx.cache.design_revision()),
4751 theta_norm,
4752 log_kappa_norm,
4753 elapsed_s,
4754 );
4755 match raw {
4756 Ok((cost, grad, hess)) => Ok(OuterEval {
4757 cost,
4758 gradient: grad,
4759 hessian: hess,
4760 inner_beta_hint: None,
4761 }),
4762 Err(err) if is_recoverable_trial_point_error(&err) => {
4770 log::debug!(
4771 "[{label}] trial point infeasible (kernel design \
4772 not constructible at theta={theta:?}): {err}; retreating",
4773 );
4774 Ok(OuterEval::infeasible(theta_dim))
4775 }
4776 Err(err) => Err(err),
4777 }
4778 };
4779
4780 let mut obj = problem.build_objective_with_eval_order(
4781 &mut ctx,
4782 |ctx: &mut &mut SpatialJointContext<'_>, theta: &Array1<f64>| {
4783 let t0 = std::time::Instant::now();
4784 let gate = ctx.nfree_skip_gate_status(theta, false, false);
4785 let resets_before = ctx.evaluator.slow_path_reset_count();
4786 let cost = ctx.eval_cost(theta);
4787 let reset_delta = ctx
4788 .evaluator
4789 .slow_path_reset_count()
4790 .saturating_sub(resets_before);
4791 if reset_delta > 0 {
4792 if !gate.shape {
4793 kphase_nfree_miss_shape.set(kphase_nfree_miss_shape.get() + reset_delta);
4794 }
4795 if gate.shape && !gate.value {
4796 kphase_nfree_miss_value.set(kphase_nfree_miss_value.get() + reset_delta);
4797 }
4798 if gate.shape && gate.value && !gate.penalty {
4799 kphase_nfree_miss_penalty.set(kphase_nfree_miss_penalty.get() + reset_delta);
4800 }
4801 if gate.shape && gate.value && gate.penalty && !gate.revision {
4802 kphase_nfree_miss_revision.set(kphase_nfree_miss_revision.get() + reset_delta);
4803 }
4804 if gate.would_skip(false) {
4805 kphase_nfree_miss_other.set(kphase_nfree_miss_other.get() + reset_delta);
4806 }
4807 }
4808 let elapsed_s = t0.elapsed().as_secs_f64();
4809 kphase_cost_calls.set(kphase_cost_calls.get() + 1);
4810 kphase_cost_total_s.set(kphase_cost_total_s.get() + elapsed_s);
4811 let (theta_norm, log_kappa_norm) = kphase_log_norms(theta, rho_dim);
4812 log::info!(
4813 "[KAPPA-PHASE] phase=cost call={} design_revision={:?} theta_norm={:.4e} log_kappa_norm={:.4e} elapsed_s={:.4}",
4814 kphase_cost_calls.get(),
4815 Some(ctx.cache.design_revision()),
4816 theta_norm,
4817 log_kappa_norm,
4818 elapsed_s,
4819 );
4820 Ok(cost)
4821 },
4822 |ctx: &mut &mut SpatialJointContext<'_>, theta: &Array1<f64>| {
4823 eval_outer(
4824 ctx,
4825 theta,
4826 if analytic_outer_hessian_available && !suppress_outer_hessian_for_nfree {
4836 OuterEvalOrder::ValueGradientHessian
4837 } else {
4838 OuterEvalOrder::ValueAndGradient
4839 },
4840 )
4841 },
4842 |ctx: &mut &mut SpatialJointContext<'_>, theta: &Array1<f64>, order: OuterEvalOrder| {
4843 eval_outer(ctx, theta, order)
4844 },
4845 Some(|ctx: &mut &mut SpatialJointContext<'_>| {
4846 ctx.reset();
4847 }),
4848 Some(|ctx: &mut &mut SpatialJointContext<'_>, theta: &Array1<f64>| {
4849 let t0 = std::time::Instant::now();
4850 let eval = ctx.eval_efs(theta);
4851 let elapsed_s = t0.elapsed().as_secs_f64();
4852 kphase_efs_calls.set(kphase_efs_calls.get() + 1);
4853 kphase_efs_total_s.set(kphase_efs_total_s.get() + elapsed_s);
4854 let (theta_norm, log_kappa_norm) = kphase_log_norms(theta, rho_dim);
4855 log::info!(
4856 "[KAPPA-PHASE] phase=efs call={} design_revision={:?} theta_norm={:.4e} log_kappa_norm={:.4e} elapsed_s={:.4}",
4857 kphase_efs_calls.get(),
4858 Some(ctx.cache.design_revision()),
4859 theta_norm,
4860 log_kappa_norm,
4861 elapsed_s,
4862 );
4863 eval
4864 }),
4865 );
4866
4867 let run_label = match kind {
4868 SpatialHyperKind::Anisotropic => "aniso-psi joint REML",
4869 SpatialHyperKind::Isotropic => "iso-kappa joint REML",
4870 };
4871 let result = problem.run(&mut obj, run_label).map_err(|e| {
4872 EstimationError::InvalidInput(format!(
4873 "{} analytic optimization failed after exhausting strategy fallbacks: {e}",
4874 kind.adjective(),
4875 ))
4876 })?;
4877 drop(obj);
4878 let kphase_total_s = kphase_optim_start.elapsed().as_secs_f64();
4879 let kphase_slow_resets = ctx
4880 .evaluator
4881 .slow_path_reset_count()
4882 .saturating_sub(kphase_slow_resets_start);
4883 let kphase_design_revision_delta = ctx
4884 .cache
4885 .design_revision()
4886 .saturating_sub(kphase_design_revision_start);
4887 let kphase_nfree_skip_touches = gam_solve::pirls::nfree_skip_row_element_touches()
4888 .saturating_sub(kphase_nfree_skip_touches_start);
4889 log::info!(
4890 "[KAPPA-PHASE-SUMMARY] n_rows={} log_kappa_dim={} n_cost={} cost_total_s={:.4} n_eval={} eval_total_s={:.4} n_efs={} efs_total_s={:.4} slow_path_resets={} design_revision_delta={} nfree_skip_row_touches={} nfree_miss_shape={} nfree_miss_value={} nfree_miss_gradient={} nfree_miss_penalty={} nfree_miss_revision={} nfree_miss_second_order={} nfree_miss_other={} optim_total_s={:.4}",
4891 data.nrows(),
4892 kphase_log_kappa_dim,
4893 kphase_cost_calls.get(),
4894 kphase_cost_total_s.get(),
4895 kphase_eval_calls.get(),
4896 kphase_eval_total_s.get(),
4897 kphase_efs_calls.get(),
4898 kphase_efs_total_s.get(),
4899 kphase_slow_resets,
4900 kphase_design_revision_delta,
4901 kphase_nfree_skip_touches,
4902 kphase_nfree_miss_shape.get(),
4903 kphase_nfree_miss_value.get(),
4904 kphase_nfree_miss_gradient.get(),
4905 kphase_nfree_miss_penalty.get(),
4906 kphase_nfree_miss_revision.get(),
4907 kphase_nfree_miss_second_order.get(),
4908 kphase_nfree_miss_other.get(),
4909 kphase_total_s,
4910 );
4911 let timing = SpatialLengthScaleOptimizationTiming {
4912 log_kappa_dim: kphase_log_kappa_dim,
4913 cost_calls: kphase_cost_calls.get(),
4914 cost_total_s: kphase_cost_total_s.get(),
4915 eval_calls: kphase_eval_calls.get(),
4916 eval_total_s: kphase_eval_total_s.get(),
4917 efs_calls: kphase_efs_calls.get(),
4918 efs_total_s: kphase_efs_total_s.get(),
4919 slow_path_resets: kphase_slow_resets,
4920 design_revision_delta: kphase_design_revision_delta,
4921 nfree_skip_row_touches: kphase_nfree_skip_touches,
4922 nfree_miss_shape: kphase_nfree_miss_shape.get(),
4923 nfree_miss_value: kphase_nfree_miss_value.get(),
4924 nfree_miss_gradient: kphase_nfree_miss_gradient.get(),
4925 nfree_miss_penalty: kphase_nfree_miss_penalty.get(),
4926 nfree_miss_revision: kphase_nfree_miss_revision.get(),
4927 nfree_miss_second_order: kphase_nfree_miss_second_order.get(),
4928 nfree_miss_other: kphase_nfree_miss_other.get(),
4929 optim_total_s: kphase_total_s,
4930 };
4931 if !result.converged {
4932 let rel_to_cost_threshold = options.tol * (1.0_f64 + result.final_value.abs());
4943 if let Some(final_grad) = result
4944 .final_grad_norm
4945 .filter(|v| v.is_finite() && *v <= rel_to_cost_threshold)
4946 {
4947 log::info!(
4948 "[{}] outer optimization hit max_iter={} but \
4949 projected gradient norm {:.3e} ≤ τ·(1+|f|) = {:.3e} \
4950 (τ={:.3e}, |f|={:.3e}); accepting iterate under the mgcv-style \
4951 relative-to-cost REML convergence criterion.",
4952 label,
4953 result.iterations,
4954 final_grad,
4955 rel_to_cost_threshold,
4956 options.tol,
4957 result.final_value.abs(),
4958 );
4959 } else if result.final_value.is_finite() {
4960 log::warn!(
4975 "[{}] {} did not converge after {} iterations \
4976 (final_objective={:.6e}, final_grad_norm={}); keeping the \
4977 frozen baseline geometry instead of aborting the fit.",
4978 label,
4979 kind.adjective(),
4980 result.iterations,
4981 result.final_value,
4982 result.final_grad_norm_report(),
4983 );
4984 return Ok((
4985 SpatialJointOutcome::NonConverged {
4986 iterations: result.iterations,
4987 final_value: result.final_value,
4988 final_grad_norm: result.final_grad_norm,
4989 },
4990 timing,
4991 ));
4992 } else {
4993 crate::bail_invalid_estim!(
4998 "{} analytic optimization diverged after {} iterations (final_objective={:.6e}, final_grad_norm={})",
4999 kind.adjective(),
5000 result.iterations,
5001 result.final_value,
5002 result.final_grad_norm_report(),
5003 );
5004 }
5005 }
5006 log::trace!(
5007 "[{}] converged in {} iterations, final_value={:.6e}, grad_norm={}",
5008 label,
5009 result.iterations,
5010 result.final_value,
5011 result.final_grad_norm_report(),
5012 );
5013 let theta_star = result.rho;
5017 Ok((
5018 SpatialJointOutcome::Optimized {
5019 theta_star,
5020 final_value: result.final_value,
5021 },
5022 timing,
5023 ))
5024}
5025
5026fn set_single_term_spatial_length_scale(
5030 term: &mut SmoothTermSpec,
5031 length_scale: f64,
5032) -> Result<(), EstimationError> {
5033 match &mut term.basis {
5034 SmoothBasisSpec::ThinPlate { spec, .. } => {
5035 spec.length_scale = length_scale;
5036 Ok(())
5037 }
5038 SmoothBasisSpec::Matern { spec, .. } => {
5039 spec.length_scale = length_scale;
5040 Ok(())
5041 }
5042 SmoothBasisSpec::Duchon { spec, .. } => {
5043 spec.length_scale = Some(length_scale);
5044 Ok(())
5045 }
5046 _ => Err(EstimationError::InvalidInput(format!(
5047 "term '{}' does not expose a spatial length scale",
5048 term.name
5049 ))),
5050 }
5051}
5052
5053fn set_single_term_spatial_aniso_log_scales(
5057 term: &mut SmoothTermSpec,
5058 eta: Vec<f64>,
5059) -> Result<(), EstimationError> {
5060 let eta = center_aniso_log_scales(&eta);
5061 match &mut term.basis {
5062 SmoothBasisSpec::Matern { spec, .. } => {
5063 spec.aniso_log_scales = Some(eta);
5064 Ok(())
5065 }
5066 SmoothBasisSpec::Duchon { spec, .. } => {
5067 spec.aniso_log_scales = Some(eta);
5068 Ok(())
5069 }
5070 _ => Err(EstimationError::InvalidInput(format!(
5071 "term '{}' does not support aniso_log_scales",
5072 term.name
5073 ))),
5074 }
5075}
5076
5077pub fn get_constant_curvature_kappa(spec: &TermCollectionSpec, term_idx: usize) -> Option<f64> {
5096 constant_curvature_term_spec(spec, term_idx).map(|cc| cc.kappa)
5097}
5098
5099pub fn constant_curvature_kappa_is_fixed(spec: &TermCollectionSpec, term_idx: usize) -> bool {
5107 constant_curvature_term_spec(spec, term_idx).is_some_and(|cc| cc.kappa_fixed)
5108}
5109
5110pub fn constant_curvature_term_indices(spec: &TermCollectionSpec) -> Vec<usize> {
5112 (0..spec.smooth_terms.len())
5113 .filter(|&idx| constant_curvature_term_spec(spec, idx).is_some())
5114 .collect()
5115}
5116
5117
5118#[derive(Debug, Clone)]
5119struct SingleSmoothTermRealization {
5120 design_local: DesignMatrix,
5121 term: SmoothTerm,
5122 dropped_penaltyinfo: Vec<DroppedPenaltyBlockInfo>,
5123}
5124
5125impl SingleSmoothTermRealization {
5126 fn active_penaltyinfo(&self) -> Vec<PenaltyInfo> {
5127 self.term
5128 .penaltyinfo_local
5129 .iter()
5130 .filter(|info| info.active)
5131 .cloned()
5132 .collect()
5133 }
5134}
5135
5136fn build_single_smooth_term_realization(
5137 data: ArrayView2<'_, f64>,
5138 termspec: &SmoothTermSpec,
5139) -> Result<SingleSmoothTermRealization, BasisError> {
5140 let raw = build_smooth_design(data, std::slice::from_ref(termspec))?;
5141 finish_single_smooth_term_realization(raw)
5142}
5143
5144fn finish_single_smooth_term_realization(
5145 raw: RawSmoothDesign,
5146) -> Result<SingleSmoothTermRealization, BasisError> {
5147 let RawSmoothDesign {
5148 term_designs,
5149 dropped_penaltyinfo,
5150 terms,
5151 ..
5152 } = raw;
5153 let term = terms.into_iter().next().ok_or_else(|| {
5154 BasisError::InvalidInput("single-term smooth build returned no term".to_string())
5155 })?;
5156 let design = term_designs.into_iter().next().ok_or_else(|| {
5157 BasisError::InvalidInput("single-term smooth build returned no term design".to_string())
5158 })?;
5159
5160 Ok(SingleSmoothTermRealization {
5161 design_local: design,
5162 term,
5163 dropped_penaltyinfo,
5164 })
5165}
5166
5167fn wrap_local_build_as_realization(
5174 mut local: LocalSmoothTermBuild,
5175 termspec: &SmoothTermSpec,
5176) -> Result<SingleSmoothTermRealization, String> {
5177 let p_local = local.dim;
5178 let lb_local = if local.box_reparam {
5179 shape_lower_bounds_local(termspec.shape, p_local)
5180 } else {
5181 None
5182 };
5183
5184 let active_count = local.penaltyinfo.iter().filter(|info| info.active).count();
5185 if active_count != local.penalties.len() {
5186 return Err(format!(
5187 "internal penalty info mismatch for term '{}': active_infos={}, penalties={}",
5188 termspec.name,
5189 active_count,
5190 local.penalties.len()
5191 ));
5192 }
5193
5194 let mut dropped_penaltyinfo = Vec::<DroppedPenaltyBlockInfo>::new();
5195 for info in local.penaltyinfo.iter().filter(|info| !info.active) {
5196 dropped_penaltyinfo.push(DroppedPenaltyBlockInfo {
5197 termname: Some(termspec.name.clone()),
5198 penalty: info.clone(),
5199 });
5200 }
5201 for info in &local.pre_dropped_penaltyinfo {
5202 dropped_penaltyinfo.push(DroppedPenaltyBlockInfo {
5203 termname: Some(termspec.name.clone()),
5204 penalty: info.clone(),
5205 });
5206 }
5207
5208 let applied_rotation: Option<gam_terms::basis::JointNullRotation> = match (
5212 local.joint_null_rotation.take(),
5213 lb_local.is_some(),
5214 local.linear_constraints.is_some(),
5215 ) {
5216 (Some(rot), false, false) => {
5217 let q = &rot.rotation;
5218 let dense = local
5219 .design
5220 .try_to_dense_by_chunks("joint-null absorption rotation (single realization)")
5221 .map_err(|e| {
5222 format!(
5223 "joint-null absorption rotation: dense conversion failed for term '{}': {}",
5224 termspec.name, e
5225 )
5226 })?;
5227 let rotated = gam_linalg::faer_ndarray::fast_ab(&dense, q);
5228 local.design = DesignMatrix::Dense(gam_linalg::matrix::DenseDesignMatrix::from(rotated));
5229 local.penalties = local
5230 .penalties
5231 .into_iter()
5232 .map(|s_local| {
5233 let qt_s = gam_linalg::faer_ndarray::fast_atb(q, &s_local);
5234 gam_linalg::faer_ndarray::fast_ab(&qt_s, q)
5235 })
5236 .collect();
5237 local.ops = vec![None; local.penalties.len()];
5238 local.kronecker_factored = None;
5239 Some(rot)
5240 }
5241 (Some(_), _, _) => None,
5242 (None, _, _) => None,
5243 };
5244
5245 let smooth_term = SmoothTerm {
5246 name: termspec.name.clone(),
5247 coeff_range: 0..p_local,
5248 shape: termspec.shape,
5249 penalties_local: local.penalties.clone(),
5250 nullspace_dims: local.nullspaces.clone(),
5251 penaltyinfo_local: local.penaltyinfo.clone(),
5252 metadata: local.metadata.clone(),
5253 lower_bounds_local: lb_local,
5254 linear_constraints_local: local.linear_constraints.clone(),
5255 kronecker_factored: local.kronecker_factored.take(),
5256 joint_null_rotation: applied_rotation,
5257 unabsorbed_global_orthogonality: None,
5260 };
5261
5262 Ok(SingleSmoothTermRealization {
5263 design_local: local.design,
5264 term: smooth_term,
5265 dropped_penaltyinfo,
5266 })
5267}
5268
5269fn freeze_geometry_from_metadata(
5280 termspec: &SmoothTermSpec,
5281 metadata: &BasisMetadata,
5282) -> Option<SmoothTermSpec> {
5283 let mut frozen = termspec.clone();
5284 match (&mut frozen.basis, metadata) {
5285 (
5286 SmoothBasisSpec::Matern {
5287 spec,
5288 input_scales: spec_scales,
5289 ..
5290 },
5291 BasisMetadata::Matern {
5292 centers,
5293 input_scales: meta_scales,
5294 identifiability_transform,
5295 nullspace_shrinkage_survived,
5296 ..
5297 },
5298 ) => {
5299 spec.center_strategy = CenterStrategy::UserProvided(centers.clone());
5300 if spec_scales.is_none()
5301 && let Some(s) = meta_scales.clone()
5302 {
5303 *spec_scales = Some(s);
5304 }
5305 if let Some(transform) = identifiability_transform.clone() {
5323 spec.identifiability = MaternIdentifiability::FrozenTransform {
5324 transform,
5325 nullspace_shrinkage_survived: Some(*nullspace_shrinkage_survived),
5326 };
5327 }
5328 Some(frozen)
5329 }
5330 (
5331 SmoothBasisSpec::Duchon {
5332 spec,
5333 input_scales: spec_scales,
5334 ..
5335 },
5336 BasisMetadata::Duchon {
5337 centers,
5338 input_scales: meta_scales,
5339 ..
5340 },
5341 ) => {
5342 spec.center_strategy = CenterStrategy::UserProvided(centers.clone());
5343 if spec_scales.is_none()
5344 && let Some(s) = meta_scales.clone()
5345 {
5346 *spec_scales = Some(s);
5347 }
5348 Some(frozen)
5349 }
5350 (
5351 SmoothBasisSpec::ThinPlate {
5352 spec,
5353 input_scales: spec_scales,
5354 ..
5355 },
5356 BasisMetadata::ThinPlate {
5357 centers,
5358 input_scales: meta_scales,
5359 ..
5360 },
5361 ) => {
5362 spec.center_strategy = CenterStrategy::UserProvided(centers.clone());
5363 if spec_scales.is_none()
5364 && let Some(s) = meta_scales.clone()
5365 {
5366 *spec_scales = Some(s);
5367 }
5368 Some(frozen)
5369 }
5370 _ => None,
5373 }
5374}
5375
5376fn rebuild_smooth_auxiliary_state(
5377 smooth: &mut SmoothDesign,
5378 dropped_penaltyinfo_by_term: &[Vec<DroppedPenaltyBlockInfo>],
5379) -> Result<(), String> {
5380 if dropped_penaltyinfo_by_term.len() != smooth.terms.len() {
5381 return Err(SmoothError::dimension_mismatch(format!(
5382 "smooth dropped-penalty cache mismatch: terms={}, dropped_sets={}",
5383 smooth.terms.len(),
5384 dropped_penaltyinfo_by_term.len()
5385 ))
5386 .into());
5387 }
5388
5389 let total_p = smooth.total_smooth_cols();
5390 let mut coefficient_lower_bounds = Array1::<f64>::from_elem(total_p, f64::NEG_INFINITY);
5391 let mut any_bounds = false;
5392 let mut linear_constraintrows: Vec<Array1<f64>> = Vec::new();
5393 let mut linear_constraint_b: Vec<f64> = Vec::new();
5394
5395 for term in &smooth.terms {
5396 let range = term.coeff_range.clone();
5397 if let Some(lb_local) = term.lower_bounds_local.as_ref() {
5398 if lb_local.len() != range.len() {
5399 return Err(SmoothError::dimension_mismatch(format!(
5400 "smooth lower-bound cache mismatch for term '{}': bounds={}, coeffs={}",
5401 term.name,
5402 lb_local.len(),
5403 range.len()
5404 ))
5405 .into());
5406 }
5407 coefficient_lower_bounds
5408 .slice_mut(s![range.clone()])
5409 .assign(lb_local);
5410 any_bounds = true;
5411 }
5412 if let Some(lin_local) = term.linear_constraints_local.as_ref() {
5413 if lin_local.a.ncols() != range.len() {
5414 return Err(SmoothError::dimension_mismatch(format!(
5415 "smooth linear-constraint cache mismatch for term '{}': cols={}, coeffs={}",
5416 term.name,
5417 lin_local.a.ncols(),
5418 range.len()
5419 ))
5420 .into());
5421 }
5422 for r in 0..lin_local.a.nrows() {
5423 let mut row = Array1::<f64>::zeros(total_p);
5424 row.slice_mut(s![range.clone()]).assign(&lin_local.a.row(r));
5425 linear_constraintrows.push(row);
5426 linear_constraint_b.push(lin_local.b[r]);
5427 }
5428 }
5429 }
5430
5431 smooth.coefficient_lower_bounds = if any_bounds {
5432 Some(coefficient_lower_bounds)
5433 } else {
5434 None
5435 };
5436 smooth.linear_constraints = if linear_constraintrows.is_empty() {
5437 None
5438 } else {
5439 let mut a = Array2::<f64>::zeros((linear_constraintrows.len(), total_p));
5440 for (i, row) in linear_constraintrows.iter().enumerate() {
5441 a.row_mut(i).assign(row);
5442 }
5443 Some(LinearInequalityConstraints {
5444 a,
5445 b: Array1::from_vec(linear_constraint_b),
5446 })
5447 };
5448 smooth.dropped_penaltyinfo = dropped_penaltyinfo_by_term
5449 .iter()
5450 .flat_map(|infos| infos.iter().cloned())
5451 .collect();
5452 Ok(())
5453}
5454
5455fn rebuild_term_collection_auxiliary_state(
5456 spec: &TermCollectionSpec,
5457 design: &mut TermCollectionDesign,
5458) -> Result<(), String> {
5459 if spec.linear_terms.len() != design.linear_ranges.len() {
5460 return Err(SmoothError::dimension_mismatch(format!(
5461 "term-collection linear bookkeeping mismatch: spec_terms={}, design_ranges={}",
5462 spec.linear_terms.len(),
5463 design.linear_ranges.len()
5464 ))
5465 .into());
5466 }
5467
5468 let p_total = design.design.ncols();
5469 let smooth_start = p_total.saturating_sub(design.smooth.total_smooth_cols());
5470 let mut coefficient_lower_bounds = Array1::<f64>::from_elem(p_total, f64::NEG_INFINITY);
5471 let mut any_bounds = false;
5472 let mut linear_constraintrows: Vec<Array1<f64>> = Vec::new();
5473 let mut linear_constraint_b: Vec<f64> = Vec::new();
5474
5475 for (linear, (_, range)) in spec.linear_terms.iter().zip(design.linear_ranges.iter()) {
5476 if range.len() != 1 {
5477 return Err(SmoothError::dimension_mismatch(format!(
5478 "linear term '{}' expected one coefficient column, found {}",
5479 linear.name,
5480 range.len()
5481 ))
5482 .into());
5483 }
5484 let col = range.start;
5485 if let Some(lb) = linear.coefficient_min {
5486 let mut row = Array1::<f64>::zeros(p_total);
5487 row[col] = 1.0;
5488 linear_constraintrows.push(row);
5489 linear_constraint_b.push(lb);
5490 }
5491 if let Some(ub) = linear.coefficient_max {
5492 let mut row = Array1::<f64>::zeros(p_total);
5493 row[col] = -1.0;
5494 linear_constraintrows.push(row);
5495 linear_constraint_b.push(-ub);
5496 }
5497 }
5498
5499 if let Some(lb_smooth) = design.smooth.coefficient_lower_bounds.as_ref() {
5500 if lb_smooth.len() != design.smooth.total_smooth_cols() {
5501 return Err(SmoothError::dimension_mismatch(format!(
5502 "smooth lower-bound width mismatch: bounds={}, smooth_cols={}",
5503 lb_smooth.len(),
5504 design.smooth.total_smooth_cols()
5505 ))
5506 .into());
5507 }
5508 coefficient_lower_bounds
5509 .slice_mut(s![
5510 smooth_start..(smooth_start + design.smooth.total_smooth_cols())
5511 ])
5512 .assign(lb_smooth);
5513 any_bounds = true;
5514 }
5515 if let Some(lin_smooth) = design.smooth.linear_constraints.as_ref() {
5516 if lin_smooth.a.ncols() != design.smooth.total_smooth_cols() {
5517 return Err(SmoothError::dimension_mismatch(format!(
5518 "smooth linear-constraint width mismatch: cols={}, smooth_cols={}",
5519 lin_smooth.a.ncols(),
5520 design.smooth.total_smooth_cols()
5521 ))
5522 .into());
5523 }
5524 let mut a_global = Array2::<f64>::zeros((lin_smooth.a.nrows(), p_total));
5525 a_global
5526 .slice_mut(s![
5527 ..,
5528 smooth_start..(smooth_start + design.smooth.total_smooth_cols())
5529 ])
5530 .assign(&lin_smooth.a);
5531 for r in 0..a_global.nrows() {
5532 linear_constraintrows.push(a_global.row(r).to_owned());
5533 linear_constraint_b.push(lin_smooth.b[r]);
5534 }
5535 }
5536
5537 let lower_bound_constraints = if any_bounds {
5538 linear_constraints_from_lower_bounds_global(&coefficient_lower_bounds)
5539 } else {
5540 None
5541 };
5542 let explicit_linear_constraints = if linear_constraintrows.is_empty() {
5543 None
5544 } else {
5545 let mut a = Array2::<f64>::zeros((linear_constraintrows.len(), p_total));
5546 for (i, row) in linear_constraintrows.iter().enumerate() {
5547 a.row_mut(i).assign(row);
5548 }
5549 Some(LinearInequalityConstraints {
5550 a,
5551 b: Array1::from_vec(linear_constraint_b),
5552 })
5553 };
5554
5555 design.coefficient_lower_bounds = if any_bounds {
5556 Some(coefficient_lower_bounds)
5557 } else {
5558 None
5559 };
5560 design.linear_constraints =
5561 merge_linear_constraints_global(explicit_linear_constraints, lower_bound_constraints);
5562 design.dropped_penaltyinfo = design.smooth.dropped_penaltyinfo.clone();
5563 Ok(())
5564}
5565
5566fn theta_values_match(left: &Array1<f64>, right: &Array1<f64>) -> bool {
5567 left.len() == right.len()
5568 && left
5569 .iter()
5570 .zip(right.iter())
5571 .all(|(&l, &r)| l.to_bits() == r.to_bits())
5572}
5573
5574fn latent_values_match(left: &Array1<f64>, right: &Array1<f64>) -> bool {
5575 theta_values_match(left, right)
5576}
5577
5578fn spatial_aniso_matches(left: Option<&[f64]>, right: Option<&[f64]>) -> bool {
5579 match (left, right) {
5580 (None, None) => true,
5581 (Some(a), Some(b)) => {
5582 a.len() == b.len()
5583 && a.iter()
5584 .zip(b.iter())
5585 .all(|(&x, &y)| x.to_bits() == y.to_bits())
5586 }
5587 _ => false,
5588 }
5589}
5590
5591fn spatial_length_scale_matches(left: Option<f64>, right: Option<f64>) -> bool {
5592 match (left, right) {
5593 (None, None) => true,
5594 (Some(a), Some(b)) => a.to_bits() == b.to_bits(),
5595 _ => false,
5596 }
5597}
5598
5599struct FrozenTermCollectionIncrementalRealizer<'d> {
5600 data: ArrayView2<'d, f64>,
5601 spec: TermCollectionSpec,
5602 design: TermCollectionDesign,
5603 fixed_blocks: Vec<DesignBlock>,
5604 dropped_penaltyinfo_by_term: Vec<Vec<DroppedPenaltyBlockInfo>>,
5605 smooth_penalty_ranges: Vec<Range<usize>>,
5606 full_penalty_ranges: Vec<Range<usize>>,
5607 basisworkspace: gam_terms::basis::BasisWorkspace,
5611 spatial_realization_geometry: Vec<Option<SmoothTermSpec>>,
5624 design_revision: u64,
5630}
5631
5632impl<'d> std::fmt::Debug for FrozenTermCollectionIncrementalRealizer<'d> {
5633 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
5634 f.debug_struct("FrozenTermCollectionIncrementalRealizer")
5635 .field("data_shape", &(self.data.nrows(), self.data.ncols()))
5636 .field("fixed_blocks", &self.fixed_blocks.len())
5637 .finish_non_exhaustive()
5638 }
5639}
5640
5641impl<'d> FrozenTermCollectionIncrementalRealizer<'d> {
5642 fn new(
5643 data: ArrayView2<'d, f64>,
5644 spec: TermCollectionSpec,
5645 design: TermCollectionDesign,
5646 ) -> Result<Self, String> {
5647 if spec.smooth_terms.len() != design.smooth.terms.len() {
5648 return Err(SmoothError::dimension_mismatch(format!(
5649 "incremental realizer smooth term mismatch: spec_terms={}, design_terms={}",
5650 spec.smooth_terms.len(),
5651 design.smooth.terms.len()
5652 ))
5653 .into());
5654 }
5655
5656 let mut smooth_cursor = 0usize;
5657 let mut smooth_penalty_ranges = Vec::with_capacity(design.smooth.terms.len());
5658 for term in &design.smooth.terms {
5659 let next = smooth_cursor + term.penalties_local.len();
5660 smooth_penalty_ranges.push(smooth_cursor..next);
5661 smooth_cursor = next;
5662 }
5663 if smooth_cursor != design.smooth.penalties.len() {
5664 return Err(SmoothError::dimension_mismatch(format!(
5665 "incremental realizer smooth penalty mismatch: ranged={}, actual={}",
5666 smooth_cursor,
5667 design.smooth.penalties.len()
5668 ))
5669 .into());
5670 }
5671
5672 let fixed_penalty_offset = design
5673 .penalties
5674 .len()
5675 .checked_sub(design.smooth.penalties.len())
5676 .ok_or_else(|| {
5677 "incremental realizer encountered invalid penalty bookkeeping".to_string()
5678 })?;
5679 let full_penalty_ranges = smooth_penalty_ranges
5680 .iter()
5681 .map(|range| (fixed_penalty_offset + range.start)..(fixed_penalty_offset + range.end))
5682 .collect::<Vec<_>>();
5683 let fixed_blocks = build_term_collection_fixed_blocks(data, &spec)
5684 .map_err(|e| format!("failed to cache fixed term-collection blocks: {e}"))?;
5685
5686 let mut dropped_penaltyinfo_by_term = Vec::with_capacity(spec.smooth_terms.len());
5687 for (term_idx, termspec) in spec.smooth_terms.iter().enumerate() {
5688 let realization =
5689 build_single_smooth_term_realization(data, termspec).map_err(|e| {
5690 format!(
5691 "failed to build cached realization for smooth term '{}' (index {}): {e}",
5692 termspec.name, term_idx
5693 )
5694 })?;
5695 let expected_cols = design.smooth.terms[term_idx].coeff_range.len();
5696 if realization.design_local.ncols() != expected_cols {
5697 return Err(SmoothError::dimension_mismatch(format!(
5698 "cached realization width mismatch for term '{}': cached_cols={}, design_cols={}",
5699 termspec.name,
5700 realization.design_local.ncols(),
5701 expected_cols
5702 ))
5703 .into());
5704 }
5705 if realization.active_penaltyinfo().len()
5706 != design.smooth.terms[term_idx].penalties_local.len()
5707 {
5708 return Err(SmoothError::dimension_mismatch(format!(
5709 "cached realization penalty mismatch for term '{}': cached_penalties={}, design_penalties={}",
5710 termspec.name,
5711 realization.active_penaltyinfo().len(),
5712 design.smooth.terms[term_idx].penalties_local.len()
5713 ))
5714 .into());
5715 }
5716 dropped_penaltyinfo_by_term.push(realization.dropped_penaltyinfo);
5717 }
5718
5719 let geometry_slots = spec.smooth_terms.len();
5720 Ok(Self {
5721 data,
5722 spec,
5723 design,
5724 fixed_blocks,
5725 dropped_penaltyinfo_by_term,
5726 smooth_penalty_ranges,
5727 full_penalty_ranges,
5728 basisworkspace: gam_terms::basis::BasisWorkspace::new(),
5729 spatial_realization_geometry: vec![None; geometry_slots],
5730 design_revision: 0,
5731 })
5732 }
5733
5734 fn design_revision(&self) -> u64 {
5735 self.design_revision
5736 }
5737
5738 fn spec(&self) -> &TermCollectionSpec {
5739 &self.spec
5740 }
5741
5742 fn design(&self) -> &TermCollectionDesign {
5743 &self.design
5744 }
5745
5746 fn supports_nfree_penalty_rekey(&self, spatial_terms: &[usize]) -> bool {
5787 if spatial_terms.len() != 1 {
5788 return false;
5789 }
5790 let term_idx = spatial_terms[0];
5791 matches!(
5792 self.design.smooth.terms.get(term_idx).map(|t| &t.metadata),
5793 Some(BasisMetadata::Duchon { .. } | BasisMetadata::ThinPlate { .. })
5794 )
5795 }
5796
5797 fn supports_nfree_gradient_only_routing(&self, spatial_terms: &[usize]) -> bool {
5806 if spatial_terms.len() != 1 {
5807 return false;
5808 }
5809 let term_idx = spatial_terms[0];
5810 matches!(
5811 self.design.smooth.terms.get(term_idx).map(|t| &t.metadata),
5812 Some(BasisMetadata::Duchon { .. } | BasisMetadata::ThinPlate { .. })
5813 )
5814 }
5815
5816 fn canonical_penalties_at_psi(
5829 &mut self,
5830 spatial_terms: &[usize],
5831 psi: &[f64],
5832 ) -> Result<(Vec<gam_terms::construction::CanonicalPenalty>, Vec<usize>), String> {
5833 if spatial_terms.len() != 1 {
5834 return Err(format!(
5835 "n-free penalty re-key requires exactly one spatial term, found {}",
5836 spatial_terms.len()
5837 ));
5838 }
5839 let term_idx = spatial_terms[0];
5840 let (ls_opt, aniso_from_psi) = spatial_term_psi_to_length_scale_and_aniso(psi);
5846 let termspec =
5849 self.spec.smooth_terms.get(term_idx).ok_or_else(|| {
5850 format!("spatial term {term_idx} out of range for n-free penalty")
5851 })?;
5852 let term = self
5853 .design
5854 .smooth
5855 .terms
5856 .get(term_idx)
5857 .ok_or_else(|| format!("realized smooth term {term_idx} out of range"))?;
5858 let p_total = self.design.design.ncols();
5861 let (locals, nullspace_dims): (Vec<Array2<f64>>, Vec<usize>) = match &term.metadata {
5862 BasisMetadata::Duchon {
5863 centers,
5864 identifiability_transform,
5865 operator_collocation_points,
5866 power,
5867 nullspace_order,
5868 aniso_log_scales,
5869 input_scales,
5870 radial_reparam,
5871 ..
5872 } => {
5873 let operator_penalties = match &termspec.basis {
5874 SmoothBasisSpec::Duchon { spec, .. } => spec.operator_penalties.clone(),
5875 _ => gam_terms::basis::DuchonOperatorPenaltySpec::default(),
5876 };
5877 let effective_ls = match input_scales.as_deref() {
5884 Some(scales) => {
5885 compensate_optional_length_scale_for_standardization(ls_opt, scales)
5886 }
5887 None => ls_opt,
5888 };
5889 gam_terms::basis::duchon_penalties_at_length_scale(
5890 centers.view(),
5891 identifiability_transform.as_ref(),
5892 operator_collocation_points.as_ref().map(|p| p.view()),
5893 &operator_penalties,
5894 *power,
5895 *nullspace_order,
5896 aniso_log_scales.as_deref(),
5897 radial_reparam.as_ref(),
5898 effective_ls,
5899 &mut self.basisworkspace,
5900 )
5901 .map_err(|e| e.to_string())?
5902 }
5903 BasisMetadata::Matern {
5904 centers,
5905 periodic,
5906 nu,
5907 include_intercept,
5908 identifiability_transform,
5909 aniso_log_scales,
5910 input_scales,
5911 ..
5912 } => {
5913 let ls = ls_opt.ok_or_else(|| {
5920 "Matérn n-free penalty re-key requires a finite length-scale".to_string()
5921 })?;
5922 let effective_ls = match input_scales.as_deref() {
5923 Some(scales) => compensate_length_scale_for_standardization(ls, scales),
5924 None => ls,
5925 };
5926 let aniso_for_penalty = aniso_from_psi.as_deref().or(aniso_log_scales.as_deref());
5927 let (penalties, nullspace_dims, _info) =
5938 matern_operator_penalty_triplet_at_length_scale(
5939 centers.view(),
5940 periodic.as_deref(),
5941 identifiability_transform.as_ref(),
5942 *nu,
5943 *include_intercept,
5944 aniso_for_penalty,
5945 effective_ls,
5946 )
5947 .map_err(|e| e.to_string())?;
5948 (penalties, nullspace_dims)
5949 }
5950 BasisMetadata::ThinPlate {
5951 centers,
5952 identifiability_transform,
5953 radial_reparam,
5954 ..
5955 } => {
5956 let ls = ls_opt.ok_or_else(|| {
5957 "thin-plate n-free penalty re-key requires a finite length-scale".to_string()
5958 })?;
5959 let double_penalty = match &termspec.basis {
5960 SmoothBasisSpec::ThinPlate { spec, .. } => spec.double_penalty,
5961 _ => false,
5962 };
5963 gam_terms::basis::thin_plate_penalties_at_length_scale(
5964 centers.view(),
5965 identifiability_transform.as_ref(),
5966 radial_reparam.as_ref(),
5967 ls,
5968 double_penalty,
5969 &mut self.basisworkspace,
5970 )
5971 .map_err(|e| e.to_string())?
5972 }
5973 other => {
5974 return Err(format!(
5975 "n-free penalty re-key unsupported for basis metadata {:?}",
5976 std::mem::discriminant(other)
5977 ));
5978 }
5979 };
5980 let templates = &self.design.penalties;
5985 if templates.len() != locals.len() {
5986 return Err(format!(
5987 "n-free penalty re-key produced {} blocks but the frozen design carries {} \
5988 — penalty topology is not ψ-stable",
5989 locals.len(),
5990 templates.len()
5991 ));
5992 }
5993 let specs: Vec<gam_solve::estimate::PenaltySpec> = templates
5994 .iter()
5995 .zip(locals.into_iter())
5996 .map(|(tmpl, local)| gam_solve::estimate::PenaltySpec::Block {
5997 local,
5998 col_range: tmpl.col_range.clone(),
5999 prior_mean: tmpl.prior_mean.clone(),
6000 structure_hint: tmpl.structure_hint.clone(),
6001 op: tmpl.op.clone(),
6002 })
6003 .collect();
6004 gam_terms::construction::canonicalize_penalty_specs(
6005 &specs,
6006 &nullspace_dims,
6007 p_total,
6008 "nfree-psi-penalty",
6009 )
6010 .map_err(|e| e.to_string())
6011 }
6012
6013 fn canonical_penalty_derivatives_at_psi(
6014 &mut self,
6015 spatial_terms: &[usize],
6016 psi: &[f64],
6017 ) -> Result<(Range<usize>, usize, Vec<Array2<f64>>), String> {
6018 if spatial_terms.len() != 1 {
6019 return Err(format!(
6020 "n-free penalty derivative re-key requires exactly one spatial term, found {}",
6021 spatial_terms.len()
6022 ));
6023 }
6024 let term_idx = spatial_terms[0];
6025 let (ls_opt, aniso_from_psi) = spatial_term_psi_to_length_scale_and_aniso(psi);
6026 let termspec = self.spec.smooth_terms.get(term_idx).ok_or_else(|| {
6027 format!("spatial term {term_idx} out of range for n-free penalty derivative")
6028 })?;
6029 let term = self
6030 .design
6031 .smooth
6032 .terms
6033 .get(term_idx)
6034 .ok_or_else(|| format!("realized smooth term {term_idx} out of range"))?;
6035 let p_total = self.design.design.ncols();
6036 let smooth_start = p_total.saturating_sub(self.design.smooth.total_smooth_cols());
6037 let global_range =
6038 (smooth_start + term.coeff_range.start)..(smooth_start + term.coeff_range.end);
6039
6040 let locals = match &term.metadata {
6041 BasisMetadata::Duchon {
6042 centers,
6043 identifiability_transform,
6044 operator_collocation_points,
6045 power,
6046 nullspace_order,
6047 aniso_log_scales,
6048 input_scales,
6049 radial_reparam,
6050 ..
6051 } => {
6052 let mut spec = match &termspec.basis {
6053 SmoothBasisSpec::Duchon { spec, .. } => spec.clone(),
6054 _ => {
6055 return Err(
6056 "Duchon n-free penalty derivative requires a Duchon term spec"
6057 .to_string(),
6058 );
6059 }
6060 };
6061 let effective_ls = match input_scales.as_deref() {
6062 Some(scales) => {
6063 compensate_optional_length_scale_for_standardization(ls_opt, scales)
6064 }
6065 None => ls_opt,
6066 };
6067 spec.length_scale = effective_ls;
6068 spec.power = *power;
6069 spec.nullspace_order = *nullspace_order;
6070 spec.aniso_log_scales = aniso_log_scales.clone();
6071 spec.radial_reparam = radial_reparam.clone();
6074 if spec.length_scale.is_none() {
6075 return Err(
6076 "Duchon n-free penalty derivative requires a hybrid length-scale"
6077 .to_string(),
6078 );
6079 }
6080 let collocation = operator_collocation_points
6081 .as_ref()
6082 .map(|points| points.view())
6083 .unwrap_or_else(|| centers.view());
6084 let (_native_sources, mut first, _native_second) =
6085 gam_terms::basis::build_duchon_native_penalty_psi_derivatives(
6086 centers.view(),
6087 &spec,
6088 identifiability_transform.as_ref(),
6089 &mut self.basisworkspace,
6090 )
6091 .map_err(|e| e.to_string())?;
6092 let (_operator_sources, operator_first, _operator_second) =
6093 gam_terms::basis::build_duchon_operator_penalty_psi_derivatives(
6094 collocation,
6095 centers.view(),
6096 &spec,
6097 identifiability_transform.as_ref(),
6098 &mut self.basisworkspace,
6099 )
6100 .map_err(|e| e.to_string())?;
6101 first.extend(operator_first);
6102 first
6103 }
6104 BasisMetadata::Matern {
6105 centers,
6106 periodic,
6107 nu,
6108 include_intercept,
6109 identifiability_transform,
6110 aniso_log_scales,
6111 input_scales,
6112 ..
6113 } => {
6114 let ls = ls_opt.ok_or_else(|| {
6115 "Matérn n-free penalty derivative requires a finite length-scale".to_string()
6116 })?;
6117 let effective_ls = match input_scales.as_deref() {
6118 Some(scales) => compensate_length_scale_for_standardization(ls, scales),
6119 None => ls,
6120 };
6121 let penalty_centers =
6122 gam_terms::basis::expand_periodic_centers(¢ers.to_owned(), periodic.as_deref())
6123 .map_err(|e| e.to_string())?;
6124 let aniso_for_penalty = aniso_from_psi.as_deref().or(aniso_log_scales.as_deref());
6125 let (first, _second) = gam_terms::basis::build_matern_operator_penalty_psi_derivatives(
6126 penalty_centers.view(),
6127 effective_ls,
6128 *nu,
6129 *include_intercept,
6130 identifiability_transform.as_ref(),
6131 aniso_for_penalty,
6132 )
6133 .map_err(|e| e.to_string())?;
6134 first
6135 }
6136 BasisMetadata::ThinPlate {
6137 centers,
6138 identifiability_transform,
6139 radial_reparam,
6140 ..
6141 } => {
6142 let ls = ls_opt.ok_or_else(|| {
6143 "thin-plate n-free penalty derivative requires a finite length-scale"
6144 .to_string()
6145 })?;
6146 let mut spec = match &termspec.basis {
6147 SmoothBasisSpec::ThinPlate { spec, .. } => spec.clone(),
6148 _ => {
6149 return Err(
6150 "thin-plate n-free penalty derivative requires a ThinPlate term spec"
6151 .to_string(),
6152 );
6153 }
6154 };
6155 spec.length_scale = ls;
6156 if spec.radial_reparam.is_none() {
6157 spec.radial_reparam = radial_reparam.clone();
6158 }
6159 let (primary, _primary_second) =
6160 gam_terms::basis::build_thin_plate_penalty_psi_derivativeswithworkspace(
6161 centers.view(),
6162 &spec,
6163 identifiability_transform.as_ref(),
6164 &mut self.basisworkspace,
6165 )
6166 .map_err(|e| e.to_string())?;
6167 if self.design.penalties.len() > 1 {
6168 vec![primary.clone(), Array2::<f64>::zeros(primary.raw_dim())]
6169 } else {
6170 vec![primary]
6171 }
6172 }
6173 other => {
6174 return Err(format!(
6175 "n-free penalty derivative re-key unsupported for basis metadata {:?}",
6176 std::mem::discriminant(other)
6177 ));
6178 }
6179 };
6180 if locals.len() != self.design.penalties.len() {
6181 return Err(format!(
6182 "n-free penalty derivative re-key produced {} blocks but the frozen design carries {} \
6183 — penalty topology is not ψ-stable",
6184 locals.len(),
6185 self.design.penalties.len()
6186 ));
6187 }
6188 Ok((global_range, p_total, locals))
6189 }
6190
6191 fn apply_log_kappa(
6192 &mut self,
6193 log_kappa: &SpatialLogKappaCoords,
6194 term_indices: &[usize],
6195 ) -> Result<(), String> {
6196 if term_indices.len() != log_kappa.dims_per_term().len() {
6197 return Err(SmoothError::dimension_mismatch(format!(
6198 "incremental realizer log-kappa term mismatch: term_indices={}, dims_per_term={}",
6199 term_indices.len(),
6200 log_kappa.dims_per_term().len()
6201 ))
6202 .into());
6203 }
6204
6205 let mut any_changed = false;
6206 for (slot, &term_idx) in term_indices.iter().enumerate() {
6207 any_changed |= self.apply_log_kappa_to_term(term_idx, log_kappa.term_slice(slot))?;
6208 }
6209
6210 if any_changed {
6211 self.refresh_full_design_operator()?;
6212 rebuild_smooth_auxiliary_state(
6213 &mut self.design.smooth,
6214 &self.dropped_penaltyinfo_by_term,
6215 )?;
6216 rebuild_term_collection_auxiliary_state(&self.spec, &mut self.design)?;
6217 self.design_revision = self.design_revision.wrapping_add(1);
6218 }
6219 Ok(())
6220 }
6221
6222 fn apply_log_kappa_to_term(&mut self, term_idx: usize, psi: &[f64]) -> Result<bool, String> {
6223 if !spatial_term_supports_hyper_optimization(&self.spec, term_idx) {
6224 return Err(SmoothError::invalid_config(format!(
6225 "incremental realizer term {term_idx} does not expose spatial hyperparameters"
6226 ))
6227 .into());
6228 }
6229 let measure_jet_term = measure_jet_term_spec(&self.spec, term_idx).is_some();
6233 let constant_curvature_term = constant_curvature_term_spec(&self.spec, term_idx).is_some();
6237 let mut next_length_scale = None;
6238 let mut next_aniso: Option<Vec<f64>> = None;
6239 if measure_jet_term {
6240 if !set_measure_jet_psi_dials(&mut self.spec, term_idx, psi)
6241 .map_err(|e| e.to_string())?
6242 {
6243 return Ok(false);
6244 }
6245 } else if constant_curvature_term {
6246 if !set_constant_curvature_kappa(&mut self.spec, term_idx, psi)
6247 .map_err(|e| e.to_string())?
6248 {
6249 return Ok(false);
6250 }
6251 } else {
6252 let current_length_scale = get_spatial_length_scale(&self.spec, term_idx);
6253 let current_aniso = get_spatial_aniso_log_scales(&self.spec, term_idx);
6254 let (ls, eta) = spatial_term_psi_to_length_scale_and_aniso(psi);
6255 next_length_scale = ls;
6256 next_aniso = eta;
6257 let same_length = spatial_length_scale_matches(current_length_scale, next_length_scale);
6258 let same_aniso = spatial_aniso_matches(current_aniso.as_deref(), next_aniso.as_deref());
6259 if same_length && same_aniso {
6260 return Ok(false);
6261 }
6262 if let Some(length_scale) = next_length_scale {
6263 set_spatial_length_scale(&mut self.spec, term_idx, length_scale)
6264 .map_err(|e| e.to_string())?;
6265 }
6266 if let Some(eta) = next_aniso.clone() {
6267 set_spatial_aniso_log_scales(&mut self.spec, term_idx, eta)
6268 .map_err(|e| e.to_string())?;
6269 }
6270 }
6271
6272 let geometry_slot = self
6283 .spatial_realization_geometry
6284 .get(term_idx)
6285 .ok_or_else(|| format!("incremental realizer geometry slot {term_idx} out of range"))?;
6286 let mut build_spec = match geometry_slot {
6287 Some(cached) => cached.clone(),
6288 None => self
6289 .spec
6290 .smooth_terms
6291 .get(term_idx)
6292 .ok_or_else(|| format!("incremental realizer smooth term {term_idx} out of range"))?
6293 .clone(),
6294 };
6295 if measure_jet_term {
6296 set_single_term_measure_jet_psi_dials(&mut build_spec, psi)
6300 .map_err(|e| e.to_string())?;
6301 } else if constant_curvature_term {
6302 set_single_term_constant_curvature_kappa(&mut build_spec, psi)
6307 .map_err(|e| e.to_string())?;
6308 } else {
6309 if let Some(length_scale) = next_length_scale {
6310 set_single_term_spatial_length_scale(&mut build_spec, length_scale)
6311 .map_err(|e| e.to_string())?;
6312 }
6313 if let Some(eta) = next_aniso {
6314 set_single_term_spatial_aniso_log_scales(&mut build_spec, eta)
6315 .map_err(|e| e.to_string())?;
6316 }
6317 }
6318
6319 let termname = build_spec.name.clone();
6320 let local = build_single_local_smooth_term(
6321 self.data,
6322 &build_spec,
6323 &mut self.basisworkspace,
6324 )
6325 .map_err(|e| {
6326 format!(
6327 "failed to rebuild smooth term '{termname}' during incremental κ realization: {e}"
6328 )
6329 })?;
6330
6331 if self.spatial_realization_geometry[term_idx].is_none()
6336 && let Some(frozen) = freeze_geometry_from_metadata(&build_spec, &local.metadata)
6337 {
6338 if let (
6350 SmoothBasisSpec::Matern {
6351 spec: frozen_spec, ..
6352 },
6353 Some(SmoothBasisSpec::Matern {
6354 spec: live_spec, ..
6355 }),
6356 ) = (
6357 &frozen.basis,
6358 self.spec
6359 .smooth_terms
6360 .get_mut(term_idx)
6361 .map(|t| &mut t.basis),
6362 ) {
6363 live_spec.identifiability = frozen_spec.identifiability.clone();
6364 live_spec.center_strategy = frozen_spec.center_strategy.clone();
6365 }
6366 self.spatial_realization_geometry[term_idx] = Some(frozen);
6367 }
6368
6369 let realization = wrap_local_build_as_realization(local, &build_spec)?;
6370 self.replace_term_realization(term_idx, realization)?;
6371 Ok(true)
6372 }
6373
6374 fn replace_term_realization(
6375 &mut self,
6376 term_idx: usize,
6377 realization: SingleSmoothTermRealization,
6378 ) -> Result<(), String> {
6379 let t_replace = std::time::Instant::now();
6380 let SingleSmoothTermRealization {
6381 design_local,
6382 term,
6383 dropped_penaltyinfo,
6384 } = realization;
6385 let SmoothTerm {
6386 name,
6387 penalties_local,
6388 nullspace_dims,
6389 penaltyinfo_local,
6390 metadata,
6391 lower_bounds_local,
6392 linear_constraints_local,
6393 joint_null_rotation,
6394 ..
6395 } = term;
6396 let coeff_range = self
6397 .design
6398 .smooth
6399 .terms
6400 .get(term_idx)
6401 .ok_or_else(|| format!("incremental realizer smooth term {term_idx} out of range"))?
6402 .coeff_range
6403 .clone();
6404 if design_local.ncols() != coeff_range.len() {
6405 return Err(SmoothError::dimension_mismatch(format!(
6406 "incremental realizer width mismatch for term {}: rebuilt_cols={}, cached_cols={}",
6407 term_idx,
6408 design_local.ncols(),
6409 coeff_range.len()
6410 ))
6411 .into());
6412 }
6413 if design_local.nrows() != self.design.design.nrows() {
6414 return Err(SmoothError::dimension_mismatch(format!(
6415 "incremental realizer row mismatch for term {}: rebuilt_rows={}, design_rows={}",
6416 term_idx,
6417 design_local.nrows(),
6418 self.design.design.nrows()
6419 ))
6420 .into());
6421 }
6422
6423 let active_penaltyinfo = penaltyinfo_local
6424 .iter()
6425 .filter(|info| info.active)
6426 .cloned()
6427 .collect::<Vec<_>>();
6428 let smooth_penalty_range = self
6429 .smooth_penalty_ranges
6430 .get(term_idx)
6431 .ok_or_else(|| {
6432 format!("incremental realizer missing smooth penalty range for term {term_idx}")
6433 })?
6434 .clone();
6435 let full_penalty_range = self
6436 .full_penalty_ranges
6437 .get(term_idx)
6438 .ok_or_else(|| {
6439 format!("incremental realizer missing full penalty range for term {term_idx}")
6440 })?
6441 .clone();
6442 if active_penaltyinfo.len() != smooth_penalty_range.len()
6443 || penalties_local.len() != smooth_penalty_range.len()
6444 || nullspace_dims.len() != smooth_penalty_range.len()
6445 {
6446 return Err(SmoothError::dimension_mismatch(format!(
6447 "incremental realizer topology changed for term '{}': penalties={}, infos={}, nullspaces={}, cached_penalties={}",
6448 name,
6449 penalties_local.len(),
6450 active_penaltyinfo.len(),
6451 nullspace_dims.len(),
6452 smooth_penalty_range.len()
6453 ))
6454 .into());
6455 }
6456
6457 self.design.smooth.term_designs[term_idx] = design_local;
6458
6459 for (offset, penalty_local) in penalties_local.iter().enumerate() {
6460 let smooth_penalty_idx = smooth_penalty_range.start + offset;
6461 let full_penalty_idx = full_penalty_range.start + offset;
6462 let nullspace_dim = nullspace_dims[offset];
6463 let penalty_info = active_penaltyinfo[offset].clone();
6464
6465 if penalty_local.nrows() != coeff_range.len()
6466 || penalty_local.ncols() != coeff_range.len()
6467 {
6468 return Err(SmoothError::dimension_mismatch(format!(
6469 "incremental realizer penalty shape mismatch for term '{}' penalty {}: \
6470 penalty is {}x{} but coeff_range has {} columns",
6471 name,
6472 offset,
6473 penalty_local.nrows(),
6474 penalty_local.ncols(),
6475 coeff_range.len()
6476 ))
6477 .into());
6478 }
6479
6480 let smooth_penalty = self
6481 .design
6482 .smooth
6483 .penalties
6484 .get_mut(smooth_penalty_idx)
6485 .ok_or_else(|| {
6486 format!(
6487 "incremental realizer smooth penalty {} out of range for term {}",
6488 smooth_penalty_idx, term_idx
6489 )
6490 })?;
6491 smooth_penalty.local.assign(penalty_local);
6494
6495 let full_bp = self
6496 .design
6497 .penalties
6498 .get_mut(full_penalty_idx)
6499 .ok_or_else(|| {
6500 format!(
6501 "incremental realizer full penalty {} out of range for term {}",
6502 full_penalty_idx, term_idx
6503 )
6504 })?;
6505 full_bp.local.assign(penalty_local);
6508
6509 self.design.smooth.nullspace_dims[smooth_penalty_idx] = nullspace_dim;
6510 self.design.nullspace_dims[full_penalty_idx] = nullspace_dim;
6511
6512 self.design.smooth.penaltyinfo[smooth_penalty_idx].global_index = smooth_penalty_idx;
6513 self.design.smooth.penaltyinfo[smooth_penalty_idx].termname = Some(name.clone());
6514 self.design.smooth.penaltyinfo[smooth_penalty_idx].penalty = penalty_info.clone();
6515
6516 self.design.penaltyinfo[full_penalty_idx].global_index = full_penalty_idx;
6517 self.design.penaltyinfo[full_penalty_idx].termname = Some(name.clone());
6518 self.design.penaltyinfo[full_penalty_idx].penalty = penalty_info;
6519 }
6520
6521 let target_term = self.design.smooth.terms.get_mut(term_idx).ok_or_else(|| {
6522 format!("incremental realizer smooth term {term_idx} disappeared during replacement")
6523 })?;
6524 target_term.penalties_local = penalties_local;
6525 target_term.nullspace_dims = nullspace_dims;
6526 target_term.penaltyinfo_local = penaltyinfo_local;
6527 target_term.metadata = metadata;
6528 target_term.lower_bounds_local = lower_bounds_local;
6529 target_term.linear_constraints_local = linear_constraints_local;
6530 target_term.joint_null_rotation = joint_null_rotation;
6531 self.dropped_penaltyinfo_by_term[term_idx] = dropped_penaltyinfo;
6532 log::info!(
6533 "[STAGE] smooth basis rebuild (term {}, '{}', cols={}): {:.3}s",
6534 term_idx,
6535 target_term.name,
6536 coeff_range.len(),
6537 t_replace.elapsed().as_secs_f64(),
6538 );
6539 Ok(())
6540 }
6541
6542 fn refresh_full_design_operator(&mut self) -> Result<(), String> {
6543 let mut blocks = Vec::<DesignBlock>::with_capacity(
6544 self.fixed_blocks.len() + self.design.smooth.term_designs.len(),
6545 );
6546 blocks.extend(self.fixed_blocks.iter().cloned());
6547 for term_design in &self.design.smooth.term_designs {
6548 blocks.push(DesignBlock::from(term_design));
6549 }
6550 self.design.design = assemble_term_collection_design_matrix(blocks)
6551 .map_err(|e| format!("failed to refresh term-collection design: {e}"))?;
6552 Ok(())
6553 }
6554}
6555
6556fn build_term_collection_fixed_blocks(
6557 data: ArrayView2<'_, f64>,
6558 spec: &TermCollectionSpec,
6559) -> Result<Vec<DesignBlock>, BasisError> {
6560 let mut blocks = Vec::<DesignBlock>::new();
6561 if !term_collection_has_one_sided_anchored_bspline(spec) {
6562 blocks.push(DesignBlock::Intercept(data.nrows()));
6563 }
6564
6565 if !spec.linear_terms.is_empty() {
6566 let mut linear_block = Array2::<f64>::zeros((data.nrows(), spec.linear_terms.len()));
6567 for (j, linear) in spec.linear_terms.iter().enumerate() {
6568 let column = linear
6572 .realized_design_column(data)
6573 .map_err(BasisError::InvalidInput)?;
6574 linear_block.column_mut(j).assign(&column);
6575 }
6576 blocks.push(DesignBlock::Dense(gam_linalg::matrix::DenseDesignMatrix::from(
6577 linear_block,
6578 )));
6579 }
6580
6581 for term in &spec.random_effect_terms {
6582 let block = build_random_effect_block(data, term)?;
6583 let re_op = RandomEffectOperator::new(block.group_ids, block.num_groups);
6584 blocks.push(DesignBlock::RandomEffect(Arc::new(re_op)));
6585 }
6586
6587 Ok(blocks)
6588}
6589
6590pub struct SpatialLengthScaleOptimizationResult<FitOut> {
6595 pub resolved_specs: Vec<TermCollectionSpec>,
6596 pub designs: Vec<TermCollectionDesign>,
6597 pub fit: FitOut,
6598 pub timing: Option<SpatialLengthScaleOptimizationTiming>,
6599}
6600
6601#[derive(Debug, Clone)]
6603pub struct ExactJointHyperSetup {
6604 rho0: Array1<f64>,
6605 rho_lower: Array1<f64>,
6606 rho_upper: Array1<f64>,
6607 log_kappa0: SpatialLogKappaCoords,
6608 log_kappa_lower: SpatialLogKappaCoords,
6609 log_kappa_upper: SpatialLogKappaCoords,
6610 auxiliary0: Array1<f64>,
6611 auxiliary_lower: Array1<f64>,
6612 auxiliary_upper: Array1<f64>,
6613}
6614
6615impl ExactJointHyperSetup {
6616 fn sanitize_rho_seed(
6617 rho0: Array1<f64>,
6618 rho_lower: &Array1<f64>,
6619 rho_upper: &Array1<f64>,
6620 ) -> Array1<f64> {
6621 Array1::from_iter(rho0.iter().enumerate().map(|(idx, &value)| {
6622 let lo = rho_lower[idx];
6623 let hi = rho_upper[idx];
6624 let fallback = 0.0_f64.clamp(lo, hi);
6625 if value.is_finite() {
6626 value.clamp(lo, hi)
6627 } else {
6628 fallback
6629 }
6630 }))
6631 }
6632
6633 pub(crate) fn new(
6634 rho0: Array1<f64>,
6635 rho_lower: Array1<f64>,
6636 rho_upper: Array1<f64>,
6637 log_kappa0: SpatialLogKappaCoords,
6638 log_kappa_lower: SpatialLogKappaCoords,
6639 log_kappa_upper: SpatialLogKappaCoords,
6640 ) -> Self {
6641 let rho0 = Self::sanitize_rho_seed(rho0, &rho_lower, &rho_upper);
6642 Self {
6643 rho0,
6644 rho_lower,
6645 rho_upper,
6646 log_kappa0,
6647 log_kappa_lower,
6648 log_kappa_upper,
6649 auxiliary0: Array1::zeros(0),
6650 auxiliary_lower: Array1::zeros(0),
6651 auxiliary_upper: Array1::zeros(0),
6652 }
6653 }
6654
6655 pub(crate) fn with_auxiliary(
6656 mut self,
6657 auxiliary0: Array1<f64>,
6658 auxiliary_lower: Array1<f64>,
6659 auxiliary_upper: Array1<f64>,
6660 ) -> Self {
6661 assert_eq!(
6662 auxiliary0.len(),
6663 auxiliary_lower.len(),
6664 "auxiliary lower bound length mismatch"
6665 );
6666 assert_eq!(
6667 auxiliary0.len(),
6668 auxiliary_upper.len(),
6669 "auxiliary upper bound length mismatch"
6670 );
6671 self.auxiliary0 = Self::sanitize_rho_seed(auxiliary0, &auxiliary_lower, &auxiliary_upper);
6672 self.auxiliary_lower = auxiliary_lower;
6673 self.auxiliary_upper = auxiliary_upper;
6674 self
6675 }
6676
6677 pub(crate) fn rho_dim(&self) -> usize {
6678 self.rho0.len()
6679 }
6680
6681 pub(crate) fn log_kappa_dim(&self) -> usize {
6682 self.log_kappa0.len()
6683 }
6684
6685 pub(crate) fn auxiliary_dim(&self) -> usize {
6686 self.auxiliary0.len()
6687 }
6688
6689 pub(crate) fn theta0(&self) -> Array1<f64> {
6690 let mut out =
6691 Array1::<f64>::zeros(self.rho_dim() + self.log_kappa_dim() + self.auxiliary_dim());
6692 out.slice_mut(s![..self.rho_dim()]).assign(&self.rho0);
6693 out.slice_mut(s![self.rho_dim()..self.rho_dim() + self.log_kappa_dim()])
6694 .assign(self.log_kappa0.as_array());
6695 out.slice_mut(s![self.rho_dim() + self.log_kappa_dim()..])
6696 .assign(&self.auxiliary0);
6697 out
6698 }
6699
6700 pub(crate) fn lower(&self) -> Array1<f64> {
6701 let mut out =
6702 Array1::<f64>::zeros(self.rho_dim() + self.log_kappa_dim() + self.auxiliary_dim());
6703 out.slice_mut(s![..self.rho_dim()]).assign(&self.rho_lower);
6704 out.slice_mut(s![self.rho_dim()..self.rho_dim() + self.log_kappa_dim()])
6705 .assign(self.log_kappa_lower.as_array());
6706 out.slice_mut(s![self.rho_dim() + self.log_kappa_dim()..])
6707 .assign(&self.auxiliary_lower);
6708 out
6709 }
6710
6711 pub(crate) fn upper(&self) -> Array1<f64> {
6712 let mut out =
6713 Array1::<f64>::zeros(self.rho_dim() + self.log_kappa_dim() + self.auxiliary_dim());
6714 out.slice_mut(s![..self.rho_dim()]).assign(&self.rho_upper);
6715 out.slice_mut(s![self.rho_dim()..self.rho_dim() + self.log_kappa_dim()])
6716 .assign(self.log_kappa_upper.as_array());
6717 out.slice_mut(s![self.rho_dim() + self.log_kappa_dim()..])
6718 .assign(&self.auxiliary_upper);
6719 out
6720 }
6721
6722 pub(crate) fn log_kappa_dims_per_term(&self) -> Vec<usize> {
6724 self.log_kappa0.dims_per_term().to_vec()
6725 }
6726}
6727
6728struct ExactJointDesignCache<'d> {
6734 realizers: Vec<FrozenTermCollectionIncrementalRealizer<'d>>,
6735 block_term_indices: Vec<Vec<usize>>,
6736 current_theta: Option<Array1<f64>>,
6737 last_cost: Option<f64>,
6738 last_eval: Option<(
6739 f64,
6740 Array1<f64>,
6741 gam_problem::HessianResult,
6742 )>,
6743 rho_dim: usize,
6744 all_dims: Vec<usize>,
6745 log_kappa_dim: usize,
6746 block_term_counts: Vec<usize>,
6747}
6748
6749impl<'d> ExactJointDesignCache<'d> {
6750 fn new(
6751 data: ArrayView2<'d, f64>,
6752 blocks: Vec<(TermCollectionSpec, TermCollectionDesign, Vec<usize>)>,
6753 rho_dim: usize,
6754 all_dims: Vec<usize>,
6755 ) -> Result<Self, String> {
6756 let n_blocks = blocks.len();
6757 let mut realizers = Vec::with_capacity(n_blocks);
6758 let mut block_term_indices = Vec::with_capacity(n_blocks);
6759 let mut block_term_counts = Vec::with_capacity(n_blocks);
6760
6761 for (spec, design, terms) in blocks {
6762 block_term_counts.push(terms.len());
6763 block_term_indices.push(terms);
6764 realizers.push(FrozenTermCollectionIncrementalRealizer::new(
6765 data, spec, design,
6766 )?);
6767 }
6768
6769 Ok(Self {
6770 realizers,
6771 block_term_indices,
6772 current_theta: None,
6773 last_cost: None,
6774 last_eval: None,
6775 rho_dim,
6776 log_kappa_dim: all_dims.iter().sum(),
6777 all_dims,
6778 block_term_counts,
6779 })
6780 }
6781
6782 fn ensure_theta(&mut self, theta: &Array1<f64>) -> Result<(), String> {
6783 if self
6784 .current_theta
6785 .as_ref()
6786 .is_some_and(|cached| theta_values_match(cached, theta))
6787 {
6788 return Ok(());
6789 }
6790
6791 let t_ensure = std::time::Instant::now();
6792 let kappa_theta_len = self.rho_dim + self.log_kappa_dim;
6793 if theta.len() < kappa_theta_len {
6794 return Err(SmoothError::dimension_mismatch(format!(
6795 "exact-joint theta length mismatch: got {}, expected at least {} (rho_dim={}, log_kappa_dim={})",
6796 theta.len(),
6797 kappa_theta_len,
6798 self.rho_dim,
6799 self.log_kappa_dim
6800 ))
6801 .into());
6802 }
6803 let theta_kappa = theta.slice(s![..kappa_theta_len]).to_owned();
6804 let full_log_kappa = SpatialLogKappaCoords::from_theta_tail_with_dims(
6805 &theta_kappa,
6806 self.rho_dim,
6807 self.all_dims.clone(),
6808 );
6809
6810 let n = self.realizers.len();
6814 let mut remaining = full_log_kappa;
6815 for block_idx in 0..n {
6816 let count = self.block_term_counts[block_idx];
6817 if block_idx < n - 1 {
6818 let (block_lk, rest) = remaining.split_at(count);
6819 self.realizers[block_idx]
6820 .apply_log_kappa(&block_lk, &self.block_term_indices[block_idx])?;
6821 remaining = rest;
6822 } else {
6823 self.realizers[block_idx]
6825 .apply_log_kappa(&remaining, &self.block_term_indices[block_idx])?;
6826 }
6827 }
6828
6829 log::info!(
6830 "[STAGE] ensure_theta (n-block, {} blocks, {} realizers): {:.3}s",
6831 n,
6832 self.realizers.len(),
6833 t_ensure.elapsed().as_secs_f64(),
6834 );
6835 self.current_theta = Some(theta.clone());
6836 self.last_cost = None;
6837 self.last_eval = None;
6838 Ok(())
6839 }
6840
6841 impl_exact_joint_theta_memo!();
6842
6843 fn store_cost_only(&mut self, theta: &Array1<f64>, cost: f64) {
6849 if self
6850 .current_theta
6851 .as_ref()
6852 .is_some_and(|cached| theta_values_match(cached, theta))
6853 {
6854 self.last_cost = Some(cost);
6855 }
6856 }
6857
6858 fn specs(&self) -> Vec<&TermCollectionSpec> {
6859 self.realizers.iter().map(|r| r.spec()).collect()
6860 }
6861
6862 fn designs(&self) -> Vec<&TermCollectionDesign> {
6863 self.realizers.iter().map(|r| r.design()).collect()
6864 }
6865
6866 fn design_revision(&self) -> u64 {
6876 self.realizers
6877 .iter()
6878 .fold(0u64, |acc, r| acc.wrapping_add(r.design_revision()))
6879 }
6880}
6881
6882pub(crate) fn seed_risk_profile_for_likelihood_family(
6883 family: &LikelihoodSpec,
6884) -> gam_problem::SeedRiskProfile {
6885 match &family.response {
6886 ResponseFamily::Gaussian => gam_problem::SeedRiskProfile::Gaussian,
6887 ResponseFamily::RoystonParmar => gam_problem::SeedRiskProfile::Survival,
6888 ResponseFamily::Binomial
6889 | ResponseFamily::Poisson
6890 | ResponseFamily::Tweedie { .. }
6891 | ResponseFamily::NegativeBinomial { .. }
6892 | ResponseFamily::Beta { .. }
6893 | ResponseFamily::Gamma => gam_problem::SeedRiskProfile::GeneralizedLinear,
6894 }
6895}
6896
6897const EXACT_JOINT_SECOND_ORDER_THETA_CAP: usize = 8;
6905
6906fn exact_joint_seed_config(
6907 risk_profile: gam_problem::SeedRiskProfile,
6908 auxiliary_dim: usize,
6909) -> gam_problem::SeedConfig {
6910 let mut config = gam_problem::SeedConfig {
6911 risk_profile,
6912 num_auxiliary_trailing: auxiliary_dim,
6913 ..Default::default()
6914 };
6915 match risk_profile {
6916 gam_problem::SeedRiskProfile::Gaussian
6917 | gam_problem::SeedRiskProfile::GaussianLocationScale => {
6918 config.max_seeds = 4;
6919 config.seed_budget = 2;
6920 }
6921 gam_problem::SeedRiskProfile::GeneralizedLinear => {
6922 config.max_seeds = 1;
6927 config.seed_budget = 1;
6928 config.screen_max_inner_iterations = 8;
6929 }
6930 gam_problem::SeedRiskProfile::Survival => {
6931 config.max_seeds = 8;
6937 config.seed_budget = 4;
6938 config.screen_max_inner_iterations = 8;
6939 }
6940 }
6941 config
6942}
6943
6944#[cfg(test)]
6945mod exact_joint_seed_config_tests {
6946 use super::*;
6947
6948 #[test]
6949 fn exact_joint_marginal_slope_profiles_get_deeper_startup_validation() {
6950 let bms = exact_joint_seed_config(gam_problem::SeedRiskProfile::GeneralizedLinear, 2);
6951 assert_eq!(bms.max_seeds, 1);
6952 assert_eq!(bms.seed_budget, 1);
6953 assert_eq!(bms.screen_max_inner_iterations, 8);
6954 assert_eq!(bms.num_auxiliary_trailing, 2);
6955
6956 let survival = exact_joint_seed_config(gam_problem::SeedRiskProfile::Survival, 3);
6957 assert_eq!(survival.max_seeds, 8);
6958 assert_eq!(survival.seed_budget, 4);
6959 assert_eq!(survival.screen_max_inner_iterations, 8);
6960 assert_eq!(survival.num_auxiliary_trailing, 3);
6961 }
6962
6963 #[test]
6964 fn exact_joint_gaussian_keeps_tight_historical_multistart_budget() {
6965 let gaussian = exact_joint_seed_config(gam_problem::SeedRiskProfile::Gaussian, 1);
6966 assert_eq!(gaussian.max_seeds, 4);
6967 assert_eq!(gaussian.seed_budget, 2);
6968 assert_eq!(
6969 gaussian.screen_max_inner_iterations,
6970 gam_problem::SeedConfig::default().screen_max_inner_iterations
6971 );
6972 assert_eq!(gaussian.num_auxiliary_trailing, 1);
6973 }
6974}
6975
6976#[cfg(test)]
6977mod wood_reference_df_tests {
6978 use super::*;
6979
6980 #[test]
6986 fn edf1_equals_two_trace_minus_trace_of_square() {
6987 let f = ndarray::array![[0.9_f64, 0.0], [0.0, 0.4]];
6991 let got = wood_reference_df(Some(&f), &(0..2)).unwrap();
6992 assert!(
6993 (got - 1.63).abs() < 1e-12,
6994 "edf1 should be 2*tr - tr(F^2) = 1.63, got {got}"
6995 );
6996 let edf = 1.3;
6999 assert!(got >= edf - 1e-12, "edf1 {got} must be >= edf {edf}");
7000 }
7001
7002 #[test]
7003 fn edf1_never_collapses_below_edf_when_offdiagonals_blow_up() {
7004 let f = ndarray::array![[0.5_f64, 40.0], [40.0, 0.5]];
7011 let tr = 1.0_f64;
7012 let got = wood_reference_df(Some(&f), &(0..2)).unwrap();
7013 assert!(
7014 got >= tr - 1e-12,
7015 "edf1 must be floored at edf (=tr={tr}) even when tr(F^2) explodes, got {got}"
7016 );
7017 assert!(got.is_finite() && got > 0.0, "edf1 must stay finite/positive");
7018 }
7019
7020 #[test]
7021 fn returns_none_on_nonpositive_or_missing_trace() {
7022 assert!(wood_reference_df(None, &(0..2)).is_none());
7025 let zero = ndarray::array![[0.0_f64, 0.0], [0.0, 0.0]];
7027 assert!(wood_reference_df(Some(&zero), &(0..2)).is_none());
7028 let f = ndarray::array![[0.5_f64, 0.0], [0.0, 0.5]];
7030 assert!(wood_reference_df(Some(&f), &(0..5)).is_none());
7031 }
7032}
7033
7034pub(crate) fn exact_joint_multistart_outer_problem(
7035 theta0: &Array1<f64>,
7036 lower: &Array1<f64>,
7037 upper: &Array1<f64>,
7038 rho_dim: usize,
7039 auxiliary_dim: usize,
7040 n_params: usize,
7041 gradient: gam_problem::Derivative,
7042 hessian: gam_problem::DeclaredHessianForm,
7043 prefer_gradient_only: bool,
7044 disable_fixed_point: bool,
7045 risk_profile: gam_problem::SeedRiskProfile,
7046 tolerance: f64,
7047 max_iter: usize,
7048 bfgs_step_cap: Option<f64>,
7057 bfgs_step_cap_psi: Option<f64>,
7058 screening_cap: Option<Arc<AtomicUsize>>,
7059 profiled_objective_size: Option<(usize, usize)>,
7080 has_constant_curvature: bool,
7089) -> gam_solve::rho_optimizer::OuterProblem {
7090 let mut seed_heuristic = theta0.to_vec();
7091 for value in &mut seed_heuristic[..rho_dim] {
7092 *value = value.exp();
7093 }
7094 let rho_ceiling = if has_constant_curvature {
7099 gam_solve::estimate::RHO_BOUND
7100 } else {
7101 12.0
7102 };
7103 let mut problem = gam_solve::rho_optimizer::OuterProblem::new(n_params)
7104 .with_gradient(gradient)
7105 .with_hessian(hessian)
7106 .with_prefer_gradient_only(prefer_gradient_only)
7107 .with_disable_fixed_point(disable_fixed_point)
7108 .with_fallback_policy(gam_solve::rho_optimizer::FallbackPolicy::Automatic)
7118 .with_psi_dim(auxiliary_dim)
7119 .with_tolerance(tolerance)
7120 .with_max_iter(max_iter)
7121 .with_bounds(lower.clone(), upper.clone())
7122 .with_initial_rho(theta0.clone())
7123 .with_bfgs_step_cap(bfgs_step_cap)
7124 .with_bfgs_step_cap_psi(bfgs_step_cap_psi)
7125 .with_seed_config({
7126 let mut sc = exact_joint_seed_config(risk_profile, auxiliary_dim);
7127 if has_constant_curvature {
7128 sc.bounds = (sc.bounds.0, rho_ceiling);
7132 }
7151 sc
7152 })
7153 .with_rho_bound(rho_ceiling)
7154 .with_heuristic_lambdas(seed_heuristic);
7155 if let Some((n_obs, p_cols)) = profiled_objective_size {
7156 problem = problem
7164 .with_objective_scale(Some(n_obs as f64))
7165 .with_problem_size(n_obs, p_cols)
7166 .with_arc_initial_regularization(Some(0.25))
7167 .with_operator_initial_trust_radius(Some(4.0));
7168 }
7169 if let Some(screening_cap) = screening_cap {
7170 problem = problem
7171 .with_screening_cap(screening_cap)
7172 .with_screen_initial_rho(true);
7173 }
7174 problem
7175}
7176
7177fn kappa_phase_failure_is_fixed_kappa_recoverable(message: &str) -> bool {
7188 message.contains("no candidate seeds passed outer startup validation")
7189 || message.contains("joint hyper rho dimension mismatch")
7190 || message.contains("objective returned a non-finite cost")
7191}
7192
7193pub fn optimize_spatial_length_scale_exact_joint<FitOut, FitFn, ExactFn, ExactEfsFn, SeedFn>(
7194 data: ArrayView2<'_, f64>,
7195 block_specs: &[TermCollectionSpec],
7196 block_term_indices: &[Vec<usize>],
7197 kappa_options: &SpatialLengthScaleOptimizationOptions,
7198 joint_setup: &ExactJointHyperSetup,
7199 seed_risk_profile: gam_problem::SeedRiskProfile,
7200 analytic_joint_gradient_available: bool,
7201 analytic_joint_hessian_available: bool,
7202 disable_fixed_point: bool,
7203 screening_cap: Option<Arc<AtomicUsize>>,
7204 outer_derivative_policy: gam_model_api::families::custom_family::OuterDerivativePolicy,
7205 mut fit_fn: FitFn,
7206 mut exact_fn: ExactFn,
7207 mut exact_efs_fn: ExactEfsFn,
7208 mut seed_inner_beta_fn: SeedFn,
7209) -> Result<SpatialLengthScaleOptimizationResult<FitOut>, String>
7210where
7211 FitOut: Clone,
7212 FitFn: FnMut(
7213 &Array1<f64>,
7214 &[TermCollectionSpec],
7215 &[TermCollectionDesign],
7216 ) -> Result<FitOut, String>,
7217 ExactFn: FnMut(
7218 &Array1<f64>,
7219 &[TermCollectionSpec],
7220 &[TermCollectionDesign],
7221 gam_solve::estimate::reml::reml_outer_engine::EvalMode,
7222 &gam_problem::outer_subsample::RowSet,
7223 ) -> Result<
7224 (
7225 f64,
7226 Array1<f64>,
7227 gam_problem::HessianResult,
7228 ),
7229 String,
7230 >,
7231 ExactEfsFn: FnMut(
7232 &Array1<f64>,
7233 &[TermCollectionSpec],
7234 &[TermCollectionDesign],
7235 ) -> Result<gam_problem::EfsEval, String>,
7236 SeedFn:
7237 FnMut(&Array1<f64>) -> Result<gam_solve::rho_optimizer::SeedOutcome, EstimationError>,
7238{
7239 let n_blocks = block_specs.len();
7240 if block_term_indices.len() != n_blocks {
7241 return Err(SmoothError::dimension_mismatch(format!(
7242 "block_specs ({}) and block_term_indices ({}) length mismatch",
7243 n_blocks,
7244 block_term_indices.len()
7245 ))
7246 .into());
7247 }
7248
7249 let log_kappa_dim = joint_setup.log_kappa_dim();
7250
7251 log::warn!(
7252 "[OUTER-FD-AUDIT/spatial-exact-joint] driver entry: aux_dim={} log_kappa_dim={} kappa_enabled={} rho_dim={} theta0_len={}",
7253 joint_setup.auxiliary_dim(),
7254 log_kappa_dim,
7255 kappa_options.enabled,
7256 joint_setup.rho_dim(),
7257 joint_setup.theta0().len()
7258 );
7259
7260 if joint_setup.auxiliary_dim() == 0 && (!kappa_options.enabled || log_kappa_dim == 0) {
7264 log::warn!(
7265 "[OUTER-FD-AUDIT/spatial-exact-joint] taking FAST path (no outer theta optimization in this driver)"
7266 );
7267 let (designs, resolved_specs) = build_term_collection_designs_and_freeze_joint(
7268 data, block_specs,
7269 )
7270 .map_err(|e| {
7271 format!("failed to build and freeze joint block designs during exact joint kappa optimization: {e}")
7272 })?;
7273 let theta0 = joint_setup.theta0();
7274
7275 let spec_refs: Vec<TermCollectionSpec> = resolved_specs.clone();
7277 let design_refs: Vec<TermCollectionDesign> = designs.clone();
7278 let fit = fit_fn(&theta0, &spec_refs, &design_refs)?;
7279 return Ok(SpatialLengthScaleOptimizationResult {
7280 resolved_specs,
7281 designs,
7282 fit,
7283 timing: None,
7284 });
7285 }
7286
7287 let theta0 = joint_setup.theta0();
7291 let lower = joint_setup.lower();
7292 let upper = joint_setup.upper();
7293 if theta0.len() < log_kappa_dim || lower.len() != theta0.len() || upper.len() != theta0.len() {
7294 return Err(SmoothError::dimension_mismatch(format!(
7295 "invalid exact joint theta setup: theta0={}, lower={}, upper={}, required_log_kappa_dim={}",
7296 theta0.len(),
7297 lower.len(),
7298 upper.len(),
7299 log_kappa_dim
7300 ))
7301 .into());
7302 }
7303 let rho_dim = joint_setup.rho_dim();
7304 let all_dims = joint_setup.log_kappa_dims_per_term();
7305
7306 let (boot_designs, best_specs) = build_term_collection_designs_and_freeze_joint(
7308 data,
7309 block_specs,
7310 )
7311 .map_err(|e| {
7312 format!(
7313 "failed to build and freeze joint block designs during exact joint kappa bootstrap: {e}"
7314 )
7315 })?;
7316 let policy_hessian_form = outer_derivative_policy.declared_hessian_form();
7326 let analytic_outer_hessian_available = analytic_joint_hessian_available
7327 && matches!(
7328 policy_hessian_form,
7329 gam_problem::DeclaredHessianForm::Either
7330 | gam_problem::DeclaredHessianForm::Dense
7331 | gam_problem::DeclaredHessianForm::Operator { .. }
7332 );
7333 let prefer_gradient_only = !analytic_outer_hessian_available;
7334
7335 let theta_dim = theta0.len();
7336 let psi_dim = theta_dim - rho_dim;
7337
7338 let cache_blocks: Vec<(TermCollectionSpec, TermCollectionDesign, Vec<usize>)> = best_specs
7340 .iter()
7341 .zip(boot_designs.iter())
7342 .zip(block_term_indices.iter())
7343 .map(|((spec, design), terms)| (spec.clone(), design.clone(), terms.clone()))
7344 .collect();
7345
7346 struct NBlockExactJointState<'d> {
7347 cache: ExactJointDesignCache<'d>,
7348 }
7349
7350 let mut state = NBlockExactJointState {
7351 cache: ExactJointDesignCache::new(data, cache_blocks, rho_dim, all_dims.clone())?,
7352 };
7353
7354 const KAPPA_PILOT_K: usize = 5_000;
7379 const KAPPA_POLISH_K: usize = 25_000;
7380 const KAPPA_POLISH_TRIGGER_N: usize = 100_000;
7381
7382 let n_total = data.nrows();
7383 let use_staged_kappa = outer_derivative_policy.should_use_staged_kappa(n_total);
7384 if use_staged_kappa {
7385 log::info!(
7386 "[KAPPA-STAGED] auto-engaging pilot+polish schedule: n={} pilot_k={} polish_k={}",
7387 n_total,
7388 KAPPA_PILOT_K,
7389 KAPPA_POLISH_K,
7390 );
7391 }
7392
7393 fn build_uniform_pilot_subsample(
7410 n_total: usize,
7411 k_target: usize,
7412 seed: u64,
7413 ) -> gam_problem::outer_subsample::OuterScoreSubsample {
7414 use gam_problem::outer_subsample::OuterScoreSubsample;
7415 let k = k_target.min(n_total);
7416 if k == 0 || n_total == 0 {
7417 return OuterScoreSubsample::from_uniform_inclusion_mask(Vec::new(), n_total, seed);
7418 }
7419 let mut mask: Vec<usize> = Vec::with_capacity(k);
7423 let mut state = seed.wrapping_add(0x9E3779B97F4A7C15);
7425 let splitmix = |s: &mut u64| -> u64 { gam_linalg::utils::splitmix64(s) };
7426 let mut taken = std::collections::HashSet::with_capacity(k);
7427 for j in (n_total - k)..n_total {
7428 let r = (splitmix(&mut state) % (j as u64 + 1)) as usize;
7429 if !taken.insert(r) {
7430 taken.insert(j);
7431 mask.push(j);
7432 } else {
7433 mask.push(r);
7434 }
7435 }
7436 mask.sort_unstable();
7437 mask.dedup();
7438 OuterScoreSubsample::from_uniform_inclusion_mask(mask, n_total, seed)
7439 }
7440
7441 let current_row_set: std::cell::RefCell<gam_problem::outer_subsample::RowSet> = if use_staged_kappa {
7442 let pilot = build_uniform_pilot_subsample(n_total, KAPPA_PILOT_K, n_total as u64);
7443 std::cell::RefCell::new(gam_problem::outer_subsample::RowSet::Subsample {
7444 rows: std::sync::Arc::clone(&pilot.rows),
7445 n_full: n_total,
7446 })
7447 } else {
7448 std::cell::RefCell::new(gam_problem::outer_subsample::RowSet::All)
7449 };
7450
7451 let exact_fn_cell = std::cell::RefCell::new(&mut exact_fn);
7452 let exact_efs_fn_cell = std::cell::RefCell::new(&mut exact_efs_fn);
7453
7454 use std::cell::Cell;
7469 let kphase_cost_calls: Cell<usize> = Cell::new(0);
7470 let kphase_cost_total_s: Cell<f64> = Cell::new(0.0);
7471 let kphase_eval_calls: Cell<usize> = Cell::new(0);
7472 let kphase_eval_total_s: Cell<f64> = Cell::new(0.0);
7473 let kphase_efs_calls: Cell<usize> = Cell::new(0);
7474 let kphase_efs_total_s: Cell<f64> = Cell::new(0.0);
7475 let kphase_optim_start = std::time::Instant::now();
7476 let kphase_log_kappa_dim = log_kappa_dim;
7477 let kphase_log_norms = |theta: &Array1<f64>| -> (f64, f64) {
7478 let theta_norm = theta.iter().map(|v| v * v).sum::<f64>().sqrt();
7479 let log_kappa_norm = if kphase_log_kappa_dim > 0 && theta.len() >= kphase_log_kappa_dim {
7480 let start = theta.len() - kphase_log_kappa_dim;
7481 theta.iter().skip(start).map(|v| v * v).sum::<f64>().sqrt()
7482 } else {
7483 0.0
7484 };
7485 (theta_norm, log_kappa_norm)
7486 };
7487
7488 use gam_solve::rho_optimizer::OuterEvalOrder;
7489 use gam_problem::{DeclaredHessianForm, Derivative, OuterEval};
7490
7491 let joint_p_cols: usize = boot_designs
7495 .iter()
7496 .map(|d| d.design.ncols())
7497 .sum::<usize>()
7498 .max(1);
7499
7500 let problem = exact_joint_multistart_outer_problem(
7501 &theta0,
7502 &lower,
7503 &upper,
7504 rho_dim,
7505 psi_dim,
7506 theta_dim,
7507 if analytic_joint_gradient_available {
7508 Derivative::Analytic
7509 } else {
7510 Derivative::Unavailable
7511 },
7512 if analytic_outer_hessian_available {
7513 DeclaredHessianForm::Either
7514 } else {
7515 DeclaredHessianForm::Unavailable
7516 },
7517 prefer_gradient_only,
7518 disable_fixed_point,
7519 seed_risk_profile,
7520 kappa_options.rel_tol.max(1e-6),
7521 kappa_options.max_outer_iter.max(1),
7522 Some(5.0),
7524 Some(kappa_options.log_step.clamp(0.25, 1.0)),
7526 screening_cap.clone(),
7527 Some((n_total, joint_p_cols)),
7530 block_specs
7533 .iter()
7534 .any(|s| !constant_curvature_term_indices(s).is_empty()),
7535 );
7536
7537 fn collect_specs(cache: &ExactJointDesignCache<'_>) -> Vec<TermCollectionSpec> {
7539 cache.specs().into_iter().cloned().collect()
7540 }
7541 fn collect_designs(cache: &ExactJointDesignCache<'_>) -> Vec<TermCollectionDesign> {
7542 cache.designs().into_iter().cloned().collect()
7543 }
7544
7545 let result = {
7546 let eval_outer = |ctx: &mut &mut NBlockExactJointState<'_>,
7547 theta: &Array1<f64>,
7548 order: OuterEvalOrder|
7549 -> Result<OuterEval, EstimationError> {
7550 if let Some((cost, grad, hess)) = ctx.cache.memoized_eval(theta) {
7551 let cached_satisfies_order = match order {
7552 OuterEvalOrder::Value => true,
7553 OuterEvalOrder::ValueAndGradient => true,
7554 OuterEvalOrder::ValueGradientHessian => hess.is_analytic(),
7555 };
7556 if cached_satisfies_order {
7557 if !cost.is_finite() {
7558 return Ok(OuterEval::infeasible(theta.len()));
7559 }
7560 if grad.iter().any(|v| !v.is_finite()) {
7573 return Ok(OuterEval::infeasible(theta.len()));
7574 }
7575 return Ok(OuterEval {
7576 cost,
7577 gradient: grad,
7578 hessian: hess,
7579 inner_beta_hint: None,
7580 });
7581 }
7582 }
7583 if let Err(err) = ctx.cache.ensure_theta(theta) {
7584 log::warn!(
7585 "[OUTER] n-block exact-joint spatial: ensure_theta failed during gradient evaluation: {err}"
7586 );
7587 return Ok(OuterEval::infeasible(theta.len()));
7588 }
7589 let design_revision = Some(ctx.cache.design_revision());
7590 let specs = collect_specs(&ctx.cache);
7591 let designs = collect_designs(&ctx.cache);
7592 let clamped = outer_derivative_policy.order_for_evaluation(order);
7600 let need_hessian = matches!(clamped, OuterEvalOrder::ValueGradientHessian)
7601 && analytic_outer_hessian_available;
7602 let eval_mode = if need_hessian {
7603 gam_solve::estimate::reml::reml_outer_engine::EvalMode::ValueGradientHessian
7604 } else {
7605 gam_solve::estimate::reml::reml_outer_engine::EvalMode::ValueAndGradient
7606 };
7607 let t0 = std::time::Instant::now();
7608 let result = {
7609 let row_set_borrow = current_row_set.borrow();
7610 (*exact_fn_cell.borrow_mut())(theta, &specs, &designs, eval_mode, &row_set_borrow)
7611 };
7612 let elapsed_s = t0.elapsed().as_secs_f64();
7613 kphase_eval_calls.set(kphase_eval_calls.get() + 1);
7614 kphase_eval_total_s.set(kphase_eval_total_s.get() + elapsed_s);
7615 let (theta_norm, log_kappa_norm) = kphase_log_norms(theta);
7616 log::info!(
7617 "[KAPPA-PHASE] phase=eval_outer call={} order={:?} design_revision={:?} theta_norm={:.4e} log_kappa_norm={:.4e} elapsed_s={:.4}",
7618 kphase_eval_calls.get(),
7619 order,
7620 design_revision,
7621 theta_norm,
7622 log_kappa_norm,
7623 elapsed_s,
7624 );
7625 match result {
7626 Ok((cost, grad, hess)) => {
7627 ctx.cache.store_eval((cost, grad.clone(), hess.clone()));
7628 if !cost.is_finite() {
7629 return Ok(OuterEval::infeasible(theta.len()));
7630 }
7631 if grad.iter().any(|v| !v.is_finite()) {
7644 return Ok(OuterEval::infeasible(theta.len()));
7645 }
7646 Ok(OuterEval {
7647 cost,
7648 gradient: grad,
7649 hessian: hess,
7650 inner_beta_hint: None,
7651 })
7652 }
7653 Err(err) => {
7654 log::warn!(
7655 "[OUTER] n-block exact-joint spatial: exact evaluation failed: {err}"
7656 );
7657 Ok(OuterEval::infeasible(theta.len()))
7658 }
7659 }
7660 };
7661
7662 let obj = problem.build_objective_with_eval_order(
7663 &mut state,
7664 |ctx: &mut &mut NBlockExactJointState<'_>, theta: &Array1<f64>| {
7665 if let Some(cost) = ctx.cache.memoized_cost(theta) {
7666 return Ok(cost);
7667 }
7668 if let Err(err) = ctx.cache.ensure_theta(theta) {
7669 log::warn!(
7670 "[OUTER] n-block exact-joint spatial: ensure_theta failed during cost evaluation: {err}"
7671 );
7672 return Ok(f64::INFINITY);
7673 }
7674 let design_revision = Some(ctx.cache.design_revision());
7675 let specs = collect_specs(&ctx.cache);
7676 let designs = collect_designs(&ctx.cache);
7677 let t0 = std::time::Instant::now();
7684 let result = {
7685 let row_set_borrow = current_row_set.borrow();
7686 (*exact_fn_cell.borrow_mut())(
7687 theta,
7688 &specs,
7689 &designs,
7690 gam_solve::estimate::reml::reml_outer_engine::EvalMode::ValueOnly,
7691 &row_set_borrow,
7692 )
7693 };
7694 let elapsed_s = t0.elapsed().as_secs_f64();
7695 kphase_cost_calls.set(kphase_cost_calls.get() + 1);
7696 kphase_cost_total_s.set(kphase_cost_total_s.get() + elapsed_s);
7697 let (theta_norm, log_kappa_norm) = kphase_log_norms(theta);
7698 log::info!(
7699 "[KAPPA-PHASE] phase=cost call={} design_revision={:?} theta_norm={:.4e} log_kappa_norm={:.4e} elapsed_s={:.4}",
7700 kphase_cost_calls.get(),
7701 design_revision,
7702 theta_norm,
7703 log_kappa_norm,
7704 elapsed_s,
7705 );
7706 match result {
7707 Ok((cost, _grad, _hess)) => {
7708 ctx.cache.store_cost_only(theta, cost);
7714 Ok(cost)
7715 }
7716 Err(err) => {
7717 log::warn!(
7718 "[OUTER] n-block exact-joint spatial: exact cost evaluation failed: {err}"
7719 );
7720 Ok(f64::INFINITY)
7721 }
7722 }
7723 },
7724 |ctx: &mut &mut NBlockExactJointState<'_>, theta: &Array1<f64>| {
7725 eval_outer(
7726 ctx,
7727 theta,
7728 if analytic_outer_hessian_available {
7729 OuterEvalOrder::ValueGradientHessian
7730 } else {
7731 OuterEvalOrder::ValueAndGradient
7732 },
7733 )
7734 },
7735 |ctx: &mut &mut NBlockExactJointState<'_>,
7736 theta: &Array1<f64>,
7737 order: OuterEvalOrder| { eval_outer(ctx, theta, order) },
7738 None::<fn(&mut &mut NBlockExactJointState<'_>)>,
7739 Some(
7740 |ctx: &mut &mut NBlockExactJointState<'_>, theta: &Array1<f64>| {
7741 ctx.cache
7742 .ensure_theta(theta)
7743 .map_err(EstimationError::InvalidInput)?;
7744 let design_revision = Some(ctx.cache.design_revision());
7745 let specs = collect_specs(&ctx.cache);
7746 let designs = collect_designs(&ctx.cache);
7747 let t0 = std::time::Instant::now();
7748 let eval_result = (*exact_efs_fn_cell.borrow_mut())(
7749 theta,
7750 &specs,
7751 &designs,
7752 );
7753 let elapsed_s = t0.elapsed().as_secs_f64();
7754 kphase_efs_calls.set(kphase_efs_calls.get() + 1);
7755 kphase_efs_total_s.set(kphase_efs_total_s.get() + elapsed_s);
7756 let (theta_norm, log_kappa_norm) = kphase_log_norms(theta);
7757 log::info!(
7758 "[KAPPA-PHASE] phase=efs call={} design_revision={:?} theta_norm={:.4e} log_kappa_norm={:.4e} elapsed_s={:.4}",
7759 kphase_efs_calls.get(),
7760 design_revision,
7761 theta_norm,
7762 log_kappa_norm,
7763 elapsed_s,
7764 );
7765 let eval = eval_result.map_err(EstimationError::RemlOptimizationFailed)?;
7766 Ok(eval)
7767 },
7768 ),
7769 );
7770 let mut obj = obj.with_seed_inner_state(
7771 move |_ctx: &mut &mut NBlockExactJointState<'_>, beta: &Array1<f64>| {
7772 (seed_inner_beta_fn)(beta)
7773 },
7774 );
7775
7776 match problem.run(&mut obj, "n-block exact-joint spatial") {
7777 Ok(result) => result,
7778 Err(e) => {
7779 let message = e.to_string();
7780 if kappa_phase_failure_is_fixed_kappa_recoverable(&message) {
7800 drop(obj);
7801 log::warn!(
7802 "[KAPPA-PHASE] length-scale optimization could not validate any seed \
7803 ({message}); falling back to a FIXED bootstrap κ (skipping κ \
7804 optimization) and fitting there — a real model at the initial \
7805 length-scale rather than raising (gam#787/#860)."
7806 );
7807 let (designs, resolved_specs) =
7808 build_term_collection_designs_and_freeze_joint(data, block_specs).map_err(
7809 |build_err| {
7810 format!(
7811 "fixed-κ fallback failed to build and freeze joint block \
7812 designs after κ optimization could not validate a seed \
7813 ({message}): {build_err}"
7814 )
7815 },
7816 )?;
7817 let fixed_theta0 = joint_setup.theta0();
7818 let spec_refs: Vec<TermCollectionSpec> = resolved_specs.clone();
7819 let design_refs: Vec<TermCollectionDesign> = designs.clone();
7820 let fit = fit_fn(&fixed_theta0, &spec_refs, &design_refs)?;
7821 return Ok(SpatialLengthScaleOptimizationResult {
7822 resolved_specs,
7823 designs,
7824 fit,
7825 timing: None,
7826 });
7827 }
7828 return Err(message);
7829 }
7830 }
7831 }; let kphase_total_s = kphase_optim_start.elapsed().as_secs_f64();
7841 log::info!(
7842 "[KAPPA-PHASE-SUMMARY] log_kappa_dim={} n_cost={} cost_total_s={:.4} n_eval={} eval_total_s={:.4} n_efs={} efs_total_s={:.4} optim_total_s={:.4}",
7843 kphase_log_kappa_dim,
7844 kphase_cost_calls.get(),
7845 kphase_cost_total_s.get(),
7846 kphase_eval_calls.get(),
7847 kphase_eval_total_s.get(),
7848 kphase_efs_calls.get(),
7849 kphase_efs_total_s.get(),
7850 kphase_total_s,
7851 );
7852 let timing = SpatialLengthScaleOptimizationTiming {
7853 log_kappa_dim: kphase_log_kappa_dim,
7854 cost_calls: kphase_cost_calls.get(),
7855 cost_total_s: kphase_cost_total_s.get(),
7856 eval_calls: kphase_eval_calls.get(),
7857 eval_total_s: kphase_eval_total_s.get(),
7858 efs_calls: kphase_efs_calls.get(),
7859 efs_total_s: kphase_efs_total_s.get(),
7860 slow_path_resets: 0,
7861 design_revision_delta: 0,
7862 nfree_skip_row_touches: 0,
7863 nfree_miss_shape: 0,
7864 nfree_miss_value: 0,
7865 nfree_miss_gradient: 0,
7866 nfree_miss_penalty: 0,
7867 nfree_miss_revision: 0,
7868 nfree_miss_second_order: 0,
7869 nfree_miss_other: 0,
7870 optim_total_s: kphase_total_s,
7871 };
7872
7873 let theta_star = result.rho;
7874
7875 if use_staged_kappa && n_total >= KAPPA_POLISH_TRIGGER_N {
7892 let polish = build_uniform_pilot_subsample(
7893 n_total,
7894 KAPPA_POLISH_K,
7895 (n_total as u64).wrapping_add(0xA5A5A5A5),
7896 );
7897 *current_row_set.borrow_mut() = gam_problem::outer_subsample::RowSet::Subsample {
7898 rows: std::sync::Arc::clone(&polish.rows),
7899 n_full: n_total,
7900 };
7901 log::info!(
7902 "[KAPPA-STAGED] rotating to polish subsample: k={} at theta_star",
7903 polish.rows.len(),
7904 );
7905 state.cache.ensure_theta(&theta_star)?;
7909 let (polish_cost, polish_grad, _) = {
7910 let specs = collect_specs(&state.cache);
7911 let designs = collect_designs(&state.cache);
7912 let row_set_borrow = current_row_set.borrow();
7913 exact_fn(
7914 &theta_star,
7915 &specs,
7916 &designs,
7917 gam_solve::estimate::reml::reml_outer_engine::EvalMode::ValueAndGradient,
7918 &row_set_borrow,
7919 )?
7920 };
7921 if !polish_cost.is_finite() || polish_grad.iter().any(|value| !value.is_finite()) {
7922 return Err(
7923 "polish subsample exact-joint evaluation produced non-finite objective pieces"
7924 .to_string(),
7925 );
7926 }
7927 }
7928 *current_row_set.borrow_mut() = gam_problem::outer_subsample::RowSet::All;
7929 if use_staged_kappa {
7930 log::info!(
7931 "[KAPPA-STAGED] rotating to full data for final coefficient fit (n={})",
7932 n_total,
7933 );
7934 }
7935
7936 state.cache.ensure_theta(&theta_star)?;
7937
7938 let resolved_specs: Vec<TermCollectionSpec> = collect_specs(&state.cache);
7939 let designs: Vec<TermCollectionDesign> = collect_designs(&state.cache);
7940
7941 let fit = fit_fn(&theta_star, &resolved_specs, &designs)?;
7942
7943 for spec in &resolved_specs {
7944 log_spatial_aniso_scales(spec);
7945 }
7946
7947 Ok(SpatialLengthScaleOptimizationResult {
7948 resolved_specs,
7949 designs,
7950 fit,
7951 timing: Some(timing),
7952 })
7953}
7954
7955fn try_exact_joint_latent_coord_optimization(
7956 data: ArrayView2<'_, f64>,
7957 y: ArrayView1<'_, f64>,
7958 weights: ArrayView1<'_, f64>,
7959 offset: ArrayView1<'_, f64>,
7960 resolvedspec: &TermCollectionSpec,
7961 best: &FittedTermCollection,
7962 family: LikelihoodSpec,
7963 options: &FitOptions,
7964 latent: &StandardLatentCoordConfig,
7965) -> Result<FittedTermCollectionWithSpec, EstimationError> {
7966 use gam_solve::rho_optimizer::OuterEvalOrder;
7967 use gam_problem::{DeclaredHessianForm, Derivative, OuterEval};
7968
7969 let rho_dim = best.fit.lambdas.len();
7970 let latent_flat_dim = latent.values.len();
7971 if latent_flat_dim == 0 {
7972 crate::bail_invalid_estim!(
7973 "latent-coordinate optimization requires a non-empty latent block"
7974 );
7975 }
7976 let direct_hypers =
7977 latent_coord_initial_direct_hypers(latent.values.id_mode(), latent.values.latent_dim())?;
7978 let analytic_rho_count = latent
7979 .analytic_penalties
7980 .as_ref()
7981 .map_or(0, |registry| registry.total_rho_count());
7982 let latent_coord_ext_dim = latent_flat_dim + analytic_rho_count + direct_hypers.len();
7983
7984 let mut theta0 = Array1::<f64>::zeros(rho_dim + latent_coord_ext_dim);
7985 theta0
7986 .slice_mut(s![..rho_dim])
7987 .assign(&best.fit.lambdas.mapv(f64::ln));
7988 theta0
7989 .slice_mut(s![rho_dim..rho_dim + latent_flat_dim])
7990 .assign(latent.values.as_flat());
7991 if !direct_hypers.is_empty() {
7992 let direct_start = rho_dim + latent_flat_dim + analytic_rho_count;
7993 theta0
7994 .slice_mut(s![direct_start..direct_start + direct_hypers.len()])
7995 .assign(&direct_hypers);
7996 }
7997
7998 let mut lower = Array1::<f64>::from_elem(theta0.len(), -12.0);
7999 let mut upper = Array1::<f64>::from_elem(theta0.len(), 12.0);
8000 let latent_bound = latent
8001 .values
8002 .as_flat()
8003 .iter()
8004 .fold(1.0_f64, |acc, &v| acc.max(v.abs()))
8005 + 10.0;
8006 for axis in rho_dim..rho_dim + latent_flat_dim {
8007 lower[axis] = -latent_bound;
8008 upper[axis] = latent_bound;
8009 }
8010
8011 struct LatentJointContext<'d> {
8012 rho_dim: usize,
8013 cache: SingleBlockLatentCoordDesignCache,
8014 evaluator: gam_solve::estimate::ExternalJointHyperEvaluator<'d>,
8015 }
8016
8017 impl<'d> LatentJointContext<'d> {
8018 fn eval_full(
8019 &mut self,
8020 theta: &Array1<f64>,
8021 order: OuterEvalOrder,
8022 ) -> Result<
8023 (
8024 f64,
8025 Array1<f64>,
8026 gam_problem::HessianResult,
8027 ),
8028 EstimationError,
8029 > {
8030 if let Some(eval) = self.cache.memoized_eval(theta) {
8031 return Ok(eval);
8032 }
8033 self.cache
8034 .ensure_theta(theta)
8035 .map_err(EstimationError::InvalidInput)?;
8036 let hyper_dirs = self
8037 .cache
8038 .hyper_dirs()
8039 .map_err(EstimationError::InvalidInput)?;
8040 let design_revision = Some(self.cache.design_revision());
8041 let registry_for_key = self.cache.analytic_penalties();
8042 self.evaluator
8043 .set_analytic_penalty_registry(registry_for_key.as_deref());
8044 let mut eval = evaluate_joint_reml_outer_eval_at_theta(
8045 &mut self.evaluator,
8046 self.cache.design(),
8047 theta,
8048 self.rho_dim,
8049 hyper_dirs,
8050 None,
8051 order,
8052 design_revision,
8053 )?;
8054 let latent = self.cache.latent().map_err(EstimationError::InvalidInput)?;
8055 if let Some(registry) = registry_for_key {
8056 let mut registry = registry.as_ref().clone();
8057 registry.apply_weight_schedules(
8058 gam_solve::estimate::reml::outer_eval::current_outer_iter() as usize,
8059 );
8060 add_analytic_penalty_objective_to_eval(
8061 theta,
8062 self.rho_dim,
8063 latent.as_ref(),
8064 ®istry,
8065 &mut eval,
8066 )?;
8067 }
8068 add_latent_id_objective_to_eval(
8069 theta,
8070 self.rho_dim,
8071 self.cache.analytic_penalty_rho_count(),
8072 latent.as_ref(),
8073 &mut eval,
8074 )?;
8075 self.cache.store_eval(eval.clone());
8076 Ok(eval)
8077 }
8078
8079 fn eval_efs(
8080 &mut self,
8081 theta: &Array1<f64>,
8082 ) -> Result<gam_problem::EfsEval, EstimationError> {
8083 self.cache
8084 .ensure_theta(theta)
8085 .map_err(EstimationError::InvalidInput)?;
8086 let hyper_dirs = self
8087 .cache
8088 .hyper_dirs()
8089 .map_err(EstimationError::InvalidInput)?;
8090 let registry_for_key = self.cache.analytic_penalties();
8091 self.evaluator
8092 .set_analytic_penalty_registry(registry_for_key.as_deref());
8093 let mut efs = evaluate_joint_reml_efs_at_theta(
8094 &mut self.evaluator,
8095 self.cache.design(),
8096 theta,
8097 self.rho_dim,
8098 hyper_dirs,
8099 None,
8100 Some(self.cache.design_revision()),
8101 )?;
8102 if let Some(registry) = registry_for_key {
8103 let mut registry = registry.as_ref().clone();
8104 registry.apply_weight_schedules(
8105 gam_solve::estimate::reml::outer_eval::current_outer_iter() as usize,
8106 );
8107 let latent = self.cache.latent().map_err(EstimationError::InvalidInput)?;
8108 let contribution = analytic_penalty_objective_contribution(
8109 theta,
8110 self.rho_dim,
8111 latent.as_ref(),
8112 ®istry,
8113 )?;
8114 efs.cost += contribution.cost;
8115 if let (Some(psi_gradient), Some(psi_indices)) =
8116 (efs.psi_gradient.as_mut(), efs.psi_indices.as_ref())
8117 {
8118 if psi_gradient.len() != psi_indices.len() {
8119 crate::bail_invalid_estim!(
8120 "latent-coordinate analytic penalty EFS psi gradient length mismatch: gradient={}, indices={}",
8121 psi_gradient.len(),
8122 psi_indices.len()
8123 );
8124 }
8125 for (local_idx, &theta_idx) in psi_indices.iter().enumerate() {
8126 psi_gradient[local_idx] += contribution.gradient[theta_idx];
8127 }
8128 }
8129 }
8130 Ok(efs)
8131 }
8132
8133 fn eval_cost(&mut self, theta: &Array1<f64>) -> f64 {
8134 if let Some(cost) = self.cache.memoized_cost(theta) {
8135 return cost;
8136 }
8137 if self.cache.ensure_theta(theta).is_err() {
8138 return f64::INFINITY;
8139 }
8140 let design_revision = Some(self.cache.design_revision());
8141 let registry_for_key = self.cache.analytic_penalties();
8142 self.evaluator
8143 .set_analytic_penalty_registry(registry_for_key.as_deref());
8144 let result = {
8145 let design = self.cache.design();
8146 self.evaluator.evaluate_cost_only(
8147 &design.design,
8148 &design.penalties,
8149 &design.nullspace_dims,
8150 design.linear_constraints.clone(),
8151 theta,
8152 self.rho_dim,
8153 None,
8154 "latent-coordinate-joint cost-only",
8155 design_revision,
8156 )
8157 };
8158 match result {
8159 Ok(cost) => {
8160 let latent = match self.cache.latent() {
8161 Ok(latent) => latent,
8162 Err(_) => return f64::INFINITY,
8163 };
8164 let contribution = match latent_id_objective_contribution(
8165 theta,
8166 self.rho_dim,
8167 self.cache.analytic_penalty_rho_count(),
8168 latent.as_ref(),
8169 ) {
8170 Ok(contribution) => contribution,
8171 Err(_) => return f64::INFINITY,
8172 };
8173 let cost = cost + contribution.cost;
8174 let cost = if let Some(registry) = registry_for_key {
8175 let mut registry = registry.as_ref().clone();
8176 registry.apply_weight_schedules(
8177 gam_solve::estimate::reml::outer_eval::current_outer_iter()
8178 as usize,
8179 );
8180 match analytic_penalty_objective_contribution(
8181 theta,
8182 self.rho_dim,
8183 latent.as_ref(),
8184 ®istry,
8185 ) {
8186 Ok(contribution) => cost + contribution.cost,
8187 Err(_) => return f64::INFINITY,
8188 }
8189 } else {
8190 cost
8191 };
8192 self.cache.store_cost(cost);
8193 cost
8194 }
8195 Err(_) => f64::INFINITY,
8196 }
8197 }
8198 }
8199
8200 let mut ctx = LatentJointContext {
8201 rho_dim,
8202 cache: SingleBlockLatentCoordDesignCache::new(
8203 data.to_owned(),
8204 resolvedspec.clone(),
8205 best.design.clone(),
8206 latent,
8207 rho_dim,
8208 )
8209 .map_err(EstimationError::InvalidInput)?,
8210 evaluator: gam_solve::estimate::ExternalJointHyperEvaluator::new(
8211 y,
8212 weights,
8213 &best.design.design,
8214 offset,
8215 &best.design.penalties,
8216 &external_opts_for_design(&family, &best.design, options),
8217 "latent-coordinate-joint",
8218 )?,
8219 };
8220 let registry_for_key = ctx.cache.analytic_penalties();
8221 ctx.evaluator
8222 .set_analytic_penalty_registry(registry_for_key.as_deref());
8223 ctx.evaluator
8224 .set_persistent_latent_values_fingerprint(latent.values.id_mode());
8225 if let Some(cached_t) = ctx
8226 .evaluator
8227 .load_persistent_latent_values(latent.values.n_obs(), latent.values.latent_dim())
8228 {
8229 let cached_t: Array2<f64> = cached_t;
8230 for (dst, src) in theta0
8231 .slice_mut(s![rho_dim..rho_dim + latent_flat_dim])
8232 .iter_mut()
8233 .zip(cached_t.iter())
8234 {
8235 *dst = *src;
8236 }
8237 }
8238
8239 let problem = exact_joint_multistart_outer_problem(
8240 &theta0,
8241 &lower,
8242 &upper,
8243 rho_dim,
8244 latent_coord_ext_dim,
8245 theta0.len(),
8246 Derivative::Analytic,
8247 DeclaredHessianForm::Unavailable,
8248 false,
8249 false,
8250 seed_risk_profile_for_likelihood_family(&family),
8251 options.tol,
8252 options.max_iter.max(1),
8253 Some(5.0),
8254 Some(0.5),
8255 None,
8256 Some((data.nrows(), best.design.design.ncols().max(1))),
8259 !constant_curvature_term_indices(resolvedspec).is_empty(),
8262 );
8263
8264 let eval_outer = |ctx: &mut &mut LatentJointContext<'_>,
8265 theta: &Array1<f64>,
8266 order: OuterEvalOrder|
8267 -> Result<OuterEval, EstimationError> {
8268 let (cost, gradient, hessian) = ctx.eval_full(theta, order)?;
8269 Ok(OuterEval {
8270 cost,
8271 gradient,
8272 hessian,
8273 inner_beta_hint: None,
8274 })
8275 };
8276
8277 let result = {
8278 let mut obj = problem.build_objective_with_eval_order(
8279 &mut ctx,
8280 |ctx: &mut &mut LatentJointContext<'_>, theta: &Array1<f64>| Ok(ctx.eval_cost(theta)),
8281 |ctx: &mut &mut LatentJointContext<'_>, theta: &Array1<f64>| {
8282 eval_outer(ctx, theta, OuterEvalOrder::ValueAndGradient)
8283 },
8284 |ctx: &mut &mut LatentJointContext<'_>, theta: &Array1<f64>, order: OuterEvalOrder| {
8285 eval_outer(ctx, theta, order)
8286 },
8287 Some(|ctx: &mut &mut LatentJointContext<'_>| {
8288 ctx.cache.reset();
8289 }),
8290 Some(|ctx: &mut &mut LatentJointContext<'_>, theta: &Array1<f64>| ctx.eval_efs(theta)),
8291 );
8292
8293 problem
8294 .run(&mut obj, "latent-coordinate joint REML")
8295 .map_err(|e| {
8296 EstimationError::InvalidInput(format!(
8297 "latent-coordinate joint optimization failed after exhausting strategy fallbacks: {e}"
8298 ))
8299 })?
8300 };
8301 if !result.converged {
8302 crate::bail_invalid_estim!(
8303 "latent-coordinate joint optimization did not converge after {} iterations (final_objective={:.6e}, final_grad_norm={})",
8304 result.iterations,
8305 result.final_value,
8306 result.final_grad_norm_report(),
8307 );
8308 }
8309
8310 let theta_star = result.rho;
8311 let rho_star = theta_star.slice(s![..rho_dim]).mapv(f64::exp);
8312 let mut final_data = data.to_owned();
8313 let flat_t = theta_star
8314 .slice(s![rho_dim..rho_dim + latent_flat_dim])
8315 .to_owned();
8316 let mut fitted_latent_values =
8317 Array2::<f64>::zeros((latent.values.n_obs(), latent.values.latent_dim()));
8318 for n in 0..latent.values.n_obs() {
8319 for axis in 0..latent.values.latent_dim() {
8320 let value = flat_t[n * latent.values.latent_dim() + axis];
8321 fitted_latent_values[[n, axis]] = value;
8322 final_data[[n, latent.feature_cols[axis]]] = value;
8323 }
8324 }
8325 let optimized = fit_term_collection_forspecwith_heuristic_lambdas(
8326 final_data.view(),
8327 y,
8328 weights,
8329 offset,
8330 resolvedspec,
8331 rho_star.as_slice(),
8332 family,
8333 options,
8334 )?;
8335 ctx.evaluator
8336 .store_persistent_latent_values(&fitted_latent_values);
8337 let mut fit = optimized.fit;
8338 fit.reml_score = result.final_value;
8339 fit.penalized_objective = result.final_value;
8340 Ok(FittedTermCollectionWithSpec {
8341 fit,
8342 design: optimized.design,
8343 resolvedspec: resolvedspec.clone(),
8344 adaptive_diagnostics: optimized.adaptive_diagnostics,
8345 kappa_timing: None,
8346 })
8347}
8348
8349pub fn fit_term_collectionwith_latent_coord_optimization(
8350 data: ArrayView2<'_, f64>,
8351 y: Array1<f64>,
8352 weights: Array1<f64>,
8353 offset: Array1<f64>,
8354 spec: &TermCollectionSpec,
8355 latent: &StandardLatentCoordConfig,
8356 family: LikelihoodSpec,
8357 options: &FitOptions,
8358) -> Result<FittedTermCollectionWithSpec, EstimationError> {
8359 let n = data.nrows();
8360 if !(y.len() == n && weights.len() == n && offset.len() == n) {
8361 crate::bail_invalid_estim!(
8362 "fit_term_collectionwith_latent_coord_optimization row mismatch: n={}, y={}, weights={}, offset={}",
8363 n,
8364 y.len(),
8365 weights.len(),
8366 offset.len()
8367 );
8368 }
8369 let best = fit_term_collection_forspec(
8370 data,
8371 y.view(),
8372 weights.view(),
8373 offset.view(),
8374 spec,
8375 family.clone(),
8376 options,
8377 )?;
8378 let resolvedspec = freeze_term_collection_from_design(spec, &best.design)?;
8379 try_exact_joint_latent_coord_optimization(
8380 data,
8381 y.view(),
8382 weights.view(),
8383 offset.view(),
8384 &resolvedspec,
8385 &best,
8386 family,
8387 options,
8388 latent,
8389 )
8390}
8391
8392pub fn fit_term_collectionwith_spatial_length_scale_optimization(
8393 data: ArrayView2<'_, f64>,
8394 y: Array1<f64>,
8395 weights: Array1<f64>,
8396 offset: Array1<f64>,
8397 spec: &TermCollectionSpec,
8398 family: LikelihoodSpec,
8399 options: &FitOptions,
8400 kappa_options: &SpatialLengthScaleOptimizationOptions,
8401) -> Result<FittedTermCollectionWithSpec, EstimationError> {
8402 let mut resolvedspec = spec.clone();
8418 let spatial_terms = spatial_length_scale_term_indices(&resolvedspec);
8419 let n = data.nrows();
8420 if !(y.len() == n && weights.len() == n && offset.len() == n) {
8421 crate::bail_invalid_estim!(
8422 "fit_term_collectionwith_spatial_length_scale_optimization row mismatch: n={}, y={}, weights={}, offset={}",
8423 n,
8424 y.len(),
8425 weights.len(),
8426 offset.len()
8427 );
8428 }
8429 if !kappa_options.enabled || spatial_terms.is_empty() {
8430 let out = fit_term_collection_forspec(
8431 data,
8432 y.view(),
8433 weights.view(),
8434 offset.view(),
8435 &resolvedspec,
8436 family,
8437 options,
8438 )?;
8439 let resolvedspec = freeze_term_collection_from_design(&resolvedspec, &out.design)?;
8440 return Ok(FittedTermCollectionWithSpec {
8441 fit: out.fit,
8442 design: out.design,
8443 resolvedspec,
8444 adaptive_diagnostics: out.adaptive_diagnostics,
8445 kappa_timing: None,
8446 });
8447 }
8448 if kappa_options.max_outer_iter == 0 {
8449 crate::bail_invalid_estim!("spatial kappa optimization requires max_outer_iter >= 1");
8450 }
8451 if !(kappa_options.log_step.is_finite() && kappa_options.log_step > 0.0) {
8452 crate::bail_invalid_estim!("spatial kappa optimization requires log_step > 0");
8453 }
8454 if !(kappa_options.min_length_scale.is_finite()
8455 && kappa_options.max_length_scale.is_finite()
8456 && kappa_options.min_length_scale > 0.0
8457 && kappa_options.max_length_scale >= kappa_options.min_length_scale)
8458 {
8459 crate::bail_invalid_estim!(
8460 "spatial kappa optimization requires valid positive length_scale bounds"
8461 );
8462 }
8463
8464 let pilot_threshold = kappa_options.pilot_subsample_threshold;
8465 if pilot_threshold > 0 && n > pilot_threshold * 2 {
8466 log::info!(
8467 "[spatial-kappa] n={n} exceeds pilot threshold {}; using pilot geometry only for deterministic anisotropy initialization",
8468 pilot_threshold * 2,
8469 );
8470 apply_spatial_anisotropy_pilot_initializer(
8471 data,
8472 &mut resolvedspec,
8473 &spatial_terms,
8474 pilot_threshold,
8475 kappa_options,
8476 );
8477 }
8478
8479 apply_response_aware_anisotropy_seed(data, y.view(), &mut resolvedspec, &spatial_terms);
8488
8489 for term_idx in constant_curvature_term_indices(&resolvedspec) {
8507 if constant_curvature_kappa_is_fixed(&resolvedspec, term_idx) {
8512 continue;
8513 }
8514 if let Some(kappa_seed) =
8515 select_constant_curvature_kappa_sign_seed(data, y.view(), &resolvedspec, term_idx)
8516 && kappa_seed != 0.0
8517 && let Some(SmoothBasisSpec::ConstantCurvature { spec: cc, .. }) =
8518 resolvedspec.smooth_terms.get_mut(term_idx).map(|t| &mut t.basis)
8519 {
8520 log::info!(
8521 "[#1464] pinned CC term {term_idx} baseline κ to κ-fair scan value {kappa_seed} \
8522 (raw profiled REML is sign-blind; scan is authoritative for the sign)"
8523 );
8524 cc.kappa = kappa_seed;
8525 }
8526 }
8527
8528 let baseline_options = superseded_fit_options(options);
8529 let mut best = fit_term_collection_forspec(
8530 data,
8531 y.view(),
8532 weights.view(),
8533 offset.view(),
8534 &resolvedspec,
8535 family.clone(),
8536 &baseline_options,
8537 )?;
8538 resolvedspec = freeze_term_collection_from_design(&resolvedspec, &best.design)?;
8539 let mut spatial_terms = spatial_length_scale_term_indices(&resolvedspec);
8549 sync_aniso_contrasts_from_metadata(&mut resolvedspec, &best.design.smooth);
8553 let mut prescan_improved = false;
8560 if !spatial_terms.is_empty() {
8561 let baseline_score = fit_score(&best.fit);
8562 let range_overrides = prescan_isotropic_spatial_range_seed(
8563 data,
8564 y.view(),
8565 weights.view(),
8566 offset.view(),
8567 &resolvedspec,
8568 baseline_score,
8569 &family,
8570 &baseline_options,
8571 kappa_options,
8572 &spatial_terms,
8573 )?;
8574 if !range_overrides.is_empty() {
8575 prescan_improved = true;
8576 for (term_idx, length_scale) in range_overrides {
8577 set_spatial_length_scale(&mut resolvedspec, term_idx, length_scale)?;
8578 }
8579 best = fit_term_collection_forspec(
8583 data,
8584 y.view(),
8585 weights.view(),
8586 offset.view(),
8587 &resolvedspec,
8588 family.clone(),
8589 &baseline_options,
8590 )?;
8591 resolvedspec = freeze_term_collection_from_design(&resolvedspec, &best.design)?;
8592 spatial_terms = spatial_length_scale_term_indices(&resolvedspec);
8596 }
8597 }
8598 if spatial_terms.is_empty() {
8599 let fitted = fit_term_collection_forspecwith_heuristic_lambdas(
8600 data,
8601 y.view(),
8602 weights.view(),
8603 offset.view(),
8604 &resolvedspec,
8605 best.fit.lambdas.as_slice(),
8606 family,
8607 options,
8608 )?;
8609 return Ok(FittedTermCollectionWithSpec {
8610 fit: fitted.fit,
8611 design: fitted.design,
8612 resolvedspec,
8613 adaptive_diagnostics: fitted.adaptive_diagnostics,
8614 kappa_timing: None,
8615 });
8616 }
8617 let initial_score = fit_score(&best.fit);
8618 if !initial_score.is_finite() {
8619 log::debug!("[spatial-kappa] initial profiled score is non-finite");
8620 }
8621 let seed_length_scales: Vec<(usize, f64)> = spatial_terms
8628 .iter()
8629 .filter_map(|&t| get_spatial_length_scale(&resolvedspec, t).map(|ls| (t, ls)))
8630 .collect();
8631 let joint_result = try_exact_joint_spatial_length_scale_optimization(
8632 data,
8633 y.view(),
8634 weights.view(),
8635 offset.view(),
8636 &resolvedspec,
8637 &best,
8638 family.clone(),
8639 options,
8640 kappa_options,
8641 &spatial_terms,
8642 )
8643 .map(|opt| {
8644 opt.map(|fit| {
8645 let score = fit_score(&fit.fit);
8646 (fit, score)
8647 })
8648 });
8649 let exact_joint = if prescan_improved && !matches!(joint_result, Ok(Some(_))) {
8659 let reason = match &joint_result {
8660 Err(e) => format!("error: {e}"),
8661 _ => "unavailable".to_string(),
8662 };
8663 log::info!(
8664 "[spatial-kappa] #1074 joint polish yielded no usable candidate \
8665 ({reason}); returning the multi-start pre-scan geometry (REML {initial_score:.5})"
8666 );
8667 FittedTermCollectionWithSpec {
8668 fit: best.fit,
8669 design: best.design,
8670 resolvedspec,
8671 adaptive_diagnostics: best.adaptive_diagnostics,
8672 kappa_timing: None,
8673 }
8674 } else {
8675 require_successful_spatial_optimization_result(initial_score, joint_result)?
8676 };
8677
8678 let exact_joint = {
8705 let primary_score = fit_score(&exact_joint.fit);
8706 let improved = primary_score.is_finite()
8707 && initial_score.is_finite()
8708 && primary_score < initial_score - 1e-7 * initial_score.abs().max(1.0);
8709 let base_spec = exact_joint.resolvedspec.clone();
8714 let geometry_unchanged = !seed_length_scales.is_empty()
8717 && seed_length_scales.iter().all(|&(t, seed_ls)| {
8718 get_spatial_length_scale(&base_spec, t)
8719 .is_some_and(|ls| (ls - seed_ls).abs() <= 1e-6 * seed_ls.abs().max(1.0))
8720 });
8721 let eligible = !improved
8722 && geometry_unchanged
8723 && !has_aniso_terms(&base_spec, &spatial_terms)
8724 && constant_curvature_term_indices(&base_spec).is_empty()
8725 && spatial_terms
8726 .iter()
8727 .any(|&t| get_spatial_length_scale(&base_spec, t).is_some());
8728 if eligible {
8729 log::info!(
8730 "[spatial-kappa] #1688 joint solve stalled at REML {primary_score:.5} \
8731 (no improvement over baseline {initial_score:.5}); running ψ-window \
8732 multistart rescue across {} seeds",
8733 JOINT_RESTART_WINDOW_FRACTIONS.len(),
8734 );
8735 let mut best_fit = exact_joint;
8736 let mut best_score = primary_score;
8738 for &fraction in JOINT_RESTART_WINDOW_FRACTIONS.iter() {
8739 match joint_solve_from_window_fraction(
8740 data,
8741 y.view(),
8742 weights.view(),
8743 offset.view(),
8744 &base_spec,
8745 &spatial_terms,
8746 fraction,
8747 &family,
8748 options,
8749 &baseline_options,
8750 kappa_options,
8751 ) {
8752 Ok(Some((candidate, score))) => {
8753 if score.is_finite()
8754 && (!best_score.is_finite()
8755 || score < best_score - 1e-7 * best_score.abs().max(1.0))
8756 {
8757 log::info!(
8758 "[spatial-kappa] #1688 multistart seed (ψ-window {fraction:.2}) \
8759 reached REML {score:.5}, improving on {best_score:.5}",
8760 );
8761 best_score = score;
8762 best_fit = candidate;
8763 }
8764 }
8765 Ok(None) => {}
8767 Err(e) => {
8771 log::info!(
8772 "[spatial-kappa] #1688 multistart seed (ψ-window {fraction:.2}) \
8773 failed ({e}); skipping"
8774 );
8775 }
8776 }
8777 }
8778 best_fit
8779 } else {
8780 exact_joint
8781 }
8782 };
8783
8784 log_spatial_aniso_scales(&exact_joint.resolvedspec);
8785 Ok(exact_joint)
8786}
8787
8788#[derive(Clone, Debug)]
8794pub struct CurvatureInference {
8795 pub term_idx: usize,
8797 pub kappa_hat: f64,
8800 pub ci: gam_geometry::curvature_estimand::KappaProfileCi,
8802 pub flatness: gam_geometry::curvature_estimand::FlatnessTest,
8806}
8807
8808pub fn curvature_inference_forspec(
8826 data: ArrayView2<'_, f64>,
8827 y: ArrayView1<'_, f64>,
8828 weights: ArrayView1<'_, f64>,
8829 offset: ArrayView1<'_, f64>,
8830 resolvedspec: &TermCollectionSpec,
8831 term_idx: usize,
8832 family: LikelihoodSpec,
8833 options: &FitOptions,
8834 level: f64,
8835) -> Result<CurvatureInference, EstimationError> {
8836 let kappa_hat = get_constant_curvature_kappa(resolvedspec, term_idx).ok_or_else(|| {
8837 EstimationError::InvalidInput(format!(
8838 "curvature_inference_forspec: term {term_idx} is not a constant-curvature smooth"
8839 ))
8840 })?;
8841 let (kappa_min, kappa_max) = constant_curvature_kappa_bounds(data, resolvedspec, term_idx);
8842
8843 let cc_fair_inputs: Option<(Array2<f64>, gam_terms::basis::ConstantCurvatureBasisSpec)> =
8868 if kappa_hat < 0.0 {
8869 match resolvedspec.smooth_terms.get(term_idx).map(|t| &t.basis) {
8870 Some(SmoothBasisSpec::ConstantCurvature {
8871 feature_cols, spec, ..
8872 }) => select_columns(data, feature_cols)
8873 .ok()
8874 .map(|x| (x, spec.clone())),
8875 _ => None,
8876 }
8877 } else {
8878 None
8879 };
8880
8881 let v_p_cache: std::cell::RefCell<std::collections::HashMap<u64, f64>> =
8886 std::cell::RefCell::new(std::collections::HashMap::new());
8887 let v_p = |kappa: f64| -> Result<f64, String> {
8888 if !kappa.is_finite() {
8889 return Err(format!("V_p probed a non-finite κ = {kappa}"));
8890 }
8891 let key = kappa.to_bits();
8892 if let Some(&cached) = v_p_cache.borrow().get(&key) {
8893 return Ok(cached);
8894 }
8895 let score = if let Some((x_term, base_spec)) = &cc_fair_inputs {
8896 let mut probe_spec = base_spec.clone();
8897 probe_spec.kappa = kappa;
8898 gam_terms::basis::constant_curvature_kappa_fair_sign_score(x_term.view(), y, &probe_spec)
8899 .map_err(|e| format!("κ-fair criterion at κ={kappa} failed: {e}"))?
8900 } else {
8901 fixed_kappa_profiled_reml_score(
8902 data,
8903 y,
8904 weights,
8905 offset,
8906 resolvedspec,
8907 term_idx,
8908 kappa,
8909 family.clone(),
8910 options,
8911 )
8912 .map_err(|e| format!("V_p fixed-κ fit at κ={kappa} failed: {e}"))?
8913 };
8914 v_p_cache.borrow_mut().insert(key, score);
8915 Ok(score)
8916 };
8917
8918 let h = (1e-3 * (kappa_max - kappa_min)).max(1e-4);
8922 let v_pp = match (v_p(kappa_hat + h), v_p(kappa_hat), v_p(kappa_hat - h)) {
8923 (Ok(vp), Ok(v0), Ok(vm)) => (vp - 2.0 * v0 + vm) / (h * h),
8924 _ => f64::NAN, };
8926
8927 let ci = gam_geometry::curvature_estimand::profile_ci_walk(
8928 &v_p, kappa_hat, v_pp, kappa_min, kappa_max, level, 1e-4,
8929 )
8930 .map_err(EstimationError::InvalidInput)?;
8931 let flatness = gam_geometry::curvature_estimand::flatness_lr_test(&v_p, kappa_hat)
8932 .map_err(EstimationError::InvalidInput)?;
8933
8934 Ok(CurvatureInference {
8935 term_idx,
8936 kappa_hat,
8937 ci,
8938 flatness,
8939 })
8940}
8941
8942#[derive(Clone, Copy, Debug, PartialEq, Eq)]
8945pub enum SmoothLrCorrection {
8946 LawleyLrEstimatedLambda,
8950 LawleyLrFixedLambda,
8955 None,
8959}
8960
8961impl SmoothLrCorrection {
8962 pub fn label(self) -> &'static str {
8964 match self {
8965 SmoothLrCorrection::LawleyLrEstimatedLambda => "lawley_lr_estimated_lambda",
8966 SmoothLrCorrection::LawleyLrFixedLambda => "lawley_lr_fixed_lambda",
8967 SmoothLrCorrection::None => "none",
8968 }
8969 }
8970}
8971
8972#[derive(Clone, Debug)]
8978pub struct SmoothTermLrInference {
8979 pub name: String,
8981 pub term_idx: usize,
8983 pub statistic_lr: f64,
8986 pub ref_df: f64,
8989 pub bartlett_factor: f64,
8992 pub bartlett_factor_conditional: Option<f64>,
8996 pub rho_variation_shift: Option<f64>,
8999 pub statistic_corrected: f64,
9001 pub p_value_uncorrected: f64,
9003 pub p_value_corrected: f64,
9006 pub material: bool,
9014 pub correction: SmoothLrCorrection,
9016}
9017
9018pub const SMOOTH_LR_MATERIAL_THRESHOLD: f64 = 0.10;
9022
9023fn fitted_rho_penalty_components(
9029 penalties: &[BlockwisePenalty],
9030 lambdas: &[f64],
9031 p_total: usize,
9032) -> Result<Vec<gam_terms::inference::lawley::RhoPenaltyComponent>, EstimationError> {
9033 if penalties.len() != lambdas.len() {
9034 return Err(EstimationError::InvalidInput(format!(
9035 "smooth_term_lr_inference: penalty/lambda count mismatch ({} penalties, {} lambdas)",
9036 penalties.len(),
9037 lambdas.len()
9038 )));
9039 }
9040 let mut components = Vec::with_capacity(penalties.len());
9041 for (idx, (penalty, &lambda)) in penalties.iter().zip(lambdas.iter()).enumerate() {
9042 if !(lambda.is_finite() && lambda >= 0.0) {
9043 return Err(EstimationError::InvalidInput(format!(
9044 "smooth_term_lr_inference: lambda[{idx}] is invalid: {lambda}"
9045 )));
9046 }
9047 let r = &penalty.col_range;
9048 if r.end > p_total {
9049 return Err(EstimationError::InvalidInput(format!(
9050 "smooth_term_lr_inference: penalty[{idx}] range {:?} exceeds coefficient dimension {p_total}",
9051 r
9052 )));
9053 }
9054 let mut s_component = Array2::<f64>::zeros((p_total, p_total));
9055 s_component
9056 .slice_mut(s![r.start..r.end, r.start..r.end])
9057 .scaled_add(lambda, &penalty.local);
9058 components.push(gam_terms::inference::lawley::RhoPenaltyComponent { s_component });
9059 }
9060 Ok(components)
9061}
9062
9063pub fn smooth_term_lr_inference_forspec(
9108 data: ArrayView2<'_, f64>,
9109 y: ArrayView1<'_, f64>,
9110 weights: ArrayView1<'_, f64>,
9111 offset: ArrayView1<'_, f64>,
9112 resolvedspec: &TermCollectionSpec,
9113 family: LikelihoodSpec,
9114 options: &FitOptions,
9115) -> Result<Vec<SmoothTermLrInference>, EstimationError> {
9116 use gam_terms::inference::lawley::{
9117 LAWLEY_PAIR_MATRIX_MAX_ROWS, known_scale_expected_jets_with_dispersion,
9118 lawley_lr_bartlett_factor, lawley_lr_mean_shift_with_rho_variation,
9119 };
9120
9121 let n = data.nrows();
9122 let full = fit_term_collection_forspec(
9125 data,
9126 y,
9127 weights,
9128 offset,
9129 resolvedspec,
9130 family.clone(),
9131 options,
9132 )?;
9133 let ll_full = full.fit.log_likelihood;
9134 let p_total = full.design.design.ncols();
9135 let lambdas = full.fit.lambdas.as_slice().ok_or_else(|| {
9136 EstimationError::InvalidInput(
9137 "smooth_term_lr_inference: non-contiguous lambda vector".to_string(),
9138 )
9139 })?;
9140 let s_lambda = weighted_blockwise_penalty_sum(&full.design.penalties, lambdas, p_total);
9141 let rho_penalty_components =
9142 fitted_rho_penalty_components(&full.design.penalties, lambdas, p_total)?;
9143 let rho_covariance = full.fit.artifacts.rho_covariance.as_ref().filter(|cov| {
9144 cov.nrows() == rho_penalty_components.len() && cov.ncols() == rho_penalty_components.len()
9145 });
9146 let full_design_dense = full.design.design.to_dense();
9148 let influence = full.fit.coefficient_influence();
9149 let family_disp = lawley_dispersion_for_family(&family, &full.fit);
9150
9151 let mut penalty_cursor = full.design.leading_penalty_blocks_before_smooth();
9155 let mut out = Vec::<SmoothTermLrInference>::new();
9156 for (term_idx, design_term) in full.design.smooth.terms.iter().enumerate() {
9157 let k = design_term.penalties_local.len();
9158 let block_start = penalty_cursor;
9159 penalty_cursor += k;
9160 if design_term.shape != ShapeConstraint::None {
9163 continue;
9164 }
9165 let coeff_range = design_term.coeff_range.clone();
9166 if coeff_range.start >= coeff_range.end || coeff_range.end > p_total {
9167 continue;
9168 }
9169 let edf = full.fit.per_term_edf(coeff_range.clone(), block_start, k);
9181 let null_dim = design_term.wald_unpenalized_dim();
9201 let edf_floor = (null_dim.max(1)) as f64;
9253 let untrusted_edf_collapse = !full.fit.outer_converged && edf < edf_floor;
9254 let unconverged_dim_floor = if untrusted_edf_collapse {
9255 coeff_range.len() as f64
9256 } else {
9257 0.0
9258 };
9259 let rho_uncertainty_df = wps_block_uncertainty_df(
9260 full.fit.weighted_gram(),
9261 full.fit.smoothing_correction(),
9262 &coeff_range,
9263 family_disp,
9264 );
9265 let ref_df = (wood_reference_df(influence, &coeff_range)
9266 .unwrap_or(0.0)
9267 .max(edf)
9268 + rho_uncertainty_df)
9269 .max(null_dim as f64)
9270 .max(unconverged_dim_floor)
9271 .max(1.0);
9272 if !(ref_df.is_finite() && ref_df > 0.0) {
9273 continue;
9274 }
9275
9276 let mut null_spec = resolvedspec.clone();
9279 let Some(spec_pos) = null_spec
9280 .smooth_terms
9281 .iter()
9282 .position(|t| t.name == design_term.name)
9283 else {
9284 continue;
9285 };
9286 null_spec.smooth_terms.remove(spec_pos);
9287 let null_fit = fit_term_collection_forspec(
9288 data,
9289 y,
9290 weights,
9291 offset,
9292 &null_spec,
9293 family.clone(),
9294 options,
9295 );
9296 let (statistic_lr, eta_null) = match null_fit {
9297 Ok(null) if null.fit.log_likelihood.is_finite() => {
9298 let w = (2.0 * (ll_full - null.fit.log_likelihood)).max(0.0);
9299 let mut eta = null.design.design.dot(&null.fit.beta);
9303 eta += &offset;
9304 (w, Some(eta))
9305 }
9306 _ => (f64::NAN, None),
9307 };
9308
9309 let chi2 = statrs::distribution::ChiSquared::new(ref_df).ok();
9310 let p_uncorrected = match (chi2.as_ref(), statistic_lr.is_finite()) {
9311 (Some(dist), true) => {
9312 use statrs::distribution::ContinuousCDF;
9313 (1.0 - dist.cdf(statistic_lr)).clamp(0.0, 1.0)
9314 }
9315 _ => f64::NAN,
9316 };
9317
9318 let mut bartlett_factor = 1.0;
9322 let mut bartlett_factor_conditional = None;
9323 let mut rho_variation_shift = None;
9324 let mut statistic_corrected = statistic_lr;
9325 let mut p_corrected = p_uncorrected;
9326 let mut correction = SmoothLrCorrection::None;
9327 if let (Some(eta), true, true) = (
9328 eta_null.as_ref(),
9329 statistic_lr.is_finite(),
9330 n <= LAWLEY_PAIR_MATRIX_MAX_ROWS,
9331 ) {
9332 let kappas: Option<Vec<_>> = (0..n)
9333 .map(|i| {
9334 known_scale_expected_jets_with_dispersion(&family, eta[i], family_disp)
9335 .and_then(|jets| jets.kappas().ok())
9336 })
9337 .collect();
9338 if let (Some(kappas), Some(dist)) = (kappas, chi2.as_ref()) {
9339 let fixed_factor = lawley_lr_bartlett_factor(
9340 full_design_dense.view(),
9341 &kappas,
9342 Some(s_lambda.view()),
9343 coeff_range.clone(),
9344 ref_df,
9345 );
9346 if let Ok(c_cond) = fixed_factor
9347 && c_cond.is_finite()
9348 && c_cond > 0.0
9349 {
9350 let mut c_applied = c_cond;
9351 correction = SmoothLrCorrection::LawleyLrFixedLambda;
9352 if let Some(cov) = rho_covariance
9353 && let Ok(total_shift) = lawley_lr_mean_shift_with_rho_variation(
9354 full_design_dense.view(),
9355 &kappas,
9356 s_lambda.view(),
9357 coeff_range.clone(),
9358 &rho_penalty_components,
9359 cov.view(),
9360 )
9361 {
9362 let mean_w = ref_df + total_shift;
9363 if let Some(c_est) =
9364 gam_terms::inference::higher_order::bartlett_factor_from_mean(
9365 mean_w, ref_df,
9366 )
9367 && c_est.is_finite()
9368 && c_est > 0.0
9369 {
9370 let conditional_shift = (c_cond - 1.0) * ref_df;
9371 c_applied = c_est;
9372 bartlett_factor_conditional = Some(c_cond);
9373 rho_variation_shift = Some(total_shift - conditional_shift);
9374 correction = SmoothLrCorrection::LawleyLrEstimatedLambda;
9375 }
9376 }
9377 use statrs::distribution::ContinuousCDF;
9378 bartlett_factor = c_applied;
9379 statistic_corrected = statistic_lr / c_applied;
9380 p_corrected = (1.0 - dist.cdf(statistic_corrected)).clamp(0.0, 1.0);
9381 }
9382 }
9383 }
9384
9385 let material = match correction {
9391 SmoothLrCorrection::LawleyLrEstimatedLambda
9392 | SmoothLrCorrection::LawleyLrFixedLambda => {
9393 let factor_move = (bartlett_factor - 1.0).abs();
9394 let p_denom = p_uncorrected.max(p_corrected).max(f64::MIN_POSITIVE);
9395 let p_move = if p_uncorrected.is_finite() && p_corrected.is_finite() {
9396 (p_corrected - p_uncorrected).abs() / p_denom
9397 } else {
9398 0.0
9399 };
9400 factor_move > SMOOTH_LR_MATERIAL_THRESHOLD || p_move > SMOOTH_LR_MATERIAL_THRESHOLD
9401 }
9402 SmoothLrCorrection::None => false,
9403 };
9404
9405 out.push(SmoothTermLrInference {
9406 name: design_term.name.clone(),
9407 term_idx,
9408 statistic_lr,
9409 ref_df,
9410 bartlett_factor,
9411 bartlett_factor_conditional,
9412 rho_variation_shift,
9413 statistic_corrected,
9414 p_value_uncorrected: p_uncorrected,
9415 p_value_corrected: p_corrected,
9416 material,
9417 correction,
9418 });
9419 }
9420 Ok(out)
9421}
9422
9423fn lawley_dispersion_for_family(family: &LikelihoodSpec, fit: &UnifiedFitResult) -> f64 {
9426 match family.response {
9427 gam_spec::ResponseFamily::Gaussian => {
9428 let sd = fit.standard_deviation;
9429 (sd * sd).max(f64::MIN_POSITIVE)
9430 }
9431 gam_spec::ResponseFamily::Gamma => {
9432 let shape = fit.standard_deviation;
9433 if shape.is_finite() && shape > 0.0 {
9434 1.0 / shape
9435 } else {
9436 1.0
9437 }
9438 }
9439 _ => 1.0,
9440 }
9441}
9442
9443fn wps_block_uncertainty_df(
9444 weighted_gram: Option<&Array2<f64>>,
9445 smoothing_correction: Option<&Array2<f64>>,
9446 coeff_range: &Range<usize>,
9447 phi: f64,
9448) -> f64 {
9449 let (Some(xwx), Some(corr)) = (weighted_gram, smoothing_correction) else {
9450 return 0.0;
9451 };
9452 let (start, end) = (coeff_range.start, coeff_range.end);
9453 if start >= end
9454 || end > xwx.nrows()
9455 || end > xwx.ncols()
9456 || end > corr.nrows()
9457 || end > corr.ncols()
9458 || !(phi.is_finite() && phi > 0.0)
9459 {
9460 return 0.0;
9461 }
9462
9463 let mut trace = 0.0;
9464 for i in start..end {
9465 for j in start..end {
9466 trace += xwx[[i, j]] * corr[[j, i]];
9467 }
9468 }
9469 trace /= phi;
9470 if trace.is_finite() && trace > 0.0 {
9471 trace
9472 } else {
9473 0.0
9474 }
9475}
9476
9477fn wood_reference_df(influence: Option<&Array2<f64>>, coeff_range: &Range<usize>) -> Option<f64> {
9501 let f = influence?;
9502 let (start, end) = (coeff_range.start, coeff_range.end);
9503 if start >= end || end > f.nrows() || end > f.ncols() {
9504 return None;
9505 }
9506 let block = f.slice(s![start..end, start..end]);
9507 let tr = (0..block.nrows()).map(|i| block[[i, i]]).sum::<f64>();
9508 let tr2 = block.dot(&block).diag().sum();
9509 (tr.is_finite() && tr2.is_finite() && tr > 0.0)
9510 .then(|| (2.0 * tr - tr2).max(tr).max(1e-12))
9511}