Flexible implementation of point estimation

README.md

@@ -100,7 +100,8 @@ Check out some of our walk-through notebooks below. We are actively working on p
 5. [Hyperparameter optimization](examples/Hyperparameter_Optimization.ipynb)
 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...
+8. [Rapid iteration with point estimation and expert statistics for Lotka-Volterra dynamics](examples/Lotka_Volterra_point_estimation_and_expert_stats.ipynb)
+9. More coming soon...
 
 ## Documentation \& Help
 

bayesflow/__init__.py

@@ -11,8 +11,9 @@
     workflows,
     utils,
 )
+
 from .adapters import Adapter
-from .approximators import ContinuousApproximator
+from .approximators import ContinuousApproximator, PointApproximator
 from .datasets import OfflineDataset, OnlineDataset, DiskDataset
 from .simulators import make_simulator
 from .workflows import BasicWorkflow

bayesflow/approximators/__init__.py

@@ -1,5 +1,6 @@
 from .approximator import Approximator
 from .continuous_approximator import ContinuousApproximator
+from .point_approximator import PointApproximator
 from .model_comparison_approximator import ModelComparisonApproximator
 
 from ..utils._docs import _add_imports_to_all

bayesflow/approximators/continuous_approximator.py

@@ -11,7 +11,7 @@
 from bayesflow.adapters import Adapter
 from bayesflow.networks import InferenceNetwork, SummaryNetwork
 from bayesflow.types import Tensor
-from bayesflow.utils import filter_kwargs, logging, split_arrays
+from bayesflow.utils import filter_kwargs, logging, split_arrays, squeeze_inner_estimates_dict
 from .approximator import Approximator
 
 
@@ -120,6 +120,8 @@ def compute_metrics(
             else:
                 inference_conditions = keras.ops.concatenate([inference_conditions, summary_outputs], axis=-1)
 
+        # Force a conversion to Tensor
+        inference_variables = keras.tree.map_structure(keras.ops.convert_to_tensor, inference_variables)
         inference_metrics = self.inference_network.compute_metrics(
             inference_variables, conditions=inference_conditions, stage=stage
         )
@@ -205,6 +207,44 @@ def get_config(self):
 
         return base_config | config
 
+    def estimate(
+        self,
+        conditions: dict[str, np.ndarray],
+        split: bool = False,
+        estimators: dict[str, callable] = None,
+        num_samples: int = 1000,
+        **kwargs,
+    ) -> dict[str, dict[str, np.ndarray]]:
+        estimators = estimators or {}
+        estimators = (
+            dict(
+                mean=lambda x, axis: dict(value=np.mean(x, keepdims=True, axis=axis)),
+                median=lambda x, axis: dict(value=np.median(x, keepdims=True, axis=axis)),
+                quantiles=lambda x, axis: dict(value=np.moveaxis(np.quantile(x, q=[0.1, 0.5, 0.9], axis=axis), 0, 1)),
+            )
+            | estimators
+        )
+
+        samples = self.sample(num_samples=num_samples, conditions=conditions, split=split, **kwargs)
+
+        estimates = {
+            variable_name: {
+                estimator_name: func(samples[variable_name], axis=1) for estimator_name, func in estimators.items()
+            }
+            for variable_name in samples.keys()
+        }
+
+        # remove unnecessary nesting
+        estimates = {
+            variable_name: {
+                outer_key: squeeze_inner_estimates_dict(estimates[variable_name][outer_key])
+                for outer_key in estimates[variable_name].keys()
+            }
+            for variable_name in estimates.keys()
+        }
+
+        return estimates
+
     def sample(
         self,
         *,

bayesflow/approximators/point_approximator.py

@@ -0,0 +1,154 @@
+import keras
+import numpy as np
+from keras.saving import (
+    register_keras_serializable as serializable,
+)
+
+from bayesflow.types import Tensor
+from bayesflow.utils import filter_kwargs, split_arrays, squeeze_inner_estimates_dict
+from .continuous_approximator import ContinuousApproximator
+
+
+@serializable(package="bayesflow.approximators")
+class PointApproximator(ContinuousApproximator):
+    """
+    A workflow for fast amortized point estimation of a conditional distribution.
+
+    The distribution is approximated by point estimators, parameterized by a feed-forward `PointInferenceNetwork`.
+    Conditions can be compressed by an optional `SummaryNetwork` or used directly as input to the inference network.
+    """
+
+    def estimate(
+        self,
+        conditions: dict[str, np.ndarray],
+        split: bool = False,
+        **kwargs,
+    ) -> dict[str, dict[str, np.ndarray]]:
+        conditions = self._prepare_conditions(conditions, **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)
+        # Reorder the nested dictionary so that original variable names are at the top.
+        estimates = self._reorder_estimates(estimates)
+        # Remove unnecessary nesting.
+        estimates = self._squeeze_estimates(estimates)
+
+        return estimates
+
+    def sample(
+        self,
+        *,
+        num_samples: int,
+        conditions: dict[str, np.ndarray],
+        split: bool = False,
+        **kwargs,
+    ) -> dict[str, np.ndarray]:
+        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 _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 _estimate output dictionary."""
+        estimates = keras.tree.map_structure(keras.ops.convert_to_numpy, estimates)
+        processed = {}
+        for score_key, score_val in estimates.items():
+            processed[score_key] = {}
+            for head_key, estimate in score_val.items():
+                adapted = self.adapter(
+                    {"inference_variables": estimate},
+                    inverse=True,
+                    strict=False,
+                    **kwargs,
+                )
+                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]]]:
+        """Reorders the nested dictionary so that the inference variable names become the top-level keys."""
+        # Grab the variable names from one sample inner dictionary.
+        sample_inner = next(iter(next(iter(estimates.values())).values()))
+        variable_names = sample_inner.keys()
+        reordered = {}
+        for variable in variable_names:
+            reordered[variable] = {}
+            for score_key, inner_dict in estimates.items():
+                reordered[variable][score_key] = {inner_key: value[variable] for inner_key, value in inner_dict.items()}
+        return reordered
+
+    def _squeeze_estimates(
+        self, estimates: dict[str, dict[str, dict[str, np.ndarray]]]
+    ) -> dict[str, dict[str, np.ndarray]]:
+        """Squeezes each inner estimate dictionary to remove unnecessary nesting."""
+        squeezed = {}
+        for variable, variable_estimates in estimates.items():
+            squeezed[variable] = {
+                score_key: squeeze_inner_estimates_dict(inner_estimate)
+                for score_key, inner_estimate in variable_estimates.items()
+            }
+        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,
+        summary_variables: Tensor = None,
+        **kwargs,
+    ) -> dict[str, dict[str, 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 inference_conditions is None:
+                inference_conditions = summary_outputs
+            else:
+                inference_conditions = keras.ops.concatenate([inference_conditions, summary_outputs], axis=1)
+
+        return self.inference_network(
+            conditions=inference_conditions,
+            **filter_kwargs(kwargs, self.inference_network.call),
+        )

bayesflow/diagnostics/__init__.py

@@ -2,6 +2,7 @@
 
 from .plots import (
     calibration_ecdf,
+    calibration_ecdf_from_quantiles,
     calibration_histogram,
     loss,
     mc_calibration,
@@ -10,6 +11,7 @@
     pairs_posterior,
     pairs_samples,
     recovery,
+    recovery_from_estimates,
     z_score_contraction,
 )
 

bayesflow/diagnostics/plots/__init__.py

@@ -1,4 +1,5 @@
 from .calibration_ecdf import calibration_ecdf
+from .calibration_ecdf_from_quantiles import calibration_ecdf_from_quantiles
 from .calibration_histogram import calibration_histogram
 from .loss import loss
 from .mc_calibration import mc_calibration
@@ -7,4 +8,5 @@
 from .pairs_posterior import pairs_posterior
 from .pairs_samples import pairs_samples
 from .recovery import recovery
+from .recovery_from_estimates import recovery_from_estimates
 from .z_score_contraction import z_score_contraction

bayesflow/diagnostics/plots/calibration_ecdf.py

@@ -163,12 +163,12 @@ def calibration_ecdf(
             plot_data["axes"].flat[j].plot(xx, yy, color=rank_ecdf_color, alpha=0.95, label="Rank ECDF")
 
     # Compute uniform ECDF and bands
-    alpha, z, L, H = simultaneous_ecdf_bands(estimates.shape[0], **kwargs.pop("ecdf_bands_kwargs", {}))
+    alpha, z, L, U = simultaneous_ecdf_bands(estimates.shape[0], **kwargs.pop("ecdf_bands_kwargs", {}))
 
     # Difference, if specified
     if difference:
         L -= z
-        H -= z
+        U -= z
         ylab = "ECDF Difference"
     else:
         ylab = "ECDF"
@@ -182,7 +182,7 @@ def calibration_ecdf(
         titles = ["Stacked ECDFs"]
 
     for ax, title in zip(plot_data["axes"].flat, titles):
-        ax.fill_between(z, L, H, color=fill_color, alpha=0.2, label=rf"{int((1 - alpha) * 100)}$\%$ Confidence Bands")
+        ax.fill_between(z, L, U, color=fill_color, alpha=0.2, label=rf"{int((1 - alpha) * 100)}$\%$ Confidence Bands")
         ax.legend(fontsize=legend_fontsize)
         ax.set_title(title, fontsize=title_fontsize)