merge upstream/dev into doc-autosummary

Author: vpratz

.github/workflows/tests.yaml

@@ -46,6 +46,7 @@ jobs:
   test:
     runs-on: ${{ matrix.os }}
     strategy:
+      fail-fast: false
       matrix:
         os: [ubuntu-latest, windows-latest]
         python-version: ["3.10", "3.11"]

README.md

@@ -16,14 +16,18 @@ fueled by continuous progress in generative AI and Bayesian inference.
 
 ## Conceptual Overview
 
+<div align="center">
+<picture>
+  <source media="(prefers-color-scheme: dark)" srcset="./img/bayesflow_landing_dark.jpg">
+  <source media="(prefers-color-scheme: light)" srcset="./img/bayesflow_landing_light.jpg">
+  <img alt="dsd" src="./img/bayesflow_landing_dark.jpg">
+</picture>
+</div>
+
 A cornerstone idea of amortized Bayesian inference is to employ generative
 neural networks for parameter estimation, model comparison, and model validation
 when working with intractable simulators whose behavior as a whole is too
-complex to be described analytically. The figure below presents a higher-level
-overview of neurally bootstrapped Bayesian inference.
-
-<img src="https://github.com/bayesflow-org/bayesflow/blob/master/img/high_level_framework.png?raw=true" width=80% height=80%>
-
+complex to be described analytically.
 
 ## Disclaimer
 
@@ -91,9 +95,10 @@ Check out some of our walk-through notebooks below. We are actively working on p
 
 1. [Two moons starter toy example](examples/TwoMoons_StarterNotebook.ipynb)
 2. [Linear regression](examples/Linear_Regression.ipynb)
-3. [Hyperparameter optimization](examples/Hyperparameter_Optimization.ipynb)
-4. [Bayesian experimental design](examples/Bayesian_Experimental_Design.ipynb)
-5. Coming soon...
+3. [Bayesian experimental design](examples/Bayesian_Experimental_Design.ipynb)
+4. [SIR model with custom summary network](examples/SIR_PosteriorEstimation.ipynb)
+5. [Hyperparameter optimization](examples/Hyperparameter_Optimization.ipynb)
+6. Coming soon...
 
 ## Documentation \& Help
 

bayesflow/__init__.py

@@ -7,9 +7,11 @@
     distributions,
     networks,
     simulators,
+    workflows,
     utils,
 )
 
+from .workflows import BasicWorkflow
 from .approximators import ContinuousApproximator
 from .adapters import Adapter
 from .datasets import OfflineDataset, OnlineDataset, DiskDataset

bayesflow/adapters/adapter.py

@@ -9,11 +9,13 @@
 
 from .transforms import (
     AsSet,
+    AsTimeSeries,
     Broadcast,
     Concatenate,
     Constrain,
     ConvertDType,
     Drop,
+    ExpandDims,
     FilterTransform,
     Keep,
     LambdaTransform,
@@ -111,24 +113,33 @@ def as_set(self, keys: str | Sequence[str]):
         self.transforms.append(transform)
         return self
 
-    def broadcast(self, keys: str | Sequence[str], *, expand_scalars: bool = True):
+    def as_time_series(self, keys: str | Sequence[str]):
         if isinstance(keys, str):
             keys = [keys]
 
-        transform = MapTransform({key: Broadcast(expand_scalars=expand_scalars) for key in keys})
+        transform = MapTransform({key: AsTimeSeries() for key in keys})
+        self.transforms.append(transform)
+        return self
+
+    def broadcast(
+        self, keys: str | Sequence[str], *, to: str, expand: str | int | tuple = "left", exclude: int | tuple = -1
+    ):
+        if isinstance(keys, str):
+            keys = [keys]
+
+        transform = Broadcast(keys, to=to, expand=expand, exclude=exclude)
         self.transforms.append(transform)
         return self
 
     def clear(self):
         self.transforms = []
         return self
 
-    def concatenate(self, keys: Sequence[str], *, into: str, axis: int = -1):
+    def concatenate(self, keys: str | Sequence[str], *, into: str, axis: int = -1):
         if isinstance(keys, str):
-            # this is a common mistake, and also passes the type checker since str is a sequence of characters
-            raise ValueError("Keys must be a sequence of strings. To rename a single key, use the `rename` method.")
-
-        transform = Concatenate(keys, into=into, axis=axis)
+            transform = Rename(keys, to_key=into)
+        else:
+            transform = Concatenate(keys, into=into, axis=axis)
         self.transforms.append(transform)
         return self
 
@@ -177,6 +188,14 @@ def drop(self, keys: str | Sequence[str]):
         self.transforms.append(transform)
         return self
 
+    def expand_dims(self, keys: str | Sequence[str], *, axis: int | tuple):
+        if isinstance(keys, str):
+            keys = [keys]
+
+        transform = ExpandDims(keys, axis=axis)
+        self.transforms.append(transform)
+        return self
+
     def keep(self, keys: str | Sequence[str]):
         if isinstance(keys, str):
             keys = [keys]

bayesflow/adapters/transforms/__init__.py

@@ -1,10 +1,12 @@
 from .as_set import AsSet
+from .as_time_series import AsTimeSeries
 from .broadcast import Broadcast
 from .concatenate import Concatenate
 from .constrain import Constrain
 from .convert_dtype import ConvertDType
 from .drop import Drop
 from .elementwise_transform import ElementwiseTransform
+from .expand_dims import ExpandDims
 from .filter_transform import FilterTransform
 from .keep import Keep
 from .lambda_transform import LambdaTransform

bayesflow/adapters/transforms/as_set.py

@@ -4,6 +4,27 @@
 
 
 class AsSet(ElementwiseTransform):
+    """
+    The `.as_set(["x", "y"])` transform indicates that both `x` and `y` are treated as sets.
+    That is, their values will be treated as *exchangable* such that they will imply
+    the same inference regardless of the values' order.
+    This is useful, for example, in a linear regression context where we can index
+    the observations in arbitrary order and always get the same regression line.
+
+    Currently, all this transform does is to ensure that the variable
+    arrays are at least 3D. The 2rd dimension is treated as the
+    set dimension and the 3rd dimension as the data dimension.
+    In the future, the transform will have more advanced behavior
+    to better ensure the correct treatment of sets.
+
+    Useage:
+
+    adapter = (
+        bf.Adapter()
+        .as_set(["x", "y"])
+        )
+    """
+
     def forward(self, data: np.ndarray, **kwargs) -> np.ndarray:
         return np.atleast_3d(data)
 

bayesflow/adapters/transforms/as_time_series.py

@@ -0,0 +1,32 @@
+import numpy as np
+
+from .elementwise_transform import ElementwiseTransform
+
+
+class AsTimeSeries(ElementwiseTransform):
+    """
+    The `.as_time_series` transform can be used to indicate that
+    variables shall be treated as time series.
+
+    Currently, all this transformation does is to ensure that the variable
+    arrays are at least 3D. The 2rd dimension is treated as the
+    time series dimension and the 3rd dimension as the data dimension.
+    In the future, the transform will have more advanced behavior
+    to better ensure the correct treatment of time series data.
+
+    Useage:
+
+    adapter = (
+        bf.Adapter()
+        .as_time_series(["x", "y"])
+        )
+    """
+
+    def forward(self, data: np.ndarray, **kwargs) -> np.ndarray:
+        return np.atleast_3d(data)
+
+    def inverse(self, data: np.ndarray, **kwargs) -> np.ndarray:
+        if data.shape[2] == 1:
+            return np.squeeze(data, axis=2)
+
+        return data

bayesflow/adapters/transforms/broadcast.py

@@ -1,51 +1,125 @@
+from collections.abc import Sequence
 import numpy as np
 
-from .elementwise_transform import ElementwiseTransform
+from keras.saving import (
+    deserialize_keras_object as deserialize,
+    register_keras_serializable as serializable,
+    serialize_keras_object as serialize,
+)
 
+from .transform import Transform
 
-class Broadcast(ElementwiseTransform):
+
+@serializable(package="bayesflow.adapters")
+class Broadcast(Transform):
     """
-    Broadcasts array to a given batch size.
+    Broadcasts arrays or scalars to the shape of a given other array.
+
+    Parameters:
+
+    expand: Where should new dimensions be added to match the number of dimensions in `to`?
+    Can be "left", "right", or an integer or tuple containing the indices of the new dimensions.
+    The latter is needed if we want to include a dimension in the middle, which will be required
+    for more advanced cases. By default we expand left.
+
+    exclude: Which dimensions (of the dimensions after expansion) should retain their size,
+    rather than being broadcasted to the corresponding dimension size of `to`?
+    By default we exclude the last dimension (usually the data dimension) from broadcasting the size.
+
     Examples:
-        >>> bc = Broadcast()
-        >>> bc(np.array(5), batch_size=3)
-        array([[5], [5], [5]])
-        >>> bc(np.array(5), batch_size=3).shape
-        (3, 1)
-        >>> bc(np.array([1, 2, 3]), batch_size=3)
-        array([[1, 2, 3], [1, 2, 3], [1, 2, 3]])
-        >>> bc(np.array([1, 2, 3]), batch_size=3).shape
-        (3, 3)
-
-        You can opt out of expanding scalars:
-        >>> bc = Broadcast(expand_scalars=False)
-        >>> bc(np.array(5), batch_size=3)
-        np.array([5, 5, 5])
-        >>> bc(np.array(5), batch_size=3).shape
-        (3,)
+        shape (1, ) array:
+        >>> a = np.array((1,))
+        shape (2, 3) array:
+        >>> b = np.array([[1, 2, 3], [4, 5, 6]])
+        shape (2, 2, 3) array:
+        >>> c = np.array([[[1, 2, 3], [4, 5, 6]], [[4, 5, 6], [1, 2, 3]]])
+        >>> dat = dict(a=a, b=b, c=c)
+
+        >>> bc = bf.adapters.transforms.Broadcast("a", to="b")
+        >>> new_dat = bc.forward(dat)
+        >>> new_dat["a"].shape
+        (2, 1)
+
+        >>> bc = bf.adapters.transforms.Broadcast("a", to="b", exclude=None)
+        >>> new_dat = bc.forward(dat)
+        >>> new_dat["a"].shape
+        (2, 3)
+
+        >>> bc = bf.adapters.transforms.Broadcast("b", to="c", expand=1)
+        >>> new_dat = bc.forward(dat)
+        >>> new_dat["b"].shape
+        (2, 2, 3)
 
     It is recommended to precede this transform with a :class:`bayesflow.adapters.transforms.ToArray` transform.
     """
 
-    def __init__(self, *, expand_scalars: bool = True):
+    def __init__(self, keys: Sequence[str], *, to: str, expand: str | int | tuple = "left", exclude: int | tuple = -1):
         super().__init__()
+        self.keys = keys
+        self.to = to
 
-        self.expand_scalars = expand_scalars
+        if isinstance(expand, int):
+            expand = (expand,)
 
-    # noinspection PyMethodOverriding
-    def forward(self, data: np.ndarray, *, batch_size: int, **kwargs):
-        data = np.repeat(data[None], batch_size, axis=0)
+        self.expand = expand
 
-        if self.expand_scalars and data.ndim == 1:
-            data = data[:, None]
+        if isinstance(exclude, int):
+            exclude = (exclude,)
 
-        return data
+        self.exclude = exclude
+
+    @classmethod
+    def from_config(cls, config: dict, custom_objects=None) -> "Broadcast":
+        return cls(
+            keys=deserialize(config["keys"], custom_objects),
+            to=deserialize(config["to"], custom_objects),
+            expand=deserialize(config["expand"], custom_objects),
+            exclude=deserialize(config["exclude"], custom_objects),
+        )
+
+    def get_config(self) -> dict:
+        return {
+            "keys": serialize(self.keys),
+            "to": serialize(self.to),
+            "expand": serialize(self.expand),
+            "exclude": serialize(self.exclude),
+        }
 
     # noinspection PyMethodOverriding
-    def inverse(self, data: np.ndarray, **kwargs) -> np.ndarray:
-        data = data[0]
+    def forward(self, data: dict[str, np.ndarray], **kwargs) -> dict[str, np.ndarray]:
+        target_shape = data[self.to].shape
+
+        data = data.copy()
+
+        for k in self.keys:
+            # ensure that .shape is defined
+            data[k] = np.asarray(data[k])
+            len_diff = len(target_shape) - len(data[k].shape)
+
+            if self.expand == "left":
+                data[k] = np.expand_dims(data[k], axis=tuple(np.arange(0, len_diff)))
+            elif self.expand == "right":
+                data[k] = np.expand_dims(data[k], axis=tuple(-np.arange(1, len_diff + 1)))
+            elif isinstance(self.expand, tuple):
+                if len(self.expand) is not len_diff:
+                    raise ValueError("Length of `expand` must match the length difference of the involed arrays.")
+                data[k] = np.expand_dims(data[k], axis=self.expand)
 
-        if self.expand_scalars:
-            data = np.squeeze(data, axis=0)
+            new_shape = target_shape
+            if self.exclude is not None:
+                new_shape = np.array(new_shape, dtype=int)
+                old_shape = np.array(data[k].shape, dtype=int)
+                exclude = list(self.exclude)
+                new_shape[exclude] = old_shape[exclude]
+                new_shape = tuple(new_shape)
 
+            data[k] = np.broadcast_to(data[k], new_shape)
+
+        return data
+
+    # noinspection PyMethodOverriding
+    def inverse(self, data: dict[str, np.ndarray], **kwargs) -> dict[str, np.ndarray]:
+        # TODO: add inverse
+        # we will likely never actually need the inverse broadcasting in practice
+        # so adding this method is not high priority
         return data

bayesflow/adapters/transforms/concatenate.py

@@ -12,7 +12,23 @@
 
 @serializable(package="bayesflow.adapters")
 class Concatenate(Transform):
-    """Concatenate multiple arrays into a new key."""
+    """Concatenate multiple arrays into a new key. Used to specify how data variables should be treated by the network.
+
+    Parameters:
+        keys: Input a list of strings, where the strings are the names of data variables.
+        into: A string telling the network how to use the variables named in keys.
+        axis: integer specifing along which axis to concatonate the keys. The last axis is used by default.
+
+    Example:
+    Suppose you have a simulator that generates variables "beta" and "sigma" from priors and then observation
+    variables "x" and "y". We can then use concatonate in the following way
+
+    adapter = (
+        bf.Adapter()
+            .concatenate(["beta", "sigma"], into="inference_variables")
+            .concatenate(["x", "y"], into="summary_variables")
+     )
+    """
 
     def __init__(self, keys: Sequence[str], *, into: str, axis: int = -1):
         self.keys = keys