Improve workflow tests with multiple summary nets / approximators

Author: stefanradev93

tests/test_workflows/conftest.py

@@ -1,15 +1,52 @@
 import pytest
 
+import keras
 
-@pytest.fixture()
-def inference_network():
-    from bayesflow.networks import CouplingFlow
+from bayesflow.utils.serialization import serializable
 
-    return CouplingFlow(depth=2)
 
+@pytest.fixture(params=["coupling_flow", "flow_matching"])
+def inference_network(request):
+    if request.param == "coupling_flow":
+        from bayesflow.networks import CouplingFlow
 
-@pytest.fixture()
-def summary_network():
-    from bayesflow.networks import TimeSeriesTransformer
+        return CouplingFlow(depth=2)
 
-    return TimeSeriesTransformer(embed_dims=(8, 8), mlp_widths=(32, 32), mlp_depths=(1, 1))
+    elif request.param == "flow_matching":
+        from bayesflow.networks import FlowMatching
+
+        return FlowMatching(subnet_kwargs=dict(widths=(32, 32)), use_optimal_transport=False)
+
+
+@pytest.fixture(params=["time_series_transformer", "fusion_transformer", "time_series_network", "custom"])
+def summary_network(request):
+    if request.param == "time_series_transformer":
+        from bayesflow.networks import TimeSeriesTransformer
+
+        return TimeSeriesTransformer(embed_dims=(8, 8), mlp_widths=(16, 8), mlp_depths=(1, 1))
+
+    elif request.param == "fusion_transformer":
+        from bayesflow.networks import FusionTransformer
+
+        return FusionTransformer(
+            embed_dims=(8, 8), mlp_widths=(8, 16), mlp_depths=(2, 1), template_dim=8, bidirectional=False
+        )
+
+    elif request.param == "time_series_network":
+        from bayesflow.networks import TimeSeriesNetwork
+
+        return TimeSeriesNetwork(filters=4, skip_steps=2)
+
+    elif request.param == "custom":
+        from bayesflow.networks import SummaryNetwork
+
+        @serializable
+        class Custom(SummaryNetwork):
+            def __init__(self, **kwargs):
+                super().__init__(**kwargs)
+                self.inner = keras.Sequential([keras.layers.LSTM(8), keras.layers.Dense(4)])
+
+            def call(self, x, **kwargs):
+                return self.inner(x, training=kwargs.get("stage") == "training")
+
+        return Custom()

tests/test_workflows/test_basic_workflow.py

@@ -1,21 +1,34 @@
+import os
+
+import keras
+
 import bayesflow as bf
 
 
-def test_basic_workflow(inference_network, summary_network):
+def test_basic_workflow(tmp_path, inference_network, summary_network):
     workflow = bf.BasicWorkflow(
         inference_network=inference_network,
         summary_network=summary_network,
         inference_variables=["parameters"],
         summary_variables=["observables"],
         simulator=bf.simulators.SIR(),
+        checkpoint_filepath=str(tmp_path),
     )
 
-    history = workflow.fit_online(epochs=2, batch_size=32, num_batches_per_epoch=2)
-    plots = workflow.plot_default_diagnostics(test_data=50, num_samples=50)
-    metrics = workflow.compute_default_diagnostics(test_data=50, num_samples=50, variable_names=["p1", "p2"])
+    # Ensure metrics work fine
+    history = workflow.fit_online(epochs=4, batch_size=8, num_batches_per_epoch=2, verbose=0)
+    plots = workflow.plot_default_diagnostics(test_data=50, num_samples=25)
+    metrics = workflow.compute_default_diagnostics(test_data=50, num_samples=25, variable_names=["p1", "p2"])
 
     assert "loss" in list(history.history.keys())
-    assert len(history.history["loss"]) == 2
+    assert len(history.history["loss"]) == 4
     assert list(plots.keys()) == ["losses", "recovery", "calibration_ecdf", "z_score_contraction"]
     assert list(metrics.columns) == ["p1", "p2"]
     assert metrics.values.shape == (3, 2)
+
+    # Ensure saving and loading from workflow works fine
+    loaded_approximator = keras.saving.load_model(os.path.join(str(tmp_path), "model.keras"))
+
+    # Get samples
+    samples = loaded_approximator.sample(conditions=workflow.simulate(5), num_samples=3)
+    assert samples["parameters"].shape == (5, 3, 2)