Merge branch 'fix-docs-link' into updated-v1-to-v2-notebook

Author: vpratz

.github/workflows/tests.yaml

@@ -7,12 +7,10 @@ on:
     branches:
       - main
       - dev
-      - update-workflows
   push:
     branches:
       - main
       - dev
-      - update-workflows
 
 defaults:
   run:
@@ -72,7 +70,13 @@ jobs:
 
       - name: Run Tests
         run: |
-          pytest -x
+          pytest -x -m "not slow"
+
+      - name: Run Slow Tests
+        # run all slow tests only on manual trigger
+        if: github.event_name == 'workflow_dispatch'
+        run: |
+          pytest -m "slow"
 
       - name: Upload test results to Codecov
         if: ${{ !cancelled() }}

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)
 ```
@@ -54,7 +54,7 @@ For an in-depth exposition, check out our walkthrough notebooks below.
 1. [Linear regression starter example](examples/Linear_Regression_Starter.ipynb)
 2. [From ABC to BayesFlow](examples/From_ABC_to_BayesFlow.ipynb)
 3. [Two moons starter example](examples/Two_Moons_Starter.ipynb)
-4. [Rapid iteration with point estimators](examples/Lotka_Volterra_point_estimation_and_expert_stats.ipynb)
+4. [Rapid iteration with point estimators](examples/Lotka_Volterra_Point_Estimation_and_Expert_Stats.ipynb)
 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)

bayesflow/__init__.py

@@ -7,8 +7,9 @@
     experimental,
     networks,
     simulators,
-    workflows,
     utils,
+    workflows,
+    wrappers,
 )
 
 from .adapters import Adapter
@@ -40,8 +41,12 @@ def setup():
         torch.autograd.set_grad_enabled(False)
 
         logging.warning(
+            "\n"
             "When using torch backend, we need to disable autograd by default to avoid excessive memory usage. Use\n"
+            "\n"
             "with torch.enable_grad():\n"
+            "    ...\n"
+            "\n"
             "in contexts where you need gradients (e.g. custom training loops)."
         )
 

bayesflow/adapters/adapter.py

@@ -1,6 +1,7 @@
-from collections.abc import MutableSequence, Sequence
+from collections.abc import MutableSequence, Sequence, Mapping
 
 import numpy as np
+
 from keras.saving import (
     deserialize_keras_object as deserialize,
     register_keras_serializable as serializable,
@@ -121,16 +122,16 @@ def inverse(self, data: dict[str, np.ndarray], **kwargs) -> dict[str, any]:
 
         return data
 
-    def __call__(self, data: dict[str, any], *, inverse: bool = False, **kwargs) -> dict[str, np.ndarray]:
+    def __call__(self, data: Mapping[str, any], *, inverse: bool = False, **kwargs) -> dict[str, np.ndarray]:
         """Apply the transforms in the given direction.
 
         Parameters
         ----------
-        data : dict
+        data : Mapping[str, any]
             The data to be transformed.
         inverse : bool, optional
             If False, apply the forward transform, else apply the inverse transform (default False).
-        **kwargs : dict
+        **kwargs
             Additional keyword arguments passed to each transform.
 
         Returns
@@ -233,28 +234,25 @@ def __len__(self):
 
     def apply(
         self,
+        include: str | Sequence[str] = None,
         *,
         forward: np.ufunc | str,
         inverse: np.ufunc | str = None,
         predicate: Predicate = None,
-        include: str | Sequence[str] = None,
         exclude: str | Sequence[str] = None,
         **kwargs,
     ):
         """Append a :py:class:`~transforms.NumpyTransform` to the adapter.
 
         Parameters
         ----------
-        forward: callable, no lambda
-            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.
-            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.
+        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
@@ -263,12 +261,6 @@ def apply(
             Names of variables to exclude from the transform.
         **kwargs : dict
             Additional keyword arguments passed to the transform.
-
-        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.
         """
         transform = FilterTransform(
             transform_constructor=NumpyTransform,
@@ -388,6 +380,7 @@ def convert_dtype(
         exclude: str | Sequence[str] = None,
     ):
         """Append a :py:class:`~transforms.ConvertDType` transform to the adapter.
+        See also :py:meth:`~bayesflow.adapters.Adapter.map_dtype`.
 
         Parameters
         ----------
@@ -525,6 +518,24 @@ def log(self, keys: str | Sequence[str], *, p1: bool = False):
         self.transforms.append(transform)
         return self
 
+    def map_dtype(self, keys: str | Sequence[str], to_dtype: str):
+        """Append a :py:class:`~transforms.ConvertDType` transform to the adapter.
+        See also :py:meth:`~bayesflow.adapters.Adapter.convert_dtype`.
+
+        Parameters
+        ----------
+        keys : str or Sequence of str
+            The names of the variables to transform.
+        to_dtype : str
+            Target dtype
+        """
+        if isinstance(keys, str):
+            keys = [keys]
+
+        transform = MapTransform({key: ConvertDType(to_dtype) for key in keys})
+        self.transforms.append(transform)
+        return self
+
     def one_hot(self, keys: str | Sequence[str], num_classes: int):
         """Append a :py:class:`~transforms.OneHot` transform to the adapter.
 
@@ -555,6 +566,24 @@ def rename(self, from_key: str, to_key: str):
         self.transforms.append(Rename(from_key, to_key))
         return self
 
+    def scale(self, keys: str | Sequence[str], by: float | np.ndarray):
+        from .transforms import Scale
+
+        if isinstance(keys, str):
+            keys = [keys]
+
+        self.transforms.append(MapTransform({key: Scale(scale=by) for key in keys}))
+        return self
+
+    def shift(self, keys: str | Sequence[str], by: float | np.ndarray):
+        from .transforms import Shift
+
+        if isinstance(keys, str):
+            keys = [keys]
+
+        self.transforms.append(MapTransform({key: Shift(shift=by) for key in keys}))
+        return self
+
     def sqrt(self, keys: str | Sequence[str]):
         """Append an :py:class:`~transforms.Sqrt` transform to the adapter.
 
@@ -572,9 +601,9 @@ def sqrt(self, keys: str | Sequence[str]):
 
     def standardize(
         self,
+        include: str | Sequence[str] = None,
         *,
         predicate: Predicate = None,
-        include: str | Sequence[str] = None,
         exclude: str | Sequence[str] = None,
         **kwargs,
     ):
@@ -603,9 +632,9 @@ def standardize(
 
     def to_array(
         self,
+        include: str | Sequence[str] = None,
         *,
         predicate: Predicate = None,
-        include: str | Sequence[str] = None,
         exclude: str | Sequence[str] = None,
         **kwargs,
     ):

bayesflow/adapters/transforms/__init__.py

@@ -14,6 +14,8 @@
 from .numpy_transform import NumpyTransform
 from .one_hot import OneHot
 from .rename import Rename
+from .scale import Scale
+from .shift import Shift
 from .sqrt import Sqrt
 from .standardize import Standardize
 from .to_array import ToArray

bayesflow/adapters/transforms/filter_transform.py

@@ -29,10 +29,10 @@ class FilterTransform(Transform):
 
     def __init__(
         self,
+        include: str | Sequence[str] = None,
         *,
         transform_constructor: Callable[..., ElementwiseTransform],
         predicate: Predicate = None,
-        include: str | Sequence[str] = None,
         exclude: str | Sequence[str] = None,
         **kwargs,
     ):

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/scale.py

@@ -0,0 +1,27 @@
+from keras.saving import (
+    deserialize_keras_object as deserialize,
+    register_keras_serializable as serializable,
+    serialize_keras_object as serialize,
+)
+import numpy as np
+
+from .elementwise_transform import ElementwiseTransform
+
+
+@serializable(package="bayesflow.adapters")
+class Scale(ElementwiseTransform):
+    def __init__(self, scale: np.typing.ArrayLike):
+        self.scale = np.array(scale)
+
+    @classmethod
+    def from_config(cls, config: dict, custom_objects=None) -> "ElementwiseTransform":
+        return cls(scale=deserialize(config["scale"]))
+
+    def get_config(self) -> dict:
+        return {"scale": serialize(self.scale)}
+
+    def forward(self, data: np.ndarray, **kwargs) -> np.ndarray:
+        return data * self.scale
+
+    def inverse(self, data: np.ndarray, **kwargs) -> np.ndarray:
+        return data / self.scale

bayesflow/adapters/transforms/shift.py

@@ -0,0 +1,27 @@
+from keras.saving import (
+    deserialize_keras_object as deserialize,
+    register_keras_serializable as serializable,
+    serialize_keras_object as serialize,
+)
+import numpy as np
+
+from .elementwise_transform import ElementwiseTransform
+
+
+@serializable(package="bayesflow.adapters")
+class Shift(ElementwiseTransform):
+    def __init__(self, shift: np.typing.ArrayLike):
+        self.shift = np.array(shift)
+
+    @classmethod
+    def from_config(cls, config: dict, custom_objects=None) -> "ElementwiseTransform":
+        return cls(shift=deserialize(config["shift"]))
+
+    def get_config(self) -> dict:
+        return {"shift": serialize(self.shift)}
+
+    def forward(self, data: np.ndarray, **kwargs) -> np.ndarray:
+        return data + self.shift
+
+    def inverse(self, data: np.ndarray, **kwargs) -> np.ndarray:
+        return data - self.shift

bayesflow/approximators/approximator.py

@@ -1,6 +1,9 @@
-import keras
+from collections.abc import Mapping
+
 import multiprocessing as mp
 
+import keras
+
 from bayesflow.adapters import Adapter
 from bayesflow.datasets import OnlineDataset
 from bayesflow.simulators import Simulator
@@ -19,7 +22,7 @@ def build_adapter(cls, **kwargs) -> Adapter:
         # implemented by each respective architecture
         raise NotImplementedError
 
-    def build_from_data(self, data: dict[str, any]) -> None:
+    def build_from_data(self, data: Mapping[str, any]) -> None:
         self.compute_metrics(**data, stage="training")
         self.built = True
 
@@ -72,7 +75,7 @@ def fit(self, *, dataset: keras.utils.PyDataset = None, simulator: Simulator = N
             A dataset containing simulations for training. If provided, `simulator` must be None.
         simulator : Simulator, optional
             A simulator used to generate a dataset. If provided, `dataset` must be None.
-        **kwargs : dict
+        **kwargs
             Additional keyword arguments passed to `keras.Model.fit()`, including (see also `build_dataset`):
 
             batch_size : int or None, default='auto'