fix for column pollution

Author: alxlyj

python/src/robyn/modeling/ridge/ridge_data_builder.py

@@ -12,19 +12,32 @@ def __init__(self, mmm_data, featurized_mmm_data):
         self.logger = logging.getLogger(__name__)
 
     def _prepare_data(self, params: Dict[str, float]) -> Tuple[pd.DataFrame, pd.Series]:
-        # Get the dependent variable
-        # Check if 'dep_var' is in columns
-        if "dep_var" in self.featurized_mmm_data.dt_mod.columns:
-            # Rename 'dep_var' to the specified value
-            self.featurized_mmm_data.dt_mod = self.featurized_mmm_data.dt_mod.rename(
-                columns={"dep_var": self.mmm_data.mmmdata_spec.dep_var}
-            )
-        y = self.featurized_mmm_data.dt_mod[self.mmm_data.mmmdata_spec.dep_var]
-
-        # Select all columns except the dependent variable
-        X = self.featurized_mmm_data.dt_mod.drop(
-            columns=[self.mmm_data.mmmdata_spec.dep_var]
-        )
+        """Prepare data for ridge regression, handling dependent variable and excluding date columns"""
+        # Get the dependent variable, handling both possible column names
+        dep_var = self.mmm_data.mmmdata_spec.dep_var
+        if dep_var not in self.featurized_mmm_data.dt_mod.columns:
+            # If dep_var column doesn't exist, try 'dep_var'
+            if 'dep_var' in self.featurized_mmm_data.dt_mod.columns:
+                # Rename 'dep_var' to the specified value
+                self.featurized_mmm_data.dt_mod = self.featurized_mmm_data.dt_mod.rename(
+                    columns={"dep_var": dep_var}
+                )
+                y = self.featurized_mmm_data.dt_mod[dep_var]
+            else:
+                raise KeyError(f"Could not find dependent variable column. Expected either '{dep_var}' or 'dep_var' in columns: {self.featurized_mmm_data.dt_mod.columns.tolist()}")
+        else:
+            y = self.featurized_mmm_data.dt_mod[dep_var]
+
+        # Select all columns except the dependent variable and date columns
+        exclude_cols = ['ds']  # Always exclude 'ds' if present
+        if dep_var in self.featurized_mmm_data.dt_mod.columns:
+            exclude_cols.append(dep_var)
+        if 'dep_var' in self.featurized_mmm_data.dt_mod.columns:
+            exclude_cols.append('dep_var')
+        
+        # Only drop columns that actually exist in the dataframe
+        exclude_cols = [col for col in exclude_cols if col in self.featurized_mmm_data.dt_mod.columns]
+        X = self.featurized_mmm_data.dt_mod.drop(columns=exclude_cols)
 
         # Convert date columns to numeric (number of days since the earliest date)
         date_columns = X.select_dtypes(include=["datetime64", "object"]).columns