Plot comparisons

bambi/models.py

@@ -542,7 +542,6 @@ def set_alias(self, aliases):
                     assert component_name in self.distributional_components
                     component = self.distributional_components[component_name]
                     for name, alias in component_aliases.items():
-
                         is_used = False
 
                         if name in component.terms:
@@ -660,7 +659,7 @@ def plot_priors(
         unobserved_rvs_names = []
         flat_rvs = []
         for unobserved in self.backend.model.unobserved_RVs:
-            if "Flat" in unobserved.__str__():
+            if "Flat" in str(unobserved):
                 flat_rvs.append(unobserved.name)
             else:
                 unobserved_rvs_names.append(unobserved.name)
@@ -751,7 +750,15 @@ def prior_predictive(self, draws=500, var_names=None, omit_offsets=True, random_
 
         return idata
 
-    def predict(self, idata, kind="mean", data=None, inplace=True, include_group_specific=True):
+    def predict(
+        self,
+        idata,
+        kind="mean",
+        data=None,
+        inplace=True,
+        include_group_specific=True,
+        sample_new_groups=False,
+    ):
         """Predict method for Bambi models
 
         Obtains in-sample and out-of-sample predictions from a fitted Bambi model.
@@ -769,16 +776,22 @@ def predict(self, idata, kind="mean", data=None, inplace=True, include_group_spe
         data : pandas.DataFrame or None
             An optional data frame with values for the predictors that are used to obtain
             out-of-sample predictions. If omitted, the original dataset is used.
-        include_group_specific : bool
-            If ``True`` make predictions including the group specific effects. Otherwise,
-            predictions are made with common effects only (i.e. group specific are set
-            to zero).
         inplace : bool
             If ``True`` it will modify ``idata`` in-place. Otherwise, it will return a copy of
             ``idata`` with the predictions added. If ``kind="mean"``, a new variable ending in
             ``"_mean"`` is added to the ``posterior`` group. If ``kind="pps"``, it appends a
             ``posterior_predictive`` group to ``idata``. If any of these already exist, it will be
             overwritten.
+        include_group_specific : bool
+            Determines if predictions incorporate group-specific effects. If ``False``, predictions
+            are made with common effects only (i.e. group specific are set to zero). Defaults to
+            ``True``.
+        sample_new_groups : bool
+            Specifies if it is allowed to obtain predictions for new groups of group-specific terms.
+            When ``True``, each posterior sample for the new groups is drawn from the posterior
+            draws of a randomly selected existing group. Since different groups may be selected at
+            each draw, the end result represents the variation across existing groups.
+            The method implemented is quivalent to `sample_new_levels="uncertainty"` in brms.
 
         Returns
         -------
@@ -806,7 +819,9 @@ def predict(self, idata, kind="mean", data=None, inplace=True, include_group_spe
                 else:
                     var_name = f"{response_aliased_name}_{name}"
 
-            means_dict[var_name] = component.predict(idata, data, include_group_specific, hsgp_dict)
+            means_dict[var_name] = component.predict(
+                idata, data, include_group_specific, hsgp_dict, sample_new_groups
+            )
 
             # Drop var/dim if already present. Needed for out-of-sample predictions.
             if var_name in idata.posterior.data_vars:

bambi/plots/__init__.py

@@ -1,3 +1,5 @@
-from .plot_cap import create_cap_data, plot_cap
+from bambi.plots.effects import comparisons, predictions
+from bambi.plots.plotting import plot_cap, plot_comparison
 
-__all__ = ["create_cap_data", "plot_cap"]
+
+__all__ = ["comparisons", "predictions", "plot_cap", "plot_comparison"]

bambi/plots/create_data.py

@@ -0,0 +1,129 @@
+import itertools
+
+import numpy as np
+import pandas as pd
+
+from bambi.models import Model
+from bambi.plots.utils import (
+    ConditionalInfo,
+    ContrastInfo,
+    enforce_dtypes,
+    get_covariates,
+    get_model_covariates,
+    make_group_panel_values,
+    make_main_values,
+    set_default_values,
+)
+
+
+def create_cap_data(model: Model, covariates: dict) -> pd.DataFrame:
+    """Create data for a Conditional Adjusted Predictions
+
+    Parameters
+    ----------
+    model : bambi.Model
+        An instance of a Bambi model
+    covariates : dict
+        A dictionary of length between one and three.
+        Keys must be taken from ("horizontal", "color", "panel").
+        The values indicate the names of variables.
+
+    Returns
+    -------
+    pandas.DataFrame
+        The data for the Conditional Adjusted Predictions dataframe and or
+        plotting.
+    """
+    data = model.data
+    covariates = get_covariates(covariates)
+    main, group, panel = covariates.main, covariates.group, covariates.panel
+
+    # Obtain data for main variable
+    main_values = make_main_values(data[main])
+    data_dict = {main: main_values}
+
+    # Obtain data for group and panel variables if not None
+    data_dict = make_group_panel_values(data, data_dict, main, group, panel, kind="predictions")
+    data_dict = set_default_values(model, data_dict, kind="predictions")
+    return enforce_dtypes(data, pd.DataFrame(data_dict))
+
+
+def create_comparisons_data(
+    condition: ConditionalInfo, contrast: ContrastInfo, user_passed: bool = False
+) -> pd.DataFrame:
+    """Create data for a Conditional Adjusted Comparisons
+
+    Parameters
+    ----------
+    condition: ConditionalInfo
+        A dataclass instance containing the model, contrast, and conditional
+        covariates to be used in the comparisons.
+    contrast: ContrastInfo
+        A dataclass instance containing the model, and contrast name and values.
+    user_passed: bool, optional
+        Whether the user passed their own 'conditional' data. Defaults to False.
+
+    Returns
+    -------
+    pd.DataFrame
+        The data for the Conditional Adjusted Comparisons dataframe and or
+        plotting.
+    """
+
+    def _grid_level(condition: ConditionalInfo, contrast: ContrastInfo):
+        """
+        Creates the data for grid-level contrasts by using the covariates passed
+        into the `conditional` arg. Values for the grid are either: (1) computed
+        using a equally spaced grid, mean, and or mode (depending on the covariate
+        dtype), and (2) a user specified value or range of values.
+        """
+        covariates = get_covariates(condition.covariates)
+
+        if user_passed:
+            data_dict = {**condition.conditional}
+        else:
+            main_values = make_main_values(condition.model.data[covariates.main])
+            data_dict = {covariates.main: main_values}
+            data_dict = make_group_panel_values(
+                condition.model.data,
+                data_dict,
+                covariates.main,
+                covariates.group,
+                covariates.panel,
+                kind="comparison",
+            )
+
+        data_dict[contrast.name] = contrast.values
+        comparison_data = set_default_values(condition.model, data_dict, kind="comparison")
+        # use cartesian product (cross join) to create contrasts
+        keys, values = zip(*comparison_data.items())
+        contrast_dict = [dict(zip(keys, v)) for v in itertools.product(*values)]
+
+        return enforce_dtypes(condition.model.data, pd.DataFrame(contrast_dict))
+
+    def _unit_level(contrast: ContrastInfo):
+        """
+        Creates the data for unit-level contrasts by using the observed (empirical)
+        data. All covariates in the model are included in the data, except for the
+        contrast predictor. The contrast predictor is replaced with either: (1) the
+        default contrast value, or (2) the user specified contrast value.
+        """
+        covariates = get_model_covariates(contrast.model)
+        df = contrast.model.data[covariates].drop(labels=contrast.name, axis=1)
+
+        contrast_vals = np.array(contrast.values)[..., None]
+        contrast_vals = np.repeat(contrast_vals, contrast.model.data.shape[0], axis=1)
+
+        contrast_df_dict = {}
+        for idx, value in enumerate(contrast_vals):
+            contrast_df_dict[f"contrast_{idx}"] = df.copy()
+            contrast_df_dict[f"contrast_{idx}"][contrast.name] = value
+
+        return pd.concat(contrast_df_dict.values())
+
+    if not condition.conditional:
+        df = _unit_level(contrast)
+    else:
+        df = _grid_level(condition, contrast)
+
+    return df