add test

Author: Kucharssim

tests/test_simulators/conftest.py

@@ -167,6 +167,32 @@ def likelihood(mu, n):
     return make_simulator([prior, likelihood], meta_fn=context)
 
 
+@pytest.fixture(params=["drop", np.nan])
+def multimodel(request):
+    from bayesflow.simulators import make_simulator, ModelComparisonSimulator
+
+    rng = np.random.default_rng()
+
+    def prior_1():
+        return dict(w=rng.uniform())
+
+    def prior_2():
+        return dict(c=rng.uniform())
+
+    def model_1(w):
+        return dict(x=w)
+
+    def model_2(c):
+        return dict(x=c)
+
+    simulator_1 = make_simulator([prior_1, model_1])
+    simulator_2 = make_simulator([prior_2, model_2])
+
+    simulator = ModelComparisonSimulator(simulators=[simulator_1, simulator_2], key_conflicts=request.param)
+
+    return simulator
+
+
 @pytest.fixture()
 def fixed_n():
     return 5

tests/test_simulators/test_simulators.py

@@ -47,3 +47,13 @@ def test_fixed_sample(composite_gaussian, batch_size, fixed_n, fixed_mu):
     assert samples["mu"].shape == (batch_size, 1)
     assert np.all(samples["mu"] == fixed_mu)
     assert samples["y"].shape == (batch_size, fixed_n)
+
+
+def test_multimodel_sample(multimodel, batch_size):
+    samples = multimodel.sample(batch_size)
+
+    if multimodel.key_conflicts == "drop":
+        assert set(samples) == {"x", "model_indices"}
+    else:
+        assert set(samples) == {"x", "model_indices", "c", "w"}
+        assert np.sum(np.isnan(samples["c"])) + np.sum(np.isnan(samples["w"])) == batch_size