Allow families to implement a custom create_extra_pps_coord() (#688)

* add create_extra_pps_coord method

* undo change

* remove argument

* proper hasattr

* add another print

* change name

* Remove prints

* Update changelog

* Fix test. Update to pytensor >= 2.12.3 and pymc>= 5.5.0

bambi/families/family.py

@@ -186,8 +186,13 @@ def posterior_predictive(self, model, posterior, **kwargs):
 
         coord_names = ["chain", "draw", response_aliased_name + "_obs"]
         is_multivariate = hasattr(model.family, "KIND") and model.family.KIND == "Multivariate"
+
         if is_multivariate:
             coord_names.append(response_aliased_name + "_dim")
+        elif hasattr(model.family, "create_extra_pps_coord"):
+            new_coords = model.family.create_extra_pps_coord()
+            coord_names.append(response_aliased_name + "_dim")
+            output_coords_all[response_aliased_name + "_dim"] = new_coords
 
         output_coords = {}
         for coord_name in coord_names:

docs/CHANGELOG.md

@@ -5,6 +5,7 @@
 ### New features
 
 * Implement new families `"ordinal"` and `"sratio"` for modeling of ordinal responses (#678)
+* Allow families to implement a custom `create_extra_pps_coord()` (#688)
 
 ### Maintenance and fixes
 

pyproject.toml

@@ -20,8 +20,8 @@ dependencies = [
     "graphviz",
     "numpy>1.22",
     "pandas>=1.0.0",
-    "pymc>=5.4.0",
-    "pytensor>=2.12.1",
+    "pymc>=5.5.0",
+    "pytensor>=2.12.3",
     "scipy>=1.7.0",
 ]
 

tests/test_aliases.py

@@ -96,4 +96,4 @@ def test_alias_equal_to_name(my_data):
     model = bmb.Model("y ~ 1 + x", my_data)
     model.set_alias({"sigma": "sigma"})
     idata = model.fit(tune=100, draws=100)
-    set(idata.posterior.data_vars) == {"Intercept", "y_mean", "x", "sigma"}
+    set(idata.posterior.data_vars) == {"Intercept", "y_mean", "x", "sigma"}

tests/test_built_models.py

@@ -646,13 +646,8 @@ def test_potentials():
 
     pot0 = model.backend.model.potentials[0].get_parents()[0]
     pot1 = model.backend.model.potentials[1].get_parents()[0]
-    assert pot0.__str__() == (
-        "Elemwise{switch,no_inplace}(Elemwise{lt,no_inplace}.0, " "Intercept, TensorConstant{-inf})"
-    )
-    assert pot1.__str__() == (
-        "Elemwise{switch,no_inplace}(Elemwise{gt,no_inplace}.0, "
-        "TensorConstant{0}, TensorConstant{-inf})"
-    )
+    assert pot0.__str__() == "Switch(Lt.0, Intercept, -inf)"
+    assert pot1.__str__() == "Switch(Gt.0, 0, -inf)"
 
 
 def test_binomial_regression():
@@ -688,7 +683,9 @@ def test_categorical_family(inhaler):
 
 
 def test_categorical_family_varying_intercept(inhaler):
-    model = bmb.Model("rating ~ period + carry + treat + (1|subject)", inhaler, family="categorical")
+    model = bmb.Model(
+        "rating ~ period + carry + treat + (1|subject)", inhaler, family="categorical"
+    )
     model.fit(draws=10, tune=10)
 
 
@@ -888,7 +885,7 @@ def test_zero_inflated_negativebinomial():
 
 
 def test_hurlde_families():
-    df = pd.DataFrame({"y": pm.draw(pm.HurdlePoisson.dist(0.5, mu=3.5), 1000)}) 
+    df = pd.DataFrame({"y": pm.draw(pm.HurdlePoisson.dist(0.5, mu=3.5), 1000)})
     model = bmb.Model("y ~ 1", df, family="hurdle_poisson")
     idata = model.fit()
     model.predict(idata, kind="pps")
@@ -908,6 +905,7 @@ def test_hurlde_families():
     idata = model.fit()
     model.predict(idata, kind="pps")
 
+
 @pytest.mark.parametrize(
     "family, link",
     [
@@ -917,11 +915,11 @@ def test_hurlde_families():
         ("sratio", "logit"),
         ("sratio", "probit"),
         ("sratio", "cloglog"),
-    ]
+    ],
 )
 def test_ordinal_families(inhaler, family, link):
     data = inhaler.copy()
-    data["carry"] = pd.Categorical(data["carry"]) # To have both numeric and categoric predictors
+    data["carry"] = pd.Categorical(data["carry"])  # To have both numeric and categoric predictors
     model = bmb.Model("rating ~ period + carry + treat", data, family=family, link=link)
     idata = model.fit(tune=100, draws=100)
     model.predict(idata, kind="pps")
@@ -932,13 +930,13 @@ def test_ordinal_families(inhaler, family, link):
 def test_cumulative_family_priors(inhaler):
     priors = {
         "threshold": bmb.Prior(
-            "Normal", 
-            mu=[-0.5, 0, 0.5], 
-            sigma=1.5, 
-            transform=pm.distributions.transforms.univariate_ordered
+            "Normal",
+            mu=[-0.5, 0, 0.5],
+            sigma=1.5,
+            transform=pm.distributions.transforms.univariate_ordered,
         )
     }
     model = bmb.Model(
         "rating ~ period + carry + treat", inhaler, family="cumulative", priors=priors
     )
-    model.fit(tune=100, draws=100)
+    model.fit(tune=100, draws=100)

tests/test_plots.py

@@ -82,6 +82,7 @@ def test_with_panel(mtcars, pps):
     # Horizontal: categorical. Group: categorical
     plot_cap(model, idata, {"horizontal": "gear", "panel": "cyl"}, pps=pps)
 
+
 @pytest.mark.parametrize("pps", [False, True])
 def test_with_group_and_panel(mtcars, pps):
     model, idata = mtcars
@@ -124,50 +125,44 @@ def test_fig_kwargs(mtcars, pps):
 def test_use_hdi(mtcars, pps):
     model, idata = mtcars
     plot_cap(
-        model, 
-        idata, 
-        {"horizontal": "hp", "color": "cyl", "panel": "gear"}, 
-        pps=pps,
-        use_hdi=False
+        model, idata, {"horizontal": "hp", "color": "cyl", "panel": "gear"}, pps=pps, use_hdi=False
     )
 
 
 @pytest.mark.parametrize("pps", [False, True])
 def test_hdi_prob(mtcars, pps):
     model, idata = mtcars
     plot_cap(
-        model, 
-        idata,
-        {"horizontal": "hp", "color": "cyl", "panel": "gear"}, 
-        pps=pps,
-        hdi_prob=0.9
+        model, idata, {"horizontal": "hp", "color": "cyl", "panel": "gear"}, pps=pps, hdi_prob=0.9
     )
 
     with pytest.raises(
         ValueError, match="'hdi_prob' must be greater than 0 and smaller than 1. It is 1.1."
     ):
         plot_cap(
-            model, 
-            idata, 
-            {"horizontal": "hp", "color": "cyl", "panel": "gear"}, 
+            model,
+            idata,
+            {"horizontal": "hp", "color": "cyl", "panel": "gear"},
             pps=pps,
-            hdi_prob=1.1)
+            hdi_prob=1.1,
+        )
 
     with pytest.raises(
         ValueError, match="'hdi_prob' must be greater than 0 and smaller than 1. It is -0.1."
     ):
         plot_cap(
-            model, 
-            idata, 
-            {"horizontal": "hp", "color": "cyl", "panel": "gear"}, 
+            model,
+            idata,
+            {"horizontal": "hp", "color": "cyl", "panel": "gear"},
             pps=pps,
-            hdi_prob=-0.1)
+            hdi_prob=-0.1,
+        )
 
 
 @pytest.mark.parametrize("pps", [False, True])
 def test_legend(mtcars, pps):
     model, idata = mtcars
-    plot_cap(model, idata, ["hp"], pps=pps,legend=False)
+    plot_cap(model, idata, ["hp"], pps=pps, legend=False)
 
 
 @pytest.mark.parametrize("pps", [False, True])
@@ -199,6 +194,6 @@ def test_multiple_outputs(pps):
     model = bmb.Model(formula, data_gamma, family="gamma")
     idata = model.fit(tune=100, draws=100, random_seed=1234)
     # Test default target
-    plot_cap(model, idata,  "x", pps=pps)
+    plot_cap(model, idata, "x", pps=pps)
     # Test user supplied target argument
     plot_cap(model, idata, "x", "alpha", pps=pps)

tests/test_predict.py

@@ -228,7 +228,9 @@ def test_predict_categorical(inhaler):
     model.predict(idata, data=inhaler.iloc[:20, :])
     assert np.allclose(idata.posterior["rating_mean"].values.sum(-1), 1)
 
-    model = bmb.Model("rating ~ period + carry + treat + (1|subject)", inhaler, family="categorical")
+    model = bmb.Model(
+        "rating ~ period + carry + treat + (1|subject)", inhaler, family="categorical"
+    )
     idata = model.fit(tune=100, draws=100)
 
     model.predict(idata)
@@ -424,4 +426,4 @@ def test_posterior_predictive_beta_binomial():
     model.predict(idata, kind="pps")
 
     n = data["n"].to_numpy()
-    assert np.all(idata.posterior_predictive["prop(y, n)"].values <= n[np.newaxis, np.newaxis, :])
+    assert np.all(idata.posterior_predictive["prop(y, n)"].values <= n[np.newaxis, np.newaxis, :])

tests/test_priors.py

@@ -264,17 +264,23 @@ def test_set_response_prior():
     priors = {"sigma": bmb.Prior("Uniform", lower=0, upper=50)}
     model = bmb.Model("y ~ x", data)
     model.set_priors(priors)
-    assert model.constant_components["sigma"].prior == bmb.Prior("Uniform", False, lower=0, upper=50)
+    assert model.constant_components["sigma"].prior == bmb.Prior(
+        "Uniform", False, lower=0, upper=50
+    )
 
     priors = {"alpha": bmb.Prior("Uniform", lower=1, upper=20)}
     model = bmb.Model("y ~ x", data, family="negativebinomial")
     model.set_priors(priors)
-    assert model.constant_components["alpha"].prior == bmb.Prior("Uniform", False, lower=1, upper=20)
+    assert model.constant_components["alpha"].prior == bmb.Prior(
+        "Uniform", False, lower=1, upper=20
+    )
 
     priors = {"alpha": bmb.Prior("Uniform", lower=0, upper=50)}
     model = bmb.Model("y ~ x", data, family="gamma")
     model.set_priors(priors)
-    assert model.constant_components["alpha"].prior == bmb.Prior("Uniform", False, lower=0, upper=50)
+    assert model.constant_components["alpha"].prior == bmb.Prior(
+        "Uniform", False, lower=0, upper=50
+    )
 
 
 def test_prior_shape():

tests/test_utils.py

@@ -8,6 +8,7 @@
 from bambi.backend.pymc import probit, cloglog
 from bambi.transformations import censored
 
+
 def test_listify():
     assert listify(None) == []
     assert listify([1, 2, 3]) == [1, 2, 3]