Merge branch 'main' into summary

Author: drbenvincent

.github/workflows/rtd-link-preview.yml

@@ -0,0 +1,16 @@
+name: Read the Docs Pull Request Preview
+on:
+  pull_request_target:
+    types:
+      - opened
+
+permissions:
+  pull-requests: write
+
+jobs:
+  documentation-links:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: readthedocs/actions/preview@v1
+        with:
+          project-slug: "causalpy"

.gitignore

@@ -7,6 +7,8 @@ _build
 build/
 dist/
 docs/_build/
+docs/build/
+docs/jupyter_execute/
 *.vscode
 .coverage
 *.jupyterlab-workspace

.pre-commit-config.yaml

@@ -22,9 +22,10 @@ repos:
         exclude_types: [svg]
       - id: check-yaml
       - id: check-added-large-files
+        exclude: &exclude_pattern 'iv_weak_instruments.ipynb'
         args: ["--maxkb=1500"]
   - repo: https://github.com/astral-sh/ruff-pre-commit
-    rev: v0.4.8
+    rev: v0.4.9
     hooks:
       # Run the linter
       - id: ruff

.readthedocs.yaml

@@ -7,9 +7,9 @@ version: 2
 
 # Set the version of Python and other tools you might need
 build:
-  os: ubuntu-20.04
+  os: ubuntu-lts-latest
   tools:
-    python: "3.10"
+    python: "3.11"
     # You can also specify other tool versions:
     # nodejs: "16"
     # rust: "1.55"

Makefile

@@ -10,8 +10,6 @@ lint:
 check_lint:
 	ruff check .
 	ruff format --diff --check .
-	nbqa black --check .
-	nbqa ruff .
 	interrogate .
 
 doctest:

README.md

@@ -34,7 +34,7 @@ pip install git+https://github.com/pymc-labs/CausalPy.git
 
 ```python
 import causalpy as cp
-
+import matplotlib.pyplot as plt
 
 # Import and process data
 df = (
@@ -57,6 +57,8 @@ fig, ax = result.plot();
 
 # Get a results summary
 result.summary()
+
+plt.show()
 ```
 
 ## Roadmap

causalpy/data/datasets.py

@@ -35,6 +35,7 @@
     "geolift1": {"filename": "geolift1.csv"},
     "risk": {"filename": "AJR2001.csv"},
     "nhefs": {"filename": "nhefs.csv"},
+    "schoolReturns": {"filename": "schoolingReturns.csv"},
 }
 
 

causalpy/data/schoolingReturns.csv

@@ -1450,7 +1450,7 @@ def __init__(
                 "mus": [self.ols_beta_first_params, self.ols_beta_second_params],
                 "sigmas": [1, 1],
                 "eta": 2,
-                "lkj_sd": 2,
+                "lkj_sd": 1,
             }
         self.priors = priors
         self.model.fit(

causalpy/pymc_experiments.py

@@ -303,8 +303,8 @@ class InstrumentalVariableRegression(ModelBuilder):
     ...                  "mus": [[-2,4], [0.5, 3]],
     ...                  "sigmas": [1, 1],
     ...                  "eta": 2,
-    ...                  "lkj_sd": 2,
-    ...              })
+    ...                  "lkj_sd": 1,
+    ...              }, None)
     Inference data...
     """
 
@@ -340,7 +340,7 @@ def build_model(self, X, Z, y, t, coords, priors):
                 sigma=priors["sigmas"][1],
                 dims="covariates",
             )
-            sd_dist = pm.HalfCauchy.dist(beta=priors["lkj_sd"], shape=2)
+            sd_dist = pm.Exponential.dist(priors["lkj_sd"], shape=2)
             chol, corr, sigmas = pm.LKJCholeskyCov(
                 name="chol_cov",
                 eta=priors["eta"],
@@ -366,24 +366,52 @@ def build_model(self, X, Z, y, t, coords, priors):
                 shape=(X.shape[0], 2),
             )
 
-    def fit(self, X, Z, y, t, coords, priors):
-        """Draw samples from posterior, prior predictive, and posterior predictive
-        distributions.
+    def sample_predictive_distribution(self, ppc_sampler="jax"):
+        """Function to sample the Multivariate Normal posterior predictive
+        Likelihood term in the IV class. This can be slow without
+        using the JAX sampler compilation method. If using the
+        JAX sampler it will sample only the posterior predictive distribution.
+        If using the PYMC sampler if will sample both the prior
+        and posterior predictive distributions."""
+        random_seed = self.sample_kwargs.get("random_seed", None)
+
+        if ppc_sampler == "jax":
+            with self:
+                self.idata.extend(
+                    pm.sample_posterior_predictive(
+                        self.idata,
+                        random_seed=random_seed,
+                        compile_kwargs={"mode": "JAX"},
+                    )
+                )
+        elif ppc_sampler == "pymc":
+            with self:
+                self.idata.extend(pm.sample_prior_predictive(random_seed=random_seed))
+                self.idata.extend(
+                    pm.sample_posterior_predictive(
+                        self.idata,
+                        random_seed=random_seed,
+                    )
+                )
+
+    def fit(self, X, Z, y, t, coords, priors, ppc_sampler=None):
+        """Draw samples from posterior distribution and potentially
+        from the prior and posterior predictive distributions. The
+        fit call can take values for the
+        ppc_sampler = ['jax', 'pymc', None]
+        We default to None, so the user can determine if they wish
+        to spend time sampling the posterior predictive distribution
+        independently.
         """
 
         # Ensure random_seed is used in sample_prior_predictive() and
         # sample_posterior_predictive() if provided in sample_kwargs.
-        random_seed = self.sample_kwargs.get("random_seed", None)
+        # Use JAX for ppc sampling of multivariate likelihood
 
         self.build_model(X, Z, y, t, coords, priors)
         with self:
             self.idata = pm.sample(**self.sample_kwargs)
-            self.idata.extend(pm.sample_prior_predictive(random_seed=random_seed))
-            self.idata.extend(
-                pm.sample_posterior_predictive(
-                    self.idata, progressbar=False, random_seed=random_seed
-                )
-            )
+        self.sample_predictive_distribution(ppc_sampler=ppc_sampler)
         return self.idata