add tests for model comparison approximator

Author: vpratz

tests/test_approximators/test_model_comparison_approximator/__init__.py

@@ -0,0 +1,74 @@
+import pytest
+import numpy as np
+
+
+@pytest.fixture
+def simulator():
+    from bayesflow import make_simulator
+    from bayesflow.simulators import ModelComparisonSimulator
+
+    def context(batch_shape, n=None):
+        if n is None:
+            n = np.random.randint(2, 5)
+        return dict(n=n)
+
+    def prior_null():
+        return dict(mu=0.0)
+
+    def prior_alternative():
+        mu = np.random.normal(loc=0, scale=1)
+        return dict(mu=mu)
+
+    def likelihood(n, mu):
+        x = np.random.normal(loc=mu, scale=1, size=n)
+        return dict(x=x)
+
+    simulator_null = make_simulator([prior_null, likelihood])
+    simulator_alternative = make_simulator([prior_alternative, likelihood])
+    return ModelComparisonSimulator(
+        simulators=[simulator_null, simulator_alternative],
+        use_mixed_batches=True,
+        shared_simulator=context,
+    )
+
+
+@pytest.fixture
+def adapter():
+    from bayesflow import Adapter
+
+    return (
+        Adapter()
+        .sqrt("n")
+        .broadcast("n", to="x")
+        .as_set("x")
+        .rename("n", "classifier_conditions")
+        .rename("x", "summary_variables")
+        .drop("mu")
+        .convert_dtype("float64", "float32")
+    )
+
+
+@pytest.fixture
+def summary_network():
+    from bayesflow.networks import DeepSet
+
+    return DeepSet(summary_dim=2, depth=1)
+
+
+@pytest.fixture
+def classifier_network():
+    from bayesflow.networks import MLP
+
+    return MLP(widths=[32, 32])
+
+
+@pytest.fixture
+def approximator(adapter, classifier_network, summary_network, simulator):
+    from bayesflow.approximators import ModelComparisonApproximator
+
+    return ModelComparisonApproximator(
+        num_models=len(simulator.simulators),
+        classifier_network=classifier_network,
+        adapter=adapter,
+        summary_network=summary_network,
+    )

tests/test_approximators/test_model_comparison_approximator/conftest.py

@@ -0,0 +1,82 @@
+import keras
+import numpy as np
+import io
+from contextlib import redirect_stdout
+
+from tests.utils import assert_models_equal
+
+
+def test_build(approximator, train_dataset):
+    assert approximator.built is False
+
+    data_shapes = keras.tree.map_structure(keras.ops.shape, train_dataset[0])
+    approximator.build(data_shapes)
+
+    assert approximator.built is True
+    assert approximator.classifier_network.built is True
+    if approximator.summary_network is not None:
+        assert approximator.summary_network.built is True
+
+
+def test_build_adapter():
+    from bayesflow.approximators import ModelComparisonApproximator
+
+    _ = ModelComparisonApproximator.build_adapter(
+        num_models=2,
+        classifier_conditions=["foo", "bar"],
+        summary_variables=["observables"],
+        model_index_name=["indices"],
+    )
+
+
+def test_build_dataset(simulator, adapter):
+    from bayesflow.approximators import ModelComparisonApproximator
+    from bayesflow.datasets import OnlineDataset
+
+    dataset = ModelComparisonApproximator.build_dataset(
+        simulator=simulator,
+        memory_budget="20 KiB",
+        num_batches=2,
+        num_models=2,
+        classifier_conditions="foo",
+        summary_variables=["x1", "x2"],
+    )
+    assert isinstance(dataset, OnlineDataset)
+
+
+def test_fit(approximator, train_dataset, validation_dataset):
+    approximator.compile(optimizer="AdamW")
+    num_epochs = 1
+
+    # Capture ostream and train model
+    with io.StringIO() as stream:
+        with redirect_stdout(stream):
+            approximator.fit(dataset=train_dataset, validation_data=validation_dataset, epochs=num_epochs)
+
+        output = stream.getvalue()
+    # check that the loss is shown
+    assert "loss" in output
+
+
+def test_save_and_load(tmp_path, approximator, train_dataset, validation_dataset):
+    # to save, the model must be built
+    data_shapes = keras.tree.map_structure(keras.ops.shape, train_dataset[0])
+    approximator.build(data_shapes)
+    approximator.compute_metrics(**train_dataset[0])
+
+    keras.saving.save_model(approximator, tmp_path / "model.keras")
+    loaded = keras.saving.load_model(tmp_path / "model.keras")
+
+    assert_models_equal(approximator, loaded)
+
+
+def test_predict(approximator, train_dataset, simulator):
+    data_shapes = keras.tree.map_structure(keras.ops.shape, train_dataset[0])
+    approximator.build(data_shapes)
+    approximator.compute_metrics(**train_dataset[0])
+
+    num_conditions = 2
+    conditions = simulator.sample(num_conditions)
+    output = approximator.predict(conditions=conditions)
+    assert isinstance(output, np.ndarray)
+    assert output.shape[0] == num_conditions