adjust signatures, extend tests

Author: vpratz

bayesflow/approximators/model_comparison_approximator.py

@@ -106,18 +106,16 @@ def build_dataset(
     def compile(
         self,
         *args,
-        classifier_metrics: Sequence[keras.Metric] = None,
-        summary_metrics: Sequence[keras.Metric] = None,
         **kwargs,
     ):
-        if classifier_metrics:
+        if "classifier_metrics" in kwargs:
             warnings.warn(
                 "Supplying classifier metrics to the approximator is no longer supported. "
                 "Please pass the metrics directly to the network using the metrics parameter.",
                 DeprecationWarning,
             )
 
-        if summary_metrics:
+        if "summary_metrics" in kwargs:
             warnings.warn(
                 "Supplying summary metrics to the approximator is no longer supported. "
                 "Please pass the metrics directly to the network using the metrics parameter.",
@@ -166,8 +164,10 @@ def compute_metrics(
             classifier_metrics |= {
                 metric.name: metric(model_indices, predictions) for metric in self.classifier_network.metrics
             }
-
-        loss = classifier_metrics.get("loss", keras.ops.zeros(())) + 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()}

bayesflow/networks/inference_network.py

@@ -9,7 +9,7 @@
 
 @serializable("bayesflow.networks")
 class InferenceNetwork(keras.Layer):
-    def __init__(self, base_distribution: str = "normal", *, metrics: Sequence[keras.Metric] = None, **kwargs):
+    def __init__(self, base_distribution: str = "normal", *, metrics: Sequence[keras.Metric] | None = None, **kwargs):
         self.custom_metrics = metrics
         super().__init__(**layer_kwargs(kwargs))
         self.base_distribution = find_distribution(base_distribution)

bayesflow/networks/summary_network.py

@@ -10,7 +10,7 @@
 
 @serializable("bayesflow.networks")
 class SummaryNetwork(keras.Layer):
-    def __init__(self, base_distribution: str = None, *, metrics: Sequence[keras.Metric] = None, **kwargs):
+    def __init__(self, base_distribution: str = None, *, metrics: Sequence[keras.Metric] | None = None, **kwargs):
         self.custom_metrics = metrics
         super().__init__(**layer_kwargs(kwargs))
         self.base_distribution = find_distribution(base_distribution)

tests/test_approximators/conftest.py

@@ -20,8 +20,11 @@ def summary_network():
 @pytest.fixture()
 def inference_network():
     from bayesflow.networks import CouplingFlow
+    from bayesflow.metrics import RootMeanSquaredError
 
-    return CouplingFlow(subnet="mlp", depth=2, subnet_kwargs=dict(widths=(32, 32)))
+    return CouplingFlow(
+        subnet="mlp", depth=2, subnet_kwargs=dict(widths=(32, 32)), metrics=[RootMeanSquaredError(name="rmse")]
+    )
 
 
 @pytest.fixture()

tests/test_approximators/test_fit.py

@@ -49,3 +49,7 @@ def test_loss_progress(approximator, train_dataset, validation_dataset):
     # check that the shown loss is not nan or zero
     assert re.search(r"\bnan\b", output) is None, "found nan in output"
     assert re.search(r"\bloss: 0\.0000e\+00\b", output) is None, "found zero loss in output"
+
+    # check that additional metric is present
+    assert "val_rmse/inference_rmse" in output, "custom metric (RMSE) not shown"
+    assert re.search(r"\bval_rmse/inference_rmse: \d+\.\d+", output) is not None, "custom metric not correctly shown"

tests/test_approximators/test_model_comparison_approximator/conftest.py

@@ -51,15 +51,17 @@ def adapter():
 @pytest.fixture
 def summary_network():
     from bayesflow.networks import DeepSet
+    from bayesflow.metrics import RootMeanSquaredError
 
-    return DeepSet(summary_dim=2, depth=1)
+    return DeepSet(summary_dim=2, depth=1, base_distribution="normal", metrics=[RootMeanSquaredError(name="rmse")])
 
 
 @pytest.fixture
 def classifier_network():
     from bayesflow.networks import MLP
+    from keras.metrics import CategoricalAccuracy
 
-    return MLP(widths=[32, 32])
+    return MLP(widths=[32, 32], metrics=[CategoricalAccuracy(name="categorical_accuracy")])
 
 
 @pytest.fixture

tests/test_approximators/test_model_comparison_approximator/test_model_comparison_approximator.py

@@ -55,7 +55,10 @@ def test_fit(approximator, train_dataset, validation_dataset):
 
         output = stream.getvalue()
     # check that the loss is shown
-    assert "loss" in output
+    assert "loss/summary_loss" in output
+    assert "loss/classifier_loss" in output
+    assert "val_categorical_accuracy/classifier_categorical_accuracy" in output
+    assert "val_rmse/summary_rmse" in output
 
 
 def test_save_and_load(tmp_path, approximator, train_dataset, validation_dataset):