Merge branch 'dev' into feat-diffusion-model

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.
@@ -245,8 +251,9 @@ Depending on your needs, you might not want to upgrade yet if one of the followi
   with the new version. Loading models from version 1.x in version 2.0+ is not supported.
 - You require a feature that was not ported to BayesFlow 2.0+ yet. To our knowledge,
   this applies to:
-  * Two-level/Hierarchical models: `TwoLevelGenerativeModel`, `TwoLevelPrior`.
-  * Sensitivity analysis: functionality from the `bayesflow.sensitivity` module.
+  * Two-level/Hierarchical models (planned for version 2.1): `TwoLevelGenerativeModel`, `TwoLevelPrior`.
+  * Sensitivity analysis (partially discontinued): functionality from the `bayesflow.sensitivity` module. This is still
+    possible, but we do no longer offer a special module for it. We plan to add a tutorial on this, see [#455](https://github.com/bayesflow-org/bayesflow/issues/455).
   * MCMC (discontinued): The `bayesflow.mcmc` module. We are considering other options
     to enable the use of BayesFlow in an MCMC setting.
   * Networks: `EvidentialNetwork`.

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

@@ -25,11 +25,13 @@
     Standardize,
     ToArray,
     Transform,
+    RandomSubsample,
+    Take,
 )
 from .transforms.filter_transform import Predicate
 
 
-@serializable
+@serializable("bayesflow.adapters")
 class Adapter(MutableSequence[Transform]):
     """
     Defines an adapter to apply various transforms to data.
@@ -79,7 +81,9 @@ def get_config(self) -> dict:
 
         return serialize(config)
 
-    def forward(self, data: dict[str, any], *, stage: str = "inference", **kwargs) -> dict[str, np.ndarray]:
+    def forward(
+        self, data: dict[str, any], *, stage: str = "inference", log_det_jac: bool = False, **kwargs
+    ) -> dict[str, np.ndarray] | tuple[dict[str, np.ndarray], dict[str, np.ndarray]]:
         """Apply the transforms in the forward direction.
 
         Parameters
@@ -88,22 +92,33 @@ def forward(self, data: dict[str, any], *, stage: str = "inference", **kwargs) -
             The data to be transformed.
         stage : str, one of ["training", "validation", "inference"]
             The stage the function is called in.
+        log_det_jac: bool, optional
+            Whether to return the log determinant of the Jacobian of the transforms.
         **kwargs : dict
             Additional keyword arguments passed to each transform.
 
         Returns
         -------
-        dict
-            The transformed data.
+        dict | tuple[dict, dict]
+            The transformed data or tuple of transformed data and log determinant of the Jacobian.
         """
         data = data.copy()
+        if not log_det_jac:
+            for transform in self.transforms:
+                data = transform(data, stage=stage, **kwargs)
+            return data
 
+        log_det_jac = {}
         for transform in self.transforms:
-            data = transform(data, stage=stage, **kwargs)
+            transformed_data = transform(data, stage=stage, **kwargs)
+            log_det_jac = transform.log_det_jac(data, log_det_jac, **kwargs)
+            data = transformed_data
 
-        return data
+        return data, log_det_jac
 
-    def inverse(self, data: dict[str, np.ndarray], *, stage: str = "inference", **kwargs) -> dict[str, any]:
+    def inverse(
+        self, data: dict[str, np.ndarray], *, stage: str = "inference", log_det_jac: bool = False, **kwargs
+    ) -> dict[str, np.ndarray] | tuple[dict[str, np.ndarray], dict[str, np.ndarray]]:
         """Apply the transforms in the inverse direction.
 
         Parameters
@@ -112,24 +127,32 @@ def inverse(self, data: dict[str, np.ndarray], *, stage: str = "inference", **kw
             The data to be transformed.
         stage : str, one of ["training", "validation", "inference"]
             The stage the function is called in.
+        log_det_jac: bool, optional
+            Whether to return the log determinant of the Jacobian of the transforms.
         **kwargs : dict
             Additional keyword arguments passed to each transform.
 
         Returns
         -------
-        dict
-            The transformed data.
+        dict | tuple[dict, dict]
+            The transformed data or tuple of transformed data and log determinant of the Jacobian.
         """
         data = data.copy()
+        if not log_det_jac:
+            for transform in reversed(self.transforms):
+                data = transform(data, stage=stage, inverse=True, **kwargs)
+            return data
 
+        log_det_jac = {}
         for transform in reversed(self.transforms):
             data = transform(data, stage=stage, inverse=True, **kwargs)
+            log_det_jac = transform.log_det_jac(data, log_det_jac, inverse=True, **kwargs)
 
-        return data
+        return data, log_det_jac
 
     def __call__(
         self, data: Mapping[str, any], *, inverse: bool = False, stage="inference", **kwargs
-    ) -> dict[str, np.ndarray]:
+    ) -> dict[str, np.ndarray] | tuple[dict[str, np.ndarray], dict[str, np.ndarray]]:
         """Apply the transforms in the given direction.
 
         Parameters
@@ -145,8 +168,8 @@ def __call__(
 
         Returns
         -------
-        dict
-            The transformed data.
+        dict | tuple[dict, dict]
+            The transformed data or tuple of transformed data and log determinant of the Jacobian.
         """
         if inverse:
             return self.inverse(data, stage=stage, **kwargs)
@@ -644,6 +667,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.
 
@@ -720,7 +765,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(
@@ -733,6 +778,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

@@ -23,6 +23,8 @@
 from .to_array import ToArray
 from .to_dict import ToDict
 from .transform import Transform
+from .random_subsample import RandomSubsample
+from .take import Take
 
 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.
 
@@ -115,3 +115,37 @@ def extra_repr(self) -> str:
             result += f", axis={self.axis}"
 
         return result
+
+    def log_det_jac(
+        self,
+        data: dict[str, np.ndarray],
+        log_det_jac: dict[str, np.ndarray],
+        *,
+        strict: bool = False,
+        inverse: bool = False,
+        **kwargs,
+    ) -> dict[str, np.ndarray]:
+        # copy to avoid side effects
+        log_det_jac = log_det_jac.copy()
+
+        if inverse:
+            if log_det_jac.get(self.into) is not None:
+                raise ValueError(
+                    "Cannot obtain an inverse Jacobian of concatenation. "
+                    "Transform your variables before you concatenate."
+                )
+
+            return log_det_jac
+
+        required_keys = set(self.keys)
+        available_keys = set(log_det_jac.keys())
+        common_keys = available_keys & required_keys
+
+        if len(common_keys) == 0:
+            return log_det_jac
+
+        parts = [log_det_jac.pop(key) for key in common_keys]
+
+        log_det_jac[self.into] = sum(parts)
+
+        return log_det_jac