Merge branch 'standardize_in_approx' of https://github.com/bayesflow-org/bayesflow into standardize_in_approx

Author: stefanradev93

README.md

@@ -51,10 +51,10 @@ complex to be described analytically.
 
 ## Install
 
-You can install the latest stable version from PyPI using:
+We currently support Python 3.10 to 3.12. You can install the latest stable version from PyPI using:
 
 ```bash
-pip install bayesflow
+pip install "bayesflow>=2.0"
 ```
 
 If you want the latest features, you can install from source:
@@ -134,7 +134,8 @@ Many examples from [Bayesian Cognitive Modeling: A Practical Course](https://bay
 5. [SIR model with custom summary network](examples/SIR_Posterior_Estimation.ipynb)
 6. [Bayesian experimental design](examples/Bayesian_Experimental_Design.ipynb)
 7. [Simple model comparison example](examples/One_Sample_TTest.ipynb)
-8. [Moving from BayesFlow v1.1 to v2.0](examples/From_BayesFlow_1.1_to_2.0.ipynb)
+8. [Likelihood estimation](examples/Likelihood_Estimation.ipynb)
+9. [Moving from BayesFlow v1.1 to v2.0](examples/From_BayesFlow_1.1_to_2.0.ipynb)
 
 More tutorials are always welcome! Please consider making a pull request if you have a cool application that you want to contribute.
 

bayesflow/approximators/continuous_approximator.py

@@ -409,18 +409,10 @@ def sample(
         dict[str, np.ndarray]
             Dictionary containing generated samples with the same keys as `conditions`.
         """
-
-        # Apply adapter transforms to raw simulated / real quantities
-        conditions = self.adapter(conditions, strict=False, stage="inference", **kwargs)
-
-        # Ensure only keys relevant for sampling are present in the conditions dictionary
+        # Adapt, optionally standardize and convert conditions to tensor.
+        conditions = self._prepare_data(conditions, **kwargs)
+        # Remove any superfluous keys, just retain actual conditions.  # TODO: is this necessary?
         conditions = {k: v for k, v in conditions.items() if k in ContinuousApproximator.CONDITION_KEYS}
-        conditions = keras.tree.map_structure(keras.ops.convert_to_tensor, conditions)
-
-        # Optionally standardize conditions
-        for key in ContinuousApproximator.CONDITION_KEYS:
-            if key in conditions and key in self.standardize:
-                conditions[key] = self.standardize_layers[key](conditions[key])
 
         # Sample and undo optional standardization
         samples = self._sample(num_samples=num_samples, **conditions, **kwargs)
@@ -438,6 +430,51 @@ def sample(
             samples = split_arrays(samples, axis=-1)
         return samples
 
+    def _prepare_data(
+        self, data: Mapping[str, np.ndarray], log_det_jac: bool = False, **kwargs
+    ) -> dict[str, Tensor] | tuple[dict[str, Tensor], dict[str, Tensor]]:
+        """
+        Adapts, optionally standardizes, and converts the data to tensors to prepare it for the inference network.
+
+        Deals with data that represents only conditions, or only inference_variables or both.
+        """
+        # TODO:
+        # * [ ] better docstring
+
+        # Adapt, and optionally keep track of ldj of transformations to inference_variables.
+        adapted = self.adapter(data, strict=False, stage="inference", log_det_jac=log_det_jac, **kwargs)
+        if log_det_jac:
+            data, log_det_jac_adapter = adapted
+            log_det_jac_inference_variables = log_det_jac_adapter.get("inference_variables", 0.0)
+        else:
+            data = adapted
+
+        # Optionally standardize conditions, if they are part of data.
+        conditions = {k: v for k, v in data.items() if k in ContinuousApproximator.CONDITION_KEYS}
+        for key, value in conditions.items():
+            if key in self.standardize and key in data.keys():
+                data[key] = self.standardize_layers[key](value)
+
+        # Optionally standardize inference variables, if they are part of data.
+        if "inference_variables" in data.keys() and "inference_variables" in self.standardize:
+            standardized = self.standardize_layers["inference_variables"](
+                data["inference_variables"], log_det_jac=log_det_jac
+            )
+
+            # Optionally keep track of appropriate log_det_jac.
+            if log_det_jac:
+                data["inference_variables"], log_det_std = standardized
+                log_det_jac_inference_variables += keras.ops.convert_to_numpy(log_det_std)
+            else:
+                data["inference_variables"] = standardized
+
+        # Convert to tensor and return.
+        data = keras.tree.map_structure(keras.ops.convert_to_tensor, data)
+        if log_det_jac:
+            return data, log_det_jac
+        else:
+            return data
+
     def _sample(
         self,
         num_samples: int,
@@ -517,24 +554,14 @@ def log_prob(self, data: Mapping[str, np.ndarray], **kwargs) -> np.ndarray:
         np.ndarray
             Log-probabilities of the distribution `p(inference_variables | inference_conditions, h(summary_conditions))`
         """
-        data, log_det_jac = self.adapter(data, strict=False, stage="inference", log_det_jac=True, **kwargs)
-        log_det_jac = log_det_jac.get("inference_variables", 0.0)
-
-        # Optionally standardize conditions
-        for key in ContinuousApproximator.CONDITION_KEYS:
-            if key in data and key in self.standardize:
-                data[key] = self.standardize_layers[key](data[key])
+        # Adapt, optionally standardize and convert to tensor. Keep track of log_det_jac.
+        data, log_det_jac = self._prepare_data(data, log_det_jac=True, **kwargs)
 
-        # Optionally standardize inference variables
-        if "inference_variables" in self.standardize:
-            data["inference_variables"], log_det_std = self.standardize_layers["inference_variables"](
-                data["inference_variables"], log_det_jac=True
-            )
-            log_det_jac += keras.ops.convert_to_numpy(log_det_std)
-
-        data = keras.tree.map_structure(keras.ops.convert_to_tensor, data)
+        # Pass data to networks and convert back to numpy array.
         log_prob = self._log_prob(**data, **kwargs)
         log_prob = keras.ops.convert_to_numpy(log_prob)
+
+        # Change of variables formula.
         log_prob = log_prob + log_det_jac
 
         return log_prob

bayesflow/approximators/point_approximator.py

@@ -55,8 +55,10 @@ def estimate(
             Each estimator output (i.e., dictionary value that is not itself a dictionary) is an array
             of shape (num_datasets, point_estimate_size, variable_block_size).
         """
-
-        conditions = self._prepare_conditions(conditions, **kwargs)
+        # Adapt, optionally standardize and convert conditions to tensor.
+        conditions = self._prepare_data(conditions, **kwargs)
+        # Remove any superfluous keys, just retain actual conditions.  # TODO: is this necessary?
+        conditions = {k: v for k, v in conditions.items() if k in ContinuousApproximator.CONDITION_KEYS}
 
         estimates = self._estimate(**conditions, **kwargs)
         estimates = self._apply_inverse_adapter_to_estimates(estimates, **kwargs)
@@ -110,9 +112,19 @@ def sample(
             Each output (i.e., dictionary value that is not itself a dictionary) is an array
             of shape (num_datasets, num_samples, variable_block_size).
         """
-        conditions = self._prepare_conditions(conditions, **kwargs)
+        # Adapt, optionally standardize and convert conditions to tensor.
+        conditions = self._prepare_data(conditions, **kwargs)
+        # Remove any superfluous keys, just retain actual conditions.  # TODO: is this necessary?
+        conditions = {k: v for k, v in conditions.items() if k in ContinuousApproximator.CONDITION_KEYS}
 
+        # Sample and undo optional standardization
         samples = self._sample(num_samples, **conditions, **kwargs)
+
+        if "inference_variables" in self.standardize:
+            for score_key in samples.keys():
+                samples[score_key] = self.standardize_layers["inference_variables"](samples[score_key], forward=False)
+
+        samples = {"inference_variables": samples}
         samples = self._apply_inverse_adapter_to_samples(samples, **kwargs)
 
         if split:
@@ -152,20 +164,20 @@ def log_prob(
 
             Log-probabilities have shape (num_datasets,).
         """
-        return super().log_prob(data=data, **kwargs)
+        # Adapt, optionally standardize and convert to tensor. Keep track of log_det_jac
+        data, log_det_jac = self._prepare_data(data, log_det_jac=True, **kwargs)
 
-    def _prepare_conditions(self, conditions: Mapping[str, np.ndarray], **kwargs) -> dict[str, Tensor]:
-        """Adapts, optionally standardizes, and converts the conditions to tensors."""
+        # Pass data to networks and convert back to numpy array
+        log_prob = self._log_prob(**data, **kwargs)
+        log_prob = keras.tree.map_structure(keras.ops.convert_to_numpy, log_prob)
 
-        conditions = self.adapter(conditions, strict=False, stage="inference", **kwargs)
-        conditions = {k: v for k, v in conditions.items() if k in ContinuousApproximator.CONDITION_KEYS}
+        # Change of variables formula, respecting log_prob to be a dictionary
+        if log_det_jac is not None:
+            log_prob = keras.tree.map_structure(lambda x: x + log_det_jac, log_prob)
 
-        # Optionally standardize conditions
-        for key, value in conditions.items():
-            if key in self.standardize:
-                conditions[key] = self.standardize_layers[key](value)
+        log_prob = PointApproximator._squeeze_parametric_score_major_dict(log_prob)
 
-        return keras.tree.map_structure(keras.ops.convert_to_tensor, conditions)
+        return log_prob
 
     def _apply_inverse_adapter_to_estimates(
         self, estimates: Mapping[str, Mapping[str, Tensor]], **kwargs

bayesflow/datasets/disk_dataset.py

@@ -37,6 +37,7 @@ def __init__(
         adapter: Adapter | None,
         stage: str = "training",
         augmentations: Mapping[str, Callable] | Callable = None,
+        shuffle: bool = True,
         **kwargs,
     ):
         """
@@ -67,6 +68,8 @@ def __init__(
 
             Note - augmentations are applied before the adapter is called and are generally
             transforms that you only want to apply during training.
+        shuffle : bool, optional
+            Whether to shuffle the dataset at initialization and at the end of each epoch. Default is True.
         **kwargs
             Additional keyword arguments passed to the base `PyDataset`.
         """
@@ -79,8 +82,9 @@ def __init__(
         self.stage = stage
 
         self.augmentations = augmentations
-
-        self.shuffle()
+        self._shuffle = shuffle
+        if self._shuffle:
+            self.shuffle()
 
     def __getitem__(self, item) -> dict[str, np.ndarray]:
         if not 0 <= item < self.num_batches:
@@ -108,7 +112,8 @@ def __getitem__(self, item) -> dict[str, np.ndarray]:
         return batch
 
     def on_epoch_end(self):
-        self.shuffle()
+        if self._shuffle:
+            self.shuffle()
 
     @property
     def num_batches(self):

bayesflow/datasets/offline_dataset.py

@@ -24,6 +24,7 @@ def __init__(
         *,
         stage: str = "training",
         augmentations: Mapping[str, Callable] | Callable = None,
+        shuffle: bool = True,
         **kwargs,
     ):
         """
@@ -51,6 +52,8 @@ def __init__(
 
             Note - augmentations are applied before the adapter is called and are generally
             transforms that you only want to apply during training.
+        shuffle : bool, optional
+            Whether to shuffle the dataset at initialization and at the end of each epoch. Default is True.
         **kwargs
             Additional keyword arguments passed to the base `PyDataset`.
         """
@@ -69,8 +72,9 @@ def __init__(
         self.indices = np.arange(self.num_samples, dtype="int64")
 
         self.augmentations = augmentations
-
-        self.shuffle()
+        self._shuffle = shuffle
+        if self._shuffle:
+            self.shuffle()
 
     def __getitem__(self, item: int) -> dict[str, np.ndarray]:
         """
@@ -122,7 +126,8 @@ def num_batches(self) -> int | None:
         return int(np.ceil(self.num_samples / self.batch_size))
 
     def on_epoch_end(self) -> None:
-        self.shuffle()
+        if self._shuffle:
+            self.shuffle()
 
     def shuffle(self) -> None:
         """Shuffle the dataset in-place."""

bayesflow/experimental/free_form_flow/free_form_flow.py

@@ -218,10 +218,10 @@ def decode(z):
             return self.decode(z, conditions, training=stage == "training")
 
         # VJP computation
-        z, vjp_fn = vjp(encode, x)
+        z, vjp_fn = vjp(encode, x, return_output=True)
         v1 = vjp_fn(v)[0]
         # JVP computation
-        x_pred, v2 = jvp(decode, (z,), (v,))
+        x_pred, v2 = jvp(decode, (z,), (v,), return_output=True)
 
         # equivalent: surrogate = ops.matmul(ops.stop_gradient(v2[:, None]), v1[:, :, None])[:, 0, 0]
         surrogate = ops.sum((ops.stop_gradient(v2) * v1), axis=-1)

tests/test_approximators/conftest.py

@@ -65,7 +65,7 @@ def point_inference_network_with_multiple_parametric_scores():
 
 
 @pytest.fixture()
-def point_approximator(adapter, point_inference_network, summary_network):
+def point_approximator_with_single_parametric_score(adapter, point_inference_network, summary_network):
     from bayesflow import PointApproximator
 
     return PointApproximator(
@@ -89,7 +89,18 @@ def point_approximator_with_multiple_parametric_scores(
 
 
 @pytest.fixture(
-    params=["continuous_approximator", "point_approximator", "point_approximator_with_multiple_parametric_scores"],
+    params=["point_approximator_with_single_parametric_score", "point_approximator_with_multiple_parametric_scores"]
+)
+def point_approximator(request):
+    return request.getfixturevalue(request.param)
+
+
+@pytest.fixture(
+    params=[
+        "continuous_approximator",
+        "point_approximator_with_single_parametric_score",
+        "point_approximator_with_multiple_parametric_scores",
+    ],
     scope="function",
 )
 def approximator(request):

tests/test_approximators/test_point_approximators/test_log_prob.py

@@ -0,0 +1,30 @@
+import numpy as np
+from bayesflow.scores import ParametricDistributionScore
+from tests.utils import check_combination_simulator_adapter
+
+
+def test_approximator_log_prob(point_approximator, simulator, batch_size, num_samples, adapter):
+    check_combination_simulator_adapter(simulator, adapter)
+
+    data = simulator.sample((batch_size,))
+
+    batch = adapter(data)
+    point_approximator.build_from_data(batch)
+
+    log_prob = point_approximator.log_prob(data=data)
+    parametric_scores = [
+        score
+        for score in point_approximator.inference_network.scores.values()
+        if isinstance(score, ParametricDistributionScore)
+    ]
+
+    if len(parametric_scores) > 1:
+        assert isinstance(log_prob, dict)
+        for score_key, score_log_prob in log_prob.items():
+            assert isinstance(score_log_prob, np.ndarray)
+            assert score_log_prob.shape == (batch_size,)
+
+    # If only one score is available, the outer nesting should be dropped.
+    else:
+        assert isinstance(log_prob, np.ndarray)
+        assert log_prob.shape == (batch_size,)

tests/test_approximators/test_point_approximators/test_sample.py

@@ -20,8 +20,6 @@ def test_approximator_sample(point_approximator, simulator, batch_size, num_samp
 
     assert isinstance(samples, dict)
 
-    print(keras.tree.map_structure(keras.ops.shape, samples))
-
     # Expect doubly nested sample dictionary if more than one samplable score is available.
     scores_for_sampling = [
         score

tests/test_networks/test_inference_networks.py

@@ -150,3 +150,12 @@ def test_save_and_load(tmp_path, inference_network, random_samples, random_condi
     loaded = keras.saving.load_model(tmp_path / "model.keras")
 
     assert_layers_equal(inference_network, loaded)
+
+
+def test_compute_metrics(inference_network, random_samples, random_conditions):
+    xz_shape = keras.ops.shape(random_samples)
+    conditions_shape = keras.ops.shape(random_conditions) if random_conditions is not None else None
+    inference_network.build(xz_shape, conditions_shape)
+
+    metrics = inference_network.compute_metrics(random_samples, conditions=random_conditions)
+    assert "loss" in metrics