Update docs

Author: jerrymhuang

bayesflow/diagnostics/plots/pairs_posterior.py

@@ -39,25 +39,27 @@ def pairs_posterior(
         Optional true parameter values that have generated the observed dataset.
     priors       : np.ndarray of shape (n_prior_draws, n_params) or None, optional (default: None)
         Optional prior samples obtained from the prior.
-    dataset_id: Optional ID of the dataset for whose posterior the pairs plot shall be generated.
-        Should only be specified if estimates contains posterior draws from multiple datasets.
+    dataset_id: Optional ID of the dataset for whose posterior the pair plots shall be generated.
+        Should only be specified if estimates contain posterior draws from multiple datasets.
     variable_keys       : list or None, optional, default: None
        Select keys from the dictionary provided in samples.
        By default, select all keys.
     variable_names       : list or None, optional, default: None
         The parameter names for nice plot titles. Inferred if None
     height            : float, optional, default: 3
-        The height of the pairplot
+        The height of the pair plots
     label_fontsize    : int, optional, default: 14
         The font size of the x and y-label texts (parameter names)
     tick_fontsize     : int, optional, default: 12
-        The font size of the axis ticklabels
+        The font size of the axis tick labels
     legend_fontsize   : int, optional, default: 16
         The font size of the legend text
     post_color        : str, optional, default: '#132a70'
         The color for the posterior histograms and KDEs
     prior_color      : str, optional, default: gray
         The color for the optional prior histograms and KDEs
+    target_color      : str, optional, default: red
+        The color for the optional true parameter lines and points
     alpha             : float in [0, 1], optional, default: 0.9
         The opacity of the posterior plots
 
@@ -83,7 +85,7 @@ def pairs_posterior(
         variable_names=variable_names,
     )
 
-    # dicts_to_arrays will keep dataset axis even if it is of length 1
+    # dicts_to_arrays will keep the dataset axis even if it is of length 1
     # however, pairs plotting requires the dataset axis to be removed
     estimates_shape = plot_data["estimates"].shape
     if len(estimates_shape) == 3 and estimates_shape[0] == 1:

bayesflow/diagnostics/plots/pairs_samples.py

@@ -40,13 +40,18 @@ def pairs_samples(
     height      : float, optional, default: 2.5
         The height of the pair plot
     color       : str, optional, default : '#8f2727'
-        The color of the plot
+        The primary color of the plot
     alpha       : float in [0, 1], optional, default: 0.9
         The opacity of the plot
+    label       : str, optional, default: "Posterior"
+        Label for the dataset to plot
     label_fontsize    : int, optional, default: 14
         The font size of the x and y-label texts (parameter names)
     tick_fontsize     : int, optional, default: 12
-        The font size of the axis ticklabels
+        The font size of the axis tick labels
+    show_single_legend : bool, optional, default: False
+        Optional toggle for the user to choose whether a single dataset
+        should also display legend
     **kwargs    : dict, optional
         Additional keyword arguments passed to the sns.PairGrid constructor
     """
@@ -85,12 +90,20 @@ def _pairs_samples(
     show_single_legend: bool = False,
     **kwargs,
 ) -> sns.PairGrid:
-    # internal version of pairs_samples creating the seaborn plot
+    """
+    Internal version of pairs_samples creating the seaborn PairPlot
+    for both a single dataset and multiple datasets.
 
-    # Parameters
-    # ----------
-    # plot_data   : output of bayesflow.utils.dict_utils.dicts_to_arrays
-    # other arguments are documented in pairs_samples
+    Parameters
+    ----------
+    plot_data   : output of bayesflow.utils.dict_utils.dicts_to_arrays
+        Formatted data to plot from the sample dataset
+    color2      : str, optional, default: 'gray'
+        Secondary color for the pair plots.
+        This is the color used for the prior draws.
+
+    Other arguments are documented in pairs_samples
+    """
 
     estimates_shape = plot_data["estimates"].shape
     if len(estimates_shape) != 2:
@@ -144,7 +157,7 @@ def _pairs_samples(
         common_norm=False,
     )
 
-    # add scatterplots to the upper diagonal
+    # add scatter plots to the upper diagonal
     g.map_upper(sns.scatterplot, alpha=0.6, s=40, edgecolor="k", color=color, lw=0)
 
     # add KDEs to the lower diagonal
@@ -168,7 +181,7 @@ def _pairs_samples(
             g.axes[i, j].tick_params(axis="both", which="major", labelsize=tick_fontsize)
             g.axes[i, j].tick_params(axis="both", which="minor", labelsize=tick_fontsize)
 
-        # adjust font size of labels
+        # adjust the font size of labels
         # the labels themselves remain the same as before, i.e., variable_names
         g.axes[i, 0].set_ylabel(variable_names[i], fontsize=label_fontsize)
         g.axes[dim - 1, i].set_xlabel(variable_names[i], fontsize=label_fontsize)
@@ -196,7 +209,7 @@ def _pairs_samples(
 
 def histplot_twinx(x, **kwargs):
     """
-    # create a histogram plot on a twin y axis
+    # create a histogram plot on a twin y-axis
     # this ensures that the y scaling of the diagonal plots
     # in independent of the y scaling of the off-diagonal plots
 

bayesflow/utils/plot_utils.py

@@ -67,7 +67,7 @@ def prepare_plot_data(
     )
     check_estimates_prior_shapes(plot_data["estimates"], plot_data["targets"])
 
-    # store variable information at top level for easy access
+    # store variable information at the top level for easy access
     variable_names = plot_data["estimates"].variable_names
     num_variables = len(variable_names)
     plot_data["variable_names"] = variable_names
@@ -249,7 +249,7 @@ def prettify_subplots(axes: np.ndarray, num_subplots: int, tick: bool = True, ti
 
 def make_quadratic(ax: plt.Axes, x_data: np.ndarray, y_data: np.ndarray):
     """
-    Utility to make a subplots quadratic in order to avoid visual illusions
+    Utility to make subplots quadratic to avoid visual illusions
     in, e.g., recovery plots.
     """
 
@@ -269,7 +269,7 @@ def make_quadratic(ax: plt.Axes, x_data: np.ndarray, y_data: np.ndarray):
 
 def gradient_line(x, y, c=None, cmap: str = "viridis", lw: float = 2.0, alpha: float = 1, ax=None):
     """
-    Plot a 1D line with color gradient determined by `c` (same shape as x and y).
+    Plot a 1D line with a color gradient determined by `c` (same shape as x and y).
     """
     if ax is None:
         ax = plt.gca()
@@ -304,7 +304,7 @@ def gradient_legend(ax, label, cmap, norm, loc="upper right"):
     - loc: legend location (default 'upper right')
     """
 
-    # Custom dummy handle to represent the gradient
+    # Custom placeholder handle to represent the gradient
     class _GradientSwatch(Rectangle):
         pass
 
@@ -361,8 +361,35 @@ def add_gradient_plot(
 
 
 def create_legends(
-    g, plot_data, color, color2, label: str = "Posterior", show_single_legend: bool = False, fontsize: int = 14
+    g,
+    plot_data: dict,
+    color: str | tuple = "#132a70",
+    color2: str | tuple = "gray",
+    label: str = "Posterior",
+    show_single_legend: bool = False,
+    legend_fontsize: int = 14,
 ):
+    """
+    Helper function to create legends for pairplots.
+
+    Parameters
+    ----------
+    g : sns.PairGrid
+        Seaborn object for the pair plots
+    plot_data   : output of bayesflow.utils.dict_utils.dicts_to_arrays
+        Formatted data to plot from the sample dataset
+    color       : str, optional, default : '#8f2727'
+        The primary color of the plot
+    color2      : str, optional, default: 'gray'
+        The secondary color for the plot
+    label       : str, optional, default: "Posterior"
+        Label for the dataset to plot
+    show_single_legend : bool, optional, default: False
+        Optional toggle for the user to choose whether a single dataset
+        should also display legend
+    legend_fontsize    : int, optional, default: 14
+        fontsize for the legend
+    """
     handles = []
     labels = []
 
@@ -391,5 +418,5 @@ def create_legends(
             loc="center left",
             bbox_to_anchor=(1, 0.5),
             frameon=False,
-            fontsize=fontsize,
+            fontsize=legend_fontsize,
         )