[IV 212] experimenting with nicer plot

Signed-off-by: Nathaniel <[email protected]>

Author: NathanielF

causalpy/pymc_experiments.py

@@ -917,13 +917,15 @@ def __init__(
         self.get_naive_OLS_fit()
         self.get_2SLS_fit()
 
-        # fit the model to the observed (pre-intervention) data
+        # fit the model to the data
         COORDS = {"instruments": self.labels_instruments, "covariates": self.labels}
         self.coords = COORDS
         if priors is None:
             priors = {
                 "mus": [self.ols_beta_first_params, self.ols_beta_second_params],
                 "sigmas": [1, 1],
+                "eta": 2,
+                "lkj_sd": 2,
             }
         self.priors = priors
         self.model.fit(
@@ -950,5 +952,5 @@ def get_naive_OLS_fit(self):
         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])
-        self.ols_beta_params = beta_params
+        self.ols_beta_params = dict(zip(self._x_design_info.column_names, beta_params))
         self.ols_reg = ols_reg

causalpy/pymc_models.py

@@ -136,7 +136,8 @@ def build_model(self, X, Z, y, t, coords, priors):
         :param priors: An optional dictionary of priors for the mus and
                       sigmas of both regressions
 
-        e.g priors = {'mus': [[10, 0], [2, 0]], 'sigmas': [[1, 1], [1, 1]]}
+        e.g priors = {'mus': [[10, 0], [2, 0]], 'sigmas': [[1, 1], [1, 1]]
+                      'eta': 2, 'lkj_sd': 2}
 
         """
 
@@ -155,16 +156,18 @@ def build_model(self, X, Z, y, t, coords, priors):
                 sigma=priors["sigmas"][1],
                 dims="covariates",
             )
-            sd_dist = pm.HalfCauchy.dist(beta=2, shape=2)
+            sd_dist = pm.HalfCauchy.dist(beta=priors["lkj_sd"], shape=2)
             chol, corr, sigmas = pm.LKJCholeskyCov(
-                name="chol_cov", eta=2, n=2, sd_dist=sd_dist
+                name="chol_cov", eta=priors["eta"], n=2, sd_dist=sd_dist
             )
             # compute and store the covariance matrix
             pm.Deterministic(name="cov", var=pt.dot(l=chol, r=chol.T))
 
             # --- Parameterization ---
             mu_y = pm.Deterministic(name="mu_y", var=pm.math.dot(X, beta_z))
+            # focal regression
             mu_t = pm.Deterministic(name="mu_t", var=pm.math.dot(Z, beta_t))
+            # instrumental regression
             mu = pm.Deterministic(name="mu", var=pt.stack(tensors=(mu_y, mu_t), axis=1))
 
             # --- Likelihood ---