Merge pull request #13 from alexhallam/11-make-c-for-categorical-data

alexhallam · web-flow · commit 33861cfedda9 · 2025-09-09T21:31:55.000-04:00
update with categoricals
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -1,6 +1,6 @@
 # Changelog
 
-## [0.2.6] - 2024-12-19
+## [0.2.7] - 2024-12-19
 
 ### ✨ Added
 
@@ -9,6 +9,10 @@
 - **Comprehensive Interaction Metadata**: Each variable now includes detailed interaction information showing all interactions it participates in
 - **Canonical Expansion**: Multi-way interactions now generate all possible combinations (2-way, 3-way, etc.) as per R/Wilkinson notation standards
 - **Test Example**: Added `test_multiway_interaction.rs` example demonstrating 2-way, 3-way, and 4-way interactions
+- **Categorical Function Support**: Added `c()` function for categorical variables with reference level specification
+- **Named Arguments**: Added support for named arguments in function calls (e.g., `ref=treatment`)
+- **Categorical Role**: Added new `Categorical` role to `VariableRole` enum for categorical variables
+- **Test Example**: Added `test_categorical_function.rs` example demonstrating categorical function usage
 
 ### 🔧 Improved
 
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/Cargo.toml b/Cargo.toml
@@ -1,7 +1,7 @@
 [package]
 description = "High-performance modern Wilkinson's formula parsing for statistical models. Parses R-style formulas into structured JSON metadata supporting linear models, mixed effects, and complex statistical specifications."
 name = "fiasto"
-version = "0.2.6"
+version = "0.2.7"
 edition = "2021"
 authors = ["Alex Hallam <alexhallam6.28@gmail.com>"]
 license = "MIT"
diff --git a/examples/test_categorical_function.rs b/examples/test_categorical_function.rs
@@ -0,0 +1,85 @@
+use fiasto::parse_formula;
+
+fn main() -> Result<(), Box<dyn std::error::Error>> {
+    println!("Testing Categorical Function c() in Fiasto");
+    println!("==========================================");
+    println!();
+
+    // Test 1: Basic categorical function without reference level
+    println!("=== Test 1: Basic categorical function ===");
+    let formula1 = "y ~ c(treatment)";
+    println!("Formula: {}", formula1);
+    match parse_formula(formula1) {
+        Ok(result) => {
+            println!("✓ Parsed successfully!");
+            println!("{}", serde_json::to_string_pretty(&result)?);
+        }
+        Err(e) => {
+            println!("✗ Error parsing formula: {}", e);
+        }
+    }
+    println!();
+
+    // Test 2: Categorical function with reference level (unquoted)
+    println!("=== Test 2: Categorical with reference level (unquoted) ===");
+    let formula2 = "y ~ c(treatment, ref=control)";
+    println!("Formula: {}", formula2);
+    match parse_formula(formula2) {
+        Ok(result) => {
+            println!("✓ Parsed successfully!");
+            println!("{}", serde_json::to_string_pretty(&result)?);
+        }
+        Err(e) => {
+            println!("✗ Error parsing formula: {}", e);
+        }
+    }
+    println!();
+
+    // Test 3: Categorical function with reference level (quoted)
+    println!("=== Test 3: Categorical with reference level (quoted) ===");
+    let formula3 = r#"y ~ c(group, ref="group1")"#;
+    println!("Formula: {}", formula3);
+    match parse_formula(formula3) {
+        Ok(result) => {
+            println!("✓ Parsed successfully!");
+            println!("{}", serde_json::to_string_pretty(&result)?);
+        }
+        Err(e) => {
+            println!("✗ Error parsing formula: {}", e);
+        }
+    }
+    println!();
+
+    // Test 4: Categorical function with other variables
+    println!("=== Test 4: Categorical with other variables ===");
+    let formula4 = "y ~ x1 + c(treatment, ref=control) + x2";
+    println!("Formula: {}", formula4);
+    match parse_formula(formula4) {
+        Ok(result) => {
+            println!("✓ Parsed successfully!");
+            println!("{}", serde_json::to_string_pretty(&result)?);
+        }
+        Err(e) => {
+            println!("✗ Error parsing formula: {}", e);
+        }
+    }
+    println!();
+
+    // Test 5: Multiple categorical variables
+    println!("=== Test 5: Multiple categorical variables ===");
+    let formula5 = "y ~ c(treatment, ref=control) + c(group, ref=\"group1\")";
+    println!("Formula: {}", formula5);
+    match parse_formula(formula5) {
+        Ok(result) => {
+            println!("✓ Parsed successfully!");
+            println!("{}", serde_json::to_string_pretty(&result)?);
+        }
+        Err(e) => {
+            println!("✗ Error parsing formula: {}", e);
+        }
+    }
+    println!();
+
+    println!("All tests completed!");
+    Ok(())
+}
diff --git a/examples/test_factor_function.rs b/examples/test_factor_function.rs
@@ -0,0 +1,85 @@
+use fiasto::parse_formula;
+
+fn main() -> Result<(), Box<dyn std::error::Error>> {
+    println!("Testing Factor Function factor() in Fiasto");
+    println!("==========================================");
+    println!();
+
+    // Test 1: Basic factor function without reference level
+    println!("=== Test 1: Basic factor function ===");
+    let formula1 = "y ~ factor(treatment)";
+    println!("Formula: {}", formula1);
+    match parse_formula(formula1) {
+        Ok(result) => {
+            println!("✓ Parsed successfully!");
+            println!("{}", serde_json::to_string_pretty(&result)?);
+        }
+        Err(e) => {
+            println!("✗ Error parsing formula: {}", e);
+        }
+    }
+    println!();
+
+    // Test 2: Factor function with reference level (unquoted)
+    println!("=== Test 2: Factor with reference level (unquoted) ===");
+    let formula2 = "y ~ factor(treatment, ref=control)";
+    println!("Formula: {}", formula2);
+    match parse_formula(formula2) {
+        Ok(result) => {
+            println!("✓ Parsed successfully!");
+            println!("{}", serde_json::to_string_pretty(&result)?);
+        }
+        Err(e) => {
+            println!("✗ Error parsing formula: {}", e);
+        }
+    }
+    println!();
+
+    // Test 3: Factor function with reference level (quoted)
+    println!("=== Test 3: Factor with reference level (quoted) ===");
+    let formula3 = r#"y ~ factor(group, ref="group1")"#;
+    println!("Formula: {}", formula3);
+    match parse_formula(formula3) {
+        Ok(result) => {
+            println!("✓ Parsed successfully!");
+            println!("{}", serde_json::to_string_pretty(&result)?);
+        }
+        Err(e) => {
+            println!("✗ Error parsing formula: {}", e);
+        }
+    }
+    println!();
+
+    // Test 4: Factor function with other variables
+    println!("=== Test 4: Factor with other variables ===");
+    let formula4 = "y ~ x1 + factor(treatment, ref=control) + x2";
+    println!("Formula: {}", formula4);
+    match parse_formula(formula4) {
+        Ok(result) => {
+            println!("✓ Parsed successfully!");
+            println!("{}", serde_json::to_string_pretty(&result)?);
+        }
+        Err(e) => {
+            println!("✗ Error parsing formula: {}", e);
+        }
+    }
+    println!();
+
+    // Test 5: Compare factor() and c() functions
+    println!("=== Test 5: Compare factor() and c() functions ===");
+    let formula5 = "y ~ factor(treatment, ref=control) + c(group, ref=\"group1\")";
+    println!("Formula: {}", formula5);
+    match parse_formula(formula5) {
+        Ok(result) => {
+            println!("✓ Parsed successfully!");
+            println!("{}", serde_json::to_string_pretty(&result)?);
+        }
+        Err(e) => {
+            println!("✗ Error parsing formula: {}", e);
+        }
+    }
+    println!();
+
+    println!("All tests completed!");
+    Ok(())
+}
diff --git a/src/internal/ast.rs b/src/internal/ast.rs
@@ -229,6 +229,13 @@ pub enum Argument {
     /// - `"group_id"` → `Argument::String("group_id")`
     String(String),
 
+    /// A named argument (key=value)
+    ///
+    /// # Examples
+    /// - `ref=treatment` → `Argument::Named("ref", "treatment")`
+    /// - `level=high` → `Argument::Named("level", "high")`
+    Named(String, String),
+
     /// A boolean value
     ///
     /// # Examples
diff --git a/src/internal/data_structures.rs b/src/internal/data_structures.rs
@@ -112,6 +112,13 @@ pub enum VariableRole {
     /// - `x1_x2` for interaction `x1:x2`
     /// - `x1_x2_x3` for interaction `x1:x2:x3`
     InteractionTerm,
+
+    /// A categorical variable with reference level specification
+    ///
+    /// # Examples
+    /// - `c(treatment, ref=control)` for categorical treatment with control as reference
+    /// - `c(group, ref="group1")` for categorical group with "group1" as reference
+    Categorical,
 }
 
 /// A transformation applied to a variable
diff --git a/src/internal/lexer.rs b/src/internal/lexer.rs
@@ -211,6 +211,10 @@ pub enum Token {
     #[token("factor")]
     Factor,
 
+    /// Categorical variable with reference level: `c(x, ref=level)`
+    #[token("c", priority = 3)]
+    C,
+
     /// Scaling transformation: `scale(x)`
     #[token("scale")]
     Scale,
diff --git a/src/internal/meta_builder.rs b/src/internal/meta_builder.rs
@@ -486,6 +486,12 @@ impl MetaBuilder {
 
     /// Adds a function/transformation term
     pub fn push_function_term(&mut self, fname: &str, args: &[Argument]) {
+        // Special handling for categorical functions
+        if fname == "c" || fname == "factor" {
+            self.push_categorical_term_with_name(fname, args);
+            return;
+        }
+
         let base_ident = args.iter().find_map(|a| match a {
             Argument::Ident(s) => Some(s.as_str()),
             _ => None,
@@ -510,6 +516,47 @@ impl MetaBuilder {
         }
     }
 
+    /// Handles categorical variables with reference level specification
+    fn push_categorical_term_with_name(&mut self, fname: &str, args: &[Argument]) {
+        // Extract the variable name (first argument)
+        let var_name = args.iter().find_map(|a| match a {
+            Argument::Ident(s) => Some(s.as_str()),
+            _ => None,
+        });
+
+        if let Some(var_name) = var_name {
+            self.ensure_variable(var_name);
+
+            // Add both Categorical and FixedEffect roles
+            self.add_role(var_name, VariableRole::Categorical);
+            self.add_role(var_name, VariableRole::FixedEffect);
+
+            // Extract reference level from named arguments
+            let ref_level = args.iter().find_map(|a| match a {
+                Argument::Named(key, value) if key == "ref" => Some(value.clone()),
+                _ => None,
+            });
+
+            // Create transformation info with reference level
+            let mut parameters = self.extract_function_parameters(fname, args);
+            if let Some(ref_level) = ref_level {
+                if let serde_json::Value::Object(ref mut params_map) = parameters {
+                    params_map.insert("ref".to_string(), serde_json::Value::String(ref_level));
+                }
+            }
+
+            let generates_columns = self.generate_transformation_columns(fname, args);
+
+            let transformation = Transformation {
+                function: fname.to_string(),
+                parameters,
+                generates_columns,
+            };
+
+            self.add_transformation(var_name, transformation);
+        }
+    }
+
     /// Handles random effects with variable-centric approach
     pub fn push_random_effect(&mut self, random_effect: &RandomEffect) {
         self.is_random_effects_model = true;
@@ -674,6 +721,10 @@ impl MetaBuilder {
             "log" => {
                 // No additional parameters for log
             }
+            "factor" => {
+                // Handle factor function parameters (same as c function)
+                // Parameters are handled by the generic case below
+            }
             _ => {
                 // Generic parameter handling
                 for (i, arg) in args.iter().enumerate() {
@@ -683,6 +734,11 @@ impl MetaBuilder {
                         Argument::String(s) => serde_json::Value::String(s.clone()),
                         Argument::Boolean(b) => serde_json::Value::Bool(*b),
                         Argument::Ident(s) => serde_json::Value::String(s.clone()),
+                        Argument::Named(key, value) => {
+                            // For named arguments, use the key directly
+                            params.insert(key.clone(), serde_json::Value::String(value.clone()));
+                            continue; // Skip the generic arg_N handling
+                        }
                     };
                     params.insert(key, value);
                 }
@@ -713,6 +769,11 @@ impl MetaBuilder {
                 }
             }
             "log" => vec![format!("{}_log", base_name)],
+            "c" | "factor" => {
+                // For categorical variables, we generate dummy variables for each level
+                // The reference level is excluded (handled by the ref parameter)
+                vec![format!("{}_categorical", base_name)]
+            }
             _ => vec![format!("{}_{}", base_name, fname)],
         }
     }
diff --git a/src/internal/parse_arg.rs b/src/internal/parse_arg.rs
@@ -77,7 +77,41 @@ pub fn parse_arg<'a>(
         match tok {
             Token::ColumnName => {
                 crate::internal::next::next(tokens, pos);
-                Ok(Argument::Ident(slice.to_string()))
+
+                // Check if this is a named argument (key=value)
+                if crate::internal::peek::peek(tokens, *pos)
+                    .map(|(t, _)| matches!(t, Token::Equal))
+                    .unwrap_or(false)
+                {
+                    let key = slice.to_string();
+                    crate::internal::next::next(tokens, pos); // Skip the equals sign
+
+                    // Parse the value
+                    if let Some((value_tok, value_slice)) =
+                        crate::internal::peek::peek(tokens, *pos).cloned()
+                    {
+                        match value_tok {
+                            Token::ColumnName => {
+                                crate::internal::next::next(tokens, pos);
+                                Ok(Argument::Named(key, value_slice.to_string()))
+                            }
+                            Token::StringLiteral => {
+                                crate::internal::next::next(tokens, pos);
+                                // Remove quotes from string literal
+                                let value = value_slice.trim_matches('"').to_string();
+                                Ok(Argument::Named(key, value))
+                            }
+                            _ => Err(ParseError::Unexpected {
+                                expected: "column name or string literal",
+                                found: Some(value_tok),
+                            }),
+                        }
+                    } else {
+                        Err(ParseError::Eoi)
+                    }
+                } else {
+                    Ok(Argument::Ident(slice.to_string()))
+                }
             }
             Token::Integer => {
                 crate::internal::next::next(tokens, pos);
@@ -87,6 +121,12 @@ pub fn parse_arg<'a>(
                 crate::internal::next::next(tokens, pos);
                 Ok(Argument::Integer(1))
             }
+            Token::StringLiteral => {
+                crate::internal::next::next(tokens, pos);
+                // Remove quotes from string literal
+                let value = slice.trim_matches('"').to_string();
+                Ok(Argument::String(value))
+            }
             _ => Err(ParseError::Unexpected {
                 expected: "argument",
                 found: Some(tok),
diff --git a/src/internal/parse_term.rs b/src/internal/parse_term.rs
diff --git a/temp_test.rs b/temp_test.rs
