Add jax as dependency re-configure IV class fit call

Author: NathanielF

causalpy/pymc_experiments.py

@@ -1453,7 +1453,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_models.py

@@ -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,27 +366,51 @@ def build_model(self, X, Z, y, t, coords, priors):
                 shape=(X.shape[0], 2),
             )
 
-    def fit(self, X, Z, y, t, coords, priors, sample_ppc=True):
-        """Draw samples from posterior, prior predictive, and posterior predictive
-        distributions.
-        """
-
-        # Ensure random_seed is used in sample_prior_predictive() and
-        # sample_posterior_predictive() if provided in sample_kwargs.
-        # Use JAX for ppc sampling of multivariate likelihood
+    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)
 
-        self.build_model(X, Z, y, t, coords, priors)
-        with self:
-            self.idata = pm.sample(**self.sample_kwargs)
-            if sample_ppc:
+        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":
+            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.
+        # 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.sample_predictive_distribution(ppc_sampler=ppc_sampler)
         return self.idata
 
 

docs/source/_static/interrogate_badge.svg

@@ -15,3 +15,4 @@ dependencies:
   - seaborn>=0.11.2
   - statsmodels
   - xarray>=v2022.11.0
+  - jax