Skip to main content

gam_config/
fit_request_document.rs

1use serde::{Deserialize, Serialize};
2use serde_json::Value as JsonValue;
3use std::collections::BTreeMap;
4use std::path::PathBuf;
5
6/// Stable identity of the serialized fit-request document.
7pub const FIT_REQUEST_SCHEMA: &str = "gam.fit-request";
8
9/// Current fit-request schema version.
10pub const FIT_REQUEST_SCHEMA_VERSION: u32 = 1;
11
12/// A complete, frontend-neutral formula fit request.
13///
14/// Training data is intentionally not embedded: Rust callers supply a
15/// dataset, Python supplies an in-memory table/array, and the CLI
16/// supplies a dataset path. Everything that changes the fitted model belongs in
17/// this document.
18#[derive(Clone, Debug, Deserialize, PartialEq, Serialize)]
19#[serde(deny_unknown_fields)]
20pub struct FitRequestDocument {
21    pub schema: String,
22    pub schema_version: u32,
23    pub formula: String,
24    #[serde(default)]
25    pub config: FitRequestConfigDocument,
26}
27
28impl FitRequestDocument {
29    pub fn new(
30        formula: impl Into<String>,
31        config: FitRequestConfigDocument,
32    ) -> Result<Self, String> {
33        let document = Self {
34            schema: FIT_REQUEST_SCHEMA.to_string(),
35            schema_version: FIT_REQUEST_SCHEMA_VERSION,
36            formula: formula.into(),
37            config,
38        };
39        document.validate()?;
40        Ok(document)
41    }
42
43    pub fn from_json(raw: &str) -> Result<Self, String> {
44        let document = serde_json::from_str::<Self>(raw)
45            .map_err(|error| format!("invalid fit request document: {error}"))?;
46        document.validate()?;
47        Ok(document)
48    }
49
50    pub fn to_canonical_json(&self) -> Result<String, String> {
51        self.validate()?;
52        serde_json::to_string(self)
53            .map_err(|error| format!("failed to serialize fit request document: {error}"))
54    }
55
56    fn validate(&self) -> Result<(), String> {
57        if self.schema != FIT_REQUEST_SCHEMA {
58            return Err(format!(
59                "fit request schema must be '{FIT_REQUEST_SCHEMA}', got {:?}",
60                self.schema
61            ));
62        }
63        if self.schema_version != FIT_REQUEST_SCHEMA_VERSION {
64            return Err(format!(
65                "unsupported fit request schema_version {}; expected {}",
66                self.schema_version, FIT_REQUEST_SCHEMA_VERSION
67            ));
68        }
69        if self.formula.trim().is_empty() {
70            return Err("fit request formula must be non-empty".to_string());
71        }
72        Ok(())
73    }
74}
75
76/// Serializable model configuration shared by Rust, Python, and the CLI.
77///
78/// Optional fields mean "use the core [`gam_models::fit_orchestration::FitConfig`]
79/// default". The document deliberately has one spelling for each concept; the
80/// parser does not carry aliases or legacy wire formats.
81#[derive(Clone, Debug, Default, Deserialize, PartialEq, Serialize)]
82#[serde(deny_unknown_fields)]
83pub struct FitRequestConfigDocument {
84    #[serde(skip_serializing_if = "Option::is_none")]
85    pub adaptive_regularization: Option<bool>,
86    #[serde(skip_serializing_if = "Option::is_none")]
87    pub baseline_makeham: Option<f64>,
88    #[serde(skip_serializing_if = "Option::is_none")]
89    pub baseline_rate: Option<f64>,
90    #[serde(skip_serializing_if = "Option::is_none")]
91    pub baseline_scale: Option<f64>,
92    #[serde(skip_serializing_if = "Option::is_none")]
93    pub baseline_shape: Option<f64>,
94    #[serde(skip_serializing_if = "Option::is_none")]
95    pub baseline_target: Option<String>,
96    #[serde(skip_serializing_if = "Option::is_none")]
97    pub ctn_stage1: Option<CtnStage1Document>,
98    #[serde(skip_serializing_if = "Option::is_none")]
99    pub expectile_tau: Option<f64>,
100    #[serde(skip_serializing_if = "Option::is_none")]
101    pub family: Option<String>,
102    #[serde(skip_serializing_if = "Option::is_none")]
103    pub firth: Option<bool>,
104    #[serde(skip_serializing_if = "Option::is_none")]
105    pub flexible_link: Option<bool>,
106    #[serde(skip_serializing_if = "Option::is_none")]
107    pub frailty_kind: Option<String>,
108    #[serde(skip_serializing_if = "Option::is_none")]
109    pub frailty_sd: Option<f64>,
110    #[serde(skip_serializing_if = "Option::is_none")]
111    pub gpu: Option<String>,
112    #[serde(skip_serializing_if = "Option::is_none")]
113    pub group_metadata: Option<BTreeMap<String, JsonValue>>,
114    #[serde(skip_serializing_if = "Option::is_none")]
115    pub hazard_loading: Option<String>,
116    #[serde(skip_serializing_if = "Option::is_none")]
117    pub latent_coordinates: Option<LatentCoordinatesDocument>,
118    #[serde(skip_serializing_if = "Option::is_none")]
119    pub link: Option<String>,
120    #[serde(skip_serializing_if = "Option::is_none")]
121    pub logslope_formula: Option<String>,
122    #[serde(skip_serializing_if = "Option::is_none")]
123    pub negative_binomial_theta: Option<f64>,
124    #[serde(skip_serializing_if = "Option::is_none")]
125    pub noise_formula: Option<String>,
126    #[serde(skip_serializing_if = "Option::is_none")]
127    pub noise_offset: Option<String>,
128    #[serde(skip_serializing_if = "Option::is_none")]
129    pub offset: Option<String>,
130    #[serde(skip_serializing_if = "Option::is_none")]
131    pub outer_max_iter: Option<usize>,
132    #[serde(skip_serializing_if = "Option::is_none")]
133    pub analytic_penalties: Option<AnalyticPenaltiesDocument>,
134    #[serde(skip_serializing_if = "Option::is_none")]
135    pub pilot_subsample_threshold: Option<usize>,
136    #[serde(skip_serializing_if = "Option::is_none")]
137    pub precision_hyperpriors: Option<BTreeMap<String, PrecisionHyperpriorDocument>>,
138    #[serde(skip_serializing_if = "Option::is_none")]
139    /// Whether to precompute the distribution-free conformal substrates (#942
140    /// jackknife+, #1098 exact full-conformal) at fit time and persist them on
141    /// the saved model. Omit to keep the default of precomputing whenever the
142    /// fit is eligible; `false` skips both.
143    ///
144    /// Measured on `y ~ s(x1,k=6) + s(x2,k=6)` (#2633): the two substrates are
145    /// 94% of a saved Gaussian model at n=20,000 (10.2 MB of 10.85 MB) and grow
146    /// linearly with the training rows. Rebuilding both costs ~5.6 ms, 0.3% of
147    /// the fit, and stays under half a second out to p=253. So turning this off
148    /// yields a ~16x smaller model (10.85 MB -> ~0.65 MB at n=20,000).
149    ///
150    /// It is opt-OUT because rebuilding needs the training design AND response
151    /// back, which a saved model deliberately does not carry: a model shipped to
152    /// a host that never sees the training data must keep them or it cannot
153    /// produce a conformal interval at all. Turn it off when the caller retains
154    /// its training data, fits in batch, or never asks for conformal intervals.
155    pub precompute_conformal: Option<bool>,
156    /// Explicit root for cross-process warm starts. Omit to disable on-disk
157    /// persistence. The path is used exactly as supplied; no temp/cache
158    /// discovery or environment fallback is performed.
159    #[serde(skip_serializing_if = "Option::is_none")]
160    pub persistent_warm_start_root: Option<PathBuf>,
161    #[serde(skip_serializing_if = "Option::is_none")]
162    pub ridge_lambda: Option<f64>,
163    #[serde(skip_serializing_if = "Option::is_none")]
164    pub scale_dimensions: Option<bool>,
165    #[serde(skip_serializing_if = "Option::is_none")]
166    pub logslope_time_degree: Option<usize>,
167    /// Number of B-spline basis functions on the `log t` margin of the
168    /// survival marginal-slope log-slope block (gam#2765, gam#2767). Omitted =
169    /// a slope that does not move along follow-up.
170    #[serde(skip_serializing_if = "Option::is_none")]
171    pub logslope_time_k: Option<usize>,
172    #[serde(skip_serializing_if = "Option::is_none")]
173    pub sigma_time_degree: Option<usize>,
174    #[serde(skip_serializing_if = "Option::is_none")]
175    pub sigma_time_k: Option<usize>,
176    #[serde(skip_serializing_if = "Option::is_none")]
177    pub smooth_descriptors: Option<SmoothDescriptorsDocument>,
178    #[serde(skip_serializing_if = "Option::is_none")]
179    pub survival_distribution: Option<String>,
180    #[serde(skip_serializing_if = "Option::is_none")]
181    pub survival_likelihood: Option<String>,
182    /// Explicit centering anchor for the survival baseline time basis, in the
183    /// data's own time units. Omit to let the fit pick it from the likelihood
184    /// mode and the truncation shape of the data — the robust interior median
185    /// exit for marginal-slope and for any genuinely left-truncated dataset
186    /// (#751/#1790), the earliest entry age otherwise.
187    ///
188    /// The CLI's `--survival-time-anchor` declares a conflict with `--request` on
189    /// the premise that this document carries the complete scientific model
190    /// configuration; until #2631 the document had no field for it, so the
191    /// premise was false.
192    #[serde(skip_serializing_if = "Option::is_none")]
193    pub survival_time_anchor: Option<f64>,
194    #[serde(skip_serializing_if = "Option::is_none")]
195    pub threshold_time_degree: Option<usize>,
196    #[serde(skip_serializing_if = "Option::is_none")]
197    pub threshold_time_k: Option<usize>,
198    #[serde(skip_serializing_if = "Option::is_none")]
199    pub time_basis: Option<String>,
200    #[serde(skip_serializing_if = "Option::is_none")]
201    pub time_degree: Option<usize>,
202    #[serde(skip_serializing_if = "Option::is_none")]
203    pub time_num_internal_knots: Option<usize>,
204    #[serde(skip_serializing_if = "Option::is_none")]
205    pub time_smooth_lambda: Option<f64>,
206    /// Container type of the caller's training table (`"pandas"`, `"polars"`,
207    /// `"pyarrow"`, `"numpy"`, ...), passed through opaquely into the saved
208    /// model payload for the predict-time output-container fallback (#394).
209    #[serde(skip_serializing_if = "Option::is_none")]
210    pub training_table_kind: Option<String>,
211    #[serde(skip_serializing_if = "Option::is_none")]
212    pub transformation_normal: Option<bool>,
213    #[serde(skip_serializing_if = "Option::is_none")]
214    pub weights: Option<String>,
215    #[serde(skip_serializing_if = "Option::is_none")]
216    pub z_column: Option<String>,
217}
218
219#[derive(Clone, Debug, Deserialize, PartialEq, Serialize)]
220#[serde(deny_unknown_fields)]
221pub struct PrecisionHyperpriorDocument {
222    pub shape: f64,
223    pub rate: f64,
224}
225
226#[derive(Clone, Debug, Default, Deserialize, PartialEq, Serialize)]
227#[serde(transparent)]
228pub struct LatentCoordinatesDocument(pub BTreeMap<String, LatentCoordinateDocument>);
229
230impl LatentCoordinatesDocument {
231    pub fn to_json_value(&self) -> Result<JsonValue, String> {
232        for (symbol, coordinate) in &self.0 {
233            if symbol.trim().is_empty() {
234                return Err("latent_coordinates keys must be non-empty symbols".to_string());
235            }
236            if coordinate.n == 0 || coordinate.d == 0 {
237                return Err(format!(
238                    "latent_coordinates['{symbol}'] requires positive n and d"
239                ));
240            }
241            if coordinate
242                .name
243                .as_deref()
244                .is_some_and(|name| name.trim().is_empty())
245            {
246                return Err(format!(
247                    "latent_coordinates['{symbol}'].name must be non-empty"
248                ));
249            }
250        }
251        serde_json::to_value(self)
252            .map_err(|error| format!("failed to serialize latent coordinates: {error}"))
253    }
254}
255
256#[derive(Clone, Debug, Deserialize, PartialEq, Serialize)]
257#[serde(deny_unknown_fields)]
258pub struct LatentCoordinateDocument {
259    pub n: usize,
260    pub d: usize,
261    #[serde(default, skip_serializing_if = "Option::is_none")]
262    pub name: Option<String>,
263    #[serde(default, skip_serializing_if = "Option::is_none")]
264    pub init: Option<JsonValue>,
265    #[serde(default, skip_serializing_if = "Option::is_none")]
266    pub manifold: Option<JsonValue>,
267    #[serde(default, skip_serializing_if = "Option::is_none")]
268    pub retraction: Option<JsonValue>,
269    #[serde(default, skip_serializing_if = "Option::is_none")]
270    pub aux_prior: Option<JsonValue>,
271    #[serde(default, skip_serializing_if = "Option::is_none")]
272    pub dim_selection: Option<JsonValue>,
273    #[serde(default, skip_serializing_if = "Option::is_none")]
274    pub aux_outcome: Option<JsonValue>,
275    #[serde(default, skip_serializing_if = "Option::is_none")]
276    pub id_mode: Option<String>,
277}
278
279#[derive(Clone, Debug, Default, Deserialize, PartialEq, Serialize)]
280#[serde(transparent)]
281pub struct AnalyticPenaltiesDocument(pub Vec<JsonValue>);
282
283impl AnalyticPenaltiesDocument {
284    pub fn to_json_value(&self) -> Result<JsonValue, String> {
285        for (index, descriptor) in self.0.iter().enumerate() {
286            let descriptor = descriptor
287                .as_object()
288                .ok_or_else(|| format!("analytic_penalties[{index}] must be an object"))?;
289            if !descriptor.get("target").is_some_and(JsonValue::is_string) {
290                return Err(format!(
291                    "analytic_penalties[{index}].target must be a latent-coordinate name"
292                ));
293            }
294        }
295        serde_json::to_value(self)
296            .map_err(|error| format!("failed to serialize analytic penalties: {error}"))
297    }
298}
299
300#[derive(Clone, Debug, Default, Deserialize, PartialEq, Serialize)]
301#[serde(transparent)]
302pub struct SmoothDescriptorsDocument(pub BTreeMap<String, JsonValue>);
303
304impl SmoothDescriptorsDocument {
305    pub fn to_json_value(&self) -> Result<JsonValue, String> {
306        for (symbol, descriptor) in &self.0 {
307            if symbol.trim().is_empty() {
308                return Err("smooth_descriptors keys must be non-empty symbols".to_string());
309            }
310            if !descriptor.is_object() {
311                return Err(format!("smooth_descriptors['{symbol}'] must be an object"));
312            }
313        }
314        serde_json::to_value(self)
315            .map_err(|error| format!("failed to serialize smooth descriptors: {error}"))
316    }
317}
318
319#[derive(Clone, Debug, Deserialize, PartialEq, Serialize)]
320#[serde(deny_unknown_fields)]
321pub struct CtnStage1Document {
322    pub response_column: String,
323    pub covariate_formula_rhs: String,
324    #[serde(default, skip_serializing_if = "Option::is_none")]
325    pub config: Option<CtnStage1ConfigDocument>,
326    #[serde(default, skip_serializing_if = "Option::is_none")]
327    pub weight_column: Option<String>,
328    #[serde(default, skip_serializing_if = "Option::is_none")]
329    pub offset_column: Option<String>,
330}
331
332#[derive(Clone, Debug, Default, Deserialize, PartialEq, Serialize)]
333#[serde(deny_unknown_fields)]
334pub struct CtnStage1ConfigDocument {
335    #[serde(default, skip_serializing_if = "Option::is_none")]
336    pub response_degree: Option<usize>,
337    #[serde(default, skip_serializing_if = "Option::is_none")]
338    pub response_num_internal_knots: Option<usize>,
339    #[serde(default, skip_serializing_if = "Option::is_none")]
340    pub response_penalty_order: Option<usize>,
341    #[serde(default, skip_serializing_if = "Option::is_none")]
342    pub response_extra_penalty_orders: Option<Vec<usize>>,
343    #[serde(default, skip_serializing_if = "Option::is_none")]
344    pub double_penalty: Option<bool>,
345}
346
347#[cfg(test)]
348mod tests {
349    use super::*;
350    use serde_json::json;
351
352    #[test]
353    fn canonical_document_round_trips_identically() {
354        let document = FitRequestDocument::new(
355            "y ~ duchon(x)",
356            FitRequestConfigDocument {
357                ctn_stage1: Some(CtnStage1Document {
358                    response_column: "dose".to_string(),
359                    covariate_formula_rhs: "s(age)".to_string(),
360                    config: Some(CtnStage1ConfigDocument {
361                        response_degree: Some(4),
362                        response_penalty_order: Some(2),
363                        ..CtnStage1ConfigDocument::default()
364                    }),
365                    weight_column: Some("case_weight".to_string()),
366                    offset_column: None,
367                }),
368                latent_coordinates: Some(
369                    serde_json::from_value(json!({
370                        "x": {"d": 2, "init": "pca", "n": 12, "name": "x"}
371                    }))
372                    .unwrap(),
373                ),
374                analytic_penalties: Some(AnalyticPenaltiesDocument(vec![json!(
375                    {"kind": "orthogonality", "target": "x", "weight": 1.0}
376                )])),
377                precision_hyperpriors: Some(BTreeMap::from([(
378                    "s(x):roughness".to_string(),
379                    PrecisionHyperpriorDocument {
380                        shape: 2.0,
381                        rate: 0.5,
382                    },
383                )])),
384                persistent_warm_start_root: Some(PathBuf::from("warm-start-fixture")),
385                smooth_descriptors: Some(
386                    serde_json::from_value(json!({
387                        "x": {"centers": 8, "kind": "duchon", "vars": ["x"]}
388                    }))
389                    .unwrap(),
390                ),
391                ..FitRequestConfigDocument::default()
392            },
393        )
394        .unwrap();
395
396        let encoded = document.to_canonical_json().unwrap();
397        let decoded = FitRequestDocument::from_json(&encoded).unwrap();
398        assert_eq!(decoded, document);
399        assert_eq!(decoded.to_canonical_json().unwrap(), encoded);
400    }
401
402    #[test]
403    fn parser_rejects_another_schema_or_version() {
404        let wrong_schema = r#"{"schema":"other","schema_version":1,"formula":"y ~ x","config":{}}"#;
405        assert!(
406            FitRequestDocument::from_json(wrong_schema)
407                .unwrap_err()
408                .contains("schema must be")
409        );
410
411        let wrong_version =
412            r#"{"schema":"gam.fit-request","schema_version":2,"formula":"y ~ x","config":{}}"#;
413        assert!(
414            FitRequestDocument::from_json(wrong_version)
415                .unwrap_err()
416                .contains("unsupported fit request schema_version")
417        );
418    }
419
420    #[test]
421    fn parser_rejects_removed_topology_selector_descriptor() {
422        let legacy = r#"{
423            "schema": "gam.fit-request",
424            "schema_version": 1,
425            "formula": "y ~ x",
426            "config": {"topology_auto_selector": {"candidates": ["circle"]}}
427        }"#;
428        let error = FitRequestDocument::from_json(legacy)
429            .expect_err("removed no-op selector field must not be silently ignored");
430        assert!(
431            error.contains("unknown field `topology_auto_selector`"),
432            "{error}"
433        );
434    }
435}