Skip to main content

gam_config/
lib.rs

1use gam_inference::formula_dsl::parse_link_choice;
2use gam_inference::model::GroupMetadata;
3use gam_models::fit_orchestration::descriptors::build_analytic_penalty_registry_from_descriptors;
4use gam_models::fit_orchestration::{CtnStage1Recipe, FitConfig};
5use gam_models::survival::location_scale::residual_distribution_inverse_link;
6use gam_models::survival::lognormal_kernel::{FrailtyScale, FrailtySpec, HazardLoading};
7use gam_models::survival::parse_survival_distribution;
8use gam_models::survival::parse_survival_likelihood_mode;
9use gam_models::transformation_normal::TransformationNormalConfig;
10use gam_problem::types::{
11    InverseLink, LinkComponent, LinkFunction, MixtureLinkSpec, SasLinkSpec, StandardLink,
12};
13use gam_solve::mixture_link::{state_from_beta_logisticspec, state_from_sasspec, state_fromspec};
14use ndarray::Array1;
15
16mod fit_request_document;
17
18pub use fit_request_document::{
19    AnalyticPenaltiesDocument, CtnStage1ConfigDocument, CtnStage1Document, FIT_REQUEST_SCHEMA,
20    FIT_REQUEST_SCHEMA_VERSION, FitRequestConfigDocument, FitRequestDocument,
21    LatentCoordinateDocument, LatentCoordinatesDocument, PrecisionHyperpriorDocument,
22    SmoothDescriptorsDocument,
23};
24
25#[derive(Clone, Copy, Debug, Eq, PartialEq)]
26pub enum CliFrailtyKind {
27    GaussianShift,
28    HazardMultiplier,
29}
30
31#[derive(Clone, Copy, Debug, Eq, PartialEq)]
32pub enum CliHazardLoading {
33    Full,
34    LoadedVsUnloaded,
35}
36
37const DEFAULT_LEARNED_FRAILTY_SCALE: FrailtyScale = FrailtyScale::Learned { initial_sigma: 0.5 };
38
39impl CtnStage1Document {
40    fn into_recipe(self) -> Result<CtnStage1Recipe, String> {
41        let mut config = TransformationNormalConfig::default();
42        if let Some(overrides) = self.config {
43            if let Some(value) = overrides.response_degree {
44                config.response_degree = value;
45            }
46            if let Some(value) = overrides.response_num_internal_knots {
47                config.response_num_internal_knots = value;
48            }
49            if let Some(value) = overrides.response_penalty_order {
50                config.response_penalty_order = value;
51            }
52            if let Some(value) = overrides.response_extra_penalty_orders {
53                config.response_extra_penalty_orders = value;
54            }
55            if let Some(value) = overrides.double_penalty {
56                config.double_penalty = value;
57            }
58        }
59        if config.response_degree == 0 {
60            return Err("ctn_stage1.config.response_degree must be >= 1".to_string());
61        }
62        if config.response_num_internal_knots < 2 {
63            return Err("ctn_stage1.config.response_num_internal_knots must be >= 2".to_string());
64        }
65        if config.response_penalty_order == 0
66            || config
67                .response_extra_penalty_orders
68                .iter()
69                .any(|order| *order == 0)
70        {
71            return Err("ctn_stage1 response penalty orders must be >= 1".to_string());
72        }
73        CtnStage1Recipe::new(
74            &self.response_column,
75            &self.covariate_formula_rhs,
76            config,
77            self.weight_column.as_deref(),
78            self.offset_column.as_deref(),
79        )
80    }
81}
82
83#[derive(Clone, Debug)]
84pub struct ResolvedFitRequest {
85    pub formula: String,
86    pub fit_config: FitConfig,
87}
88
89pub struct SurvivalInverseLinkInput<'a> {
90    pub link: Option<&'a str>,
91    pub mixture_rho: Option<&'a str>,
92    pub sas_init: Option<&'a str>,
93    pub beta_logistic_init: Option<&'a str>,
94    pub survival_distribution: &'a str,
95}
96
97pub fn parse_fit_request_json(request_json: &str) -> Result<ResolvedFitRequest, String> {
98    resolve_fit_request_document(FitRequestDocument::from_json(request_json)?)
99}
100
101/// Validate a fit request through the production resolver and serialize it in
102/// the one canonical byte representation used by every frontend.
103pub fn canonicalize_fit_request_json(request_json: &str) -> Result<String, String> {
104    let request = FitRequestDocument::from_json(request_json)?;
105    resolve_fit_request_document(request.clone())?;
106    request.to_canonical_json()
107}
108
109pub fn resolve_fit_request_document(
110    request: FitRequestDocument,
111) -> Result<ResolvedFitRequest, String> {
112    let formula = request.formula;
113    let fit_config = resolve_fit_request_config(request.config)?;
114    Ok(ResolvedFitRequest {
115        formula,
116        fit_config,
117    })
118}
119
120/// Parse the canonical config object used by non-formula helper APIs.
121/// Formula fit entry points must use [`FitRequestDocument`] instead.
122pub fn parse_fit_config_json(config_json: Option<&str>) -> Result<FitConfig, String> {
123    let config = match config_json {
124        Some(raw) if !raw.trim().is_empty() => {
125            serde_json::from_str::<FitRequestConfigDocument>(raw)
126                .map_err(|error| format!("invalid fit config object: {error}"))?
127        }
128        _ => FitRequestConfigDocument::default(),
129    };
130    resolve_fit_request_config(config)
131}
132
133pub fn resolve_fit_request_config(
134    json_config: FitRequestConfigDocument,
135) -> Result<FitConfig, String> {
136    let mut fit_config = FitConfig::default();
137    fit_config.group_metadata = json_config.group_metadata.and_then(nonempty_group_metadata);
138    if let Some(training_table_kind) = json_config.training_table_kind {
139        fit_config.training_table_kind = training_table_kind;
140    }
141    fit_config.penalty_block_gamma_priors =
142        parse_precision_hyperpriors(json_config.precision_hyperpriors)?;
143    let latent_coordinates = json_config
144        .latent_coordinates
145        .as_ref()
146        .map(|coordinates| coordinates.to_json_value())
147        .transpose()?;
148    let analytic_penalties = json_config
149        .analytic_penalties
150        .as_ref()
151        .map(|penalties| penalties.to_json_value())
152        .transpose()?;
153    build_analytic_penalty_registry_from_descriptors(
154        latent_coordinates.as_ref(),
155        analytic_penalties.as_ref(),
156    )?;
157    fit_config.latents = latent_coordinates;
158    fit_config.analytic_penalties = analytic_penalties;
159    fit_config.smooth_overrides = json_config
160        .smooth_descriptors
161        .as_ref()
162        .map(|descriptors| descriptors.to_json_value())
163        .transpose()?;
164    fit_config.family = json_config.family;
165    fit_config.negative_binomial_theta = json_config.negative_binomial_theta;
166    fit_config.expectile_tau = json_config.expectile_tau;
167    fit_config.offset_column = json_config.offset;
168    fit_config.weight_column = json_config.weights;
169    if let Some(ridge_lambda) = json_config.ridge_lambda {
170        fit_config.ridge_lambda = ridge_lambda;
171    }
172    if let Some(flag) = json_config.transformation_normal {
173        fit_config.transformation_normal = flag;
174    }
175    // `survival_likelihood` is `Option<String>` end to end (#2301): pass the
176    // caller's choice straight through — `None` (unset) stays unset so the
177    // `Surv(...)` seam resolves the one canonical default, and `Some(mode)`
178    // carries the explicit request (including onto a non-survival response,
179    // where it is a typed rejection).
180    fit_config.survival_likelihood = json_config.survival_likelihood;
181    // Passed through unvalidated on purpose: `FitConfig::resolve()` below is the
182    // canonical validation seam, so the anchor is checked in exactly one place
183    // and a direct Rust caller cannot bypass what this document layer enforces.
184    fit_config.survival_time_anchor = json_config.survival_time_anchor;
185    if let Some(distribution) = json_config.survival_distribution {
186        fit_config.survival_distribution = distribution;
187    }
188    if let Some(target) = json_config.baseline_target {
189        fit_config.baseline_target = target;
190    }
191    if let Some(value) = json_config.baseline_scale {
192        fit_config.baseline_scale = Some(value);
193    }
194    if let Some(value) = json_config.baseline_shape {
195        fit_config.baseline_shape = Some(value);
196    }
197    if let Some(value) = json_config.baseline_rate {
198        fit_config.baseline_rate = Some(value);
199    }
200    if let Some(value) = json_config.baseline_makeham {
201        fit_config.baseline_makeham = Some(value);
202    }
203    if let Some(value) = json_config.time_basis {
204        fit_config.time_basis = value;
205    }
206    if let Some(value) = json_config.time_degree {
207        fit_config.time_degree = value;
208    }
209    if let Some(value) = json_config.time_num_internal_knots {
210        fit_config.time_num_internal_knots = value;
211    }
212    if let Some(value) = json_config.time_smooth_lambda {
213        fit_config.time_smooth_lambda = value;
214    }
215    fit_config.threshold_time_k = json_config.threshold_time_k;
216    if let Some(value) = json_config.threshold_time_degree {
217        fit_config.threshold_time_degree = value;
218    }
219    fit_config.sigma_time_k = json_config.sigma_time_k;
220    if let Some(value) = json_config.sigma_time_degree {
221        fit_config.sigma_time_degree = value;
222    }
223    fit_config.logslope_time_k = json_config.logslope_time_k;
224    if let Some(value) = json_config.logslope_time_degree {
225        fit_config.logslope_time_degree = value;
226    }
227    fit_config.z_column = json_config.z_column;
228    if let Some(formula) = json_config.logslope_formula {
229        fit_config.logslope_formula = Some(formula);
230    }
231    if let Some(stage1) = json_config.ctn_stage1 {
232        fit_config.ctn_stage1 = Some(stage1.into_recipe()?);
233    }
234    fit_config.link = json_config.link;
235    if let Some(flag) = json_config.flexible_link {
236        fit_config.flexible_link = flag;
237    }
238    if let Some(flag) = json_config.precompute_conformal {
239        fit_config.precompute_conformal = Some(flag);
240    }
241    if let Some(flag) = json_config.scale_dimensions {
242        fit_config.scale_dimensions = flag;
243    }
244    if let Some(value) = json_config.pilot_subsample_threshold {
245        fit_config.spatial_optimization.pilot_subsample_threshold = value;
246    }
247    if let Some(flag) = json_config.adaptive_regularization {
248        fit_config.adaptive_regularization = Some(flag);
249    }
250    if let Some(formula) = json_config.noise_formula {
251        fit_config.noise_formula = Some(formula);
252    }
253    if let Some(column) = json_config.noise_offset {
254        fit_config.noise_offset_column = Some(column);
255    }
256    if let Some(flag) = json_config.firth {
257        fit_config.firth = flag;
258    }
259    if let Some(value) = json_config.outer_max_iter {
260        fit_config.outer_max_iter = Some(value);
261    }
262    if let Some(root) = json_config.persistent_warm_start_root {
263        fit_config = fit_config.with_persistent_warm_start_root(root);
264    }
265    if let Some(raw_gpu) = json_config.gpu {
266        fit_config.gpu_policy = parse_gpu_policy(&raw_gpu)?;
267    }
268    fit_config.frailty = parse_json_frailty_spec(
269        json_config.frailty_kind,
270        json_config.frailty_sd,
271        json_config.hazard_loading,
272    )?;
273    fit_config = fit_config.resolve()?;
274    Ok(fit_config)
275}
276
277pub fn resolve_cli_frailty_spec(
278    frailty_kind: Option<CliFrailtyKind>,
279    frailty_sd: Option<f64>,
280    hazard_loading: Option<CliHazardLoading>,
281    context: &str,
282) -> Result<FrailtySpec, String> {
283    let resolve_scale = || -> Result<FrailtyScale, String> {
284        match frailty_sd {
285            None => Ok(DEFAULT_LEARNED_FRAILTY_SCALE),
286            Some(sigma) => {
287                if !sigma.is_finite() || sigma < 0.0 {
288                    return Err(format!(
289                        "{context} requires a finite --frailty-sd >= 0, got {sigma}"
290                    ));
291                }
292                Ok(FrailtyScale::Fixed { sigma })
293            }
294        }
295    };
296
297    match frailty_kind {
298        None => {
299            if frailty_sd.is_some() || hazard_loading.is_some() {
300                return Err(format!(
301                    "{context} requires --frailty-kind when --frailty-sd or --hazard-loading is provided"
302                ));
303            }
304            Ok(FrailtySpec::None)
305        }
306        Some(CliFrailtyKind::GaussianShift) => {
307            if hazard_loading.is_some() {
308                return Err(format!(
309                    "{context} does not accept --hazard-loading with --frailty-kind gaussian-shift"
310                ));
311            }
312            Ok(FrailtySpec::GaussianShift {
313                scale: resolve_scale()?,
314            })
315        }
316        Some(CliFrailtyKind::HazardMultiplier) => Ok(FrailtySpec::HazardMultiplier {
317            scale: resolve_scale()?,
318            loading: hazard_loading.map(cli_hazard_loading).ok_or_else(|| {
319                format!("{context} requires --hazard-loading with --frailty-kind hazard-multiplier")
320            })?,
321        }),
322    }
323}
324
325pub fn parse_survival_likelihood_cli(raw: &str) -> Result<String, String> {
326    let normalized = raw.trim().to_ascii_lowercase();
327    parse_survival_likelihood_mode(&normalized)?;
328    Ok(normalized)
329}
330
331pub fn parse_baseline_target_cli(raw: &str) -> Result<String, String> {
332    let normalized = raw.trim().to_ascii_lowercase();
333    match normalized.as_str() {
334        "linear" | "weibull" | "gompertz" | "gompertz-makeham" => Ok(normalized),
335        other => Err(format!(
336            "unsupported --baseline-target '{other}'; use linear|weibull|gompertz|gompertz-makeham"
337        )),
338    }
339}
340
341pub fn parse_comma_f64(v: &str, label: &str) -> Result<Vec<f64>, String> {
342    let mut out = Vec::new();
343    for part in v.split(',') {
344        let t = part.trim();
345        if t.is_empty() {
346            continue;
347        }
348        let parsed = t
349            .parse::<f64>()
350            .map_err(|err| format!("{label} contains non-numeric value '{t}': {err}"))?;
351        if !parsed.is_finite() {
352            return Err(format!("{label} contains non-finite value '{t}'"));
353        }
354        out.push(parsed);
355    }
356    Ok(out)
357}
358
359pub fn effective_link_to_standard(
360    link: LinkFunction,
361    context: &str,
362) -> Result<StandardLink, String> {
363    StandardLink::try_from(link).map_err(|_| {
364        format!(
365            "{context}: state-bearing link `{}` must be routed through `InverseLink::Sas` / `InverseLink::BetaLogistic`, not `Standard(_)`",
366            link.name()
367        )
368    })
369}
370
371pub fn parse_survival_inverse_link(
372    input: SurvivalInverseLinkInput<'_>,
373) -> Result<InverseLink, String> {
374    if let Some(raw) = input.link {
375        let name = raw.trim().to_ascii_lowercase();
376        if name == "loglog" || name == "cauchit" {
377            // `loglog` and `cauchit` have no scalar `LinkFunction`/`StandardLink`
378            // representative, but the blended-link kernels implement their inverse
379            // link and derivative jets exactly (`LinkComponent::LogLog` /
380            // `LinkComponent::Cauchit`). Represent a survival `--link loglog` /
381            // `--link cauchit` as a single-component mixture: it carries weight 1.0
382            // with no free mixing logits, so it evaluates as exactly that link and
383            // flows end-to-end through the fully-wired `InverseLink::Mixture` survival
384            // path (prepare/construct/row-kernel/predict).
385            if input.sas_init.is_some() {
386                return Err("--sas-init requires --link sas".to_string());
387            }
388            if input.beta_logistic_init.is_some() {
389                return Err("--beta-logistic-init requires --link beta-logistic".to_string());
390            }
391            if input.mixture_rho.is_some() {
392                return Err(
393                    "--mixture-rho requires survival --link blended(...)/mixture(...)".to_string(),
394                );
395            }
396            let component = if name == "loglog" {
397                LinkComponent::LogLog
398            } else {
399                LinkComponent::Cauchit
400            };
401            return state_fromspec(&MixtureLinkSpec {
402                components: vec![component],
403                initial_rho: Array1::zeros(0),
404            })
405            .map(InverseLink::Mixture)
406            .map_err(|e| format!("invalid survival {name} link state: {e}"));
407        }
408    }
409    let choice = parse_link_choice(input.link, false).map_err(|err| {
410        let err = err.to_string();
411        if let Some(raw) = input.link {
412            let name = raw.trim().to_ascii_lowercase();
413            if err.starts_with("unsupported --link ") || err.starts_with("unsupported link type ") {
414                return format!(
415                    "unsupported survival --link '{name}'; {}",
416                    survival_link_usage()
417                );
418            }
419        }
420        err
421    })?;
422    if let Some(choice) = choice {
423        if let Some(components) = choice.mixture_components {
424            if input.sas_init.is_some() || input.beta_logistic_init.is_some() {
425                return Err(
426                    "survival blended(...) link does not accept --sas-init/--beta-logistic-init"
427                        .to_string(),
428                );
429            }
430            let expected = components.len().saturating_sub(1);
431            let initial_rho = if let Some(raw) = input.mixture_rho {
432                let vals = parse_comma_f64(raw, "--mixture-rho")?;
433                if vals.len() != expected {
434                    return Err(format!(
435                        "--mixture-rho expects {expected} values for blended({})",
436                        components
437                            .iter()
438                            .map(|component| component.name())
439                            .collect::<Vec<_>>()
440                            .join(",")
441                    ));
442                }
443                Array1::from_vec(vals)
444            } else {
445                Array1::zeros(expected)
446            };
447            return state_fromspec(&MixtureLinkSpec {
448                components,
449                initial_rho,
450            })
451            .map(InverseLink::Mixture)
452            .map_err(|e| format!("invalid survival blended link state: {e}"));
453        }
454
455        if input.mixture_rho.is_some() {
456            return Err(
457                "--mixture-rho requires survival --link blended(...)/mixture(...)".to_string(),
458            );
459        }
460        match choice.link {
461            LinkFunction::Sas => {
462                if input.beta_logistic_init.is_some() {
463                    return Err("--beta-logistic-init requires --link beta-logistic".to_string());
464                }
465                let (epsilon, log_delta) = if let Some(raw) = input.sas_init {
466                    let vals = parse_comma_f64(raw, "--sas-init")?;
467                    if vals.len() != 2 {
468                        return Err(format!(
469                            "--sas-init expects two values: epsilon,log_delta (got {})",
470                            vals.len()
471                        ));
472                    }
473                    (vals[0], vals[1])
474                } else {
475                    (0.0, 0.0)
476                };
477                state_from_sasspec(SasLinkSpec {
478                    initial_epsilon: epsilon,
479                    initial_log_delta: log_delta,
480                })
481                .map(InverseLink::Sas)
482                .map_err(|e| format!("invalid survival SAS link state: {e}"))
483            }
484            LinkFunction::BetaLogistic => {
485                if input.sas_init.is_some() {
486                    return Err("--sas-init requires --link sas".to_string());
487                }
488                let (epsilon, delta) = if let Some(raw) = input.beta_logistic_init {
489                    let vals = parse_comma_f64(raw, "--beta-logistic-init")?;
490                    if vals.len() != 2 {
491                        return Err(format!(
492                            "--beta-logistic-init expects two values: epsilon,delta (got {})",
493                            vals.len()
494                        ));
495                    }
496                    (vals[0], vals[1])
497                } else {
498                    (0.0, 0.0)
499                };
500                state_from_beta_logisticspec(SasLinkSpec {
501                    initial_epsilon: epsilon,
502                    initial_log_delta: delta,
503                })
504                .map(InverseLink::BetaLogistic)
505                .map_err(|e| format!("invalid survival Beta-Logistic link state: {e}"))
506            }
507            LinkFunction::Log => Err(format!(
508                "unsupported survival --link 'log'; {}",
509                survival_link_usage()
510            )),
511            other => {
512                if input.sas_init.is_some() {
513                    return Err("--sas-init requires --link sas".to_string());
514                }
515                if input.beta_logistic_init.is_some() {
516                    return Err("--beta-logistic-init requires --link beta-logistic".to_string());
517                }
518                Ok(InverseLink::Standard(effective_link_to_standard(
519                    other,
520                    "survival inverse link",
521                )?))
522            }
523        }
524    } else {
525        if input.mixture_rho.is_some() {
526            return Err("--mixture-rho requires --link blended(...)/mixture(...)".to_string());
527        }
528        if input.sas_init.is_some() {
529            return Err("--sas-init requires --link sas".to_string());
530        }
531        if input.beta_logistic_init.is_some() {
532            return Err("--beta-logistic-init requires --link beta-logistic".to_string());
533        }
534        let dist = parse_survival_distribution(input.survival_distribution)?;
535        Ok(residual_distribution_inverse_link(dist))
536    }
537}
538
539fn parse_json_frailty_spec(
540    frailty_kind: Option<String>,
541    frailty_sd: Option<f64>,
542    hazard_loading: Option<String>,
543) -> Result<FrailtySpec, String> {
544    if let Some(kind) = frailty_kind {
545        let trimmed = kind.trim().to_ascii_lowercase();
546        let scale = frailty_sd
547            .map(|sigma| FrailtyScale::Fixed { sigma })
548            .unwrap_or(DEFAULT_LEARNED_FRAILTY_SCALE);
549        let hazard_loading = hazard_loading
550            .as_ref()
551            .map(|raw| raw.trim().to_ascii_lowercase());
552        let frailty = match trimmed.as_str() {
553            "none" | "" => {
554                if frailty_sd.is_some() || hazard_loading.is_some() {
555                    return Err(
556                        "frailty_kind='none' does not accept frailty_sd or hazard_loading"
557                            .to_string(),
558                    );
559                }
560                FrailtySpec::None
561            }
562            "hazard-multiplier" => {
563                let loading = match hazard_loading.as_deref() {
564                    Some("full") | None => HazardLoading::Full,
565                    Some("loaded-vs-unloaded") => HazardLoading::LoadedVsUnloaded,
566                    Some(other) => {
567                        return Err(format!(
568                            "unknown hazard_loading '{other}'; supported: 'full', 'loaded-vs-unloaded'"
569                        ));
570                    }
571                };
572                FrailtySpec::HazardMultiplier { scale, loading }
573            }
574            "gaussian-shift" => {
575                if hazard_loading.is_some() {
576                    return Err(
577                        "hazard_loading is valid only with frailty_kind='hazard-multiplier'"
578                            .to_string(),
579                    );
580                }
581                FrailtySpec::GaussianShift { scale }
582            }
583            other => {
584                return Err(format!(
585                    "unknown frailty_kind '{other}'; supported: 'none', 'hazard-multiplier', 'gaussian-shift'"
586                ));
587            }
588        };
589        frailty.validate().map_err(|err| err.to_string())?;
590        Ok(frailty)
591    } else if frailty_sd.is_some() || hazard_loading.is_some() {
592        Err("frailty_kind is required when frailty_sd or hazard_loading is provided".to_string())
593    } else {
594        Ok(FrailtySpec::None)
595    }
596}
597
598fn cli_hazard_loading(loading: CliHazardLoading) -> HazardLoading {
599    match loading {
600        CliHazardLoading::Full => HazardLoading::Full,
601        CliHazardLoading::LoadedVsUnloaded => HazardLoading::LoadedVsUnloaded,
602    }
603}
604
605fn parse_precision_hyperpriors(
606    precision_hyperpriors: Option<std::collections::BTreeMap<String, PrecisionHyperpriorDocument>>,
607) -> Result<Vec<(String, f64, f64)>, String> {
608    let mut out = Vec::with_capacity(precision_hyperpriors.as_ref().map_or(0, |map| map.len()));
609    for (label, prior) in precision_hyperpriors.unwrap_or_default() {
610        if label.trim().is_empty() {
611            return Err("precision_hyperpriors keys must be non-empty".to_string());
612        }
613        if !prior.shape.is_finite() || prior.shape <= 0.0 {
614            return Err(format!(
615                "precision_hyperpriors['{label}'].shape must be finite and > 0"
616            ));
617        }
618        if !prior.rate.is_finite() || prior.rate < 0.0 {
619            return Err(format!(
620                "precision_hyperpriors['{label}'].rate must be finite and >= 0"
621            ));
622        }
623        out.push((label, prior.shape, prior.rate));
624    }
625    Ok(out)
626}
627
628fn nonempty_group_metadata(metadata: GroupMetadata) -> Option<GroupMetadata> {
629    if metadata.is_empty() {
630        None
631    } else {
632        Some(metadata)
633    }
634}
635
636fn parse_gpu_policy(raw_gpu: &str) -> Result<gam_gpu::GpuPolicy, String> {
637    gam_gpu::GpuPolicy::parse(raw_gpu).ok_or_else(|| {
638        format!(
639            "invalid gpu policy '{}'; supported values are auto, off, required",
640            raw_gpu
641        )
642    })
643}
644
645fn survival_link_usage() -> &'static str {
646    "use identity|logit|probit|cloglog|loglog|cauchit|sas|beta-logistic|blended(...)/mixture(...) or flexible(...)"
647}
648
649#[cfg(test)]
650mod tests {
651    use super::*;
652    use gam_models::survival::lognormal_kernel::FrailtySpec;
653    use serde_json::{Value, json};
654
655    struct ParityCase {
656        name: &'static str,
657        cli: FitConfig,
658        json: Value,
659    }
660
661    fn base_cli() -> FitConfig {
662        FitConfig::default()
663    }
664
665    fn resolved_cli(input: FitConfig) -> Result<FitConfig, String> {
666        input.resolve()
667    }
668
669    fn resolved_json(config: Value) -> Result<FitConfig, String> {
670        let config = serde_json::from_value::<FitRequestConfigDocument>(config)
671            .map_err(|error| format!("invalid test fit config: {error}"))?;
672        let request = FitRequestDocument::new("y ~ x", config)?;
673        resolve_fit_request_document(request).map(|resolved| {
674            assert_eq!(resolved.formula, "y ~ x");
675            resolved.fit_config
676        })
677    }
678
679    fn canonical_fit_config(config: FitConfig) -> String {
680        format!("{config:#?}")
681    }
682
683    /// #2633: the conformal-precompute switch must reach `FitConfig` through the
684    /// shared wire document, which is the single path BOTH front ends use — the
685    /// CLI maps `--precompute-conformal` into this document and the Python FFI
686    /// parses the same JSON key. A knob only reachable from Rust would be the
687    /// front-end parity gap this campaign exists to remove.
688    #[test]
689    fn precompute_conformal_threads_from_the_wire_document_2633() {
690        // Absent means "use the default", which is to precompute. It must stay
691        // `None` rather than being materialized into `Some(true)`, so the core
692        // default remains the single source of truth for the behaviour.
693        let defaulted = resolved_json(json!({})).expect("empty config resolves");
694        assert_eq!(
695            defaulted.precompute_conformal, None,
696            "omitting the key must leave the core default untouched"
697        );
698
699        let off = resolved_json(json!({"precompute_conformal": false}))
700            .expect("precompute_conformal=false resolves");
701        assert_eq!(
702            off.precompute_conformal,
703            Some(false),
704            "an explicit false must reach FitConfig, or the substrates are still precomputed"
705        );
706
707        let on = resolved_json(json!({"precompute_conformal": true}))
708            .expect("precompute_conformal=true resolves");
709        assert_eq!(on.precompute_conformal, Some(true));
710    }
711
712    #[test]
713    fn persistent_warm_start_is_disabled_by_default_and_preserves_explicit_root_2639() {
714        let defaulted = resolved_json(json!({})).expect("empty config resolves");
715        assert!(
716            defaulted.persistent_warm_start_store.is_none(),
717            "omitting the root must leave persistence disabled"
718        );
719
720        let exact_root = std::path::PathBuf::from("caller-owned/../warm-root");
721        let configured = resolved_json(json!({
722            "persistent_warm_start_root": exact_root
723        }))
724        .expect("an explicit persistence root resolves")
725        .persistent_warm_start_store
726        .expect("the root must become a store capability");
727        assert_eq!(
728            configured.root(),
729            exact_root,
730            "configuration must not canonicalize or relocate the caller's root"
731        );
732
733        let empty = resolved_json(json!({"persistent_warm_start_root": ""}))
734            .expect_err("an empty persistence root is not an explicit location");
735        assert!(empty.contains("persistent_warm_start_root must not be empty"));
736    }
737
738    #[test]
739    fn frailty_resolvers_preserve_fixed_vs_learned_scale_mode() {
740        assert_eq!(
741            resolve_cli_frailty_spec(Some(CliFrailtyKind::GaussianShift), Some(0.3), None, "test",)
742                .unwrap(),
743            FrailtySpec::GaussianShift {
744                scale: FrailtyScale::Fixed { sigma: 0.3 },
745            }
746        );
747        assert_eq!(
748            resolve_cli_frailty_spec(Some(CliFrailtyKind::GaussianShift), None, None, "test",)
749                .unwrap(),
750            FrailtySpec::GaussianShift {
751                scale: DEFAULT_LEARNED_FRAILTY_SCALE,
752            }
753        );
754        assert_eq!(
755            parse_json_frailty_spec(
756                Some("hazard-multiplier".to_string()),
757                None,
758                Some("full".to_string()),
759            )
760            .unwrap(),
761            FrailtySpec::HazardMultiplier {
762                scale: DEFAULT_LEARNED_FRAILTY_SCALE,
763                loading: HazardLoading::Full,
764            }
765        );
766    }
767
768    #[test]
769    fn rich_request_document_resolves_every_frontend_parity_field() {
770        let request = FitRequestDocument::new(
771            "y ~ duchon(z)",
772            FitRequestConfigDocument {
773                ctn_stage1: Some(CtnStage1Document {
774                    response_column: "dose".to_string(),
775                    covariate_formula_rhs: "s(age)".to_string(),
776                    config: Some(CtnStage1ConfigDocument {
777                        response_degree: Some(4),
778                        response_num_internal_knots: Some(9),
779                        response_penalty_order: Some(2),
780                        response_extra_penalty_orders: Some(vec![1, 3]),
781                        double_penalty: Some(false),
782                    }),
783                    weight_column: Some("case_weight".to_string()),
784                    offset_column: Some("stage1_offset".to_string()),
785                }),
786                precision_hyperpriors: Some(std::collections::BTreeMap::from([(
787                    "duchon(z):roughness".to_string(),
788                    PrecisionHyperpriorDocument {
789                        shape: 2.5,
790                        rate: 0.75,
791                    },
792                )])),
793                latent_coordinates: Some(
794                    serde_json::from_value(json!({
795                        "z": {"n": 20, "d": 2, "name": "z", "init": "pca"}
796                    }))
797                    .unwrap(),
798                ),
799                analytic_penalties: Some(AnalyticPenaltiesDocument(vec![json!({
800                    "kind": "orthogonality",
801                    "target": "z",
802                    "weight": 1.25
803                })])),
804                smooth_descriptors: Some(
805                    serde_json::from_value(json!({
806                        "z": {"kind": "duchon", "vars": ["z"], "centers": 8}
807                    }))
808                    .unwrap(),
809                ),
810                ..FitRequestConfigDocument::default()
811            },
812        )
813        .unwrap();
814
815        let canonical = request.to_canonical_json().unwrap();
816        assert_eq!(
817            canonicalize_fit_request_json(&canonical).unwrap(),
818            canonical
819        );
820        let resolved = parse_fit_request_json(&canonical).unwrap();
821        assert_eq!(resolved.formula, "y ~ duchon(z)");
822        let config = resolved.fit_config;
823        let stage1 = config.ctn_stage1.unwrap();
824        assert_eq!(stage1.response_column, "dose");
825        assert_eq!(stage1.covariate_formula_rhs, "s(age)");
826        assert_eq!(stage1.config.response_degree, 4);
827        assert_eq!(stage1.config.response_num_internal_knots, 9);
828        assert_eq!(stage1.config.response_extra_penalty_orders, vec![1, 3]);
829        assert_eq!(
830            config.penalty_block_gamma_priors,
831            vec![("duchon(z):roughness".to_string(), 2.5, 0.75)]
832        );
833        assert_eq!(config.latents.unwrap()["z"]["d"], json!(2));
834        assert_eq!(config.analytic_penalties.unwrap()[0]["target"], "z");
835        assert_eq!(config.smooth_overrides.unwrap()["z"]["kind"], "duchon");
836    }
837
838    #[test]
839    fn rich_request_rejects_invalid_prior_and_order_dependent_penalty_target() {
840        let invalid_prior = FitRequestDocument::new(
841            "y ~ x",
842            FitRequestConfigDocument {
843                precision_hyperpriors: Some(std::collections::BTreeMap::from([(
844                    "x".to_string(),
845                    PrecisionHyperpriorDocument {
846                        shape: 0.0,
847                        rate: 1.0,
848                    },
849                )])),
850                ..FitRequestConfigDocument::default()
851            },
852        )
853        .unwrap();
854        assert!(
855            resolve_fit_request_document(invalid_prior)
856                .unwrap_err()
857                .contains("shape must be finite and > 0")
858        );
859
860        let numeric_target = FitRequestDocument::new(
861            "y ~ s(z)",
862            FitRequestConfigDocument {
863                latent_coordinates: Some(
864                    serde_json::from_value(json!({"z": {"n": 4, "d": 1}})).unwrap(),
865                ),
866                analytic_penalties: Some(AnalyticPenaltiesDocument(vec![json!({
867                    "kind": "orthogonality",
868                    "target": 0
869                })])),
870                ..FitRequestConfigDocument::default()
871            },
872        )
873        .unwrap();
874        assert!(
875            resolve_fit_request_document(numeric_target)
876                .unwrap_err()
877                .contains("target must be a latent-coordinate name")
878        );
879    }
880
881    #[test]
882    fn cli_shaped_and_json_wire_config_resolution_match() {
883        let cases = vec![
884            ParityCase {
885                name: "family and link selection",
886                cli: {
887                    let mut input = base_cli();
888                    input.family = Some("binomial".to_string());
889                    input.link = Some("probit".to_string());
890                    input.flexible_link = true;
891                    input
892                },
893                json: json!({
894                    "family": "binomial",
895                    "link": "probit",
896                    "flexible_link": true
897                }),
898            },
899            ParityCase {
900                name: "offset weights ridge and noise offset columns",
901                cli: {
902                    let mut input = base_cli();
903                    input.offset_column = Some("eta_offset".to_string());
904                    input.weight_column = Some("case_weight".to_string());
905                    input.noise_offset_column = Some("sigma_offset".to_string());
906                    input.ridge_lambda = 0.125;
907                    input
908                },
909                json: json!({
910                    "offset": "eta_offset",
911                    "weights": "case_weight",
912                    "noise_offset": "sigma_offset",
913                    "ridge_lambda": 0.125
914                }),
915            },
916            ParityCase {
917                name: "weibull survival likelihood and baseline scale shape",
918                cli: {
919                    let mut input = base_cli();
920                    input.survival_likelihood = Some("weibull".to_string());
921                    input.baseline_target = "weibull".to_string();
922                    input.baseline_scale = Some(2.5);
923                    input.baseline_shape = Some(1.75);
924                    input
925                },
926                json: json!({
927                    "survival_likelihood": "weibull",
928                    "baseline_target": "weibull",
929                    "baseline_scale": 2.5,
930                    "baseline_shape": 1.75
931                }),
932            },
933            ParityCase {
934                name: "transformation survival gompertz makeham baseline",
935                cli: {
936                    let mut input = base_cli();
937                    input.survival_likelihood = Some("transformation".to_string());
938                    input.baseline_target = "gompertz-makeham".to_string();
939                    input.baseline_shape = Some(1.2);
940                    input.baseline_rate = Some(0.04);
941                    input.baseline_makeham = Some(0.01);
942                    input
943                },
944                json: json!({
945                    "survival_likelihood": "transformation",
946                    "baseline_target": "gompertz-makeham",
947                    "baseline_shape": 1.2,
948                    "baseline_rate": 0.04,
949                    "baseline_makeham": 0.01
950                }),
951            },
952            ParityCase {
953                name: "survival likelihood values are canonicalized",
954                cli: {
955                    let mut input = base_cli();
956                    input.survival_likelihood = Some("TRANSFORMATION".to_string());
957                    input
958                },
959                json: json!({
960                    "survival_likelihood": "Transformation"
961                }),
962            },
963            ParityCase {
964                name: "noise formula logslope z column and scale dimensions",
965                cli: {
966                    let mut input = base_cli();
967                    input.noise_formula = Some("~ s(age) + treatment".to_string());
968                    input.logslope_formula = Some("~ s(dose)".to_string());
969                    input.z_column = Some("dose".to_string());
970                    input.scale_dimensions = true;
971                    input
972                },
973                json: json!({
974                    "noise_formula": "~ s(age) + treatment",
975                    "logslope_formula": "~ s(dose)",
976                    "z_column": "dose",
977                    "scale_dimensions": true
978                }),
979            },
980            ParityCase {
981                name: "firth transformation normal outer iterations and adaptive regularization",
982                cli: {
983                    let mut input = base_cli();
984                    input.firth = true;
985                    input.transformation_normal = true;
986                    input.outer_max_iter = Some(7);
987                    input.adaptive_regularization = Some(true);
988                    input
989                },
990                json: json!({
991                    "firth": true,
992                    "transformation_normal": true,
993                    "outer_max_iter": 7,
994                    "adaptive_regularization": true
995                }),
996            },
997            ParityCase {
998                name: "gpu policy toggle",
999                cli: {
1000                    let mut input = base_cli();
1001                    input.gpu_policy = gam_gpu::GpuPolicy::Off;
1002                    input
1003                },
1004                json: json!({
1005                    "gpu": "off"
1006                }),
1007            },
1008            ParityCase {
1009                name: "hazard multiplier frailty fields",
1010                cli: {
1011                    let mut input = base_cli();
1012                    input.frailty = FrailtySpec::HazardMultiplier {
1013                        scale: FrailtyScale::Fixed { sigma: 0.35 },
1014                        loading: HazardLoading::LoadedVsUnloaded,
1015                    };
1016                    input
1017                },
1018                json: json!({
1019                    "frailty_kind": "hazard-multiplier",
1020                    "frailty_sd": 0.35,
1021                    "hazard_loading": "loaded-vs-unloaded"
1022                }),
1023            },
1024            ParityCase {
1025                name: "gaussian shift frailty fields",
1026                cli: {
1027                    let mut input = base_cli();
1028                    input.frailty = FrailtySpec::GaussianShift {
1029                        scale: FrailtyScale::Fixed { sigma: 0.2 },
1030                    };
1031                    input
1032                },
1033                json: json!({
1034                    "frailty_kind": "gaussian-shift",
1035                    "frailty_sd": 0.2
1036                }),
1037            },
1038        ];
1039
1040        for case in cases {
1041            let cli = resolved_cli(case.cli)
1042                .unwrap_or_else(|err| panic!("{}: CLI-shaped config failed: {err}", case.name));
1043            let json = resolved_json(case.json)
1044                .unwrap_or_else(|err| panic!("{}: JSON wire config failed: {err}", case.name));
1045            assert_eq!(
1046                canonical_fit_config(cli),
1047                canonical_fit_config(json),
1048                "{}",
1049                case.name
1050            );
1051        }
1052    }
1053
1054    #[test]
1055    fn cli_shaped_and_json_wire_config_resolution_rejections_match() {
1056        let cases = vec![
1057            ParityCase {
1058                name: "negative ridge lambda",
1059                cli: {
1060                    let mut input = base_cli();
1061                    input.ridge_lambda = -1.0;
1062                    input
1063                },
1064                json: json!({
1065                    "ridge_lambda": -1.0
1066                }),
1067            },
1068            ParityCase {
1069                name: "linear baseline rejects shape",
1070                cli: {
1071                    let mut input = base_cli();
1072                    input.baseline_shape = Some(1.1);
1073                    input
1074                },
1075                json: json!({
1076                    "baseline_shape": 1.1
1077                }),
1078            },
1079            ParityCase {
1080                name: "weibull likelihood rejects gompertz target",
1081                cli: {
1082                    let mut input = base_cli();
1083                    input.survival_likelihood = Some("weibull".to_string());
1084                    input.baseline_target = "gompertz".to_string();
1085                    input
1086                },
1087                json: json!({
1088                    "survival_likelihood": "weibull",
1089                    "baseline_target": "gompertz"
1090                }),
1091            },
1092        ];
1093
1094        for case in cases {
1095            let cli = resolved_cli(case.cli).expect_err(case.name);
1096            let json = resolved_json(case.json).expect_err(case.name);
1097            assert_eq!(cli, json, "{}", case.name);
1098        }
1099    }
1100
1101    // ── parse_comma_f64 ───────────────────────────────────────────────────
1102
1103    #[test]
1104    fn parse_comma_f64_empty_string_returns_empty_vec() {
1105        assert_eq!(parse_comma_f64("", "x").unwrap(), Vec::<f64>::new());
1106        assert_eq!(parse_comma_f64("   ", "x").unwrap(), Vec::<f64>::new());
1107    }
1108
1109    #[test]
1110    fn parse_comma_f64_single_value() {
1111        assert_eq!(parse_comma_f64("3.14", "x").unwrap(), vec![3.14]);
1112    }
1113
1114    #[test]
1115    fn parse_comma_f64_multiple_values_with_spaces() {
1116        let result = parse_comma_f64("1.0, 2.5, -3.0", "x").unwrap();
1117        assert_eq!(result, vec![1.0, 2.5, -3.0]);
1118    }
1119
1120    #[test]
1121    fn parse_comma_f64_non_numeric_returns_error() {
1122        let err = parse_comma_f64("1.0, bad, 3.0", "--vals").unwrap_err();
1123        assert!(err.contains("--vals"), "error should name the label: {err}");
1124        assert!(
1125            err.contains("bad"),
1126            "error should name the bad token: {err}"
1127        );
1128    }
1129
1130    #[test]
1131    fn parse_comma_f64_infinity_returns_error() {
1132        let err = parse_comma_f64("inf", "--vals").unwrap_err();
1133        assert!(
1134            err.contains("non-finite"),
1135            "error should say non-finite: {err}"
1136        );
1137    }
1138
1139    #[test]
1140    fn parse_comma_f64_nan_returns_error() {
1141        let err = parse_comma_f64("nan", "--vals").unwrap_err();
1142        assert!(
1143            err.contains("non-finite"),
1144            "error should say non-finite: {err}"
1145        );
1146    }
1147
1148    // ── parse_survival_likelihood_cli ─────────────────────────────────────
1149
1150    #[test]
1151    fn parse_survival_likelihood_cli_valid_values() {
1152        assert_eq!(
1153            parse_survival_likelihood_cli("transformation").unwrap(),
1154            "transformation"
1155        );
1156        assert_eq!(parse_survival_likelihood_cli("weibull").unwrap(), "weibull");
1157        // case-insensitive
1158        assert_eq!(parse_survival_likelihood_cli("WEIBULL").unwrap(), "weibull");
1159        assert_eq!(
1160            parse_survival_likelihood_cli("Transformation").unwrap(),
1161            "transformation"
1162        );
1163    }
1164
1165    #[test]
1166    fn parse_survival_likelihood_cli_invalid_returns_error() {
1167        assert!(parse_survival_likelihood_cli("lognormal").is_err());
1168        assert!(parse_survival_likelihood_cli("").is_err());
1169    }
1170
1171    // ── parse_baseline_target_cli ─────────────────────────────────────────
1172
1173    #[test]
1174    fn parse_baseline_target_cli_valid_values() {
1175        for target in &["linear", "weibull", "gompertz", "gompertz-makeham"] {
1176            assert_eq!(
1177                parse_baseline_target_cli(target).unwrap(),
1178                *target,
1179                "should accept '{target}'"
1180            );
1181        }
1182        // trimmed and lowercased
1183        assert_eq!(parse_baseline_target_cli("  Weibull  ").unwrap(), "weibull");
1184    }
1185
1186    #[test]
1187    fn parse_baseline_target_cli_invalid_returns_error() {
1188        let err = parse_baseline_target_cli("cox").unwrap_err();
1189        assert!(
1190            err.contains("cox"),
1191            "error should name the bad value: {err}"
1192        );
1193    }
1194}