add integration tests for regression kink design

Author: drbenvincent

causalpy/tests/test_integration_pymc_examples.py

@@ -1,3 +1,4 @@
+import numpy as np
 import pandas as pd
 import pytest
 
@@ -6,6 +7,18 @@
 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():
     """
@@ -217,6 +230,77 @@ def test_rd_drinking():
     assert len(result.idata.posterior.coords["draw"]) == sample_kwargs["draws"]
 
 
+@pytest.mark.integration
+def test_rkink():
+    """
+    Test Regression Kink design.
+
+    Loads data and checks:
+    1. data is a dataframe
+    2. pymc_experiments.RegressionKink returns correct type
+    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
+    result = 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,
+    )
+    assert isinstance(df, pd.DataFrame)
+    assert isinstance(result, cp.pymc_experiments.RegressionKink)
+    assert len(result.idata.posterior.coords["chain"]) == sample_kwargs["chains"]
+    assert len(result.idata.posterior.coords["draw"]) == sample_kwargs["draws"]
+
+
+@pytest.mark.integration
+def test_rkink_bandwidth():
+    """
+    Test Regression Kink experiment with bandwidth parameter.
+
+    Generates synthetic data and checks:
+    1. data is a dataframe
+    2. pymc_experiments.RegressionKink returns correct type
+    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
+    result = 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.3,
+    )
+    assert isinstance(df, pd.DataFrame)
+    assert isinstance(result, cp.pymc_experiments.RegressionKink)
+    assert len(result.idata.posterior.coords["chain"]) == sample_kwargs["chains"]
+    assert len(result.idata.posterior.coords["draw"]) == sample_kwargs["draws"]
+
+
 @pytest.mark.integration
 def test_its():
     """