Handle multivariate responses with HSGP

bambi/backend/model_components.py

@@ -53,7 +53,7 @@ def build(self, pymc_backend, bmb_model):
             self.build_intercept(bmb_model)
             self.build_offsets()
             self.build_common_terms(pymc_backend, bmb_model)
-            self.build_hsgp_terms(pymc_backend)
+            self.build_hsgp_terms(bmb_model, pymc_backend)
             self.build_group_specific_terms(pymc_backend, bmb_model)
 
     def build_intercept(self, bmb_model):
@@ -109,7 +109,7 @@ def build_common_terms(self, pymc_backend, bmb_model):
             # Add term to linear predictor
             self.output += pt.dot(data, coefs)
 
-    def build_hsgp_terms(self, pymc_backend):
+    def build_hsgp_terms(self, bmb_model, pymc_backend):
         """Add HSGP (Hilbert-Space Gaussian Process approximation) terms to the PyMC model.
 
         The linear predictor 'X @ b + Z @ u' can be augmented with non-parametric HSGP terms
@@ -120,7 +120,7 @@ def build_hsgp_terms(self, pymc_backend):
             for name, values in hsgp_term.coords.items():
                 if name not in pymc_backend.model.coords:
                     pymc_backend.model.add_coords({name: values})
-            self.output += hsgp_term.build()
+            self.output += hsgp_term.build(bmb_model)
 
     def build_group_specific_terms(self, pymc_backend, bmb_model):
         """Add group-specific (random or varying) terms to the PyMC model.

bambi/backend/terms.py

@@ -12,7 +12,7 @@
     make_weighted_distribution,
     GP_KERNELS,
 )
-from bambi.families.multivariate import MultivariateFamily, Multinomial, DirichletMultinomial
+from bambi.families.multivariate import MultivariateFamily
 from bambi.families.univariate import Categorical, Cumulative, StoppingRatio
 from bambi.priors import Prior
 
@@ -234,22 +234,16 @@ def build(self, pymc_backend, bmb_model):
         # Auxiliary parameters and data
         kwargs = {"observed": data, "dims": ("__obs__",)}
 
-        if isinstance(
-            self.family,
-            (
-                MultivariateFamily,
-                Categorical,
-                Cumulative,
-                StoppingRatio,
-                Multinomial,
-                DirichletMultinomial,
-            ),
-        ):
+        if isinstance(self.family, (MultivariateFamily, Categorical, Cumulative, StoppingRatio)):
             response_term = bmb_model.response_component.term
             response_name = response_term.alias or response_term.name
             dim_name = response_name + "_dim"
             pymc_backend.model.add_coords({dim_name: response_term.levels})
             dims = ("__obs__", dim_name)
+
+            # For multivariate families, the outcome variable has two dimensions too.
+            if isinstance(self.family, MultivariateFamily):
+                kwargs["dims"] = dims
         else:
             dims = ("__obs__",)
 
@@ -447,7 +441,7 @@ def __init__(self, term):
             if self.term.by_levels is not None:
                 self.coords[f"{self.term.alias}_by"] = self.coords.pop(f"{self.term.name}_by")
 
-    def build(self):
+    def build(self, spec):
         # Get the name of the term
         label = self.name
 
@@ -507,6 +501,19 @@ def build(self):
             phi = phi.eval()
 
         # Build weights coefficient
+        # Handle the case where the outcome is multivariate
+        if isinstance(spec.family, (MultivariateFamily, Categorical)):
+            # Append the dims of the response variables to the coefficient and contribution dims
+            # In general:
+            # coeff_dims: ('weights_dim', ) -> ('weights_dim', f'{response}_dim')
+            # contribution_dims: ('__obs__', ) -> ('__obs__', f'{response}_dim')
+            response_dims = tuple(spec.response_component.term.coords)
+            coeff_dims = coeff_dims + response_dims
+            contribution_dims = contribution_dims + response_dims
+
+            # Append a dimension to sqrt_psd: ('weights_dim', ) -> ('weights_dim', 1)
+            sqrt_psd = sqrt_psd[:, np.newaxis]
+
         if self.term.centered:
             coeffs = pm.Normal(f"{label}_weights", sigma=sqrt_psd, dims=coeff_dims)
         else:

bambi/families/family.py

@@ -158,6 +158,9 @@ def posterior_predictive(self, model, posterior, **kwargs):
                 pm.Truncated.dist(response_dist.dist(**kwargs), lower=lower, upper=upper)
             )
         else:
+            # for k, v in kwargs.items():
+            #     print(k, v.shape)
+            # kwargs["n"] = np.ones(shape=20).astype(int)
             output_array = pm.draw(response_dist.dist(**kwargs))
 
         return xr.DataArray(output_array, coords=coords)