Sample from parametric scoring rules; more refactoring in PointApproximator

Author: han-ol

bayesflow/approximators/point_approximator.py

@@ -27,12 +27,9 @@ def estimate(
         if not self.built:
             raise AssertionError("PointApproximator needs to be built before predicting with it.")
 
-        # Prepare the input conditions.
         conditions = self._prepare_conditions(conditions, **kwargs)
-        # Run the internal estimation and convert the output to numpy.
-        estimates = self._run_inference(conditions, **kwargs)
-        # Postprocess the inference output with the inverse adapter.
-        estimates = self._apply_inverse_adapter(estimates, **kwargs)
+        estimates = self._estimate(**conditions, **kwargs)
+        estimates = self._apply_inverse_adapter_to_estimates(estimates, **kwargs)
         # Optionally split the arrays along the last axis.
         if split:
             estimates = split_arrays(estimates, axis=-1)
@@ -43,25 +40,40 @@ def estimate(
 
         return estimates
 
+    def sample(
+        self,
+        *,
+        num_samples: int,
+        conditions: dict[str, np.ndarray],
+        split: bool = False,
+        **kwargs,
+    ) -> dict[str, np.ndarray]:
+        if not self.built:
+            raise AssertionError("This model needs to be built before using it for sampling.")
+
+        conditions = self._prepare_conditions(conditions, **kwargs)
+        samples = self._sample(num_samples, **conditions, **kwargs)
+        samples = self._apply_inverse_adapter_to_samples(samples, **kwargs)
+        # Optionally split the arrays along the last axis.
+        if split:
+            samples = split_arrays(samples, axis=-1)
+        # Squeeze samples if there's only one key-value pair.
+        samples = self._squeeze_samples(samples)
+
+        return samples
+
     def _prepare_conditions(self, conditions: dict[str, np.ndarray], **kwargs) -> dict[str, Tensor]:
         """Adapts and converts the conditions to tensors."""
         conditions = self.adapter(conditions, strict=False, stage="inference", **kwargs)
         return keras.tree.map_structure(keras.ops.convert_to_tensor, conditions)
 
-    def _run_inference(self, conditions: dict[str, Tensor], **kwargs) -> dict[str, dict[str, np.ndarray]]:
-        """Runs the internal _estimate function and converts the result to numpy arrays."""
-        # Run the estimation.
-        inference_output = self._estimate(**conditions, **kwargs)
-        # Wrap the result in a dict and convert to numpy.
-        wrapped_output = {"inference_variables": inference_output}
-        return keras.tree.map_structure(keras.ops.convert_to_numpy, wrapped_output)
-
-    def _apply_inverse_adapter(
-        self, estimates: dict[str, dict[str, np.ndarray]], **kwargs
+    def _apply_inverse_adapter_to_estimates(
+        self, estimates: dict[str, dict[str, Tensor]], **kwargs
     ) -> dict[str, dict[str, dict[str, np.ndarray]]]:
-        """Applies the inverse adapter on each inner element of the inference outputs."""
+        """Applies the inverse adapter on each inner element of the _estimate output dictionary."""
+        estimates = keras.tree.map_structure(keras.ops.convert_to_numpy, estimates)
         processed = {}
-        for score_key, score_val in estimates["inference_variables"].items():
+        for score_key, score_val in estimates.items():
             processed[score_key] = {}
             for head_key, estimate in score_val.items():
                 adapted = self.adapter(
@@ -73,6 +85,21 @@ def _apply_inverse_adapter(
                 processed[score_key][head_key] = adapted
         return processed
 
+    def _apply_inverse_adapter_to_samples(
+        self, samples: dict[str, Tensor], **kwargs
+    ) -> dict[str, dict[str, np.ndarray]]:
+        """Applies the inverse adapter to a dictionary of samples."""
+        samples = keras.tree.map_structure(keras.ops.convert_to_numpy, samples)
+        processed = {}
+        for score_key, samples in samples.items():
+            processed[score_key] = self.adapter(
+                {"inference_variables": samples},
+                inverse=True,
+                strict=False,
+                **kwargs,
+            )
+        return processed
+
     def _reorder_estimates(
         self, estimates: dict[str, dict[str, dict[str, np.ndarray]]]
     ) -> dict[str, dict[str, dict[str, np.ndarray]]]:
@@ -99,6 +126,12 @@ def _squeeze_estimates(
             }
         return squeezed
 
+    def _squeeze_samples(self, samples: dict[str, np.ndarray]) -> np.ndarray or dict[str, np.ndarray]:
+        """Squeezes the samples dictionary to just the value if there is only one key-value pair."""
+        if len(samples) == 1:
+            return next(iter(samples.values()))  # Extract and return the only item's value
+        return samples
+
     def _estimate(
         self,
         inference_conditions: Tensor = None,

bayesflow/networks/point_inference_network.py

@@ -165,8 +165,29 @@ def compute_metrics(self, x: Tensor, conditions: Tensor = None, stage: str = "tr
 
     # WIP: untested draft of sample method
     @allow_batch_size
-    def sample(self, batch_shape: Shape, conditions: Tensor = None, **kwargs) -> dict[str, Tensor]:
-        output = self.subnet(conditions)
+    def sample(self, batch_shape: Shape, conditions: Tensor = None) -> dict[str, Tensor]:
+        """
+        Parameters
+        ----------
+        batch_shape : tuple,
+            Expected dimensions depend on `conditions`
+            - conditional sampling: (batch_size, num_samples) if `conditions` is a tensor
+              of shape (batch_size, num_samples)
+            - unconditional sampling: (num_samples,) if `conditions` is None
+        conditions : Tensor or None, default None
+            Optional inference conditions. If `conditions` is not given, the method will return unconditional samples.
+
+        Returns
+        -------
+        samples : dict[str, Tensor]
+            Samples for every parametric scoring rule. Dict values have shape (batch_size, num_samples, num_variables)
+            or (num_samples, num_variables) for conditional or unconditional sampling respectively.
+        """
+        if conditions is None:  # unconditional estimation uses a fixed input vector
+            conditions = keras.ops.ones(batch_shape, dtype="float32").reshape(1, -1, 1)
+
+        # conditions are duplicated along axis 1 num_sample times
+        output = self.subnet(conditions[:, 0, :])
         samples = {}
 
         for score_key, score in self.scores.items():

tests/test_approximators/conftest.py

@@ -6,6 +6,11 @@ def batch_size():
     return 8
 
 
+@pytest.fixture()
+def num_samples():
+    return 100
+
+
 @pytest.fixture()
 def summary_network():
     return None
@@ -32,18 +37,32 @@ def continuous_approximator(adapter, inference_network, summary_network):
 @pytest.fixture()
 def point_inference_network():
     from bayesflow.networks import PointInferenceNetwork
-    from bayesflow.scores import NormedDifferenceScore, QuantileScore
+    from bayesflow.scores import NormedDifferenceScore, QuantileScore, MultivariateNormalScore
 
     return PointInferenceNetwork(
         scores=dict(
             mean=NormedDifferenceScore(k=2),
             quantiles=QuantileScore(q=[0.1, 0.5, 0.9]),
+            mvn=MultivariateNormalScore(),
         ),
         subnet="mlp",
         subnet_kwargs=dict(widths=(32, 32)),
     )
 
 
+@pytest.fixture()
+def point_inference_network_with_multiple_parametric_scores():
+    from bayesflow.networks import PointInferenceNetwork
+    from bayesflow.scores import MultivariateNormalScore
+
+    return PointInferenceNetwork(
+        scores=dict(
+            mvn1=MultivariateNormalScore(),
+            mvn2=MultivariateNormalScore(),
+        ),
+    )
+
+
 @pytest.fixture()
 def point_approximator(adapter, point_inference_network, summary_network):
     from bayesflow import PointApproximator
@@ -55,8 +74,23 @@ def point_approximator(adapter, point_inference_network, summary_network):
     )
 
 
-# @pytest.fixture(params=["continuous_approximator"], scope="function")
-@pytest.fixture(params=["continuous_approximator", "point_approximator"], scope="function")
+@pytest.fixture()
+def point_approximator_with_multiple_parametric_scores(
+    adapter, point_inference_network_with_multiple_parametric_scores, summary_network
+):
+    from bayesflow import PointApproximator
+
+    return PointApproximator(
+        adapter=adapter,
+        inference_network=point_inference_network_with_multiple_parametric_scores,
+        summary_network=summary_network,
+    )
+
+
+@pytest.fixture(
+    params=["continuous_approximator", "point_approximator", "point_approximator_with_multiple_parametric_scores"],
+    scope="function",
+)
 def approximator(request):
     return request.getfixturevalue(request.param)
 

tests/test_approximators/test_point_approximators/test_sample.py

@@ -0,0 +1,35 @@
+import keras
+import numpy as np
+from bayesflow.scores import ParametricDistributionScore
+
+
+def test_approximator_sample(point_approximator, simulator, batch_size, num_samples, adapter):
+    data = simulator.sample((batch_size,))
+
+    batch = adapter(data)
+    point_approximator.build_from_data(batch)
+
+    samples = point_approximator.sample(num_samples=num_samples, conditions=data)
+
+    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
+        for score in point_approximator.inference_network.scores.values()
+        if isinstance(score, ParametricDistributionScore)
+    ]
+
+    if len(scores_for_sampling) > 1:
+        for score_key, score_samples in samples.items():
+            for variable, variable_estimates in score_samples.items():
+                assert isinstance(variable_estimates, np.ndarray)
+                assert variable_estimates.shape[:-1] == (batch_size, num_samples)
+
+    # If only one score is available, the outer nesting should be dropped.
+    else:
+        for variable, variable_estimates in samples.items():
+            assert isinstance(variable_estimates, np.ndarray)
+            assert variable_estimates.shape[:-1] == (batch_size, num_samples)