Merge remote-tracking branch 'upstream/dev' into standardize_in_approx

Author: vpratz

bayesflow/approximators/__init__.py

@@ -1,4 +1,4 @@
-"""
+r"""
 A collection of :py:class:`~bayesflow.approximators.Approximator`\ s, which embody the inference task and the
 neural network components used to perform it.
 """

bayesflow/datasets/__init__.py

@@ -1,4 +1,4 @@
-"""
+r"""
 A collection of `keras.utils.PyDataset <https://keras.io/api/utils/python_utils/#pydataset-class>`__\ s, which
 wrap your data-generating process (i.e., your :py:class:`~bayesflow.simulators.Simulator`) and thus determine the
 effective training strategy (e.g., online or offline).

bayesflow/diagnostics/plots/pairs_posterior.py

@@ -7,6 +7,7 @@
 import seaborn as sns
 
 from bayesflow.utils.dict_utils import dicts_to_arrays
+from bayesflow.utils.plot_utils import create_legends
 
 from .pairs_samples import _pairs_samples
 
@@ -21,6 +22,7 @@ def pairs_posterior(
     height: int = 3,
     post_color: str | tuple = "#132a70",
     prior_color: str | tuple = "gray",
+    target_color: str | tuple = "red",
     alpha: float = 0.9,
     label_fontsize: int = 14,
     tick_fontsize: int = 12,
@@ -37,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
 
@@ -81,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:
@@ -109,14 +113,30 @@ def pairs_posterior(
         # Create DataFrame with variable names as columns
         g.data = pd.DataFrame(targets, columns=targets.variable_names)
         g.data["_source"] = "True Parameter"
-        g.map_diag(plot_true_params)
+        g.map_diag(plot_true_params_as_lines, color=target_color)
+        g.map_offdiag(plot_true_params_as_points, color=target_color)
+
+        create_legends(
+            g,
+            plot_data,
+            color=post_color,
+            color2=prior_color,
+            legend_fontsize=legend_fontsize,
+            show_single_legend=False,
+        )
 
     return g
 
 
-def plot_true_params(x, hue=None, **kwargs):
-    """Custom function to plot true parameters on the diagonal."""
+def plot_true_params_as_lines(x, hue=None, color=None, **kwargs):
+    """Custom function to plot true parameters on the diagonal as dashed lines."""
     # hue needs to be added to handle the case of plotting both posterior and prior
     param = x.iloc[0]  # Get the single true value for the diagonal
     # only plot on the diagonal a vertical line for the true parameter
-    plt.axvline(param, color="black", linestyle="--")
+    plt.axvline(param, color=color, linestyle="--")
+
+
+def plot_true_params_as_points(x, y, color=None, marker="x", **kwargs):
+    """Custom function to plot true parameters on the off-diagonal as a single point."""
+    if len(x) > 0 and len(y) > 0:
+        plt.scatter(x.iloc[0], y.iloc[0], color=color, marker=marker, **kwargs)

bayesflow/diagnostics/plots/pairs_samples.py

@@ -8,6 +8,7 @@
 
 from bayesflow.utils import logging
 from bayesflow.utils.dict_utils import dicts_to_arrays
+from bayesflow.utils.plot_utils import create_legends
 
 
 def pairs_samples(
@@ -17,8 +18,10 @@ def pairs_samples(
     height: float = 2.5,
     color: str | tuple = "#132a70",
     alpha: float = 0.9,
+    label: str = "Posterior",
     label_fontsize: int = 14,
     tick_fontsize: int = 12,
+    show_single_legend: bool = False,
     **kwargs,
 ) -> sns.PairGrid:
     """
@@ -37,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
     """
@@ -59,8 +67,11 @@ def pairs_samples(
         height=height,
         color=color,
         alpha=alpha,
+        label=label,
         label_fontsize=label_fontsize,
         tick_fontsize=tick_fontsize,
+        show_single_legend=show_single_legend,
+        **kwargs,
     )
 
     return g
@@ -72,17 +83,27 @@ def _pairs_samples(
     color: str | tuple = "#132a70",
     color2: str | tuple = "gray",
     alpha: float = 0.9,
+    label: str = "Posterior",
     label_fontsize: int = 14,
     tick_fontsize: int = 12,
     legend_fontsize: int = 14,
+    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:
@@ -136,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
@@ -146,11 +167,6 @@ def _pairs_samples(
         logging.exception("KDE failed due to the following exception:\n" + repr(e) + "\nSubstituting scatter plot.")
         g.map_lower(sns.scatterplot, alpha=0.6, s=40, edgecolor="k", color=color, lw=0)
 
-    # need to add legend here such that colors are recognized
-    if plot_data["priors"] is not None:
-        g.add_legend(fontsize=legend_fontsize, loc="center right")
-        g._legend.set_title(None)
-
     # Generate grids
     dim = g.axes.shape[0]
     for i in range(dim):
@@ -165,32 +181,48 @@ 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)
 
+    # need to add legend here such that colors are recognized
+    # if plot_data["priors"] is not None:
+    #     g.add_legend(fontsize=legend_fontsize, loc="center right")
+    #     g._legend.set_title(None)
+
+    create_legends(
+        g,
+        plot_data,
+        color=color,
+        color2=color2,
+        legend_fontsize=legend_fontsize,
+        label=label,
+        show_single_legend=show_single_legend,
+    )
+
     # Return figure
     g.tight_layout()
 
     return g
 
 
-# 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
 def histplot_twinx(x, **kwargs):
-    # Create a twin axis
-    ax2 = plt.gca().twinx()
+    """
+    # 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
 
+    Parameters
+    ----------
+    x : np.ndarray
+        Data to be plotted.
+    """
     # create a histogram on the twin axis
-    sns.histplot(x, **kwargs, ax=ax2)
+    sns.histplot(x, legend=False, **kwargs)
 
     # make the twin axis invisible
     plt.gca().spines["right"].set_visible(False)
     plt.gca().spines["top"].set_visible(False)
-    ax2.set_ylabel("")
-    ax2.set_yticks([])
-    ax2.set_yticklabels([])
 
     return None

bayesflow/distributions/__init__.py

@@ -1,4 +1,4 @@
-"""
+r"""
 A collection of :py:class:`~bayesflow.distributions.Distribution`\ s,
 which represent the latent space for :py:class:`~bayesflow.networks.InferenceNetwork`\ s
 or the summary space of :py:class:`~bayesflow.networks.SummaryNetwork`\ s.

bayesflow/experimental/free_form_flow/free_form_flow.py

@@ -37,7 +37,7 @@ class FreeFormFlow(InferenceNetwork):
         "activation": "mish",
         "kernel_initializer": "he_normal",
         "residual": True,
-        "dropout": 0.05,
+        "dropout": 0.0,
         "spectral_normalization": False,
     }
 
@@ -46,7 +46,7 @@ class FreeFormFlow(InferenceNetwork):
         "activation": "mish",
         "kernel_initializer": "he_normal",
         "residual": True,
-        "dropout": 0.05,
+        "dropout": 0.0,
         "spectral_normalization": False,
     }
 
@@ -219,7 +219,7 @@ def decode(z):
 
         # VJP computation
         z, vjp_fn = vjp(encode, x, return_output=True)
-        v1 = vjp_fn(v)[0]
+        v1 = vjp_fn(v)
         # JVP computation
         x_pred, v2 = jvp(decode, (z,), (v,), return_output=True)
 

bayesflow/metrics/__init__.py

@@ -1,4 +1,4 @@
-"""
+r"""
 A collection of `keras.Metric <https://keras.io/api/metrics/base_metric/#metric-class>`__\ s for evaluating the
 performance of models.
 """

bayesflow/networks/__init__.py

@@ -1,4 +1,4 @@
-"""
+r"""
 A rich collection of neural network architectures for use in :py:class:`~bayesflow.approximators.Approximator`\ s.
 
 The module features inference networks (IN), summary networks (SN), as well as general purpose networks.

bayesflow/scores/scoring_rule.py

@@ -165,7 +165,7 @@ def get_head(self, key: str, output_shape: Shape) -> keras.Sequential:
         return keras.Sequential([subnet, dense, reshape, link])
 
     def score(self, estimates: dict[str, Tensor], targets: Tensor, weights: Tensor) -> Tensor:
-        """Scores a batch of probabilistic estimates of distributions based on samples
+        r"""Scores a batch of probabilistic estimates of distributions based on samples
         of the corresponding distributions.
 
         Parameters