Tests for ApproximatorEnsemble

Author: han-ol

tests/test_approximators/test_approximator_ensemble/__init__.py

@@ -0,0 +1,142 @@
+import pytest
+import numpy as np
+from tests.utils import check_combination_simulator_adapter
+
+
+@pytest.fixture()
+def train_dataset_for_ensemble(batch_size, adapter, simulator):
+    check_combination_simulator_adapter(simulator, adapter)
+
+    from bayesflow import OfflineEnsembleDataset
+
+    num_batches = 4
+    data = simulator.sample((num_batches * batch_size,))
+    return OfflineEnsembleDataset(
+        num_ensemble=2, data=data, adapter=adapter, batch_size=batch_size, workers=4, max_queue_size=num_batches
+    )
+
+
+@pytest.fixture()
+def continuous_and_point_approximator_ensemble(
+    continuous_approximator, point_approximator_with_single_parametric_score
+):
+    from bayesflow import ApproximatorEnsemble
+
+    return ApproximatorEnsemble(
+        dict(cont_approx=continuous_approximator, point_approx=point_approximator_with_single_parametric_score)
+    )
+
+
+@pytest.fixture(
+    params=[
+        "continuous_and_point_approximator_ensemble",
+    ],
+    scope="function",
+)
+def continuous_approximator_ensemble(request):
+    return request.getfixturevalue(request.param)
+
+
+@pytest.fixture
+def model_comparison_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 model_comparison_train_dataset_for_ensemble(batch_size, adapter, simulator):
+    check_combination_simulator_adapter(simulator, adapter)
+
+    from bayesflow import OfflineEnsembleDataset
+
+    num_batches = 4
+    data = simulator.sample((num_batches * batch_size,))
+    return OfflineEnsembleDataset(
+        num_ensemble=2, data=data, adapter=adapter, batch_size=batch_size, workers=4, max_queue_size=num_batches
+    )
+
+
+@pytest.fixture
+def model_comparison_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 basic_model_comparison_ensemble(model_comparison_adapter):
+    from bayesflow.approximators import ModelComparisonApproximator, ApproximatorEnsemble
+    from bayesflow.networks import DeepSet, MLP
+
+    classifier_network = MLP(widths=[32, 32])
+
+    summary_network = DeepSet(summary_dim=2, depth=1)
+
+    approx_1 = ModelComparisonApproximator(
+        num_models=2,
+        classifier_network=classifier_network,
+        adapter=model_comparison_adapter,
+        summary_network=summary_network,
+    )
+    approx_2 = ModelComparisonApproximator(
+        num_models=2,
+        classifier_network=classifier_network,
+        adapter=model_comparison_adapter,
+        summary_network=summary_network,
+    )
+
+    return ApproximatorEnsemble(dict(approx_1=approx_1, approx_2=approx_2))
+
+
+@pytest.fixture(
+    params=[
+        "basic_model_comparison_ensemble",
+    ],
+    scope="function",
+)
+def model_comparison_approximator_ensemble(request):
+    return request.getfixturevalue(request.param)
+
+
+@pytest.fixture(
+    params=[
+        "continuous_and_point_approximator_ensemble",
+        "basic_model_comparison_ensemble",
+    ],
+    scope="function",
+)
+def approximator_ensemble(request):
+    return request.getfixturevalue(request.param)

tests/test_approximators/test_approximator_ensemble/conftest.py

@@ -0,0 +1,41 @@
+import keras
+from tests.utils import check_combination_simulator_adapter
+
+
+def test_build_continuous(continuous_approximator_ensemble, simulator, batch_size, adapter):
+    check_combination_simulator_adapter(simulator, adapter)
+
+    num_batches = 4
+    data = simulator.sample((num_batches * batch_size,))
+
+    batch = adapter(data)
+    batch = keras.tree.map_structure(keras.ops.convert_to_tensor, batch)
+    batch_shapes = keras.tree.map_structure(keras.ops.shape, batch)
+    print(batch_shapes)
+    continuous_approximator_ensemble.build(batch_shapes)
+
+    for member in continuous_approximator_ensemble.approximators.values():
+        for layer in member.standardize_layers.values():
+            assert layer.built
+            for count in layer.count:
+                assert count == 0.0
+
+
+def test_build_model_comparison(
+    model_comparison_approximator_ensemble, model_comparison_simulator, batch_size, model_comparison_adapter
+):
+    check_combination_simulator_adapter(model_comparison_simulator, model_comparison_adapter)
+
+    num_batches = 4
+    data = model_comparison_simulator.sample((num_batches * batch_size,))
+
+    batch = model_comparison_adapter(data)
+    batch = keras.tree.map_structure(keras.ops.convert_to_tensor, batch)
+    batch_shapes = keras.tree.map_structure(keras.ops.shape, batch)
+    model_comparison_approximator_ensemble.build(batch_shapes)
+
+    for member in model_comparison_approximator_ensemble.approximators.values():
+        for layer in member.standardize_layers.values():
+            assert layer.built
+            for count in layer.count:
+                assert count == 0.0

tests/test_approximators/test_approximator_ensemble/test_build.py

@@ -0,0 +1,21 @@
+import keras
+from tests.utils import check_combination_simulator_adapter
+
+
+def test_approximator_estimate(continuous_approximator_ensemble, simulator, batch_size, adapter):
+    check_combination_simulator_adapter(simulator, adapter)
+
+    num_batches = 4
+    data = simulator.sample((num_batches * batch_size,))
+
+    batch = adapter(data)
+    batch = keras.tree.map_structure(keras.ops.convert_to_tensor, batch)
+    batch_shapes = keras.tree.map_structure(keras.ops.shape, batch)
+    continuous_approximator_ensemble.build(batch_shapes)
+
+    estimates = continuous_approximator_ensemble.estimate(data)
+
+    assert isinstance(estimates, dict)
+
+    for estimates_value in estimates.values():
+        assert isinstance(estimates_value, dict)

tests/test_approximators/test_approximator_ensemble/test_estimate.py

@@ -0,0 +1,21 @@
+import keras
+import numpy as np
+from tests.utils import check_combination_simulator_adapter
+
+
+def test_approximator_log_prob(continuous_approximator_ensemble, simulator, batch_size, adapter):
+    check_combination_simulator_adapter(simulator, adapter)
+
+    num_batches = 4
+    data = simulator.sample((num_batches * batch_size,))
+
+    batch = adapter(data)
+    batch = keras.tree.map_structure(keras.ops.convert_to_tensor, batch)
+    batch_shapes = keras.tree.map_structure(keras.ops.shape, batch)
+    continuous_approximator_ensemble.build(batch_shapes)
+
+    log_prob = continuous_approximator_ensemble.log_prob(data=data)
+    assert isinstance(log_prob, dict)
+
+    for log_prob_value in log_prob.values():
+        assert isinstance(log_prob_value, (np.ndarray, dict))

tests/test_approximators/test_approximator_ensemble/test_log_prob.py

@@ -0,0 +1,41 @@
+import keras
+import numpy as np
+from tests.utils import check_combination_simulator_adapter
+
+
+# def test_predict(model_comparison_approximator_ensemble, model_comparison_train_dataset_for_ensemble, simulator):
+#     approximator_ensemble = model_comparison_approximator_ensemble
+#     data_shapes = keras.tree.map_structure(keras.ops.shape, model_comparison_train_dataset_for_ensemble[0])
+#     approximator_ensemble.build(data_shapes)
+#     approximator_ensemble.compute_metrics(**model_comparison_train_dataset_for_ensemble[0])
+#
+#     num_conditions = 2
+#     conditions = simulator.sample(num_conditions)
+#     predictions = approximator_ensemble.predict(conditions=conditions)
+#
+#     for predictions_value in predictions.values():
+#         assert isinstance(predictions_value, np.ndarray)
+#         assert predictions_value.shape[0] == num_conditions
+
+
+def test_predict_model_comparison(
+    model_comparison_approximator_ensemble, model_comparison_simulator, batch_size, model_comparison_adapter
+):
+    check_combination_simulator_adapter(model_comparison_simulator, model_comparison_adapter)
+
+    num_batches = 4
+    data = model_comparison_simulator.sample((num_batches * batch_size,))
+
+    batch = model_comparison_adapter(data)
+    batch = keras.tree.map_structure(keras.ops.convert_to_tensor, batch)
+    batch_shapes = keras.tree.map_structure(keras.ops.shape, batch)
+    model_comparison_approximator_ensemble.build(batch_shapes)
+
+    num_conditions = 2
+    conditions = model_comparison_simulator.sample(num_conditions)
+    print(conditions)
+    predictions = model_comparison_approximator_ensemble.predict(conditions=conditions)
+
+    for predictions_value in predictions.values():
+        assert isinstance(predictions_value, np.ndarray)
+        assert predictions_value.shape[0] == num_conditions

tests/test_approximators/test_approximator_ensemble/test_predict.py

@@ -0,0 +1,21 @@
+import keras
+from tests.utils import check_combination_simulator_adapter
+
+
+def test_approximator_sample(continuous_approximator_ensemble, simulator, batch_size, adapter):
+    check_combination_simulator_adapter(simulator, adapter)
+
+    num_batches = 4
+    data = simulator.sample((num_batches * batch_size,))
+
+    batch = adapter(data)
+    batch = keras.tree.map_structure(keras.ops.convert_to_tensor, batch)
+    batch_shapes = keras.tree.map_structure(keras.ops.shape, batch)
+    continuous_approximator_ensemble.build(batch_shapes)
+
+    samples = continuous_approximator_ensemble.sample(num_samples=2, conditions=data)
+
+    assert isinstance(samples, dict)
+
+    for samples_value in samples.values():
+        assert isinstance(samples_value, dict)