Repository: arviz-devs/arviz

Issue:
Add group argument to plot pair
## Tell us about it
Add a group argument (like the one in plot density) to plot pair.

Handling on divergences in this case may require extra modifications to plotting code.



Discussion:
@OriolAbril I'd like to work on this issue. Kindly guide me through it. Can you explain more about the divergence param.
`posterior` --> `sample_stats`

`prior` --> `sample_stats_prior`
@ahartikainen Thanks.

Resolving patch:
diff --git a/arviz/plots/backends/bokeh/pairplot.py b/arviz/plots/backends/bokeh/pairplot.py
--- a/arviz/plots/backends/bokeh/pairplot.py
+++ b/arviz/plots/backends/bokeh/pairplot.py
@@ -16,7 +16,7 @@
 
 def plot_pair(
     ax,
-    _posterior,
+    infdata_group,
     numvars,
     figsize,
     textsize,
@@ -46,7 +46,7 @@ def plot_pair(
     if numvars == 2:
         (figsize, _, _, _, _, _) = _scale_fig_size(figsize, textsize, numvars - 1, numvars - 1)
 
-        source_dict = dict(zip(flat_var_names, [list(post) for post in _posterior]))
+        source_dict = dict(zip(flat_var_names, [list(post) for post in infdata_group]))
 
         if divergences:
             divergenve_name = "divergences_{}".format(str(uuid4()))
@@ -82,8 +82,8 @@ def plot_pair(
                 ax.circle(flat_var_names[0], flat_var_names[1], source=source)
         elif kind == "kde":
             plot_kde(
-                _posterior[0],
-                _posterior[1],
+                infdata_group[0],
+                infdata_group[1],
                 contour=contour,
                 fill_last=fill_last,
                 ax=ax,
@@ -92,7 +92,7 @@ def plot_pair(
                 show=False,
             )
         else:
-            ax.hexbin(_posterior[0], _posterior[1], size=0.5)
+            ax.hexbin(infdata_group[0], infdata_group[1], size=0.5)
             ax.grid.visible = False
 
         if divergences:
@@ -150,10 +150,10 @@ def plot_pair(
 
         tmp_flat_var_names = None
         if len(flat_var_names) == len(list(set(flat_var_names))):
-            source_dict = dict(zip(flat_var_names, [list(post) for post in _posterior]))
+            source_dict = dict(zip(flat_var_names, [list(post) for post in infdata_group]))
         else:
             tmp_flat_var_names = ["{}__{}".format(name, str(uuid4())) for name in flat_var_names]
-            source_dict = dict(zip(tmp_flat_var_names, [list(post) for post in _posterior]))
+            source_dict = dict(zip(tmp_flat_var_names, [list(post) for post in infdata_group]))
         if divergences:
             divergenve_name = "divergences_{}".format(str(uuid4()))
             source_dict[divergenve_name] = (
@@ -190,8 +190,8 @@ def plot_pair(
                         ax[j, i].circle(var1, var2, source=source)
 
                 elif kind == "kde":
-                    var1_kde = _posterior[i]
-                    var2_kde = _posterior[j + 1]
+                    var1_kde = infdata_group[i]
+                    var2_kde = infdata_group[j + 1]
                     plot_kde(
                         var1_kde,
                         var2_kde,
@@ -205,8 +205,8 @@ def plot_pair(
                     )
 
                 else:
-                    var1_hexbin = _posterior[i]
-                    var2_hexbin = _posterior[j + 1]
+                    var1_hexbin = infdata_group[i]
+                    var2_hexbin = infdata_group[j + 1]
                     ax[j, i].grid.visible = False
                     ax[j, i].hexbin(var1_hexbin, var2_hexbin, size=0.5)
 
diff --git a/arviz/plots/backends/matplotlib/pairplot.py b/arviz/plots/backends/matplotlib/pairplot.py
--- a/arviz/plots/backends/matplotlib/pairplot.py
+++ b/arviz/plots/backends/matplotlib/pairplot.py
@@ -14,7 +14,7 @@
 
 def plot_pair(
     ax,
-    _posterior,
+    infdata_group,
     numvars,
     figsize,
     textsize,
@@ -48,11 +48,11 @@ def plot_pair(
             fig, ax = plt.subplots(figsize=figsize, **backend_kwargs)
 
         if kind == "scatter":
-            ax.plot(_posterior[0], _posterior[1], **plot_kwargs)
+            ax.plot(infdata_group[0], infdata_group[1], **plot_kwargs)
         elif kind == "kde":
             plot_kde(
-                _posterior[0],
-                _posterior[1],
+                infdata_group[0],
+                infdata_group[1],
                 contour=contour,
                 fill_last=fill_last,
                 ax=ax,
@@ -60,7 +60,7 @@ def plot_pair(
             )
         else:
             hexbin = ax.hexbin(
-                _posterior[0], _posterior[1], mincnt=1, gridsize=gridsize, **plot_kwargs
+                infdata_group[0], infdata_group[1], mincnt=1, gridsize=gridsize, **plot_kwargs
             )
             ax.grid(False)
 
@@ -70,7 +70,9 @@ def plot_pair(
 
         if divergences:
             ax.plot(
-                _posterior[0][diverging_mask], _posterior[1][diverging_mask], **divergences_kwargs
+                infdata_group[0][diverging_mask],
+                infdata_group[1][diverging_mask],
+                **divergences_kwargs
             )
 
         ax.set_xlabel("{}".format(flat_var_names[0]), fontsize=ax_labelsize, wrap=True)
@@ -99,14 +101,14 @@ def plot_pair(
             fig, ax = plt.subplots(numvars - 1, numvars - 1, figsize=figsize, **backend_kwargs)
         hexbin_values = []
         for i in range(0, numvars - 1):
-            var1 = _posterior[i]
+            var1 = infdata_group[i]
 
             for j in range(0, numvars - 1):
                 if j < i:
                     ax[j, i].axis("off")
                     continue
 
-                var2 = _posterior[j + 1]
+                var2 = infdata_group[j + 1]
 
                 if kind == "scatter":
                     ax[j, i].plot(var1, var2, **plot_kwargs)
diff --git a/arviz/plots/pairplot.py b/arviz/plots/pairplot.py
--- a/arviz/plots/pairplot.py
+++ b/arviz/plots/pairplot.py
@@ -9,6 +9,7 @@
 
 def plot_pair(
     data,
+    group="posterior",
     var_names=None,
     coords=None,
     figsize=None,
@@ -34,6 +35,8 @@ def plot_pair(
     data : obj
         Any object that can be converted to an az.InferenceData object
         Refer to documentation of az.convert_to_dataset for details
+    group : str, optional
+        Specifies which InferenceData group should be plotted.  Defaults to 'posterior'.
     var_names : list of variable names
         Variables to be plotted, if None all variable are plotted
     coords : mapping, optional
@@ -55,7 +58,8 @@ def plot_pair(
     fill_last : bool
         If True fill the last contour of the 2D KDE plot. Defaults to True.
     divergences : Boolean
-        If True divergences will be plotted in a different color
+        If True divergences will be plotted in a different color, only if group is either 'prior'
+        or 'posterior'.
     colorbar : bool
         If True a colorbar will be included as part of the plot (Defaults to False).
         Only works when kind=hexbin
@@ -142,18 +146,27 @@ def plot_pair(
 
     # Get posterior draws and combine chains
     data = convert_to_inference_data(data)
-    posterior_data = convert_to_dataset(data, group="posterior")
-    var_names = _var_names(var_names, posterior_data)
-    flat_var_names, _posterior = xarray_to_ndarray(
-        get_coords(posterior_data, coords), var_names=var_names, combined=True
+    grouped_data = convert_to_dataset(data, group=group)
+    var_names = _var_names(var_names, grouped_data)
+    flat_var_names, infdata_group = xarray_to_ndarray(
+        get_coords(grouped_data, coords), var_names=var_names, combined=True
     )
 
     divergent_data = None
     diverging_mask = None
+
+    # Assigning divergence group based on group param
+    if group == "posterior":
+        divergent_group = "sample_stats"
+    elif group == "prior":
+        divergent_group = "sample_stats_prior"
+    else:
+        divergences = False
+
     # Get diverging draws and combine chains
     if divergences:
-        if hasattr(data, "sample_stats") and hasattr(data.sample_stats, "diverging"):
-            divergent_data = convert_to_dataset(data, group="sample_stats")
+        if hasattr(data, divergent_group) and hasattr(getattr(data, divergent_group), "diverging"):
+            divergent_data = convert_to_dataset(data, group=divergent_group)
             _, diverging_mask = xarray_to_ndarray(
                 divergent_data, var_names=("diverging",), combined=True
             )
@@ -164,12 +177,12 @@ def plot_pair(
                 "Divergences data not found, plotting without divergences. "
                 "Make sure the sample method provides divergences data and "
                 "that it is present in the `diverging` field of `sample_stats` "
-                "or set divergences=False",
+                "or `sample_stats_prior` or set divergences=False",
                 SyntaxWarning,
             )
 
     if gridsize == "auto":
-        gridsize = int(len(_posterior[0]) ** 0.35)
+        gridsize = int(len(infdata_group[0]) ** 0.35)
 
     numvars = len(flat_var_names)
 
@@ -178,7 +191,7 @@ def plot_pair(
 
     pairplot_kwargs = dict(
         ax=ax,
-        _posterior=_posterior,
+        infdata_group=infdata_group,
         numvars=numvars,
         figsize=figsize,
         textsize=textsize,


Test patch: