From 6bee429ab09a82d893bfa919507b163a90a41ad1 Mon Sep 17 00:00:00 2001 From: Jason Preszler Date: Fri, 18 Aug 2023 16:08:53 -0700 Subject: [PATCH 01/23] Issue 129: increase docstring coverage, now at 86% --- causalpy/custom_exceptions.py | 11 +- causalpy/data/simulate_data.py | 3 + causalpy/plot_utils.py | 4 + causalpy/pymc_experiments.py | 28 +++-- causalpy/pymc_models.py | 1 + causalpy/skl_experiments.py | 7 ++ causalpy/skl_models.py | 9 ++ causalpy/tests/conftest.py | 1 + causalpy/tests/test_data_loading.py | 4 + .../tests/test_integration_pymc_examples.py | 117 ++++++++++++++++++ .../tests/test_integration_skl_examples.py | 59 +++++++++ causalpy/tests/test_pymc_models.py | 37 ++++++ docs/source/_static/interrogate_badge.svg | 8 +- 13 files changed, 274 insertions(+), 15 deletions(-) diff --git a/causalpy/custom_exceptions.py b/causalpy/custom_exceptions.py index e73dc604..c3f14549 100644 --- a/causalpy/custom_exceptions.py +++ b/causalpy/custom_exceptions.py @@ -1,8 +1,13 @@ +""" +Custom Exceptions for CausalPy. +""" + + class BadIndexException(Exception): """Custom exception used when we have a mismatch in types between the dataframe index and an event, typically a treatment or intervention.""" - def __init__(self, message): + def __init__(self, message: str): self.message = message @@ -10,12 +15,12 @@ class FormulaException(Exception): """Exception raised given when there is some error in a user-provided model formula""" - def __init__(self, message): + def __init__(self, message: str): self.message = message class DataException(Exception): """Exception raised given when there is some error in user-provided dataframe""" - def __init__(self, message): + def __init__(self, message: str): self.message = message diff --git a/causalpy/data/simulate_data.py b/causalpy/data/simulate_data.py index 38bb8376..b752fbbe 100644 --- a/causalpy/data/simulate_data.py +++ b/causalpy/data/simulate_data.py @@ -11,6 +11,9 @@ def _smoothed_gaussian_random_walk( gaussian_random_walk_mu, gaussian_random_walk_sigma, N, lowess_kwargs ): + """ + Generates Gaussian random walk data and applies LOWESS + """ x = np.arange(N) y = norm(gaussian_random_walk_mu, gaussian_random_walk_sigma).rvs(N).cumsum() filtered = lowess(y, x, **lowess_kwargs) diff --git a/causalpy/plot_utils.py b/causalpy/plot_utils.py index a402ddf9..52d44bbc 100644 --- a/causalpy/plot_utils.py +++ b/causalpy/plot_utils.py @@ -1,3 +1,7 @@ +""" +Plotting utility functions. +""" + from typing import Any, Dict, Optional, Tuple, Union import arviz as az diff --git a/causalpy/pymc_experiments.py b/causalpy/pymc_experiments.py index 68da0638..fb394ad8 100644 --- a/causalpy/pymc_experiments.py +++ b/causalpy/pymc_experiments.py @@ -1,3 +1,13 @@ +""" +Experiment routines for PyMC models. + +Includes: +1. ExperimentalDesign base class +2. Pre-Post Fit +3. Synthetic Control +4. Difference in differences +5. Regression Discontinuity +""" import warnings from typing import Optional, Union @@ -36,7 +46,7 @@ def idata(self): """Access to the InferenceData object""" return self.model.idata - def print_coefficients(self): + def print_coefficients(self) -> None: """Prints the model coefficients""" print("Model coefficients:") coeffs = az.extract(self.idata.posterior, var_names="beta") @@ -236,7 +246,7 @@ def plot(self, counterfactual_label="Counterfactual", **kwargs): return (fig, ax) - def summary(self): + def summary(self) -> None: """Print text output summarising the results""" print(f"{self.expt_type:=^80}") @@ -524,13 +534,14 @@ def _plot_causal_impact_arrow(self, ax): va="center", ) - def _causal_impact_summary_stat(self): + def _causal_impact_summary_stat(self) -> str: + """Computes the mean and 94% credible interval bounds for the causal impact.""" percentiles = self.causal_impact.quantile([0.03, 1 - 0.03]).values ci = r"$CI_{94\%}$" + f"[{percentiles[0]:.2f}, {percentiles[1]:.2f}]" causal_impact = f"{self.causal_impact.mean():.2f}, " return f"Causal impact = {causal_impact + ci}" - def summary(self): + def summary(self) -> None: """Print text output summarising the results""" print(f"{self.expt_type:=^80}") @@ -716,7 +727,7 @@ def plot(self): ) return (fig, ax) - def summary(self): + def summary(self) -> None: """Print text output summarising the results""" print(f"{self.expt_type:=^80}") @@ -795,7 +806,7 @@ def __init__( # ================================================================ - def _input_validation(self): + def _input_validation(self) -> None: """Validate the input data and model formula for correctness""" if not _series_has_2_levels(self.data[self.group_variable_name]): raise DataException( @@ -856,13 +867,14 @@ def plot(self): ax[1].set(title="Estimated treatment effect") return fig, ax - def _causal_impact_summary_stat(self): + def _causal_impact_summary_stat(self) -> str: + """Computes the mean and 94% credible interval bounds for the causal impact.""" percentiles = self.causal_impact.quantile([0.03, 1 - 0.03]).values ci = r"$CI_{94\%}$" + f"[{percentiles[0]:.2f}, {percentiles[1]:.2f}]" causal_impact = f"{self.causal_impact.mean():.2f}, " return f"Causal impact = {causal_impact + ci}" - def summary(self): + def summary(self) -> None: """Print text output summarising the results""" print(f"{self.expt_type:=^80}") diff --git a/causalpy/pymc_models.py b/causalpy/pymc_models.py index 33833b91..ad54a94a 100644 --- a/causalpy/pymc_models.py +++ b/causalpy/pymc_models.py @@ -40,6 +40,7 @@ def build_model(self, X, y, coords) -> None: raise NotImplementedError("This method must be implemented by a subclass") def _data_setter(self, X) -> None: + """Set data for the model.""" with self.model: pm.set_data({"X": X}) diff --git a/causalpy/skl_experiments.py b/causalpy/skl_experiments.py index 67505c46..b9209c4a 100644 --- a/causalpy/skl_experiments.py +++ b/causalpy/skl_experiments.py @@ -1,3 +1,6 @@ +""" +Experiments for Scikit-Learn models +""" import warnings from typing import Optional @@ -78,6 +81,7 @@ def __init__( self.post_impact_cumulative = np.cumsum(self.post_impact) def plot(self, counterfactual_label="Counterfactual", **kwargs): + """Plot experiment results""" fig, ax = plt.subplots(3, 1, sharex=True, figsize=(7, 8)) ax[0].plot(self.datapre.index, self.pre_y, "k.") @@ -140,9 +144,11 @@ def plot(self, counterfactual_label="Counterfactual", **kwargs): return (fig, ax) def get_coeffs(self): + """Returns model coefficients""" return np.squeeze(self.model.coef_) def plot_coeffs(self): + """Plots coefficient bar plot""" df = pd.DataFrame( {"predictor variable": self.labels, "ols_coef": self.get_coeffs()} ) @@ -463,6 +469,7 @@ def _is_treated(self, x): return np.greater_equal(x, self.treatment_threshold) def plot(self): + """Plot results""" fig, ax = plt.subplots() # Plot raw data sns.scatterplot( diff --git a/causalpy/skl_models.py b/causalpy/skl_models.py index 1649c094..d2d0cae2 100644 --- a/causalpy/skl_models.py +++ b/causalpy/skl_models.py @@ -1,3 +1,9 @@ +""" +Scikit-Learn Models + +Includes: +1. Weighted Proportion +""" from functools import partial import numpy as np @@ -18,9 +24,11 @@ class WeightedProportion(LinearModel, RegressorMixin): """ def loss(self, W, X, y): + """Compute root mean squared loss with data X, weights W, and predictor y""" return np.sqrt(np.mean((y - np.dot(X, W.T)) ** 2)) def fit(self, X, y): + """Fit model on data X with predictor y""" w_start = [1 / X.shape[1]] * X.shape[1] coef_ = fmin_slsqp( partial(self.loss, X=X, y=y), @@ -34,4 +42,5 @@ def fit(self, X, y): return self def predict(self, X): + """Predict results for data X""" return np.dot(X, self.coef_.T) diff --git a/causalpy/tests/conftest.py b/causalpy/tests/conftest.py index 48845fc5..d652cc8f 100644 --- a/causalpy/tests/conftest.py +++ b/causalpy/tests/conftest.py @@ -4,5 +4,6 @@ @pytest.fixture(scope="session") def rng() -> np.random.Generator: + """Random number generator that can persist through a pytest session""" seed: int = sum(map(ord, "causalpy")) return np.random.default_rng(seed=seed) diff --git a/causalpy/tests/test_data_loading.py b/causalpy/tests/test_data_loading.py index 554684d6..5d562aa8 100644 --- a/causalpy/tests/test_data_loading.py +++ b/causalpy/tests/test_data_loading.py @@ -19,6 +19,10 @@ @pytest.mark.parametrize("dataset_name", tests) def test_data_loading(dataset_name): + """ + Checks that test data can be loaded into data frames and that there are no + missing values in any column. + """ df = cp.load_data(dataset_name) assert isinstance(df, pd.DataFrame) # Check that there are no missing values in any column diff --git a/causalpy/tests/test_integration_pymc_examples.py b/causalpy/tests/test_integration_pymc_examples.py index 53c4e361..7b43dc64 100644 --- a/causalpy/tests/test_integration_pymc_examples.py +++ b/causalpy/tests/test_integration_pymc_examples.py @@ -8,6 +8,15 @@ @pytest.mark.integration def test_did(): + """ + Test Difference in Differences (DID) PyMC experiment. + + Loads data and checks: + 1. data is a dataframe + 2. pymc_experiements.DifferenceInDifferences 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 + """ df = cp.load_data("did") result = cp.pymc_experiments.DifferenceInDifferences( df, @@ -27,6 +36,18 @@ def test_did(): @pytest.mark.integration def test_did_banks_simple(): + """ + Test simple Differences In Differences Experiment on the 'banks' data set. + + formula="bib ~ 1 + district * post_treatment" + + Loads, transforms data and checks: + 1. data is a dataframe + 2. pymc_experiements.DifferenceInDifferences 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 + + """ treatment_time = 1930.5 df = ( cp.load_data("banks") @@ -67,6 +88,18 @@ def test_did_banks_simple(): @pytest.mark.integration def test_did_banks_multi(): + """ + Test multiple regression Differences In Differences Experiment on the 'banks' + data set. + + formula="bib ~ 1 + year + district + post_treatment + district:post_treatment" + + Loads, transforms data and checks: + 1. data is a dataframe + 2. pymc_experiements.DifferenceInDifferences 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 + """ treatment_time = 1930.5 df = ( cp.load_data("banks") @@ -106,6 +139,15 @@ def test_did_banks_multi(): @pytest.mark.integration def test_rd(): + """ + Test Regression Discontinuity experiment. + + Loads data and checks: + 1. data is a dataframe + 2. pymc_experiments.RegressionDiscontinuity 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 + """ df = cp.load_data("rd") result = cp.pymc_experiments.RegressionDiscontinuity( df, @@ -122,6 +164,15 @@ def test_rd(): @pytest.mark.integration def test_rd_bandwidth(): + """ + Test Regression Discontinuity experiment with bandwidth parameter. + + Loads data and checks: + 1. data is a dataframe + 2. pymc_experiments.RegressionDiscontinuity 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 + """ df = cp.load_data("rd") result = cp.pymc_experiments.RegressionDiscontinuity( df, @@ -139,6 +190,15 @@ def test_rd_bandwidth(): @pytest.mark.integration def test_rd_drinking(): + """ + Test Regression Discontinuity experiment on drinking age data. + + Loads data and checks: + 1. data is a dataframe + 2. pymc_experiments.RegressionDiscontinuity 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 + """ df = ( cp.load_data("drinking") .rename(columns={"agecell": "age"}) @@ -159,6 +219,15 @@ def test_rd_drinking(): @pytest.mark.integration def test_its(): + """ + Test Interrupted Time-Series experiment. + + Loads data and checks: + 1. data is a dataframe + 2. pymc_experiments.SyntheticControl 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 + """ df = ( cp.load_data("its") .assign(date=lambda x: pd.to_datetime(x["date"])) @@ -179,6 +248,16 @@ def test_its(): @pytest.mark.integration def test_its_covid(): + """ + Test Interrupted Time-Series experiment on COVID data. + + Loads data and checks: + 1. data is a dataframe + 2. pymc_experiments.InterruptedtimeSeries 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 + """ + df = ( cp.load_data("covid") .assign(date=lambda x: pd.to_datetime(x["date"])) @@ -199,6 +278,16 @@ def test_its_covid(): @pytest.mark.integration def test_sc(): + """ + Test Synthetic Control experiment. + + Loads data and checks: + 1. data is a dataframe + 2. pymc_experiments.SyntheticControl 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 + """ + df = cp.load_data("sc") treatment_time = 70 result = cp.pymc_experiments.SyntheticControl( @@ -215,6 +304,16 @@ def test_sc(): @pytest.mark.integration def test_sc_brexit(): + """ + Test Synthetic Control experiment on Brexit data. + + Loads data and checks: + 1. data is a dataframe + 2. pymc_experiments.SyntheticControl 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 + """ + df = ( cp.load_data("brexit") .assign(Time=lambda x: pd.to_datetime(x["Time"])) @@ -243,6 +342,15 @@ def test_sc_brexit(): @pytest.mark.integration def test_ancova(): + """ + Test Pre-PostNEGD experiment on anova1 data. + + Loads data and checks: + 1. data is a dataframe + 2. pymc_experiments.PrePostNEGD 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 + """ df = cp.load_data("anova1") result = cp.pymc_experiments.PrePostNEGD( df, @@ -259,6 +367,15 @@ def test_ancova(): @pytest.mark.integration def test_geolift1(): + """ + Test Synthetic Control experiment on geo lift data. + + Loads data and checks: + 1. data is a dataframe + 2. pymc_experiments.SyntheticControl 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 + """ df = ( cp.load_data("geolift1") .assign(time=lambda x: pd.to_datetime(x["time"])) diff --git a/causalpy/tests/test_integration_skl_examples.py b/causalpy/tests/test_integration_skl_examples.py index c5295613..4c863c79 100644 --- a/causalpy/tests/test_integration_skl_examples.py +++ b/causalpy/tests/test_integration_skl_examples.py @@ -9,6 +9,14 @@ @pytest.mark.integration def test_did(): + """ + Test Difference in Differences (DID) Sci-Kit Learn experiment. + + Loads data and checks: + 1. data is a dataframe + 2. skl_experiements.DifferenceInDifferences returns correct type + """ + data = cp.load_data("did") result = cp.skl_experiments.DifferenceInDifferences( data, @@ -25,6 +33,13 @@ def test_did(): @pytest.mark.integration def test_rd_drinking(): + """ + Test Regression Discontinuity Sci-Kit Learn experiment on drinking age data. + + Loads data and checks: + 1. data is a dataframe + 2. skl_experiements.RegressionDiscontinuity returns correct type + """ df = ( cp.load_data("drinking") .rename(columns={"agecell": "age"}) @@ -44,6 +59,14 @@ def test_rd_drinking(): @pytest.mark.integration def test_its(): + """ + Test Interrupted Time Series Sci-Kit Learn experiment. + + Loads data and checks: + 1. data is a dataframe + 2. skl_experiements.SyntheticControl returns correct type + """ + df = ( cp.load_data("its") .assign(date=lambda x: pd.to_datetime(x["date"])) @@ -62,6 +85,13 @@ def test_its(): @pytest.mark.integration def test_sc(): + """ + Test Synthetic Control Sci-Kit Learn experiment. + + Loads data and checks: + 1. data is a dataframe + 2. skl_experiements.SyntheticControl returns correct type + """ df = cp.load_data("sc") treatment_time = 70 result = cp.skl_experiments.SyntheticControl( @@ -76,6 +106,13 @@ def test_sc(): @pytest.mark.integration def test_rd_linear_main_effects(): + """ + Test Regression Discontinuity Sci-Kit Learn experiment main effects. + + Loads data and checks: + 1. data is a dataframe + 2. skl_experiements.RegressionDiscontinuity returns correct type + """ data = cp.load_data("rd") result = cp.skl_experiments.RegressionDiscontinuity( data, @@ -90,6 +127,14 @@ def test_rd_linear_main_effects(): @pytest.mark.integration def test_rd_linear_main_effects_bandwidth(): + """ + Test Regression Discontinuity Sci-Kit Learn experiment, main effects with + bandwidth parameter. + + Loads data and checks: + 1. data is a dataframe + 2. skl_experiements.RegressionDiscontinuity returns correct type + """ data = cp.load_data("rd") result = cp.skl_experiments.RegressionDiscontinuity( data, @@ -105,6 +150,13 @@ def test_rd_linear_main_effects_bandwidth(): @pytest.mark.integration def test_rd_linear_with_interaction(): + """ + Test Regression Discontinuity Sci-Kit Learn experiment with interaction. + + Loads data and checks: + 1. data is a dataframe + 2. skl_experiements.RegressionDiscontinuity returns correct type + """ data = cp.load_data("rd") result = cp.skl_experiments.RegressionDiscontinuity( data, @@ -119,6 +171,13 @@ def test_rd_linear_with_interaction(): @pytest.mark.integration def test_rd_linear_with_gaussian_process(): + """ + Test Regression Discontinuity Sci-Kit Learn experiment with Gaussian process model. + + Loads data and checks: + 1. data is a dataframe + 2. skl_experiements.RegressionDiscontinuity returns correct type + """ data = cp.load_data("rd") kernel = 1.0 * ExpSineSquared(1.0, 5.0) + WhiteKernel(1e-1) result = cp.skl_experiments.RegressionDiscontinuity( diff --git a/causalpy/tests/test_pymc_models.py b/causalpy/tests/test_pymc_models.py index 24e7d102..d535ba65 100644 --- a/causalpy/tests/test_pymc_models.py +++ b/causalpy/tests/test_pymc_models.py @@ -11,7 +11,16 @@ class MyToyModel(ModelBuilder): + """ + A subclass of ModelBuilder with a simple regression model for use in tests. + """ + def build_model(self, X, y, coords): + """ + Required .build_model() method of a ModelBuilder subclass + + This is a basic 1-variable linear regression model for use in tests. + """ with self: X_ = pm.MutableData(name="X", value=X) y_ = pm.MutableData(name="y", value=y) @@ -22,7 +31,17 @@ def build_model(self, X, y, coords): class TestModelBuilder: + """ + Related tests that check aspects of ModelBuilder objects. + """ + def test_init(self): + """ + Test initialization. + + Creates ModelBuilder() object and checks that idata is None and no sample + kwargs are specified. + """ mb = ModelBuilder() assert mb.idata is None assert mb.sample_kwargs == {} @@ -37,12 +56,20 @@ def test_init(self): argnames="X", argvalues=[np.ones(2), None], ids=["X-array", "X-None"] ) def test_model_builder(self, X, y, coords) -> None: + """ + Tests that a ModelBuilder() object without a .build_model() method raises + appropriate exception. + """ with pytest.raises( NotImplementedError, match="This method must be implemented by a subclass" ): ModelBuilder().build_model(X=X, y=y, coords=coords) def test_fit_build_not_implemented(self): + """ + Tests that a ModelBuilder() object without a .fit() method raises appropriate + exception. + """ with pytest.raises( NotImplementedError, match="This method must be implemented by a subclass" ): @@ -54,6 +81,16 @@ def test_fit_build_not_implemented(self): ids=["None-coords", "dict-coords"], ) def test_fit_predict(self, coords, rng) -> None: + """ + Test fit and predict methods on MyToyModel. + + Generates normal data, fits the model, makes predictions, scores the model + then: + 1. checks that model.idata is az.InferenceData type + 2. checks that beta, sigma, mu, and y_hat can be extract from idata + 3. checks score is a pandas series of the correct shape + 4. checks that predictions are az.InferenceData type + """ X = rng.normal(loc=0, scale=1, size=(20, 2)) y = rng.normal(loc=0, scale=1, size=(20,)) model = MyToyModel(sample_kwargs={"chains": 2, "draws": 2}) diff --git a/docs/source/_static/interrogate_badge.svg b/docs/source/_static/interrogate_badge.svg index aa7bbc1d..4567996b 100644 --- a/docs/source/_static/interrogate_badge.svg +++ b/docs/source/_static/interrogate_badge.svg @@ -1,10 +1,10 @@ - interrogate: 52.4% + interrogate: 86.0% - + @@ -12,8 +12,8 @@ interrogate interrogate - 52.4% - 52.4% + 86.0% + 86.0% From 809277e7d2e7ae522a49acb25db45e13c2105fe6 Mon Sep 17 00:00:00 2001 From: Jason Preszler Date: Tue, 22 Aug 2023 15:26:08 -0700 Subject: [PATCH 02/23] Examples for pymc experiments and most models --- causalpy/data/simulate_data.py | 47 +++++-- causalpy/plot_utils.py | 17 ++- causalpy/pymc_experiments.py | 149 +++++++++++++++++++++- causalpy/pymc_models.py | 138 ++++++++++++++++++-- docs/source/_static/interrogate_badge.svg | 6 +- 5 files changed, 328 insertions(+), 29 deletions(-) diff --git a/causalpy/data/simulate_data.py b/causalpy/data/simulate_data.py index b752fbbe..e9d66b13 100644 --- a/causalpy/data/simulate_data.py +++ b/causalpy/data/simulate_data.py @@ -29,12 +29,13 @@ def generate_synthetic_control_data( lowess_kwargs=default_lowess_kwargs, ): """ - Example: - >> import pathlib - >> df, weightings_true = generate_synthetic_control_data( - treatment_time=treatment_time - ) - >> df.to_csv(pathlib.Path.cwd() / 'synthetic_control.csv', index=False) + Example + -------- + >>> import pathlib + >>> df, weightings_true = generate_synthetic_control_data( + ... treatment_time=treatment_time + ... ) + >>> df.to_csv(pathlib.Path.cwd() / 'synthetic_control.csv', index=False) """ # 1. Generate non-treated variables @@ -73,6 +74,7 @@ def generate_synthetic_control_data( def generate_time_series_data( N=100, treatment_time=70, beta_temp=-1, beta_linear=0.5, beta_intercept=3 ): + """ """ x = np.arange(0, 100, 1) df = pd.DataFrame( { @@ -102,6 +104,7 @@ def generate_time_series_data( def generate_time_series_data_seasonal(treatment_time): + """ """ dates = pd.date_range( start=pd.to_datetime("2010-01-01"), end=pd.to_datetime("2020-01-01"), freq="M" ) @@ -149,6 +152,13 @@ def generate_time_series_data_simple(treatment_time, slope=0.0): def generate_did(): + """ + Generate Difference in Differences data + + Example + -------- + >>> df = generate_did() + """ # true parameters control_intercept = 1 treat_intercept_delta = 0.25 @@ -194,10 +204,13 @@ def generate_regression_discontinuity_data( N=100, true_causal_impact=0.5, true_treatment_threshold=0.0 ): """ - Example use: - >> import pathlib - >> df = generate_regression_discontinuity_data(true_treatment_threshold=0.5) - >> df.to_csv(pathlib.Path.cwd() / 'regression_discontinuity.csv', index=False) + Generate regression discontinuity example data + + Example + -------- + >>> import pathlib + >>> df = generate_regression_discontinuity_data(true_treatment_threshold=0.5) + >>> df.to_csv(pathlib.Path.cwd() / 'regression_discontinuity.csv', index=False) """ def is_treated(x): @@ -217,6 +230,20 @@ def impact(x): def generate_ancova_data( N=200, pre_treatment_means=np.array([10, 12]), treatment_effect=2, sigma=1 ): + """ + Generate ANCOVA eample data + + Example + -------- + >>> import pathlib + >>> df = generate_ancova_data( + ... N=200, + ... pre_treatment_threshold=np.array([10, 12]), + ... treatment_effect=2, + ... sigma=1 + ... ) + >>> df.to_csv(pathlib.Path.cwd() / 'ancova_data.csv', index=False) + """ group = np.random.choice(2, size=N) pre = np.random.normal(loc=pre_treatment_means[group]) post = pre + treatment_effect * group + np.random.normal(size=N) * sigma diff --git a/causalpy/plot_utils.py b/causalpy/plot_utils.py index 52d44bbc..915ff857 100644 --- a/causalpy/plot_utils.py +++ b/causalpy/plot_utils.py @@ -21,7 +21,22 @@ def plot_xY( hdi_prob: float = 0.94, label: Union[str, None] = None, ) -> Tuple[Line2D, PolyCollection]: - """Utility function to plot HDI intervals.""" + """ + Utility function to plot HDI intervals. + + :param x: + Pandas datetime index or numpy array of x-axis values + :param y: + Xarray data array of y-axis data + :param ax: + Matplotlib ax object + :param plot_hdi_kwargs: + Dictionary of keyword arguments passed to ax.plot() + :param hdi_prob: + The size of the HDI, default is 0.94 + :param label: + The plot label + """ if plot_hdi_kwargs is None: plot_hdi_kwargs = {} diff --git a/causalpy/pymc_experiments.py b/causalpy/pymc_experiments.py index fb394ad8..812113de 100644 --- a/causalpy/pymc_experiments.py +++ b/causalpy/pymc_experiments.py @@ -30,7 +30,7 @@ class ExperimentalDesign: - """Base class""" + """Base class for other experiment types""" model = None expt_type = None @@ -43,7 +43,7 @@ def __init__(self, model=None, **kwargs): @property def idata(self): - """Access to the InferenceData object""" + """Access to the models InferenceData object""" return self.model.idata def print_coefficients(self) -> None: @@ -66,8 +66,32 @@ def print_coefficients(self) -> None: class PrePostFit(ExperimentalDesign): - """A class to analyse quasi-experiments where parameter estimation is based on just - the pre-intervention data.""" + """ + A class to analyse quasi-experiments where parameter estimation is based on just + the pre-intervention data. + + :param data: + A pandas data frame + :param treatment_time: + The time when treatment occured, should be in reference to the data index + :param formula: + A statistical model formula + :param model: + A PyMC model + + Example + -------- + >>> sc = cp.load_data("sc") + >>> seed = 42 + >>> result = cp.pymc_experiments.PrePostFit( + ... sc, + ... treatment_time, + ... formula="actual ~ 0 + a + b + c + d + e + f + g", + ... model=cp.pymc_models.WeightedSumFitter( + ... sample_kwargs={"target_accept": 0.95, "random_seed": seed} + ... ), + ... ) + """ def __init__( self, @@ -256,13 +280,64 @@ def summary(self) -> None: class InterruptedTimeSeries(PrePostFit): - """Interrupted time series analysis""" + """ + A wrapper around PrePostFit class + + :param data: + A pandas data frame + :param treatment_time: + The time when treatment occured, should be in reference to the data index + :param formula: + A statistical model formula + :param model: + A PyMC model + + Example + -------- + >>> df = ( + ... cp.load_data("its") + ... .assign(date=lambda x: pd.to_datetime(x["date"])) + ... .set_index("date") + ... ) + >>> treatment_time = pd.to_datetime("2017-01-01") + >>> seed = 42 + >>> result = cp.pymc_experiments.InterruptedTimeSeries( + ... df, + ... treatment_time, + ... formula="y ~ 1 + t + C(month)", + ... model=cp.pymc_models.LinearRegression(sample_kwargs={"random_seed": seed}), + ... ) + """ expt_type = "Interrupted Time Series" class SyntheticControl(PrePostFit): - """A wrapper around the PrePostFit class""" + """A wrapper around the PrePostFit class + + :param data: + A pandas data frame + :param treatment_time: + The time when treatment occured, should be in reference to the data index + :param formula: + A statistical model formula + :param model: + A PyMC model + + Example + -------- + >>> df = cp.load_data("sc") + >>> treatment_time = 70 + >>> seed = 42 + >>> result = cp.pymc_experiments.SyntheticControl( + ... df, + ... treatment_time, + ... formula="actual ~ 0 + a + b + c + d + e + f + g", + ... model=cp.pymc_models.WeightedSumFitter( + ... sample_kwargs={"target_accept": 0.95, "random_seed": seed} + ... ), + ... ) + """ expt_type = "Synthetic Control" @@ -285,6 +360,28 @@ class DifferenceInDifferences(ExperimentalDesign): There is no pre/post intervention data distinction for DiD, we fit all the data available. + :param data: + A pandas data frame + :param formula: + A statistical model formula + :param time_variable_name: + Name of the data column for the time variable + :param group_variable_name: + Name of the data column for the group variable + :param model: + A PyMC model for difference in differences + + Example + -------- + >>> df = cp.load_data("did") + >>> seed = 42 + >>> result = cp.pymc_experiments.DifferenceInDifferences( + ... df, + ... formula="y ~ 1 + group*post_treatment", + ... time_variable_name="t", + ... group_variable_name="group", + ... model=cp.pymc_models.LinearRegression(sample_kwargs={"random_seed": seed}), + ... ) """ @@ -572,6 +669,18 @@ class RegressionDiscontinuity(ExperimentalDesign): :param bandwidth: Data outside of the bandwidth (relative to the discontinuity) is not used to fit the model. + + Example + -------- + >>> df = cp.load_data("rd") + >>> seed = 42 + >>> result = cp.pymc_experiments.RegressionDiscontinuity( + ... df, + ... formula="y ~ 1 + x + treated + x:treated", + ... model=cp.pymc_models.LinearRegression(sample_kwargs={"random_seed": seed}), + ... treatment_threshold=0.5, + ... ) + """ def __init__( @@ -742,7 +851,33 @@ def summary(self) -> None: class PrePostNEGD(ExperimentalDesign): - """A class to analyse data from pretest/posttest designs""" + """ + A class to analyse data from pretest/posttest designs + + :param data: + A pandas data frame + :param formula: + A statistical model formula + :param group_variable_name: + Name of the column in data for the group variable + :param pretreatment_variable_name: + Name of the column in data for the pretreatment variable + :param model: + A PyMC model + + Example + -------- + >>> df = cp.load_data("anova1") + >>> seed = 42 + >>> result = cp.pymc_experiments.PrePostNEGD( + ... df, + ... formula="post ~ 1 + C(group) + pre", + ... group_variable_name="group", + ... pretreatment_variable_name="pre", + ... model=cp.pymc_models.LinearRegression(sample_kwargs={"random_seed": seed}), + ... ) + + """ def __init__( self, diff --git a/causalpy/pymc_models.py b/causalpy/pymc_models.py index ad54a94a..befa6399 100644 --- a/causalpy/pymc_models.py +++ b/causalpy/pymc_models.py @@ -1,3 +1,8 @@ +""" +Defines generic PyMC ModelBuilder class and subclasses for +* WeightedSumFitter model for Synthetic Control experiments +* LinearRegression model +""" from typing import Any, Dict, Optional import arviz as az @@ -11,6 +16,13 @@ class ModelBuilder(pm.Model): """ This is a wrapper around pm.Model to give scikit-learn like API. + + Public Methods + -------- + * build_model: must be implemented by subclasses + * fit: populates idata attribute + * predict: returns predictions on new data + * score: returns Bayesian R^2 """ def __init__(self, sample_kwargs: Optional[Dict[str, Any]] = None): @@ -40,13 +52,57 @@ def build_model(self, X, y, coords) -> None: raise NotImplementedError("This method must be implemented by a subclass") def _data_setter(self, X) -> None: - """Set data for the model.""" + """ + Set data for the model. + + This method is used internally to register new data for the model for + prediction. + """ with self.model: pm.set_data({"X": X}) def fit(self, X, y, coords: Optional[Dict[str, Any]] = None) -> None: - """Draw samples from posterior, prior predictive, and posterior predictive - distributions. + """Draw samples fromposterior, prior predictive, and posterior predictive + distributions, placing them in the model's idata attribute. + + Example + ------- + >>> import numpy as np + >>> import pymc as pm + >>> from causalpy.pymc_models import ModelBuilder + >>> class MyToyModel(ModelBuilder): + ... def build_model(self, X, y, coords): + ... with self: + ... X_ = pm.MutableData(name="X", value=X) + ... y_ = pm.MutableData(name="y", value=y) + ... beta = pm.Normal("beta", mu=0, sigma=1, shape=X_.shape[1]) + ... sigma = pm.HalfNormal("sigma", sigma=1) + ... mu = pm.Deterministic("mu", pm.math.dot(X_, beta)) + ... pm.Normal("y_hat", mu=mu, sigma=sigma, observed=y_) + + >>> rng = np.random.default_rng(seed=42) + >>> X = rng.normal(loc=0, scale=1, size=(20, 2)) + >>> y = rng.normal(loc=0, scale=1, size=(20,)) + >>> model = MyToyModel(sample_kwargs={"chains": 2, "draws": 2}) + >>> model.fit(X, y) + Only 2 samples in chain. + Auto-assigning NUTS sampler... + Initializing NUTS using jitter+adapt_diag... + Multiprocess sampling (2 chains in 4 jobs) + NUTS: [beta, sigma] + Sampling 2 chains for 1_000 tune and 2 draw iterations (2_000 + 4 draws total) + took 0 seconds.gences] + The number of samples is too small to check convergence reliably. + Sampling: [beta, sigma, y_hat] + Sampling: [y_hat] + Inference data with groups: + > posterior + > posterior_predictive + > sample_stats + > prior + > prior_predictive + > observed_data + > constant_data """ self.build_model(X, y, coords) with self.model: @@ -58,7 +114,25 @@ def fit(self, X, y, coords: Optional[Dict[str, Any]] = None) -> None: return self.idata def predict(self, X): - """Predict data given input data `X`""" + """ + Predict data given input data `X` + + Results in KeyError if model hasn't been fit. + + Example + ------- + # Assumes `model` has been initialized and .fit() has been run, + # see ModelBuilder().fit() for example + + >>> X_new = rng.normal(loc=0, scale=1, size=(20,2)) + >>> model.predict(X_new) + Sampling: [beta, y_hat] + Inference data with groups: + > posterior_predictive + > observed_data + > constant_data + """ + self._data_setter(X) with self.model: # sample with new input data post_pred = pm.sample_posterior_predictive( @@ -74,6 +148,14 @@ def score(self, X, y) -> pd.Series: The Bayesian :math:`R^2` is not the same as the traditional coefficient of determination, https://en.wikipedia.org/wiki/Coefficient_of_determination. + Example + -------- + # Assuming `model` has been fit + >>> model.score(X, y) # X, y are random data here + Sampling: [y_hat] + r2 0.352251 + r2_std 0.051624 + dtype: float64 """ yhat = self.predict(X) yhat = az.extract( @@ -86,10 +168,33 @@ def score(self, X, y) -> pd.Series: class WeightedSumFitter(ModelBuilder): - """Used for synthetic control experiments""" + """ + Used for synthetic control experiments + + Defines model: + y ~ Normal(mu, sigma) + sigma ~ HalfNormal(1) + mu = X * beta + beta ~ Dirichlet(1,...,1) + + Public Methods + --------------- + * build_model + """ def build_model(self, X, y, coords): - """Defines the PyMC model""" + """ + Defines the PyMC model: + + y ~ Normal(mu, sigma) + sigma ~ HalfNormal(1) + mu = X * beta + beta ~ Dirichlet(1,...,1) + + Example + -------- + + """ with self: self.add_coords(coords) n_predictors = X.shape[1] @@ -107,10 +212,27 @@ def build_model(self, X, y, coords): class LinearRegression(ModelBuilder): - """Custom PyMC model for linear regression""" + """ + Custom PyMC model for linear regression + + Public Methods + --------------- + * build_model + """ def build_model(self, X, y, coords): - """Defines the PyMC model""" + """ + Defines the PyMC model + + y ~ Normal(mu, sigma) + mu = X * beta + beta ~ Normal(0, 50) + sigma ~ HalfNormal(1) + + Example + -------- + + """ with self: self.add_coords(coords) X = pm.MutableData("X", X, dims=["obs_ind", "coeffs"]) diff --git a/docs/source/_static/interrogate_badge.svg b/docs/source/_static/interrogate_badge.svg index 4567996b..f334c8ee 100644 --- a/docs/source/_static/interrogate_badge.svg +++ b/docs/source/_static/interrogate_badge.svg @@ -1,5 +1,5 @@ - interrogate: 86.0% + interrogate: 88.2% @@ -12,8 +12,8 @@ interrogate interrogate - 86.0% - 86.0% + 88.2% + 88.2% From fd9261368ee0a526690387bde363ab9fc7814a0d Mon Sep 17 00:00:00 2001 From: Jason Preszler Date: Tue, 22 Aug 2023 16:09:30 -0700 Subject: [PATCH 03/23] increased interrogate threshold, just need examples in skl experiements and some pymc models --- causalpy/data/datasets.py | 3 ++ causalpy/data/simulate_data.py | 54 +++++++++++++++++++++-- causalpy/tests/conftest.py | 6 +++ causalpy/tests/test_data_loading.py | 3 ++ causalpy/tests/test_utils.py | 4 ++ causalpy/utils.py | 3 ++ docs/source/_static/interrogate_badge.svg | 8 ++-- pyproject.toml | 2 +- 8 files changed, 74 insertions(+), 9 deletions(-) diff --git a/causalpy/data/datasets.py b/causalpy/data/datasets.py index 5b7c1d39..97014345 100644 --- a/causalpy/data/datasets.py +++ b/causalpy/data/datasets.py @@ -1,3 +1,6 @@ +""" +Functions to load example datasets +""" import pathlib import pandas as pd diff --git a/causalpy/data/simulate_data.py b/causalpy/data/simulate_data.py index e9d66b13..5dcdca42 100644 --- a/causalpy/data/simulate_data.py +++ b/causalpy/data/simulate_data.py @@ -1,3 +1,6 @@ +""" +Functions that generate data sets used in examples +""" import numpy as np import pandas as pd from scipy.stats import dirichlet, gamma, norm, uniform @@ -13,6 +16,15 @@ def _smoothed_gaussian_random_walk( ): """ Generates Gaussian random walk data and applies LOWESS + + :param gaussian_random_walk_mu: + Mean of the random walk + :param gaussian_random_walk_sigma: + Standard deviation of the random walk + :param N: + Length of the random walk + :param lowess_kwargs: + Keyword argument dictionary passed to statsmodels lowess """ x = np.arange(N) y = norm(gaussian_random_walk_mu, gaussian_random_walk_sigma).rvs(N).cumsum() @@ -29,13 +41,24 @@ def generate_synthetic_control_data( lowess_kwargs=default_lowess_kwargs, ): """ + Generates data for synthetic control example. + + :param N: + Number fo data points + :param treatment_time: + Index where treatment begins in the generated data frame + :param grw_mu: + Mean of Gaussian Random Walk + :param grw_sigma: + Standard deviation of Gaussian Random Walk + :lowess_kwargs: + Keyword argument dictionary passed to statsmodels lowess + Example -------- - >>> import pathlib >>> df, weightings_true = generate_synthetic_control_data( ... treatment_time=treatment_time ... ) - >>> df.to_csv(pathlib.Path.cwd() / 'synthetic_control.csv', index=False) """ # 1. Generate non-treated variables @@ -74,7 +97,21 @@ def generate_synthetic_control_data( def generate_time_series_data( N=100, treatment_time=70, beta_temp=-1, beta_linear=0.5, beta_intercept=3 ): - """ """ + """ + Generates interrupted time series example data + + :param N: + Length of the time series + :param treatment_time: + Index of when treatment begins + :param beta_temp: + The temperature coefficient + :param beta_linear: + The linear coefficient + :param beta_intercept: + The intercept + + """ x = np.arange(0, 100, 1) df = pd.DataFrame( { @@ -104,7 +141,9 @@ def generate_time_series_data( def generate_time_series_data_seasonal(treatment_time): - """ """ + """ + Generates 10 years of monthly data with seasonality + """ dates = pd.date_range( start=pd.to_datetime("2010-01-01"), end=pd.to_datetime("2020-01-01"), freq="M" ) @@ -170,6 +209,7 @@ def generate_did(): def outcome( t, control_intercept, treat_intercept_delta, trend, Δ, group, post_treatment ): + """Compute the outcome of each unit""" return ( control_intercept + (treat_intercept_delta * group) @@ -214,9 +254,11 @@ def generate_regression_discontinuity_data( """ def is_treated(x): + """Check if x was treated""" return np.greater_equal(x, true_treatment_threshold) def impact(x): + """Assign true_causal_impact to all treaated entries""" y = np.zeros(len(x)) y[is_treated(x)] = true_causal_impact return y @@ -263,6 +305,10 @@ def generate_geolift_data(): causal_impact = 0.2 def create_series(n=52, amplitude=1, length_scale=2): + """ + Returns numpy tile with generated seasonality data repeated over + multiple years + """ return np.tile( generate_seasonality(n=n, amplitude=amplitude, length_scale=2) + 3, n_years ) diff --git a/causalpy/tests/conftest.py b/causalpy/tests/conftest.py index d652cc8f..92a2a77e 100644 --- a/causalpy/tests/conftest.py +++ b/causalpy/tests/conftest.py @@ -1,3 +1,9 @@ +""" +CausalPy Test Configuration + +Functions: +* rng: random number generator with session level scope +""" import numpy as np import pytest diff --git a/causalpy/tests/test_data_loading.py b/causalpy/tests/test_data_loading.py index 5d562aa8..343c2c22 100644 --- a/causalpy/tests/test_data_loading.py +++ b/causalpy/tests/test_data_loading.py @@ -1,3 +1,6 @@ +""" +Tests that example data can be loaded into data frames. +""" import pandas as pd import pytest diff --git a/causalpy/tests/test_utils.py b/causalpy/tests/test_utils.py index fb972cb6..cad4c312 100644 --- a/causalpy/tests/test_utils.py +++ b/causalpy/tests/test_utils.py @@ -1,3 +1,7 @@ +""" +Tests for utility functions +""" + import pandas as pd from causalpy.utils import _is_variable_dummy_coded, _series_has_2_levels diff --git a/causalpy/utils.py b/causalpy/utils.py index 50e79694..850cc0ab 100644 --- a/causalpy/utils.py +++ b/causalpy/utils.py @@ -1,3 +1,6 @@ +""" +Utility functions +""" import pandas as pd diff --git a/docs/source/_static/interrogate_badge.svg b/docs/source/_static/interrogate_badge.svg index f334c8ee..377961aa 100644 --- a/docs/source/_static/interrogate_badge.svg +++ b/docs/source/_static/interrogate_badge.svg @@ -1,10 +1,10 @@ - interrogate: 88.2% + interrogate: 97.1% - + @@ -12,8 +12,8 @@ interrogate interrogate - 88.2% - 88.2% + 97.1% + 97.1% diff --git a/pyproject.toml b/pyproject.toml index 2de41c98..687e2b82 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -114,7 +114,7 @@ ignore-module = false ignore-nested-functions = false ignore-nested-classes = true ignore-setters = false -fail-under = 40 +fail-under = 85 exclude = ["setup.py", "docs", "build"] ignore-regex = ["^get$", "^mock_.*", ".*BaseClass.*"] # possible values: 0 (minimal output), 1 (-v), 2 (-vv) From 8a9800f744f452a6c0b58746522d4357a174ec09 Mon Sep 17 00:00:00 2001 From: Jason Preszler Date: Wed, 23 Aug 2023 15:23:32 -0700 Subject: [PATCH 04/23] pymc models and experiments done and docs checked --- causalpy/pymc_experiments.py | 248 ++++++++++++++++++++++++++++++++--- causalpy/pymc_models.py | 138 +++++++++++++------ 2 files changed, 329 insertions(+), 57 deletions(-) diff --git a/causalpy/pymc_experiments.py b/causalpy/pymc_experiments.py index 812113de..c98095a2 100644 --- a/causalpy/pymc_experiments.py +++ b/causalpy/pymc_experiments.py @@ -1,13 +1,16 @@ """ Experiment routines for PyMC models. -Includes: -1. ExperimentalDesign base class -2. Pre-Post Fit -3. Synthetic Control -4. Difference in differences -5. Regression Discontinuity +- ExperimentalDesign base class +- Pre-Post Fit +- Interrupted Time Series +- Synthetic Control +- Difference in differences +- Regression Discontinuity +- Pretest/Posttest Nonequivalent Group Design + """ + import warnings from typing import Optional, Union @@ -30,7 +33,11 @@ class ExperimentalDesign: - """Base class for other experiment types""" + """ + Base class for other experiment types + + See subclasses for examples of most methods + """ model = None expt_type = None @@ -43,11 +50,63 @@ def __init__(self, model=None, **kwargs): @property def idata(self): - """Access to the models InferenceData object""" + """ + Access to the models InferenceData object + + Example + -------- + If `result` is the result of the Difference in Differences experiment example + + >>> result.idata + Inference data with groups: + > posterior + > posterior_predictive + > sample_stats + > prior + > prior_predictive + > observed_data + > constant_data + >>> result.idata.posterior + + Dimensions: (chain: 4, draw: 1000, coeffs: 4, obs_ind: 40) + Coordinates: + * chain (chain) int64 0 1 2 3 + * draw (draw) int64 0 1 2 3 4 5 6 7 8 ... 992 993 994 995 996 997 998 + 999 + * coeffs (coeffs) None: - """Prints the model coefficients""" + """ + Prints the model coefficients + + Example + -------- + If `result` is from the Difference in Differences experiment example + + >>> result.print_coefficients() + Model coefficients: + Intercept 1.08, 94% HDI [1.03, 1.13] + post_treatment[T.True] 0.98, 94% HDI [0.91, 1.06] + group 0.16, 94% HDI [0.09, 0.23] + group:post_treatment[T.True] 0.51, 94% HDI [0.41, 0.61] + sigma 0.08, 94% HDI [0.07, 0.10] + """ print("Model coefficients:") coeffs = az.extract(self.idata.posterior, var_names="beta") # Note: f"{name: <30}" pads the name with spaces so that we have alignment of @@ -82,6 +141,7 @@ class PrePostFit(ExperimentalDesign): Example -------- >>> sc = cp.load_data("sc") + >>> treatment_time = 70 >>> seed = 42 >>> result = cp.pymc_experiments.PrePostFit( ... sc, @@ -91,6 +151,17 @@ class PrePostFit(ExperimentalDesign): ... sample_kwargs={"target_accept": 0.95, "random_seed": seed} ... ), ... ) + Auto-assigning NUTS sampler... + Initializing NUTS using jitter+adapt_diag... + Multiprocess sampling (4 chains in 4 jobs) + NUTS: [beta, sigma] + Sampling 4 chains for 1_000 tune and 1_000 draw iterations + (4_000 + 4_000 draws total) took 11 seconds. + Sampling: [beta, sigma, y_hat] + Sampling: [y_hat] + Sampling: [y_hat] + Sampling: [y_hat] + Sampling: [y_hat] """ def __init__( @@ -105,6 +176,8 @@ def __init__( self._input_validation(data, treatment_time) self.treatment_time = treatment_time + # set experiment type - usually done in subclasses + self.expt_type = "Pre-Post Fit" # split data in to pre and post intervention self.datapre = data[data.index <= self.treatment_time] self.datapost = data[data.index > self.treatment_time] @@ -171,7 +244,14 @@ def _input_validation(self, data, treatment_time): ) def plot(self, counterfactual_label="Counterfactual", **kwargs): - """Plot the results""" + """ + Plot the results + + Example + -------- + >>> result.plot() + + """ fig, ax = plt.subplots(3, 1, sharex=True, figsize=(7, 8)) # TOP PLOT -------------------------------------------------- @@ -271,7 +351,24 @@ def plot(self, counterfactual_label="Counterfactual", **kwargs): return (fig, ax) def summary(self) -> None: - """Print text output summarising the results""" + """ + Print text output summarising the results + + Example + --------- + >>> result.summary() + ===============================Synthetic Control=============================== + Formula: actual ~ 0 + a + b + c + d + e + f + g + Model coefficients: + a 0.33, 94% HDI [0.30, 0.38] + b 0.05, 94% HDI [0.01, 0.09] + c 0.31, 94% HDI [0.26, 0.35] + d 0.06, 94% HDI [0.01, 0.10] + e 0.02, 94% HDI [0.00, 0.06] + f 0.20, 94% HDI [0.12, 0.26] + g 0.04, 94% HDI [0.00, 0.08] + sigma 0.26, 94% HDI [0.22, 0.30] + """ print(f"{self.expt_type:=^80}") print(f"Formula: {self.formula}") @@ -307,6 +404,17 @@ class InterruptedTimeSeries(PrePostFit): ... formula="y ~ 1 + t + C(month)", ... model=cp.pymc_models.LinearRegression(sample_kwargs={"random_seed": seed}), ... ) + Auto-assigning NUTS sampler... + Initializing NUTS using jitter+adapt_diag... + Multiprocess sampling (4 chains in 4 jobs) + NUTS: [beta, sigma] + Sampling 4 chains for 1_000 tune and 1_000 draw iterations + (4_000 + 4_000 draws total) took 3 seconds. + Sampling: [beta, sigma, y_hat] + Sampling: [y_hat] + Sampling: [y_hat] + Sampling: [y_hat] + Sampling: [y_hat] """ expt_type = "Interrupted Time Series" @@ -337,6 +445,17 @@ class SyntheticControl(PrePostFit): ... sample_kwargs={"target_accept": 0.95, "random_seed": seed} ... ), ... ) + Auto-assigning NUTS sampler... + Initializing NUTS using jitter+adapt_diag... + Multiprocess sampling (4 chains in 4 jobs) + NUTS: [beta, sigma] + Sampling 4 chains for 1_000 tune and 1_000 draw iterations + (4_000 + 4_000 draws total) took 11 seconds. + Sampling: [beta, sigma, y_hat] + Sampling: [y_hat] + Sampling: [y_hat] + Sampling: [y_hat] + Sampling: [y_hat] """ expt_type = "Synthetic Control" @@ -382,7 +501,17 @@ class DifferenceInDifferences(ExperimentalDesign): ... group_variable_name="group", ... model=cp.pymc_models.LinearRegression(sample_kwargs={"random_seed": seed}), ... ) - + Auto-assigning NUTS sampler... + Initializing NUTS using jitter+adapt_diag... + Multiprocess sampling (4 chains in 4 jobs) + NUTS: [beta, sigma] + Sampling 4 chains for 1_000 tune and 1_000 draw iterations + (4_000 + 4_000 draws total) took 1 seconds. + Sampling: [beta, sigma, y_hat] + Sampling: [y_hat] + Sampling: [y_hat] + Sampling: [y_hat] + Sampling: [y_hat] """ def __init__( @@ -503,6 +632,12 @@ def _input_validation(self): def plot(self): """Plot the results. Creating the combined mean + HDI legend entries is a bit involved. + + Example + -------- + Assuming `result` is the result of a DiD experiment: + + >>> result.plot() """ fig, ax = plt.subplots() @@ -639,7 +774,25 @@ def _causal_impact_summary_stat(self) -> str: return f"Causal impact = {causal_impact + ci}" def summary(self) -> None: - """Print text output summarising the results""" + """ + Print text output summarising the results + + Example + -------- + Assuming `result` is a DiD experiment + + >>> result.summary() + ==========================Difference in Differences========================= + Formula: y ~ 1 + group*post_treatment + Results: + Causal impact = 0.51, $CI_{94%}$[0.41, 0.61] + Model coefficients: + Intercept 1.08, 94% HDI [1.03, 1.13] + post_treatment[T.True] 0.98, 94% HDI [0.91, 1.06] + group 0.16, 94% HDI [0.09, 0.23] + group:post_treatment[T.True] 0.51, 94% HDI [0.41, 0.61] + sigma 0.08, 94% HDI [0.07, 0.10] + """ print(f"{self.expt_type:=^80}") print(f"Formula: {self.formula}") @@ -680,7 +833,17 @@ class RegressionDiscontinuity(ExperimentalDesign): ... model=cp.pymc_models.LinearRegression(sample_kwargs={"random_seed": seed}), ... treatment_threshold=0.5, ... ) - + Auto-assigning NUTS sampler... + Initializing NUTS using jitter+adapt_diag... + Multiprocess sampling (4 chains in 4 jobs) + NUTS: [beta, sigma] + Sampling 4 chains for 1_000 tune and 1_000 draw iterations + (4_000 + 4_000 draws total) took 2 seconds. + Sampling: [beta, sigma, y_hat] + Sampling: [y_hat] + Sampling: [y_hat] + Sampling: [y_hat] + Sampling: [y_hat] """ def __init__( @@ -791,7 +954,13 @@ def _is_treated(self, x): return np.greater_equal(x, self.treatment_threshold) def plot(self): - """Plot the results""" + """ + Plot the results + + Example + -------- + >>> result.plot() + """ fig, ax = plt.subplots() # Plot raw data sns.scatterplot( @@ -837,7 +1006,25 @@ def plot(self): return (fig, ax) def summary(self) -> None: - """Print text output summarising the results""" + """ + Print text output summarising the results + + Example + -------- + >>> result.summary() + ============================Regression Discontinuity========================== + Formula: y ~ 1 + x + treated + x:treated + Running variable: x + Threshold on running variable: 0.5 + Results: + Discontinuity at threshold = 0.92 + Model coefficients: + Intercept 0.09, 94% HDI [0.00, 0.17] + treated[T.True] 2.48, 94% HDI [1.66, 3.27] + x 1.32, 94% HDI [1.14, 1.50] + x:treated[T.True] -3.12, 94% HDI [-4.17, -2.05] + sigma 0.35, 94% HDI [0.31, 0.41] + """ print(f"{self.expt_type:=^80}") print(f"Formula: {self.formula}") @@ -876,7 +1063,16 @@ class PrePostNEGD(ExperimentalDesign): ... pretreatment_variable_name="pre", ... model=cp.pymc_models.LinearRegression(sample_kwargs={"random_seed": seed}), ... ) - + Auto-assigning NUTS sampler... + Initializing NUTS using jitter+adapt_diag... + Multiprocess sampling (4 chains in 4 jobs) + NUTS: [beta, sigma] + Sampling 4 chains for 1_000 tune and 1_000 draw iterations + (4_000 + 4_000 draws total) took 3 seconds. + Sampling: [beta, sigma, y_hat] + Sampling: [y_hat] + Sampling: [y_hat] + Sampling: [y_hat] """ def __init__( @@ -1010,7 +1206,23 @@ def _causal_impact_summary_stat(self) -> str: return f"Causal impact = {causal_impact + ci}" def summary(self) -> None: - """Print text output summarising the results""" + """ + Print text output summarising the results + + Example + -------- + >>> result.summary() + =================Pretest/posttest Nonequivalent Group Design================ + Formula: post ~ 1 + C(group) + pre + Results: + Causal impact = 1.89, $CI_{94%}$[1.70, 2.07] + Model coefficients: + Intercept -0.46, 94% HDI [-1.17, 0.22] + C(group)[T.1] 1.89, 94% HDI [1.70, 2.07] + pre 1.05, 94% HDI [0.98, 1.12] + sigma 0.51, 94% HDI [0.46, 0.56] + + """ print(f"{self.expt_type:=^80}") print(f"Formula: {self.formula}") diff --git a/causalpy/pymc_models.py b/causalpy/pymc_models.py index befa6399..757ef597 100644 --- a/causalpy/pymc_models.py +++ b/causalpy/pymc_models.py @@ -1,7 +1,12 @@ """ Defines generic PyMC ModelBuilder class and subclasses for -* WeightedSumFitter model for Synthetic Control experiments -* LinearRegression model + +- WeightedSumFitter model for Synthetic Control experiments +- LinearRegression model + +Models are intended to be used from inside an experiment +class (see pymc_experiments.py). This is why the examples require some extra +manipulation input data, often to ensure `y` has the correct shape. """ from typing import Any, Dict, Optional @@ -18,11 +23,11 @@ class ModelBuilder(pm.Model): This is a wrapper around pm.Model to give scikit-learn like API. Public Methods - -------- - * build_model: must be implemented by subclasses - * fit: populates idata attribute - * predict: returns predictions on new data - * score: returns Bayesian R^2 + --------------- + - build_model: must be implemented by subclasses + - fit: populates idata attribute + - predict: returns predictions on new data + - score: returns Bayesian R^2 """ def __init__(self, sample_kwargs: Optional[Dict[str, Any]] = None): @@ -35,7 +40,7 @@ def __init__(self, sample_kwargs: Optional[Dict[str, Any]] = None): self.sample_kwargs = sample_kwargs if sample_kwargs is not None else {} def build_model(self, X, y, coords) -> None: - """Build the model. + """Build the model, must be implemented by subclass. Example ------- @@ -65,6 +70,9 @@ def fit(self, X, y, coords: Optional[Dict[str, Any]] = None) -> None: """Draw samples fromposterior, prior predictive, and posterior predictive distributions, placing them in the model's idata attribute. + .. note:: + Calls the build_model method + Example ------- >>> import numpy as np @@ -79,7 +87,6 @@ def fit(self, X, y, coords: Optional[Dict[str, Any]] = None) -> None: ... sigma = pm.HalfNormal("sigma", sigma=1) ... mu = pm.Deterministic("mu", pm.math.dot(X_, beta)) ... pm.Normal("y_hat", mu=mu, sigma=sigma, observed=y_) - >>> rng = np.random.default_rng(seed=42) >>> X = rng.normal(loc=0, scale=1, size=(20, 2)) >>> y = rng.normal(loc=0, scale=1, size=(20,)) @@ -117,12 +124,13 @@ def predict(self, X): """ Predict data given input data `X` - Results in KeyError if model hasn't been fit. + .. caution:: + Results in KeyError if model hasn't been fit. Example ------- - # Assumes `model` has been initialized and .fit() has been run, - # see ModelBuilder().fit() for example + Assumes `model` has been initialized and .fit() has been run, + see ModelBuilder().fit() for example. >>> X_new = rng.normal(loc=0, scale=1, size=(20,2)) >>> model.predict(X_new) @@ -150,7 +158,8 @@ def score(self, X, y) -> pd.Series: Example -------- - # Assuming `model` has been fit + Assuming `model` has been fit + >>> model.score(X, y) # X, y are random data here Sampling: [y_hat] r2 0.352251 @@ -171,28 +180,49 @@ class WeightedSumFitter(ModelBuilder): """ Used for synthetic control experiments - Defines model: - y ~ Normal(mu, sigma) - sigma ~ HalfNormal(1) - mu = X * beta - beta ~ Dirichlet(1,...,1) + .. note: Generally, the `.fit()` method should be rather than calling + `.build_model()` directly. - Public Methods - --------------- - * build_model + Defines the PyMC model: + + - y ~ Normal(mu, sigma) + - sigma ~ HalfNormal(1) + - mu = X * beta + - beta ~ Dirichlet(1,...,1) + + Example + -------- + >>> sc = cp.load_data("sc") + >>> X = sc[['a', 'b', 'c', 'd', 'e', 'f', 'g']] + >>> y = np.asarray(sc['actual']).reshape((sc.shape[0], 1)) + >>> wsf = WeightedSumFitter() + >>> wsf.fit(X,y) + Auto-assigning NUTS sampler... + Initializing NUTS using jitter+adapt_diag... + Multiprocess sampling (4 chains in 4 jobs) + NUTS: [beta, sigma] + Sampling 4 chains for 1_000 tune and 1_000 draw iterations + (4_000 + 4_000 draws total) took 3 seconds. + Sampling: [beta, sigma, y_hat] + Sampling: [y_hat] + Inference data with groups: + > posterior + > posterior_predictive + > sample_stats + > prior + > prior_predictive + > observed_data + > constant_data """ def build_model(self, X, y, coords): """ Defines the PyMC model: - y ~ Normal(mu, sigma) - sigma ~ HalfNormal(1) - mu = X * beta - beta ~ Dirichlet(1,...,1) - - Example - -------- + - y ~ Normal(mu, sigma) + - sigma ~ HalfNormal(1) + - mu = X * beta + - beta ~ Dirichlet(1,...,1) """ with self: @@ -215,23 +245,53 @@ class LinearRegression(ModelBuilder): """ Custom PyMC model for linear regression - Public Methods - --------------- - * build_model + .. note: Generally, the `.fit()` method should be rather than calling + `.build_model()` directly. + + Defines the PyMC model + + - y ~ Normal(mu, sigma) + - mu = X * beta + - beta ~ Normal(0, 50) + - sigma ~ HalfNormal(1) + + Example + -------- + >>> rd = cp.load_data("rd") + >>> X = rd[["x", "treated"]] + >>> y = np.asarray(rd["y"]).reshape((rd["y"].shape[0],1)) + >>> lr = LinearRegression() + >>> lr.fit(X, y, coords={ + 'coeffs': ['x', 'treated'], + 'obs_indx': np.arange(rd.shape[0]) + } + ) + Auto-assigning NUTS sampler... + Initializing NUTS using jitter+adapt_diag... + Multiprocess sampling (4 chains in 4 jobs) + NUTS: [beta, sigma] + Sampling 4 chains for 1_000 tune and 1_000 draw iterations ( + 4_000 + 4_000 draws total) took 1 seconds. + Sampling: [beta, sigma, y_hat] + Sampling: [y_hat] + Inference data with groups: + > posterior + > posterior_predictive + > sample_stats + > prior + > prior_predictive + > observed_data + > constant_data """ def build_model(self, X, y, coords): """ Defines the PyMC model - y ~ Normal(mu, sigma) - mu = X * beta - beta ~ Normal(0, 50) - sigma ~ HalfNormal(1) - - Example - -------- - + - y ~ Normal(mu, sigma) + - mu = X * beta + - beta ~ Normal(0, 50) + - sigma ~ HalfNormal(1) """ with self: self.add_coords(coords) From 110505aa262167caf8c3f20ee97ace6a4481f019 Mon Sep 17 00:00:00 2001 From: Jason Preszler Date: Wed, 23 Aug 2023 16:52:12 -0700 Subject: [PATCH 05/23] Finished docstrings and checked documentation results --- causalpy/pymc_models.py | 1 + causalpy/skl_experiments.py | 162 ++++++++++++++++++++++++++++++++++-- causalpy/skl_models.py | 27 +++++- 3 files changed, 182 insertions(+), 8 deletions(-) diff --git a/causalpy/pymc_models.py b/causalpy/pymc_models.py index 757ef597..da922e4c 100644 --- a/causalpy/pymc_models.py +++ b/causalpy/pymc_models.py @@ -7,6 +7,7 @@ Models are intended to be used from inside an experiment class (see pymc_experiments.py). This is why the examples require some extra manipulation input data, often to ensure `y` has the correct shape. + """ from typing import Any, Dict, Optional diff --git a/causalpy/skl_experiments.py b/causalpy/skl_experiments.py index b9209c4a..3471fdc5 100644 --- a/causalpy/skl_experiments.py +++ b/causalpy/skl_experiments.py @@ -1,5 +1,12 @@ """ Experiments for Scikit-Learn models + +- ExperimentalDesign: base class for skl experiments +- PrePostFit: base class for synthetic control and interrupted time series +- SyntheticControl +- InterruptedTimeSeries +- DifferenceInDifferences +- RegressionDiscontinuity """ import warnings from typing import Optional @@ -27,8 +34,33 @@ def __init__(self, model=None, **kwargs): class PrePostFit(ExperimentalDesign): - """A class to analyse quasi-experiments where parameter estimation is based on just - the pre-intervention data.""" + """ + A class to analyse quasi-experiments where parameter estimation is based on just + the pre-intervention data. + + :param data: + A pandas data frame + :param treatment_time: + The index or time value of when treatment begins + :param formula: + A statistical model formula + :param model: + An sklearn model object + + Example + -------- + >>> from sklearn.linear_model import LinearRegression + >>> import causalpy as cp + >>> df = cp.load_data("sc") + >>> treatment_time = 70 + >>> result = cp.skl_experiments.PrePostFit( + ... df, + ... treatment_time, + ... formula="actual ~ 0 + a + b + c + d + e + f + g", + ... model = cp.skl_models.WeightedProportion() + ... ) + + """ def __init__( self, @@ -144,7 +176,16 @@ def plot(self, counterfactual_label="Counterfactual", **kwargs): return (fig, ax) def get_coeffs(self): - """Returns model coefficients""" + """ + Returns model coefficients + + Example + -------- + >>> result.get_coeffs() + array([3.97370896e-01, 1.53881980e-01, 4.48747123e-01, 1.04639857e-16, + 0.00000000e+00, 0.00000000e+00, 2.92931287e-16]) + + """ return np.squeeze(self.model.coef_) def plot_coeffs(self): @@ -161,13 +202,68 @@ def plot_coeffs(self): class InterruptedTimeSeries(PrePostFit): - """Interrupted time series analysis""" + """ + Interrupted time series analysis, a wrapper around the PrePostFit class + + :param data: + A pandas data frame + :param treatment_time: + The index or time value of when treatment begins + :param formula: + A statistical model formula + :param model: + An sklearn model object + + Example + -------- + >>> from sklearn.linear_model import LinearRegression + >>> import pandas as pd + >>> import causalpy as cp + >>> df = ( + ... cp.load_data("its") + ... .assign(date=lambda x: pd.to_datetime(x["date"])) + ... .set_index("date") + ... ) + >>> treatment_time = pd.to_datetime("2017-01-01") + >>> result = cp.skl_experiments.InterruptedTimeSeries( + ... df, + ... treatment_time, + ... formula="y ~ 1 + t + C(month)", + ... model = LinearRegression() + ... ) + + """ expt_type = "Interrupted Time Series" class SyntheticControl(PrePostFit): - """A wrapper around the PrePostFit class""" + """ + A wrapper around the PrePostFit class + + :param data: + A pandas data frame + :param treatment_time: + The index or time value of when treatment begins + :param formula: + A statistical model formula + :param model: + An sklearn model object + + Example + -------- + >>> from sklearn.linear_model import LinearRegression + >>> import causalpy as cp + >>> df = cp.load_data("sc") + >>> treatment_time = 70 + >>> result = cp.skl_experiments.SyntheticControl( + ... df, + ... treatment_time, + ... formula="actual ~ 0 + a + b + c + d + e + f + g", + ... model = cp.skl_models.WeightedProportion() + ... ) + + """ def plot(self, plot_predictors=False, **kwargs): """Plot the results""" @@ -187,6 +283,32 @@ class DifferenceInDifferences(ExperimentalDesign): There is no pre/post intervention data distinction for DiD, we fit all the data available. + + :param data: + A pandas data frame + :param formula: + A statistical model formula + :param time_variable_name: + Name of the data column for the time variable + :param group_variable_name: + Name of the data column for the group variable + :param model: + A PyMC model for difference in differences + + Example + -------- + >>> df = cp.load_data("did") + >>> seed = 42 + >>> result = cp.skl_experiments.DifferenceInDifferences( + ... data, + ... formula="y ~ 1 + group*post_treatment", + ... time_variable_name="t", + ... group_variable_name="group", + ... treated=1, + ... untreated=0, + ... model=LinearRegression(), + ... ) + """ def __init__( @@ -373,6 +495,17 @@ class RegressionDiscontinuity(ExperimentalDesign): :param bandwidth: Data outside of the bandwidth (relative to the discontinuity) is not used to fit the model. + + Example + -------- + >>> data = cp.load_data("rd") + >>> result = cp.skl_experiments.RegressionDiscontinuity( + ... data, + ... formula="y ~ 1 + x + treated", + ... model=LinearRegression(), + ... treatment_threshold=0.5, + ... ) + """ def __init__( @@ -503,7 +636,24 @@ def plot(self): return (fig, ax) def summary(self): - """Print text output summarising the results""" + """ + Print text output summarising the results + + Example + -------- + >>> result.summary() + Difference in Differences experiment + Formula: y ~ 1 + x + treated + Running variable: x + Threshold on running variable: 0.5 + Results: + Discontinuity at threshold = 0.19 + Model coefficients: + Intercept 0.0 + treated[T.True] 0.19034196317793994 + x 1.229600855360073 + + """ print("Difference in Differences experiment") print(f"Formula: {self.formula}") print(f"Running variable: {self.running_variable_name}") diff --git a/causalpy/skl_models.py b/causalpy/skl_models.py index d2d0cae2..c7acc6c3 100644 --- a/causalpy/skl_models.py +++ b/causalpy/skl_models.py @@ -1,8 +1,8 @@ """ Scikit-Learn Models -Includes: -1. Weighted Proportion +- Weighted Proportion + """ from functools import partial @@ -21,6 +21,29 @@ class WeightedProportion(LinearModel, RegressorMixin): Inspiration taken from this blog post https://towardsdatascience.com/understanding-synthetic-control-methods-dd9a291885a1 + + Example + -------- + >>> rng = np.random.default_rng(seed=42) + >>> X = rng.normal(loc=0, scale=1, size=(20,2)) + >>> y = rng.normal(loc=0, scale=1, size=(20,)) + >>> wp.fit(X, y) + WeightedProportion() + >>> wp.coef_ + array([[0.36719946, 0.63280054]]) + >>> X_new = rng.normal(loc=0, scale=1, size=(10,2)) + >>> wp.predict(X_new) + array([[-0.8298643 ], + [ 0.43072465], + [ 0.76319257], + [-0.42062812], + [ 0.1939908 ], + [-1.18557609], + [-0.0230188 ], + [ 0.48923816], + [-0.05656294], + [ 0.0339618 ]]) + """ def loss(self, W, X, y): From adb94c5cee540a323e9dc08d16db50c71123fc8d Mon Sep 17 00:00:00 2001 From: Jason Preszler Date: Thu, 24 Aug 2023 07:32:35 -0700 Subject: [PATCH 06/23] small fixes related to pr comments --- causalpy/pymc_models.py | 5 ++++- causalpy/skl_experiments.py | 1 - 2 files changed, 4 insertions(+), 2 deletions(-) diff --git a/causalpy/pymc_models.py b/causalpy/pymc_models.py index da922e4c..97fe6875 100644 --- a/causalpy/pymc_models.py +++ b/causalpy/pymc_models.py @@ -5,10 +5,13 @@ - LinearRegression model Models are intended to be used from inside an experiment -class (see pymc_experiments.py). This is why the examples require some extra +class (see `pymc_experiments.py +`_). +This is why the examples require some extra manipulation input data, often to ensure `y` has the correct shape. """ + from typing import Any, Dict, Optional import arviz as az diff --git a/causalpy/skl_experiments.py b/causalpy/skl_experiments.py index 3471fdc5..68854758 100644 --- a/causalpy/skl_experiments.py +++ b/causalpy/skl_experiments.py @@ -298,7 +298,6 @@ class DifferenceInDifferences(ExperimentalDesign): Example -------- >>> df = cp.load_data("did") - >>> seed = 42 >>> result = cp.skl_experiments.DifferenceInDifferences( ... data, ... formula="y ~ 1 + group*post_treatment", From 5c2870e64db2bcdcc096283305ab29eb816213b0 Mon Sep 17 00:00:00 2001 From: Jason Preszler Date: Wed, 30 Aug 2023 20:45:54 -0700 Subject: [PATCH 07/23] dcotest on pymc_models good except rng caused score differences --- causalpy/pymc_models.py | 169 +++++++++++++++++++--------------------- 1 file changed, 82 insertions(+), 87 deletions(-) diff --git a/causalpy/pymc_models.py b/causalpy/pymc_models.py index 97fe6875..8d2b922b 100644 --- a/causalpy/pymc_models.py +++ b/causalpy/pymc_models.py @@ -48,15 +48,17 @@ def build_model(self, X, y, coords) -> None: Example ------- + >>> import pymc as pm + >>> from causalpy.pymc_models import ModelBuilder >>> class CausalPyModel(ModelBuilder): - >>> def build_model(self, X, y): - >>> with self: - >>> X_ = pm.MutableData(name="X", value=X) - >>> y_ = pm.MutableData(name="y", value=y) - >>> beta = pm.Normal("beta", mu=0, sigma=1, shape=X_.shape[1]) - >>> sigma = pm.HalfNormal("sigma", sigma=1) - >>> mu = pm.Deterministic("mu", pm.math.dot(X_, beta)) - >>> pm.Normal("y_hat", mu=mu, sigma=sigma, observed=y_) + ... def build_model(self, X, y): + ... with self: + ... X_ = pm.MutableData(name="X", value=X) + ... y_ = pm.MutableData(name="y", value=y) + ... beta = pm.Normal("beta", mu=0, sigma=1, shape=X_.shape[1]) + ... sigma = pm.HalfNormal("sigma", sigma=1) + ... mu = pm.Deterministic("mu", pm.math.dot(X_, beta)) + ... pm.Normal("y_hat", mu=mu, sigma=sigma, observed=y_) """ raise NotImplementedError("This method must be implemented by a subclass") @@ -83,37 +85,22 @@ def fit(self, X, y, coords: Optional[Dict[str, Any]] = None) -> None: >>> import pymc as pm >>> from causalpy.pymc_models import ModelBuilder >>> class MyToyModel(ModelBuilder): - ... def build_model(self, X, y, coords): - ... with self: - ... X_ = pm.MutableData(name="X", value=X) - ... y_ = pm.MutableData(name="y", value=y) - ... beta = pm.Normal("beta", mu=0, sigma=1, shape=X_.shape[1]) - ... sigma = pm.HalfNormal("sigma", sigma=1) - ... mu = pm.Deterministic("mu", pm.math.dot(X_, beta)) - ... pm.Normal("y_hat", mu=mu, sigma=sigma, observed=y_) + ... def build_model(self, X, y, coords): + ... with self: + ... X_ = pm.MutableData(name="X", value=X) + ... y_ = pm.MutableData(name="y", value=y) + ... beta = pm.Normal("beta", mu=0, sigma=1, shape=X_.shape[1]) + ... sigma = pm.HalfNormal("sigma", sigma=1) + ... mu = pm.Deterministic("mu", pm.math.dot(X_, beta)) + ... pm.Normal("y_hat", mu=mu, sigma=sigma, observed=y_) >>> rng = np.random.default_rng(seed=42) >>> X = rng.normal(loc=0, scale=1, size=(20, 2)) >>> y = rng.normal(loc=0, scale=1, size=(20,)) - >>> model = MyToyModel(sample_kwargs={"chains": 2, "draws": 2}) - >>> model.fit(X, y) - Only 2 samples in chain. - Auto-assigning NUTS sampler... - Initializing NUTS using jitter+adapt_diag... - Multiprocess sampling (2 chains in 4 jobs) - NUTS: [beta, sigma] - Sampling 2 chains for 1_000 tune and 2 draw iterations (2_000 + 4 draws total) - took 0 seconds.gences] - The number of samples is too small to check convergence reliably. - Sampling: [beta, sigma, y_hat] - Sampling: [y_hat] - Inference data with groups: - > posterior - > posterior_predictive - > sample_stats - > prior - > prior_predictive - > observed_data - > constant_data + >>> model = MyToyModel( + ... sample_kwargs={"chains": 2, "draws": 2, "progressbar": False} + ... ) + >>> model.fit(X, y) # doctest: +ELLIPSIS + Inference ... """ self.build_model(X, y, coords) with self.model: @@ -133,16 +120,30 @@ def predict(self, X): Example ------- - Assumes `model` has been initialized and .fit() has been run, - see ModelBuilder().fit() for example. - + >>> import causalpy as cp + >>> import numpy as np + >>> import pymc as pm + >>> from causalpy.pymc_models import ModelBuilder + >>> class MyToyModel(ModelBuilder): + ... def build_model(self, X, y, coords): + ... with self: + ... X_ = pm.MutableData(name="X", value=X) + ... y_ = pm.MutableData(name="y", value=y) + ... beta = pm.Normal("beta", mu=0, sigma=1, shape=X_.shape[1]) + ... sigma = pm.HalfNormal("sigma", sigma=1) + ... mu = pm.Deterministic("mu", pm.math.dot(X_, beta)) + ... pm.Normal("y_hat", mu=mu, sigma=sigma, observed=y_) + >>> rng = np.random.default_rng(seed=42) + >>> X = rng.normal(loc=0, scale=1, size=(20, 2)) + >>> y = rng.normal(loc=0, scale=1, size=(20,)) + >>> model = MyToyModel( + ... sample_kwargs={"chains": 2, "draws": 2, "progressbar": False} + ... ) + >>> model.fit(X, y) # doctest: +ELLIPSIS + Inference... >>> X_new = rng.normal(loc=0, scale=1, size=(20,2)) - >>> model.predict(X_new) - Sampling: [beta, y_hat] - Inference data with groups: - > posterior_predictive - > observed_data - > constant_data + >>> model.predict(X_new) # doctest: +ELLIPSIS + Inference... """ self._data_setter(X) @@ -162,9 +163,28 @@ def score(self, X, y) -> pd.Series: Example -------- - Assuming `model` has been fit - - >>> model.score(X, y) # X, y are random data here + >>> import causalpy as cp + >>> import numpy as np + >>> import pymc as pm + >>> from causalpy.pymc_models import ModelBuilder + >>> class MyToyModel(ModelBuilder): + ... def build_model(self, X, y, coords): + ... with self: + ... X_ = pm.MutableData(name="X", value=X) + ... y_ = pm.MutableData(name="y", value=y) + ... beta = pm.Normal("beta", mu=0, sigma=1, shape=X_.shape[1]) + ... sigma = pm.HalfNormal("sigma", sigma=1) + ... mu = pm.Deterministic("mu", pm.math.dot(X_, beta)) + ... pm.Normal("y_hat", mu=mu, sigma=sigma, observed=y_) + >>> rng = np.random.default_rng(seed=42) + >>> X = rng.normal(loc=0, scale=1, size=(20, 2)) + >>> y = rng.normal(loc=0, scale=1, size=(20,)) + >>> model = MyToyModel( + ... sample_kwargs={"chains": 2, "draws": 2, "progressbar": False} + ... ) + >>> model.fit(X, y) # doctest: +ELLIPSIS + Inference... + >>> model.score(X, y) Sampling: [y_hat] r2 0.352251 r2_std 0.051624 @@ -196,27 +216,14 @@ class WeightedSumFitter(ModelBuilder): Example -------- + >>> import causalpy as cp + >>> import numpy as np + >>> from causalpy.pymc_models import WeightedSumFitter >>> sc = cp.load_data("sc") >>> X = sc[['a', 'b', 'c', 'd', 'e', 'f', 'g']] >>> y = np.asarray(sc['actual']).reshape((sc.shape[0], 1)) - >>> wsf = WeightedSumFitter() - >>> wsf.fit(X,y) - Auto-assigning NUTS sampler... - Initializing NUTS using jitter+adapt_diag... - Multiprocess sampling (4 chains in 4 jobs) - NUTS: [beta, sigma] - Sampling 4 chains for 1_000 tune and 1_000 draw iterations - (4_000 + 4_000 draws total) took 3 seconds. - Sampling: [beta, sigma, y_hat] - Sampling: [y_hat] - Inference data with groups: - > posterior - > posterior_predictive - > sample_stats - > prior - > prior_predictive - > observed_data - > constant_data + >>> wsf = WeightedSumFitter(sample_kwargs={"progressbar": False}) + >>> _ = wsf.fit(X,y) """ def build_model(self, X, y, coords): @@ -261,31 +268,19 @@ class LinearRegression(ModelBuilder): Example -------- + >>> import causalpy as cp + >>> import numpy as np + >>> from causalpy.pymc_models import LinearRegression >>> rd = cp.load_data("rd") >>> X = rd[["x", "treated"]] >>> y = np.asarray(rd["y"]).reshape((rd["y"].shape[0],1)) - >>> lr = LinearRegression() + >>> lr = LinearRegression(sample_kwargs={"progressbar": False}) >>> lr.fit(X, y, coords={ - 'coeffs': ['x', 'treated'], - 'obs_indx': np.arange(rd.shape[0]) - } - ) - Auto-assigning NUTS sampler... - Initializing NUTS using jitter+adapt_diag... - Multiprocess sampling (4 chains in 4 jobs) - NUTS: [beta, sigma] - Sampling 4 chains for 1_000 tune and 1_000 draw iterations ( - 4_000 + 4_000 draws total) took 1 seconds. - Sampling: [beta, sigma, y_hat] - Sampling: [y_hat] - Inference data with groups: - > posterior - > posterior_predictive - > sample_stats - > prior - > prior_predictive - > observed_data - > constant_data + ... 'coeffs': ['x', 'treated'], + ... 'obs_indx': np.arange(rd.shape[0]) + ... }, + ... ) # doctest: +ELLIPSIS + Inference... """ def build_model(self, X, y, coords): From 2220ed0243b21133da0045992a4247553d7be2d0 Mon Sep 17 00:00:00 2001 From: Jason Preszler Date: Thu, 31 Aug 2023 08:18:06 -0700 Subject: [PATCH 08/23] doctest good on pymc experiments except for summaries --- causalpy/pymc_experiments.py | 197 ++++++++++++++++------------------- causalpy/pymc_models.py | 6 +- 2 files changed, 92 insertions(+), 111 deletions(-) diff --git a/causalpy/pymc_experiments.py b/causalpy/pymc_experiments.py index c98095a2..f2c3002e 100644 --- a/causalpy/pymc_experiments.py +++ b/causalpy/pymc_experiments.py @@ -55,38 +55,22 @@ def idata(self): Example -------- - If `result` is the result of the Difference in Differences experiment example - - >>> result.idata - Inference data with groups: - > posterior - > posterior_predictive - > sample_stats - > prior - > prior_predictive - > observed_data - > constant_data - >>> result.idata.posterior + >>> import causalpy as cp + >>> df = cp.load_data("did") + >>> seed = 42 + >>> result = cp.pymc_experiments.DifferenceInDifferences( + ... df, + ... formula="y ~ 1 + group*post_treatment", + ... time_variable_name="t", + ... group_variable_name="group", + ... model=cp.pymc_models.LinearRegression( + ... sample_kwargs={"random_seed": seed, "progressbar": False}), + ... ) + >>> result.idata # doctest: +ELLIPSIS + Inference data... + >>> result.idata.posterior # doctest: +ELLIPSIS - Dimensions: (chain: 4, draw: 1000, coeffs: 4, obs_ind: 40) - Coordinates: - * chain (chain) int64 0 1 2 3 - * draw (draw) int64 0 1 2 3 4 5 6 7 8 ... 992 993 994 995 996 997 998 - 999 - * coeffs (coeffs) None: Example -------- - If `result` is from the Difference in Differences experiment example - + >>> import causalpy as cp + >>> df = cp.load_data("did") + >>> seed = 42 + >>> result = cp.pymc_experiments.DifferenceInDifferences( + ... df, + ... formula="y ~ 1 + group*post_treatment", + ... time_variable_name="t", + ... group_variable_name="group", + ... model=cp.pymc_models.LinearRegression( + ... sample_kwargs={"random_seed": seed, "progressbar": False}), + ... ) >>> result.print_coefficients() Model coefficients: Intercept 1.08, 94% HDI [1.03, 1.13] @@ -140,6 +133,7 @@ class PrePostFit(ExperimentalDesign): Example -------- + >>> import causalpy as cp >>> sc = cp.load_data("sc") >>> treatment_time = 70 >>> seed = 42 @@ -148,20 +142,13 @@ class PrePostFit(ExperimentalDesign): ... treatment_time, ... formula="actual ~ 0 + a + b + c + d + e + f + g", ... model=cp.pymc_models.WeightedSumFitter( - ... sample_kwargs={"target_accept": 0.95, "random_seed": seed} + ... sample_kwargs={ + ... "target_accept": 0.95, + ... "random_seed": seed, + ... "progressbar": False + ... } ... ), ... ) - Auto-assigning NUTS sampler... - Initializing NUTS using jitter+adapt_diag... - Multiprocess sampling (4 chains in 4 jobs) - NUTS: [beta, sigma] - Sampling 4 chains for 1_000 tune and 1_000 draw iterations - (4_000 + 4_000 draws total) took 11 seconds. - Sampling: [beta, sigma, y_hat] - Sampling: [y_hat] - Sampling: [y_hat] - Sampling: [y_hat] - Sampling: [y_hat] """ def __init__( @@ -249,8 +236,7 @@ def plot(self, counterfactual_label="Counterfactual", **kwargs): Example -------- - >>> result.plot() - + >>> result.plot() # doctest: +SKIP """ fig, ax = plt.subplots(3, 1, sharex=True, figsize=(7, 8)) @@ -356,6 +342,22 @@ def summary(self) -> None: Example --------- + >>> import causalpy as cp + >>> sc = cp.load_data("sc") + >>> treatment_time = 70 + >>> seed = 42 + >>> result = cp.pymc_experiments.PrePostFit( + ... sc, + ... treatment_time, + ... formula="actual ~ 0 + a + b + c + d + e + f + g", + ... model=cp.pymc_models.WeightedSumFitter( + ... sample_kwargs={ + ... "target_accept": 0.95, + ... "random_seed": seed, + ... "progressbar": False, + ... } + ... ), + ... ) >>> result.summary() ===============================Synthetic Control=============================== Formula: actual ~ 0 + a + b + c + d + e + f + g @@ -391,6 +393,7 @@ class InterruptedTimeSeries(PrePostFit): Example -------- + >>> import causalpy as cp >>> df = ( ... cp.load_data("its") ... .assign(date=lambda x: pd.to_datetime(x["date"])) @@ -402,19 +405,14 @@ class InterruptedTimeSeries(PrePostFit): ... df, ... treatment_time, ... formula="y ~ 1 + t + C(month)", - ... model=cp.pymc_models.LinearRegression(sample_kwargs={"random_seed": seed}), + ... model=cp.pymc_models.LinearRegression( + ... sample_kwargs={ + ... "target_accept": 0.95, + ... "random_seed": seed, + ... "progressbar": False, + ... } + ... ) ... ) - Auto-assigning NUTS sampler... - Initializing NUTS using jitter+adapt_diag... - Multiprocess sampling (4 chains in 4 jobs) - NUTS: [beta, sigma] - Sampling 4 chains for 1_000 tune and 1_000 draw iterations - (4_000 + 4_000 draws total) took 3 seconds. - Sampling: [beta, sigma, y_hat] - Sampling: [y_hat] - Sampling: [y_hat] - Sampling: [y_hat] - Sampling: [y_hat] """ expt_type = "Interrupted Time Series" @@ -434,6 +432,7 @@ class SyntheticControl(PrePostFit): Example -------- + >>> import causalpy as cp >>> df = cp.load_data("sc") >>> treatment_time = 70 >>> seed = 42 @@ -442,20 +441,13 @@ class SyntheticControl(PrePostFit): ... treatment_time, ... formula="actual ~ 0 + a + b + c + d + e + f + g", ... model=cp.pymc_models.WeightedSumFitter( - ... sample_kwargs={"target_accept": 0.95, "random_seed": seed} + ... sample_kwargs={ + ... "target_accept": 0.95, + ... "random_seed": seed, + ... "progressbar": False, + ... } ... ), ... ) - Auto-assigning NUTS sampler... - Initializing NUTS using jitter+adapt_diag... - Multiprocess sampling (4 chains in 4 jobs) - NUTS: [beta, sigma] - Sampling 4 chains for 1_000 tune and 1_000 draw iterations - (4_000 + 4_000 draws total) took 11 seconds. - Sampling: [beta, sigma, y_hat] - Sampling: [y_hat] - Sampling: [y_hat] - Sampling: [y_hat] - Sampling: [y_hat] """ expt_type = "Synthetic Control" @@ -492,6 +484,7 @@ class DifferenceInDifferences(ExperimentalDesign): Example -------- + >>> import causalpy as cp >>> df = cp.load_data("did") >>> seed = 42 >>> result = cp.pymc_experiments.DifferenceInDifferences( @@ -499,19 +492,14 @@ class DifferenceInDifferences(ExperimentalDesign): ... formula="y ~ 1 + group*post_treatment", ... time_variable_name="t", ... group_variable_name="group", - ... model=cp.pymc_models.LinearRegression(sample_kwargs={"random_seed": seed}), + ... model=cp.pymc_models.LinearRegression( + ... sample_kwargs={ + ... "target_accept": 0.95, + ... "random_seed": seed, + ... "progressbar": False, + ... } + ... ) ... ) - Auto-assigning NUTS sampler... - Initializing NUTS using jitter+adapt_diag... - Multiprocess sampling (4 chains in 4 jobs) - NUTS: [beta, sigma] - Sampling 4 chains for 1_000 tune and 1_000 draw iterations - (4_000 + 4_000 draws total) took 1 seconds. - Sampling: [beta, sigma, y_hat] - Sampling: [y_hat] - Sampling: [y_hat] - Sampling: [y_hat] - Sampling: [y_hat] """ def __init__( @@ -637,7 +625,7 @@ def plot(self): -------- Assuming `result` is the result of a DiD experiment: - >>> result.plot() + >>> result.plot() # doctest: +SKIP """ fig, ax = plt.subplots() @@ -825,25 +813,21 @@ class RegressionDiscontinuity(ExperimentalDesign): Example -------- + >>> import causalpy as cp >>> df = cp.load_data("rd") >>> seed = 42 >>> result = cp.pymc_experiments.RegressionDiscontinuity( ... df, ... formula="y ~ 1 + x + treated + x:treated", - ... model=cp.pymc_models.LinearRegression(sample_kwargs={"random_seed": seed}), + ... model=cp.pymc_models.LinearRegression( + ... sample_kwargs={ + ... "target_accept": 0.95, + ... "random_seed": seed, + ... "progressbar": False, + ... }, + ... ), ... treatment_threshold=0.5, ... ) - Auto-assigning NUTS sampler... - Initializing NUTS using jitter+adapt_diag... - Multiprocess sampling (4 chains in 4 jobs) - NUTS: [beta, sigma] - Sampling 4 chains for 1_000 tune and 1_000 draw iterations - (4_000 + 4_000 draws total) took 2 seconds. - Sampling: [beta, sigma, y_hat] - Sampling: [y_hat] - Sampling: [y_hat] - Sampling: [y_hat] - Sampling: [y_hat] """ def __init__( @@ -959,7 +943,7 @@ def plot(self): Example -------- - >>> result.plot() + >>> result.plot() # doctest: +SKIP """ fig, ax = plt.subplots() # Plot raw data @@ -1054,6 +1038,7 @@ class PrePostNEGD(ExperimentalDesign): Example -------- + >>> import causalpy as cp >>> df = cp.load_data("anova1") >>> seed = 42 >>> result = cp.pymc_experiments.PrePostNEGD( @@ -1061,18 +1046,14 @@ class PrePostNEGD(ExperimentalDesign): ... formula="post ~ 1 + C(group) + pre", ... group_variable_name="group", ... pretreatment_variable_name="pre", - ... model=cp.pymc_models.LinearRegression(sample_kwargs={"random_seed": seed}), + ... model=cp.pymc_models.LinearRegression( + ... sample_kwargs={ + ... "target_accept": 0.95, + ... "random_seed": seed, + ... "progressbar": False, + ... } + ... ) ... ) - Auto-assigning NUTS sampler... - Initializing NUTS using jitter+adapt_diag... - Multiprocess sampling (4 chains in 4 jobs) - NUTS: [beta, sigma] - Sampling 4 chains for 1_000 tune and 1_000 draw iterations - (4_000 + 4_000 draws total) took 3 seconds. - Sampling: [beta, sigma, y_hat] - Sampling: [y_hat] - Sampling: [y_hat] - Sampling: [y_hat] """ def __init__( diff --git a/causalpy/pymc_models.py b/causalpy/pymc_models.py index 8d2b922b..d786fc37 100644 --- a/causalpy/pymc_models.py +++ b/causalpy/pymc_models.py @@ -180,14 +180,14 @@ def score(self, X, y) -> pd.Series: >>> X = rng.normal(loc=0, scale=1, size=(20, 2)) >>> y = rng.normal(loc=0, scale=1, size=(20,)) >>> model = MyToyModel( - ... sample_kwargs={"chains": 2, "draws": 2, "progressbar": False} + ... sample_kwargs={"chains": 2, "draws": 200, "progressbar": False} ... ) >>> model.fit(X, y) # doctest: +ELLIPSIS Inference... >>> model.score(X, y) Sampling: [y_hat] - r2 0.352251 - r2_std 0.051624 + r2 0.376489 + r2_std 0.081305 dtype: float64 """ yhat = self.predict(X) From 5d6c8eb0074fb3e371b56d8732d03a2fb7d5aab8 Mon Sep 17 00:00:00 2001 From: Jason Preszler Date: Fri, 1 Sep 2023 09:33:27 -0700 Subject: [PATCH 09/23] doctest clean and added to github actions --- .github/workflows/ci.yml | 4 ++ causalpy/data/simulate_data.py | 8 +++- causalpy/pymc_experiments.py | 82 +++++++++++++++++++++++++++------- causalpy/pymc_models.py | 26 +++++------ causalpy/skl_experiments.py | 45 +++++++++++++------ causalpy/skl_models.py | 15 ++----- pyproject.toml | 1 + 7 files changed, 124 insertions(+), 57 deletions(-) diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 8e8295ea..db618756 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -34,6 +34,10 @@ jobs: uses: actions/setup-python@v3 with: python-version: ${{ matrix.python-version }} + - name: Run doctests + run: | + pip install -e .[test] + pytest --doctest-modules causalpy/ - name: Run tests run: | pip install -e .[test] diff --git a/causalpy/data/simulate_data.py b/causalpy/data/simulate_data.py index 5dcdca42..e99f2eac 100644 --- a/causalpy/data/simulate_data.py +++ b/causalpy/data/simulate_data.py @@ -56,8 +56,9 @@ def generate_synthetic_control_data( Example -------- + >>> from causalpy.data.simulate_data import generate_synthetic_control_data >>> df, weightings_true = generate_synthetic_control_data( - ... treatment_time=treatment_time + ... treatment_time=70 ... ) """ @@ -196,6 +197,7 @@ def generate_did(): Example -------- + >>> from causalpy.data.simulate_data import generate_did >>> df = generate_did() """ # true parameters @@ -249,6 +251,7 @@ def generate_regression_discontinuity_data( Example -------- >>> import pathlib + >>> from causalpy.data.simulate_data import generate_regression_discontinuity_data >>> df = generate_regression_discontinuity_data(true_treatment_threshold=0.5) >>> df.to_csv(pathlib.Path.cwd() / 'regression_discontinuity.csv', index=False) """ @@ -278,9 +281,10 @@ def generate_ancova_data( Example -------- >>> import pathlib + >>> from causalpy.data.simulate_data import generate_ancova_data >>> df = generate_ancova_data( ... N=200, - ... pre_treatment_threshold=np.array([10, 12]), + ... pre_treatment_means=np.array([10, 12]), ... treatment_effect=2, ... sigma=1 ... ) diff --git a/causalpy/pymc_experiments.py b/causalpy/pymc_experiments.py index f2c3002e..9e49336d 100644 --- a/causalpy/pymc_experiments.py +++ b/causalpy/pymc_experiments.py @@ -359,7 +359,7 @@ def summary(self) -> None: ... ), ... ) >>> result.summary() - ===============================Synthetic Control=============================== + ==================================Pre-Post Fit================================== Formula: actual ~ 0 + a + b + c + d + e + f + g Model coefficients: a 0.33, 94% HDI [0.30, 0.38] @@ -757,7 +757,7 @@ def _plot_causal_impact_arrow(self, ax): def _causal_impact_summary_stat(self) -> str: """Computes the mean and 94% credible interval bounds for the causal impact.""" percentiles = self.causal_impact.quantile([0.03, 1 - 0.03]).values - ci = r"$CI_{94\%}$" + f"[{percentiles[0]:.2f}, {percentiles[1]:.2f}]" + ci = "$CI_{94%}$" + f"[{percentiles[0]:.2f}, {percentiles[1]:.2f}]" causal_impact = f"{self.causal_impact.mean():.2f}, " return f"Causal impact = {causal_impact + ci}" @@ -767,16 +767,31 @@ def summary(self) -> None: Example -------- - Assuming `result` is a DiD experiment - + >>> import causalpy as cp + >>> df = cp.load_data("did") + >>> seed = 42 + >>> result = cp.pymc_experiments.DifferenceInDifferences( + ... df, + ... formula="y ~ 1 + group*post_treatment", + ... time_variable_name="t", + ... group_variable_name="group", + ... model=cp.pymc_models.LinearRegression( + ... sample_kwargs={ + ... "target_accept": 0.95, + ... "random_seed": seed, + ... "progressbar": False, + ... } + ... ) + ... ) >>> result.summary() - ==========================Difference in Differences========================= + ===========================Difference in Differences============================ Formula: y ~ 1 + group*post_treatment + Results: Causal impact = 0.51, $CI_{94%}$[0.41, 0.61] Model coefficients: Intercept 1.08, 94% HDI [1.03, 1.13] - post_treatment[T.True] 0.98, 94% HDI [0.91, 1.06] + post_treatment[T.True] 0.98, 94% HDI [0.92, 1.05] group 0.16, 94% HDI [0.09, 0.23] group:post_treatment[T.True] 0.51, 94% HDI [0.41, 0.61] sigma 0.08, 94% HDI [0.07, 0.10] @@ -995,19 +1010,35 @@ def summary(self) -> None: Example -------- + >>> import causalpy as cp + >>> df = cp.load_data("rd") + >>> seed = 42 + >>> result = cp.pymc_experiments.RegressionDiscontinuity( + ... df, + ... formula="y ~ 1 + x + treated + x:treated", + ... model=cp.pymc_models.LinearRegression( + ... sample_kwargs={ + ... "target_accept": 0.95, + ... "random_seed": seed, + ... "progressbar": False, + ... }, + ... ), + ... treatment_threshold=0.5, + ... ) >>> result.summary() - ============================Regression Discontinuity========================== + ============================Regression Discontinuity============================ Formula: y ~ 1 + x + treated + x:treated Running variable: x Threshold on running variable: 0.5 + Results: - Discontinuity at threshold = 0.92 + Discontinuity at threshold = 0.91 Model coefficients: - Intercept 0.09, 94% HDI [0.00, 0.17] - treated[T.True] 2.48, 94% HDI [1.66, 3.27] + Intercept 0.09, 94% HDI [-0.00, 0.17] + treated[T.True] 2.45, 94% HDI [1.66, 3.28] x 1.32, 94% HDI [1.14, 1.50] - x:treated[T.True] -3.12, 94% HDI [-4.17, -2.05] - sigma 0.35, 94% HDI [0.31, 0.41] + x:treated[T.True] -3.08, 94% HDI [-4.17, -2.05] + sigma 0.36, 94% HDI [0.31, 0.41] """ print(f"{self.expt_type:=^80}") @@ -1182,7 +1213,7 @@ def plot(self): def _causal_impact_summary_stat(self) -> str: """Computes the mean and 94% credible interval bounds for the causal impact.""" percentiles = self.causal_impact.quantile([0.03, 1 - 0.03]).values - ci = r"$CI_{94\%}$" + f"[{percentiles[0]:.2f}, {percentiles[1]:.2f}]" + ci = r"$CI_{94%}$" + f"[{percentiles[0]:.2f}, {percentiles[1]:.2f}]" causal_impact = f"{self.causal_impact.mean():.2f}, " return f"Causal impact = {causal_impact + ci}" @@ -1192,14 +1223,31 @@ def summary(self) -> None: Example -------- + >>> import causalpy as cp + >>> df = cp.load_data("anova1") + >>> seed = 42 + >>> result = cp.pymc_experiments.PrePostNEGD( + ... df, + ... formula="post ~ 1 + C(group) + pre", + ... group_variable_name="group", + ... pretreatment_variable_name="pre", + ... model=cp.pymc_models.LinearRegression( + ... sample_kwargs={ + ... "target_accept": 0.95, + ... "random_seed": seed, + ... "progressbar": False, + ... } + ... ) + ... ) >>> result.summary() - =================Pretest/posttest Nonequivalent Group Design================ + ==================Pretest/posttest Nonequivalent Group Design=================== Formula: post ~ 1 + C(group) + pre + Results: - Causal impact = 1.89, $CI_{94%}$[1.70, 2.07] + Causal impact = 1.88, $CI_{94%}$[1.69, 2.07] Model coefficients: - Intercept -0.46, 94% HDI [-1.17, 0.22] - C(group)[T.1] 1.89, 94% HDI [1.70, 2.07] + Intercept -0.47, 94% HDI [-1.16, 0.24] + C(group)[T.1] 1.88, 94% HDI [1.69, 2.07] pre 1.05, 94% HDI [0.98, 1.12] sigma 0.51, 94% HDI [0.46, 0.56] diff --git a/causalpy/pymc_models.py b/causalpy/pymc_models.py index d786fc37..e6b57143 100644 --- a/causalpy/pymc_models.py +++ b/causalpy/pymc_models.py @@ -99,7 +99,7 @@ def fit(self, X, y, coords: Optional[Dict[str, Any]] = None) -> None: >>> model = MyToyModel( ... sample_kwargs={"chains": 2, "draws": 2, "progressbar": False} ... ) - >>> model.fit(X, y) # doctest: +ELLIPSIS + >>> model.fit(X, y) Inference ... """ self.build_model(X, y, coords) @@ -139,10 +139,10 @@ def predict(self, X): >>> model = MyToyModel( ... sample_kwargs={"chains": 2, "draws": 2, "progressbar": False} ... ) - >>> model.fit(X, y) # doctest: +ELLIPSIS + >>> model.fit(X, y) Inference... >>> X_new = rng.normal(loc=0, scale=1, size=(20,2)) - >>> model.predict(X_new) # doctest: +ELLIPSIS + >>> model.predict(X_new) Inference... """ @@ -177,17 +177,16 @@ def score(self, X, y) -> pd.Series: ... mu = pm.Deterministic("mu", pm.math.dot(X_, beta)) ... pm.Normal("y_hat", mu=mu, sigma=sigma, observed=y_) >>> rng = np.random.default_rng(seed=42) - >>> X = rng.normal(loc=0, scale=1, size=(20, 2)) - >>> y = rng.normal(loc=0, scale=1, size=(20,)) + >>> X = rng.normal(loc=0, scale=1, size=(200, 2)) + >>> y = rng.normal(loc=0, scale=1, size=(200,)) >>> model = MyToyModel( - ... sample_kwargs={"chains": 2, "draws": 200, "progressbar": False} + ... sample_kwargs={"chains": 2, "draws": 2000, "progressbar": False} ... ) - >>> model.fit(X, y) # doctest: +ELLIPSIS + >>> model.fit(X, y) Inference... - >>> model.score(X, y) - Sampling: [y_hat] - r2 0.376489 - r2_std 0.081305 + >>> round(model.score(X, y),2) # using round() to simplify doctest + r2 0.34 + r2_std 0.02 dtype: float64 """ yhat = self.predict(X) @@ -223,7 +222,8 @@ class WeightedSumFitter(ModelBuilder): >>> X = sc[['a', 'b', 'c', 'd', 'e', 'f', 'g']] >>> y = np.asarray(sc['actual']).reshape((sc.shape[0], 1)) >>> wsf = WeightedSumFitter(sample_kwargs={"progressbar": False}) - >>> _ = wsf.fit(X,y) + >>> wsf.fit(X,y) + Inference ... """ def build_model(self, X, y, coords): @@ -279,7 +279,7 @@ class LinearRegression(ModelBuilder): ... 'coeffs': ['x', 'treated'], ... 'obs_indx': np.arange(rd.shape[0]) ... }, - ... ) # doctest: +ELLIPSIS + ... ) Inference... """ diff --git a/causalpy/skl_experiments.py b/causalpy/skl_experiments.py index 68854758..f088f6b1 100644 --- a/causalpy/skl_experiments.py +++ b/causalpy/skl_experiments.py @@ -59,7 +59,6 @@ class PrePostFit(ExperimentalDesign): ... formula="actual ~ 0 + a + b + c + d + e + f + g", ... model = cp.skl_models.WeightedProportion() ... ) - """ def __init__( @@ -181,10 +180,18 @@ def get_coeffs(self): Example -------- + >>> from sklearn.linear_model import LinearRegression + >>> import causalpy as cp + >>> df = cp.load_data("sc") + >>> treatment_time = 70 + >>> result = cp.skl_experiments.PrePostFit( + ... df, + ... treatment_time, + ... formula="actual ~ 0 + a + b + c + d + e + f + g", + ... model = cp.skl_models.WeightedProportion() + ... ) >>> result.get_coeffs() - array([3.97370896e-01, 1.53881980e-01, 4.48747123e-01, 1.04639857e-16, - 0.00000000e+00, 0.00000000e+00, 2.92931287e-16]) - + array(...) """ return np.squeeze(self.model.coef_) @@ -262,7 +269,6 @@ class SyntheticControl(PrePostFit): ... formula="actual ~ 0 + a + b + c + d + e + f + g", ... model = cp.skl_models.WeightedProportion() ... ) - """ def plot(self, plot_predictors=False, **kwargs): @@ -293,13 +299,15 @@ class DifferenceInDifferences(ExperimentalDesign): :param group_variable_name: Name of the data column for the group variable :param model: - A PyMC model for difference in differences + An skl model for difference in differences Example -------- + >>> import causalpy as cp + >>> from sklearn.linear_model import LinearRegression >>> df = cp.load_data("did") >>> result = cp.skl_experiments.DifferenceInDifferences( - ... data, + ... df, ... formula="y ~ 1 + group*post_treatment", ... time_variable_name="t", ... group_variable_name="group", @@ -307,7 +315,6 @@ class DifferenceInDifferences(ExperimentalDesign): ... untreated=0, ... model=LinearRegression(), ... ) - """ def __init__( @@ -497,6 +504,8 @@ class RegressionDiscontinuity(ExperimentalDesign): Example -------- + >>> import causalpy as cp + >>> from sklearn.linear_model import LinearRegression >>> data = cp.load_data("rd") >>> result = cp.skl_experiments.RegressionDiscontinuity( ... data, @@ -504,7 +513,6 @@ class RegressionDiscontinuity(ExperimentalDesign): ... model=LinearRegression(), ... treatment_threshold=0.5, ... ) - """ def __init__( @@ -640,18 +648,27 @@ def summary(self): Example -------- - >>> result.summary() + >>> import causalpy as cp + >>> from sklearn.linear_model import LinearRegression + >>> data = cp.load_data("rd") + >>> result = cp.skl_experiments.RegressionDiscontinuity( + ... data, + ... formula="y ~ 1 + x + treated", + ... model=LinearRegression(), + ... treatment_threshold=0.5, + ... ) + >>> result.summary() # doctest: +NORMALIZE_WHITESPACE Difference in Differences experiment Formula: y ~ 1 + x + treated Running variable: x Threshold on running variable: 0.5 + Results: Discontinuity at threshold = 0.19 Model coefficients: - Intercept 0.0 - treated[T.True] 0.19034196317793994 - x 1.229600855360073 - + Intercept 0.0 + treated[T.True] 0.19034196317793994 + x 1.229600855360073 """ print("Difference in Differences experiment") print(f"Formula: {self.formula}") diff --git a/causalpy/skl_models.py b/causalpy/skl_models.py index c7acc6c3..f9adcdd3 100644 --- a/causalpy/skl_models.py +++ b/causalpy/skl_models.py @@ -24,26 +24,19 @@ class WeightedProportion(LinearModel, RegressorMixin): Example -------- + >>> import numpy as np + >>> from causalpy.skl_models import WeightedProportion >>> rng = np.random.default_rng(seed=42) >>> X = rng.normal(loc=0, scale=1, size=(20,2)) >>> y = rng.normal(loc=0, scale=1, size=(20,)) + >>> wp = WeightedProportion() >>> wp.fit(X, y) WeightedProportion() >>> wp.coef_ array([[0.36719946, 0.63280054]]) >>> X_new = rng.normal(loc=0, scale=1, size=(10,2)) >>> wp.predict(X_new) - array([[-0.8298643 ], - [ 0.43072465], - [ 0.76319257], - [-0.42062812], - [ 0.1939908 ], - [-1.18557609], - [-0.0230188 ], - [ 0.48923816], - [-0.05656294], - [ 0.0339618 ]]) - + array(...) """ def loss(self, W, X, y): diff --git a/pyproject.toml b/pyproject.toml index 687e2b82..2a7a6c49 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -96,6 +96,7 @@ addopts = [ "--strict-config", "--cov=causalpy", "--cov-report=term-missing", + "--doctest-modules", ] testpaths = "causalpy/tests" markers = [ From e392335d0b9bfbc130b9655c03fc96ca92debae3 Mon Sep 17 00:00:00 2001 From: Jason Preszler Date: Mon, 4 Sep 2023 16:11:26 -0700 Subject: [PATCH 10/23] Added IV examples after rebasing --- causalpy/pymc_experiments.py | 45 +++++++++++++++++++++++ causalpy/pymc_models.py | 37 ++++++++++++++++++- causalpy/tests/test_input_validation.py | 3 ++ causalpy/version.py | 1 + docs/source/_static/interrogate_badge.svg | 6 +-- 5 files changed, 88 insertions(+), 4 deletions(-) diff --git a/causalpy/pymc_experiments.py b/causalpy/pymc_experiments.py index 9e49336d..3f17c910 100644 --- a/causalpy/pymc_experiments.py +++ b/causalpy/pymc_experiments.py @@ -1301,6 +1301,40 @@ class InstrumentalVariable(ExperimentalDesign): "lkj_sd": 2, } + Example + -------- + >>> import pandas as pd + >>> import causalpy as cp + >>> from causalpy.pymc_experiments import InstrumentalVariable + >>> from causalpy.pymc_models import InstrumentalVariableRegression + >>> import numpy as np + >>> N = 100 + >>> e1 = np.random.normal(0, 3, N) + >>> e2 = np.random.normal(0, 1, N) + >>> Z = np.random.uniform(0, 1, N) + >>> ## Ensure the endogeneity of the the treatment variable + >>> X = -1 + 4 * Z + e2 + 2 * e1 + >>> y = 2 + 3 * X + 3 * e1 + >>> test_data = pd.DataFrame({"y": y, "X": X, "Z": Z}) + >>> sample_kwargs = { + ... "tune": 10, + ... "draws": 20, + ... "chains": 4, + ... "cores": 4, + ... "target_accept": 0.95, + ... "progressbar": False, + ... } + >>> instruments_formula = "X ~ 1 + Z" + >>> formula = "y ~ 1 + X" + >>> instruments_data = test_data[["X", "Z"]] + >>> data = test_data[["y", "X"]] + >>> iv = InstrumentalVariable( + ... instruments_data=instruments_data, + ... data=data, + ... instruments_formula=instruments_formula, + ... formula=formula, + ... model=InstrumentalVariableRegression(sample_kwargs=sample_kwargs), + ... ) """ def __init__( @@ -1355,6 +1389,12 @@ def __init__( ) def get_2SLS_fit(self): + """ + Two Stage Least Squares Fit + + This function is called by the experiment, results are used for + priors if none are provided. + """ first_stage_reg = sk_lin_reg().fit(self.Z, self.t) fitted_Z_values = first_stage_reg.predict(self.Z) X2 = self.data.copy(deep=True) @@ -1371,6 +1411,11 @@ def get_2SLS_fit(self): self.second_stage_reg = second_stage_reg def get_naive_OLS_fit(self): + """ + Naive Ordinary Least Squares + + This function is called by the experiment. + """ ols_reg = sk_lin_reg().fit(self.X, self.y) beta_params = list(ols_reg.coef_[0][1:]) beta_params.insert(0, ols_reg.intercept_[0]) diff --git a/causalpy/pymc_models.py b/causalpy/pymc_models.py index e6b57143..b179be63 100644 --- a/causalpy/pymc_models.py +++ b/causalpy/pymc_models.py @@ -303,7 +303,42 @@ def build_model(self, X, y, coords): class InstrumentalVariableRegression(ModelBuilder): - """Custom PyMC model for instrumental linear regression""" + """Custom PyMC model for instrumental linear regression + + Example + -------- + >>> import causalpy as cp + >>> import numpy as np + >>> from causalpy.pymc_models import InstrumentalVariableRegression + >>> N = 10 + >>> e1 = np.random.normal(0, 3, N) + >>> e2 = np.random.normal(0, 1, N) + >>> Z = np.random.uniform(0, 1, N) + >>> ## Ensure the endogeneity of the the treatment variable + >>> X = -1 + 4 * Z + e2 + 2 * e1 + >>> y = 2 + 3 * X + 3 * e1 + >>> t = X.reshape(10,1) + >>> y = y.reshape(10,1) + >>> Z = np.asarray([[1, Z[i]] for i in range(0,10)]) + >>> X = np.asarray([[1, X[i]] for i in range(0,10)]) + >>> COORDS = {'instruments': ['Intercept', 'Z'], 'covariates': ['Intercept', 'X']} + >>> sample_kwargs = { + ... "tune": 5, + ... "draws": 10, + ... "chains": 2, + ... "cores": 2, + ... "target_accept": 0.95, + ... "progressbar": False, + ... } + >>> iv_reg = InstrumentalVariableRegression(sample_kwargs=sample_kwargs) + >>> iv_reg.fit(X, Z,y, t, COORDS, { + ... "mus": [[-2,4], [0.5, 3]], + ... "sigmas": [1, 1], + ... "eta": 2, + ... "lkj_sd": 2, + ... }) + Inference data... + """ def build_model(self, X, Z, y, t, coords, priors): """Specify model with treatment regression and focal regression data and priors diff --git a/causalpy/tests/test_input_validation.py b/causalpy/tests/test_input_validation.py index 04461fb5..9442c8bc 100644 --- a/causalpy/tests/test_input_validation.py +++ b/causalpy/tests/test_input_validation.py @@ -1,3 +1,5 @@ +"""Input validation tests""" + import pandas as pd import pytest @@ -205,6 +207,7 @@ def test_rd_validation_treated_is_dummy(): def test_iv_treatment_var_is_present(): + """Test the treatment variable is present for Instrumental Variable experiment""" data = pd.DataFrame({"x": [1, 2, 3], "y": [2, 4, 5]}) instruments_formula = "risk ~ 1 + logmort0" formula = "loggdp ~ 1 + risk" diff --git a/causalpy/version.py b/causalpy/version.py index 3dc1f76b..0700568d 100644 --- a/causalpy/version.py +++ b/causalpy/version.py @@ -1 +1,2 @@ +"""CausalPy Version""" __version__ = "0.1.0" diff --git a/docs/source/_static/interrogate_badge.svg b/docs/source/_static/interrogate_badge.svg index 377961aa..3cd04abb 100644 --- a/docs/source/_static/interrogate_badge.svg +++ b/docs/source/_static/interrogate_badge.svg @@ -1,5 +1,5 @@ - interrogate: 97.1% + interrogate: 97.2% @@ -12,8 +12,8 @@ interrogate interrogate - 97.1% - 97.1% + 97.2% + 97.2% From 76d970a95e714220818cd7aaa4652a3049f74a85 Mon Sep 17 00:00:00 2001 From: Jason Preszler Date: Tue, 5 Sep 2023 08:02:49 -0700 Subject: [PATCH 11/23] add statsmodels to dependencies --- pyproject.toml | 1 + 1 file changed, 1 insertion(+) diff --git a/pyproject.toml b/pyproject.toml index 2a7a6c49..118fd599 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -38,6 +38,7 @@ dependencies = [ "scikit-learn>=1", "scipy", "seaborn>=0.11.2", + "statsmodels", "xarray>=v2022.11.0", ] From b5e6dc675adb26aaaaebcf79490bce0ded91f1c3 Mon Sep 17 00:00:00 2001 From: Jason Preszler Date: Tue, 5 Sep 2023 13:29:11 -0700 Subject: [PATCH 12/23] increase draws and decrease precision on summaries --- causalpy/pymc_experiments.py | 58 ++++++++++++++++++++---------------- causalpy/skl_experiments.py | 6 ++-- 2 files changed, 36 insertions(+), 28 deletions(-) diff --git a/causalpy/pymc_experiments.py b/causalpy/pymc_experiments.py index 3f17c910..ad4869ab 100644 --- a/causalpy/pymc_experiments.py +++ b/causalpy/pymc_experiments.py @@ -90,15 +90,19 @@ def print_coefficients(self) -> None: ... time_variable_name="t", ... group_variable_name="group", ... model=cp.pymc_models.LinearRegression( - ... sample_kwargs={"random_seed": seed, "progressbar": False}), + ... sample_kwargs={ + ... "draws": 2000, + ... "random_seed": seed, + ... "progressbar": False + ... }), ... ) - >>> result.print_coefficients() + >>> result.print_coefficients() # doctest: +NUMBER Model coefficients: - Intercept 1.08, 94% HDI [1.03, 1.13] - post_treatment[T.True] 0.98, 94% HDI [0.91, 1.06] - group 0.16, 94% HDI [0.09, 0.23] - group:post_treatment[T.True] 0.51, 94% HDI [0.41, 0.61] - sigma 0.08, 94% HDI [0.07, 0.10] + Intercept 1.0, 94% HDI [1.0, 1.1] + post_treatment[T.True] 0.9, 94% HDI [0.9, 1.0] + group 0.1, 94% HDI [0.0, 0.2] + group:post_treatment[T.True] 0.5, 94% HDI [0.4, 0.6] + sigma 0.0, 94% HDI [0.0, 0.1] """ print("Model coefficients:") coeffs = az.extract(self.idata.posterior, var_names="beta") @@ -352,22 +356,23 @@ def summary(self) -> None: ... formula="actual ~ 0 + a + b + c + d + e + f + g", ... model=cp.pymc_models.WeightedSumFitter( ... sample_kwargs={ + ... "draws": 2000, ... "target_accept": 0.95, ... "random_seed": seed, ... "progressbar": False, ... } ... ), ... ) - >>> result.summary() + >>> result.summary() # doctest: +NUMBER ==================================Pre-Post Fit================================== Formula: actual ~ 0 + a + b + c + d + e + f + g Model coefficients: - a 0.33, 94% HDI [0.30, 0.38] + a 0.34, 94% HDI [0.30, 0.38] b 0.05, 94% HDI [0.01, 0.09] c 0.31, 94% HDI [0.26, 0.35] d 0.06, 94% HDI [0.01, 0.10] e 0.02, 94% HDI [0.00, 0.06] - f 0.20, 94% HDI [0.12, 0.26] + f 0.19, 94% HDI [0.11, 0.26] g 0.04, 94% HDI [0.00, 0.08] sigma 0.26, 94% HDI [0.22, 0.30] """ @@ -777,6 +782,7 @@ def summary(self) -> None: ... group_variable_name="group", ... model=cp.pymc_models.LinearRegression( ... sample_kwargs={ + ... "draws": 2000, ... "target_accept": 0.95, ... "random_seed": seed, ... "progressbar": False, @@ -788,12 +794,12 @@ def summary(self) -> None: Formula: y ~ 1 + group*post_treatment Results: - Causal impact = 0.51, $CI_{94%}$[0.41, 0.61] + Causal impact = 0.51, $CI_{94%}$[0.41, 0.60] Model coefficients: - Intercept 1.08, 94% HDI [1.03, 1.13] - post_treatment[T.True] 0.98, 94% HDI [0.92, 1.05] + Intercept 1.08, 94% HDI [1.03, 1.12] + post_treatment[T.True] 0.99, 94% HDI [0.92, 1.05] group 0.16, 94% HDI [0.09, 0.23] - group:post_treatment[T.True] 0.51, 94% HDI [0.41, 0.61] + group:post_treatment[T.True] 0.51, 94% HDI [0.41, 0.60] sigma 0.08, 94% HDI [0.07, 0.10] """ @@ -1018,6 +1024,7 @@ def summary(self) -> None: ... formula="y ~ 1 + x + treated + x:treated", ... model=cp.pymc_models.LinearRegression( ... sample_kwargs={ + ... "draws": 2000, ... "target_accept": 0.95, ... "random_seed": seed, ... "progressbar": False, @@ -1035,9 +1042,9 @@ def summary(self) -> None: Discontinuity at threshold = 0.91 Model coefficients: Intercept 0.09, 94% HDI [-0.00, 0.17] - treated[T.True] 2.45, 94% HDI [1.66, 3.28] + treated[T.True] 2.45, 94% HDI [1.64, 3.28] x 1.32, 94% HDI [1.14, 1.50] - x:treated[T.True] -3.08, 94% HDI [-4.17, -2.05] + x:treated[T.True] -3.09, 94% HDI [-4.16, -2.03] sigma 0.36, 94% HDI [0.31, 0.41] """ @@ -1233,23 +1240,24 @@ def summary(self) -> None: ... pretreatment_variable_name="pre", ... model=cp.pymc_models.LinearRegression( ... sample_kwargs={ + ... "draws": 2000, ... "target_accept": 0.95, ... "random_seed": seed, ... "progressbar": False, ... } ... ) ... ) - >>> result.summary() + >>> result.summary() # doctest: +NUMBER ==================Pretest/posttest Nonequivalent Group Design=================== Formula: post ~ 1 + C(group) + pre Results: - Causal impact = 1.88, $CI_{94%}$[1.69, 2.07] + Causal impact = 1.8, $CI_{94%}$[1.6, 2.0] Model coefficients: - Intercept -0.47, 94% HDI [-1.16, 0.24] - C(group)[T.1] 1.88, 94% HDI [1.69, 2.07] - pre 1.05, 94% HDI [0.98, 1.12] - sigma 0.51, 94% HDI [0.46, 0.56] + Intercept -0.4, 94% HDI [-1.2, 0.2] + C(group)[T.1] 1.8, 94% HDI [1.6, 2.0] + pre 1.0, 94% HDI [0.9, 1.1] + sigma 0.5, 94% HDI [0.4, 0.5] """ @@ -1317,9 +1325,9 @@ class InstrumentalVariable(ExperimentalDesign): >>> y = 2 + 3 * X + 3 * e1 >>> test_data = pd.DataFrame({"y": y, "X": X, "Z": Z}) >>> sample_kwargs = { - ... "tune": 10, - ... "draws": 20, - ... "chains": 4, + ... "tune": 1, + ... "draws": 5, + ... "chains": 1, ... "cores": 4, ... "target_accept": 0.95, ... "progressbar": False, diff --git a/causalpy/skl_experiments.py b/causalpy/skl_experiments.py index f088f6b1..8be30d4e 100644 --- a/causalpy/skl_experiments.py +++ b/causalpy/skl_experiments.py @@ -657,7 +657,7 @@ def summary(self): ... model=LinearRegression(), ... treatment_threshold=0.5, ... ) - >>> result.summary() # doctest: +NORMALIZE_WHITESPACE + >>> result.summary() # doctest: +NORMALIZE_WHITESPACE,+NUMBER Difference in Differences experiment Formula: y ~ 1 + x + treated Running variable: x @@ -667,8 +667,8 @@ def summary(self): Discontinuity at threshold = 0.19 Model coefficients: Intercept 0.0 - treated[T.True] 0.19034196317793994 - x 1.229600855360073 + treated[T.True] 0.19 + x 1.23 """ print("Difference in Differences experiment") print(f"Formula: {self.formula}") From 235972196550861d6105a94c663a891322236d65 Mon Sep 17 00:00:00 2001 From: Jason Preszler Date: Wed, 6 Sep 2023 10:43:38 -0700 Subject: [PATCH 13/23] more precision reduction --- causalpy/pymc_experiments.py | 26 +++++++++++++------------- 1 file changed, 13 insertions(+), 13 deletions(-) diff --git a/causalpy/pymc_experiments.py b/causalpy/pymc_experiments.py index ad4869ab..81cd6db7 100644 --- a/causalpy/pymc_experiments.py +++ b/causalpy/pymc_experiments.py @@ -371,9 +371,9 @@ def summary(self) -> None: b 0.05, 94% HDI [0.01, 0.09] c 0.31, 94% HDI [0.26, 0.35] d 0.06, 94% HDI [0.01, 0.10] - e 0.02, 94% HDI [0.00, 0.06] - f 0.19, 94% HDI [0.11, 0.26] - g 0.04, 94% HDI [0.00, 0.08] + e 0.0, 94% HDI [0.0, 0.0] + f 0.1, 94% HDI [0.1, 0.2] + g 0.0, 94% HDI [0.0, 0.0] sigma 0.26, 94% HDI [0.22, 0.30] """ @@ -789,18 +789,18 @@ def summary(self) -> None: ... } ... ) ... ) - >>> result.summary() + >>> result.summary() # doctest: +NUMBER ===========================Difference in Differences============================ Formula: y ~ 1 + group*post_treatment Results: - Causal impact = 0.51, $CI_{94%}$[0.41, 0.60] + Causal impact = 0.5, $CI_{94%}$[0.4, 0.6] Model coefficients: - Intercept 1.08, 94% HDI [1.03, 1.12] - post_treatment[T.True] 0.99, 94% HDI [0.92, 1.05] - group 0.16, 94% HDI [0.09, 0.23] - group:post_treatment[T.True] 0.51, 94% HDI [0.41, 0.60] - sigma 0.08, 94% HDI [0.07, 0.10] + Intercept 1.0, 94% HDI [1.0, 1.1] + post_treatment[T.True] 0.9, 94% HDI [0.9, 1.0] + group 0.1, 94% HDI [0.0, 0.2] + group:post_treatment[T.True] 0.5, 94% HDI [0.4, 0.6] + sigma 0.0, 94% HDI [0.0, 0.1] """ print(f"{self.expt_type:=^80}") @@ -1032,7 +1032,7 @@ def summary(self) -> None: ... ), ... treatment_threshold=0.5, ... ) - >>> result.summary() + >>> result.summary() # doctest: +NUMBER ============================Regression Discontinuity============================ Formula: y ~ 1 + x + treated + x:treated Running variable: x @@ -1041,8 +1041,8 @@ def summary(self) -> None: Results: Discontinuity at threshold = 0.91 Model coefficients: - Intercept 0.09, 94% HDI [-0.00, 0.17] - treated[T.True] 2.45, 94% HDI [1.64, 3.28] + Intercept 0.0, 94% HDI [0.0, 0.1] + treated[T.True] 2.4, 94% HDI [1.6, 3.2] x 1.32, 94% HDI [1.14, 1.50] x:treated[T.True] -3.09, 94% HDI [-4.16, -2.03] sigma 0.36, 94% HDI [0.31, 0.41] From 7ea2c112407017ba8dcffe02c4345b171434e53f Mon Sep 17 00:00:00 2001 From: Jason Preszler Date: Wed, 6 Sep 2023 13:07:38 -0700 Subject: [PATCH 14/23] remove unneeded ellipsis options --- causalpy/pymc_experiments.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/causalpy/pymc_experiments.py b/causalpy/pymc_experiments.py index 81cd6db7..69f0f25a 100644 --- a/causalpy/pymc_experiments.py +++ b/causalpy/pymc_experiments.py @@ -66,9 +66,9 @@ def idata(self): ... model=cp.pymc_models.LinearRegression( ... sample_kwargs={"random_seed": seed, "progressbar": False}), ... ) - >>> result.idata # doctest: +ELLIPSIS + >>> result.idata Inference data... - >>> result.idata.posterior # doctest: +ELLIPSIS + >>> result.idata.posterior Dimensions... """ From 11aa92547b7f18a791fd3a531c9f9c0c23635240 Mon Sep 17 00:00:00 2001 From: Jason Preszler Date: Sat, 9 Sep 2023 09:57:54 -0700 Subject: [PATCH 15/23] fix formula rendering --- causalpy/pymc_models.py | 64 ++++++++++++------- .../tests/test_integration_pymc_examples.py | 4 +- 2 files changed, 44 insertions(+), 24 deletions(-) diff --git a/causalpy/pymc_models.py b/causalpy/pymc_models.py index b179be63..8bb02078 100644 --- a/causalpy/pymc_models.py +++ b/causalpy/pymc_models.py @@ -5,7 +5,8 @@ - LinearRegression model Models are intended to be used from inside an experiment -class (see `pymc_experiments.py +class (see :py:mod:`pymc_experiments.py<.pymc_experiments.py>` old +link `pymc_experiments.py `_). This is why the examples require some extra manipulation input data, often to ensure `y` has the correct shape. @@ -31,7 +32,7 @@ class ModelBuilder(pm.Model): - build_model: must be implemented by subclasses - fit: populates idata attribute - predict: returns predictions on new data - - score: returns Bayesian R^2 + - score: returns Bayesian :math: `R^2` """ def __init__(self, sample_kwargs: Optional[Dict[str, Any]] = None): @@ -208,10 +209,15 @@ class WeightedSumFitter(ModelBuilder): Defines the PyMC model: - - y ~ Normal(mu, sigma) - - sigma ~ HalfNormal(1) - - mu = X * beta - - beta ~ Dirichlet(1,...,1) + .. math:: + + sigma \sim HalfNormal(1) + + beta \sim Dirichlet(1,...,1) + + mu = X * beta + + y \sim Normal(mu, sigma) Example -------- @@ -224,18 +230,23 @@ class WeightedSumFitter(ModelBuilder): >>> wsf = WeightedSumFitter(sample_kwargs={"progressbar": False}) >>> wsf.fit(X,y) Inference ... - """ + """ # noqa: W605 def build_model(self, X, y, coords): """ Defines the PyMC model: - - y ~ Normal(mu, sigma) - - sigma ~ HalfNormal(1) - - mu = X * beta - - beta ~ Dirichlet(1,...,1) + .. math:: - """ + sigma \sim HalfNormal(1) + + beta \sim Dirichlet(1,...,1) + + mu = X * beta + + y \sim Normal(mu, sigma) + + """ # noqa: W605 with self: self.add_coords(coords) n_predictors = X.shape[1] @@ -261,10 +272,14 @@ class LinearRegression(ModelBuilder): Defines the PyMC model - - y ~ Normal(mu, sigma) - - mu = X * beta - - beta ~ Normal(0, 50) - - sigma ~ HalfNormal(1) + .. math:: + beta \sim Normal(0, 50) + + sigma \sim HalfNormal(1) + + mu = X * beta + + y \sim Normal(mu, sigma) Example -------- @@ -281,17 +296,22 @@ class LinearRegression(ModelBuilder): ... }, ... ) Inference... - """ + """ # noqa: W605 def build_model(self, X, y, coords): """ Defines the PyMC model - - y ~ Normal(mu, sigma) - - mu = X * beta - - beta ~ Normal(0, 50) - - sigma ~ HalfNormal(1) - """ + .. math:: + beta \sim Normal(0, 50) + + sigma \sim HalfNormal(1) + + mu = X * beta + + y \sim Normal(mu, sigma) + + """ # noqa: W605 with self: self.add_coords(coords) X = pm.MutableData("X", X, dims=["obs_ind", "coeffs"]) diff --git a/causalpy/tests/test_integration_pymc_examples.py b/causalpy/tests/test_integration_pymc_examples.py index 7b43dc64..59911a0e 100644 --- a/causalpy/tests/test_integration_pymc_examples.py +++ b/causalpy/tests/test_integration_pymc_examples.py @@ -39,7 +39,7 @@ def test_did_banks_simple(): """ Test simple Differences In Differences Experiment on the 'banks' data set. - formula="bib ~ 1 + district * post_treatment" + :code: `formula="bib ~ 1 + district * post_treatment"` Loads, transforms data and checks: 1. data is a dataframe @@ -92,7 +92,7 @@ def test_did_banks_multi(): Test multiple regression Differences In Differences Experiment on the 'banks' data set. - formula="bib ~ 1 + year + district + post_treatment + district:post_treatment" + :code: `formula="bib ~ 1 + year + district + post_treatment + district:post_treatment"` # noqa: E501 Loads, transforms data and checks: 1. data is a dataframe From 53f0fd0f7dd0003cb8e2447e1cdd4af99cf312d1 Mon Sep 17 00:00:00 2001 From: Jason Preszler Date: Sun, 10 Sep 2023 13:23:34 -0700 Subject: [PATCH 16/23] fixed link and removed awkward `round` call --- causalpy/pymc_models.py | 13 ++++++++----- 1 file changed, 8 insertions(+), 5 deletions(-) diff --git a/causalpy/pymc_models.py b/causalpy/pymc_models.py index 8bb02078..be93c188 100644 --- a/causalpy/pymc_models.py +++ b/causalpy/pymc_models.py @@ -5,9 +5,7 @@ - LinearRegression model Models are intended to be used from inside an experiment -class (see :py:mod:`pymc_experiments.py<.pymc_experiments.py>` old -link `pymc_experiments.py -`_). +class (see :doc:`PyMC experiments`). This is why the examples require some extra manipulation input data, often to ensure `y` has the correct shape. @@ -181,11 +179,16 @@ def score(self, X, y) -> pd.Series: >>> X = rng.normal(loc=0, scale=1, size=(200, 2)) >>> y = rng.normal(loc=0, scale=1, size=(200,)) >>> model = MyToyModel( - ... sample_kwargs={"chains": 2, "draws": 2000, "progressbar": False} + ... sample_kwargs={ + ... "chains": 2, + ... "draws": 2000, + ... "progressbar": False, + ... "random_seed": rng + ... } ... ) >>> model.fit(X, y) Inference... - >>> round(model.score(X, y),2) # using round() to simplify doctest + >>> model.score(X, y) # doctest: +NUMBER r2 0.34 r2_std 0.02 dtype: float64 From df5ae6227ede4cc0a8b43dfa4479183d795c615f Mon Sep 17 00:00:00 2001 From: "Benjamin T. Vincent" Date: Mon, 11 Sep 2023 20:58:47 +0100 Subject: [PATCH 17/23] fix minor formatting error --- docs/source/examples.rst | 2 +- docs/source/index.rst | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/docs/source/examples.rst b/docs/source/examples.rst index 9f3afcb5..b70a6300 100644 --- a/docs/source/examples.rst +++ b/docs/source/examples.rst @@ -54,7 +54,7 @@ Regression Discontinuity Instrumental Variables Regression -======================== +================================= .. toctree:: :titlesonly: diff --git a/docs/source/index.rst b/docs/source/index.rst index de86d5d4..bba56599 100644 --- a/docs/source/index.rst +++ b/docs/source/index.rst @@ -110,7 +110,7 @@ Interrupted time series analysis is appropriate when you have a time series of o .. image:: _static/interrupted_time_series_pymc.svg Instrumental Variable Regression -""""""""""""""""""""""" +"""""""""""""""""""""""""""""""" Instrumental Variable regression is an appropriate technique when you wish to estimate the treatment effect of some variable on another, but are concerned that the treatment variable is endogenous in the system of interest i.e. correlated with the errors. In this case an "instrument" variable can be used in a regression context to disentangle treatment effect due to the threat of confounding due to endogeneity. .. image:: _static/iv_reg2.png From 8cc283e0251fb23eaba4d4175919a282330da47b Mon Sep 17 00:00:00 2001 From: Jason Preszler Date: Tue, 12 Sep 2023 10:35:32 -0700 Subject: [PATCH 18/23] remove redundencies, clean up some wording --- causalpy/data/simulate_data.py | 2 +- causalpy/pymc_experiments.py | 221 ++++++++------------------------- causalpy/pymc_models.py | 170 +++++++------------------ causalpy/skl_experiments.py | 59 +++------ 4 files changed, 115 insertions(+), 337 deletions(-) diff --git a/causalpy/data/simulate_data.py b/causalpy/data/simulate_data.py index e99f2eac..a2c9fa20 100644 --- a/causalpy/data/simulate_data.py +++ b/causalpy/data/simulate_data.py @@ -46,7 +46,7 @@ def generate_synthetic_control_data( :param N: Number fo data points :param treatment_time: - Index where treatment begins in the generated data frame + Index where treatment begins in the generated dataframe :param grw_mu: Mean of Gaussian Random Walk :param grw_sigma: diff --git a/causalpy/pymc_experiments.py b/causalpy/pymc_experiments.py index 69f0f25a..95fdafe4 100644 --- a/causalpy/pymc_experiments.py +++ b/causalpy/pymc_experiments.py @@ -52,25 +52,6 @@ def __init__(self, model=None, **kwargs): def idata(self): """ Access to the models InferenceData object - - Example - -------- - >>> import causalpy as cp - >>> df = cp.load_data("did") - >>> seed = 42 - >>> result = cp.pymc_experiments.DifferenceInDifferences( - ... df, - ... formula="y ~ 1 + group*post_treatment", - ... time_variable_name="t", - ... group_variable_name="group", - ... model=cp.pymc_models.LinearRegression( - ... sample_kwargs={"random_seed": seed, "progressbar": False}), - ... ) - >>> result.idata - Inference data... - >>> result.idata.posterior - - Dimensions... """ return self.model.idata @@ -127,7 +108,7 @@ class PrePostFit(ExperimentalDesign): the pre-intervention data. :param data: - A pandas data frame + A pandas dataframe :param treatment_time: The time when treatment occured, should be in reference to the data index :param formula: @@ -153,6 +134,18 @@ class PrePostFit(ExperimentalDesign): ... } ... ), ... ) + >>> result.summary() # doctest: +NUMBER + ==================================Pre-Post Fit================================== + Formula: actual ~ 0 + a + b + c + d + e + f + g + Model coefficients: + a 0.33, 94% HDI [0.30, 0.38] + b 0.05, 94% HDI [0.01, 0.09] + c 0.31, 94% HDI [0.26, 0.35] + d 0.06, 94% HDI [0.01, 0.10] + e 0.02, 94% HDI [0.00, 0.06] + f 0.20, 94% HDI [0.12, 0.26] + g 0.04, 94% HDI [0.00, 0.08] + sigma 0.26, 94% HDI [0.22, 0.30] """ def __init__( @@ -237,10 +230,6 @@ def _input_validation(self, data, treatment_time): def plot(self, counterfactual_label="Counterfactual", **kwargs): """ Plot the results - - Example - -------- - >>> result.plot() # doctest: +SKIP """ fig, ax = plt.subplots(3, 1, sharex=True, figsize=(7, 8)) @@ -343,38 +332,6 @@ def plot(self, counterfactual_label="Counterfactual", **kwargs): def summary(self) -> None: """ Print text output summarising the results - - Example - --------- - >>> import causalpy as cp - >>> sc = cp.load_data("sc") - >>> treatment_time = 70 - >>> seed = 42 - >>> result = cp.pymc_experiments.PrePostFit( - ... sc, - ... treatment_time, - ... formula="actual ~ 0 + a + b + c + d + e + f + g", - ... model=cp.pymc_models.WeightedSumFitter( - ... sample_kwargs={ - ... "draws": 2000, - ... "target_accept": 0.95, - ... "random_seed": seed, - ... "progressbar": False, - ... } - ... ), - ... ) - >>> result.summary() # doctest: +NUMBER - ==================================Pre-Post Fit================================== - Formula: actual ~ 0 + a + b + c + d + e + f + g - Model coefficients: - a 0.34, 94% HDI [0.30, 0.38] - b 0.05, 94% HDI [0.01, 0.09] - c 0.31, 94% HDI [0.26, 0.35] - d 0.06, 94% HDI [0.01, 0.10] - e 0.0, 94% HDI [0.0, 0.0] - f 0.1, 94% HDI [0.1, 0.2] - g 0.0, 94% HDI [0.0, 0.0] - sigma 0.26, 94% HDI [0.22, 0.30] """ print(f"{self.expt_type:=^80}") @@ -388,7 +345,7 @@ class InterruptedTimeSeries(PrePostFit): A wrapper around PrePostFit class :param data: - A pandas data frame + A pandas dataframe :param treatment_time: The time when treatment occured, should be in reference to the data index :param formula: @@ -427,7 +384,7 @@ class SyntheticControl(PrePostFit): """A wrapper around the PrePostFit class :param data: - A pandas data frame + A pandas dataframe :param treatment_time: The time when treatment occured, should be in reference to the data index :param formula: @@ -477,7 +434,7 @@ class DifferenceInDifferences(ExperimentalDesign): There is no pre/post intervention data distinction for DiD, we fit all the data available. :param data: - A pandas data frame + A pandas dataframe :param formula: A statistical model formula :param time_variable_name: @@ -505,6 +462,18 @@ class DifferenceInDifferences(ExperimentalDesign): ... } ... ) ... ) + >>> result.summary() # doctest: +NUMBER + ===========================Difference in Differences============================ + Formula: y ~ 1 + group*post_treatment + + Results: + Causal impact = 0.5, $CI_{94%}$[0.4, 0.6] + Model coefficients: + Intercept 1.0, 94% HDI [1.0, 1.1] + post_treatment[T.True] 0.9, 94% HDI [0.9, 1.0] + group 0.1, 94% HDI [0.0, 0.2] + group:post_treatment[T.True] 0.5, 94% HDI [0.4, 0.6] + sigma 0.0, 94% HDI [0.0, 0.1] """ def __init__( @@ -625,12 +594,6 @@ def _input_validation(self): def plot(self): """Plot the results. Creating the combined mean + HDI legend entries is a bit involved. - - Example - -------- - Assuming `result` is the result of a DiD experiment: - - >>> result.plot() # doctest: +SKIP """ fig, ax = plt.subplots() @@ -769,38 +732,6 @@ def _causal_impact_summary_stat(self) -> str: def summary(self) -> None: """ Print text output summarising the results - - Example - -------- - >>> import causalpy as cp - >>> df = cp.load_data("did") - >>> seed = 42 - >>> result = cp.pymc_experiments.DifferenceInDifferences( - ... df, - ... formula="y ~ 1 + group*post_treatment", - ... time_variable_name="t", - ... group_variable_name="group", - ... model=cp.pymc_models.LinearRegression( - ... sample_kwargs={ - ... "draws": 2000, - ... "target_accept": 0.95, - ... "random_seed": seed, - ... "progressbar": False, - ... } - ... ) - ... ) - >>> result.summary() # doctest: +NUMBER - ===========================Difference in Differences============================ - Formula: y ~ 1 + group*post_treatment - - Results: - Causal impact = 0.5, $CI_{94%}$[0.4, 0.6] - Model coefficients: - Intercept 1.0, 94% HDI [1.0, 1.1] - post_treatment[T.True] 0.9, 94% HDI [0.9, 1.0] - group 0.1, 94% HDI [0.0, 0.2] - group:post_treatment[T.True] 0.5, 94% HDI [0.4, 0.6] - sigma 0.0, 94% HDI [0.0, 0.1] """ print(f"{self.expt_type:=^80}") @@ -849,6 +780,20 @@ class RegressionDiscontinuity(ExperimentalDesign): ... ), ... treatment_threshold=0.5, ... ) + >>> result.summary() # doctest: +NUMBER + ============================Regression Discontinuity============================ + Formula: y ~ 1 + x + treated + x:treated + Running variable: x + Threshold on running variable: 0.5 + + Results: + Discontinuity at threshold = 0.91 + Model coefficients: + Intercept 0.09, 94% HDI [-0.00, 0.17] + treated[T.True] 2.45, 94% HDI [1.66, 3.28] + x 1.32, 94% HDI [1.14, 1.50] + x:treated[T.True] -3.08, 94% HDI [-4.17, -2.05] + sigma 0.36, 94% HDI [0.31, 0.41] """ def __init__( @@ -961,10 +906,6 @@ def _is_treated(self, x): def plot(self): """ Plot the results - - Example - -------- - >>> result.plot() # doctest: +SKIP """ fig, ax = plt.subplots() # Plot raw data @@ -1013,39 +954,6 @@ def plot(self): def summary(self) -> None: """ Print text output summarising the results - - Example - -------- - >>> import causalpy as cp - >>> df = cp.load_data("rd") - >>> seed = 42 - >>> result = cp.pymc_experiments.RegressionDiscontinuity( - ... df, - ... formula="y ~ 1 + x + treated + x:treated", - ... model=cp.pymc_models.LinearRegression( - ... sample_kwargs={ - ... "draws": 2000, - ... "target_accept": 0.95, - ... "random_seed": seed, - ... "progressbar": False, - ... }, - ... ), - ... treatment_threshold=0.5, - ... ) - >>> result.summary() # doctest: +NUMBER - ============================Regression Discontinuity============================ - Formula: y ~ 1 + x + treated + x:treated - Running variable: x - Threshold on running variable: 0.5 - - Results: - Discontinuity at threshold = 0.91 - Model coefficients: - Intercept 0.0, 94% HDI [0.0, 0.1] - treated[T.True] 2.4, 94% HDI [1.6, 3.2] - x 1.32, 94% HDI [1.14, 1.50] - x:treated[T.True] -3.09, 94% HDI [-4.16, -2.03] - sigma 0.36, 94% HDI [0.31, 0.41] """ print(f"{self.expt_type:=^80}") @@ -1064,7 +972,7 @@ class PrePostNEGD(ExperimentalDesign): A class to analyse data from pretest/posttest designs :param data: - A pandas data frame + A pandas dataframe :param formula: A statistical model formula :param group_variable_name: @@ -1092,6 +1000,17 @@ class PrePostNEGD(ExperimentalDesign): ... } ... ) ... ) + >>> result.summary() # doctest: +NUMBER + ==================Pretest/posttest Nonequivalent Group Design=================== + Formula: post ~ 1 + C(group) + pre + + Results: + Causal impact = 1.8, $CI_{94%}$[1.6, 2.0] + Model coefficients: + Intercept -0.4, 94% HDI [-1.2, 0.2] + C(group)[T.1] 1.8, 94% HDI [1.6, 2.0] + pre 1.0, 94% HDI [0.9, 1.1] + sigma 0.5, 94% HDI [0.4, 0.5] """ def __init__( @@ -1227,38 +1146,6 @@ def _causal_impact_summary_stat(self) -> str: def summary(self) -> None: """ Print text output summarising the results - - Example - -------- - >>> import causalpy as cp - >>> df = cp.load_data("anova1") - >>> seed = 42 - >>> result = cp.pymc_experiments.PrePostNEGD( - ... df, - ... formula="post ~ 1 + C(group) + pre", - ... group_variable_name="group", - ... pretreatment_variable_name="pre", - ... model=cp.pymc_models.LinearRegression( - ... sample_kwargs={ - ... "draws": 2000, - ... "target_accept": 0.95, - ... "random_seed": seed, - ... "progressbar": False, - ... } - ... ) - ... ) - >>> result.summary() # doctest: +NUMBER - ==================Pretest/posttest Nonequivalent Group Design=================== - Formula: post ~ 1 + C(group) + pre - - Results: - Causal impact = 1.8, $CI_{94%}$[1.6, 2.0] - Model coefficients: - Intercept -0.4, 94% HDI [-1.2, 0.2] - C(group)[T.1] 1.8, 94% HDI [1.6, 2.0] - pre 1.0, 94% HDI [0.9, 1.1] - sigma 0.5, 94% HDI [0.4, 0.5] - """ print(f"{self.expt_type:=^80}") diff --git a/causalpy/pymc_models.py b/causalpy/pymc_models.py index be93c188..834b9c66 100644 --- a/causalpy/pymc_models.py +++ b/causalpy/pymc_models.py @@ -30,7 +30,43 @@ class ModelBuilder(pm.Model): - build_model: must be implemented by subclasses - fit: populates idata attribute - predict: returns predictions on new data - - score: returns Bayesian :math: `R^2` + - score: returns Bayesian :math:`R^2` + + Example + ------- + >>> import causalpy as cp + >>> import numpy as np + >>> import pymc as pm + >>> from causalpy.pymc_models import ModelBuilder + >>> class MyToyModel(ModelBuilder): + ... def build_model(self, X, y, coords): + ... with self: + ... X_ = pm.MutableData(name="X", value=X) + ... y_ = pm.MutableData(name="y", value=y) + ... beta = pm.Normal("beta", mu=0, sigma=1, shape=X_.shape[1]) + ... sigma = pm.HalfNormal("sigma", sigma=1) + ... mu = pm.Deterministic("mu", pm.math.dot(X_, beta)) + ... pm.Normal("y_hat", mu=mu, sigma=sigma, observed=y_) + >>> rng = np.random.default_rng(seed=42) + >>> X = rng.normal(loc=0, scale=1, size=(20, 2)) + >>> y = rng.normal(loc=0, scale=1, size=(20,)) + >>> model = MyToyModel( + ... sample_kwargs={ + ... "chains": 2, + ... "draws": 2000, + ... "progressbar": False, + ... "random_seed": rng, + ... } + ... ) + >>> model.fit(X, y) + Inference... + >>> X_new = rng.normal(loc=0, scale=1, size=(20,2)) + >>> model.predict(X_new) + Inference... + >>> model.score(X, y) # doctest: +NUMBER + r2 0.38 + r2_std 0.08 + dtype: float64 """ def __init__(self, sample_kwargs: Optional[Dict[str, Any]] = None): @@ -43,22 +79,7 @@ def __init__(self, sample_kwargs: Optional[Dict[str, Any]] = None): self.sample_kwargs = sample_kwargs if sample_kwargs is not None else {} def build_model(self, X, y, coords) -> None: - """Build the model, must be implemented by subclass. - - Example - ------- - >>> import pymc as pm - >>> from causalpy.pymc_models import ModelBuilder - >>> class CausalPyModel(ModelBuilder): - ... def build_model(self, X, y): - ... with self: - ... X_ = pm.MutableData(name="X", value=X) - ... y_ = pm.MutableData(name="y", value=y) - ... beta = pm.Normal("beta", mu=0, sigma=1, shape=X_.shape[1]) - ... sigma = pm.HalfNormal("sigma", sigma=1) - ... mu = pm.Deterministic("mu", pm.math.dot(X_, beta)) - ... pm.Normal("y_hat", mu=mu, sigma=sigma, observed=y_) - """ + """Build the model, must be implemented by subclass.""" raise NotImplementedError("This method must be implemented by a subclass") def _data_setter(self, X) -> None: @@ -74,32 +95,6 @@ def _data_setter(self, X) -> None: def fit(self, X, y, coords: Optional[Dict[str, Any]] = None) -> None: """Draw samples fromposterior, prior predictive, and posterior predictive distributions, placing them in the model's idata attribute. - - .. note:: - Calls the build_model method - - Example - ------- - >>> import numpy as np - >>> import pymc as pm - >>> from causalpy.pymc_models import ModelBuilder - >>> class MyToyModel(ModelBuilder): - ... def build_model(self, X, y, coords): - ... with self: - ... X_ = pm.MutableData(name="X", value=X) - ... y_ = pm.MutableData(name="y", value=y) - ... beta = pm.Normal("beta", mu=0, sigma=1, shape=X_.shape[1]) - ... sigma = pm.HalfNormal("sigma", sigma=1) - ... mu = pm.Deterministic("mu", pm.math.dot(X_, beta)) - ... pm.Normal("y_hat", mu=mu, sigma=sigma, observed=y_) - >>> rng = np.random.default_rng(seed=42) - >>> X = rng.normal(loc=0, scale=1, size=(20, 2)) - >>> y = rng.normal(loc=0, scale=1, size=(20,)) - >>> model = MyToyModel( - ... sample_kwargs={"chains": 2, "draws": 2, "progressbar": False} - ... ) - >>> model.fit(X, y) - Inference ... """ self.build_model(X, y, coords) with self.model: @@ -117,32 +112,6 @@ def predict(self, X): .. caution:: Results in KeyError if model hasn't been fit. - Example - ------- - >>> import causalpy as cp - >>> import numpy as np - >>> import pymc as pm - >>> from causalpy.pymc_models import ModelBuilder - >>> class MyToyModel(ModelBuilder): - ... def build_model(self, X, y, coords): - ... with self: - ... X_ = pm.MutableData(name="X", value=X) - ... y_ = pm.MutableData(name="y", value=y) - ... beta = pm.Normal("beta", mu=0, sigma=1, shape=X_.shape[1]) - ... sigma = pm.HalfNormal("sigma", sigma=1) - ... mu = pm.Deterministic("mu", pm.math.dot(X_, beta)) - ... pm.Normal("y_hat", mu=mu, sigma=sigma, observed=y_) - >>> rng = np.random.default_rng(seed=42) - >>> X = rng.normal(loc=0, scale=1, size=(20, 2)) - >>> y = rng.normal(loc=0, scale=1, size=(20,)) - >>> model = MyToyModel( - ... sample_kwargs={"chains": 2, "draws": 2, "progressbar": False} - ... ) - >>> model.fit(X, y) - Inference... - >>> X_new = rng.normal(loc=0, scale=1, size=(20,2)) - >>> model.predict(X_new) - Inference... """ self._data_setter(X) @@ -160,38 +129,6 @@ def score(self, X, y) -> pd.Series: The Bayesian :math:`R^2` is not the same as the traditional coefficient of determination, https://en.wikipedia.org/wiki/Coefficient_of_determination. - Example - -------- - >>> import causalpy as cp - >>> import numpy as np - >>> import pymc as pm - >>> from causalpy.pymc_models import ModelBuilder - >>> class MyToyModel(ModelBuilder): - ... def build_model(self, X, y, coords): - ... with self: - ... X_ = pm.MutableData(name="X", value=X) - ... y_ = pm.MutableData(name="y", value=y) - ... beta = pm.Normal("beta", mu=0, sigma=1, shape=X_.shape[1]) - ... sigma = pm.HalfNormal("sigma", sigma=1) - ... mu = pm.Deterministic("mu", pm.math.dot(X_, beta)) - ... pm.Normal("y_hat", mu=mu, sigma=sigma, observed=y_) - >>> rng = np.random.default_rng(seed=42) - >>> X = rng.normal(loc=0, scale=1, size=(200, 2)) - >>> y = rng.normal(loc=0, scale=1, size=(200,)) - >>> model = MyToyModel( - ... sample_kwargs={ - ... "chains": 2, - ... "draws": 2000, - ... "progressbar": False, - ... "random_seed": rng - ... } - ... ) - >>> model.fit(X, y) - Inference... - >>> model.score(X, y) # doctest: +NUMBER - r2 0.34 - r2_std 0.02 - dtype: float64 """ yhat = self.predict(X) yhat = az.extract( @@ -207,7 +144,7 @@ class WeightedSumFitter(ModelBuilder): """ Used for synthetic control experiments - .. note: Generally, the `.fit()` method should be rather than calling + .. note: Generally, the `.fit()` method should be used rather than calling `.build_model()` directly. Defines the PyMC model: @@ -237,19 +174,8 @@ class WeightedSumFitter(ModelBuilder): def build_model(self, X, y, coords): """ - Defines the PyMC model: - - .. math:: - - sigma \sim HalfNormal(1) - - beta \sim Dirichlet(1,...,1) - - mu = X * beta - - y \sim Normal(mu, sigma) - - """ # noqa: W605 + Defines the PyMC model + """ with self: self.add_coords(coords) n_predictors = X.shape[1] @@ -270,7 +196,7 @@ class LinearRegression(ModelBuilder): """ Custom PyMC model for linear regression - .. note: Generally, the `.fit()` method should be rather than calling + .. note: Generally, the `.fit()` method should be used rather than calling `.build_model()` directly. Defines the PyMC model @@ -304,17 +230,7 @@ class LinearRegression(ModelBuilder): def build_model(self, X, y, coords): """ Defines the PyMC model - - .. math:: - beta \sim Normal(0, 50) - - sigma \sim HalfNormal(1) - - mu = X * beta - - y \sim Normal(mu, sigma) - - """ # noqa: W605 + """ with self: self.add_coords(coords) X = pm.MutableData("X", X, dims=["obs_ind", "coeffs"]) diff --git a/causalpy/skl_experiments.py b/causalpy/skl_experiments.py index 8be30d4e..9db34736 100644 --- a/causalpy/skl_experiments.py +++ b/causalpy/skl_experiments.py @@ -1,7 +1,7 @@ """ Experiments for Scikit-Learn models -- ExperimentalDesign: base class for skl experiments +- ExperimentalDesign: base class for scikit-learn experiments - PrePostFit: base class for synthetic control and interrupted time series - SyntheticControl - InterruptedTimeSeries @@ -45,7 +45,7 @@ class PrePostFit(ExperimentalDesign): :param formula: A statistical model formula :param model: - An sklearn model object + An scikit-learn model object Example -------- @@ -59,6 +59,8 @@ class PrePostFit(ExperimentalDesign): ... formula="actual ~ 0 + a + b + c + d + e + f + g", ... model = cp.skl_models.WeightedProportion() ... ) + >>> result.get_coeffs() + array(...) """ def __init__( @@ -177,21 +179,6 @@ def plot(self, counterfactual_label="Counterfactual", **kwargs): def get_coeffs(self): """ Returns model coefficients - - Example - -------- - >>> from sklearn.linear_model import LinearRegression - >>> import causalpy as cp - >>> df = cp.load_data("sc") - >>> treatment_time = 70 - >>> result = cp.skl_experiments.PrePostFit( - ... df, - ... treatment_time, - ... formula="actual ~ 0 + a + b + c + d + e + f + g", - ... model = cp.skl_models.WeightedProportion() - ... ) - >>> result.get_coeffs() - array(...) """ return np.squeeze(self.model.coef_) @@ -299,7 +286,7 @@ class DifferenceInDifferences(ExperimentalDesign): :param group_variable_name: Name of the data column for the group variable :param model: - An skl model for difference in differences + An scikit-learn model for difference in differences Example -------- @@ -513,6 +500,18 @@ class RegressionDiscontinuity(ExperimentalDesign): ... model=LinearRegression(), ... treatment_threshold=0.5, ... ) + >>> result.summary() # doctest: +NORMALIZE_WHITESPACE,+NUMBER + Difference in Differences experiment + Formula: y ~ 1 + x + treated + Running variable: x + Threshold on running variable: 0.5 + + Results: + Discontinuity at threshold = 0.19 + Model coefficients: + Intercept 0.0 + treated[T.True] 0.19 + x 1.23 """ def __init__( @@ -645,30 +644,6 @@ def plot(self): def summary(self): """ Print text output summarising the results - - Example - -------- - >>> import causalpy as cp - >>> from sklearn.linear_model import LinearRegression - >>> data = cp.load_data("rd") - >>> result = cp.skl_experiments.RegressionDiscontinuity( - ... data, - ... formula="y ~ 1 + x + treated", - ... model=LinearRegression(), - ... treatment_threshold=0.5, - ... ) - >>> result.summary() # doctest: +NORMALIZE_WHITESPACE,+NUMBER - Difference in Differences experiment - Formula: y ~ 1 + x + treated - Running variable: x - Threshold on running variable: 0.5 - - Results: - Discontinuity at threshold = 0.19 - Model coefficients: - Intercept 0.0 - treated[T.True] 0.19 - x 1.23 """ print("Difference in Differences experiment") print(f"Formula: {self.formula}") From 8a28509ef6ccaf8e6bbb3a69af46eb9f76d4587f Mon Sep 17 00:00:00 2001 From: Jason Preszler Date: Wed, 13 Sep 2023 15:42:39 -0700 Subject: [PATCH 19/23] fix test failure and doc anomalies --- causalpy/pymc_experiments.py | 21 ++++++++++++--------- causalpy/pymc_models.py | 18 ++++++++++-------- 2 files changed, 22 insertions(+), 17 deletions(-) diff --git a/causalpy/pymc_experiments.py b/causalpy/pymc_experiments.py index 95fdafe4..d0e1318f 100644 --- a/causalpy/pymc_experiments.py +++ b/causalpy/pymc_experiments.py @@ -128,6 +128,7 @@ class PrePostFit(ExperimentalDesign): ... formula="actual ~ 0 + a + b + c + d + e + f + g", ... model=cp.pymc_models.WeightedSumFitter( ... sample_kwargs={ + ... "draws": 2000, ... "target_accept": 0.95, ... "random_seed": seed, ... "progressbar": False @@ -138,12 +139,12 @@ class PrePostFit(ExperimentalDesign): ==================================Pre-Post Fit================================== Formula: actual ~ 0 + a + b + c + d + e + f + g Model coefficients: - a 0.33, 94% HDI [0.30, 0.38] + a 0.34, 94% HDI [0.30, 0.38] b 0.05, 94% HDI [0.01, 0.09] c 0.31, 94% HDI [0.26, 0.35] d 0.06, 94% HDI [0.01, 0.10] e 0.02, 94% HDI [0.00, 0.06] - f 0.20, 94% HDI [0.12, 0.26] + f 0.19, 94% HDI [0.11, 0.26] g 0.04, 94% HDI [0.00, 0.08] sigma 0.26, 94% HDI [0.22, 0.30] """ @@ -773,6 +774,7 @@ class RegressionDiscontinuity(ExperimentalDesign): ... formula="y ~ 1 + x + treated + x:treated", ... model=cp.pymc_models.LinearRegression( ... sample_kwargs={ + ... "draws": 2000, ... "target_accept": 0.95, ... "random_seed": seed, ... "progressbar": False, @@ -790,9 +792,9 @@ class RegressionDiscontinuity(ExperimentalDesign): Discontinuity at threshold = 0.91 Model coefficients: Intercept 0.09, 94% HDI [-0.00, 0.17] - treated[T.True] 2.45, 94% HDI [1.66, 3.28] + treated[T.True] 2.45, 94% HDI [1.64, 3.28] x 1.32, 94% HDI [1.14, 1.50] - x:treated[T.True] -3.08, 94% HDI [-4.17, -2.05] + x:treated[T.True] -3.09, 94% HDI [-4.16, -2.03] sigma 0.36, 94% HDI [0.31, 0.41] """ @@ -1224,12 +1226,13 @@ class InstrumentalVariable(ExperimentalDesign): >>> instruments_data = test_data[["X", "Z"]] >>> data = test_data[["y", "X"]] >>> iv = InstrumentalVariable( - ... instruments_data=instruments_data, - ... data=data, - ... instruments_formula=instruments_formula, - ... formula=formula, - ... model=InstrumentalVariableRegression(sample_kwargs=sample_kwargs), + ... instruments_data=instruments_data, + ... data=data, + ... instruments_formula=instruments_formula, + ... formula=formula, + ... model=InstrumentalVariableRegression(sample_kwargs=sample_kwargs), ... ) + """ def __init__( diff --git a/causalpy/pymc_models.py b/causalpy/pymc_models.py index 834b9c66..4b10d734 100644 --- a/causalpy/pymc_models.py +++ b/causalpy/pymc_models.py @@ -64,8 +64,8 @@ class ModelBuilder(pm.Model): >>> model.predict(X_new) Inference... >>> model.score(X, y) # doctest: +NUMBER - r2 0.38 - r2_std 0.08 + r2 0.3 + r2_std 0.0 dtype: float64 """ @@ -144,8 +144,9 @@ class WeightedSumFitter(ModelBuilder): """ Used for synthetic control experiments - .. note: Generally, the `.fit()` method should be used rather than calling - `.build_model()` directly. + .. note:: + Generally, the `.fit()` method should be used rather than + calling `.build_model()` directly. Defines the PyMC model: @@ -196,8 +197,9 @@ class LinearRegression(ModelBuilder): """ Custom PyMC model for linear regression - .. note: Generally, the `.fit()` method should be used rather than calling - `.build_model()` directly. + .. note: + Generally, the `.fit()` method should be used rather than + calling `.build_model()` directly. Defines the PyMC model @@ -291,8 +293,8 @@ def build_model(self, X, Z, y, t, coords, priors): instruments and covariates :param priors: An optional dictionary of priors for the mus and sigmas of both regressions - - :code:`priors = {"mus": [0, 0], "sigmas": [1, 1], "eta": 2, "lkj_sd": 2}` + :code:`priors = {"mus": [0, 0], "sigmas": [1, 1], + "eta": 2, "lkj_sd": 2}` """ From e2760bcc4c6286eaf88e09ce01e9ee075bd8c689 Mon Sep 17 00:00:00 2001 From: Jason Preszler Date: Wed, 13 Sep 2023 16:01:27 -0700 Subject: [PATCH 20/23] reduce precision on reg. discont. --- causalpy/pymc_experiments.py | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/causalpy/pymc_experiments.py b/causalpy/pymc_experiments.py index d0e1318f..1c3b39ea 100644 --- a/causalpy/pymc_experiments.py +++ b/causalpy/pymc_experiments.py @@ -143,7 +143,7 @@ class PrePostFit(ExperimentalDesign): b 0.05, 94% HDI [0.01, 0.09] c 0.31, 94% HDI [0.26, 0.35] d 0.06, 94% HDI [0.01, 0.10] - e 0.02, 94% HDI [0.00, 0.06] + e 0.02, 94% HDI [0.00, 0.07] f 0.19, 94% HDI [0.11, 0.26] g 0.04, 94% HDI [0.00, 0.08] sigma 0.26, 94% HDI [0.22, 0.30] @@ -791,11 +791,11 @@ class RegressionDiscontinuity(ExperimentalDesign): Results: Discontinuity at threshold = 0.91 Model coefficients: - Intercept 0.09, 94% HDI [-0.00, 0.17] - treated[T.True] 2.45, 94% HDI [1.64, 3.28] - x 1.32, 94% HDI [1.14, 1.50] - x:treated[T.True] -3.09, 94% HDI [-4.16, -2.03] - sigma 0.36, 94% HDI [0.31, 0.41] + Intercept 0.0, 94% HDI [0.0, 0.1] + treated[T.True] 2.4, 94% HDI [1.6, 3.2] + x 1.3, 94% HDI [1.1, 1.5] + x:treated[T.True] -3.0, 94% HDI [-4.1, -2.0] + sigma 0.3, 94% HDI [0.3, 0.4] """ def __init__( From 34d795d12a9d5512c4d1ae6dbe4f9616b13f9280 Mon Sep 17 00:00:00 2001 From: Jason Preszler Date: Thu, 14 Sep 2023 14:18:44 -0700 Subject: [PATCH 21/23] add make doctest and instructions to contributing --- CONTRIBUTING.md | 14 ++++++++++++++ Makefile | 4 ++++ 2 files changed, 18 insertions(+) diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index 29a345f9..2939100e 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -121,6 +121,20 @@ We recommend that your contribution complies with the following guidelines befor - All public methods must have informative docstrings with sample usage when appropriate. +- Example usage in docstrings is tested via doctest, which can be run via + + ```bash + make doctest + ``` + +- Doctest can also be run directly via pytest, which can be helpful to run only specific tests during development. The following commands run all doctests, only doctests in the pymc_models module, and only the doctests for the `ModelBuilder` class in pymc_models: + + ```bash + pytest --doctest-modules causalpy/ + pytest --doctest-modules causalpy/pymc_models.py + pytest --doctest-modules causalpy/pmyc_models.py::causalpy.pymc_models.ModelBuilder + ``` + - To indicate a work in progress please mark the PR as `draft`. Drafts may be useful to (1) indicate you are working on something to avoid duplicated work, (2) request broad review of functionality or API, or (3) seek collaborators. - All other tests pass when everything is rebuilt from scratch. Tests can be run with: diff --git a/Makefile b/Makefile index b2b562e9..e9ee13f7 100644 --- a/Makefile +++ b/Makefile @@ -17,6 +17,10 @@ check_lint: nbqa isort --check-only . interrogate . +doctest: + pip install causalpy[test] + pytest --doctest-modules causalpy/ + test: pip install causalpy[test] pytest From f37d3b34f69f8d245602f616e324f63388e99e42 Mon Sep 17 00:00:00 2001 From: Jason Preszler Date: Thu, 14 Sep 2023 14:37:41 -0700 Subject: [PATCH 22/23] turn down precision on PrePostFit doctest --- causalpy/pymc_experiments.py | 16 ++++++++-------- 1 file changed, 8 insertions(+), 8 deletions(-) diff --git a/causalpy/pymc_experiments.py b/causalpy/pymc_experiments.py index 1c3b39ea..502c0d41 100644 --- a/causalpy/pymc_experiments.py +++ b/causalpy/pymc_experiments.py @@ -139,14 +139,14 @@ class PrePostFit(ExperimentalDesign): ==================================Pre-Post Fit================================== Formula: actual ~ 0 + a + b + c + d + e + f + g Model coefficients: - a 0.34, 94% HDI [0.30, 0.38] - b 0.05, 94% HDI [0.01, 0.09] - c 0.31, 94% HDI [0.26, 0.35] - d 0.06, 94% HDI [0.01, 0.10] - e 0.02, 94% HDI [0.00, 0.07] - f 0.19, 94% HDI [0.11, 0.26] - g 0.04, 94% HDI [0.00, 0.08] - sigma 0.26, 94% HDI [0.22, 0.30] + a 0.3, 94% HDI [0.3, 0.3] + b 0.0, 94% HDI [0.0, 0.0] + c 0.3, 94% HDI [0.2, 0.3] + d 0.0, 94% HDI [0.0, 0.1] + e 0.0, 94% HDI [0.0, 0.0] + f 0.1, 94% HDI [0.1, 0.2] + g 0.0, 94% HDI [0.0, 0.0] + sigma 0.2, 94% HDI [0.2, 0.3] """ def __init__( From c80d78e093e73ce68da9c4f36cdea2a5eecfea9a Mon Sep 17 00:00:00 2001 From: Jason Preszler Date: Thu, 14 Sep 2023 14:46:49 -0700 Subject: [PATCH 23/23] skip doctests that write simulated data to csv files --- causalpy/data/simulate_data.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/causalpy/data/simulate_data.py b/causalpy/data/simulate_data.py index a2c9fa20..c54a88ca 100644 --- a/causalpy/data/simulate_data.py +++ b/causalpy/data/simulate_data.py @@ -253,7 +253,8 @@ def generate_regression_discontinuity_data( >>> import pathlib >>> from causalpy.data.simulate_data import generate_regression_discontinuity_data >>> df = generate_regression_discontinuity_data(true_treatment_threshold=0.5) - >>> df.to_csv(pathlib.Path.cwd() / 'regression_discontinuity.csv', index=False) + >>> df.to_csv(pathlib.Path.cwd() / 'regression_discontinuity.csv', + ... index=False) # doctest: +SKIP """ def is_treated(x): @@ -288,7 +289,8 @@ def generate_ancova_data( ... treatment_effect=2, ... sigma=1 ... ) - >>> df.to_csv(pathlib.Path.cwd() / 'ancova_data.csv', index=False) + >>> df.to_csv(pathlib.Path.cwd() / 'ancova_data.csv', + ... index=False) # doctest: +SKIP """ group = np.random.choice(2, size=N) pre = np.random.normal(loc=pre_treatment_means[group])