test exceptions are raised for the new input validation

Author: drbenvincent

causalpy/pymc_experiments.py

@@ -12,7 +12,7 @@
 """
 
 import warnings
-from typing import Optional, Union
+from typing import Union
 
 import arviz as az
 import matplotlib.pyplot as plt
@@ -794,7 +794,7 @@ def __init__(
         model=None,
         running_variable_name: str = "x",
         epsilon: float = 0.001,
-        bandwidth: Optional[float] = None,
+        bandwidth: float = np.inf,
         **kwargs,
     ):
         super().__init__(model=model, **kwargs)
@@ -807,7 +807,7 @@ def __init__(
         self.bandwidth = bandwidth
         self._input_validation()
 
-        if self.bandwidth is not None:
+        if self.bandwidth is not np.inf:
             fmin = self.treatment_threshold - self.bandwidth
             fmax = self.treatment_threshold + self.bandwidth
             filtered_data = self.data.query(f"{fmin} <= x <= {fmax}")
@@ -836,7 +836,7 @@ def __init__(
         self.score = self.model.score(X=self.X, y=self.y)
 
         # get the model predictions of the observed data
-        if self.bandwidth is not None:
+        if self.bandwidth is not np.inf:
             xi = np.linspace(fmin, fmax, 200)
         else:
             xi = np.linspace(
@@ -988,7 +988,7 @@ def __init__(
         model=None,
         running_variable_name: str = "x",
         epsilon: float = 0.001,
-        bandwidth: Optional[float] = None,
+        bandwidth: float = np.inf,
         **kwargs,
     ):
         super().__init__(model=model, **kwargs)
@@ -1001,7 +1001,7 @@ def __init__(
         self.bandwidth = bandwidth
         self._input_validation()
 
-        if self.bandwidth is not None:
+        if self.bandwidth is not np.inf:
             fmin = self.kink_point - self.bandwidth
             fmax = self.kink_point + self.bandwidth
             filtered_data = self.data.query(f"{fmin} <= x <= {fmax}")
@@ -1027,7 +1027,7 @@ def __init__(
         self.score = self.model.score(X=self.X, y=self.y)
 
         # get the model predictions of the observed data
-        if self.bandwidth is not None:
+        if self.bandwidth is not np.inf:
             xi = np.linspace(fmin, fmax, 200)
         else:
             xi = np.linspace(

causalpy/tests/test_input_validation.py

@@ -1,5 +1,6 @@
 """Input validation tests"""
 
+import numpy as np
 import pandas as pd
 import pytest
 
@@ -223,3 +224,83 @@ def test_iv_treatment_var_is_present():
                 sample_kwargs=sample_kwargs
             ),
         )
+
+
+# Regression kink design
+
+
+def setup_regression_kink_data(kink):
+    """Set up data for regression kink design tests"""
+    # define parameters for data generation
+    seed = 42
+    rng = np.random.default_rng(seed)
+    N = 50
+    kink = 0.5
+    beta = [0, -1, 0, 2, 0]
+    sigma = 0.05
+    # generate data
+    x = rng.uniform(-1, 1, N)
+    y = reg_kink_function(x, beta, kink) + rng.normal(0, sigma, N)
+    return pd.DataFrame({"x": x, "y": y, "treated": x >= kink})
+
+
+def reg_kink_function(x, beta, kink):
+    """Utility function for regression kink design. Returns a piecewise linear function
+    evaluated at x with a kink at kink and parameters beta"""
+    return (
+        beta[0]
+        + beta[1] * x
+        + beta[2] * x**2
+        + beta[3] * (x - kink) * (x >= kink)
+        + beta[4] * (x - kink) ** 2 * (x >= kink)
+    )
+
+
+def test_rkink_bandwidth_check():
+    """Test that we get exceptions when bandwidth parameter is <= 0"""
+    with pytest.raises(ValueError):
+        kink = 0.5
+        df = setup_regression_kink_data(kink)
+        _ = cp.pymc_experiments.RegressionKink(
+            df,
+            formula=f"y ~ 1 + x + I((x-{kink})*treated)",
+            model=cp.pymc_models.LinearRegression(sample_kwargs=sample_kwargs),
+            kink_point=kink,
+            bandwidth=0,
+        )
+
+    with pytest.raises(ValueError):
+        kink = 0.5
+        df = setup_regression_kink_data(kink)
+        _ = cp.pymc_experiments.RegressionKink(
+            df,
+            formula=f"y ~ 1 + x + I((x-{kink})*treated)",
+            model=cp.pymc_models.LinearRegression(sample_kwargs=sample_kwargs),
+            kink_point=kink,
+            bandwidth=-1,
+        )
+
+
+def test_rkink_epsilon_check():
+    """Test that we get exceptions when epsilon parameter is <= 0"""
+    with pytest.raises(ValueError):
+        kink = 0.5
+        df = setup_regression_kink_data(kink)
+        _ = cp.pymc_experiments.RegressionKink(
+            df,
+            formula=f"y ~ 1 + x + I((x-{kink})*treated)",
+            model=cp.pymc_models.LinearRegression(sample_kwargs=sample_kwargs),
+            kink_point=kink,
+            epsilon=0,
+        )
+
+    with pytest.raises(ValueError):
+        kink = 0.5
+        df = setup_regression_kink_data(kink)
+        _ = cp.pymc_experiments.RegressionKink(
+            df,
+            formula=f"y ~ 1 + x + I((x-{kink})*treated)",
+            model=cp.pymc_models.LinearRegression(sample_kwargs=sample_kwargs),
+            kink_point=kink,
+            epsilon=-1,
+        )

causalpy/tests/test_integration_pymc_examples.py

@@ -7,18 +7,6 @@
 sample_kwargs = {"tune": 20, "draws": 20, "chains": 2, "cores": 2}
 
 
-def reg_kink_function(x, beta, kink):
-    """Utility function for regression kink design. Returns a piecewise linear function
-    evaluated at x with a kink at kink and parameters beta"""
-    return (
-        beta[0]
-        + beta[1] * x
-        + beta[2] * x**2
-        + beta[3] * (x - kink) * (x >= kink)
-        + beta[4] * (x - kink) ** 2 * (x >= kink)
-    )
-
-
 @pytest.mark.integration
 def test_did():
     """
@@ -230,6 +218,33 @@ def test_rd_drinking():
     assert len(result.idata.posterior.coords["draw"]) == sample_kwargs["draws"]
 
 
+def setup_regression_kink_data(kink):
+    """Set up data for regression kink design tests"""
+    # define parameters for data generation
+    seed = 42
+    rng = np.random.default_rng(seed)
+    N = 50
+    kink = 0.5
+    beta = [0, -1, 0, 2, 0]
+    sigma = 0.05
+    # generate data
+    x = rng.uniform(-1, 1, N)
+    y = reg_kink_function(x, beta, kink) + rng.normal(0, sigma, N)
+    return pd.DataFrame({"x": x, "y": y, "treated": x >= kink})
+
+
+def reg_kink_function(x, beta, kink):
+    """Utility function for regression kink design. Returns a piecewise linear function
+    evaluated at x with a kink at kink and parameters beta"""
+    return (
+        beta[0]
+        + beta[1] * x
+        + beta[2] * x**2
+        + beta[3] * (x - kink) * (x >= kink)
+        + beta[4] * (x - kink) ** 2 * (x >= kink)
+    )
+
+
 @pytest.mark.integration
 def test_rkink():
     """
@@ -241,18 +256,8 @@ def test_rkink():
     3. the correct number of MCMC chains exists in the posterior inference data
     4. the correct number of MCMC draws exists in the posterior inference data
     """
-    # define parameters for data generation
-    seed = 42
-    rng = np.random.default_rng(seed)
-    N = 50
     kink = 0.5
-    beta = [0, -1, 0, 2, 0]
-    sigma = 0.05
-    # generate data
-    x = rng.uniform(-1, 1, N)
-    y = reg_kink_function(x, beta, kink) + rng.normal(0, sigma, N)
-    df = pd.DataFrame({"x": x, "y": y, "treated": x >= kink})
-    # run experiment
+    df = setup_regression_kink_data(kink)
     result = cp.pymc_experiments.RegressionKink(
         df,
         formula=f"y ~ 1 + x + I((x-{kink})*treated)",
@@ -276,18 +281,8 @@ def test_rkink_bandwidth():
     3. the correct number of MCMC chains exists in the posterior inference data
     4. the correct number of MCMC draws exists in the posterior inference data
     """
-    # define parameters for data generation
-    seed = 42
-    rng = np.random.default_rng(seed)
-    N = 50
     kink = 0.5
-    beta = [0, -1, 0, 2, 0]
-    sigma = 0.05
-    # generate data
-    x = rng.uniform(-1, 1, N)
-    y = reg_kink_function(x, beta, kink) + rng.normal(0, sigma, N)
-    df = pd.DataFrame({"x": x, "y": y, "treated": x >= kink})
-    # run experiment
+    df = setup_regression_kink_data(kink)
     result = cp.pymc_experiments.RegressionKink(
         df,
         formula=f"y ~ 1 + x + I((x-{kink})*treated)",