Merge pull request #460 from bayesflow-org/dev

v2.0.3

Author: LarsKue

README.md

@@ -15,6 +15,25 @@ 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.
+While it shares the same overall goals with the 1.x versions, the API is not compatible.
+
+> [!CAUTION]
+> A few features, most notably hierarchical models, have not been ported to BayesFlow 2.0+
+> yet. We are working on those features and plan to add them soon. You can find the complete
+> list in the [FAQ](#faq) below.
+
+The [Moving from BayesFlow v1.1 to v2.0](examples/From_BayesFlow_1.1_to_2.0.ipynb) guide
+highlights how concepts and classes relate between the two versions.
+
 ## Conceptual Overview
 
 <div align="center">
@@ -216,11 +235,48 @@ while the old version was based on TensorFlow.
 
 -------------
 
+**Question:**
+Should I switch to BayesFlow 2.0+ now? Are there features that are still missing?
+
+**Answer:**
+In general, we recommend to switch, as the new version is easier to use and will continue
+to receive improvements and new features. However, a few features are still missing, so you
+might want to wait until everything you need has been ported to BayesFlow 2.0+.
+
+Depending on your needs, you might not want to upgrade yet if one of the following applies:
+
+- You have an ongoing project that uses BayesFlow 1.x, and you do not want to allocate
+  time for migrating it to the new API.
+- You have already trained models in BayesFlow 1.x, that you do not want to re-train
+  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 (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`.
+  * Model misspecification detection: MMD test in the summary space (see #384).
+
+If you encounter any functionality that is missing and not listed here, please let us
+know by opening an issue.
+
+-------------
+
 **Question:**
 I still need the old BayesFlow for some of my projects. How can I install it?
 
 **Answer:**
 You can find and install the old Bayesflow version via the `stable-legacy` branch on GitHub.
+The corresponding [documentation](https://bayesflow.org/stable-legacy/index.html) can be
+accessed by selecting the "stable-legacy" entry in the version picker of the documentation.
+
+You can also install the latest version of BayesFlow v1.x from PyPI using
+
+```
+pip install "bayesflow<2.0"
+```
 
 -------------
 

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

@@ -29,7 +29,7 @@
 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 +79,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 +90,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 +125,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 +166,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)

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

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.
@@ -87,6 +87,11 @@ def constrain(x):
 
                     def unconstrain(x):
                         return inverse_sigmoid((x - lower) / (upper - lower))
+
+                    def ldj(x):
+                        x = (x - lower) / (upper - lower)
+                        return -np.log(x) - np.log1p(-x) - np.log(upper - lower)
+
                 case str() as name:
                     raise ValueError(f"Unsupported method name for double bounded constraint: '{name}'.")
                 case other:
@@ -101,13 +106,22 @@ def constrain(x):
 
                     def unconstrain(x):
                         return inverse_softplus(x - lower)
+
+                    def ldj(x):
+                        x = x - lower
+                        return x - np.log(np.exp(x) - 1)
+
                 case "exp" | "log":
 
                     def constrain(x):
                         return np.exp(x) + lower
 
                     def unconstrain(x):
                         return np.log(x - lower)
+
+                    def ldj(x):
+                        return -np.log(x - lower)
+
                 case str() as name:
                     raise ValueError(f"Unsupported method name for single bounded constraint: '{name}'.")
                 case other:
@@ -122,13 +136,21 @@ def constrain(x):
 
                     def unconstrain(x):
                         return -inverse_softplus(-(x - upper))
+
+                    def ldj(x):
+                        x = -(x - upper)
+                        return x - np.log(np.exp(x) - 1)
+
                 case "exp" | "log":
 
                     def constrain(x):
                         return -np.exp(-x) + upper
 
                     def unconstrain(x):
                         return -np.log(-x + upper)
+
+                    def ldj(x):
+                        return -np.log(-x + upper)
                 case str() as name:
                     raise ValueError(f"Unsupported method name for single bounded constraint: '{name}'.")
                 case other:
@@ -142,6 +164,7 @@ def unconstrain(x):
 
         self.constrain = constrain
         self.unconstrain = unconstrain
+        self.ldj = ldj
 
         # do this last to avoid serialization issues
         match inclusive:
@@ -178,3 +201,9 @@ def forward(self, data: np.ndarray, **kwargs) -> np.ndarray:
     def inverse(self, data: np.ndarray, **kwargs) -> np.ndarray:
         # inverse means network space -> data space, so constrain the data
         return self.constrain(data)
+
+    def log_det_jac(self, data: np.ndarray, inverse: bool = False, **kwargs) -> np.ndarray:
+        ldj = self.ldj(data)
+        if inverse:
+            ldj = -ldj
+        return np.sum(ldj, axis=tuple(range(1, ldj.ndim)))

bayesflow/adapters/transforms/convert_dtype.py

@@ -5,7 +5,7 @@
 from .elementwise_transform import ElementwiseTransform
 
 
-@serializable
+@serializable("bayesflow.adapters")
 class ConvertDType(ElementwiseTransform):
     """
     Default transform used to convert all floats from float64 to float32 to be in line with keras framework.

bayesflow/adapters/transforms/drop.py

@@ -5,7 +5,7 @@
 from .transform import Transform
 
 
-@serializable
+@serializable("bayesflow.adapters")
 class Drop(Transform):
     """
     Transform to drop variables from further calculation.
@@ -46,3 +46,6 @@ def inverse(self, data: dict[str, any], **kwargs) -> dict[str, any]:
 
     def extra_repr(self) -> str:
         return "[" + ", ".join(map(repr, self.keys)) + "]"
+
+    def log_det_jac(self, data: dict[str, any], log_det_jac: dict[str, any], inverse: bool = False, **kwargs):
+        return self.inverse(data=log_det_jac) if inverse else self.forward(data=log_det_jac)