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