Fix get_model_covariates() utility function (bambinos#801)

GStechschulte · Apr 14, 2024 · dd43450 · dd43450
1 parent e3a4393
commit dd43450
Show file tree

Hide file tree

Showing 5 changed files with 30 additions and 9 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -13,6 +13,7 @@
 ### Maintenance and fixes
 
 * Fix bug in predictions with models using HSGP (#780)
+* Fix `get_model_covariates()` utility function (#801)
 
 ### Documentation
 

diff --git a/bambi/interpret/utils.py b/bambi/interpret/utils.py
@@ -258,6 +258,9 @@ def get_model_covariates(model: Model) -> np.ndarray:
 
     flatten_covariates = [item for sublist in covariates for item in sublist]
 
+    # Don't include non-covariate names (#797)
+    flatten_covariates = [name for name in flatten_covariates if name in model.data]
+
     return np.unique(flatten_covariates)
 
 

diff --git a/pyproject.toml b/pyproject.toml
@@ -21,7 +21,7 @@ dependencies = [
     "formulae>=0.5.3",
     "graphviz",
     "pandas>=1.0.0",
-    "pymc>=5.12.0",
+    "pymc>=5.12.0,<5.13.0",
 ]
 
 [project.optional-dependencies]

diff --git a/tests/test_interpret.py b/tests/test_interpret.py
@@ -2,12 +2,14 @@
 This module contains tests for the helper functions of the 'interpret' sub-package.
 Tests here do not test any of the plotting functionality.
 """
+
 import numpy as np
 import pandas as pd
 import pytest
 
 import bambi as bmb
 from bambi.interpret.helpers import data_grid, select_draws
+from bambi.interpret.utils import get_model_covariates
 
 
 CHAINS = 4
@@ -190,3 +192,18 @@ def test_select_draws_no_effect(request, mtcars, condition):
         assert draws.shape == (CHAINS, DRAWS, 14)
     elif id == "3":
         assert draws.shape == (CHAINS, DRAWS, 2)
+
+
+# ------------------------------------------------------------------------------------------------ #
+#                                         Tests for utils                                          #
+# ------------------------------------------------------------------------------------------------ #
+
+
+def test_get_model_covariates():
+    """Tests `get_model_covariates()` does not include non-covariate names"""
+    # See issue 797
+    df = pd.DataFrame({"y": np.arange(10), "x": np.random.normal(size=10)})
+    knots = np.linspace(np.min(df["x"]), np.max(df["x"]), 4 + 2)[1:-1]
+    formula = "y ~ 1 + bs(x, degree=3, knots=knots)"
+    model = bmb.Model(formula, df)
+    assert set(get_model_covariates(model)) == {"x"}
diff --git a/tests/test_models.py b/tests/test_models.py
@@ -192,13 +192,13 @@ def test_cell_means_parameterization(self, crossed_data):
     def test_2_factors_saturated(self, crossed_data):
         model = bmb.Model("Y ~ threecats*fourcats", crossed_data)
         idata = self.fit(model)
-        assert list(idata.posterior.data_vars) == [
+        assert set(idata.posterior.data_vars) == {
             "Intercept",
             "threecats",
             "fourcats",
             "threecats:fourcats",
             "Y_sigma",
-        ]
+        }
         assert list(idata.posterior["threecats_dim"].values) == ["b", "c"]
         assert list(idata.posterior["fourcats_dim"].values) == ["b", "c", "d"]
         assert list(idata.posterior["threecats:fourcats_dim"].values) == [
@@ -214,12 +214,12 @@ def test_2_factors_saturated(self, crossed_data):
     def test_2_factors_no_intercept(self, crossed_data):
         model = bmb.Model("Y ~ 0 + threecats*fourcats", crossed_data)
         idata = self.fit(model)
-        assert list(idata.posterior.data_vars) == [
+        assert set(idata.posterior.data_vars) == {
             "threecats",
             "fourcats",
             "threecats:fourcats",
             "Y_sigma",
-        ]
+        }
         assert list(idata.posterior["threecats_dim"].values) == ["a", "b", "c"]
         assert list(idata.posterior["fourcats_dim"].values) == ["b", "c", "d"]
         assert list(idata.posterior["threecats:fourcats_dim"].values) == [
@@ -235,7 +235,7 @@ def test_2_factors_no_intercept(self, crossed_data):
     def test_2_factors_cell_means(self, crossed_data):
         model = bmb.Model("Y ~ 0 + threecats:fourcats", crossed_data)
         idata = self.fit(model)
-        assert list(idata.posterior.data_vars) == ["threecats:fourcats", "Y_sigma"]
+        assert set(idata.posterior.data_vars) == {"threecats:fourcats", "Y_sigma"}
         assert list(idata.posterior["threecats:fourcats_dim"].values) == [
             "a, a",
             "a, b",
@@ -255,7 +255,7 @@ def test_2_factors_cell_means(self, crossed_data):
     def test_cell_means_with_covariate(self, crossed_data):
         model = bmb.Model("Y ~ 0 + threecats + continuous", crossed_data)
         idata = self.fit(model)
-        assert list(idata.posterior.data_vars) == ["threecats", "continuous", "Y_sigma"]
+        assert set(idata.posterior.data_vars) == {"threecats", "continuous", "Y_sigma"}
         assert list(idata.posterior["threecats_dim"].values) == ["a", "b", "c"]
         self.predict_oos(model, idata)
 
@@ -477,15 +477,15 @@ def test_group_specific_categorical_interaction(self, crossed_data):
         idata = self.fit(model)
         self.predict_oos(model, idata)
 
-        assert list(idata.posterior.data_vars) == [
+        assert set(idata.posterior.data_vars) == {
             "Intercept",
             "continuous",
             "Y_sigma",
             "1|site_sigma",
             "threecats:fourcats|site_sigma",
             "1|site",
             "threecats:fourcats|site",
-        ]
+        }
         assert list(idata.posterior["threecats:fourcats|site"].coords) == [
             "chain",
             "draw",