1use std::collections::{BTreeMap, HashMap};
2
3use pest::Parser;
4use pest::iterators::Pair;
5use pest_derive::Parser;
6
7use crate::smooth::BoundedCoefficientPriorSpec;
8use crate::term_builder::{MARGINAL_SLOPE_Z_ALIAS, marginal_slope_z_alias_is_live};
9use gam_problem::types::{
10 InverseLink, LikelihoodSpec, LinkComponent, LinkFunction, StandardLink, WigglePenaltyConfig,
11};
12
13#[derive(Parser)]
14#[grammar_inline = r#"
15WHITESPACE = _{ " " | "\t" | NEWLINE }
16
17top_function_call = { SOI ~ function_call ~ EOI }
18top_expr = { SOI ~ expr ~ EOI }
19formula = { SOI ~ expr ~ "~" ~ rhs ~ EOI }
20rhs = { term ~ ("+" ~ term)* }
21term = { expr }
22
23expr = { sum }
24sum = { product ~ (add_op ~ product)* }
25add_op = { "+" | "-" }
26product = { interact ~ (mul_op ~ interact)* }
27mul_op = { "*" | "/" }
28interact = { power ~ (interact_op ~ power)* }
29interact_op = { ":" }
30power = { unary ~ (pow_op ~ unary)* }
31pow_op = { "^" }
32unary = { unary_op* ~ primary }
33unary_op = _{ "+" | "-" }
34
35primary = { function_call | list_lit | tuple_lit | ident | number | string_lit | "(" ~ expr ~ ")" }
36list_lit = @{ "[" ~ (!"]" ~ ANY)* ~ "]" }
37tuple_lit = @{ "(" ~ (!("," | ")") ~ ANY)+ ~ "," ~ (!")" ~ ANY)* ~ ")" }
38function_call = { ident ~ "(" ~ arg_list? ~ ")" }
39arg_list = { arg ~ ("," ~ arg)* }
40arg = { named_arg | expr }
41named_arg = { ident ~ "=" ~ expr }
42
43ident = @{ ident_start ~ ident_continue* }
44ident_start = _{ ASCII_ALPHA | "_" }
45ident_continue = _{ ASCII_ALPHANUMERIC | "_" | "." }
46
47number = @{
48 "-"?
49 ~ (ASCII_DIGIT+ ~ ("." ~ ASCII_DIGIT*)? | "." ~ ASCII_DIGIT+)
50 ~ (("e" | "E") ~ ("+" | "-")? ~ ASCII_DIGIT+)?
51}
52
53string_lit = @{ "\"" ~ (!"\"" ~ ANY)* ~ "\"" | "'" ~ (!"'" ~ ANY)* ~ "'" }
54"#]
55struct FormulaParser;
56
57#[derive(Clone, Debug, PartialEq, Eq)]
58pub struct FormulaDslParse {
59 pub response_expr: String,
60 pub rhs_terms: Vec<String>,
61}
62
63#[derive(Clone, Debug, PartialEq, Eq)]
64pub enum CallArgSpec {
65 Positional(String),
66 Named { key: String, value: String },
67}
68
69#[derive(Clone, Debug, PartialEq, Eq)]
70pub struct FunctionCallSpec {
71 pub name: String,
72 pub args: Vec<CallArgSpec>,
73}
74
75#[derive(Clone, Debug, PartialEq, Eq)]
86pub enum FormulaDslError {
87 ParseError { reason: String },
91 UnknownIdentifier { reason: String },
95 InvalidArgument { reason: String },
98 IncompatibleTerm { reason: String },
102 MalformedConfig { reason: String },
106}
107
108gam_linalg::impl_reason_error_boilerplate! {
109 FormulaDslError {
110 ParseError,
111 UnknownIdentifier,
112 InvalidArgument,
113 IncompatibleTerm,
114 MalformedConfig,
115 }
116}
117
118impl From<String> for FormulaDslError {
124 fn from(reason: String) -> Self {
125 FormulaDslError::ParseError { reason }
126 }
127}
128
129pub fn parse_formula_dsl(formula: &str) -> Result<FormulaDslParse, String> {
130 validate_balanced_delimiters(formula, "invalid formula syntax")?;
131 let mut parsed =
132 FormulaParser::parse(Rule::formula, formula).map_err(|e| FormulaDslError::ParseError {
133 reason: format!("invalid formula syntax: {e}"),
134 })?;
135 let formula_pair = parsed.next().ok_or_else(|| FormulaDslError::ParseError {
136 reason: "invalid formula syntax: empty parse".to_string(),
137 })?;
138
139 let mut response_expr: Option<String> = None;
140 let mut rhs_terms: Option<Vec<String>> = None;
141
142 for part in formula_pair.into_inner() {
143 if part.as_rule() == Rule::rhs {
144 rhs_terms = Some(extract_rhs_terms(part)?);
145 } else if part.as_rule() == Rule::expr && response_expr.is_none() {
146 response_expr = Some(part.as_str().trim().to_string());
147 }
148 }
149
150 let response_expr = response_expr.ok_or_else(|| FormulaDslError::ParseError {
151 reason: "invalid formula: missing response expression".to_string(),
152 })?;
153 let rhs_terms = rhs_terms.ok_or_else(|| FormulaDslError::ParseError {
154 reason: "invalid formula: missing RHS terms".to_string(),
155 })?;
156 if rhs_terms.is_empty() {
157 return Err(FormulaDslError::ParseError {
158 reason: "formula has no usable terms".to_string(),
159 }
160 .into());
161 }
162
163 Ok(FormulaDslParse {
164 response_expr,
165 rhs_terms,
166 })
167}
168
169fn delimiter_balance_error(prefix: &str) -> String {
170 format!("{prefix}: unbalanced parentheses or quotes")
171}
172
173fn validate_balanced_delimiters(input: &str, prefix: &str) -> Result<(), String> {
177 let mut stack = Vec::<char>::new();
178 let mut in_single = false;
179 let mut in_double = false;
180
181 for ch in input.chars() {
182 let quoted = in_single || in_double;
183 if ch == '\'' && !in_double {
184 in_single = !in_single;
185 } else if ch == '"' && !in_single {
186 in_double = !in_double;
187 } else if !quoted && matches!(ch, '(' | '[' | '{') {
188 stack.push(ch);
189 } else if !quoted && matches!(ch, ')' | ']' | '}') {
190 let expected = match ch {
191 ')' => '(',
192 ']' => '[',
193 _ => '{',
195 };
196 if stack.pop() != Some(expected) {
197 return Err(FormulaDslError::ParseError {
198 reason: delimiter_balance_error(prefix),
199 }
200 .into());
201 }
202 }
203 }
204
205 if in_single || in_double || !stack.is_empty() {
206 return Err(FormulaDslError::ParseError {
207 reason: delimiter_balance_error(prefix),
208 }
209 .into());
210 }
211 Ok(())
212}
213
214fn extract_rhs_terms(rhs: Pair<'_, Rule>) -> Result<Vec<String>, String> {
215 let mut out = Vec::new();
216 let mut depth = 0_i32;
217 let mut in_single = false;
218 let mut in_double = false;
219 let mut start = 0_usize;
220 let mut last_significant: Option<char> = None;
233 let text = rhs.as_str();
234 let bytes = text.as_bytes();
235 for (idx, &b) in bytes.iter().enumerate() {
236 let ch = b as char;
237 let quoted = in_single || in_double;
238 if ch == '\'' && !in_double {
239 in_single = !in_single;
240 } else if ch == '"' && !in_single {
241 in_double = !in_double;
242 } else if !quoted && matches!(ch, '(' | '[' | '{') {
243 depth += 1;
244 } else if !quoted && matches!(ch, ')' | ']' | '}') && depth > 0 {
245 depth -= 1;
246 } else if ch == '+'
247 && !quoted
248 && depth == 0
249 && !matches!(
250 last_significant,
251 None | Some(':' | '*' | '/' | '^' | '+' | '-')
252 )
253 {
254 let term = text[start..idx].trim();
255 if term.is_empty() {
256 return Err(FormulaDslError::ParseError {
257 reason: "formula RHS contains an empty term".to_string(),
258 }
259 .into());
260 }
261 out.push(term.to_string());
262 start = idx + 1;
263 }
264 if !ch.is_ascii_whitespace() {
265 last_significant = Some(ch);
266 }
267 }
268 if in_single || in_double || depth != 0 {
269 return Err(FormulaDslError::ParseError {
270 reason: "formula RHS has unbalanced quotes or parentheses".to_string(),
271 }
272 .into());
273 }
274 let tail = text[start..].trim();
275 if tail.is_empty() {
276 return Err(FormulaDslError::ParseError {
277 reason: "formula RHS contains an empty term".to_string(),
278 }
279 .into());
280 }
281 out.push(tail.to_string());
282 Ok(out)
283}
284
285type WrAtomList = Vec<String>;
310
311fn expand_wr_term(raw: &str) -> Result<Vec<WrAtomList>, String> {
312 let mut parsed = FormulaParser::parse(Rule::top_expr, raw).map_err(|e| {
313 FormulaDslError::ParseError {
314 reason: format!("invalid term syntax in `{raw}`: {e}"),
315 }
316 .to_string()
317 })?;
318 let top = parsed.next().ok_or_else(|| {
319 FormulaDslError::ParseError {
320 reason: format!("invalid term syntax in `{raw}`: empty parse"),
321 }
322 .to_string()
323 })?;
324 let expr = top
325 .into_inner()
326 .find(|p| p.as_rule() == Rule::expr)
327 .ok_or_else(|| {
328 FormulaDslError::ParseError {
329 reason: format!("invalid term syntax in `{raw}`: missing expr"),
330 }
331 .to_string()
332 })?;
333 let interactions = expand_expr(expr, raw)?;
334 let normalized: Vec<WrAtomList> = interactions
335 .into_iter()
336 .map(normalize_interaction)
337 .collect();
338 let mut seen = std::collections::BTreeSet::<Vec<String>>::new();
341 let mut out = Vec::<WrAtomList>::new();
342 for term in normalized {
343 let key = term.clone();
344 if seen.insert(key) {
345 out.push(term);
346 }
347 }
348 Ok(out)
349}
350
351fn normalize_interaction(mut atoms: WrAtomList) -> WrAtomList {
352 atoms.sort();
353 atoms.dedup();
354 atoms
355}
356
357fn expand_expr(pair: Pair<'_, Rule>, raw: &str) -> Result<Vec<WrAtomList>, String> {
358 match pair.as_rule() {
359 Rule::expr => {
360 let inner = pair.into_inner().next().ok_or_else(|| {
361 FormulaDslError::ParseError {
362 reason: format!("invalid term syntax in `{raw}`: empty expr"),
363 }
364 .to_string()
365 })?;
366 expand_expr(inner, raw)
367 }
368 Rule::sum => {
369 let mut iter = pair.into_inner();
370 let first = iter.next().ok_or_else(|| {
371 FormulaDslError::ParseError {
372 reason: format!("invalid term syntax in `{raw}`: empty sum"),
373 }
374 .to_string()
375 })?;
376 let mut acc = expand_expr(first, raw)?;
377 while let Some(op) = iter.next() {
378 if op.as_rule() != Rule::add_op {
379 return Err(FormulaDslError::ParseError {
380 reason: format!(
381 "invalid term syntax in `{raw}`: expected add operator, got `{:?}`",
382 op.as_rule()
383 ),
384 }
385 .into());
386 }
387 let op_str = op.as_str().trim();
388 let operand = iter.next().ok_or_else(|| {
389 FormulaDslError::ParseError {
390 reason: format!("invalid term syntax in `{raw}`: dangling `{op_str}`"),
391 }
392 .to_string()
393 })?;
394 if op_str == "-" {
395 return Err(FormulaDslError::IncompatibleTerm {
396 reason: format!(
397 "binary `-` is not supported inside a formula term in `{raw}` \
398 (use multiple `+` terms or drop the unwanted predictor explicitly)"
399 ),
400 }
401 .into());
402 }
403 let mut rhs = expand_expr(operand, raw)?;
404 acc.append(&mut rhs);
405 }
406 Ok(acc)
407 }
408 Rule::product => {
409 let mut iter = pair.into_inner();
410 let first = iter.next().ok_or_else(|| {
411 FormulaDslError::ParseError {
412 reason: format!("invalid term syntax in `{raw}`: empty product"),
413 }
414 .to_string()
415 })?;
416 let mut acc = expand_expr(first, raw)?;
417 while let Some(op) = iter.next() {
418 if op.as_rule() != Rule::mul_op {
419 return Err(FormulaDslError::ParseError {
420 reason: format!(
421 "invalid term syntax in `{raw}`: expected `*` or `/`, got `{:?}`",
422 op.as_rule()
423 ),
424 }
425 .into());
426 }
427 let op_str = op.as_str().trim();
428 let operand = iter.next().ok_or_else(|| {
429 FormulaDslError::ParseError {
430 reason: format!("invalid term syntax in `{raw}`: dangling `{op_str}`"),
431 }
432 .to_string()
433 })?;
434 let rhs = expand_expr(operand, raw)?;
435 acc = match op_str {
436 "*" => wr_cross(acc, rhs),
437 "/" => wr_nest(acc, rhs),
438 other => {
439 return Err(FormulaDslError::ParseError {
440 reason: format!(
441 "invalid term syntax in `{raw}`: unrecognized mul operator `{other}`"
442 ),
443 }
444 .into());
445 }
446 };
447 }
448 Ok(acc)
449 }
450 Rule::interact => {
451 let mut iter = pair.into_inner();
452 let first = iter.next().ok_or_else(|| {
453 FormulaDslError::ParseError {
454 reason: format!("invalid term syntax in `{raw}`: empty interact"),
455 }
456 .to_string()
457 })?;
458 let mut acc = expand_expr(first, raw)?;
459 while let Some(op) = iter.next() {
460 if op.as_rule() != Rule::interact_op {
461 return Err(FormulaDslError::ParseError {
462 reason: format!(
463 "invalid term syntax in `{raw}`: expected `:`, got `{:?}`",
464 op.as_rule()
465 ),
466 }
467 .into());
468 }
469 let operand = iter.next().ok_or_else(|| {
470 FormulaDslError::ParseError {
471 reason: format!("invalid term syntax in `{raw}`: dangling `:`"),
472 }
473 .to_string()
474 })?;
475 let rhs = expand_expr(operand, raw)?;
476 acc = wr_interact(acc, rhs, raw)?;
477 }
478 Ok(acc)
479 }
480 Rule::power => {
481 let mut iter = pair.into_inner();
482 let first = iter.next().ok_or_else(|| {
483 FormulaDslError::ParseError {
484 reason: format!("invalid term syntax in `{raw}`: empty power"),
485 }
486 .to_string()
487 })?;
488 let base = expand_expr(first, raw)?;
489 let Some(op) = iter.next() else {
490 return Ok(base);
491 };
492 if op.as_rule() != Rule::pow_op {
493 return Err(FormulaDslError::ParseError {
494 reason: format!(
495 "invalid term syntax in `{raw}`: expected `^`, got `{:?}`",
496 op.as_rule()
497 ),
498 }
499 .into());
500 }
501 let exponent_pair = iter.next().ok_or_else(|| {
502 FormulaDslError::ParseError {
503 reason: format!("invalid term syntax in `{raw}`: dangling `^`"),
504 }
505 .to_string()
506 })?;
507 let exp_text = exponent_pair.as_str().trim();
508 let n: usize = exp_text.parse().map_err(|_| {
509 FormulaDslError::ParseError {
510 reason: format!(
511 "invalid term syntax in `{raw}`: `^` exponent must be a positive integer, got `{exp_text}`"
512 ),
513 }
514 .to_string()
515 })?;
516 if n == 0 {
517 return Err(FormulaDslError::ParseError {
518 reason: format!(
519 "invalid term syntax in `{raw}`: `^0` is not a meaningful formula expansion"
520 ),
521 }
522 .into());
523 }
524 if let Some(extra_op) = iter.next() {
525 let extra = extra_op.as_str().trim();
526 return Err(FormulaDslError::ParseError {
527 reason: format!(
528 "invalid term syntax in `{raw}`: chained `^` operators are not supported; \
529 use one positive integer exponent, got another `{extra}`"
530 ),
531 }
532 .into());
533 }
534 Ok(wr_power(base, n))
535 }
536 Rule::unary => {
537 let unary_start = pair.as_span().start();
538 for inner in pair.into_inner() {
539 if inner.as_rule() == Rule::primary {
540 let prefix_len = inner.as_span().start().saturating_sub(unary_start);
541 if prefix_len > 0 {
542 let prefix = &raw[unary_start..inner.as_span().start()];
543 if prefix.chars().any(|ch| matches!(ch, '+' | '-')) {
544 return Err(FormulaDslError::IncompatibleTerm {
545 reason: format!(
546 "unary `+`/`-` is not supported inside a formula term in `{raw}`"
547 ),
548 }
549 .into());
550 }
551 }
552 return expand_expr(inner, raw);
553 }
554 }
555 Err(FormulaDslError::ParseError {
556 reason: format!("invalid term syntax in `{raw}`: empty unary"),
557 }
558 .into())
559 }
560 Rule::primary => {
561 let span = pair.as_str().trim().to_string();
562 let inner = pair.into_inner().next();
563 match inner {
564 Some(child) if child.as_rule() == Rule::expr => expand_expr(child, raw),
565 Some(child) => Ok(vec![vec![child.as_str().trim().to_string()]]),
566 None => Ok(vec![vec![span]]),
567 }
568 }
569 _ => Err(FormulaDslError::ParseError {
570 reason: format!(
571 "invalid term syntax in `{raw}`: unexpected node `{:?}`",
572 pair.as_rule()
573 ),
574 }
575 .into()),
576 }
577}
578
579fn wr_cross(left: Vec<WrAtomList>, right: Vec<WrAtomList>) -> Vec<WrAtomList> {
580 let mut out = Vec::with_capacity(left.len() + right.len() + left.len() * right.len());
582 out.extend(left.iter().cloned());
583 out.extend(right.iter().cloned());
584 for l in &left {
585 for r in &right {
586 let mut merged: WrAtomList = l.iter().cloned().chain(r.iter().cloned()).collect();
587 merged.sort();
588 merged.dedup();
589 out.push(merged);
590 }
591 }
592 out
593}
594
595fn wr_nest(left: Vec<WrAtomList>, right: Vec<WrAtomList>) -> Vec<WrAtomList> {
596 let mut left_atoms: WrAtomList = left.iter().flatten().cloned().collect();
604 left_atoms.sort();
605 left_atoms.dedup();
606
607 let mut out = Vec::with_capacity(left.len() + right.len());
608 out.extend(left.iter().cloned());
609 for r in &right {
610 let mut merged: WrAtomList = left_atoms
611 .iter()
612 .cloned()
613 .chain(r.iter().cloned())
614 .collect();
615 merged.sort();
616 merged.dedup();
617 out.push(merged);
618 }
619 out
620}
621
622fn wr_interact(
623 left: Vec<WrAtomList>,
624 right: Vec<WrAtomList>,
625 raw: &str,
626) -> Result<Vec<WrAtomList>, String> {
627 let mut out = Vec::with_capacity(left.len() * right.len());
629 for l in &left {
630 for r in &right {
631 let combined: WrAtomList = l.iter().cloned().chain(r.iter().cloned()).collect();
632 for atom in &combined {
636 if atom.contains('(') {
637 return Err(FormulaDslError::IncompatibleTerm {
638 reason: format!(
639 "interaction operator `:` with function-call atom is not supported in `{raw}`. \
640 Use te(...) for smooth interactions or group()/factor() with a separate \
641 interaction strategy for categorical effects."
642 ),
643 }
644 .into());
645 }
646 }
647 let mut merged = combined;
648 merged.sort();
649 merged.dedup();
650 out.push(merged);
651 }
652 }
653 Ok(out)
654}
655
656fn wr_power(base: Vec<WrAtomList>, n: usize) -> Vec<WrAtomList> {
657 if base.is_empty() {
662 return Vec::new();
663 }
664 let m = base.len();
665 let mut out = Vec::<WrAtomList>::new();
666 let max_size = n.min(m);
667 for size in 1..=max_size {
668 let mut indices: Vec<usize> = (0..size).collect();
670 loop {
671 let mut merged = WrAtomList::new();
672 for &i in &indices {
673 merged.extend(base[i].iter().cloned());
674 }
675 merged.sort();
676 merged.dedup();
677 out.push(merged);
678 let mut k = size;
680 while k > 0 {
681 k -= 1;
682 if indices[k] != k + m - size {
683 indices[k] += 1;
684 for j in (k + 1)..size {
685 indices[j] = indices[j - 1] + 1;
686 }
687 break;
688 }
689 if k == 0 {
690 k = usize::MAX;
691 break;
692 }
693 }
694 if k == usize::MAX {
695 break;
696 }
697 }
698 }
699 out
700}
701
702fn is_exact_ident(raw: &str) -> bool {
703 let mut chars = raw.chars();
704 let Some(first) = chars.next() else {
705 return false;
706 };
707 if !first.is_ascii_alphabetic() && first != '_' {
708 return false;
709 }
710 chars.all(|ch| ch.is_ascii_alphanumeric() || ch == '_' || ch == '.')
711}
712
713pub fn parse_function_call(input: &str) -> Result<FunctionCallSpec, String> {
714 validate_balanced_delimiters(input, "invalid function call syntax")?;
715 let mut parsed = FormulaParser::parse(Rule::top_function_call, input).map_err(|e| {
716 FormulaDslError::ParseError {
717 reason: format!("invalid function call syntax: {e}"),
718 }
719 })?;
720 let top = parsed.next().ok_or_else(|| FormulaDslError::ParseError {
721 reason: "invalid function call syntax: empty parse".to_string(),
722 })?;
723 let call = top
724 .into_inner()
725 .find(|p| p.as_rule() == Rule::function_call)
726 .ok_or_else(|| FormulaDslError::ParseError {
727 reason: "invalid function call syntax: missing call".to_string(),
728 })?;
729 parse_call_pair(call)
730}
731
732fn parse_call_pair(call: Pair<'_, Rule>) -> Result<FunctionCallSpec, String> {
733 let mut name: Option<String> = None;
734 let mut args = Vec::<CallArgSpec>::new();
735 for part in call.into_inner() {
736 if part.as_rule() == Rule::ident {
737 if name.is_none() {
738 name = Some(part.as_str().trim().to_string());
739 }
740 } else if part.as_rule() == Rule::arg_list {
741 for a in part.into_inner() {
742 if a.as_rule() != Rule::arg {
743 continue;
744 }
745 let mut a_inner = a.into_inner();
746 let Some(first) = a_inner.next() else {
747 continue;
748 };
749 if first.as_rule() == Rule::named_arg {
750 let mut ni = first.into_inner();
751 let key = ni
752 .next()
753 .ok_or_else(|| FormulaDslError::ParseError {
754 reason: "invalid named argument key".to_string(),
755 })?
756 .as_str()
757 .trim()
758 .to_ascii_lowercase();
759 let value = ni
760 .next()
761 .ok_or_else(|| FormulaDslError::ParseError {
762 reason: "invalid named argument value".to_string(),
763 })?
764 .as_str()
765 .trim()
766 .to_string();
767 args.push(CallArgSpec::Named { key, value });
768 } else if first.as_rule() == Rule::expr {
769 args.push(CallArgSpec::Positional(first.as_str().trim().to_string()));
770 }
771 }
772 }
773 }
774 let name = name.ok_or_else(|| FormulaDslError::ParseError {
775 reason: "invalid function call: missing name".to_string(),
776 })?;
777 Ok(FunctionCallSpec { name, args })
778}
779
780#[cfg(test)]
781mod tests {
782 use super::{
783 CallArgSpec, ParsedTerm, parse_formula, parse_formula_dsl, parse_function_call,
784 parse_linkwiggle_formulaspec, parsed_term_column_names, parsed_terms_reference_column,
785 validate_marginal_slope_z_alias_exclusion, validate_marginal_slope_z_column_exclusion,
786 };
787 use std::collections::{BTreeMap, BTreeSet, HashMap};
788
789 #[test]
790 fn parsed_term_column_names_includes_by_smooth_grouping_variable() {
791 let parsed =
799 parse_formula("y ~ s(x, by=g) + z + a:b").expect("formula with a by= smooth parses");
800 let mut cols = BTreeSet::<String>::new();
801 parsed_term_column_names(&parsed.terms, &mut cols);
802 for expected in ["x", "g", "z", "a", "b"] {
803 assert!(
804 cols.contains(expected),
805 "parsed_term_column_names dropped '{expected}'; got {cols:?}"
806 );
807 }
808 assert!(
810 !cols.contains("y"),
811 "response leaked into term columns: {cols:?}"
812 );
813 }
814
815 #[test]
816 fn linkwiggle_parser_does_not_bake_in_cubic_only_restriction() {
817 for deg in [2usize, 4, 5, 10] {
827 let mut options = BTreeMap::new();
828 options.insert("degree".to_string(), deg.to_string());
829 options.insert("internal_knots".to_string(), "3".to_string());
830 let raw = format!("timewiggle(degree={deg}, internal_knots=3)");
831 let spec = parse_linkwiggle_formulaspec(&options, &raw)
832 .expect("non-cubic wiggle degree must parse at the shared layer");
833 assert_eq!(
834 spec.degree, deg,
835 "parser must carry the requested degree through verbatim"
836 );
837 }
838
839 let mut zero = BTreeMap::new();
842 zero.insert("degree".to_string(), "0".to_string());
843 zero.insert("internal_knots".to_string(), "3".to_string());
844 let err = parse_linkwiggle_formulaspec(&zero, "linkwiggle(degree=0, internal_knots=3)")
845 .expect_err("degree=0 must be rejected");
846 assert!(
847 err.contains("degree >= 1"),
848 "error should state the positive-degree lower bound, got: {err}"
849 );
850 }
851
852 #[test]
853 fn parses_nested_formula_terms() {
854 let parsed =
855 parse_formula_dsl("log(y) ~ x1 + s(log(x2 + 1), bs=\"tps\", k=10) + te(x3, x4)")
856 .expect("parse");
857 assert_eq!(parsed.response_expr, "log(y)");
858 assert_eq!(parsed.rhs_terms.len(), 3);
859 assert_eq!(parsed.rhs_terms[0], "x1");
860 assert_eq!(parsed.rhs_terms[1], "s(log(x2 + 1), bs=\"tps\", k=10)");
861 assert_eq!(parsed.rhs_terms[2], "te(x3, x4)");
862 }
863
864 #[test]
865 fn parses_cyclic_formula_aliases() {
866 let parsed = parse_formula("y ~ cyclic(theta, period_start=0, period_end=6.283)")
867 .expect("parse cyclic formula");
868 match &parsed.terms[0] {
869 super::ParsedTerm::Smooth { vars, options, .. } => {
870 assert_eq!(vars, &vec!["theta".to_string()]);
871 assert_eq!(options.get("type").map(String::as_str), Some("cyclic"));
872 assert_eq!(options.get("period_start").map(String::as_str), Some("0"));
873 }
874 other => panic!("expected cyclic smooth term, got {other:?}"),
875 }
876 }
877
878 #[test]
879 fn sphere_aliases_all_dispatch_to_intrinsic_s2_basis() {
880 for alias in ["sphere", "sos", "spherical", "s2"] {
887 let parsed = parse_formula(&format!("y ~ {alias}(lat, lon)"))
888 .unwrap_or_else(|e| panic!("parse {alias}: {e}"));
889 match &parsed.terms[0] {
890 super::ParsedTerm::Smooth { vars, options, .. } => {
891 assert_eq!(
892 vars,
893 &vec!["lat".to_string(), "lon".to_string()],
894 "{alias} should keep (lat, lon) as its variables"
895 );
896 assert_eq!(
897 options.get("type").map(String::as_str),
898 Some("sphere"),
899 "{alias} must dispatch to the intrinsic sphere basis (type=sphere)"
900 );
901 }
902 other => panic!("expected sphere smooth term for {alias}, got {other:?}"),
903 }
904 }
905 }
906
907 #[test]
908 fn parses_function_callwithnamed_and_positional_args() {
909 let call = parse_function_call("s(log(x + 1), type=\"duchon\", centers=12)").expect("call");
910 assert_eq!(call.name, "s");
911 assert_eq!(call.args.len(), 3);
912 assert_eq!(
913 call.args[0],
914 CallArgSpec::Positional("log(x + 1)".to_string())
915 );
916 assert_eq!(
917 call.args[1],
918 CallArgSpec::Named {
919 key: "type".to_string(),
920 value: "\"duchon\"".to_string()
921 }
922 );
923 }
924
925 #[test]
926 fn parses_tensor_boundary_list_options() {
927 let call = parse_function_call(
928 "te(day_of_week, hour, boundary=['periodic', 'periodic'], period=[7, 24])",
929 )
930 .expect("call");
931 assert_eq!(call.name, "te");
932 assert_eq!(call.args.len(), 4);
933 assert_eq!(
934 call.args[2],
935 CallArgSpec::Named {
936 key: "boundary".to_string(),
937 value: "['periodic', 'periodic']".to_string(),
938 }
939 );
940 }
941
942 #[test]
943 fn parse_formula_dsl_reports_unbalanced_parentheses() {
944 let err = parse_formula_dsl("y ~ s(x, k=10").expect_err("expected parse failure");
945 assert!(err.contains("unbalanced parentheses"));
946 }
947
948 #[test]
949 fn parse_function_call_reports_unbalanced_parentheses() {
950 let err = parse_function_call("s(x, k=10").expect_err("expected parse failure");
951 assert!(err.contains("unbalanced parentheses"));
952 }
953
954 #[test]
955 fn parse_formula_accepts_tuple_smooth_options() {
956 let parsed = parse_formula("z ~ te(x, y, k=(20, 20))")
957 .expect("tuple-valued smooth option should parse");
958 assert_eq!(parsed.terms.len(), 1);
959
960 let dsl = parse_formula_dsl("z ~ te(x, y, k=(20, 20))")
961 .expect("tuple-valued smooth option should parse in the DSL layer");
962 assert_eq!(dsl.rhs_terms, vec!["te(x, y, k=(20, 20))"]);
963
964 let call = parse_function_call("te(x, y, k=(20, 20))")
965 .expect("tuple-valued smooth option should parse as a function call");
966 assert_eq!(
967 call.args[2],
968 CallArgSpec::Named {
969 key: "k".to_string(),
970 value: "(20, 20)".to_string(),
971 }
972 );
973 }
974
975 #[test]
976 fn parse_formula_rejects_unsupported_top_level_rhs_expressions() {
977 for formula in ["y ~ x - z", "y ~ -x", "y ~ (x)", "y ~ x - 1"] {
986 let err = parse_formula(formula).expect_err("expected formula parse failure");
987 assert!(err.to_string().contains("unsupported top-level RHS term"));
988 }
989 }
990
991 #[test]
997 fn parse_formula_supports_wr_slash_nesting() {
998 let parsed = parse_formula("y ~ x / z").expect("`/` is supported as WR nesting");
999 assert_eq!(parsed.response, "y");
1000 assert_eq!(parsed.terms.len(), 2);
1001 let names: Vec<String> = parsed
1002 .terms
1003 .iter()
1004 .map(|t| match t {
1005 ParsedTerm::Linear { name, .. } => format!("Linear({name})"),
1006 ParsedTerm::Interaction { vars, .. } => {
1007 format!("Interaction({})", vars.join(":"))
1008 }
1009 other => format!("Other({other:?})"),
1010 })
1011 .collect();
1012 assert_eq!(
1013 names,
1014 vec!["Linear(x)".to_string(), "Interaction(x:z)".to_string()]
1015 );
1016 }
1017
1018 fn wr_term_labels(formula: &str) -> Vec<String> {
1022 let parsed = parse_formula(formula).unwrap_or_else(|e| panic!("parse {formula}: {e}"));
1023 parsed
1024 .terms
1025 .iter()
1026 .map(|t| match t {
1027 ParsedTerm::Linear { name, .. } => name.clone(),
1028 ParsedTerm::Interaction { vars, .. } => vars.join(":"),
1029 other => format!("Other({other:?})"),
1030 })
1031 .collect()
1032 }
1033
1034 #[test]
1041 fn parse_formula_chained_wr_nesting_is_hierarchical() {
1042 assert_eq!(
1044 wr_term_labels("y ~ a/b/c"),
1045 vec!["a".to_string(), "a:b".to_string(), "a:b:c".to_string()],
1046 "a/b/c must nest hierarchically with no spurious a:c"
1047 );
1048 assert_eq!(
1050 wr_term_labels("y ~ a*b/c"),
1051 vec![
1052 "a".to_string(),
1053 "b".to_string(),
1054 "a:b".to_string(),
1055 "a:b:c".to_string()
1056 ],
1057 "a*b/c nests c within the whole a*b group"
1058 );
1059 assert_eq!(
1061 wr_term_labels("y ~ x/z"),
1062 vec!["x".to_string(), "x:z".to_string()],
1063 );
1064 assert_eq!(
1066 wr_term_labels("y ~ a/b/c/d"),
1067 vec![
1068 "a".to_string(),
1069 "a:b".to_string(),
1070 "a:b:c".to_string(),
1071 "a:b:c:d".to_string()
1072 ],
1073 );
1074 }
1075
1076 #[test]
1081 fn parse_formula_supports_wr_star_crossing() {
1082 let parsed = parse_formula("y ~ x * z").expect("`*` is supported as WR crossing");
1083 assert_eq!(parsed.response, "y");
1084 assert_eq!(parsed.terms.len(), 3);
1085 let names: Vec<String> = parsed
1086 .terms
1087 .iter()
1088 .map(|t| match t {
1089 ParsedTerm::Linear { name, .. } => format!("Linear({name})"),
1090 ParsedTerm::Interaction { vars, .. } => {
1091 format!("Interaction({})", vars.join(":"))
1092 }
1093 other => format!("Other({other:?})"),
1094 })
1095 .collect();
1096 assert_eq!(
1097 names,
1098 vec![
1099 "Linear(x)".to_string(),
1100 "Linear(z)".to_string(),
1101 "Interaction(x:z)".to_string(),
1102 ]
1103 );
1104 }
1105
1106 #[test]
1107 fn parse_formula_rejects_unary_signs_inside_wr_expansion() {
1108 for formula in ["y ~ x:-z", "y ~ a*-b", "y ~ x/-z", "y ~ x:+z"] {
1109 let err = parse_formula(formula)
1110 .expect_err("WR expansion must not silently drop unary signs");
1111 let msg = err.to_string();
1112 assert!(
1113 msg.contains("unary `+`/`-` is not supported"),
1114 "unexpected error for {formula}: {msg}"
1115 );
1116 }
1117 }
1118
1119 #[test]
1120 fn parse_formula_supports_wr_power_crossing() {
1121 let parsed = parse_formula("y ~ (x + z)^2").expect("`^` is supported as WR power");
1122 assert_eq!(parsed.response, "y");
1123 assert_eq!(parsed.terms.len(), 3);
1124 let names: Vec<String> = parsed
1125 .terms
1126 .iter()
1127 .map(|t| match t {
1128 ParsedTerm::Linear { name, .. } => format!("Linear({name})"),
1129 ParsedTerm::Interaction { vars, .. } => {
1130 format!("Interaction({})", vars.join(":"))
1131 }
1132 other => format!("Other({other:?})"),
1133 })
1134 .collect();
1135 assert_eq!(
1136 names,
1137 vec![
1138 "Linear(x)".to_string(),
1139 "Linear(z)".to_string(),
1140 "Interaction(x:z)".to_string(),
1141 ]
1142 );
1143 }
1144
1145 #[test]
1146 fn parse_formula_rejects_chained_wr_power() {
1147 let err = parse_formula("y ~ (x + z)^2^3")
1148 .expect_err("chained WR powers must not silently drop later exponents");
1149 let msg = err.to_string();
1150 assert!(
1151 msg.contains("chained `^` operators are not supported"),
1152 "error should explain that chained powers are rejected, got: {msg}"
1153 );
1154 }
1155
1156 #[test]
1157 fn parsed_terms_reference_column_sees_the_by_smooth_variable() {
1158 let parsed = parse_formula("y ~ s(x, by=g)").expect("parse by-smooth");
1164 assert!(
1165 parsed_terms_reference_column(&parsed.terms, "g"),
1166 "s(x, by=g) references column g via options[\"by\"]"
1167 );
1168 assert!(parsed_terms_reference_column(&parsed.terms, "x"));
1169 assert!(!parsed_terms_reference_column(&parsed.terms, "absent"));
1170 }
1171
1172 #[test]
1173 fn marginal_slope_z_column_validator_detects_linear_and_smooth_reuse() {
1174 let main = parse_formula("y ~ x + z").expect("parse main");
1175 let logslope = parse_formula("y ~ s(z, type=duchon, centers=6)").expect("parse logslope");
1176
1177 assert!(parsed_terms_reference_column(&main.terms, "z"));
1178 assert!(parsed_terms_reference_column(&logslope.terms, "z"));
1179
1180 let err = validate_marginal_slope_z_column_exclusion(
1181 &main,
1182 &parse_formula("y ~ 1").expect("parse clean logslope"),
1183 "z",
1184 "bernoulli marginal-slope",
1185 "--logslope-formula",
1186 )
1187 .expect_err("main formula should be rejected");
1188 assert!(err.contains("cannot also appear in the main formula"));
1189
1190 let err = validate_marginal_slope_z_column_exclusion(
1191 &parse_formula("y ~ x").expect("parse clean main"),
1192 &logslope,
1193 "z",
1194 "bernoulli marginal-slope",
1195 "--logslope-formula",
1196 )
1197 .expect_err("logslope formula should be rejected");
1198 assert!(err.contains("cannot also appear in --logslope-formula"));
1199 }
1200
1201 fn frame(columns: &[&str]) -> HashMap<String, usize> {
1203 columns
1204 .iter()
1205 .enumerate()
1206 .map(|(i, c)| ((*c).to_string(), i))
1207 .collect()
1208 }
1209
1210 #[test]
1218 fn marginal_slope_alias_validator_rejects_bare_z_when_z_column_is_named_otherwise() {
1219 let cols = frame(&["y", "x", "pgs_ctn_z"]);
1220 let main = parse_formula("y ~ x + z").expect("parse main");
1221
1222 validate_marginal_slope_z_column_exclusion(
1225 &main,
1226 &parse_formula("y ~ 1").expect("parse logslope"),
1227 "pgs_ctn_z",
1228 "bernoulli marginal-slope",
1229 "logslope_formula",
1230 )
1231 .expect("literal-name validator cannot see the alias");
1232
1233 let err = validate_marginal_slope_z_alias_exclusion(
1234 &main,
1235 &cols,
1236 "pgs_ctn_z",
1237 "bernoulli marginal-slope",
1238 )
1239 .expect_err("bare `z` resolves to the reserved score column");
1240 assert!(err.contains("reserves z column 'pgs_ctn_z'"), "{err}");
1241 assert!(err.contains("cannot also appear in the main formula"), "{err}");
1242 }
1243
1244 #[test]
1246 fn marginal_slope_alias_validator_accepts_legitimate_formulas() {
1247 validate_marginal_slope_z_alias_exclusion(
1249 &parse_formula("y ~ x + w").expect("parse clean main"),
1250 &frame(&["y", "x", "w", "pgs_ctn_z"]),
1251 "pgs_ctn_z",
1252 "bernoulli marginal-slope",
1253 )
1254 .expect("a baseline without the score is legitimate");
1255
1256 validate_marginal_slope_z_alias_exclusion(
1260 &parse_formula("y ~ x + z").expect("parse main with real z"),
1261 &frame(&["y", "x", "z", "pgs_ctn_z"]),
1262 "pgs_ctn_z",
1263 "bernoulli marginal-slope",
1264 )
1265 .expect("a real `z` column is not the alias");
1266
1267 validate_marginal_slope_z_alias_exclusion(
1270 &parse_formula("y ~ x + z").expect("parse main"),
1271 &frame(&["y", "x", "z"]),
1272 "z",
1273 "bernoulli marginal-slope",
1274 )
1275 .expect("alias guard defers to the literal-name validator when they coincide");
1276 }
1277
1278 #[test]
1281 fn marginal_slope_literal_z_column_is_still_rejected() {
1282 let err = validate_marginal_slope_z_column_exclusion(
1283 &parse_formula("y ~ x + pgs_ctn_z").expect("parse main"),
1284 &parse_formula("y ~ 1").expect("parse logslope"),
1285 "pgs_ctn_z",
1286 "bernoulli marginal-slope",
1287 "logslope_formula",
1288 )
1289 .expect_err("the reserved column named in full must still be rejected");
1290 assert!(err.contains("reserves z column 'pgs_ctn_z'"), "{err}");
1291 }
1292
1293 #[test]
1294 fn logslope_surface_declarations_are_additive() {
1295 let parsed = parse_formula("y ~ s(pc1) + logslope(z2, s(pc2)) + logslope(z3, x3)")
1296 .expect("parse additive logslope surfaces");
1297 assert_eq!(parsed.terms.len(), 1);
1298 assert_eq!(parsed.logslope_surfaces.len(), 2);
1299 assert_eq!(parsed.logslope_surfaces[0].z_column, "z2");
1300 assert_eq!(parsed.logslope_surfaces[0].terms.len(), 1);
1301 assert_eq!(parsed.logslope_surfaces[1].z_column, "z3");
1302 assert_eq!(parsed.logslope_surfaces[1].terms.len(), 1);
1303 }
1304
1305 #[test]
1306 fn marginal_slope_z_column_validator_reserves_all_surface_z_columns() {
1307 let main = parse_formula("y ~ x").expect("parse main");
1308 let logslope = parse_formula("y ~ s(pc1) + logslope(z2, s(z3)) + logslope(z3, x)")
1309 .expect("parse logslope surfaces");
1310 let err = validate_marginal_slope_z_column_exclusion(
1311 &main,
1312 &logslope,
1313 "z1",
1314 "bernoulli marginal-slope",
1315 "--logslope-formula",
1316 )
1317 .expect_err("surface formula should reject another reserved z coordinate");
1318 assert!(err.contains("reserves z column 'z3'"));
1319 }
1320
1321 fn random_effect_lenient_unseen(formula: &str) -> bool {
1324 let parsed = parse_formula(formula).expect("parse random-effect formula");
1325 let re = parsed.terms.iter().find_map(|t| match t {
1326 ParsedTerm::RandomEffect { lenient_unseen, .. } => Some(*lenient_unseen),
1327 _ => None,
1328 });
1329 re.unwrap_or_else(|| panic!("{formula} did not lower to a RandomEffect term"))
1330 }
1331
1332 #[test]
1333 fn factor_wrapper_is_strict_on_unseen_levels_while_group_re_are_lenient() {
1334 assert!(
1343 !random_effect_lenient_unseen("y ~ factor(g)"),
1344 "factor(g) is a fixed categorical factor: strict (lenient_unseen=false) on unseen levels"
1345 );
1346 for lenient in ["y ~ group(g)", "y ~ re(g)", "y ~ s(g, bs=re)"] {
1347 assert!(
1348 random_effect_lenient_unseen(lenient),
1349 "{lenient} is a genuine random effect: lenient (lenient_unseen=true) on unseen levels"
1350 );
1351 }
1352 }
1353}
1354
1355#[derive(Clone, Debug)]
1360pub struct LinkWiggleFormulaSpec {
1361 pub degree: usize,
1362 pub num_internal_knots: usize,
1363 pub penalty_orders: Vec<usize>,
1364 pub double_penalty: bool,
1365}
1366
1367pub fn default_linkwiggle_formulaspec() -> LinkWiggleFormulaSpec {
1368 let cfg = WigglePenaltyConfig::cubic_triple_operator_default();
1369 LinkWiggleFormulaSpec {
1370 degree: cfg.degree,
1371 num_internal_knots: cfg.num_internal_knots,
1372 penalty_orders: cfg.penalty_orders,
1373 double_penalty: cfg.double_penalty,
1374 }
1375}
1376
1377#[derive(Clone, Debug)]
1378pub struct LinkFormulaSpec {
1379 pub link: String,
1380 pub mixture_rho: Option<String>,
1381 pub sas_init: Option<String>,
1382 pub beta_logistic_init: Option<String>,
1383}
1384
1385#[derive(Clone, Debug)]
1386pub struct SurvivalFormulaSpec {
1387 pub spec: Option<String>,
1388 pub survival_distribution: Option<String>,
1389}
1390
1391#[derive(Clone, Debug)]
1392pub struct ParsedFormula {
1393 pub response: String,
1394 pub terms: Vec<ParsedTerm>,
1395 pub logslope_surfaces: Vec<LogSlopeSurfaceSpec>,
1396 pub linkwiggle: Option<LinkWiggleFormulaSpec>,
1397 pub timewiggle: Option<LinkWiggleFormulaSpec>,
1398 pub linkspec: Option<LinkFormulaSpec>,
1399 pub survivalspec: Option<SurvivalFormulaSpec>,
1400}
1401
1402#[derive(Clone, Debug)]
1403pub struct LogSlopeSurfaceSpec {
1404 pub z_column: String,
1405 pub terms: Vec<ParsedTerm>,
1406}
1407
1408pub fn marginal_slope_logslope_surfaces(
1409 logslope_formula: &ParsedFormula,
1410 default_z_column: &str,
1411) -> Result<Vec<LogSlopeSurfaceSpec>, String> {
1412 let mut surfaces = Vec::new();
1413 if !logslope_formula.terms.is_empty() {
1414 surfaces.push(LogSlopeSurfaceSpec {
1415 z_column: default_z_column.to_string(),
1416 terms: logslope_formula.terms.clone(),
1417 });
1418 }
1419 surfaces.extend(logslope_formula.logslope_surfaces.clone());
1420 if surfaces.is_empty() {
1421 surfaces.push(LogSlopeSurfaceSpec {
1422 z_column: default_z_column.to_string(),
1423 terms: Vec::new(),
1424 });
1425 }
1426 let mut seen = std::collections::BTreeSet::<String>::new();
1427 for surface in &surfaces {
1428 if !seen.insert(surface.z_column.clone()) {
1429 return Err(FormulaDslError::IncompatibleTerm {
1430 reason: format!(
1431 "logslope formula declares z column '{}' more than once; each z coordinate needs exactly one log-slope surface",
1432 surface.z_column
1433 ),
1434 }
1435 .into());
1436 }
1437 }
1438 Ok(surfaces)
1439}
1440
1441#[derive(Clone, Debug)]
1442pub enum ParsedTerm {
1443 Linear {
1444 name: String,
1445 explicit: bool,
1446 double_penalty: bool,
1447 coefficient_min: Option<f64>,
1448 coefficient_max: Option<f64>,
1449 },
1450 BoundedLinear {
1451 name: String,
1452 min: f64,
1453 max: f64,
1454 prior: BoundedCoefficientPriorSpec,
1455 double_penalty: bool,
1456 },
1457 RandomEffect {
1458 name: String,
1459 lenient_unseen: bool,
1470 },
1471 Smooth {
1472 label: String,
1473 vars: Vec<String>,
1474 kind: SmoothKind,
1475 options: BTreeMap<String, String>,
1476 },
1477 LinkWiggle {
1478 options: BTreeMap<String, String>,
1479 },
1480 TimeWiggle {
1481 options: BTreeMap<String, String>,
1482 },
1483 LinkConfig {
1484 options: BTreeMap<String, String>,
1485 },
1486 SurvivalConfig {
1487 options: BTreeMap<String, String>,
1488 },
1489 LogSlopeSurface {
1490 z_column: String,
1491 terms: Vec<ParsedTerm>,
1492 },
1493 Interaction {
1501 vars: Vec<String>,
1502 double_penalty: bool,
1503 },
1504}
1505
1506pub fn parsed_term_column_names(
1515 terms: &[ParsedTerm],
1516 out: &mut std::collections::BTreeSet<String>,
1517) {
1518 for term in terms {
1519 match term {
1520 ParsedTerm::Linear { name, .. }
1521 | ParsedTerm::BoundedLinear { name, .. }
1522 | ParsedTerm::RandomEffect { name, .. } => {
1523 out.insert(name.clone());
1524 }
1525 ParsedTerm::Smooth { vars, options, .. } => {
1526 out.extend(vars.iter().cloned());
1527 if let Some(by) = options.get("by") {
1528 out.insert(by.clone());
1529 }
1530 }
1531 ParsedTerm::Interaction { vars, .. } => {
1532 out.extend(vars.iter().cloned());
1533 }
1534 ParsedTerm::LinkWiggle { .. }
1535 | ParsedTerm::TimeWiggle { .. }
1536 | ParsedTerm::LinkConfig { .. }
1537 | ParsedTerm::SurvivalConfig { .. } => {}
1538 ParsedTerm::LogSlopeSurface { z_column, terms } => {
1539 out.insert(z_column.clone());
1540 parsed_term_column_names(terms, out);
1541 }
1542 }
1543 }
1544}
1545
1546pub fn parsed_terms_reference_column(terms: &[ParsedTerm], column_name: &str) -> bool {
1547 terms.iter().any(|term| match term {
1548 ParsedTerm::Linear { name, .. }
1549 | ParsedTerm::BoundedLinear { name, .. }
1550 | ParsedTerm::RandomEffect { name, .. } => name == column_name,
1551 ParsedTerm::Smooth { vars, options, .. } => {
1552 vars.iter().any(|var| var == column_name)
1553 || options.get("by").is_some_and(|by| by == column_name)
1554 }
1555 ParsedTerm::Interaction { vars, .. } => vars.iter().any(|var| var == column_name),
1556 ParsedTerm::LinkWiggle { .. }
1557 | ParsedTerm::TimeWiggle { .. }
1558 | ParsedTerm::LinkConfig { .. }
1559 | ParsedTerm::SurvivalConfig { .. } => false,
1560 ParsedTerm::LogSlopeSurface { z_column, terms } => {
1561 z_column == column_name || parsed_terms_reference_column(terms, column_name)
1562 }
1563 })
1564}
1565
1566pub fn validate_marginal_slope_z_alias_exclusion(
1585 main_formula: &ParsedFormula,
1586 col_map: &HashMap<String, usize>,
1587 z_column: &str,
1588 context: &str,
1589) -> Result<(), String> {
1590 if z_column == MARGINAL_SLOPE_Z_ALIAS {
1591 return Ok(());
1593 }
1594 if !marginal_slope_z_alias_is_live(col_map, z_column) {
1595 return Ok(());
1596 }
1597 if parsed_terms_reference_column(&main_formula.terms, MARGINAL_SLOPE_Z_ALIAS) {
1598 return Err(FormulaDslError::IncompatibleTerm {
1599 reason: format!(
1600 "{context} reserves z column '{z_column}' as the auxiliary latent score, and the \
1601 bare name '{MARGINAL_SLOPE_Z_ALIAS}' in the main formula resolves to it; it \
1602 cannot also appear in the main formula. Give the baseline its own covariates \
1603 (the score's effect is carried by the log-slope formula)."
1604 ),
1605 }
1606 .into());
1607 }
1608 Ok(())
1609}
1610
1611pub fn validate_marginal_slope_z_column_exclusion(
1612 main_formula: &ParsedFormula,
1613 logslope_formula: &ParsedFormula,
1614 z_column: &str,
1615 context: &str,
1616 logslope_label: &str,
1617) -> Result<(), String> {
1618 let surfaces = marginal_slope_logslope_surfaces(logslope_formula, z_column)?;
1619 let mut reserved_z_columns = std::collections::BTreeSet::<&str>::new();
1623 reserved_z_columns.insert(z_column);
1624 reserved_z_columns.extend(surfaces.iter().map(|surface| surface.z_column.as_str()));
1625
1626 for reserved in &reserved_z_columns {
1627 if parsed_terms_reference_column(&main_formula.terms, reserved) {
1628 return Err(FormulaDslError::IncompatibleTerm {
1629 reason: format!(
1630 "{context} reserves z column '{reserved}' as the auxiliary latent score; it cannot also appear in the main formula"
1631 ),
1632 }
1633 .into());
1634 }
1635 }
1636 for reserved in &reserved_z_columns {
1637 if parsed_terms_reference_column(&logslope_formula.terms, reserved) {
1638 return Err(FormulaDslError::IncompatibleTerm {
1639 reason: format!(
1640 "{context} reserves z column '{reserved}' as the auxiliary latent score; it cannot also appear in {logslope_label}"
1641 ),
1642 }
1643 .into());
1644 }
1645 for surface in &surfaces {
1646 if parsed_terms_reference_column(&surface.terms, reserved) {
1647 return Err(FormulaDslError::IncompatibleTerm {
1648 reason: format!(
1649 "{context} reserves z column '{reserved}' as an auxiliary latent score; it cannot also appear in {logslope_label}"
1650 ),
1651 }
1652 .into());
1653 }
1654 }
1655 }
1656 Ok(())
1657}
1658
1659#[derive(Clone, Copy, Debug)]
1660pub enum SmoothKind {
1661 S,
1662 Te,
1663 T2,
1667 Ti,
1673}
1674
1675#[derive(Clone, Copy, Debug)]
1676pub enum LinkMode {
1677 Strict,
1678 Flexible,
1679}
1680
1681#[derive(Clone, Debug)]
1682pub struct LinkChoice {
1683 pub mode: LinkMode,
1684 pub link: LinkFunction,
1685 pub mixture_components: Option<Vec<LinkComponent>>,
1686}
1687
1688pub fn effectivelinkwiggle_formulaspec(
1693 formula_linkwiggle: Option<&LinkWiggleFormulaSpec>,
1694 link_choice: Option<&LinkChoice>,
1695) -> Option<LinkWiggleFormulaSpec> {
1696 formula_linkwiggle.cloned().or_else(|| {
1697 link_choice.and_then(|choice| {
1698 if matches!(choice.mode, LinkMode::Flexible) {
1699 Some(default_linkwiggle_formulaspec())
1700 } else {
1701 None
1702 }
1703 })
1704 })
1705}
1706
1707pub const fn linkname_supports_joint_wiggle(link: LinkFunction) -> bool {
1708 !matches!(link, LinkFunction::Sas | LinkFunction::BetaLogistic)
1709}
1710
1711pub const fn linkchoice_supports_joint_wiggle(choice: &LinkChoice) -> bool {
1712 match &choice.mixture_components {
1713 None => linkname_supports_joint_wiggle(choice.link),
1714 Some(_) => false,
1715 }
1716}
1717
1718pub fn require_linkchoice_supports_joint_wiggle(
1719 choice: &LinkChoice,
1720 context: &str,
1721) -> Result<(), String> {
1722 if linkchoice_supports_joint_wiggle(choice) {
1723 Ok(())
1724 } else {
1725 Err(joint_wiggle_unsupported_link_message(context))
1726 }
1727}
1728
1729pub const fn likelihood_spec_supports_joint_wiggle(likelihood: &LikelihoodSpec) -> bool {
1730 inverse_link_supports_joint_wiggle(&likelihood.link)
1731}
1732
1733pub fn require_likelihood_spec_supports_joint_wiggle(
1734 likelihood: &LikelihoodSpec,
1735 context: &str,
1736) -> Result<(), String> {
1737 if likelihood_spec_supports_joint_wiggle(likelihood) {
1738 Ok(())
1739 } else {
1740 Err(joint_wiggle_unsupported_link_message(context))
1741 }
1742}
1743
1744pub const fn inverse_link_supports_joint_wiggle(link: &InverseLink) -> bool {
1757 matches!(
1758 link,
1759 InverseLink::Standard(StandardLink::Identity)
1760 | InverseLink::Standard(StandardLink::Log)
1761 | InverseLink::Standard(StandardLink::Logit)
1762 | InverseLink::Standard(StandardLink::Probit)
1763 | InverseLink::Standard(StandardLink::CLogLog)
1764 | InverseLink::Standard(StandardLink::LogLog)
1765 | InverseLink::Standard(StandardLink::Cauchit)
1766 )
1767}
1768
1769pub fn require_inverse_link_supports_joint_wiggle(
1770 link: &InverseLink,
1771 context: &str,
1772) -> Result<(), String> {
1773 if inverse_link_supports_joint_wiggle(link) {
1774 Ok(())
1775 } else {
1776 Err(joint_wiggle_unsupported_link_message(context))
1777 }
1778}
1779
1780pub const fn binomial_inverse_link_supports_joint_wiggle(link: &InverseLink) -> bool {
1792 matches!(
1793 link,
1794 InverseLink::Standard(StandardLink::Logit)
1795 | InverseLink::Standard(StandardLink::Probit)
1796 | InverseLink::Standard(StandardLink::CLogLog)
1797 | InverseLink::Standard(StandardLink::LogLog)
1798 | InverseLink::Standard(StandardLink::Cauchit)
1799 )
1800}
1801
1802pub fn require_binomial_inverse_link_supports_joint_wiggle(
1803 link: &InverseLink,
1804 context: &str,
1805) -> Result<(), String> {
1806 if binomial_inverse_link_supports_joint_wiggle(link) {
1807 Ok(())
1808 } else {
1809 Err(FormulaDslError::IncompatibleTerm {
1810 reason: format!(
1811 "{context} does not support identity, log, latent-cloglog, SAS, BetaLogistic, or Mixture links; wiggle is only available for jointly fitted standard binomial probability links (logit/probit/cloglog/loglog/cauchit)"
1812 ),
1813 }
1814 .into())
1815 }
1816}
1817
1818pub fn joint_wiggle_unsupported_link_message(context: &str) -> String {
1819 format!(
1820 "{context} does not support latent-cloglog, SAS, BetaLogistic, or Mixture links; wiggle is only available for jointly fitted standard links"
1821 )
1822}
1823
1824pub fn option_usize(map: &BTreeMap<String, String>, key: &str) -> Option<usize> {
1829 map.get(key).and_then(|v| v.parse::<usize>().ok())
1830}
1831
1832fn validate_known_term_options(
1839 term_name: &str,
1840 options: &BTreeMap<String, String>,
1841 known: &[&str],
1842 raw: &str,
1843) -> Result<(), String> {
1844 let known_set: std::collections::BTreeSet<&&str> = known.iter().collect();
1845 for key in options.keys() {
1846 if !known_set.contains(&key.as_str()) {
1847 let known_sorted = {
1848 let mut v = known.to_vec();
1849 v.sort_unstable();
1850 v.join(", ")
1851 };
1852 let known_hint = if known.is_empty() {
1853 "no options".to_string()
1854 } else {
1855 format!("[{known_sorted}]")
1856 };
1857 return Err(FormulaDslError::InvalidArgument {
1858 reason: format!(
1859 "{term_name}() does not accept option `{key}` (in `{raw}`); known options: {known_hint}"
1860 ),
1861 }
1862 .into());
1863 }
1864 }
1865 Ok(())
1866}
1867
1868pub fn option_usize_any(map: &BTreeMap<String, String>, keys: &[&str]) -> Option<usize> {
1869 for key in keys {
1870 if let Some(v) = option_usize(map, key) {
1871 return Some(v);
1872 }
1873 }
1874 None
1875}
1876
1877pub fn option_usize_strict(
1884 map: &BTreeMap<String, String>,
1885 key: &str,
1886) -> Result<Option<usize>, String> {
1887 match map.get(key) {
1888 None => Ok(None),
1889 Some(raw) => raw.parse::<usize>().map(Some).map_err(|err| {
1890 FormulaDslError::InvalidArgument {
1891 reason: format!(
1892 "option `{key}={raw}` is not a non-negative integer; \
1893 expected a whole number >= 0: {err}"
1894 ),
1895 }
1896 .into()
1897 }),
1898 }
1899}
1900
1901pub fn option_usize_any_strict(
1904 map: &BTreeMap<String, String>,
1905 keys: &[&str],
1906) -> Result<Option<usize>, String> {
1907 for key in keys {
1908 if let Some(v) = option_usize_strict(map, key)? {
1909 return Ok(Some(v));
1910 }
1911 }
1912 Ok(None)
1913}
1914
1915pub fn option_f64(map: &BTreeMap<String, String>, key: &str) -> Option<f64> {
1916 map.get(key).and_then(|v| v.parse::<f64>().ok())
1917}
1918
1919pub fn option_f64_strict(map: &BTreeMap<String, String>, key: &str) -> Result<Option<f64>, String> {
1923 match map.get(key) {
1924 None => Ok(None),
1925 Some(raw) => match raw.parse::<f64>() {
1926 Ok(v) if v.is_finite() => Ok(Some(v)),
1927 Ok(v) => Err(FormulaDslError::InvalidArgument {
1928 reason: format!("option `{key}={raw}` parses as {v} which is not a finite number"),
1929 }
1930 .into()),
1931 Err(err) => Err(FormulaDslError::InvalidArgument {
1932 reason: format!(
1933 "option `{key}={raw}` is not a valid number; expected a finite decimal: {err}"
1934 ),
1935 }
1936 .into()),
1937 },
1938 }
1939}
1940
1941pub fn option_bool(map: &BTreeMap<String, String>, key: &str) -> Option<bool> {
1942 map.get(key)
1943 .and_then(|v| match v.trim().to_ascii_lowercase().as_str() {
1944 "true" | "1" | "yes" | "y" => Some(true),
1945 "false" | "0" | "no" | "n" => Some(false),
1946 _ => None,
1947 })
1948}
1949
1950pub fn option_bool_strict(
1956 map: &BTreeMap<String, String>,
1957 key: &str,
1958) -> Result<Option<bool>, String> {
1959 match map.get(key) {
1960 None => Ok(None),
1961 Some(raw) => match raw.trim().to_ascii_lowercase().as_str() {
1962 "true" | "1" | "yes" | "y" => Ok(Some(true)),
1963 "false" | "0" | "no" | "n" => Ok(Some(false)),
1964 _ => Err(FormulaDslError::InvalidArgument {
1965 reason: format!(
1966 "option `{key}={raw}` is not a boolean; \
1967 expected one of true/false/yes/no/1/0"
1968 ),
1969 }
1970 .into()),
1971 },
1972 }
1973}
1974
1975pub fn strip_quotes(v: &str) -> &str {
1976 let b = v.as_bytes();
1977 if b.len() >= 2
1978 && ((b[0] == b'\'' && b[b.len() - 1] == b'\'') || (b[0] == b'"' && b[b.len() - 1] == b'"'))
1979 {
1980 &v[1..v.len() - 1]
1981 } else {
1982 v
1983 }
1984}
1985
1986fn parse_linear_constraint_bounds(
1991 options: &BTreeMap<String, String>,
1992 raw: &str,
1993) -> Result<(Option<f64>, Option<f64>), String> {
1994 let min = parse_optional_f64_option_alias(options, &["min", "lower"], raw, "linear")?;
1995 let max = parse_optional_f64_option_alias(options, &["max", "upper"], raw, "linear")?;
1996 if let (Some(min), Some(max)) = (min, max)
1997 && (!min.is_finite() || !max.is_finite() || min > max)
1998 {
1999 return Err(FormulaDslError::InvalidArgument {
2000 reason: format!(
2001 "linear coefficient constraints require finite min <= max, got min={min}, max={max}: {raw}"
2002 ),
2003 }
2004 .into());
2005 }
2006 Ok((min, max))
2007}
2008
2009fn parse_required_f64_option(
2010 options: &BTreeMap<String, String>,
2011 key: &str,
2012 raw: &str,
2013) -> Result<f64, String> {
2014 let value = options
2015 .get(key)
2016 .ok_or_else(|| FormulaDslError::MalformedConfig {
2017 reason: format!("bounded() is missing required '{key}' argument: {raw}"),
2018 })?;
2019 value.parse::<f64>().map_err(|err| {
2020 FormulaDslError::InvalidArgument {
2021 reason: format!(
2022 "bounded() argument '{key}' must be a finite number, got '{}': {err}: {raw}",
2023 value
2024 ),
2025 }
2026 .into()
2027 })
2028}
2029
2030fn parse_optional_f64_option(
2031 options: &BTreeMap<String, String>,
2032 key: &str,
2033 raw: &str,
2034) -> Result<Option<f64>, String> {
2035 match options.get(key) {
2036 Some(value) => value.parse::<f64>().map(Some).map_err(|err| {
2037 FormulaDslError::InvalidArgument {
2038 reason: format!(
2039 "bounded() argument '{key}' must be a finite number, got '{}': {err}: {raw}",
2040 value
2041 ),
2042 }
2043 .into()
2044 }),
2045 None => Ok(None),
2046 }
2047}
2048
2049fn parse_optional_f64_option_alias(
2050 options: &BTreeMap<String, String>,
2051 keys: &[&str],
2052 raw: &str,
2053 fn_label: &str,
2054) -> Result<Option<f64>, String> {
2055 let mut found: Option<(&str, f64)> = None;
2056 for key in keys {
2057 if let Some(value) = options.get(*key) {
2058 let parsed = value
2059 .parse::<f64>()
2060 .map_err(|err| FormulaDslError::InvalidArgument {
2061 reason: format!(
2062 "{fn_label}() argument '{key}' must be a finite number, got '{}': {err}: {raw}",
2063 value
2064 ),
2065 })?;
2066 if found.is_some() {
2067 return Err(FormulaDslError::IncompatibleTerm {
2068 reason: format!(
2069 "{fn_label}() cannot specify both '{}' and '{}': {raw}",
2070 found.expect("present").0,
2071 key
2072 ),
2073 }
2074 .into());
2075 }
2076 found = Some((key, parsed));
2077 }
2078 }
2079 Ok(found.map(|(_, v)| v))
2080}
2081
2082fn parse_linkwiggle_penalty_orders(raw: Option<&str>) -> Result<Vec<usize>, String> {
2083 let Some(raw) = raw.map(str::trim) else {
2084 return Ok(WigglePenaltyConfig::cubic_triple_operator_default().penalty_orders);
2085 };
2086 if raw.is_empty() {
2087 return Ok(WigglePenaltyConfig::cubic_triple_operator_default().penalty_orders);
2088 }
2089 let mut out = Vec::<usize>::new();
2090 for token in raw.split(',') {
2091 let t = token.trim().to_ascii_lowercase();
2092 if t.is_empty() {
2093 continue;
2094 }
2095 match t.as_str() {
2096 "all" => {
2097 out.extend([1, 2, 3]);
2098 }
2099 "slope" | "1" => out.push(1),
2100 "curvature" | "2" => out.push(2),
2101 "curvature-change" | "curvature_change" | "3" => out.push(3),
2102 _ => {
2103 return Err(FormulaDslError::InvalidArgument {
2104 reason: format!(
2105 "invalid linkwiggle penalty_order '{t}'; use all|slope|curvature|curvature-change or 1/2/3"
2106 ),
2107 }
2108 .into());
2109 }
2110 }
2111 }
2112 if out.is_empty() {
2113 out.extend(WigglePenaltyConfig::cubic_triple_operator_default().penalty_orders);
2114 }
2115 out.sort_unstable();
2116 out.dedup();
2117 Ok(out)
2118}
2119
2120pub fn parse_linkwiggle_formulaspec(
2121 options: &BTreeMap<String, String>,
2122 raw: &str,
2123) -> Result<LinkWiggleFormulaSpec, String> {
2124 let allowed = [
2125 "degree",
2126 "internal_knots",
2127 "penalty_order",
2128 "double_penalty",
2129 ];
2130 let unknown = options
2131 .keys()
2132 .filter(|key| !allowed.contains(&key.as_str()))
2133 .cloned()
2134 .collect::<Vec<_>>();
2135 let term_name = raw.split('(').next().unwrap_or("linkwiggle");
2136 if !unknown.is_empty() {
2137 return Err(FormulaDslError::InvalidArgument {
2138 reason: format!(
2139 "{}() does not support option(s) {}: {raw}",
2140 term_name,
2141 unknown.join(", ")
2142 ),
2143 }
2144 .into());
2145 }
2146 let defaults = WigglePenaltyConfig::cubic_triple_operator_default();
2147 let degree = option_usize_strict(options, "degree")?.unwrap_or(defaults.degree);
2163 if degree < 1 {
2164 return Err(FormulaDslError::InvalidArgument {
2165 reason: format!("{term_name}() requires degree >= 1: {raw}"),
2166 }
2167 .into());
2168 }
2169 let num_internal_knots =
2170 option_usize_strict(options, "internal_knots")?.unwrap_or(defaults.num_internal_knots);
2171 if num_internal_knots == 0 {
2172 return Err(FormulaDslError::InvalidArgument {
2173 reason: format!("{term_name}() requires internal_knots > 0: {raw}"),
2174 }
2175 .into());
2176 }
2177 let penalty_order_requested = options
2197 .get("penalty_order")
2198 .is_some_and(|raw| !raw.trim().is_empty());
2199 let mut penalty_orders =
2200 parse_linkwiggle_penalty_orders(options.get("penalty_order").map(String::as_str))?;
2201 if !penalty_order_requested {
2202 let supported: Vec<usize> = penalty_orders
2203 .iter()
2204 .copied()
2205 .filter(|&order| order <= degree)
2206 .collect();
2207 if supported.is_empty() {
2208 return Err(FormulaDslError::InvalidArgument {
2209 reason: format!(
2210 "{term_name}(degree={degree}) supports no default penalty order: the default ladder {penalty_orders:?} contains no derivative order a degree-{degree} basis can penalize"
2211 ),
2212 }
2213 .into());
2214 }
2215 penalty_orders = supported;
2216 }
2217 let double_penalty =
2218 option_bool_strict(options, "double_penalty")?.unwrap_or(defaults.double_penalty);
2219 Ok(LinkWiggleFormulaSpec {
2220 degree,
2221 num_internal_knots,
2222 penalty_orders,
2223 double_penalty,
2224 })
2225}
2226
2227fn parse_link_formulaspec(
2228 options: &BTreeMap<String, String>,
2229 raw: &str,
2230) -> Result<LinkFormulaSpec, String> {
2231 let link = options
2232 .get("type")
2233 .map(|s| s.trim().to_string())
2234 .ok_or_else(|| FormulaDslError::MalformedConfig {
2235 reason: format!("link() requires type=<link-name>: {raw}"),
2236 })?;
2237 if link.is_empty() {
2238 return Err(FormulaDslError::MalformedConfig {
2239 reason: format!("link() requires a non-empty type: {raw}"),
2240 }
2241 .into());
2242 }
2243 let mixture_rho = options.get("rho").map(|s| s.trim().to_string());
2244 let sas_init = options.get("sas_init").map(|s| s.trim().to_string());
2245 let beta_logistic_init = options
2246 .get("beta_logistic_init")
2247 .map(|s| s.trim().to_string());
2248 Ok(LinkFormulaSpec {
2249 link,
2250 mixture_rho,
2251 sas_init,
2252 beta_logistic_init,
2253 })
2254}
2255
2256fn parse_survival_formulaspec(
2257 options: &BTreeMap<String, String>,
2258 raw: &str,
2259) -> Result<SurvivalFormulaSpec, String> {
2260 if options.is_empty() {
2261 return Err(FormulaDslError::MalformedConfig {
2262 reason: format!(
2263 "survmodel() requires at least one named option (e.g., spec=..., distribution=...): {raw}"
2264 ),
2265 }
2266 .into());
2267 }
2268 Ok(SurvivalFormulaSpec {
2269 spec: options.get("spec").map(|s| s.trim().to_string()),
2270 survival_distribution: options.get("distribution").map(|s| s.trim().to_string()),
2271 })
2272}
2273
2274fn parse_bounded_priorspec(
2275 options: &BTreeMap<String, String>,
2276 min: f64,
2277 max: f64,
2278 raw: &str,
2279) -> Result<BoundedCoefficientPriorSpec, String> {
2280 let prior_mode = options.get("prior").map(|s| s.to_ascii_lowercase());
2281 let pull = options.get("pull").map(|s| s.to_ascii_lowercase());
2282 let target = parse_optional_f64_option(options, "target", raw)?;
2283 let strength = parse_optional_f64_option(options, "strength", raw)?;
2284
2285 let target_mode = target.is_some() || strength.is_some();
2286 if prior_mode.is_some() && pull.is_some() {
2287 return Err(FormulaDslError::IncompatibleTerm {
2288 reason: format!("bounded() cannot combine prior=... with pull=...: {raw}"),
2289 }
2290 .into());
2291 }
2292 if prior_mode.is_some() && target_mode {
2293 return Err(FormulaDslError::IncompatibleTerm {
2294 reason: format!("bounded() cannot combine prior=... with target/strength: {raw}"),
2295 }
2296 .into());
2297 }
2298 if pull.is_some() && target_mode {
2299 return Err(FormulaDslError::IncompatibleTerm {
2300 reason: format!("bounded() cannot combine pull=... with target/strength: {raw}"),
2301 }
2302 .into());
2303 }
2304
2305 if let Some(priorname) = prior_mode {
2306 return match priorname.as_str() {
2307 "none" => Ok(BoundedCoefficientPriorSpec::None),
2308 "uniform" | "log-jacobian" | "log_jacobian" | "jacobian" => {
2309 Ok(BoundedCoefficientPriorSpec::Uniform)
2310 }
2311 "center" => Ok(BoundedCoefficientPriorSpec::Beta { a: 2.0, b: 2.0 }),
2312 _ => Err(FormulaDslError::InvalidArgument {
2313 reason: format!(
2314 "bounded() prior must currently be one of none|uniform|log-jacobian|center, got '{}': {raw}",
2315 priorname
2316 ),
2317 }
2318 .into()),
2319 };
2320 }
2321
2322 if let Some(pull_mode) = pull {
2323 return match pull_mode.as_str() {
2324 "uniform" | "log-jacobian" | "log_jacobian" | "jacobian" => {
2325 Ok(BoundedCoefficientPriorSpec::Uniform)
2326 }
2327 "center" => Ok(BoundedCoefficientPriorSpec::Beta { a: 2.0, b: 2.0 }),
2328 _ => Err(FormulaDslError::InvalidArgument {
2329 reason: format!(
2330 "bounded() pull must currently be 'uniform'/'log-jacobian' or 'center', got '{}': {raw}",
2331 pull_mode
2332 ),
2333 }
2334 .into()),
2335 };
2336 }
2337
2338 if target_mode {
2339 let targetvalue = target.ok_or_else(|| FormulaDslError::MalformedConfig {
2340 reason: format!("bounded() target is required with strength: {raw}"),
2341 })?;
2342 let strengthvalue = strength.ok_or_else(|| FormulaDslError::MalformedConfig {
2343 reason: format!("bounded() strength is required with target: {raw}"),
2344 })?;
2345 if !(min < targetvalue && targetvalue < max) {
2346 return Err(FormulaDslError::InvalidArgument {
2347 reason: format!("bounded() target must lie strictly inside ({min}, {max}): {raw}"),
2348 }
2349 .into());
2350 }
2351 if !strengthvalue.is_finite() || strengthvalue <= 0.0 {
2352 return Err(FormulaDslError::InvalidArgument {
2353 reason: format!("bounded() strength must be finite and > 0: {raw}"),
2354 }
2355 .into());
2356 }
2357 let z = (targetvalue - min) / (max - min);
2358 let a = 1.0 + strengthvalue * z;
2359 let b = 1.0 + strengthvalue * (1.0 - z);
2360 return Ok(BoundedCoefficientPriorSpec::Beta { a, b });
2361 }
2362
2363 Ok(BoundedCoefficientPriorSpec::None)
2364}
2365
2366pub fn formula_rhs_text(formula: &str) -> Result<String, String> {
2371 let parsed = parse_formula_dsl(formula)?;
2372 if parsed.rhs_terms.is_empty() {
2373 return Err(FormulaDslError::ParseError {
2374 reason: "formula right-hand side cannot be empty".to_string(),
2375 }
2376 .into());
2377 }
2378 Ok(parsed.rhs_terms.join(" + "))
2379}
2380
2381pub fn parse_surv_response(
2388 lhs: &str,
2389) -> Result<Option<(Option<String>, String, String)>, FormulaDslError> {
2390 let trimmed = lhs.trim();
2391 let call = match parse_function_call(trimmed) {
2392 Ok(call) => call,
2393 Err(_) => return Ok(None),
2394 };
2395 if !call.name.eq_ignore_ascii_case("surv") {
2396 return Ok(None);
2397 }
2398 let vars = call
2399 .args
2400 .iter()
2401 .filter_map(|arg| match arg {
2402 CallArgSpec::Positional(v) => Some(v.trim().to_string()),
2403 CallArgSpec::Named { .. } => None,
2404 })
2405 .filter(|s| !s.is_empty())
2406 .collect::<Vec<_>>();
2407 match vars.len() {
2408 2 => Ok(Some((None, vars[0].clone(), vars[1].clone()))),
2412 3 => Ok(Some((
2413 Some(vars[0].clone()),
2414 vars[1].clone(),
2415 vars[2].clone(),
2416 ))),
2417 n => Err(FormulaDslError::InvalidArgument {
2418 reason: format!(
2419 "Surv(...) expects either Surv(time, event) (right-censored) or \
2420 Surv(entry, exit, event) (left-truncated); got {n} columns"
2421 ),
2422 }),
2423 }
2424}
2425
2426pub fn parse_surv_interval_response(
2441 lhs: &str,
2442) -> Result<Option<(String, String, String)>, FormulaDslError> {
2443 let trimmed = lhs.trim();
2444 let call = match parse_function_call(trimmed) {
2445 Ok(call) => call,
2446 Err(_) => return Ok(None),
2447 };
2448 if !call.name.eq_ignore_ascii_case("survinterval") {
2449 return Ok(None);
2450 }
2451 let vars = call
2452 .args
2453 .iter()
2454 .filter_map(|arg| match arg {
2455 CallArgSpec::Positional(v) => Some(v.trim().to_string()),
2456 CallArgSpec::Named { .. } => None,
2457 })
2458 .filter(|s| !s.is_empty())
2459 .collect::<Vec<_>>();
2460 match vars.len() {
2461 3 => Ok(Some((vars[0].clone(), vars[1].clone(), vars[2].clone()))),
2462 n => Err(FormulaDslError::InvalidArgument {
2463 reason: format!(
2464 "SurvInterval(...) expects SurvInterval(L, R, event) (interval-censored, the \
2465 observed bracket T ∈ (L, R]); got {n} columns"
2466 ),
2467 }),
2468 }
2469}
2470
2471fn top_level_formula_separator(input: &str) -> Result<Option<usize>, String> {
2472 let mut depth = 0_i32;
2473 let mut in_single = false;
2474 let mut in_double = false;
2475
2476 for (idx, ch) in input.char_indices() {
2477 let quoted = in_single || in_double;
2478 if ch == '\'' && !in_double {
2479 in_single = !in_single;
2480 } else if ch == '"' && !in_single {
2481 in_double = !in_double;
2482 } else if !quoted && matches!(ch, '(' | '[' | '{') {
2483 depth += 1;
2484 } else if !quoted && matches!(ch, ')' | ']' | '}') && depth > 0 {
2485 depth -= 1;
2486 } else if ch == '~' && !quoted && depth == 0 {
2487 return Ok(Some(idx));
2488 }
2489 }
2490
2491 if in_single || in_double || depth != 0 {
2492 return Err(FormulaDslError::ParseError {
2493 reason: "invalid auxiliary formula syntax: unbalanced parentheses or quotes"
2494 .to_string(),
2495 }
2496 .into());
2497 }
2498 Ok(None)
2499}
2500
2501pub fn parse_matching_auxiliary_formula(
2502 formula: &str,
2503 response: &str,
2504 flag_name: &str,
2505) -> Result<(String, ParsedFormula), FormulaDslError> {
2506 let rhs = formula.trim();
2507 if top_level_formula_separator(rhs)?.is_some() {
2508 return Err(FormulaDslError::InvalidArgument {
2509 reason: format!(
2510 "{flag_name} expects only the terms after '~', not a full 'response ~ terms' formula; use {flag_name} 's(x)' instead of {flag_name} 'y ~ s(x)' (or pass '1' for an intercept-only noise model)"
2511 ),
2512 });
2513 }
2514 let parsed_formula = parse_formula(&format!("{response} ~ {rhs}"))?;
2515 Ok((rhs.to_string(), parsed_formula))
2516}
2517
2518pub fn validate_auxiliary_formula_controls(
2519 parsed_formula: &ParsedFormula,
2520 flag_name: &str,
2521) -> Result<(), String> {
2522 if parsed_formula.linkwiggle.is_some() {
2523 return Err(FormulaDslError::IncompatibleTerm {
2524 reason: format!(
2525 "linkwiggle(...) is only supported in the main formula, not {flag_name}"
2526 ),
2527 }
2528 .into());
2529 }
2530 if parsed_formula.timewiggle.is_some() {
2531 return Err(FormulaDslError::IncompatibleTerm {
2532 reason: format!(
2533 "timewiggle(...) is only supported in the main survival formula, not {flag_name}"
2534 ),
2535 }
2536 .into());
2537 }
2538 if parsed_formula.linkspec.is_some() {
2539 return Err(FormulaDslError::IncompatibleTerm {
2540 reason: format!("link(...) is only supported in the main formula, not {flag_name}"),
2541 }
2542 .into());
2543 }
2544 if parsed_formula.survivalspec.is_some() {
2545 return Err(FormulaDslError::IncompatibleTerm {
2546 reason: format!(
2547 "survmodel(...) is only supported in the main survival formula, not {flag_name}"
2548 ),
2549 }
2550 .into());
2551 }
2552 if !parsed_formula.logslope_surfaces.is_empty() && flag_name != "--logslope-formula" {
2553 return Err(FormulaDslError::IncompatibleTerm {
2554 reason: format!(
2555 "logslope(...) is only supported in --logslope-formula, not {flag_name}"
2556 ),
2557 }
2558 .into());
2559 }
2560 Ok::<(), _>(())
2561}
2562
2563pub fn parse_formula(formula: &str) -> Result<ParsedFormula, FormulaDslError> {
2564 let parsed_dsl =
2565 parse_formula_dsl(formula).map_err(|reason| FormulaDslError::ParseError { reason })?;
2566 let lhs = parsed_dsl.response_expr.trim();
2567 if lhs.is_empty() {
2568 return Err(FormulaDslError::ParseError {
2569 reason: "formula response (left-hand side) cannot be empty".to_string(),
2570 });
2571 }
2572 let mut terms = Vec::<ParsedTerm>::new();
2573 let mut linkwiggle: Option<LinkWiggleFormulaSpec> = None;
2574 let mut timewiggle: Option<LinkWiggleFormulaSpec> = None;
2575 let mut linkspec: Option<LinkFormulaSpec> = None;
2576 let mut survivalspec: Option<SurvivalFormulaSpec> = None;
2577 let mut logslope_surfaces = Vec::<LogSlopeSurfaceSpec>::new();
2578 let mut seen_term_keys: std::collections::BTreeSet<String> = std::collections::BTreeSet::new();
2583 let mut expanded_terms = Vec::<String>::new();
2584 for raw in parsed_dsl.rhs_terms {
2585 let trimmed = raw.trim();
2586 if trimmed.is_empty() {
2587 expanded_terms.push(String::new());
2588 continue;
2589 }
2590 let is_call = parse_function_call(trimmed).is_ok();
2595 let needs_expansion = !is_call
2596 && trimmed
2597 .chars()
2598 .scan(0i32, |depth, ch| {
2599 let d_before = *depth;
2600 if matches!(ch, '(' | '[' | '{') {
2601 *depth += 1;
2602 } else if matches!(ch, ')' | ']' | '}') && *depth > 0 {
2603 *depth -= 1;
2604 }
2605 Some((d_before, ch))
2606 })
2607 .any(|(d, ch)| d == 0 && matches!(ch, ':' | '*' | '/' | '^'));
2608 if needs_expansion {
2609 for atoms in
2610 expand_wr_term(trimmed).map_err(|reason| FormulaDslError::ParseError { reason })?
2611 {
2612 if atoms.is_empty() {
2613 continue;
2614 }
2615 expanded_terms.push(atoms.join(":"));
2616 }
2617 } else {
2618 expanded_terms.push(trimmed.to_string());
2619 }
2620 }
2621
2622 for raw in expanded_terms {
2623 let t = raw.trim();
2624 if t.is_empty() || t == "1" {
2625 continue;
2626 }
2627 if t == "0" || t == "-1" {
2628 return Err(FormulaDslError::IncompatibleTerm {
2629 reason: "formula terms '0'/'-1' (intercept removal) are not supported yet"
2630 .to_string(),
2631 });
2632 }
2633 let key: String = {
2637 let mut acc = String::with_capacity(t.len());
2638 let mut in_single = false;
2639 let mut in_double = false;
2640 for ch in t.chars() {
2641 match ch {
2642 '\'' if !in_double => {
2643 in_single = !in_single;
2644 acc.push(ch);
2645 }
2646 '"' if !in_single => {
2647 in_double = !in_double;
2648 acc.push(ch);
2649 }
2650 c if c.is_whitespace() && !in_single && !in_double => {}
2651 _ => acc.push(ch),
2652 }
2653 }
2654 acc
2655 };
2656 if !seen_term_keys.insert(key.clone()) {
2657 return Err(FormulaDslError::IncompatibleTerm {
2658 reason: format!(
2659 "formula `{formula}` lists term `{t}` more than once. \
2660 Duplicate terms produce a rank-deficient design; \
2661 keep one copy or differentiate them (e.g. distinct k=, bs= options)."
2662 ),
2663 });
2664 }
2665 match parse_term(t)? {
2666 ParsedTerm::LinkWiggle { options } => {
2667 if linkwiggle.is_some() {
2668 return Err(FormulaDslError::IncompatibleTerm {
2669 reason: "formula can include at most one linkwiggle(...) term".to_string(),
2670 });
2671 }
2672 linkwiggle = Some(parse_linkwiggle_formulaspec(&options, t)?);
2673 }
2674 ParsedTerm::TimeWiggle { options } => {
2675 if timewiggle.is_some() {
2676 return Err(FormulaDslError::IncompatibleTerm {
2677 reason: "formula can include at most one timewiggle(...) term".to_string(),
2678 });
2679 }
2680 timewiggle = Some(parse_linkwiggle_formulaspec(&options, t)?);
2681 }
2682 ParsedTerm::LinkConfig { options } => {
2683 if linkspec.is_some() {
2684 return Err(FormulaDslError::IncompatibleTerm {
2685 reason: "formula can include at most one link(...) term".to_string(),
2686 });
2687 }
2688 linkspec = Some(parse_link_formulaspec(&options, t)?);
2689 }
2690 ParsedTerm::SurvivalConfig { options } => {
2691 if survivalspec.is_some() {
2692 return Err(FormulaDslError::IncompatibleTerm {
2693 reason: "formula can include at most one survmodel(...) term".to_string(),
2694 });
2695 }
2696 survivalspec = Some(parse_survival_formulaspec(&options, t)?);
2697 }
2698 ParsedTerm::LogSlopeSurface { z_column, terms } => {
2699 logslope_surfaces.push(LogSlopeSurfaceSpec { z_column, terms });
2700 }
2701 other => terms.push(other),
2702 }
2703 }
2704 if lhs.chars().all(|c| c.is_alphanumeric() || c == '_')
2710 && !lhs.is_empty()
2711 && parsed_terms_reference_column(&terms, lhs)
2712 {
2713 return Err(FormulaDslError::IncompatibleTerm {
2714 reason: format!(
2715 "formula `{formula}` uses response column `{lhs}` as its own predictor. \
2716 This fits y as a function of itself and is almost certainly a typo. \
2717 Drop the term that mentions `{lhs}` from the right-hand side."
2718 ),
2719 });
2720 }
2721 Ok(ParsedFormula {
2722 response: lhs.to_string(),
2723 terms,
2724 logslope_surfaces,
2725 linkwiggle,
2726 timewiggle,
2727 linkspec,
2728 survivalspec,
2729 })
2730}
2731
2732pub fn parse_term(raw: &str) -> Result<ParsedTerm, String> {
2733 fn split_call_args(call: &FunctionCallSpec) -> (Vec<String>, BTreeMap<String, String>) {
2734 let mut vars = Vec::<String>::new();
2735 let mut options = BTreeMap::<String, String>::new();
2736 for arg in &call.args {
2737 match arg {
2738 CallArgSpec::Positional(v) => vars.push(v.trim().to_string()),
2739 CallArgSpec::Named { key, value } => {
2740 options.insert(key.to_ascii_lowercase(), strip_quotes(value).to_string());
2741 }
2742 }
2743 }
2744 (vars, options)
2745 }
2746
2747 if raw.contains(':')
2751 && !raw.contains('(')
2752 && raw.split(':').all(|piece| is_exact_ident(piece.trim()))
2753 {
2754 let vars: Vec<String> = raw
2755 .split(':')
2756 .map(|piece| piece.trim().to_string())
2757 .collect();
2758 if vars.len() >= 2 {
2759 let mut sorted = vars.clone();
2760 sorted.sort();
2761 sorted.dedup();
2762 if sorted.len() != vars.len() {
2763 return Err(FormulaDslError::IncompatibleTerm {
2764 reason: format!(
2765 "interaction term `{raw}` references the same variable more than once"
2766 ),
2767 }
2768 .into());
2769 }
2770 return Ok(ParsedTerm::Interaction {
2771 vars: sorted,
2772 double_penalty: false,
2773 });
2774 }
2775 }
2776
2777 if let Ok(call) = parse_function_call(raw) {
2781 let name = call.name.to_ascii_lowercase();
2782 let (vars, mut options) = split_call_args(&call);
2783 match name.as_str() {
2784 "constrain" | "constraint" | "box" => {
2785 if vars.len() != 1 {
2786 return Err(FormulaDslError::InvalidArgument {
2787 reason: format!(
2788 "constrain()/constraint()/box() expects exactly one variable: {raw}"
2789 ),
2790 }
2791 .into());
2792 }
2793 validate_known_term_options(
2794 "constrain",
2795 &options,
2796 &["min", "lower", "max", "upper", "double_penalty"],
2797 raw,
2798 )?;
2799 let (coefficient_min, coefficient_max) =
2800 parse_linear_constraint_bounds(&options, raw)?;
2801 if coefficient_min.is_none() && coefficient_max.is_none() {
2802 return Err(FormulaDslError::MalformedConfig {
2803 reason: format!(
2804 "constrain()/constraint()/box() requires at least one of min/lower/max/upper: {raw}"
2805 ),
2806 }
2807 .into());
2808 }
2809 return Ok(ParsedTerm::Linear {
2810 name: vars[0].clone(),
2811 explicit: true,
2812 double_penalty: option_bool_strict(&options, "double_penalty")?
2813 .unwrap_or(false),
2814 coefficient_min,
2815 coefficient_max,
2816 });
2817 }
2818 "nonnegative" | "nonnegative_coef" => {
2819 if vars.len() != 1 {
2820 return Err(FormulaDslError::InvalidArgument {
2821 reason: format!("nonnegative() expects exactly one variable: {raw}"),
2822 }
2823 .into());
2824 }
2825 validate_known_term_options("nonnegative", &options, &["double_penalty"], raw)?;
2826 return Ok(ParsedTerm::Linear {
2827 name: vars[0].clone(),
2828 explicit: true,
2829 double_penalty: option_bool_strict(&options, "double_penalty")?
2830 .unwrap_or(false),
2831 coefficient_min: Some(0.0),
2832 coefficient_max: None,
2833 });
2834 }
2835 "nonpositive" | "nonpositive_coef" => {
2836 if vars.len() != 1 {
2837 return Err(FormulaDslError::InvalidArgument {
2838 reason: format!("nonpositive() expects exactly one variable: {raw}"),
2839 }
2840 .into());
2841 }
2842 validate_known_term_options("nonpositive", &options, &["double_penalty"], raw)?;
2843 return Ok(ParsedTerm::Linear {
2844 name: vars[0].clone(),
2845 explicit: true,
2846 double_penalty: option_bool_strict(&options, "double_penalty")?
2847 .unwrap_or(false),
2848 coefficient_min: None,
2849 coefficient_max: Some(0.0),
2850 });
2851 }
2852 "bounded" => {
2853 if vars.len() != 1 {
2854 return Err(FormulaDslError::InvalidArgument {
2855 reason: format!("bounded() expects exactly one variable: {raw}"),
2856 }
2857 .into());
2858 }
2859 validate_known_term_options(
2860 "bounded",
2861 &options,
2862 &[
2863 "min",
2864 "max",
2865 "prior",
2866 "pull",
2867 "target",
2868 "strength",
2869 "double_penalty",
2870 ],
2871 raw,
2872 )?;
2873 let min = parse_required_f64_option(&options, "min", raw)?;
2874 let max = parse_required_f64_option(&options, "max", raw)?;
2875 if !min.is_finite() || !max.is_finite() || min >= max {
2876 return Err(FormulaDslError::InvalidArgument {
2877 reason: format!(
2878 "bounded() requires finite min < max, got min={min}, max={max}: {raw}"
2879 ),
2880 }
2881 .into());
2882 }
2883 let prior = parse_bounded_priorspec(&options, min, max, raw)?;
2884 return Ok(ParsedTerm::BoundedLinear {
2885 name: vars[0].clone(),
2886 min,
2887 max,
2888 prior,
2889 double_penalty: option_bool_strict(&options, "double_penalty")?
2897 .unwrap_or(false),
2898 });
2899 }
2900 "group" | "re" | "factor" => {
2901 if vars.len() != 1 {
2902 return Err(FormulaDslError::InvalidArgument {
2903 reason: format!(
2904 "{name}() expects exactly one variable, got '{}': {raw}",
2905 vars.join(",")
2906 ),
2907 }
2908 .into());
2909 }
2910 let lenient_unseen = name != "factor";
2918 return Ok(ParsedTerm::RandomEffect {
2919 name: vars[0].clone(),
2920 lenient_unseen,
2921 });
2922 }
2923 "tensor" | "interaction" | "te" => {
2924 if vars.len() < 2 {
2925 return Err(FormulaDslError::InvalidArgument {
2926 reason: format!(
2927 "tensor()/interaction()/te() requires at least two variables: {raw}"
2928 ),
2929 }
2930 .into());
2931 }
2932 return Ok(ParsedTerm::Smooth {
2933 label: raw.to_string(),
2934 vars,
2935 kind: SmoothKind::Te,
2936 options,
2937 });
2938 }
2939 "t2" => {
2940 if vars.len() < 2 {
2941 return Err(FormulaDslError::InvalidArgument {
2942 reason: format!("t2() requires at least two variables: {raw}"),
2943 }
2944 .into());
2945 }
2946 return Ok(ParsedTerm::Smooth {
2947 label: raw.to_string(),
2948 vars,
2949 kind: SmoothKind::T2,
2950 options,
2951 });
2952 }
2953 "ti" => {
2954 if vars.len() < 2 {
2961 return Err(FormulaDslError::InvalidArgument {
2962 reason: format!("ti() requires at least two variables: {raw}"),
2963 }
2964 .into());
2965 }
2966 return Ok(ParsedTerm::Smooth {
2967 label: raw.to_string(),
2968 vars,
2969 kind: SmoothKind::Ti,
2970 options,
2971 });
2972 }
2973 "fs" | "sz" => {
2974 if vars.len() != 2 {
2975 return Err(format!("{}() expects exactly two variables: {raw}", name));
2976 }
2977 options.insert("bs".to_string(), name.clone());
2978 return Ok(ParsedTerm::Smooth {
2979 label: raw.to_string(),
2980 vars,
2981 kind: SmoothKind::S,
2982 options,
2983 });
2984 }
2985 "thinplate" | "thin_plate" | "tps" => {
2986 if vars.len() < 2 {
2987 return Err(FormulaDslError::InvalidArgument {
2988 reason: format!(
2989 "thinplate()/thin_plate()/tps() requires at least two variables: {raw}"
2990 ),
2991 }
2992 .into());
2993 }
2994 options.insert("type".to_string(), "tps".to_string());
2995 return Ok(ParsedTerm::Smooth {
2996 label: raw.to_string(),
2997 vars,
2998 kind: SmoothKind::S,
2999 options,
3000 });
3001 }
3002 "smooth" | "s" | "cyclic" | "periodic" | "cc" | "cp" => {
3003 if vars.is_empty() {
3004 return Err(FormulaDslError::InvalidArgument {
3005 reason: format!("smooth()/s() requires at least one variable: {raw}"),
3006 }
3007 .into());
3008 }
3009 let bs_is_re = options
3015 .get("bs")
3016 .or_else(|| options.get("type"))
3017 .map(|v| {
3018 v.trim()
3019 .trim_matches(|c| c == '\'' || c == '"')
3020 .to_ascii_lowercase()
3021 })
3022 .as_deref()
3023 == Some("re");
3024 if bs_is_re && vars.len() == 1 {
3025 return Ok(ParsedTerm::RandomEffect {
3028 name: vars[0].clone(),
3029 lenient_unseen: true,
3030 });
3031 }
3032 if matches!(name.as_str(), "cyclic" | "periodic" | "cc" | "cp") {
3033 options.insert("type".to_string(), "cyclic".to_string());
3034 }
3035 if matches!(name.as_str(), "fs" | "sz") {
3036 options.insert("bs".to_string(), name.clone());
3037 }
3038 return Ok(ParsedTerm::Smooth {
3039 label: raw.to_string(),
3040 vars,
3041 kind: SmoothKind::S,
3042 options,
3043 });
3044 }
3045 "sphere" | "sos" | "spherical" | "s2" => {
3046 if vars.len() != 2 {
3053 return Err(FormulaDslError::InvalidArgument {
3054 reason: format!(
3055 "{name}() expects exactly two variables: latitude and longitude; got {} in {raw}",
3056 vars.len()
3057 ),
3058 }
3059 .into());
3060 }
3061 options.insert("type".to_string(), "sphere".to_string());
3062 return Ok(ParsedTerm::Smooth {
3063 label: raw.to_string(),
3064 vars,
3065 kind: SmoothKind::S,
3066 options,
3067 });
3068 }
3069 "mjs" | "measurejet" | "measure_jet" | "web" => {
3070 if vars.is_empty() {
3076 return Err(FormulaDslError::InvalidArgument {
3077 reason: format!("{name}() requires at least one variable: {raw}"),
3078 }
3079 .into());
3080 }
3081 options.insert("type".to_string(), "measurejet".to_string());
3082 return Ok(ParsedTerm::Smooth {
3083 label: raw.to_string(),
3084 vars,
3085 kind: SmoothKind::S,
3086 options,
3087 });
3088 }
3089 "curv" | "curvature" | "constant_curvature" | "mkappa" => {
3090 if vars.is_empty() {
3096 return Err(FormulaDslError::InvalidArgument {
3097 reason: format!("{name}() requires at least one variable: {raw}"),
3098 }
3099 .into());
3100 }
3101 options.insert("type".to_string(), "curvature".to_string());
3102 return Ok(ParsedTerm::Smooth {
3103 label: raw.to_string(),
3104 vars,
3105 kind: SmoothKind::S,
3106 options,
3107 });
3108 }
3109 "matern" => {
3110 if vars.is_empty() {
3111 return Err(FormulaDslError::InvalidArgument {
3112 reason: format!("matern() requires at least one variable: {raw}"),
3113 }
3114 .into());
3115 }
3116 options.insert("type".to_string(), "matern".to_string());
3117 return Ok(ParsedTerm::Smooth {
3118 label: raw.to_string(),
3119 vars,
3120 kind: SmoothKind::S,
3121 options,
3122 });
3123 }
3124 "duchon" => {
3125 if vars.is_empty() {
3126 return Err(FormulaDslError::InvalidArgument {
3127 reason: format!("duchon() requires at least one variable: {raw}"),
3128 }
3129 .into());
3130 }
3131 if option_bool(&options, "cyclic").unwrap_or(false)
3132 || option_bool(&options, "periodic").unwrap_or(false)
3133 {
3134 options.insert("cyclic".to_string(), "true".to_string());
3135 }
3136 options.insert("type".to_string(), "duchon".to_string());
3137 return Ok(ParsedTerm::Smooth {
3138 label: raw.to_string(),
3139 vars,
3140 kind: SmoothKind::S,
3141 options,
3142 });
3143 }
3144 "pca" => {
3145 if vars.is_empty() {
3146 return Err(FormulaDslError::InvalidArgument {
3147 reason: format!("pca() requires at least one variable: {raw}"),
3148 }
3149 .into());
3150 }
3151 options.insert("type".to_string(), "pca".to_string());
3152 return Ok(ParsedTerm::Smooth {
3153 label: raw.to_string(),
3154 vars,
3155 kind: SmoothKind::S,
3156 options,
3157 });
3158 }
3159 "linkwiggle" => {
3160 if !vars.is_empty() {
3161 return Err(FormulaDslError::InvalidArgument {
3162 reason: format!(
3163 "linkwiggle() takes named options only; positional args are not supported: {raw}"
3164 ),
3165 }
3166 .into());
3167 }
3168 return Ok(ParsedTerm::LinkWiggle { options });
3169 }
3170 "timewiggle" => {
3171 if !vars.is_empty() {
3172 return Err(FormulaDslError::InvalidArgument {
3173 reason: format!(
3174 "timewiggle() takes named options only; positional args are not supported: {raw}"
3175 ),
3176 }
3177 .into());
3178 }
3179 return Ok(ParsedTerm::TimeWiggle { options });
3180 }
3181 "link" => {
3182 if !vars.is_empty() {
3183 return Err(FormulaDslError::InvalidArgument {
3184 reason: format!(
3185 "link() takes named options only; positional args are not supported: {raw}"
3186 ),
3187 }
3188 .into());
3189 }
3190 return Ok(ParsedTerm::LinkConfig { options });
3191 }
3192 "survmodel" => {
3193 if !vars.is_empty() {
3194 return Err(FormulaDslError::InvalidArgument {
3195 reason: format!(
3196 "survmodel() takes named options only; positional args are not supported: {raw}"
3197 ),
3198 }
3199 .into());
3200 }
3201 return Ok(ParsedTerm::SurvivalConfig { options });
3202 }
3203 "logslope" | "log_slope" | "log_slope_surface" => {
3204 validate_known_term_options("logslope", &options, &[], raw)?;
3205 if vars.len() < 2 {
3206 return Err(FormulaDslError::InvalidArgument {
3207 reason: format!(
3208 "logslope() expects a z column followed by one or more RHS terms; add one logslope(z, ...) declaration per vector-z coordinate: {raw}"
3209 ),
3210 }
3211 .into());
3212 }
3213 let z_column = vars[0].trim();
3214 if !is_exact_ident(z_column) {
3215 return Err(FormulaDslError::InvalidArgument {
3216 reason: format!(
3217 "logslope() z column must be a bare column name, got `{z_column}` in {raw}"
3218 ),
3219 }
3220 .into());
3221 }
3222 let rhs = vars[1..].join(" + ");
3223 let parsed = parse_formula(&format!("__logslope__ ~ {rhs}"))?;
3224 if !parsed.logslope_surfaces.is_empty() {
3225 return Err(FormulaDslError::IncompatibleTerm {
3226 reason: format!(
3227 "logslope() declarations cannot be nested inside another logslope(): {raw}"
3228 ),
3229 }
3230 .into());
3231 }
3232 validate_auxiliary_formula_controls(&parsed, "logslope()")?;
3233 return Ok(ParsedTerm::LogSlopeSurface {
3234 z_column: z_column.to_string(),
3235 terms: parsed.terms,
3236 });
3237 }
3238 "linear" => {
3239 if vars.len() != 1 {
3240 return Err(FormulaDslError::InvalidArgument {
3241 reason: format!("linear() expects exactly one variable: {raw}"),
3242 }
3243 .into());
3244 }
3245 validate_known_term_options(
3246 "linear",
3247 &options,
3248 &["min", "lower", "max", "upper", "double_penalty"],
3249 raw,
3250 )?;
3251 let (coefficient_min, coefficient_max) =
3252 parse_linear_constraint_bounds(&options, raw)?;
3253 let double_penalty =
3254 option_bool_strict(&options, "double_penalty")?.unwrap_or(false);
3255 if vars[0].contains(':') {
3256 if coefficient_min.is_some() || coefficient_max.is_some() {
3257 return Err(FormulaDslError::IncompatibleTerm {
3258 reason: format!(
3259 "linear() coefficient bounds are not supported on an interaction: {raw}"
3260 ),
3261 }
3262 .into());
3263 }
3264 let mut interaction_vars = vars[0]
3265 .split(':')
3266 .map(str::trim)
3267 .map(str::to_string)
3268 .collect::<Vec<_>>();
3269 if interaction_vars.len() < 2
3270 || interaction_vars.iter().any(|var| !is_exact_ident(var))
3271 {
3272 return Err(FormulaDslError::InvalidArgument {
3273 reason: format!(
3274 "linear() interaction must contain at least two bare column names: {raw}"
3275 ),
3276 }
3277 .into());
3278 }
3279 interaction_vars.sort();
3280 let original_len = interaction_vars.len();
3281 interaction_vars.dedup();
3282 if interaction_vars.len() != original_len {
3283 return Err(FormulaDslError::IncompatibleTerm {
3284 reason: format!(
3285 "linear() interaction references the same variable more than once: {raw}"
3286 ),
3287 }
3288 .into());
3289 }
3290 return Ok(ParsedTerm::Interaction {
3291 vars: interaction_vars,
3292 double_penalty,
3293 });
3294 }
3295 return Ok(ParsedTerm::Linear {
3296 name: vars[0].clone(),
3297 explicit: true,
3298 double_penalty,
3299 coefficient_min,
3300 coefficient_max,
3301 });
3302 }
3303 _ => {
3304 return Err(format!(
3305 "unknown term function `{name}` in '{raw}'. Supported: bounded(), linear(), constrain()/constraint()/box(), nonnegative(), nonpositive(), smooth()/s(), cyclic()/periodic()/cc()/cp(), thinplate()/thin_plate()/tps(), tensor()/interaction()/te(), t2(), ti(), fs(), sz(), group()/re()/factor(), sphere()/sos()/spherical(), s2(), matern(), duchon(), pca(), logslope()/log_slope(), linkwiggle(), timewiggle(), link(), survmodel()"
3306 ));
3307 }
3308 }
3309 }
3310
3311 let ident = raw.trim();
3312 if !is_exact_ident(ident) {
3313 return Err(FormulaDslError::UnknownIdentifier {
3314 reason: format!("unsupported top-level RHS term: {raw}"),
3315 }
3316 .into());
3317 }
3318
3319 Ok(ParsedTerm::Linear {
3320 name: ident.to_string(),
3321 explicit: false,
3322 double_penalty: false,
3323 coefficient_min: None,
3324 coefficient_max: None,
3325 })
3326}
3327
3328pub fn parse_link_choice(
3333 raw: Option<&str>,
3334 flexible_flag: bool,
3335) -> Result<Option<LinkChoice>, FormulaDslError> {
3336 if raw.is_none() && !flexible_flag {
3337 return Ok(None);
3338 }
3339 let Some(v) = raw else {
3340 return Ok(Some(LinkChoice {
3341 mode: LinkMode::Flexible,
3342 link: LinkFunction::Probit,
3343 mixture_components: None,
3344 }));
3345 };
3346 let t = v.trim().to_ascii_lowercase();
3347 if let Some(inner) = t
3348 .strip_prefix("flexible(")
3349 .and_then(|s| s.strip_suffix(')'))
3350 {
3351 if let Some(components_inner) = inner
3352 .strip_prefix("blended(")
3353 .and_then(|s| s.strip_suffix(')'))
3354 .or_else(|| {
3355 inner
3356 .strip_prefix("mixture(")
3357 .and_then(|s| s.strip_suffix(')'))
3358 })
3359 {
3360 parse_link_component_list(components_inner)?;
3361 return Err(FormulaDslError::IncompatibleTerm {
3362 reason:
3363 "flexible(...) does not support blended(...)/mixture(...) links; wiggle is only supported for jointly fit standard links"
3364 .to_string(),
3365 });
3366 }
3367 let link = parse_linkname(inner)?;
3368 if !linkname_supports_joint_wiggle(link) {
3369 return Err(FormulaDslError::IncompatibleTerm {
3370 reason:
3371 "flexible(...) does not support sas/beta-logistic links; wiggle is only supported for jointly fit standard links"
3372 .to_string(),
3373 });
3374 }
3375 return Ok(Some(LinkChoice {
3376 mode: LinkMode::Flexible,
3377 link,
3378 mixture_components: None,
3379 }));
3380 }
3381 if let Some(inner) = t
3382 .strip_prefix("blended(")
3383 .and_then(|s| s.strip_suffix(')'))
3384 .or_else(|| t.strip_prefix("mixture(").and_then(|s| s.strip_suffix(')')))
3385 {
3386 if flexible_flag {
3387 return Err(FormulaDslError::IncompatibleTerm {
3388 reason:
3389 "--flexible-link cannot be combined with --link blended(...)/mixture(...); blended inverse links are not flexible-link mode"
3390 .to_string(),
3391 });
3392 }
3393 let components = parse_link_component_list(inner)?;
3394 return Ok(Some(LinkChoice {
3395 mode: LinkMode::Strict,
3396 link: LinkFunction::Logit,
3397 mixture_components: Some(components),
3398 }));
3399 }
3400
3401 let link = parse_linkname(&t)?;
3402 if flexible_flag && !linkname_supports_joint_wiggle(link) {
3403 return Err(FormulaDslError::IncompatibleTerm {
3404 reason:
3405 "--flexible-link does not support sas/beta-logistic links; wiggle is only supported for jointly fit standard links"
3406 .to_string(),
3407 });
3408 }
3409 Ok(Some(LinkChoice {
3410 mode: if flexible_flag {
3411 LinkMode::Flexible
3412 } else {
3413 LinkMode::Strict
3414 },
3415 link,
3416 mixture_components: None,
3417 }))
3418}
3419
3420pub fn parse_linkname(v: &str) -> Result<LinkFunction, FormulaDslError> {
3421 match v.trim() {
3422 "identity" => Ok(LinkFunction::Identity),
3423 "log" => Ok(LinkFunction::Log),
3424 "logit" | "binomial-logit" => Ok(LinkFunction::Logit),
3425 "probit" | "binomial-probit" => Ok(LinkFunction::Probit),
3426 "cloglog" | "binomial-cloglog" => Ok(LinkFunction::CLogLog),
3427 "loglog" => Ok(LinkFunction::LogLog),
3428 "cauchit" => Ok(LinkFunction::Cauchit),
3429 "sas" => Ok(LinkFunction::Sas),
3430 "beta-logistic" => Ok(LinkFunction::BetaLogistic),
3431 other => Err(FormulaDslError::UnknownIdentifier {
3432 reason: format!(
3433 "unsupported link type '{other}'; \
3434 use one of identity|log|logit|probit|cloglog|loglog|cauchit|binomial-logit|binomial-probit|binomial-cloglog|sas|beta-logistic|blended(...)/mixture(...) or flexible(...). \
3435 Both `--link <type>` (CLI flag) and `link(type=<type>)` (formula term) accept the same set."
3436 ),
3437 }),
3438 }
3439}
3440
3441pub fn parse_link_component(v: &str) -> Result<LinkComponent, String> {
3442 match v.trim() {
3443 "logit" => Ok(LinkComponent::Logit),
3444 "probit" => Ok(LinkComponent::Probit),
3445 "cloglog" => Ok(LinkComponent::CLogLog),
3446 "loglog" => Ok(LinkComponent::LogLog),
3447 "cauchit" => Ok(LinkComponent::Cauchit),
3448 other => Err(FormulaDslError::UnknownIdentifier {
3449 reason: format!(
3450 "unsupported blended-link component '{other}'; use probit|logit|cloglog|loglog|cauchit"
3451 ),
3452 }
3453 .into()),
3454 }
3455}
3456
3457pub fn parse_link_component_list(v: &str) -> Result<Vec<LinkComponent>, String> {
3458 let mut out = Vec::new();
3459 for part in v.split(',') {
3460 let trimmed = part.trim();
3461 if trimmed.is_empty() {
3462 continue;
3463 }
3464 let comp = parse_link_component(trimmed)?;
3465 if out.contains(&comp) {
3466 return Err(FormulaDslError::IncompatibleTerm {
3467 reason: "blended(...) cannot contain duplicate components".to_string(),
3468 }
3469 .into());
3470 }
3471 out.push(comp);
3472 }
3473 if out.len() < 2 {
3474 return Err(FormulaDslError::InvalidArgument {
3475 reason: "blended(...) requires at least two components".to_string(),
3476 }
3477 .into());
3478 }
3479 Ok(out)
3480}