Merge branch 'dev' into adapater_nan

# Conflicts:
#	bayesflow/adapters/adapter.py
#	bayesflow/adapters/transforms/__init__.py

Author: arrjon

.gitignore

@@ -39,3 +39,6 @@ docs/
 
 # MacOS
 .DS_Store
+
+# Rproj
+.Rproj.user

README.md

@@ -15,6 +15,12 @@ It provides users and researchers with:
 BayesFlow (version 2+) is designed to be a flexible and efficient tool that enables rapid statistical inference
 fueled by continuous progress in generative AI and Bayesian inference.
 
+> [!IMPORTANT]
+> As the 2.0 version introduced many new features, we still have to make breaking changes from time to time.
+> This especially concerns **saving and loading** of models. We aim to stabilize this from the 2.1 release onwards.
+> Until then, consider pinning your BayesFlow 2.0 installation to an exact version, or re-training after an update
+> for less costly models.
+
 ## Important Note for Existing Users
 
 You are currently looking at BayesFlow 2.0+, which is a complete rewrite of the library.

bayesflow/__init__.py

@@ -50,17 +50,11 @@ def setup():
             "in contexts where you need gradients (e.g. custom training loops)."
         )
 
+    # dynamically add __version__ attribute
+    from importlib.metadata import version
 
-# dynamically add version dunder variable
-try:
-    from importlib.metadata import version, PackageNotFoundError
+    globals()["__version__"] = version("bayesflow")
 
-    __version__ = version(__name__)
-except PackageNotFoundError:
-    __version__ = "2.0.0"
-finally:
-    del version
-    del PackageNotFoundError
 
 # call and clean up namespace
 setup()

bayesflow/adapters/adapter.py

@@ -14,23 +14,28 @@
     Drop,
     ExpandDims,
     FilterTransform,
+    Group,
     Keep,
     Log,
     MapTransform,
     NumpyTransform,
     OneHot,
     Rename,
     SerializableCustomTransform,
+    Squeeze,
     Sqrt,
     Standardize,
     ToArray,
     Transform,
+    Ungroup,
+    RandomSubsample,
+    Take,
     NanToNum,
 )
 from .transforms.filter_transform import Predicate
 
 
-@serializable
+@serializable("bayesflow.adapters")
 class Adapter(MutableSequence[Transform]):
     """
     Defines an adapter to apply various transforms to data.
@@ -599,6 +604,52 @@ def expand_dims(self, keys: str | Sequence[str], *, axis: int | tuple):
         self.transforms.append(transform)
         return self
 
+    def group(self, keys: Sequence[str], into: str, *, prefix: str = ""):
+        """Append a :py:class:`~transforms.Group` transform to the adapter.
+
+        Groups the given variables as a dictionary in the key `into`. As most transforms do
+        not support nested structures, this should usually be the last transform in the adapter.
+
+        Parameters
+        ----------
+        keys : Sequence of str
+            The names of the variables to group together.
+        into : str
+            The name of the variable to store the grouped variables in.
+        prefix : str, optional
+            An optional common prefix of the variable names before grouping, which will be removed after grouping.
+
+        Raises
+        ------
+        ValueError
+            If a prefix is specified, but a provided key does not start with the prefix.
+        """
+        if isinstance(keys, str):
+            keys = [keys]
+
+        transform = Group(keys=keys, into=into, prefix=prefix)
+        self.transforms.append(transform)
+        return self
+
+    def ungroup(self, key: str, *, prefix: str = ""):
+        """Append an :py:class:`~transforms.Ungroup` transform to the adapter.
+
+        Ungroups the the variables in `key` from a dictionary into individual entries. Most transforms do
+        not support nested structures, so this can be used to flatten a nested structure.
+        The nesting can be re-established after the transforms using the :py:meth:`group` method.
+
+        Parameters
+        ----------
+        key : str
+            The name of the variable to ungroup. The corresponding variable has to be a dictionary.
+        prefix : str, optional
+            An optional common prefix that will be added to the ungrouped variable names. This can be necessary
+            to avoid duplicate names.
+        """
+        transform = Ungroup(key=key, prefix=prefix)
+        self.transforms.append(transform)
+        return self
+
     def keep(self, keys: str | Sequence[str]):
         """Append a :py:class:`~transforms.Keep` transform to the adapter.
 
@@ -666,6 +717,28 @@ def one_hot(self, keys: str | Sequence[str], num_classes: int):
         self.transforms.append(transform)
         return self
 
+    def random_subsample(self, key: str, *, sample_size: int | float, axis: int = -1):
+        """
+        Append a :py:class:`~transforms.RandomSubsample` transform to the adapter.
+
+        Parameters
+        ----------
+        key : str or Sequence of str
+            The name of the variable to subsample.
+        sample_size : int or float
+            The number of samples to draw, or a fraction between 0 and 1 of the total number of samples to draw.
+        axis: int, optional
+            Which axis to draw samples over. The last axis is used by default.
+        """
+
+        if not isinstance(key, str):
+            raise TypeError("Can only subsample one batch entry at a time.")
+
+        transform = MapTransform({key: RandomSubsample(sample_size=sample_size, axis=axis)})
+
+        self.transforms.append(transform)
+        return self
+
     def rename(self, from_key: str, to_key: str):
         """Append a :py:class:`~transforms.Rename` transform to the adapter.
 
@@ -709,6 +782,24 @@ def split(self, key: str, *, into: Sequence[str], indices_or_sections: int | Seq
 
         return self
 
+    def squeeze(self, keys: str | Sequence[str], *, axis: int | tuple):
+        """Append a :py:class:`~transforms.Squeeze` transform to the adapter.
+
+        Parameters
+        ----------
+        keys : str or Sequence of str
+            The names of the variables to squeeze.
+        axis : int or tuple
+            The axis to squeeze. As the number of batch dimensions might change, we advise using negative
+            numbers (i.e., indexing from the end instead of the start).
+        """
+        if isinstance(keys, str):
+            keys = [keys]
+
+        transform = MapTransform({key: Squeeze(axis=axis) for key in keys})
+        self.transforms.append(transform)
+        return self
+
     def sqrt(self, keys: str | Sequence[str]):
         """Append an :py:class:`~transforms.Sqrt` transform to the adapter.
 
@@ -742,7 +833,7 @@ def standardize(
             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
+        **kwargs :
             Additional keyword arguments passed to the transform.
         """
         transform = FilterTransform(
@@ -755,6 +846,42 @@ def standardize(
         self.transforms.append(transform)
         return self
 
+    def take(
+        self,
+        include: str | Sequence[str] = None,
+        *,
+        indices: Sequence[int],
+        axis: int = -1,
+        predicate: Predicate = None,
+        exclude: str | Sequence[str] = None,
+    ):
+        """
+        Append a :py:class:`~transforms.Take` transform to the adapter.
+
+        Parameters
+        ----------
+        include : str or Sequence of str, optional
+            Names of variables to include in the transform.
+        indices : Sequence of int
+            Which indices to take from the data.
+        axis : int, optional
+            Which axis to take from. The last axis is used by default.
+        predicate : Predicate, optional
+            Function that indicates which variables should be transformed.
+        exclude : str or Sequence of str, optional
+            Names of variables to exclude from the transform.
+        """
+        transform = FilterTransform(
+            transform_constructor=Take,
+            predicate=predicate,
+            include=include,
+            exclude=exclude,
+            indices=indices,
+            axis=axis,
+        )
+        self.transforms.append(transform)
+        return self
+
     def to_array(
         self,
         include: str | Sequence[str] = None,

bayesflow/adapters/transforms/__init__.py

@@ -8,6 +8,7 @@
 from .elementwise_transform import ElementwiseTransform
 from .expand_dims import ExpandDims
 from .filter_transform import FilterTransform
+from .group import Group
 from .keep import Keep
 from .log import Log
 from .map_transform import MapTransform
@@ -18,11 +19,15 @@
 from .serializable_custom_transform import SerializableCustomTransform
 from .shift import Shift
 from .split import Split
+from .squeeze import Squeeze
 from .sqrt import Sqrt
 from .standardize import Standardize
 from .to_array import ToArray
 from .to_dict import ToDict
 from .transform import Transform
+from .random_subsample import RandomSubsample
+from .take import Take
+from .ungroup import Ungroup
 from .nan_to_num import NanToNum
 
 from ...utils._docs import _add_imports_to_all

bayesflow/adapters/transforms/as_set.py

@@ -5,7 +5,7 @@
 from .elementwise_transform import ElementwiseTransform
 
 
-@serializable
+@serializable("bayesflow.adapters")
 class AsSet(ElementwiseTransform):
     """The `.as_set(["x", "y"])` transform indicates that both `x` and `y` are treated as sets.
 

bayesflow/adapters/transforms/as_time_series.py

@@ -5,7 +5,7 @@
 from .elementwise_transform import ElementwiseTransform
 
 
-@serializable
+@serializable("bayesflow.adapters")
 class AsTimeSeries(ElementwiseTransform):
     """The `.as_time_series` transform can be used to indicate that variables shall be treated as time series.
 

bayesflow/adapters/transforms/broadcast.py

@@ -6,7 +6,7 @@
 from .transform import Transform
 
 
-@serializable
+@serializable("bayesflow.adapters")
 class Broadcast(Transform):
     """
     Broadcasts arrays or scalars to the shape of a given other array.

bayesflow/adapters/transforms/concatenate.py

@@ -7,7 +7,7 @@
 from .transform import Transform
 
 
-@serializable
+@serializable("bayesflow.adapters")
 class Concatenate(Transform):
     """Concatenate multiple arrays into a new key. Used to specify how data variables should be treated by the network.
 

bayesflow/adapters/transforms/constrain.py

@@ -11,7 +11,7 @@
 from .elementwise_transform import ElementwiseTransform
 
 
-@serializable
+@serializable("bayesflow.adapters")
 class Constrain(ElementwiseTransform):
     """
     Constrains neural network predictions of a data variable to specified bounds.