1use super::*;
2
3#[derive(Clone)]
4pub(crate) struct TransformationExactGeometryCache {
5 pub(crate) key: Vec<u64>,
6 pub(crate) covariate_spec_resolved: TermCollectionSpec,
7 pub(crate) covariate_design: TermCollectionDesign,
8 pub(crate) family: TransformationNormalFamily,
9 pub(crate) blocks: Vec<ParameterBlockSpec>,
10 pub(crate) hyper_layout: SharedCustomFamilyHyperLayout,
11}
12
13impl TransformationExactGeometryCache {
14 pub(crate) fn update_block_log_lambdas(
15 &mut self,
16 log_lambdas: &Array1<f64>,
17 ) -> Result<(), String> {
18 let spec = self
19 .blocks
20 .first_mut()
21 .ok_or_else(|| "missing transformation block spec".to_string())?;
22 if log_lambdas.len() != spec.initial_log_lambdas.len() {
23 return Err(TransformationNormalError::InvalidInput {
24 reason: format!(
25 "transformation final fit rho length mismatch: got {}, expected {}",
26 log_lambdas.len(),
27 spec.initial_log_lambdas.len()
28 ),
29 }
30 .into());
31 }
32 gam_problem::validate_log_strengths(log_lambdas.iter().copied()).map_err(|error| {
33 TransformationNormalError::InvalidInput {
34 reason: format!("invalid transformation smoothing strength: {error}"),
35 }
36 })?;
37 spec.initial_log_lambdas = log_lambdas.clone();
38 Ok(())
39 }
40}
41
42pub(crate) fn transformation_spatial_geometry_key(
43 spec: &TermCollectionSpec,
44 spatial_terms: &[usize],
45) -> Result<Vec<u64>, String> {
46 let mut key = Vec::new();
47 key.push(spatial_terms.len() as u64);
48 for &term_idx in spatial_terms {
49 let term = spec.smooth_terms.get(term_idx).ok_or_else(|| {
50 format!(
51 "transformation spatial geometry key term index {term_idx} out of range for {} smooth terms",
52 spec.smooth_terms.len()
53 )
54 })?;
55 key.push(term_idx as u64);
56
57 let payload = serde_json::to_vec(term).map_err(|err| {
65 format!("failed to serialize transformation spatial geometry term {term_idx}: {err}")
66 })?;
67 key.push(payload.len() as u64);
68 for chunk in payload.chunks(8) {
69 let mut bytes = [0u8; 8];
70 for (dst, src) in bytes.iter_mut().zip(chunk.iter().copied()) {
71 *dst = src;
72 }
73 key.push(u64::from_le_bytes(bytes));
74 }
75 }
76 Ok(key)
77}
78
79#[derive(Clone)]
85pub struct TransformationNormalFitResult {
86 pub family: TransformationNormalFamily,
87 pub fit: UnifiedFitResult,
88 pub covariate_spec_resolved: TermCollectionSpec,
89 pub covariate_design: TermCollectionDesign,
90 pub score_calibration: TransformationScoreCalibration,
91}
92
93pub fn fit_transformation_normal(
101 response: &Array1<f64>,
102 weights: &Array1<f64>,
103 offset: &Array1<f64>,
104 covariate_data: ArrayView2<'_, f64>,
105 covariate_spec: &TermCollectionSpec,
106 config: &TransformationNormalConfig,
107 options: &BlockwiseFitOptions,
108 kappa_options: &SpatialLengthScaleOptimizationOptions,
109 warm_start: Option<&TransformationWarmStart>,
110) -> Result<TransformationNormalFitResult, String> {
111 let mut options = options.clone();
112 let covariate_spec = covariate_spec.clone();
117
118 let boot_design = build_term_collection_design(covariate_data, &covariate_spec)
121 .map_err(|e| format!("failed to build bootstrap covariate design: {e}"))?;
122 let boot_spec = freeze_term_collection_from_design(&covariate_spec, &boot_design)
123 .map_err(|e| format!("failed to freeze bootstrap covariate spatial basis centers: {e}"))?;
124 let mut effective_config = config.clone();
125 if !config.response_num_internal_knots_pinned {
131 effective_config.response_num_internal_knots = effective_response_num_internal_knots(
132 config,
133 response.len(),
134 boot_design.design.ncols(),
135 response.view(),
136 );
137 }
138
139 let (resp_val, resp_deriv, resp_penalties, resp_knots, resp_transform) =
142 build_response_basis(response, &effective_config)?;
143
144 let realized_p_total = resp_val.ncols().saturating_mul(boot_design.design.ncols());
152 let ctn_inner_cap = CTN_INNER_MAX_CYCLES_BASE
153 .saturating_add(realized_p_total.saturating_mul(CTN_INNER_MAX_CYCLES_PER_DIM))
154 .min(CTN_INNER_MAX_CYCLES_CEILING);
155 options.inner_max_cycles = options.inner_max_cycles.min(ctn_inner_cap);
156
157 let spatial_terms = spatial_length_scale_term_indices(&covariate_spec);
159
160 if spatial_terms.is_empty() || !kappa_options.enabled {
161 let cov_design = boot_design;
165 let cov_spec_resolved = boot_spec;
166 let effective_offset = cov_design
167 .compose_offset(offset.view(), "transformation-normal fit")
168 .map_err(|error| error.to_string())?;
169
170 let family = TransformationNormalFamily::from_prebuilt_response_basis(
171 response,
172 resp_val,
173 resp_deriv,
174 resp_penalties,
175 resp_knots.clone(),
176 effective_config.response_degree,
177 resp_transform,
178 weights,
179 &effective_offset,
180 cov_design.design.clone(),
181 cov_design
182 .penalties
183 .iter()
184 .map(|bp| bp.to_penalty_matrix(cov_design.design.ncols()))
185 .collect(),
186 &effective_config,
187 warm_start,
188 )?;
189 let rho0 = family.penalty_scale_log_lambdas()?;
190 let blocks = vec![family.block_spec(&rho0)?];
191 let fit = fit_custom_family(&family, &blocks, &options)
192 .map_err(|e| format!("transformation fit failed: {e}"))?;
193 let (fit, score_calibration) = calibrate_transformation_scores(&family, fit)?;
194
195 return Ok(TransformationNormalFitResult {
196 family,
197 fit,
198 covariate_spec_resolved: cov_spec_resolved,
199 covariate_design: cov_design,
200 score_calibration,
201 });
202 }
203
204 let kappa0 = SpatialLogKappaCoords::from_length_scales_aniso(
209 &covariate_spec,
210 &spatial_terms,
211 kappa_options,
212 )
213 .reseed_from_data(
214 covariate_data,
215 &covariate_spec,
216 &spatial_terms,
217 kappa_options,
218 )
219 .map_err(|error| error.to_string())?;
220 let kappa_dims = kappa0.dims_per_term().to_vec();
221 let kappa_lower = SpatialLogKappaCoords::lower_bounds_aniso_from_data(
222 covariate_data,
223 &covariate_spec,
224 &spatial_terms,
225 &kappa_dims,
226 kappa_options,
227 )
228 .map_err(|error| error.to_string())?;
229 let kappa_upper = SpatialLogKappaCoords::upper_bounds_aniso_from_data(
230 covariate_data,
231 &covariate_spec,
232 &spatial_terms,
233 &kappa_dims,
234 kappa_options,
235 )
236 .map_err(|error| error.to_string())?;
237 let kappa0 = kappa0.clamp_to_bounds(&kappa_lower, &kappa_upper);
239
240 let analytic_psi_available =
242 build_block_spatial_psi_derivatives(covariate_data, &boot_spec, &boot_design)?.is_some();
243
244 let probe_design = build_term_collection_design(covariate_data, &boot_spec)
258 .map_err(|e| format!("failed to rebuild frozen probe covariate design: {e}"))?;
259 let probe_offset = probe_design
260 .compose_offset(offset.view(), "transformation-normal spatial probe")
261 .map_err(|error| error.to_string())?;
262
263 let probe_family = TransformationNormalFamily::from_prebuilt_response_basis(
265 response,
266 resp_val.clone(),
267 resp_deriv.clone(),
268 resp_penalties.clone(),
269 resp_knots.clone(),
270 effective_config.response_degree,
271 resp_transform.clone(),
272 weights,
273 &probe_offset,
274 probe_design.design.clone(),
275 probe_design
276 .penalties
277 .iter()
278 .map(|bp| bp.to_penalty_matrix(probe_design.design.ncols()))
279 .collect(),
280 &effective_config,
281 warm_start,
282 )?;
283 let rho0 = probe_family.penalty_scale_log_lambdas()?;
284 let probe_block = probe_family.block_spec(&rho0)?;
285 let n_penalties = probe_block.initial_log_lambdas.len();
286 log::info!(
287 "[transformation-normal] exact joint setup: rho_dim={} log_kappa_dim={} dims_per_term={:?}",
288 n_penalties,
289 kappa0.len(),
290 kappa_dims,
291 );
292 let rho_floor = -12.0;
293 let rho_lower = Array1::<f64>::from_elem(n_penalties, rho_floor);
294 let rho_upper = Array1::<f64>::from_elem(n_penalties, 12.0);
295 let probe_blocks = vec![probe_block.clone()];
296 let (_, cap_hessian) = crate::custom_family::custom_family_outer_derivatives(
297 &probe_family,
298 &probe_blocks,
299 &options,
300 );
301 let analytic_gradient = analytic_psi_available;
302 let analytic_hessian_supported = analytic_gradient && cap_hessian.is_analytic();
303 let analytic_hessian = analytic_hessian_supported;
304 if analytic_hessian {
305 log::info!(
306 "[transformation-normal] CTN exact joint analytic outer Hessian is available for spatial kappa optimization; using exact second-order outer geometry"
307 );
308 }
309
310 let (rho0_min, rho0_max) = if rho0.is_empty() {
311 (0.0, 0.0)
312 } else {
313 (
314 rho0.iter().copied().fold(f64::INFINITY, f64::min),
315 rho0.iter().copied().fold(f64::NEG_INFINITY, f64::max),
316 )
317 };
318 log::info!(
319 "[transformation-normal] skipping baseline custom-family prefit before exact joint optimization \
320 (rho_dim={}, log_kappa_dim={}, rho0_range=[{:.3}, {:.3}]); using CTN warm start and penalty-scale rho seed",
321 n_penalties,
322 kappa0.len(),
323 rho0_min,
324 rho0_max,
325 );
326
327 if !analytic_psi_available {
328 return Err(
329 "transformation-normal spatial length-scale optimization requires analytic spatial psi derivatives"
330 .to_string(),
331 );
332 }
333
334 let exact_mode_branch: RefCell<ExactCoefficientModeBranch> =
340 RefCell::new(ExactCoefficientModeBranch::default());
341
342 let joint_setup =
343 ExactJointHyperSetup::new(rho0, rho_lower, rho_upper, kappa0, kappa_lower, kappa_upper);
344
345 let rv = resp_val.clone();
347 let rd = resp_deriv.clone();
348 let rp = resp_penalties.clone();
349 let rk = resp_knots.clone();
350 let rt = resp_transform.clone();
351 let rdeg = effective_config.response_degree;
352 let cfg = effective_config.clone();
353 let ws = warm_start.cloned();
354
355 let make_family =
357 |cov_design: &TermCollectionDesign| -> Result<TransformationNormalFamily, String> {
358 let effective_offset = cov_design
359 .compose_offset(offset.view(), "transformation-normal spatial fit")
360 .map_err(|error| error.to_string())?;
361 TransformationNormalFamily::from_prebuilt_response_basis(
362 response,
363 rv.clone(),
364 rd.clone(),
365 rp.clone(),
366 rk.clone(),
367 rdeg,
368 rt.clone(),
369 weights,
370 &effective_offset,
371 cov_design.design.clone(),
372 cov_design
373 .penalties
374 .iter()
375 .map(|bp| bp.to_penalty_matrix(cov_design.design.ncols()))
376 .collect(),
377 &cfg,
378 ws.as_ref(),
379 )
380 };
381
382 let block_specs_slice = [boot_spec.clone()];
383 let block_term_indices_slice = [spatial_terms.clone()];
384 let exact_geometry_cache: RefCell<Option<TransformationExactGeometryCache>> =
385 RefCell::new(None);
386 let spatial_terms_for_cache = spatial_terms.clone();
387
388 let ensure_exact_geometry = |spec: &TermCollectionSpec,
389 design: &TermCollectionDesign,
390 rho: &Array1<f64>,
391 hyper_values: &Array1<f64>|
392 -> Result<(), String> {
393 let effective_spec = freeze_term_collection_from_design(spec, design)
394 .map_err(|e| format!("failed to freeze transformation geometry key: {e}"))?;
395 let key = transformation_spatial_geometry_key(&effective_spec, &spatial_terms_for_cache)?;
396 let needs_rebuild = exact_geometry_cache
397 .borrow()
398 .as_ref()
399 .map(|cached| cached.key != key)
400 .unwrap_or(true);
401 if !needs_rebuild {
402 let mut cache = exact_geometry_cache.borrow_mut();
403 let cached = cache
404 .as_mut()
405 .ok_or_else(|| "missing transformation exact geometry cache".to_string())?;
406 if cached.hyper_layout.values().len() != hyper_values.len()
407 || cached
408 .hyper_layout
409 .values()
410 .iter()
411 .zip(hyper_values)
412 .any(|(cached, current)| cached.to_bits() != current.to_bits())
413 {
414 return Err(
415 "transformation exact geometry key reused across distinct hypercoordinate values"
416 .to_string(),
417 );
418 }
419 return cached.update_block_log_lambdas(rho);
420 }
421
422 let geom_start = std::time::Instant::now();
423 let exact_design = build_term_collection_design(covariate_data, &effective_spec)
424 .map_err(|e| format!("failed to rebuild frozen transformation geometry: {e}"))?;
425 let family = make_family(&exact_design)?;
426 let cov_psi_derivs =
427 build_block_spatial_psi_derivatives(covariate_data, &effective_spec, &exact_design)?
428 .ok_or_else(|| {
429 "missing covariate spatial psi derivatives for transformation model".to_string()
430 })?;
431 let tensor_derivs = build_tensor_psi_derivatives(&family, &cov_psi_derivs)?;
432
433 log::debug!(
434 "[transformation-normal] rebuilt exact geometry cache for {} spatial terms in {:.3}s",
435 spatial_terms_for_cache.len(),
436 geom_start.elapsed().as_secs_f64(),
437 );
438
439 exact_geometry_cache.replace(Some(TransformationExactGeometryCache {
440 key,
441 covariate_spec_resolved: effective_spec,
442 covariate_design: exact_design,
443 blocks: vec![family.block_spec(rho)?],
444 family,
445 hyper_layout: Arc::new(CustomFamilyHyperLayout::new(
446 vec![tensor_derivs],
447 Vec::new(),
448 hyper_values.clone(),
449 )?),
450 }));
451 Ok(())
452 };
453
454 let exact_mode_candidates = |eval_mode: gam_problem::EvalMode,
455 rho: &Array1<f64>|
456 -> Vec<Option<CustomFamilyWarmStart>> {
457 let (froze, candidates) = exact_mode_branch.borrow_mut().candidates(eval_mode, rho);
458 if froze {
459 log::info!(
460 "[transformation-normal] froze deterministic exact coefficient-mode branch at the first derivative-bearing outer seed evaluation"
461 );
462 }
463 candidates
464 };
465
466 log::info!(
467 "[transformation-normal] entering exact joint outer optimization \
468 (analytic_gradient={}, analytic_hessian={})",
469 analytic_gradient,
470 analytic_hessian,
471 );
472 let outer_derivative_policy =
477 probe_family.outer_derivative_policy(&probe_blocks, joint_setup.log_kappa_dim(), &options);
478 let solved = optimize_spatial_length_scale_exact_joint(
479 covariate_data,
480 &block_specs_slice,
481 &block_term_indices_slice,
482 kappa_options,
483 &joint_setup,
484 gam_solve::seeding::SeedRiskProfile::Gaussian,
485 analytic_gradient,
486 analytic_hessian,
487 true,
494 None,
495 outer_derivative_policy,
496 |theta,
498 specs: &[TermCollectionSpec],
499 designs: &[TermCollectionDesign],
500 provenance: SpatialFitProvenance<'_, CustomFamilyJointHyperModeSelection>| {
501 let rho = theta.slice(s![..joint_setup.rho_dim()]).to_owned();
502 let hyper_values = theta.slice(s![joint_setup.rho_dim()..]).to_owned();
503 ensure_exact_geometry(&specs[0], &designs[0], &rho, &hyper_values)?;
504 let mut cache_ref = exact_geometry_cache.borrow_mut();
505 let geometry = cache_ref
506 .as_mut()
507 .ok_or_else(|| "missing transformation exact geometry cache".to_string())?;
508 let final_options = crate::outer_subsample::exact_outer_options_for_row_set(
509 &options,
510 &gam_problem::outer_subsample::RowSet::All,
511 );
512 let fit = match provenance {
513 SpatialFitProvenance::NoOuterOptimization => {
514 let warm_starts =
515 exact_mode_candidates(gam_problem::EvalMode::ValueOnly, &rho);
516 let selection = evaluate_custom_family_joint_hyper_best_mode_shared(
517 &geometry.family,
518 &geometry.blocks,
519 &final_options,
520 &rho,
521 Arc::clone(&geometry.hyper_layout),
522 &warm_starts,
523 gam_problem::EvalMode::ValueOnly,
524 )
525 .map_err(|e| format!("transformation fixed mode profile: {e}"))?;
526 log::info!(
527 "[transformation-normal] user-fixed coefficient mode selected candidate={} objective={:.16e}",
528 selection.selected_candidate,
529 selection.result.objective,
530 );
531 fit_custom_family_user_fixed_log_lambdas_from_mode_selection(
532 &geometry.family,
533 &geometry.blocks,
534 &final_options,
535 selection,
536 )
537 }
538 SpatialFitProvenance::Certified { outer, mode: selection } => {
539 log::info!(
540 "[transformation-normal] consuming certified terminal coefficient mode candidate={} objective={:.16e} without profile replay",
541 selection.selected_candidate,
542 selection.result.objective,
543 );
544 fit_custom_family_fixed_log_lambdas_from_mode_selection(
545 &geometry.family,
546 &geometry.blocks,
547 &final_options,
548 selection,
549 theta,
550 outer,
551 )
552 }
553 }
554 .map_err(|e| format!("transformation fit_fn: {e}"))?;
555 if let Some(block) = fit.block_states.first() {
556 *geometry
557 .family
558 .row_quantity_cache
559 .lock()
560 .expect("CTN row quantity cache mutex poisoned") = None;
561 let final_rows = geometry.family.row_quantities(&block.beta)?;
562 let max_abs_h = final_rows
563 .h
564 .iter()
565 .copied()
566 .map(f64::abs)
567 .fold(0.0, f64::max);
568 let cov_chunk = geometry
569 .family
570 .covariate_design
571 .try_row_chunk(0..response.len())
572 .map_err(|err| {
573 format!("final CTN covariate design validation failed: {err}")
574 })?;
575 let max_abs_cov = cov_chunk.iter().copied().map(f64::abs).fold(0.0, f64::max);
576 log::info!(
577 "[transformation-normal] final fixed-rho CTN validation: max_abs_h={:.6e} max_abs_covariate_basis={:.6e}",
578 max_abs_h,
579 max_abs_cov
580 );
581 }
582 Ok(TransformationNormalFitResult {
583 family: geometry.family.clone(),
584 fit,
585 covariate_spec_resolved: geometry.covariate_spec_resolved.clone(),
586 covariate_design: geometry.covariate_design.clone(),
587 score_calibration: TransformationScoreCalibration::finite_support_pit(),
588 })
589 },
590 |theta,
592 specs: &[TermCollectionSpec],
593 designs: &[TermCollectionDesign],
594 eval_mode,
595 row_set| {
596 let rho = theta.slice(s![..joint_setup.rho_dim()]).to_owned();
597 let hyper_values = theta.slice(s![joint_setup.rho_dim()..]).to_owned();
598 ensure_exact_geometry(&specs[0], &designs[0], &rho, &hyper_values)?;
599 let mut cache_ref = exact_geometry_cache.borrow_mut();
600 let geometry = cache_ref
601 .as_mut()
602 .ok_or_else(|| "missing transformation exact geometry cache".to_string())?;
603 let warm_starts = exact_mode_candidates(eval_mode, &rho);
604 let competing_modes = warm_starts.len() > 1;
605 let eval_options =
609 crate::outer_subsample::exact_outer_options_for_row_set(&options, row_set);
610 let selection = evaluate_custom_family_joint_hyper_best_mode_shared(
611 &geometry.family,
612 &geometry.blocks,
613 &eval_options,
614 &rho,
615 Arc::clone(&geometry.hyper_layout),
616 &warm_starts,
617 eval_mode,
618 )
619 .map_err(|e| format!("transformation exact joint mode profile: {e}"))?;
620 for (candidate_idx, rejection) in selection.rejected_candidates.iter().enumerate() {
621 if let Some(rejection) = rejection {
622 log::warn!(
623 "[transformation-normal] rejected exact coefficient-mode candidate mode_candidate={candidate_idx}: {rejection}"
624 );
625 }
626 }
627 if competing_modes {
628 for (candidate_idx, objective) in selection.screened_objectives.iter().enumerate() {
629 if let Some(objective) = objective {
630 let source = if candidate_idx == 0 {
631 "cold"
632 } else {
633 "carried"
634 };
635 log::info!(
636 "[transformation-normal] exact coefficient-mode screen source={} objective={:.16e}",
637 source,
638 objective,
639 );
640 }
641 }
642 let selected_source = if selection.selected_candidate == 0 {
643 "cold"
644 } else {
645 "carried"
646 };
647 log::info!(
648 "[transformation-normal] selected exact coefficient mode source={} objective={:.16e}",
649 selected_source,
650 selection.result.objective,
651 );
652 }
653 let objective = selection.result.objective;
654 let gradient = selection.result.gradient.clone();
655 let outer_hessian = selection.result.outer_hessian.clone();
656 exact_mode_branch
657 .borrow_mut()
658 .record_value(eval_mode, selection.result.warm_start.clone());
659
660 Ok(ExactJointEvaluation {
661 objective,
662 gradient,
663 hessian: outer_hessian,
664 mode: selection,
665 })
666 },
667 |_theta,
668 _specs: &[TermCollectionSpec],
669 _designs: &[TermCollectionDesign],
670 _row_set| {
671 Err::<ExactJointEfsEvaluation<crate::custom_family::CustomFamilyJointHyperModeSelection>, String>("transformation-normal EFS callback invoked even though fixed-point optimization is disabled for beta-dependent exact curvature".to_string())
672 },
673 |_beta: &Array1<f64>| Ok(gam_solve::rho_optimizer::SeedOutcome::NoSlot),
674 )?;
675
676 let mut fit = solved.fit;
677 let (calibrated_fit, score_calibration) =
678 calibrate_transformation_scores(&fit.family, fit.fit.clone())?;
679 fit.fit = calibrated_fit;
680 fit.score_calibration = score_calibration;
681 Ok(fit)
682}