Test approximator.sample with keep_conditions

Author: han-ol

tests/test_approximators/test_sample.py

@@ -1,3 +1,4 @@
+import numpy as np
 import keras
 from tests.utils import check_combination_simulator_adapter, check_approximator_multivariate_normal_score
 
@@ -18,3 +19,37 @@ def test_approximator_sample(approximator, simulator, batch_size, adapter):
     samples = approximator.sample(num_samples=2, conditions=data)
 
     assert isinstance(samples, dict)
+
+
+def test_approximator_sample_keep_conditions(approximator, simulator, batch_size, adapter):
+    check_combination_simulator_adapter(simulator, adapter)
+    # as long as MultivariateNormalScore is unstable, skip
+    check_approximator_multivariate_normal_score(approximator)
+
+    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)
+    approximator.build(batch_shapes)
+
+    num_samples = 2
+    samples_and_conditions = approximator.sample(num_samples=num_samples, conditions=data, keep_conditions=True)
+
+    assert isinstance(samples_and_conditions, dict)
+
+    adapted_samples_and_conditions = adapter(samples_and_conditions, strict=False)
+
+    assert any(k in adapted_samples_and_conditions for k in approximator.CONDITION_KEYS), (
+        f"adapter(approximator.sample(..., keep_conditions=True)) must return at least one of"
+        f"{approximator.CONDITION_KEYS!r}. Keys are {adapted_samples_and_conditions.keys()}."
+    )
+
+    for key, value in adapted_samples_and_conditions.items():
+        assert value.shape[:2] == (num_batches * batch_size, num_samples), (
+            f"{key} should have shape ({num_batches * batch_size}, {num_samples}, ...) but has {value.shape}."
+        )
+
+        if key in approximator.CONDITION_KEYS:
+            assert np.all(np.ptp(value, axis=1) == 0), "Not all values are the same along axis 1"