replace asserts with custom exceptions + add utility functions and associated tests

Author: drbenvincent

causalpy/custom_exceptions.py

@@ -4,3 +4,18 @@ class BadIndexException(Exception):
 
     def __init__(self, message):
         self.message = message
+
+
+class FormulaException(Exception):
+    """Exception raised given when there is some error in a user-provided model
+    formula"""
+
+    def __init__(self, message):
+        self.message = message
+
+
+class DataException(Exception):
+    """Exception raised given when there is some error in user-provided dataframe"""
+
+    def __init__(self, message):
+        self.message = message

causalpy/pymc_experiments.py

@@ -8,8 +8,13 @@
 import xarray as xr
 from patsy import build_design_matrices, dmatrices
 
-from causalpy.custom_exceptions import BadIndexException
+from causalpy.custom_exceptions import (
+    BadIndexException,
+    DataException,
+    FormulaException,
+)
 from causalpy.plot_utils import plot_xY
+from causalpy.utils import _is_variable_dummy_coded, _series_has_2_levels
 
 LEGEND_FONT_SIZE = 12
 az.style.use("arviz-darkgrid")
@@ -117,7 +122,7 @@ def __init__(
         self.post_impact_cumulative = self.post_impact.cumsum(dim="obs_ind")
 
     def _input_validation(self, data, treatment_time):
-        """Validate the input data for correctness"""
+        """Validate the input data and model formula for correctness"""
         if isinstance(data.index, pd.DatetimeIndex) and not isinstance(
             treatment_time, pd.Timestamp
         ):
@@ -361,28 +366,27 @@ def __init__(
                 self.causal_impact = self.idata.posterior["beta"].isel({"coeffs": i})
 
     def _input_validation(self):
-        """Validate the input data for correctness"""
-        assert (
-            "post_treatment" in self.formula
-        ), "A predictor called `post_treatment` should be in the dataframe"
-        assert (
-            "post_treatment" in self.data.columns
-        ), "Require a boolean column labelling observations which are `treated`"
-        # Check for `unit` in the incoming dataframe.
-        # *This is only used for plotting purposes*
-        assert (
-            "unit" in self.data.columns
-        ), """
-        Require a `unit` column to label unique units.
-        This is used for plotting purposes
-        """
-        # Check that `group_variable_name` is dummy coded. It should be 0 or 1
-        assert not set(self.data[self.group_variable_name]).difference(
-            set([0, 1])
-        ), f"""
-            The grouping variable {self.group_variable_name} should be dummy coded.
-            Consisting of 0's and 1's only.
-        """
+        """Validate the input data and model formula for correctness"""
+        if "post_treatment" not in self.formula:
+            raise FormulaException(
+                "A predictor called `post_treatment` should be in the formula"
+            )
+
+        if "post_treatment" not in self.data.columns:
+            raise DataException(
+                "Require a boolean column labelling observations which are `treated`"
+            )
+
+        if "unit" not in self.data.columns:
+            raise DataException(
+                "Require a `unit` column to label unique units. This is used for plotting purposes"  # noqa: E501
+            )
+
+        if _is_variable_dummy_coded(self.data[self.group_variable_name]) is False:
+            raise DataException(
+                f"""The grouping variable {self.group_variable_name} should be dummy
+                coded. Consisting of 0's and 1's only."""
+            )
 
     def plot(self):
         """Plot the results.
@@ -744,16 +748,17 @@ def __init__(
         # ================================================================
 
     def _input_validation(self):
-        """Validate the input data for correctness"""
+        """Validate the input data and model formula for correctness"""
         # Check that `group_variable_name` has TWO levels, representing the
         # treated/untreated. But it does not matter what the actual names of
         # the levels are.
-        assert (
-            len(pd.Categorical(self.data[self.group_variable_name]).categories) == 2
-        ), f"""
-            There must be 2 levels of the grouping variable {self.group_variable_name}
-            .I.e. the treated and untreated.
-        """
+        if not _series_has_2_levels(self.data[self.group_variable_name]):
+            raise ValueError(
+                f"""
+                There must be 2 levels of the grouping variable
+                {self.group_variable_name}. I.e. the treated and untreated.
+                """
+            )
 
     def plot(self):
         """Plot the results"""

causalpy/tests/test_utils.py

@@ -0,0 +1,22 @@
+import pandas as pd
+
+from causalpy.utils import _is_variable_dummy_coded, _series_has_2_levels
+
+
+def test_dummy_coding():
+    """Test if the function to check if a variable is dummy coded works correctly"""
+    assert _is_variable_dummy_coded(pd.Series([False, True, False, True])) is True
+    assert _is_variable_dummy_coded(pd.Series([False, True, False, "frog"])) is False
+    assert _is_variable_dummy_coded(pd.Series([0, 0, 1, 0, 1])) is True
+    assert _is_variable_dummy_coded(pd.Series([0, 0, 1, 0, 2])) is False
+    assert _is_variable_dummy_coded(pd.Series([0, 0.5, 1, 0, 1])) is False
+
+
+def test_2_level_series():
+    """Test if the function to check if a variable has 2 levels works correctly"""
+    assert _series_has_2_levels(pd.Series(["a", "a", "b"])) is True
+    assert _series_has_2_levels(pd.Series(["a", "a", "b", "c"])) is False
+    assert _series_has_2_levels(pd.Series(["coffee", "tea", "coffee"])) is True
+    assert _series_has_2_levels(pd.Series(["water", "tea", "coffee"])) is False
+    assert _series_has_2_levels(pd.Series([0, 1, 0, 1])) is True
+    assert _series_has_2_levels(pd.Series([0, 1, 0, 2])) is False

causalpy/utils.py

@@ -0,0 +1,12 @@
+import pandas as pd
+
+
+def _is_variable_dummy_coded(series: pd.Series) -> bool:
+    """Check if a data in the provided Series is dummy coded. It should be 0 or 1
+    only."""
+    return len(set(series).difference(set([0, 1]))) == 0
+
+
+def _series_has_2_levels(series: pd.Series) -> bool:
+    """Check that the variable in the provided Series has 2 levels"""
+    return len(pd.Categorical(series).categories) == 2