Bayesian Workflow with SEMs

[NathanielF]

Bayesian Workflow with SEMs

Related to proposal here
#806

• Notebook follows style guide https://docs.pymc.io/en/latest/contributing/jupyter_style.html
• PR description contains a link to the relevant issue:
  • a tracker one for existing notebooks (tracker issues have the "tracker id" label)
  • or a proposal one for new notebooks
• Check the notebook is not excluded from any pre-commit check: https://github.com/pymc-devs/pymc-examples/blob/main/.pre-commit-config.yaml

Helpful links

• https://github.com/pymc-devs/pymc-examples/blob/main/CONTRIBUTING.md

---

📚 Documentation preview 📚: https://pymc-examples--807.org.readthedocs.build/en/807/

[review-notebook-app]

Check out this pull request on 

See visual diffs & provide feedback on Jupyter Notebooks.

---

Powered by ReviewNB

[NathanielF]

Apparent indexing issue between pymc versions 5.17 --> 5.30

Indexing trick works for pymc 5.17, but breaks on 5.30

---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
Cell In[15], [line 61](vscode-notebook-cell:?execution_count=15&line=61)
     [57](vscode-notebook-cell:?execution_count=15&line=57) priors = {"lambdas": [1, 0.5], "eta": 2, "B": [0, 0.5], "tau": [0, 1]}
     [59](vscode-notebook-cell:?execution_count=15&line=59) priors_wide = {"lambdas": [1, 5], "eta": 2, "B": [0, 5], "tau": [0, 10]}
---> [61](vscode-notebook-cell:?execution_count=15&line=61) sem_model_hierarchical_tight = make_hierarchical(priors, grp_idx)
     [62](vscode-notebook-cell:?execution_count=15&line=62) sem_model_hierarchical_wide = make_hierarchical(priors_wide, grp_idx)
     [64](vscode-notebook-cell:?execution_count=15&line=64) pm.model_to_graphviz(sem_model_hierarchical_tight)

Cell In[15], [line 52](vscode-notebook-cell:?execution_count=15&line=52)
     [50](vscode-notebook-cell:?execution_count=15&line=50)         Sigma_y.append(Sigma_y_g)
     [51](vscode-notebook-cell:?execution_count=15&line=51)     Sigma_y = pt.stack(Sigma_y)
---> [52](vscode-notebook-cell:?execution_count=15&line=52)     _ = pm.MvNormal("likelihood", mu=0, cov=Sigma_y[grp_idx], dims=('obs', 'indicators'))
     [54](vscode-notebook-cell:?execution_count=15&line=54) return sem_model_hierarchical

File ~/mambaforge/envs/pymc_examples_new/lib/python3.9/site-packages/pymc/distributions/distribution.py:310, in Distribution.__new__(cls, name, rng, dims, initval, observed, total_size, transform, *args, **kwargs)
    [307](https://file+.vscode-resource.vscode-cdn.net/Users/nathanielforde/Documents/Github/pymc-examples/examples/case_studies/~/mambaforge/envs/pymc_examples_new/lib/python3.9/site-packages/pymc/distributions/distribution.py:307)     elif observed is not None:
    [308](https://file+.vscode-resource.vscode-cdn.net/Users/nathanielforde/Documents/Github/pymc-examples/examples/case_studies/~/mambaforge/envs/pymc_examples_new/lib/python3.9/site-packages/pymc/distributions/distribution.py:308)         kwargs["shape"] = tuple(observed.shape)
--> [310](https://file+.vscode-resource.vscode-cdn.net/Users/nathanielforde/Documents/Github/pymc-examples/examples/case_studies/~/mambaforge/envs/pymc_examples_new/lib/python3.9/site-packages/pymc/distributions/distribution.py:310) rv_out = cls.dist(*args, **kwargs)
    [312](https://file+.vscode-resource.vscode-cdn.net/Users/nathanielforde/Documents/Github/pymc-examples/examples/case_studies/~/mambaforge/envs/pymc_examples_new/lib/python3.9/site-packages/pymc/distributions/distribution.py:312) rv_out = model.register_rv(
    [313](https://file+.vscode-resource.vscode-cdn.net/Users/nathanielforde/Documents/Github/pymc-examples/examples/case_studies/~/mambaforge/envs/pymc_examples_new/lib/python3.9/site-packages/pymc/distributions/distribution.py:313)     rv_out,
    [314](https://file+.vscode-resource.vscode-cdn.net/Users/nathanielforde/Documents/Github/pymc-examples/examples/case_studies/~/mambaforge/envs/pymc_examples_new/lib/python3.9/site-packages/pymc/distributions/distribution.py:314)     name,
   (...)
    [319](https://file+.vscode-resource.vscode-cdn.net/Users/nathanielforde/Documents/Github/pymc-examples/examples/case_studies/~/mambaforge/envs/pymc_examples_new/lib/python3.9/site-packages/pymc/distributions/distribution.py:319)     initval=initval,
    [320](https://file+.vscode-resource.vscode-cdn.net/Users/nathanielforde/Documents/Github/pymc-examples/examples/case_studies/~/mambaforge/envs/pymc_examples_new/lib/python3.9/site-packages/pymc/distributions/distribution.py:320) )
    [322](https://file+.vscode-resource.vscode-cdn.net/Users/nathanielforde/Documents/Github/pymc-examples/examples/case_studies/~/mambaforge/envs/pymc_examples_new/lib/python3.9/site-packages/pymc/distributions/distribution.py:322) # add in pretty-printing support

File ~/mambaforge/envs/pymc_examples_new/lib/python3.9/site-packages/pymc/distributions/multivariate.py:262, in MvNormal.dist(cls, mu, cov, tau, chol, lower, **kwargs)
    [259](https://file+.vscode-resource.vscode-cdn.net/Users/nathanielforde/Documents/Github/pymc-examples/examples/case_studies/~/mambaforge/envs/pymc_examples_new/lib/python3.9/site-packages/pymc/distributions/multivariate.py:259) @classmethod
    [260](https://file+.vscode-resource.vscode-cdn.net/Users/nathanielforde/Documents/Github/pymc-examples/examples/case_studies/~/mambaforge/envs/pymc_examples_new/lib/python3.9/site-packages/pymc/distributions/multivariate.py:260) def dist(cls, mu, cov=None, tau=None, chol=None, lower=True, **kwargs):
    [261](https://file+.vscode-resource.vscode-cdn.net/Users/nathanielforde/Documents/Github/pymc-examples/examples/case_studies/~/mambaforge/envs/pymc_examples_new/lib/python3.9/site-packages/pymc/distributions/multivariate.py:261)     mu = pt.as_tensor_variable(mu)
--> [262](https://file+.vscode-resource.vscode-cdn.net/Users/nathanielforde/Documents/Github/pymc-examples/examples/case_studies/~/mambaforge/envs/pymc_examples_new/lib/python3.9/site-packages/pymc/distributions/multivariate.py:262)     cov = quaddist_matrix(cov, chol, tau, lower)
    [263](https://file+.vscode-resource.vscode-cdn.net/Users/nathanielforde/Documents/Github/pymc-examples/examples/case_studies/~/mambaforge/envs/pymc_examples_new/lib/python3.9/site-packages/pymc/distributions/multivariate.py:263)     # PyTensor is stricter about the shape of mu, than PyMC used to be
    [264](https://file+.vscode-resource.vscode-cdn.net/Users/nathanielforde/Documents/Github/pymc-examples/examples/case_studies/~/mambaforge/envs/pymc_examples_new/lib/python3.9/site-packages/pymc/distributions/multivariate.py:264)     mu = pt.broadcast_arrays(mu, cov[..., -1])[0]

File ~/mambaforge/envs/pymc_examples_new/lib/python3.9/site-packages/pymc/distributions/multivariate.py:123, in quaddist_matrix(cov, chol, tau, lower, *args, **kwargs)
    [121](https://file+.vscode-resource.vscode-cdn.net/Users/nathanielforde/Documents/Github/pymc-examples/examples/case_studies/~/mambaforge/envs/pymc_examples_new/lib/python3.9/site-packages/pymc/distributions/multivariate.py:121)     cov = pt.as_tensor_variable(cov)
    [122](https://file+.vscode-resource.vscode-cdn.net/Users/nathanielforde/Documents/Github/pymc-examples/examples/case_studies/~/mambaforge/envs/pymc_examples_new/lib/python3.9/site-packages/pymc/distributions/multivariate.py:122)     if cov.ndim != 2:
--> [123](https://file+.vscode-resource.vscode-cdn.net/Users/nathanielforde/Documents/Github/pymc-examples/examples/case_studies/~/mambaforge/envs/pymc_examples_new/lib/python3.9/site-packages/pymc/distributions/multivariate.py:123)         raise ValueError("cov must be two dimensional.")
    [124](https://file+.vscode-resource.vscode-cdn.net/Users/nathanielforde/Documents/Github/pymc-examples/examples/case_studies/~/mambaforge/envs/pymc_examples_new/lib/python3.9/site-packages/pymc/distributions/multivariate.py:124) elif tau is not None:
    [125](https://file+.vscode-resource.vscode-cdn.net/Users/nathanielforde/Documents/Github/pymc-examples/examples/case_studies/~/mambaforge/envs/pymc_examples_new/lib/python3.9/site-packages/pymc/distributions/multivariate.py:125)     tau = pt.as_tensor_variable(tau)

ValueError: cov must be two dimensional.

and similar breakage for matmul

---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
Cell In[16], line 61
     57 priors = {"lambdas": [1, 0.5], "eta": 2, "B": [0, 0.5], "tau": [0, 1]}
     59 priors_wide = {"lambdas": [1, 5], "eta": 2, "B": [0, 5], "tau": [0, 10]}
---> 61 sem_model_hierarchical_tight = make_hierarchical(priors, grp_idx)
     62 sem_model_hierarchical_wide = make_hierarchical(priors_wide, grp_idx)
     64 pm.model_to_graphviz(sem_model_hierarchical_tight)

Cell In[16], line 42, in make_hierarchical(priors, grp_idx)
     38 tau = pm.Normal(
     39     "tau", mu=priors["tau"][0], sigma=priors["tau"][1], dims=("group", "indicators")
     40 )  # observed intercepts
     41 alpha = pm.Normal("alpha", mu=0, sigma=0.5, dims=("group", "latent"))  # latent means
---> 42 M = pt.matmul(Lambda, inv_I_minus_B)
     43 mu_latent = pt.matmul(alpha[:, None, :], M.transpose(0, 2, 1))[:, :, 0]
     44 mu_y = pm.Deterministic("mu_y", tau + mu_latent)

File ~/mambaforge/envs/pymc_examples_new/lib/python3.9/site-packages/pytensor/tensor/math.py:3010, in matmul(x1, x2, dtype)
   2968 def matmul(x1: "ArrayLike", x2: "ArrayLike", dtype: Optional["DTypeLike"] = None):
   2969     """Compute the matrix product of two tensor variables.
   2970 
   2971     Parameters
   (...)
   3008         respecting the signature ``(n, k), (k, m) -> (n, m)``:
   3009     """
-> 3010     return MatMul(dtype=dtype)(x1, x2)

File ~/mambaforge/envs/pymc_examples_new/lib/python3.9/site-packages/pytensor/graph/op.py:295, in Op.__call__(self, *inputs, **kwargs)
    253 r"""Construct an `Apply` node using :meth:`Op.make_node` and return its outputs.
    254 
    255 This method is just a wrapper around :meth:`Op.make_node`.
   (...)
    292 
    293 """
    294 return_list = kwargs.pop("return_list", False)
--> 295 node = self.make_node(*inputs, **kwargs)
    297 if config.compute_test_value != "off":
    298     compute_test_value(node)

File ~/mambaforge/envs/pymc_examples_new/lib/python3.9/site-packages/pytensor/tensor/math.py:2953, in MatMul.make_node(self, a, b)
   2950 if 0 in {a.ndim, b.ndim}:
   2951     raise ValueError("inputs to `matmul` cannot be scalar.")
-> 2953 out_shape = self._get_output_shape(
   2954     a, b, (a.type.shape, b.type.shape), validate=True
   2955 )
   2956 out = TensorType(dtype=self.dtype, shape=out_shape)()
   2957 return Apply(self, [a, b], [out])

File ~/mambaforge/envs/pymc_examples_new/lib/python3.9/site-packages/pytensor/tensor/math.py:2923, in MatMul._get_output_shape(cls, x1, x2, shapes, validate)
   2921 elif x1.ndim == 2 and x2.ndim > 2:
   2922     if validate and x1_shape[-1] != x2_shape[-2]:
-> 2923         raise ValueError(
   2924             "number of columns of input 1 must be equal "
   2925             "the length of the 2nd-last dimension of input 2"
   2926         )
   2927     return x2_shape[:-2] + x1_shape[-2:-1] + x2_shape[-1:]
   2928 else:

ValueError: number of columns of input 1 must be equal the length of the 2nd-last dimension of input 2

This works in pymc 5.17

alpha = np.random.normal(0, 1, size=(2, 4))
M = np.random.normal(0, 1, size=(2, 12, 4))
inv_I_minus_B = np.random.normal(0,1, size=(2, 4, 4))
Lambda = np.random.normal(0,1, size=(12, 4))

Sigma_y = np.random.normal(0, 1, size=(2, 12, 12))
print("Sigma_y shape", Sigma_y.shape)

print(np.matmul(Lambda, inv_I_minus_B).shape)

mu_y = np.matmul(alpha[:, None, :], M.transpose(0, 2, 1))[:, 0, :]
print("Mu_y shape", mu_y.shape)

[fonnesbeck]

I would ditch the Rhat plots -- all the action is in a tiny region just above 1.0, so most of the plot is irrelevant.