Repository: sbi-dev/sbi

Issue:
improve pairplot default args
quick fixes: 
- use different color for points and samples (currently it is `C0`, same as samples)
- use different 2D marker. atn it is a rather big blue dot. a red cross would be better. 
- add option for a legend when showing multiple set of samples


Discussion:

Resolving patch:
diff --git a/sbi/analysis/plot.py b/sbi/analysis/plot.py
--- a/sbi/analysis/plot.py
+++ b/sbi/analysis/plot.py
@@ -2,6 +2,7 @@
 # under the Affero General Public License v3, see <https://www.gnu.org/licenses/>.
 
 import collections
+from logging import warn
 from typing import Any, Dict, List, Optional, Tuple, Union
 
 import matplotlib as mpl
@@ -211,7 +212,10 @@ def diag_func(row, **kwargs):
             for n, v in enumerate(samples):
                 if opts["diag"][n] == "hist":
                     plt.hist(
-                        v[:, row], color=opts["samples_colors"][n], **opts["hist_diag"]
+                        v[:, row],
+                        color=opts["samples_colors"][n],
+                        label=opts["samples_labels"][n],
+                        **opts["hist_diag"],
                     )
                 elif opts["diag"][n] == "kde":
                     density = gaussian_kde(
@@ -226,7 +230,7 @@ def diag_func(row, **kwargs):
                         ys,
                         color=opts["samples_colors"][n],
                     )
-                elif "upper" in opts.keys() and opts["upper"][n] == "scatter":
+                elif "offdiag" in opts.keys() and opts["offdiag"][n] == "scatter":
                     for single_sample in v:
                         plt.axvline(
                             single_sample[row],
@@ -280,12 +284,12 @@ def pairplot(
     ] = None,
     limits: Optional[Union[List, torch.Tensor]] = None,
     subset: Optional[List[int]] = None,
-    upper: Optional[str] = "hist",
-    diag: Optional[str] = "hist",
+    offdiag: Optional[Union[List[str], str]] = "hist",
+    diag: Optional[Union[List[str], str]] = "hist",
     figsize: Tuple = (10, 10),
     labels: Optional[List[str]] = None,
     ticks: Union[List, torch.Tensor] = [],
-    points_colors: List[str] = plt.rcParams["axes.prop_cycle"].by_key()["color"],
+    upper: Optional[str] = None,
     fig=None,
     axes=None,
     **kwargs,
@@ -305,16 +309,18 @@ def pairplot(
         subset: List containing the dimensions to plot. E.g. subset=[1,3] will plot
             plot only the 1st and 3rd dimension but will discard the 0th and 2nd (and,
             if they exist, the 4th, 5th and so on).
-        upper: Plotting style for upper diagonal, {hist, scatter, contour, cond, None}.
+        offdiag: Plotting style for upper diagonal, {hist, scatter, contour, cond,
+            None}.
+        upper: deprecated, use offdiag instead.
         diag: Plotting style for diagonal, {hist, cond, None}.
         figsize: Size of the entire figure.
         labels: List of strings specifying the names of the parameters.
         ticks: Position of the ticks.
-        points_colors: Colors of the `points`.
         fig: matplotlib figure to plot on.
         axes: matplotlib axes corresponding to fig.
-        **kwargs: Additional arguments to adjust the plot, see the source code in
-            `_get_default_opts()` in `sbi.utils.plot` for more details.
+        **kwargs: Additional arguments to adjust the plot, e.g., `samples_colors`,
+            `points_colors` and many more, see the source code in `_get_default_opts()`
+            in `sbi.analysis.plot` for details.
 
     Returns: figure and axis of posterior distribution plot
     """
@@ -332,21 +338,30 @@ def pairplot(
 
     samples, dim, limits = prepare_for_plot(samples, limits)
 
+    # checks.
+    if opts["legend"]:
+        assert len(opts["samples_labels"]) >= len(
+            samples
+        ), "Provide at least as many labels as samples."
+    if opts["upper"] is not None:
+        warn("upper is deprecated, use offdiag instead.")
+        opts["offdiag"] = opts["upper"]
+
     # Prepare diag/upper/lower
     if type(opts["diag"]) is not list:
         opts["diag"] = [opts["diag"] for _ in range(len(samples))]
-    if type(opts["upper"]) is not list:
-        opts["upper"] = [opts["upper"] for _ in range(len(samples))]
+    if type(opts["offdiag"]) is not list:
+        opts["offdiag"] = [opts["offdiag"] for _ in range(len(samples))]
     # if type(opts['lower']) is not list:
     #    opts['lower'] = [opts['lower'] for _ in range(len(samples))]
     opts["lower"] = None
 
     diag_func = get_diag_func(samples, limits, opts, **kwargs)
 
-    def upper_func(row, col, limits, **kwargs):
+    def offdiag_func(row, col, limits, **kwargs):
         if len(samples) > 0:
             for n, v in enumerate(samples):
-                if opts["upper"][n] == "hist" or opts["upper"][n] == "hist2d":
+                if opts["offdiag"][n] == "hist" or opts["offdiag"][n] == "hist2d":
                     hist, xedges, yedges = np.histogram2d(
                         v[:, col],
                         v[:, row],
@@ -368,7 +383,7 @@ def upper_func(row, col, limits, **kwargs):
                         aspect="auto",
                     )
 
-                elif opts["upper"][n] in [
+                elif opts["offdiag"][n] in [
                     "kde",
                     "kde2d",
                     "contour",
@@ -393,7 +408,7 @@ def upper_func(row, col, limits, **kwargs):
                     positions = np.vstack([X.ravel(), Y.ravel()])
                     Z = np.reshape(density(positions).T, X.shape)
 
-                    if opts["upper"][n] == "kde" or opts["upper"][n] == "kde2d":
+                    if opts["offdiag"][n] == "kde" or opts["offdiag"][n] == "kde2d":
                         plt.imshow(
                             Z,
                             extent=(
@@ -405,7 +420,7 @@ def upper_func(row, col, limits, **kwargs):
                             origin="lower",
                             aspect="auto",
                         )
-                    elif opts["upper"][n] == "contour":
+                    elif opts["offdiag"][n] == "contour":
                         if opts["contour_offdiag"]["percentile"]:
                             Z = probs2contours(Z, opts["contour_offdiag"]["levels"])
                         else:
@@ -426,14 +441,14 @@ def upper_func(row, col, limits, **kwargs):
                         )
                     else:
                         pass
-                elif opts["upper"][n] == "scatter":
+                elif opts["offdiag"][n] == "scatter":
                     plt.scatter(
                         v[:, col],
                         v[:, row],
                         color=opts["samples_colors"][n],
                         **opts["scatter_offdiag"],
                     )
-                elif opts["upper"][n] == "plot":
+                elif opts["offdiag"][n] == "plot":
                     plt.plot(
                         v[:, col],
                         v[:, row],
@@ -444,7 +459,7 @@ def upper_func(row, col, limits, **kwargs):
                     pass
 
     return _arrange_plots(
-        diag_func, upper_func, dim, limits, points, opts, fig=fig, axes=axes
+        diag_func, offdiag_func, dim, limits, points, opts, fig=fig, axes=axes
     )
 
 
@@ -459,7 +474,6 @@ def marginal_plot(
     figsize: Tuple = (10, 10),
     labels: Optional[List[str]] = None,
     ticks: Union[List, torch.Tensor] = [],
-    points_colors: List[str] = plt.rcParams["axes.prop_cycle"].by_key()["color"],
     fig=None,
     axes=None,
     **kwargs,
@@ -485,8 +499,9 @@ def marginal_plot(
         points_colors: Colors of the `points`.
         fig: matplotlib figure to plot on.
         axes: matplotlib axes corresponding to fig.
-        **kwargs: Additional arguments to adjust the plot, see the source code in
-            `_get_default_opts()` in `sbi.utils.plot` for more details.
+        **kwargs: Additional arguments to adjust the plot, e.g., `samples_colors`,
+            `points_colors` and many more, see the source code in `_get_default_opts()`
+            in `sbi.analysis.plot` for details.
 
     Returns: figure and axis of posterior distribution plot
     """
@@ -523,7 +538,6 @@ def conditional_marginal_plot(
     figsize: Tuple = (10, 10),
     labels: Optional[List[str]] = None,
     ticks: Union[List, torch.Tensor] = [],
-    points_colors: List[str] = plt.rcParams["axes.prop_cycle"].by_key()["color"],
     fig=None,
     axes=None,
     **kwargs,
@@ -559,8 +573,9 @@ def conditional_marginal_plot(
 
         fig: matplotlib figure to plot on.
         axes: matplotlib axes corresponding to fig.
-        **kwargs: Additional arguments to adjust the plot, see the source code in
-            `_get_default_opts()` in `sbi.utils.plot` for more details.
+        **kwargs: Additional arguments to adjust the plot, e.g., `samples_colors`,
+            `points_colors` and many more, see the source code in `_get_default_opts()`
+            in `sbi.analysis.plot` for details.
 
     Returns: figure and axis of posterior distribution plot
     """
@@ -596,7 +611,6 @@ def conditional_pairplot(
     figsize: Tuple = (10, 10),
     labels: Optional[List[str]] = None,
     ticks: Union[List, torch.Tensor] = [],
-    points_colors: List[str] = plt.rcParams["axes.prop_cycle"].by_key()["color"],
     fig=None,
     axes=None,
     **kwargs,
@@ -634,8 +648,9 @@ def conditional_pairplot(
 
         fig: matplotlib figure to plot on.
         axes: matplotlib axes corresponding to fig.
-        **kwargs: Additional arguments to adjust the plot, see the source code in
-            `_get_default_opts()` in `sbi.utils.plot` for more details.
+        **kwargs: Additional arguments to adjust the plot, e.g., `samples_colors`,
+            `points_colors` and many more, see the source code in `_get_default_opts()`
+            in `sbi.analysis.plot` for details.
 
     Returns: figure and axis of posterior distribution plot
     """
@@ -644,7 +659,7 @@ def conditional_pairplot(
     # Setting these is required because _pairplot_scaffold will check if opts['diag'] is
     # `None`. This would break if opts has no key 'diag'. Same for 'upper'.
     diag = "cond"
-    upper = "cond"
+    offdiag = "cond"
 
     opts = _get_default_opts()
     # update the defaults dictionary by the current values of the variables (passed by
@@ -656,7 +671,7 @@ def conditional_pairplot(
     dim, limits, eps_margins = prepare_for_conditional_plot(condition, opts)
     diag_func = get_conditional_diag_func(opts, limits, eps_margins, resolution)
 
-    def upper_func(row, col, **kwargs):
+    def offdiag_func(row, col, **kwargs):
         p_image = (
             eval_conditional_density(
                 opts["density"],
@@ -675,21 +690,21 @@ def upper_func(row, col, **kwargs):
             p_image.T,
             origin="lower",
             extent=(
-                limits[col, 0],
-                limits[col, 1],
-                limits[row, 0],
-                limits[row, 1],
+                limits[col, 0].item(),
+                limits[col, 1].item(),
+                limits[row, 0].item(),
+                limits[row, 1].item(),
             ),
             aspect="auto",
         )
 
     return _arrange_plots(
-        diag_func, upper_func, dim, limits, points, opts, fig=fig, axes=axes
+        diag_func, offdiag_func, dim, limits, points, opts, fig=fig, axes=axes
     )
 
 
 def _arrange_plots(
-    diag_func, upper_func, dim, limits, points, opts, fig=None, axes=None
+    diag_func, offdiag_func, dim, limits, points, opts, fig=None, axes=None
 ):
     """
     Arranges the plots for any function that plots parameters either in a row of 1D
@@ -700,7 +715,7 @@ def _arrange_plots(
             the plot (or the columns of a row of 1D marginals). It will be passed the
             current `row` (i.e. which parameter that is to be plotted) and the `limits`
             for all dimensions.
-        upper_func: Plotting function that will be executed for the upper-diagonal
+        offdiag_func: Plotting function that will be executed for the upper-diagonal
             elements of the plot. It will be passed the current `row` and `col` (i.e.
             which parameters are to be plotted and the `limits` for all dimensions. None
             if we are in a 1D setting.
@@ -755,7 +770,7 @@ def _arrange_plots(
         else:
             raise NotImplementedError
         rows = cols = len(subset)
-    flat = upper_func is None
+    flat = offdiag_func is None
     if flat:
         rows = 1
         opts["lower"] = None
@@ -798,7 +813,7 @@ def _arrange_plots(
             elif row == col:
                 current = "diag"
             elif row < col:
-                current = "upper"
+                current = "offdiag"
             else:
                 current = "lower"
 
@@ -876,11 +891,14 @@ def _arrange_plots(
                             extent,
                             color=opts["points_colors"][n],
                             **opts["points_diag"],
+                            label=opts["points_labels"][n],
                         )
+                if opts["legend"] and col == 0:
+                    plt.legend(**opts["legend_kwargs"])
 
             # Off-diagonals
             else:
-                upper_func(
+                offdiag_func(
                     row=row,
                     col=col,
                     limits=limits,
@@ -923,22 +941,28 @@ def _arrange_plots(
 
 def _get_default_opts():
     """Return default values for plotting specs."""
-
     return {
         # 'lower': None,     # hist/scatter/None  # TODO: implement
         # title and legend
         "title": None,
         "legend": False,
+        "legend_kwargs": {},
         # labels
-        "labels_points": [],  # for points
-        "labels_samples": [],  # for samples
-        # colors
-        "samples_colors": plt.rcParams["axes.prop_cycle"].by_key()["color"],
+        "points_labels": [f"points_{idx}" for idx in range(10)],  # for points
+        "samples_labels": [f"samples_{idx}" for idx in range(10)],  # for samples
+        # colors: take even colors for samples, odd colors for points
+        "samples_colors": plt.rcParams["axes.prop_cycle"].by_key()["color"][0::2],
+        "points_colors": plt.rcParams["axes.prop_cycle"].by_key()["color"][1::2],
         # ticks
         "tickformatter": mpl.ticker.FormatStrFormatter("%g"),
         "tick_labels": None,
         # options for hist
-        "hist_diag": {"alpha": 1.0, "bins": 50, "density": False, "histtype": "step"},
+        "hist_diag": {
+            "alpha": 1.0,
+            "bins": 50,
+            "density": False,
+            "histtype": "step",
+        },
         "hist_offdiag": {
             # 'edgecolor': 'none',
             # 'linewidth': 0.0,
@@ -962,10 +986,10 @@ def _get_default_opts():
         "points_diag": {},
         "points_offdiag": {
             "marker": ".",
-            "markersize": 20,
+            "markersize": 10,
         },
         # other options
-        "fig_bg_colors": {"upper": None, "diag": None, "lower": None},
+        "fig_bg_colors": {"offdiag": None, "diag": None, "lower": None},
         "fig_subplots_adjust": {
             "top": 0.9,
         },
@@ -1124,7 +1148,10 @@ def _sbc_rank_plot(
     if params_in_subplots:
         if fig is None or ax is None:
             fig, ax = plt.subplots(
-                num_rows, min(num_parameters, num_cols), figsize=figsize, sharey=sharey
+                num_rows,
+                min(num_parameters, num_cols),
+                figsize=figsize,
+                sharey=sharey,
             )
             ax = np.atleast_1d(ax)  # type: ignore
         else:
diff --git a/sbi/analysis/tensorboard_output.py b/sbi/analysis/tensorboard_output.py
--- a/sbi/analysis/tensorboard_output.py
+++ b/sbi/analysis/tensorboard_output.py
@@ -101,7 +101,10 @@ def plot_summary(
     plot_options.update(figsize=figsize, fontsize=fontsize)
     if fig is None or axes is None:
         fig, axes = plt.subplots(
-            1, len(tags), figsize=plot_options["figsize"], **plot_options["subplots"]
+            1,
+            len(tags),
+            figsize=plot_options["figsize"],
+            **plot_options["subplots"],
         )
     axes = np.atleast_1d(axes)  # type: ignore
 
@@ -143,10 +146,10 @@ def _get_event_data_from_log_dir(
 
     Args:
         log_dir: log dir of a tensorboard summary writer.
-        size_guidance: to avoid causing out of memory erros by loading too much data at
+        size_guidance: to avoid causing out of memory errors by loading too much data at
             once into memory. Defaults to tensorboards default size_guidance.
 
-    Returns a nested, exhaustive dictionary of all event data unter log_dir.
+    Returns a nested, exhaustive dictionary of all event data under log_dir.
 
     Based on: https://stackoverflow.com/a/45899735/7770835
     """
@@ -169,8 +172,8 @@ def _get_event_data_from_log_dir(
 
                 # ScalarEvent has three attributes, wall_time, step, and value
                 # a generic way to get data from all other EventType as for ScalarEvent,
-                # we inspect their argument signature. These events are namedtuples that
-                # can be found here:
+                # we inspect their argument signature. These events are named tuples
+                # that can be found here:
                 # https://github.com/tensorflow/tensorboard/blob/b84f3738032277894c6f3fd3e011f032a89d002c/tensorboard/backend/event_processing/event_accumulator.py#L37
                 _type = type(data[0])
                 for attribute in inspect.getfullargspec(_type).args:
diff --git a/setup.py b/setup.py
--- a/setup.py
+++ b/setup.py
@@ -17,9 +17,7 @@
 # Package meta-data.
 NAME = "sbi"
 DESCRIPTION = "Simulation-based inference."
-KEYWORDS = (
-    "bayesian parameter inference system_identification simulator PyTorch"
-)
+KEYWORDS = "bayesian parameter inference system_identification simulator PyTorch"
 URL = "https://github.com/mackelab/sbi"
 EMAIL = "sbi@mackelab.org"
 AUTHOR = """Álvaro Tejero-Cantero, Jakob H. Macke, Jan-Matthis Lückmann, Conor M.
@@ -106,9 +104,7 @@ def run(self):
             pass
 
         self.status("Building Source and Wheel (universal) distribution…")
-        os.system(
-            "{0} setup.py sdist bdist_wheel --universal".format(sys.executable)
-        )
+        os.system("{0} setup.py sdist bdist_wheel --universal".format(sys.executable))
 
         self.status("Uploading the package to PyPI via Twine…")
         os.system("twine upload dist/*")
@@ -131,9 +127,7 @@ def run(self):
     author_email=EMAIL,
     python_requires=REQUIRES_PYTHON,
     url=URL,
-    packages=find_packages(
-        exclude=["tests", "*.tests", "*.tests.*", "tests.*"]
-    ),
+    packages=find_packages(exclude=["tests", "*.tests", "*.tests.*", "tests.*"]),
     install_requires=REQUIRED,
     extras_require=EXTRAS,
     include_package_data=True,


Test patch:
diff --git a/tests/linearGaussian_snpe_test.py b/tests/linearGaussian_snpe_test.py
--- a/tests/linearGaussian_snpe_test.py
+++ b/tests/linearGaussian_snpe_test.py
@@ -77,11 +77,16 @@ def test_c2st_snpe_on_linearGaussian(
     else:
         prior = utils.BoxUniform(-2.0 * ones(num_dim), 2.0 * ones(num_dim))
         target_samples = samples_true_posterior_linear_gaussian_uniform_prior(
-            x_o, likelihood_shift, likelihood_cov, prior=prior, num_samples=num_samples
+            x_o,
+            likelihood_shift,
+            likelihood_cov,
+            prior=prior,
+            num_samples=num_samples,
         )
 
     simulator, prior = prepare_for_sbi(
-        lambda theta: linear_gaussian(theta, likelihood_shift, likelihood_cov), prior
+        lambda theta: linear_gaussian(theta, likelihood_shift, likelihood_cov),
+        prior,
     )
 
     inference = snpe_method(prior, show_progress_bars=False)
@@ -106,7 +111,12 @@ def test_c2st_snpe_on_linearGaussian(
     if prior_str == "gaussian":
         # For the Gaussian prior, we compute the KLd between ground truth and posterior.
         dkl = get_dkl_gaussian_prior(
-            posterior, x_o[0], likelihood_shift, likelihood_cov, prior_mean, prior_cov
+            posterior,
+            x_o[0],
+            likelihood_shift,
+            likelihood_cov,
+            prior_mean,
+            prior_cov,
         )
 
         max_dkl = 0.15
@@ -171,7 +181,8 @@ def test_density_estimators_on_linearGaussian(density_estimtor):
     target_samples = gt_posterior.sample((num_samples,))
 
     simulator, prior = prepare_for_sbi(
-        lambda theta: linear_gaussian(theta, likelihood_shift, likelihood_cov), prior
+        lambda theta: linear_gaussian(theta, likelihood_shift, likelihood_cov),
+        prior,
     )
 
     inference = SNPE_C(prior, density_estimator=density_estimtor)
@@ -226,13 +237,18 @@ def test_c2st_snpe_on_linearGaussian_different_dims(density_estimator="maf"):
 
     simulator, prior = prepare_for_sbi(
         lambda theta: linear_gaussian(
-            theta, likelihood_shift, likelihood_cov, num_discarded_dims=discard_dims
+            theta,
+            likelihood_shift,
+            likelihood_cov,
+            num_discarded_dims=discard_dims,
         ),
         prior,
     )
     # Test whether prior can be `None`.
     inference = SNPE_C(
-        prior=None, density_estimator=density_estimator, show_progress_bars=False
+        prior=None,
+        density_estimator=density_estimator,
+        show_progress_bars=False,
     )
 
     # type: ignore
@@ -306,7 +322,8 @@ def test_c2st_multi_round_snpe_on_linearGaussian(method_str: str):
         density_estimator = "maf"
 
     simulator, prior = prepare_for_sbi(
-        lambda theta: linear_gaussian(theta, likelihood_shift, likelihood_cov), prior
+        lambda theta: linear_gaussian(theta, likelihood_shift, likelihood_cov),
+        prior,
     )
     creation_args = dict(
         prior=prior,
@@ -366,7 +383,10 @@ def test_c2st_multi_round_snpe_on_linearGaussian(method_str: str):
         ).set_default_x(x_o)
         accept_reject_fn = get_density_thresholder(posterior1, quantile=1e-4)
         proposal = RestrictedPrior(
-            prior, accept_reject_fn, posterior=posterior1, sample_with=sample_method
+            prior,
+            accept_reject_fn,
+            posterior=posterior1,
+            sample_with=sample_method,
         )
         theta = proposal.sample((1000,))
         x = simulator(theta)
@@ -412,7 +432,8 @@ def test_api_snpe_c_posterior_correction(sample_with, mcmc_method, prior_str):
         prior = utils.BoxUniform(-2.0 * ones(num_dim), 2.0 * ones(num_dim))
 
     simulator, prior = prepare_for_sbi(
-        lambda theta: linear_gaussian(theta, likelihood_shift, likelihood_cov), prior
+        lambda theta: linear_gaussian(theta, likelihood_shift, likelihood_cov),
+        prior,
     )
     inference = SNPE_C(prior, show_progress_bars=False)
 
@@ -469,7 +490,8 @@ def test_api_force_first_round_loss(
     prior = utils.BoxUniform(-2.0 * ones(num_dim), 2.0 * ones(num_dim))
 
     simulator, prior = prepare_for_sbi(
-        lambda theta: linear_gaussian(theta, likelihood_shift, likelihood_cov), prior
+        lambda theta: linear_gaussian(theta, likelihood_shift, likelihood_cov),
+        prior,
     )
     inference = SNPE_C(prior, show_progress_bars=False)
 
@@ -559,7 +581,7 @@ def simulator(theta):
         limits=[[-2, 2], [-2, 2], [-2, 2]],
         figsize=(2, 2),
         diag="kde",
-        upper="kde",
+        offdiag="kde",
     )
 
     limits = [[-2, 2], [-2, 2], [-2, 2]]
@@ -681,7 +703,8 @@ def test_example_posterior(snpe_method: type):
         extra_kwargs = dict()
 
     simulator, prior = prepare_for_sbi(
-        lambda theta: linear_gaussian(theta, likelihood_shift, likelihood_cov), prior
+        lambda theta: linear_gaussian(theta, likelihood_shift, likelihood_cov),
+        prior,
     )
     inference = snpe_method(prior, show_progress_bars=False)
 
diff --git a/tests/plot_test.py b/tests/plot_test.py
--- a/tests/plot_test.py
+++ b/tests/plot_test.py
@@ -7,13 +7,26 @@
 from matplotlib.axes import Axes
 from matplotlib.figure import Figure
 from matplotlib.pyplot import subplots
-from torch.utils.tensorboard import SummaryWriter
+from torch.utils.tensorboard.writer import SummaryWriter
 
-from sbi.analysis import plot_summary, sbc_rank_plot
+from sbi.analysis import pairplot, plot_summary, sbc_rank_plot
 from sbi.inference import SNLE, SNPE, SNRE, prepare_for_sbi, simulate_for_sbi
 from sbi.utils import BoxUniform
 
 
+@pytest.mark.parametrize("samples", (torch.randn(100, 2), [torch.randn(100, 2)] * 2))
+@pytest.mark.parametrize("labels", (None, ["a", "b"]))
+@pytest.mark.parametrize("legend", (True, False))
+@pytest.mark.parametrize("offdiag", ("hist", "scatter"))
+@pytest.mark.parametrize("samples_labels", (["a", "b"], None))
+@pytest.mark.parametrize("points_labels", (["a", "b"], None))
+@pytest.mark.parametrize("points", (None, torch.ones(2)))
+def test_pairplot(
+    samples, labels, legend, offdiag, samples_labels, points_labels, points
+):
+    pairplot(**{k: v for k, v in locals().items() if v is not None})
+
+
 @pytest.mark.parametrize("method", (SNPE, SNLE, SNRE))
 def test_plot_summary(method, tmp_path):
     num_dim = 1
@@ -68,6 +81,9 @@ def test_sbc_rank_plot(num_parameters, num_cols, custom_figure, plot_type):
     )
     if not custom_figure:
         if num_parameters > num_cols:
-            assert ax.shape == (int(np.ceil(num_parameters / num_cols)), num_cols)
+            assert ax.shape == (
+                int(np.ceil(num_parameters / num_cols)),
+                num_cols,
+            )
         else:
             assert ax.shape == (num_parameters,)