Added IV examples after rebasing

Author: jpreszler

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])

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

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"

causalpy/version.py

@@ -1 +1,2 @@
+"""CausalPy Version"""
 __version__ = "0.1.0"