Merge branch 'dev' into allow-networks

# Conflicts:
#	bayesflow/experimental/free_form_flow/free_form_flow.py
#	bayesflow/networks/consistency_models/consistency_model.py
#	bayesflow/networks/coupling_flow/coupling_flow.py
#	bayesflow/networks/flow_matching/flow_matching.py
#	bayesflow/networks/time_series_network/skip_recurrent.py
#	bayesflow/networks/time_series_network/time_series_network.py

Author: LarsKue

README.md

@@ -37,14 +37,14 @@ Using the high-level interface is easy, as demonstrated by the minimal working e
 import bayesflow as bf
 
 workflow = bf.BasicWorkflow(
-    inference_network=bf.networks.FlowMatching(),
-    summary_network=bf.networks.TimeSeriesTransformer(),
+    inference_network=bf.networks.CouplingFlow(),
+    summary_network=bf.networks.TimeSeriesNetwork(),
     inference_variables=["parameters"],
     summary_variables=["observables"],
     simulator=bf.simulators.SIR()
 )
 
-history = workflow.fit_online(epochs=50, batch_size=32, num_batches_per_epoch=500)
+history = workflow.fit_online(epochs=15, batch_size=32, num_batches_per_epoch=200)
 
 diagnostics = workflow.plot_default_diagnostics(test_data=300)
 ```
@@ -58,6 +58,7 @@ For an in-depth exposition, check out our walkthrough notebooks below.
 5. [SIR model with custom summary network](examples/SIR_Posterior_Estimation.ipynb)
 6. [Bayesian experimental design](examples/Bayesian_Experimental_Design.ipynb)
 7. [Simple model comparison example](examples/One_Sample_TTest.ipynb)
+8. [Moving from BayesFlow v1.1 to v2.0](examples/From_BayesFlow_1.1_to_2.0.ipynb)
 
 More tutorials are always welcome! Please consider making a pull request if you have a cool application that you want to contribute.
 

bayesflow/__init__.py

@@ -7,8 +7,9 @@
     experimental,
     networks,
     simulators,
-    workflows,
     utils,
+    workflows,
+    wrappers,
 )
 
 from .adapters import Adapter

bayesflow/adapters/adapter.py

@@ -1,4 +1,4 @@
-from collections.abc import MutableSequence, Sequence, Mapping
+from collections.abc import Callable, MutableSequence, Sequence, Mapping
 
 import numpy as np
 
@@ -24,6 +24,7 @@
     NumpyTransform,
     OneHot,
     Rename,
+    SerializableCustomTransform,
     Sqrt,
     Standardize,
     ToArray,
@@ -246,16 +247,58 @@ def apply(
 
         Parameters
         ----------
-        forward: callable, no lambda
-            Function to transform the data in the forward pass.
+        forward : str or np.ufunc
+            The name of the NumPy function to use for the forward transformation.
+        inverse : str or np.ufunc, optional
+            The name of the NumPy function to use for the inverse transformation.
+            By default, the inverse is inferred from the forward argument for supported methods.
+            You can find the supported methods in
+            :py:const:`~bayesflow.adapters.transforms.NumpyTransform.INVERSE_METHODS`.
+        predicate : Predicate, optional
+            Function that indicates which variables should be transformed.
+        include : str or Sequence of str, optional
+            Names of variables to include in the transform.
+        exclude : str or Sequence of str, optional
+            Names of variables to exclude from the transform.
+        **kwargs : dict
+            Additional keyword arguments passed to the transform.
+        """
+        transform = FilterTransform(
+            transform_constructor=NumpyTransform,
+            predicate=predicate,
+            include=include,
+            exclude=exclude,
+            forward=forward,
+            inverse=inverse,
+            **kwargs,
+        )
+        self.transforms.append(transform)
+        return self
+
+    def apply_serializable(
+        self,
+        include: str | Sequence[str] = None,
+        *,
+        forward: Callable[[np.ndarray, ...], np.ndarray],
+        inverse: Callable[[np.ndarray, ...], np.ndarray],
+        predicate: Predicate = None,
+        exclude: str | Sequence[str] = None,
+        **kwargs,
+    ):
+        """Append a :py:class:`~transforms.SerializableCustomTransform` to the adapter.
+
+        Parameters
+        ----------
+        forward : function, no lambda
+            Registered serializable function to transform the data in the forward pass.
             For the adapter to be serializable, this function has to be serializable
             as well (see Notes). Therefore, only proper functions and no lambda
-            functions should be used here.
-        inverse: callable, no lambda
-            Function to transform the data in the inverse pass.
+            functions can be used here.
+        inverse : function, no lambda
+            Registered serializable function to transform the data in the inverse pass.
             For the adapter to be serializable, this function has to be serializable
             as well (see Notes). Therefore, only proper functions and no lambda
-            functions should be used here.
+            functions can be used here.
         predicate : Predicate, optional
             Function that indicates which variables should be transformed.
         include : str or Sequence of str, optional
@@ -265,14 +308,45 @@ def apply(
         **kwargs : dict
             Additional keyword arguments passed to the transform.
 
+        Raises
+        ------
+        ValueError
+            When the provided functions are not registered serializable functions.
+
         Notes
         -----
-        Important: This is only serializable if the forward and inverse functions are serializable.
-        This most likely means you will have to pass the scope that the forward and inverse functions are contained in
-        to the `custom_objects` argument of the `deserialize` function when deserializing this class.
+        Important: The forward and inverse functions have to be registered with Keras.
+        To do so, use the `@keras.saving.register_keras_serializable` decorator.
+        They must also be registered (and identical) when loading the adapter
+        at a later point in time.
+
+        Examples
+        --------
+
+        The example below shows how to use the
+        `keras.saving.register_keras_serializable` decorator to
+        register functions with Keras. Note that for this simple
+        example, one usually would use the simpler :py:meth:`apply`
+        method.
+
+        >>> import keras
+        >>>
+        >>> @keras.saving.register_keras_serializable("custom")
+        >>> def forward_fn(x):
+        >>>     return x**2
+        >>>
+        >>> @keras.saving.register_keras_serializable("custom")
+        >>> def inverse_fn(x):
+        >>>     return x**0.5
+        >>>
+        >>> adapter = bf.Adapter().apply_serializable(
+        >>>     "x",
+        >>>     forward=forward_fn,
+        >>>     inverse=inverse_fn,
+        >>> )
         """
         transform = FilterTransform(
-            transform_constructor=NumpyTransform,
+            transform_constructor=SerializableCustomTransform,
             predicate=predicate,
             include=include,
             exclude=exclude,

bayesflow/adapters/transforms/__init__.py

@@ -15,6 +15,7 @@
 from .one_hot import OneHot
 from .rename import Rename
 from .scale import Scale
+from .serializable_custom_transform import SerializableCustomTransform
 from .shift import Shift
 from .sqrt import Sqrt
 from .standardize import Standardize

bayesflow/adapters/transforms/numpy_transform.py

@@ -17,6 +17,7 @@ class NumpyTransform(ElementwiseTransform):
         The name of the NumPy function to apply in the inverse transformation.
     """
 
+    #: Dict of `np.ufunc` that support automatic selection of their inverse.
     INVERSE_METHODS = {
         np.arctan: np.tan,
         np.exp: np.log,

bayesflow/adapters/transforms/serializable_custom_transform.py

@@ -0,0 +1,183 @@
+from collections.abc import Callable
+import numpy as np
+from keras.saving import (
+    deserialize_keras_object as deserialize,
+    register_keras_serializable as serializable,
+    serialize_keras_object as serialize,
+    get_registered_name,
+    get_registered_object,
+)
+from .elementwise_transform import ElementwiseTransform
+from ...utils import filter_kwargs
+import inspect
+
+
+@serializable(package="bayesflow.adapters")
+class SerializableCustomTransform(ElementwiseTransform):
+    """
+    Transforms a parameter using a pair of registered serializable forward and inverse functions.
+
+    Parameters
+    ----------
+    forward : function, no lambda
+        Registered serializable function to transform the data in the forward pass.
+        For the adapter to be serializable, this function has to be serializable
+        as well (see Notes). Therefore, only proper functions and no lambda
+        functions can be used here.
+    inverse : function, no lambda
+        Function to transform the data in the inverse pass.
+        For the adapter to be serializable, this function has to be serializable
+        as well (see Notes). Therefore, only proper functions and no lambda
+        functions can be used here.
+
+    Raises
+    ------
+    ValueError
+        When the provided functions are not registered serializable functions.
+
+    Notes
+    -----
+    Important: The forward and inverse functions have to be registered with Keras.
+    To do so, use the `@keras.saving.register_keras_serializable` decorator.
+    They must also be registered (and identical) when loading the adapter
+    at a later point in time.
+
+    """
+
+    def __init__(
+        self,
+        *,
+        forward: Callable[[np.ndarray, ...], np.ndarray],
+        inverse: Callable[[np.ndarray, ...], np.ndarray],
+    ):
+        super().__init__()
+
+        self._check_serializable(forward, label="forward")
+        self._check_serializable(inverse, label="inverse")
+        self._forward = forward
+        self._inverse = inverse
+
+    @classmethod
+    def _check_serializable(cls, function, label=""):
+        GENERAL_EXAMPLE_CODE = (
+            "The example code below shows the structure of a correctly decorated function:\n\n"
+            "```\n"
+            "import keras\n\n"
+            "@keras.saving.register_keras_serializable('custom')\n"
+            f"def my_{label}(...):\n"
+            "    [your code goes here...]\n"
+            "```\n"
+        )
+        if function is None:
+            raise TypeError(
+                f"'{label}' must be a registered serializable function, was 'NoneType'.\n{GENERAL_EXAMPLE_CODE}"
+            )
+        registered_name = get_registered_name(function)
+        # check if function is a lambda function
+        if registered_name == "<lambda>":
+            raise ValueError(
+                f"The provided function for '{label}' is a lambda function, "
+                "which cannot be serialized. "
+                "Please provide a registered serializable function by using the "
+                "@keras.saving.register_keras_serializable decorator."
+                f"\n{GENERAL_EXAMPLE_CODE}"
+            )
+        if inspect.ismethod(function):
+            raise ValueError(
+                f"The provided value for '{label}' is a method, not a function. "
+                "Methods cannot be serialized separately from their classes. "
+                "Please provide a registered serializable function instead by "
+                "moving the functionality to a function (i.e., outside of the class) and "
+                "using the @keras.saving.register_keras_serializable decorator."
+                f"\n{GENERAL_EXAMPLE_CODE}"
+            )
+        registered_object_for_name = get_registered_object(registered_name)
+        if registered_object_for_name is None:
+            try:
+                source_max_lines = 5
+                function_source_code = inspect.getsource(function).split("\n")
+                if len(function_source_code) > source_max_lines:
+                    function_source_code = function_source_code[:source_max_lines] + ["    [...]"]
+
+                example_code = "For your provided function, this would look like this:\n\n"
+                example_code += "\n".join(
+                    ["```", "import keras\n", "@keras.saving.register_keras_serializable('custom')"]
+                    + function_source_code
+                    + ["```"]
+                )
+            except OSError:
+                example_code = GENERAL_EXAMPLE_CODE
+            raise ValueError(
+                f"The provided function for '{label}' is not registered with Keras.\n"
+                "Please register the function using the "
+                "@keras.saving.register_keras_serializable decorator.\n"
+                f"{example_code}"
+            )
+        if registered_object_for_name is not function:
+            raise ValueError(
+                f"The provided function for '{label}' does not match the function "
+                f"registered under its name '{registered_name}'. "
+                f"(registered function: {registered_object_for_name}, provided function: {function}). "
+            )
+
+    @classmethod
+    def from_config(cls, config: dict, custom_objects=None) -> "SerializableCustomTransform":
+        if get_registered_object(config["forward"]["config"], custom_objects) is None:
+            provided_function_msg = ""
+            if config["_forward_source_code"]:
+                provided_function_msg = (
+                    f"\nThe originally provided function was:\n\n```\n{config['_forward_source_code']}\n```"
+                )
+            raise TypeError(
+                "\n\nPLEASE READ HERE:\n"
+                "-----------------\n"
+                "The forward function that was provided as `forward` "
+                "is not registered with Keras, making deserialization impossible. "
+                f"Please ensure that it is registered as '{config['forward']['config']}' and identical to the original "
+                "function before loading your model."
+                f"{provided_function_msg}"
+            )
+        if get_registered_object(config["inverse"]["config"], custom_objects) is None:
+            provided_function_msg = ""
+            if config["_inverse_source_code"]:
+                provided_function_msg = (
+                    f"\nThe originally provided function was:\n\n```\n{config['_inverse_source_code']}\n```"
+                )
+            raise TypeError(
+                "\n\nPLEASE READ HERE:\n"
+                "-----------------\n"
+                "The inverse function that was provided as `inverse` "
+                "is not registered with Keras, making deserialization impossible. "
+                f"Please ensure that it is registered as '{config['inverse']['config']}' and identical to the original "
+                "function before loading your model."
+                f"{provided_function_msg}"
+            )
+        forward = deserialize(config["forward"], custom_objects)
+        inverse = deserialize(config["inverse"], custom_objects)
+        return cls(
+            forward=forward,
+            inverse=inverse,
+        )
+
+    def get_config(self) -> dict:
+        forward_source_code = inverse_source_code = None
+        try:
+            forward_source_code = inspect.getsource(self._forward)
+            inverse_source_code = inspect.getsource(self._inverse)
+        except OSError:
+            pass
+        return {
+            "forward": serialize(self._forward),
+            "inverse": serialize(self._inverse),
+            "_forward_source_code": forward_source_code,
+            "_inverse_source_code": inverse_source_code,
+        }
+
+    def forward(self, data: np.ndarray, **kwargs) -> np.ndarray:
+        # filter kwargs so that other transform args like batch_size, strict, ... are not passed through
+        kwargs = filter_kwargs(kwargs, self._forward)
+        return self._forward(data, **kwargs)
+
+    def inverse(self, data: np.ndarray, **kwargs) -> np.ndarray:
+        kwargs = filter_kwargs(kwargs, self._inverse)
+        return self._inverse(data, **kwargs)

bayesflow/diagnostics/plots/calibration_ecdf.py

@@ -19,6 +19,7 @@ def calibration_ecdf(
     figsize: Sequence[float] = None,
     label_fontsize: int = 16,
     legend_fontsize: int = 14,
+    legend_location: str = "upper right",
     title_fontsize: int = 18,
     tick_fontsize: int = 12,
     rank_ecdf_color: str = "#132a70",
@@ -184,7 +185,7 @@ def calibration_ecdf(
 
     for ax, title in zip(plot_data["axes"].flat, titles):
         ax.fill_between(z, L, U, color=fill_color, alpha=0.2, label=rf"{int((1 - alpha) * 100)}$\%$ Confidence Bands")
-        ax.legend(fontsize=legend_fontsize)
+        ax.legend(fontsize=legend_fontsize, loc=legend_location)
         ax.set_title(title, fontsize=title_fontsize)
 
     prettify_subplots(plot_data["axes"], num_subplots=plot_data["num_variables"], tick_fontsize=tick_fontsize)

bayesflow/diagnostics/plots/calibration_ecdf_from_quantiles.py

@@ -19,6 +19,7 @@ def calibration_ecdf_from_quantiles(
     figsize: Sequence[float] = None,
     label_fontsize: int = 16,
     legend_fontsize: int = 14,
+    legend_location: str = "upper right",
     title_fontsize: int = 18,
     tick_fontsize: int = 12,
     rank_ecdf_color: str = "#132a70",
@@ -173,7 +174,7 @@ def calibration_ecdf_from_quantiles(
             alpha=0.2,
             label=rf"{int((1 - alpha) * 100)}$\%$ Confidence Bands" + "\n(pointwise)",
         )
-        ax.legend(fontsize=legend_fontsize)
+        ax.legend(fontsize=legend_fontsize, loc=legend_location)
         ax.set_title(title, fontsize=title_fontsize)
 
     prettify_subplots(plot_data["axes"], num_subplots=plot_data["num_variables"], tick_fontsize=tick_fontsize)