Merge remote-tracking branch 'origin/dev' into dev

Author: LarsKue

.pre-commit-config.yaml

@@ -14,7 +14,7 @@
 repos:
 - repo: https://github.com/astral-sh/ruff-pre-commit
   # Ruff version.
-  rev: v0.4.9
+  rev: v0.9.6
   hooks:
     # Run the linter.
     - id: ruff
@@ -28,7 +28,7 @@ repos:
 #    - id: nbqa-ruff
 #      args: [--ignore=E402]  # E402: module level import not at top of file
 - repo: https://github.com/pre-commit/pre-commit-hooks
-  rev: v4.4.0
+  rev: v5.0.0
   hooks:
   # A bunch of other pre-defined hooks.
   - id: check-yaml

README.md

@@ -93,12 +93,14 @@ conda env create --file environment.yaml --name bayesflow
 
 Check out some of our walk-through notebooks below. We are actively working on porting all notebooks to the new interface so more will be available soon!
 
-1. [Two moons starter toy example](examples/TwoMoons_StarterNotebook.ipynb)
-2. [Linear regression](examples/Linear_Regression.ipynb)
-3. [Bayesian experimental design](examples/Bayesian_Experimental_Design.ipynb)
-4. [SIR model with custom summary network](examples/SIR_PosteriorEstimation.ipynb)
+1. [Linear regression starter example](examples/Linear_Regression_Starter.ipynb)
+2. [From ABC to BayesFlow](examples/From_ABC_to_BayesFlow.ipynb)
+3. [Two moons starter example](examples/Two_Moons_Starter.ipynb)
+4. [SIR model with custom summary network](examples/SIR_Posterior_Estimation.ipynb)
 5. [Hyperparameter optimization](examples/Hyperparameter_Optimization.ipynb)
-6. Coming soon...
+6. [Bayesian experimental design](examples/Bayesian_Experimental_Design.ipynb)
+7. [Simple model comparison example (One-Sample T-Test)](examples/One_Sample_TTest.ipynb)
+8. More coming soon...
 
 ## Documentation \& Help
 

bayesflow/adapters/transforms/broadcast.py

@@ -79,12 +79,19 @@ def __init__(
 
     @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=deserialize(config["expand"], custom_objects),
-            exclude=deserialize(config["exclude"], custom_objects),
-            squeeze=deserialize(config["squeeze"], custom_objects),
+            expand=expand,
+            exclude=exclude,
+            squeeze=squeeze,
         )
 
     def get_config(self) -> dict:

bayesflow/adapters/transforms/concatenate.py

@@ -30,26 +30,28 @@ class Concatenate(Transform):
      )
     """
 
-    def __init__(self, keys: Sequence[str], *, into: str, axis: int = -1):
+    def __init__(self, keys: Sequence[str], *, into: str, axis: int = -1, _indices: list | None = None):
         self.keys = keys
         self.into = into
         self.axis = axis
 
-        self.indices = None
+        self.indices = _indices
 
     @classmethod
     def from_config(cls, config: dict, custom_objects=None) -> "Concatenate":
         return cls(
             keys=deserialize(config["keys"], custom_objects),
             into=deserialize(config["into"], custom_objects),
             axis=deserialize(config["axis"], custom_objects),
+            _indices=deserialize(config["indices"], custom_objects),
         )
 
     def get_config(self) -> dict:
         return {
             "keys": serialize(self.keys),
             "into": serialize(self.into),
             "axis": serialize(self.axis),
+            "indices": serialize(self.indices),
         }
 
     def forward(self, data: dict[str, any], *, strict: bool = True, **kwargs) -> dict[str, any]:

bayesflow/adapters/transforms/filter_transform.py

@@ -4,6 +4,8 @@
 import numpy as np
 from keras.saving import (
     deserialize_keras_object as deserialize,
+    get_registered_name,
+    get_registered_object,
     register_keras_serializable as serializable,
     serialize_keras_object as serialize,
 )
@@ -79,21 +81,33 @@ def extra_repr(self) -> str:
 
     @classmethod
     def from_config(cls, config: dict, custom_objects=None) -> "Transform":
-        def transform_constructor(*args, **kwargs):
-            raise RuntimeError(
-                "Instantiating new elementwise transforms on a deserialized FilterTransform is not yet supported (and"
-                "may never be). As a work-around, you can manually register the elementwise transform constructor after"
-                "deserialization:\n"
-                "obj = deserialize(config)\n"
-                "obj.transform_constructor = MyElementwiseTransform"
-            )
-
+        transform_constructor = get_registered_object(config["transform_constructor"])
+        try:
+            kwargs = deserialize(config["kwargs"])
+        except TypeError as e:
+            if transform_constructor.__name__ == "LambdaTransform":
+                raise TypeError(
+                    "LambdaTransform (created by Adapter.apply) could not be deserialized.\n"
+                    "This is probably because the custom transform functions `forward` and "
+                    "`backward` from `Adapter.apply` were not passed as `custom_objects`.\n"
+                    "For example, if your adapter uses\n"
+                    "`Adapter.apply(forward=forward_transform, inverse=inverse_transform)`,\n"
+                    "you have to pass\n"
+                    '`custom_objects={"forward_transform": forward_transform, '
+                    '"inverse_transform": inverse_transform}`\n'
+                    "to the function you use to load the serialized object."
+                ) from e
+            raise TypeError(
+                "The transform could not be deserialized properly. "
+                "The most likely reason is that some classes or functions "
+                "are not known during deserialization. Please pass them as `custom_objects`."
+            ) from e
         instance = cls(
             transform_constructor=transform_constructor,
             predicate=deserialize(config["predicate"], custom_objects),
             include=deserialize(config["include"], custom_objects),
             exclude=deserialize(config["exclude"], custom_objects),
-            **config["kwargs"],
+            **kwargs,
         )
 
         instance.transform_map = deserialize(config["transform_map"])
@@ -102,6 +116,7 @@ def transform_constructor(*args, **kwargs):
 
     def get_config(self) -> dict:
         return {
+            "transform_constructor": get_registered_name(self.transform_constructor),
             "predicate": serialize(self.predicate),
             "include": serialize(self.include),
             "exclude": serialize(self.exclude),

bayesflow/adapters/transforms/standardize.py

@@ -41,10 +41,14 @@ def __init__(
 
     @classmethod
     def from_config(cls, config: dict, custom_objects=None) -> "Standardize":
+        # Deserialize turns tuples to lists, undo it if necessary
+        deserialized_axis = deserialize(config["axis"], custom_objects)
+        if isinstance(deserialized_axis, list):
+            deserialized_axis = tuple(deserialized_axis)
         return cls(
             mean=deserialize(config["mean"], custom_objects),
             std=deserialize(config["std"], custom_objects),
-            axis=deserialize(config["axis"], custom_objects),
+            axis=deserialized_axis,
             momentum=deserialize(config["momentum"], custom_objects),
         )
 

bayesflow/approximators/continuous_approximator.py

@@ -140,6 +140,8 @@ def sample(
         **kwargs,
     ) -> dict[str, np.ndarray]:
         conditions = self.adapter(conditions, strict=False, stage="inference", **kwargs)
+        # at inference time, inference_variables are estimated by the networks and thus ignored in conditions
+        conditions.pop("inference_variables", None)
         conditions = keras.tree.map_structure(keras.ops.convert_to_tensor, conditions)
         conditions = {"inference_variables": self._sample(num_samples=num_samples, **conditions, **kwargs)}
         conditions = keras.tree.map_structure(keras.ops.convert_to_numpy, conditions)

bayesflow/approximators/model_comparison_approximator.py

@@ -1,6 +1,7 @@
 from collections.abc import Mapping, Sequence
 
 import keras
+import numpy as np
 from keras.saving import (
     deserialize_keras_object as deserialize,
     register_keras_serializable as serializable,
@@ -198,3 +199,46 @@ def get_config(self):
         }
 
         return base_config | config
+
+    def predict(
+        self,
+        *,
+        conditions: dict[str, np.ndarray],
+        logits: bool = False,
+        **kwargs,
+    ) -> np.ndarray:
+        conditions = self.adapter(conditions, strict=False, stage="inference", **kwargs)
+        # at inference time, model_indices are predicted by the networks and thus ignored in conditions
+        conditions.pop("model_indices", None)
+        conditions = keras.tree.map_structure(keras.ops.convert_to_tensor, conditions)
+
+        output = self._predict(**conditions, **kwargs)
+
+        if not logits:
+            output = keras.ops.softmax(output)
+
+        output = keras.ops.convert_to_numpy(output)
+
+        return output
+
+    def _predict(self, classifier_conditions: Tensor = None, summary_variables: Tensor = None, **kwargs) -> Tensor:
+        if self.summary_network is None:
+            if summary_variables is not None:
+                raise ValueError("Cannot use summary variables without a summary network.")
+        else:
+            if summary_variables is None:
+                raise ValueError("Summary variables are required when a summary network is present")
+
+            summary_outputs = self.summary_network(
+                summary_variables, **filter_kwargs(kwargs, self.summary_network.call)
+            )
+
+            if classifier_conditions is None:
+                classifier_conditions = summary_outputs
+            else:
+                classifier_conditions = keras.ops.concatenate([classifier_conditions, summary_outputs], axis=1)
+
+        output = self.classifier_network(classifier_conditions)
+        output = self.logits_projector(output)
+
+        return output

bayesflow/datasets/offline_dataset.py

@@ -6,7 +6,9 @@
 
 class OfflineDataset(keras.utils.PyDataset):
     """
-    A dataset that is pre-simulated and stored in memory.
+    A dataset that is pre-simulated and stored in memory. When storing and loading data from disk, it is recommended to
+    save any pre-simulated data in raw form and create the `OfflineDataset` object only after loading in the raw data.
+    See the `DiskDataset` class for handling large datasets that are split into multiple smaller files.
     """
 
     def __init__(self, data: dict[str, np.ndarray], batch_size: int, adapter: Adapter | None, **kwargs):

bayesflow/diagnostics/metrics/__init__.py

@@ -1,3 +1,4 @@
 from .calibration_error import calibration_error
 from .posterior_contraction import posterior_contraction
 from .root_mean_squared_error import root_mean_squared_error
+from .expected_calibration_error import expected_calibration_error