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 mut covariate_spec = covariate_spec.clone();
117 let seeded = crate::fit_orchestration::drivers::seed_measure_jet_auto_ranges(
138 covariate_data,
139 response.view(),
140 weights.view(),
141 &mut covariate_spec,
142 );
143 if seeded > 0 {
144 log::info!(
145 "[#2750] screened the representer range of {seeded} auto measure-jet term(s) against \
146 the response before the transformation-normal design build"
147 );
148 }
149 let covariate_spec = covariate_spec;
150
151 let boot_design = build_term_collection_design(covariate_data, &covariate_spec)
154 .map_err(|e| format!("failed to build bootstrap covariate design: {e}"))?;
155 let boot_spec = freeze_term_collection_from_design(&covariate_spec, &boot_design)
156 .map_err(|e| format!("failed to freeze bootstrap covariate spatial basis centers: {e}"))?;
157 let mut effective_config = config.clone();
158 if !config.response_num_internal_knots_pinned {
164 effective_config.response_num_internal_knots = effective_response_num_internal_knots(
165 config,
166 response.len(),
167 boot_design.design.ncols(),
168 response.view(),
169 );
170 }
171
172 let (resp_val, resp_deriv, resp_penalties, resp_knots, resp_transform) =
175 build_response_basis(response, &effective_config)?;
176
177 let realized_p_total = resp_val.ncols().saturating_mul(boot_design.design.ncols());
185 let ctn_inner_cap = CTN_INNER_MAX_CYCLES_BASE
186 .saturating_add(realized_p_total.saturating_mul(CTN_INNER_MAX_CYCLES_PER_DIM))
187 .min(CTN_INNER_MAX_CYCLES_CEILING);
188 options.inner_max_cycles = options.inner_max_cycles.min(ctn_inner_cap);
189
190 let spatial_terms = spatial_length_scale_term_indices(&covariate_spec);
192
193 if spatial_terms.is_empty() || !kappa_options.enabled {
194 let cov_design = boot_design;
198 let cov_spec_resolved = boot_spec;
199 let effective_offset = cov_design
200 .compose_offset(offset.view(), "transformation-normal fit")
201 .map_err(|error| error.to_string())?;
202
203 let family = TransformationNormalFamily::from_prebuilt_response_basis(
204 response,
205 resp_val,
206 resp_deriv,
207 resp_penalties,
208 resp_knots.clone(),
209 effective_config.response_degree,
210 resp_transform,
211 weights,
212 &effective_offset,
213 cov_design.design.clone(),
214 cov_design
215 .penalties
216 .iter()
217 .map(|bp| bp.to_penalty_matrix(cov_design.design.ncols()))
218 .collect(),
219 &effective_config,
220 warm_start,
221 )?;
222 let rho0 = family.penalty_scale_log_lambdas()?;
223 let blocks = vec![family.block_spec(&rho0)?];
224 let fit = fit_custom_family(&family, &blocks, &options)
225 .map_err(|e| format!("transformation fit failed: {e}"))?;
226 let (fit, score_calibration) = calibrate_transformation_scores(&family, fit)?;
227
228 return Ok(TransformationNormalFitResult {
229 family,
230 fit,
231 covariate_spec_resolved: cov_spec_resolved,
232 covariate_design: cov_design,
233 score_calibration,
234 });
235 }
236
237 let kappa0 = SpatialLogKappaCoords::from_length_scales_aniso(
242 &covariate_spec,
243 &spatial_terms,
244 kappa_options,
245 )
246 .reseed_from_data(
247 covariate_data,
248 &covariate_spec,
249 &spatial_terms,
250 kappa_options,
251 )
252 .map_err(|error| error.to_string())?;
253 let kappa_dims = kappa0.dims_per_term().to_vec();
254 let kappa_lower = SpatialLogKappaCoords::lower_bounds_aniso_from_data(
255 covariate_data,
256 &covariate_spec,
257 &spatial_terms,
258 &kappa_dims,
259 kappa_options,
260 )
261 .map_err(|error| error.to_string())?;
262 let kappa_upper = SpatialLogKappaCoords::upper_bounds_aniso_from_data(
263 covariate_data,
264 &covariate_spec,
265 &spatial_terms,
266 &kappa_dims,
267 kappa_options,
268 )
269 .map_err(|error| error.to_string())?;
270 let kappa0 = kappa0.clamp_to_bounds(&kappa_lower, &kappa_upper);
272
273 let analytic_psi_available =
275 build_block_spatial_psi_derivatives(covariate_data, &boot_spec, &boot_design)?.is_some();
276
277 let probe_design = build_term_collection_design(covariate_data, &boot_spec)
291 .map_err(|e| format!("failed to rebuild frozen probe covariate design: {e}"))?;
292 let probe_offset = probe_design
293 .compose_offset(offset.view(), "transformation-normal spatial probe")
294 .map_err(|error| error.to_string())?;
295
296 let probe_family = TransformationNormalFamily::from_prebuilt_response_basis(
298 response,
299 resp_val.clone(),
300 resp_deriv.clone(),
301 resp_penalties.clone(),
302 resp_knots.clone(),
303 effective_config.response_degree,
304 resp_transform.clone(),
305 weights,
306 &probe_offset,
307 probe_design.design.clone(),
308 probe_design
309 .penalties
310 .iter()
311 .map(|bp| bp.to_penalty_matrix(probe_design.design.ncols()))
312 .collect(),
313 &effective_config,
314 warm_start,
315 )?;
316 let rho0 = probe_family.penalty_scale_log_lambdas()?;
317 let probe_block = probe_family.block_spec(&rho0)?;
318 let n_penalties = probe_block.initial_log_lambdas.len();
319 log::info!(
320 "[transformation-normal] exact joint setup: rho_dim={} log_kappa_dim={} dims_per_term={:?}",
321 n_penalties,
322 kappa0.len(),
323 kappa_dims,
324 );
325 let rho_floor = -12.0;
326 let rho_lower = Array1::<f64>::from_elem(n_penalties, rho_floor);
327 let rho_upper = Array1::<f64>::from_elem(n_penalties, 12.0);
328 let probe_blocks = vec![probe_block.clone()];
329 let (_, cap_hessian) = crate::custom_family::custom_family_outer_derivatives(
330 &probe_family,
331 &probe_blocks,
332 &options,
333 );
334 let analytic_gradient = analytic_psi_available;
335 let analytic_hessian_supported = analytic_gradient && cap_hessian.is_analytic();
336 let analytic_hessian = analytic_hessian_supported;
337 if analytic_hessian {
338 log::info!(
339 "[transformation-normal] CTN exact joint analytic outer Hessian is available for spatial kappa optimization; using exact second-order outer geometry"
340 );
341 }
342
343 let (rho0_min, rho0_max) = if rho0.is_empty() {
344 (0.0, 0.0)
345 } else {
346 (
347 rho0.iter().copied().fold(f64::INFINITY, f64::min),
348 rho0.iter().copied().fold(f64::NEG_INFINITY, f64::max),
349 )
350 };
351 log::info!(
352 "[transformation-normal] skipping baseline custom-family prefit before exact joint optimization \
353 (rho_dim={}, log_kappa_dim={}, rho0_range=[{:.3}, {:.3}]); using CTN warm start and penalty-scale rho seed",
354 n_penalties,
355 kappa0.len(),
356 rho0_min,
357 rho0_max,
358 );
359
360 if !analytic_psi_available {
361 return Err(
362 "transformation-normal spatial length-scale optimization requires analytic spatial psi derivatives"
363 .to_string(),
364 );
365 }
366
367 let exact_mode_branch: RefCell<ExactCoefficientModeBranch> =
373 RefCell::new(ExactCoefficientModeBranch::default());
374
375 let joint_setup =
376 ExactJointHyperSetup::new(rho0, rho_lower, rho_upper, kappa0, kappa_lower, kappa_upper);
377
378 let rv = resp_val.clone();
380 let rd = resp_deriv.clone();
381 let rp = resp_penalties.clone();
382 let rk = resp_knots.clone();
383 let rt = resp_transform.clone();
384 let rdeg = effective_config.response_degree;
385 let cfg = effective_config.clone();
386 let ws = warm_start.cloned();
387
388 let make_family =
390 |cov_design: &TermCollectionDesign| -> Result<TransformationNormalFamily, String> {
391 let effective_offset = cov_design
392 .compose_offset(offset.view(), "transformation-normal spatial fit")
393 .map_err(|error| error.to_string())?;
394 TransformationNormalFamily::from_prebuilt_response_basis(
395 response,
396 rv.clone(),
397 rd.clone(),
398 rp.clone(),
399 rk.clone(),
400 rdeg,
401 rt.clone(),
402 weights,
403 &effective_offset,
404 cov_design.design.clone(),
405 cov_design
406 .penalties
407 .iter()
408 .map(|bp| bp.to_penalty_matrix(cov_design.design.ncols()))
409 .collect(),
410 &cfg,
411 ws.as_ref(),
412 )
413 };
414
415 let block_specs_slice = [boot_spec.clone()];
416 let block_term_indices_slice = [spatial_terms.clone()];
417 let exact_geometry_cache: RefCell<Option<TransformationExactGeometryCache>> =
418 RefCell::new(None);
419 let spatial_terms_for_cache = spatial_terms.clone();
420
421 let ensure_exact_geometry = |spec: &TermCollectionSpec,
422 design: &TermCollectionDesign,
423 rho: &Array1<f64>,
424 hyper_values: &Array1<f64>|
425 -> Result<(), String> {
426 let effective_spec = freeze_term_collection_from_design(spec, design)
427 .map_err(|e| format!("failed to freeze transformation geometry key: {e}"))?;
428 let key = transformation_spatial_geometry_key(&effective_spec, &spatial_terms_for_cache)?;
429 let needs_rebuild = exact_geometry_cache
430 .borrow()
431 .as_ref()
432 .map(|cached| cached.key != key)
433 .unwrap_or(true);
434 if !needs_rebuild {
435 let mut cache = exact_geometry_cache.borrow_mut();
436 let cached = cache
437 .as_mut()
438 .ok_or_else(|| "missing transformation exact geometry cache".to_string())?;
439 if cached.hyper_layout.values().len() != hyper_values.len()
440 || cached
441 .hyper_layout
442 .values()
443 .iter()
444 .zip(hyper_values)
445 .any(|(cached, current)| cached.to_bits() != current.to_bits())
446 {
447 return Err(
448 "transformation exact geometry key reused across distinct hypercoordinate values"
449 .to_string(),
450 );
451 }
452 return cached.update_block_log_lambdas(rho);
453 }
454
455 let geom_start = std::time::Instant::now();
456 let exact_design = build_term_collection_design(covariate_data, &effective_spec)
457 .map_err(|e| format!("failed to rebuild frozen transformation geometry: {e}"))?;
458 let family = make_family(&exact_design)?;
459 let cov_psi_derivs =
460 build_block_spatial_psi_derivatives(covariate_data, &effective_spec, &exact_design)?
461 .ok_or_else(|| {
462 "missing covariate spatial psi derivatives for transformation model".to_string()
463 })?;
464 let tensor_derivs = build_tensor_psi_derivatives(&family, &cov_psi_derivs)?;
465
466 log::debug!(
467 "[transformation-normal] rebuilt exact geometry cache for {} spatial terms in {:.3}s",
468 spatial_terms_for_cache.len(),
469 geom_start.elapsed().as_secs_f64(),
470 );
471
472 exact_geometry_cache.replace(Some(TransformationExactGeometryCache {
473 key,
474 covariate_spec_resolved: effective_spec,
475 covariate_design: exact_design,
476 blocks: vec![family.block_spec(rho)?],
477 family,
478 hyper_layout: Arc::new(CustomFamilyHyperLayout::new(
479 vec![tensor_derivs],
480 Vec::new(),
481 hyper_values.clone(),
482 )?),
483 }));
484 Ok(())
485 };
486
487 let exact_mode_candidates = |eval_mode: gam_problem::EvalMode,
488 rho: &Array1<f64>|
489 -> Vec<Option<CustomFamilyWarmStart>> {
490 let (froze, candidates) = exact_mode_branch.borrow_mut().candidates(eval_mode, rho);
491 if froze {
492 log::info!(
493 "[transformation-normal] froze deterministic exact coefficient-mode branch at the first derivative-bearing outer seed evaluation"
494 );
495 }
496 candidates
497 };
498
499 log::info!(
500 "[transformation-normal] entering exact joint outer optimization \
501 (analytic_gradient={}, analytic_hessian={})",
502 analytic_gradient,
503 analytic_hessian,
504 );
505 let outer_derivative_policy =
510 probe_family.outer_derivative_policy(&probe_blocks, joint_setup.log_kappa_dim(), &options);
511 let solved = optimize_spatial_length_scale_exact_joint(
512 covariate_data,
513 &block_specs_slice,
514 &block_term_indices_slice,
515 kappa_options,
516 &joint_setup,
517 gam_solve::seeding::SeedRiskProfile::Gaussian,
518 analytic_gradient,
519 analytic_hessian,
520 true,
527 None,
528 outer_derivative_policy,
529 |theta,
531 specs: &[TermCollectionSpec],
532 designs: &[TermCollectionDesign],
533 provenance: SpatialFitProvenance<'_, CustomFamilyJointHyperModeSelection>| {
534 let rho = theta.slice(s![..joint_setup.rho_dim()]).to_owned();
535 let hyper_values = theta.slice(s![joint_setup.rho_dim()..]).to_owned();
536 ensure_exact_geometry(&specs[0], &designs[0], &rho, &hyper_values)?;
537 let mut cache_ref = exact_geometry_cache.borrow_mut();
538 let geometry = cache_ref
539 .as_mut()
540 .ok_or_else(|| "missing transformation exact geometry cache".to_string())?;
541 let final_options = crate::outer_subsample::exact_outer_options_for_row_set(
542 &options,
543 &gam_problem::outer_subsample::RowSet::All,
544 );
545 let fit = match provenance {
546 SpatialFitProvenance::NoOuterOptimization => {
547 let warm_starts =
548 exact_mode_candidates(gam_problem::EvalMode::ValueOnly, &rho);
549 let selection = evaluate_custom_family_joint_hyper_best_mode_shared(
550 &geometry.family,
551 &geometry.blocks,
552 &final_options,
553 &rho,
554 Arc::clone(&geometry.hyper_layout),
555 &warm_starts,
556 gam_problem::EvalMode::ValueOnly,
557 )
558 .map_err(|e| format!("transformation fixed mode profile: {e}"))?;
559 log::info!(
560 "[transformation-normal] user-fixed coefficient mode selected candidate={} objective={:.16e}",
561 selection.selected_candidate,
562 selection.result.objective,
563 );
564 fit_custom_family_user_fixed_log_lambdas_from_mode_selection(
565 &geometry.family,
566 &geometry.blocks,
567 &final_options,
568 selection,
569 )
570 }
571 SpatialFitProvenance::Certified { outer, mode: selection } => {
572 log::info!(
573 "[transformation-normal] consuming certified terminal coefficient mode candidate={} objective={:.16e} without profile replay",
574 selection.selected_candidate,
575 selection.result.objective,
576 );
577 fit_custom_family_fixed_log_lambdas_from_mode_selection(
578 &geometry.family,
579 &geometry.blocks,
580 &final_options,
581 selection,
582 theta,
583 outer,
584 )
585 }
586 }
587 .map_err(|e| format!("transformation fit_fn: {e}"))?;
588 if let Some(block) = fit.block_states.first() {
589 *geometry
590 .family
591 .row_quantity_cache
592 .lock()
593 .expect("CTN row quantity cache mutex poisoned") = None;
594 let final_rows = geometry.family.row_quantities(&block.beta)?;
595 let max_abs_h = final_rows
596 .h
597 .iter()
598 .copied()
599 .map(f64::abs)
600 .fold(0.0, f64::max);
601 let cov_chunk = geometry
602 .family
603 .covariate_design
604 .try_row_chunk(0..response.len())
605 .map_err(|err| {
606 format!("final CTN covariate design validation failed: {err}")
607 })?;
608 let max_abs_cov = cov_chunk.iter().copied().map(f64::abs).fold(0.0, f64::max);
609 log::info!(
610 "[transformation-normal] final fixed-rho CTN validation: max_abs_h={:.6e} max_abs_covariate_basis={:.6e}",
611 max_abs_h,
612 max_abs_cov
613 );
614 }
615 Ok(TransformationNormalFitResult {
616 family: geometry.family.clone(),
617 fit,
618 covariate_spec_resolved: geometry.covariate_spec_resolved.clone(),
619 covariate_design: geometry.covariate_design.clone(),
620 score_calibration: TransformationScoreCalibration::finite_support_pit(),
621 })
622 },
623 |theta,
625 specs: &[TermCollectionSpec],
626 designs: &[TermCollectionDesign],
627 eval_mode,
628 row_set,
629 _| {
630 let rho = theta.slice(s![..joint_setup.rho_dim()]).to_owned();
631 let hyper_values = theta.slice(s![joint_setup.rho_dim()..]).to_owned();
632 ensure_exact_geometry(&specs[0], &designs[0], &rho, &hyper_values)?;
633 let mut cache_ref = exact_geometry_cache.borrow_mut();
634 let geometry = cache_ref
635 .as_mut()
636 .ok_or_else(|| "missing transformation exact geometry cache".to_string())?;
637 let warm_starts = exact_mode_candidates(eval_mode, &rho);
638 let competing_modes = warm_starts.len() > 1;
639 let eval_options =
643 crate::outer_subsample::exact_outer_options_for_row_set(&options, row_set);
644 let selection = evaluate_custom_family_joint_hyper_best_mode_shared(
645 &geometry.family,
646 &geometry.blocks,
647 &eval_options,
648 &rho,
649 Arc::clone(&geometry.hyper_layout),
650 &warm_starts,
651 eval_mode,
652 )
653 .map_err(|e| format!("transformation exact joint mode profile: {e}"))?;
654 for (candidate_idx, rejection) in selection.rejected_candidates.iter().enumerate() {
655 if let Some(rejection) = rejection {
656 log::warn!(
657 "[transformation-normal] rejected exact coefficient-mode candidate mode_candidate={candidate_idx}: {rejection}"
658 );
659 }
660 }
661 if competing_modes {
662 for (candidate_idx, objective) in selection.screened_objectives.iter().enumerate() {
663 if let Some(objective) = objective {
664 let source = if candidate_idx == 0 {
665 "cold"
666 } else {
667 "carried"
668 };
669 log::info!(
670 "[transformation-normal] exact coefficient-mode screen source={} objective={:.16e}",
671 source,
672 objective,
673 );
674 }
675 }
676 let selected_source = if selection.selected_candidate == 0 {
677 "cold"
678 } else {
679 "carried"
680 };
681 log::info!(
682 "[transformation-normal] selected exact coefficient mode source={} objective={:.16e}",
683 selected_source,
684 selection.result.objective,
685 );
686 }
687 let objective = selection.result.objective;
688 let gradient = selection.result.gradient.clone();
689 let outer_hessian = selection.result.outer_hessian.clone();
690 exact_mode_branch
691 .borrow_mut()
692 .record_value(eval_mode, selection.result.warm_start.clone());
693
694 Ok(ExactJointEvaluation {
695 objective,
696 gradient,
697 hessian: outer_hessian,
698 mode: selection,
699 })
700 },
701 |_, _: &[TermCollectionSpec], _: &[TermCollectionDesign], _| {
702 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())
703 },
704 |_: &Array1<f64>| Ok(gam_solve::rho_optimizer::SeedOutcome::NoSlot),
705 )?;
706
707 let mut fit = solved.fit;
708 let (calibrated_fit, score_calibration) =
709 calibrate_transformation_scores(&fit.family, fit.fit.clone())?;
710 fit.fit = calibrated_fit;
711 fit.score_calibration = score_calibration;
712 Ok(fit)
713}