From db29b5b95352035611130a246d7a9df79f0d39cc Mon Sep 17 00:00:00 2001 From: "Benjamin T. Vincent" Date: Fri, 5 Jan 2024 11:41:21 +0000 Subject: [PATCH] add tests --- causalpy/tests/test_pymc_models.py | 35 ++++++++++++++++++++++++++++++ 1 file changed, 35 insertions(+) diff --git a/causalpy/tests/test_pymc_models.py b/causalpy/tests/test_pymc_models.py index d535ba65..481bce9e 100644 --- a/causalpy/tests/test_pymc_models.py +++ b/causalpy/tests/test_pymc_models.py @@ -123,3 +123,38 @@ def test_idata_property(): ) assert hasattr(result, "idata") assert isinstance(result.idata, az.InferenceData) + + +def test_result_reproducibility(): + """Test that we can reproduce the results from the model. We could in theory test + this with all the model and experiment types, but what is being targetted is + the ModelBuilder.fit method, so we should be safe testing with just one model. Here + we use the DifferenceInDifferences experiment class.""" + # Load the data + df = cp.load_data("did") + # Set a random seed + sample_kwargs["random_seed"] = 42 + # Calculate the result twice + result1 = cp.pymc_experiments.DifferenceInDifferences( + df, + formula="y ~ 1 + group + t + group:post_treatment", + time_variable_name="t", + group_variable_name="group", + treated=1, + untreated=0, + model=cp.pymc_models.LinearRegression(sample_kwargs=sample_kwargs), + ) + result2 = cp.pymc_experiments.DifferenceInDifferences( + df, + formula="y ~ 1 + group + t + group:post_treatment", + time_variable_name="t", + group_variable_name="group", + treated=1, + untreated=0, + model=cp.pymc_models.LinearRegression(sample_kwargs=sample_kwargs), + ) + assert np.all(result1.idata.posterior.mu == result2.idata.posterior.mu) + assert np.all(result1.idata.prior.mu == result2.idata.prior.mu) + assert np.all( + result1.idata.prior_predictive.y_hat == result2.idata.prior_predictive.y_hat + )