Merge branch 'main' into feat-diffusion-model

Author: arrjon

.github/workflows/publish.yaml

@@ -1,18 +1,53 @@
+name: Publish Python 🐍 distribution 📦 to PyPI
 
-name: Publish to PyPI.org
 on:
   release:
     types: [published]
+
 jobs:
-  pypi:
+  build:
+    name: Build distribution 📦
+    runs-on: ubuntu-latest
+
+    steps:
+    - uses: actions/checkout@v4
+      with:
+        persist-credentials: false
+    - name: Set up Python
+      uses: actions/setup-python@v5
+      with:
+        python-version: "3.x"
+    - name: Install pypa/build
+      run: >-
+        python3 -m
+        pip install
+        build
+        --user
+    - name: Build a binary wheel and a source tarball
+      run: python3 -m build
+    - name: Store the distribution packages
+      uses: actions/upload-artifact@v4
+      with:
+        name: python-package-distributions
+        path: dist/
+
+  publish-to-pypi:
+    name: >-
+      Publish Python 🐍 distribution 📦 to PyPI
+    needs:
+    - build
     runs-on: ubuntu-latest
+    environment:
+      name: pypi
+      url: https://pypi.org/p/bayesflow  # Replace <package-name> with your PyPI project name
+    permissions:
+      id-token: write  # IMPORTANT: mandatory for trusted publishing
+
     steps:
-      - name: Checkout
-        uses: actions/checkout@v4
-        with:
-          fetch-depth: 0
-      - run: python3 -m pip install -U build && python3 -m build
-      - name: Publish package
-        uses: pypa/gh-action-pypi-publish@release/v1
-        with:
-          password: ${{ secrets.PYPI_API_TOKEN }}
+    - name: Download all the dists
+      uses: actions/download-artifact@v4
+      with:
+        name: python-package-distributions
+        path: dist/
+    - name: Publish distribution 📦 to PyPI
+      uses: pypa/gh-action-pypi-publish@release/v1

README.md

@@ -64,15 +64,19 @@ More tutorials are always welcome! Please consider making a pull request if you
 
 ## Install
 
-BayesFlow v2 is not yet installable via PyPI, but you can use the following command to install the latest version of the `main` branch:
+You can install the latest stable version from PyPI using:
 
 ```bash
-pip install git+https://github.com/bayesflow-org/bayesflow.git
+pip install bayesflow
 ```
 
-If you encounter problems with this or require more control, please refer to the instructions to install from source below.
+If you want the latest features, you can install from source:
 
-Note: `pip install bayesflow` will install the v1 version of BayesFlow.
+```bash
+pip install git+https://github.com/bayesflow-org/bayesflow.git@dev
+```
+
+If you encounter problems with this or require more control, please refer to the instructions to install from source below.
 
 ### Backend
 

bayesflow/__init__.py

@@ -51,6 +51,17 @@ def setup():
         )
 
 
+# dynamically add version dunder variable
+try:
+    from importlib.metadata import version, PackageNotFoundError
+
+    __version__ = version(__name__)
+except PackageNotFoundError:
+    __version__ = "2.0.0"
+finally:
+    del version
+    del PackageNotFoundError
+
 # call and clean up namespace
 setup()
 del setup

bayesflow/adapters/adapter.py

@@ -2,11 +2,7 @@
 
 import numpy as np
 
-from keras.saving import (
-    deserialize_keras_object as deserialize,
-    register_keras_serializable as serializable,
-    serialize_keras_object as serialize,
-)
+from bayesflow.utils.serialization import deserialize, serialize, serializable
 
 from .transforms import (
     AsSet,
@@ -33,7 +29,7 @@
 from .transforms.filter_transform import Predicate
 
 
-@serializable(package="bayesflow.adapters")
+@serializable
 class Adapter(MutableSequence[Transform]):
     """
     Defines an adapter to apply various transforms to data.
@@ -74,18 +70,24 @@ def create_default(inference_variables: Sequence[str]) -> "Adapter":
 
     @classmethod
     def from_config(cls, config: dict, custom_objects=None) -> "Adapter":
-        return cls(transforms=deserialize(config["transforms"], custom_objects))
+        return cls(**deserialize(config, custom_objects=custom_objects))
 
     def get_config(self) -> dict:
-        return {"transforms": serialize(self.transforms)}
+        config = {
+            "transforms": self.transforms,
+        }
+
+        return serialize(config)
 
-    def forward(self, data: dict[str, any], **kwargs) -> dict[str, np.ndarray]:
+    def forward(self, data: dict[str, any], *, stage: str = "inference", **kwargs) -> dict[str, np.ndarray]:
         """Apply the transforms in the forward direction.
 
         Parameters
         ----------
         data : dict
             The data to be transformed.
+        stage : str, one of ["training", "validation", "inference"]
+            The stage the function is called in.
         **kwargs : dict
             Additional keyword arguments passed to each transform.
 
@@ -97,17 +99,19 @@ def forward(self, data: dict[str, any], **kwargs) -> dict[str, np.ndarray]:
         data = data.copy()
 
         for transform in self.transforms:
-            data = transform(data, **kwargs)
+            data = transform(data, stage=stage, **kwargs)
 
         return data
 
-    def inverse(self, data: dict[str, np.ndarray], **kwargs) -> dict[str, any]:
+    def inverse(self, data: dict[str, np.ndarray], *, stage: str = "inference", **kwargs) -> dict[str, any]:
         """Apply the transforms in the inverse direction.
 
         Parameters
         ----------
         data : dict
             The data to be transformed.
+        stage : str, one of ["training", "validation", "inference"]
+            The stage the function is called in.
         **kwargs : dict
             Additional keyword arguments passed to each transform.
 
@@ -119,11 +123,13 @@ def inverse(self, data: dict[str, np.ndarray], **kwargs) -> dict[str, any]:
         data = data.copy()
 
         for transform in reversed(self.transforms):
-            data = transform(data, inverse=True, **kwargs)
+            data = transform(data, stage=stage, inverse=True, **kwargs)
 
         return data
 
-    def __call__(self, data: Mapping[str, any], *, inverse: bool = False, **kwargs) -> dict[str, np.ndarray]:
+    def __call__(
+        self, data: Mapping[str, any], *, inverse: bool = False, stage="inference", **kwargs
+    ) -> dict[str, np.ndarray]:
         """Apply the transforms in the given direction.
 
         Parameters
@@ -132,6 +138,8 @@ def __call__(self, data: Mapping[str, any], *, inverse: bool = False, **kwargs)
             The data to be transformed.
         inverse : bool, optional
             If False, apply the forward transform, else apply the inverse transform (default False).
+        stage : str, one of ["training", "validation", "inference"]
+            The stage the function is called in.
         **kwargs
             Additional keyword arguments passed to each transform.
 
@@ -141,9 +149,9 @@ def __call__(self, data: Mapping[str, any], *, inverse: bool = False, **kwargs)
             The transformed data.
         """
         if inverse:
-            return self.inverse(data, **kwargs)
+            return self.inverse(data, stage=stage, **kwargs)
 
-        return self.forward(data, **kwargs)
+        return self.forward(data, stage=stage, **kwargs)
 
     def __repr__(self):
         result = ""
@@ -667,6 +675,18 @@ def shift(self, keys: str | Sequence[str], by: float | np.ndarray):
         self.transforms.append(MapTransform({key: Shift(shift=by) for key in keys}))
         return self
 
+    def split(self, key: str, *, into: Sequence[str], indices_or_sections: int | Sequence[int] = None, axis: int = -1):
+        from .transforms import Split
+
+        if isinstance(into, str):
+            transform = Rename(key, into)
+        else:
+            transform = Split(key, into, indices_or_sections, axis)
+
+        self.transforms.append(transform)
+
+        return self
+
     def sqrt(self, keys: str | Sequence[str]):
         """Append an :py:class:`~transforms.Sqrt` transform to the adapter.
 
@@ -743,3 +763,10 @@ def to_array(
         )
         self.transforms.append(transform)
         return self
+
+    def to_dict(self):
+        from .transforms import ToDict
+
+        transform = ToDict()
+        self.transforms.append(transform)
+        return self

bayesflow/adapters/transforms/__init__.py

@@ -17,9 +17,11 @@
 from .scale import Scale
 from .serializable_custom_transform import SerializableCustomTransform
 from .shift import Shift
+from .split import Split
 from .sqrt import Sqrt
 from .standardize import Standardize
 from .to_array import ToArray
+from .to_dict import ToDict
 from .transform import Transform
 
 from ...utils._docs import _add_imports_to_all

bayesflow/adapters/transforms/as_set.py

@@ -1,10 +1,11 @@
-from keras.saving import register_keras_serializable as serializable
 import numpy as np
 
+from bayesflow.utils.serialization import serializable
+
 from .elementwise_transform import ElementwiseTransform
 
 
-@serializable(package="bayesflow.adapters")
+@serializable
 class AsSet(ElementwiseTransform):
     """The `.as_set(["x", "y"])` transform indicates that both `x` and `y` are treated as sets.
 
@@ -33,9 +34,5 @@ def inverse(self, data: np.ndarray, **kwargs) -> np.ndarray:
 
         return data
 
-    @classmethod
-    def from_config(cls, config: dict, custom_objects=None) -> "AsSet":
-        return cls()
-
     def get_config(self) -> dict:
         return {}

bayesflow/adapters/transforms/as_time_series.py

@@ -1,10 +1,11 @@
 import numpy as np
-from keras.saving import register_keras_serializable as serializable
+
+from bayesflow.utils.serialization import serializable
 
 from .elementwise_transform import ElementwiseTransform
 
 
-@serializable(package="bayesflow.adapters")
+@serializable
 class AsTimeSeries(ElementwiseTransform):
     """The `.as_time_series` transform can be used to indicate that variables shall be treated as time series.
 
@@ -29,9 +30,5 @@ def inverse(self, data: np.ndarray, **kwargs) -> np.ndarray:
 
         return data
 
-    @classmethod
-    def from_config(cls, config: dict, custom_objects=None) -> "AsTimeSeries":
-        return cls()
-
     def get_config(self) -> dict:
         return {}

bayesflow/adapters/transforms/broadcast.py

@@ -1,16 +1,12 @@
 from collections.abc import Sequence
 import numpy as np
 
-from keras.saving import (
-    deserialize_keras_object as deserialize,
-    register_keras_serializable as serializable,
-    serialize_keras_object as serialize,
-)
+from bayesflow.utils.serialization import serialize, serializable
 
 from .transform import Transform
 
 
-@serializable(package="bayesflow.adapters")
+@serializable
 class Broadcast(Transform):
     """
     Broadcasts arrays or scalars to the shape of a given other array.
@@ -96,31 +92,15 @@ def __init__(
         self.exclude = exclude
         self.squeeze = squeeze
 
-    @classmethod
-    def from_config(cls, config: dict, custom_objects=None) -> "Broadcast":
-        # Deserialize turns tuples to lists, undo it if necessary
-        exclude = deserialize(config["exclude"], custom_objects)
-        exclude = tuple(exclude) if isinstance(exclude, list) else exclude
-        expand = deserialize(config["expand"], custom_objects)
-        expand = tuple(expand) if isinstance(expand, list) else expand
-        squeeze = deserialize(config["squeeze"], custom_objects)
-        squeeze = tuple(squeeze) if isinstance(squeeze, list) else squeeze
-        return cls(
-            keys=deserialize(config["keys"], custom_objects),
-            to=deserialize(config["to"], custom_objects),
-            expand=expand,
-            exclude=exclude,
-            squeeze=squeeze,
-        )
-
     def get_config(self) -> dict:
-        return {
-            "keys": serialize(self.keys),
-            "to": serialize(self.to),
-            "expand": serialize(self.expand),
-            "exclude": serialize(self.exclude),
-            "squeeze": serialize(self.squeeze),
+        config = {
+            "keys": self.keys,
+            "to": self.to,
+            "expand": self.expand,
+            "exclude": self.exclude,
+            "squeeze": self.squeeze,
         }
+        return serialize(config)
 
     # noinspection PyMethodOverriding
     def forward(self, data: dict[str, np.ndarray], **kwargs) -> dict[str, np.ndarray]: