Refactor and simplify due to standardize

stefanradev93 · stefanradev93 · commit 38f22289ea40 · 2025-06-01T14:31:49.000-04:00
diff --git a/bayesflow/approximators/continuous_approximator.py b/bayesflow/approximators/continuous_approximator.py
@@ -7,7 +7,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, squeeze_inner_estimates_dict
+from bayesflow.utils import filter_kwargs, logging, split_arrays, squeeze_inner_estimates_dict, concatenate_valid
 from bayesflow.utils.serialization import serialize, deserialize, serializable
 
 from .approximator import Approximator
@@ -180,7 +180,9 @@ def compute_metrics(
 
         summary_metrics, summary_outputs = self._compute_summary_metrics(summary_variables, stage=stage)
 
-        inference_conditions = self._combine_conditions(inference_conditions, summary_outputs, stage=stage)
+        if "inference_conditions" in self.standardize:
+            inference_conditions = self.standardize_layers["inference_conditions"](inference_conditions, stage=stage)
+        inference_conditions = concatenate_valid((inference_conditions, summary_outputs), axis=-1)
 
         inference_variables = self._prepare_inference_variables(inference_variables, stage=stage)
 
@@ -192,6 +194,7 @@ def compute_metrics(
             loss = inference_metrics["loss"] + summary_metrics["loss"]
         else:
             loss = inference_metrics.pop("loss")
+
         inference_metrics = {f"{key}/inference_{key}": value for key, value in inference_metrics.items()}
         summary_metrics = {f"{key}/summary_{key}": value for key, value in summary_metrics.items()}
 
@@ -222,21 +225,6 @@ def _prepare_inference_variables(self, inference_variables: Tensor, stage: str)
 
         return inference_variables
 
-    def _combine_conditions(
-        self, inference_conditions: Tensor | None, summary_outputs: Tensor | None, stage: str
-    ) -> Tensor:
-        """Helper function to combine direct (inference) conditions and outputs of the summary network."""
-        if inference_conditions is None:
-            return summary_outputs
-
-        if "inference_conditions" in self.standardize:
-            inference_conditions = self.standardize_layers["inference_conditions"](inference_conditions, stage=stage)
-
-        if summary_outputs is None:
-            return inference_conditions
-
-        return keras.ops.concatenate([inference_conditions, summary_outputs], axis=-1)
-
     def fit(self, *args, **kwargs):
         """
         Trains the approximator on the provided dataset or on-demand data generated from the given simulator.
@@ -457,24 +445,17 @@ def _sample(
         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.")
+        if (self.summary_network is None) != (summary_variables is None):
+            raise ValueError("Summary variables and summary network must be used together.")
 
+        if self.summary_network is not None:
             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)
+            inference_conditions = concatenate_valid((inference_conditions, summary_outputs), axis=-1)
 
         if inference_conditions is not None:
-            # conditions must always have shape (batch_size, dims)
+            # conditions must always have shape (batch_size, ...)
             batch_size = keras.ops.shape(inference_conditions)[0]
             inference_conditions = keras.ops.expand_dims(inference_conditions, axis=1)
             inference_conditions = keras.ops.broadcast_to(
@@ -485,9 +466,7 @@ def _sample(
             batch_shape = (num_samples,)
 
         return self.inference_network.sample(
-            batch_shape,
-            conditions=inference_conditions,
-            **filter_kwargs(kwargs, self.inference_network.sample),
+            batch_shape, conditions=inference_conditions, **filter_kwargs(kwargs, self.inference_network.sample)
         )
 
     def summaries(self, data: Mapping[str, np.ndarray], **kwargs) -> np.ndarray:
@@ -567,21 +546,14 @@ def _log_prob(
         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.")
+        if (self.summary_network is None) != (summary_variables is None):
+            raise ValueError("Summary variables and summary network must be used together.")
 
+        if self.summary_network is not None:
             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)
+            inference_conditions = concatenate_valid((inference_conditions, summary_outputs), axis=-1)
 
         return self.inference_network.log_prob(
             inference_variables,
diff --git a/bayesflow/approximators/model_comparison_approximator.py b/bayesflow/approximators/model_comparison_approximator.py
@@ -8,7 +8,7 @@
 from bayesflow.networks import SummaryNetwork
 from bayesflow.simulators import ModelComparisonSimulator, Simulator
 from bayesflow.types import Tensor
-from bayesflow.utils import filter_kwargs, logging
+from bayesflow.utils import filter_kwargs, logging, concatenate_valid
 from bayesflow.utils.serialization import serialize, deserialize, serializable
 
 from .approximator import Approximator
@@ -180,7 +180,10 @@ def compute_metrics(
 
         summary_metrics, summary_outputs = self._compute_summary_metrics(summary_variables, stage=stage)
 
-        classifier_conditions = self._combine_conditions(classifier_conditions, summary_outputs, stage=stage)
+        if classifier_conditions is not None and "classifier_conditions" in self.standardize:
+            classifier_conditions = self.standardize_layers["classifier_conditions"](classifier_conditions, stage=stage)
+
+        classifier_conditions = concatenate_valid((classifier_conditions, summary_outputs), axis=-1)
 
         logits = self._compute_logits(classifier_conditions)
         cross_entropy = keras.ops.mean(keras.losses.categorical_crossentropy(model_indices, logits, from_logits=True))
@@ -193,49 +196,17 @@ def compute_metrics(
                 metric.name: metric(model_indices, predictions) for metric in self.classifier_network.metrics
             }
 
-        loss = classifier_metrics.get("loss") + summary_metrics.get("loss", keras.ops.zeros(()))
+        if "loss" in summary_metrics:
+            loss = classifier_metrics["loss"] + summary_metrics["loss"]
+        else:
+            loss = classifier_metrics.pop("loss")
 
         classifier_metrics = {f"{key}/classifier_{key}": value for key, value in classifier_metrics.items()}
         summary_metrics = {f"{key}/summary_{key}": value for key, value in summary_metrics.items()}
 
         metrics = {"loss": loss} | classifier_metrics | summary_metrics
         return metrics
 
-    def _compute_summary_metrics(self, summary_variables: Tensor, stage: str) -> tuple[dict, Tensor | None]:
-        """Helper function to compute summary metrics and outputs."""
-        if self.summary_network is None:
-            return {}, None
-        if summary_variables is None:
-            raise ValueError("Summary variables are required when a summary network is present.")
-
-        if "summary_variables" in self.standardize:
-            summary_variables = self.standardize_layers["summary_variables"](summary_variables, stage=stage)
-
-        summary_metrics = self.summary_network.compute_metrics(summary_variables, stage=stage)
-        summary_outputs = summary_metrics.pop("outputs")
-        return summary_metrics, summary_outputs
-
-    def _combine_conditions(
-        self, classifier_conditions: Tensor | None, summary_outputs: Tensor | None, stage
-    ) -> Tensor:
-        """Helper to combine classifier conditions and summary outputs, if present."""
-        if classifier_conditions is None:
-            return summary_outputs
-
-        if "classifier_conditions" in self.standardize:
-            classifier_conditions = self.standardize_layers["inference_conditions"](classifier_conditions, stage=stage)
-
-        if summary_outputs is None:
-            return classifier_conditions
-
-        return keras.ops.concatenate([classifier_conditions, summary_outputs], axis=-1)
-
-    def _compute_logits(self, classifier_conditions: Tensor) -> Tensor:
-        """Helper to compute projected logits from the classifier network."""
-        logits = self.classifier_network(classifier_conditions)
-        logits = self.logits_projector(logits)
-        return logits
-
     def fit(
         self,
         *,
@@ -352,7 +323,7 @@ def predict(
         self,
         *,
         conditions: Mapping[str, np.ndarray],
-        logits: bool = False,
+        probs: bool = True,
         **kwargs,
     ) -> np.ndarray:
         """
@@ -363,15 +334,14 @@ def predict(
         ----------
         conditions : Mapping[str, np.ndarray]
             Dictionary of conditioning variables as NumPy arrays.
-        logits: bool, default=False
-            Should the posterior model probabilities be on the (unconstrained) logit space?
-            If `False`, the output is a unit simplex instead.
+        probs: bool, optional
+            A flag indicating whether model probabilities (True) or logits (False) are returned. Default is True.
         **kwargs : dict
-            Additional keyword arguments for the adapter and classification process.
+            Additional keyword arguments for the adapter and classifier.
 
         Returns
         -------
-        np.ndarray
+        outputs: np.ndarray
             Predicted posterior model probabilities given `conditions`.
         """
 
@@ -389,34 +359,7 @@ def predict(
 
         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
+        return keras.ops.convert_to_numpy(keras.ops.softmax(output) if probs else output)
 
     def summaries(self, data: Mapping[str, np.ndarray], **kwargs) -> np.ndarray:
         """
@@ -449,6 +392,40 @@ def summaries(self, data: Mapping[str, np.ndarray], **kwargs) -> np.ndarray:
 
         return summaries
 
+    def _compute_logits(self, classifier_conditions: Tensor) -> Tensor:
+        """Helper to compute projected logits from the classifier network."""
+        logits = self.classifier_network(classifier_conditions)
+        logits = self.logits_projector(logits)
+        return logits
+
+    def _predict(self, classifier_conditions: Tensor = None, summary_variables: Tensor = None, **kwargs) -> Tensor:
+        """Helper method to obtain logits from the internal classifier based on conditions."""
+        if (self.summary_network is None) != (summary_variables is None):
+            raise ValueError("Summary variables and summary network must be used together.")
+
+        if self.summary_network is not None:
+            summary_outputs = self.summary_network(
+                summary_variables, **filter_kwargs(kwargs, self.summary_network.call)
+            )
+            classifier_conditions = concatenate_valid((classifier_conditions, summary_outputs), axis=-1)
+
+        logits = self._compute_logits(classifier_conditions)
+        return logits
+
+    def _compute_summary_metrics(self, summary_variables: Tensor, stage: str) -> tuple[dict, Tensor | None]:
+        """Helper function to compute summary metrics and outputs."""
+        if self.summary_network is None:
+            return {}, None
+        if summary_variables is None:
+            raise ValueError("Summary variables are required when a summary network is present.")
+
+        if "summary_variables" in self.standardize:
+            summary_variables = self.standardize_layers["summary_variables"](summary_variables, stage=stage)
+
+        summary_metrics = self.summary_network.compute_metrics(summary_variables, stage=stage)
+        summary_outputs = summary_metrics.pop("outputs")
+        return summary_metrics, summary_outputs
+
     def _batch_size_from_data(self, data: Mapping[str, any]) -> int:
         """
         Fetches the current batch size from an input dictionary. Can only be used during training when
diff --git a/bayesflow/approximators/point_approximator.py b/bayesflow/approximators/point_approximator.py
@@ -5,7 +5,7 @@
 import keras
 
 from bayesflow.types import Tensor
-from bayesflow.utils import filter_kwargs, split_arrays, squeeze_inner_estimates_dict, logging
+from bayesflow.utils import filter_kwargs, split_arrays, squeeze_inner_estimates_dict, logging, concatenate_valid
 from bayesflow.utils.serialization import serializable
 
 from .continuous_approximator import ContinuousApproximator
@@ -57,11 +57,14 @@ def estimate(
         """
 
         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 = PointApproximator._reorder_estimates(estimates)
         # Remove unnecessary nesting.
@@ -108,9 +111,10 @@ def sample(
             of shape (num_datasets, num_samples, variable_block_size).
         """
         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:
             raise NotImplementedError("split=True is currently not supported for `PointApproximator`.")
 
@@ -148,18 +152,19 @@ def log_prob(
 
             Log-probabilities have shape (num_datasets,).
         """
-        log_prob = super().log_prob(data=data, **kwargs)
-        # Squeeze log probabilities dictionary if there's only one key-value pair.
-        log_prob = PointApproximator._squeeze_parametric_score_major_dict(log_prob)
-
-        return log_prob
+        return super().log_prob(data=data, **kwargs)
 
     def _prepare_conditions(self, conditions: Mapping[str, np.ndarray], **kwargs) -> dict[str, Tensor]:
-        """Adapts and converts the conditions to tensors."""
+        """Adapts, optionally standardizes, and converts the conditions to tensors."""
 
         conditions = self.adapter(conditions, strict=False, stage="inference", **kwargs)
         conditions = {k: v for k, v in conditions.items() if k in ContinuousApproximator.CONDITION_KEYS}
 
+        # Optionally standardize conditions
+        for key, value in conditions.items():
+            if key in self.standardize:
+                conditions[key] = self.standardize_layers[key](value)
+
         return keras.tree.map_structure(keras.ops.convert_to_tensor, conditions)
 
     def _apply_inverse_adapter_to_estimates(
@@ -233,7 +238,7 @@ def _squeeze_estimates(
     def _squeeze_parametric_score_major_dict(samples: Mapping[str, np.ndarray]) -> np.ndarray or dict[str, np.ndarray]:
         """Squeezes the 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 next(iter(samples.values()))
         return samples
 
     def _estimate(
@@ -242,21 +247,14 @@ def _estimate(
         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.")
+        if (self.summary_network is None) != (summary_variables is None):
+            raise ValueError("Summary variables and summary network must be used together.")
 
+        if self.summary_network is not None:
             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)
+            inference_conditions = concatenate_valid((inference_conditions, summary_outputs), axis=-1)
 
         return self.inference_network(
             conditions=inference_conditions,
