Merge pull request #7 from alexhallam/6-add-intercept-and-have-a-new-structure-1

fix issue

Author: alexhallam

CHANGELOG.md

@@ -1,5 +1,37 @@
 # Changelog
 
+## [Unreleased]
+
+### ✨ Added
+
+- **Identity Role for Plain Terms**: Added new `Identity` role to `VariableRole` enum for variables that appear as plain terms in formulas (e.g., `x` in `y ~ x`)
+- **Intercept Column Support**: Added automatic inclusion of `"intercept"` column in `all_generated_columns` when `has_intercept` is true
+- **Formula Order Mapping**: Added new `all_generated_columns_formula_order` field that maps formula order (1, 2, 3...) to column names
+- **Comprehensive Test Suite**: Added 4 new unit tests to verify intercept handling and formula order mapping functionality
+
+### 🔧 Improved
+
+- **Variable Role Assignment**: Plain terms now correctly receive `Identity` role instead of `FixedEffect` role
+- **Generated Columns Preservation**: Variables with `Identity` role now preserve their original column name in generated columns list
+- **Intercept Positioning**: Intercept column is automatically inserted at index 1 (after response variable) in `all_generated_columns`
+
+### 🐛 Fixed
+
+- **Issue #4**: Fixed plain terms not receiving proper `Identity` role when appearing alone in formulas
+- **Issue #6**: Fixed missing intercept column in `all_generated_columns` and added formula order mapping
+
+### 🧪 Testing
+
+- **Unit Test Coverage**: Added comprehensive test coverage for intercept and formula order functionality
+- **Regression Prevention**: Tests ensure intercept is present when `has_intercept` is true and absent when false
+- **Order Validation**: Tests verify correct column ordering in both `all_generated_columns` and `all_generated_columns_formula_order`
+
+### 🔄 Internal Changes
+
+- **MetaBuilder Enhancement**: Updated `build()` method to handle intercept insertion and formula order mapping
+- **Data Structure Updates**: Enhanced `FormulaMetaData` struct with new `all_generated_columns_formula_order` field
+- **Role Management**: Improved role assignment logic in `push_plain_term()` and `add_transformation()` methods
+
 ## [0.2.4] - 2025-09-05
 
 ### Added

src/internal/data_structures.rs

@@ -448,7 +448,15 @@ pub struct FormulaMetadataInfo {
 ///         family: Some("gaussian".to_string())
 ///     },
 ///     columns,
-///     all_generated_columns: vec!["y".to_string(), "x".to_string(), "group".to_string()]
+///     all_generated_columns: vec!["y".to_string(), "intercept".to_string(), "x".to_string(), "group".to_string()],
+///     all_generated_columns_formula_order: {
+///         let mut map = HashMap::new();
+///         map.insert("1".to_string(), "y".to_string());
+///         map.insert("2".to_string(), "intercept".to_string());
+///         map.insert("3".to_string(), "x".to_string());
+///         map.insert("4".to_string(), "group".to_string());
+///         map
+///     }
 /// };
 /// ```
 #[derive(Debug, Serialize, Deserialize, Clone)]
@@ -462,8 +470,31 @@ pub struct FormulaMetaData {
     /// Detailed information about each variable
     pub columns: HashMap<String, VariableInfo>,
 
-    /// All generated column names ordered by variable ID
+    /// All generated column names ordered by variable ID, including intercept if present
     pub all_generated_columns: Vec<String>,
+
+    /// Mapping of formula order to column names
+    ///
+    /// This field provides a mapping from formula order (as string keys "1", "2", etc.)
+    /// to the corresponding column names. The order follows the formula structure:
+    /// 1. Response variable
+    /// 2. Intercept (if present)
+    /// 3. Variables in order of appearance in the formula
+    ///
+    /// # Examples
+    ///
+    /// For formula `y ~ x + poly(x, 2) + log(z)`:
+    /// ```json
+    /// {
+    ///   "1": "y",
+    ///   "2": "intercept",
+    ///   "3": "x",
+    ///   "4": "x_poly_1",
+    ///   "5": "x_poly_2",
+    ///   "6": "z_log"
+    /// }
+    /// ```
+    pub all_generated_columns_formula_order: HashMap<String, String>,
 }
 
 // Legacy structures for backward compatibility

src/internal/meta_builder.rs

@@ -234,7 +234,7 @@ impl MetaBuilder {
     pub fn add_transformation(&mut self, name: &str, transformation: Transformation) {
         if let Some(var_info) = self.columns.get_mut(name) {
             var_info.transformations.push(transformation.clone());
-            
+
             // If the variable has an Identity role, preserve the original variable name
             // and add the transformation's generated columns
             if var_info.roles.contains(&VariableRole::Identity) {
@@ -670,10 +670,47 @@ impl MetaBuilder {
         let mut sorted_vars: Vec<_> = self.columns.values().collect();
         sorted_vars.sort_by_key(|v| v.id);
 
-        for var in sorted_vars {
+        for var in &sorted_vars {
             all_generated_columns.extend(var.generated_columns.clone());
         }
 
+        // Add intercept column if has_intercept is true
+        if has_intercept {
+            all_generated_columns.insert(1, "intercept".to_string()); // Insert after response (index 1)
+        }
+
+        // Generate all_generated_columns_formula_order mapping
+        let mut all_generated_columns_formula_order = std::collections::HashMap::new();
+        let mut order_index = 1;
+
+        // Add response variable (always first)
+        if let Some(response_var) = sorted_vars.iter().find(|v| v.id == 1) {
+            if let Some(response_col) = response_var.generated_columns.first() {
+                all_generated_columns_formula_order
+                    .insert(order_index.to_string(), response_col.clone());
+                order_index += 1;
+            }
+        }
+
+        // Add intercept if present
+        if has_intercept {
+            all_generated_columns_formula_order
+                .insert(order_index.to_string(), "intercept".to_string());
+            order_index += 1;
+        }
+
+        // Add all other variables in order
+        for var in &sorted_vars {
+            if var.id != 1 {
+                // Skip response (already added)
+                for col in &var.generated_columns {
+                    all_generated_columns_formula_order
+                        .insert(order_index.to_string(), col.clone());
+                    order_index += 1;
+                }
+            }
+        }
+
         crate::internal::data_structures::FormulaMetaData {
             formula: input.to_string(),
             metadata: FormulaMetadataInfo {
@@ -684,6 +721,7 @@ impl MetaBuilder {
             },
             columns: self.columns,
             all_generated_columns,
+            all_generated_columns_formula_order,
         }
     }
 }