added validations for interactions, test coverage expanded to test interaction terms,more generic messages

Author: Rojan Shrestha

causalpy/experiments/diff_in_diff.py

@@ -15,6 +15,8 @@
 Difference in differences
 """
 
+import re
+
 import arviz as az
 import numpy as np
 import pandas as pd
@@ -233,42 +235,21 @@ def __init__(
         return
 
     def input_validation(self):
+        # Validate formula structure and interaction interaction terms
+        self._validate_formula_interaction_terms()
+
         """Validate the input data and model formula for correctness"""
         # Check if post_treatment_variable_name is in formula
         if self.post_treatment_variable_name not in self.formula:
-            if self.post_treatment_variable_name == "post_treatment":
-                # Default case - user didn't specify custom name, so guide them to use "post_treatment"
-                raise FormulaException(
-                    "Missing 'post_treatment' in formula.\n"
-                    "Note: post_treatment_variable_name might have been set to 'post_treatment' by default.\n"
-                    "Add 'post_treatment' to formula (e.g., 'y ~ 1 + group*post_treatment').\n"
-                    "Or to use custom name, provide additional argument post_treatment_variable_name='your_post_treatment_variable_name'."
-                )
-            else:
-                # Custom case - user specified custom name, so remind them what they specified
-                raise FormulaException(
-                    f"Missing required variable '{self.post_treatment_variable_name}' in formula.\n\n"
-                    f"Since you specified post_treatment_variable_name='{self.post_treatment_variable_name}', "
-                    f"please ensure formula includes '{self.post_treatment_variable_name}'"
-                )
+            raise FormulaException(
+                f"Missing required variable '{self.post_treatment_variable_name}' in formula"
+            )
 
         # Check if post_treatment_variable_name is in data columns
         if self.post_treatment_variable_name not in self.data.columns:
-            if self.post_treatment_variable_name == "post_treatment":
-                # Default case - user didn't specify custom name, so guide them to use "post_treatment"
-                raise DataException(
-                    "Missing 'post_treatment' column in dataset.\n"
-                    "Note: post_treatment_variable_name might have been set to 'post_treatment' by default.\n"
-                    "Ensure dataset has boolean column 'post_treatment'.\n"
-                    "or to use custom name, provide additional argument post_treatment_variable_name='your_post_treatment_variable_name'."
-                )
-            else:
-                # Custom case - user specified custom name, so remind them what they specified
-                raise DataException(
-                    f"Missing required column '{self.post_treatment_variable_name}' in dataset.\n\n"
-                    f"Since you specified post_treatment_variable_name='{self.post_treatment_variable_name}', "
-                    f"please ensure dataset has boolean column named '{self.post_treatment_variable_name}'"
-                )
+            raise DataException(
+                f"Missing required column '{self.post_treatment_variable_name}' in dataset"
+            )
 
         if "unit" not in self.data.columns:
             raise DataException(
@@ -281,6 +262,61 @@ def input_validation(self):
                 coded. Consisting of 0's and 1's only."""
             )
 
+    def _get_interaction_terms(self):
+        """
+        Extract interaction terms from the formula.
+        Returns a list of interaction terms (those with '*' or ':').
+        """
+        # Define interaction indicators
+        INTERACTION_INDICATORS = ["*", ":"]
+
+        # Remove whitespace
+        formula = self.formula.replace(" ", "")
+
+        # Extract right-hand side of the formula
+        rhs = formula.split("~")[1]
+
+        # Split terms by '+' or '-' while keeping them intact
+        terms = re.split(r"(?=[+-])", rhs)
+
+        # Clean up terms and get interaction terms (those with '*' or ':')
+        interaction_terms = []
+        for term in terms:
+            # Remove leading + or - for processing
+            clean_term = term.lstrip("+-")
+            if any(indicator in clean_term for indicator in INTERACTION_INDICATORS):
+                interaction_terms.append(clean_term)
+
+        return interaction_terms
+
+    def _validate_formula_interaction_terms(self):
+        """
+        Validate that the formula contains at most one interaction term and no three-way or higher-order interactions.
+        Raises FormulaException if more than one interaction term is found or if any interaction term has more than 2 variables.
+        """
+        # Define interaction indicators
+        INTERACTION_INDICATORS = ["*", ":"]
+
+        # Get interaction terms
+        interaction_terms = self._get_interaction_terms()
+
+        # Check for interaction terms with more than 2 variables (more than one '*' or ':')
+        for term in interaction_terms:
+            total_indicators = sum(
+                term.count(indicator) for indicator in INTERACTION_INDICATORS
+            )
+            if (
+                total_indicators >= 2
+            ):  # 3 or more variables (e.g., a*b*c or a:b:c has 2 symbols)
+                raise FormulaException(
+                    f"Formula contains interaction term with more than 2 variables: {term}. Only two-way interactions are allowed."
+                )
+
+        if len(interaction_terms) > 1:
+            raise FormulaException(
+                f"Formula contains more than 1 interaction term: {interaction_terms}. Maximum of 1 allowed."
+            )
+
     def summary(self, round_to=None) -> None:
         """Print summary of main results and model coefficients.
 

causalpy/tests/test_input_validation.py

@@ -30,18 +30,29 @@
 
 
 def test_did_validation_post_treatment_formula():
-    """Test that we get a FormulaException if do not include post_treatment in the
-    formula"""
+    """Test that we get a FormulaException for invalid formulas and missing post_treatment variables"""
     df = pd.DataFrame(
         {
             "group": [0, 0, 1, 1],
             "t": [0, 1, 0, 1],
             "unit": [0, 0, 1, 1],
             "post_treatment": [0, 1, 0, 1],
+            "male": [0, 1, 0, 1],  # Additional variable for testing
             "y": [1, 2, 3, 4],
         }
     )
 
+    df_with_custom = pd.DataFrame(
+        {
+            "group": [0, 0, 1, 1],
+            "t": [0, 1, 0, 1],
+            "unit": [0, 0, 1, 1],
+            "custom_post": [0, 1, 0, 1],  # Custom column name
+            "y": [1, 2, 3, 4],
+        }
+    )
+
+    # Test 1: Missing post_treatment variable in formula
     with pytest.raises(FormulaException):
         _ = cp.DifferenceInDifferences(
             df,
@@ -51,6 +62,7 @@ def test_did_validation_post_treatment_formula():
             model=cp.pymc_models.LinearRegression(sample_kwargs=sample_kwargs),
         )
 
+    # Test 2: Missing post_treatment variable in formula (duplicate test)
     with pytest.raises(FormulaException):
         _ = cp.DifferenceInDifferences(
             df,
@@ -60,6 +72,88 @@ def test_did_validation_post_treatment_formula():
             model=cp.pymc_models.LinearRegression(sample_kwargs=sample_kwargs),
         )
 
+    # Test 3: Custom post_treatment_variable_name but formula uses different name
+    with pytest.raises(FormulaException):
+        _ = cp.DifferenceInDifferences(
+            df_with_custom,
+            formula="y ~ 1 + group*post_treatment",  # Formula uses 'post_treatment'
+            time_variable_name="t",
+            group_variable_name="group",
+            post_treatment_variable_name="custom_post",  # But user specifies 'custom_post'
+            model=cp.pymc_models.LinearRegression(sample_kwargs=sample_kwargs),
+        )
+
+    # Test 4: Default post_treatment_variable_name but formula uses different name
+    with pytest.raises(FormulaException):
+        _ = cp.DifferenceInDifferences(
+            df,
+            formula="y ~ 1 + group*custom_post",  # Formula uses 'custom_post'
+            time_variable_name="t",
+            group_variable_name="group",
+            # post_treatment_variable_name defaults to "post_treatment"
+            model=cp.pymc_models.LinearRegression(sample_kwargs=sample_kwargs),
+        )
+
+    # Test 5: Repeated interaction terms (should be invalid)
+    with pytest.raises(FormulaException):
+        _ = cp.DifferenceInDifferences(
+            df,
+            formula="y ~ 1 + group + group*post_treatment + group*post_treatment",
+            time_variable_name="t",
+            group_variable_name="group",
+            model=cp.pymc_models.LinearRegression(sample_kwargs=sample_kwargs),
+        )
+
+    # Test 6: Three-way interactions using * (should be invalid)
+    with pytest.raises(FormulaException):
+        _ = cp.DifferenceInDifferences(
+            df,
+            formula="y ~ 1 + group + group*post_treatment*male",
+            time_variable_name="t",
+            group_variable_name="group",
+            model=cp.pymc_models.LinearRegression(sample_kwargs=sample_kwargs),
+        )
+
+    # Test 7: Three-way interactions using : (should be invalid)
+    with pytest.raises(FormulaException):
+        _ = cp.DifferenceInDifferences(
+            df,
+            formula="y ~ 1 + group + group:post_treatment:male",
+            time_variable_name="t",
+            group_variable_name="group",
+            model=cp.pymc_models.LinearRegression(sample_kwargs=sample_kwargs),
+        )
+
+    # Test 8: Multiple different interaction terms using * (should be invalid)
+    with pytest.raises(FormulaException):
+        _ = cp.DifferenceInDifferences(
+            df,
+            formula="y ~ 1 + group + group*post_treatment + group*male",
+            time_variable_name="t",
+            group_variable_name="group",
+            model=cp.pymc_models.LinearRegression(sample_kwargs=sample_kwargs),
+        )
+
+    # Test 9: Multiple different interaction terms using : (should be invalid)
+    with pytest.raises(FormulaException):
+        _ = cp.DifferenceInDifferences(
+            df,
+            formula="y ~ 1 + group + group:post_treatment + group:male",
+            time_variable_name="t",
+            group_variable_name="group",
+            model=cp.pymc_models.LinearRegression(sample_kwargs=sample_kwargs),
+        )
+
+    # Test 10: Mixed issues - multiple terms + three-way interaction (should be invalid)
+    with pytest.raises(FormulaException):
+        _ = cp.DifferenceInDifferences(
+            df,
+            formula="y ~ 1 + group + group*post_treatment + group:post_treatment:male",
+            time_variable_name="t",
+            group_variable_name="group",
+            model=cp.pymc_models.LinearRegression(sample_kwargs=sample_kwargs),
+        )
+
 
 def test_did_validation_post_treatment_data():
     """Test that we get a DataException if do not include post_treatment in the data"""
@@ -91,6 +185,27 @@ def test_did_validation_post_treatment_data():
             model=cp.pymc_models.LinearRegression(sample_kwargs=sample_kwargs),
         )
 
+    # Test 2: Custom post_treatment_variable_name but column doesn't exist in data
+    df_with_post = pd.DataFrame(
+        {
+            "group": [0, 0, 1, 1],
+            "t": [0, 1, 0, 1],
+            "unit": [0, 0, 1, 1],
+            "post_treatment": [0, 1, 0, 1],  # Data has 'post_treatment'
+            "y": [1, 2, 3, 4],
+        }
+    )
+
+    with pytest.raises(DataException):
+        _ = cp.DifferenceInDifferences(
+            df_with_post,
+            formula="y ~ 1 + group*custom_post",  # Formula uses 'custom_post'
+            time_variable_name="t",
+            group_variable_name="group",
+            post_treatment_variable_name="custom_post",  # User specifies 'custom_post'
+            model=cp.pymc_models.LinearRegression(sample_kwargs=sample_kwargs),
+        )
+
 
 def test_did_validation_unit_data():
     """Test that we get a DataException if do not include unit in the data"""