Adding tests

Author: sjedhoff

tests/test_diagnostics/test_diagnostics_metrics.py

@@ -35,6 +35,13 @@ def test_metric_calibration_error(random_estimates, random_targets, var_names):
     assert out["values"].shape == (random_estimates["sigma"].shape[-1],)
     assert out["variable_names"] == ["sigma"]
 
+    # test quantities
+    test_quantities = {
+        r"$\beta_1 + \beta_2$": lambda data: np.sum(data["beta"], axis=-1),
+        r"$\beta_1 \cdot \beta_2$": lambda data: np.prod(data["beta"], axis=-1),
+    }
+    out = bf.diagnostics.metrics.calibration_error(random_estimates, random_targets, test_quantities=test_quantities)
+    assert out["values"].shape[0] == len(test_quantities) + num_variables(random_estimates)
 
 def test_posterior_contraction(random_estimates, random_targets):
     # basic functionality: automatic variable names
@@ -47,6 +54,16 @@ def test_posterior_contraction(random_estimates, random_targets):
     out = bf.diagnostics.metrics.posterior_contraction(random_estimates, random_targets, aggregation=None)
     assert out["values"].shape == (random_estimates["sigma"].shape[0], num_variables(random_estimates))
 
+    # test quantities
+    test_quantities = {
+        r"$\beta_1 + \beta_2$": lambda data: np.sum(data["beta"], axis=-1),
+        r"$\beta_1 \cdot \beta_2$": lambda data: np.prod(data["beta"], axis=-1),
+    }
+    out = bf.diagnostics.metrics.posterior_contraction(
+        random_estimates, random_targets, test_quantities=test_quantities
+    )
+    assert out["values"].shape[0] == len(test_quantities) + num_variables(random_estimates)
+
 
 def test_root_mean_squared_error(random_estimates, random_targets):
     # basic functionality: automatic variable names
@@ -56,6 +73,16 @@ def test_root_mean_squared_error(random_estimates, random_targets):
     assert out["metric_name"] == "NRMSE"
     assert out["variable_names"] == ["beta_0", "beta_1", "sigma"]
 
+    # test quantities
+    test_quantities = {
+        r"$\beta_1 + \beta_2$": lambda data: np.sum(data["beta"], axis=-1),
+        r"$\beta_1 \cdot \beta_2$": lambda data: np.prod(data["beta"], axis=-1),
+    }
+    out = bf.diagnostics.metrics.root_mean_squared_error(
+        random_estimates, random_targets, test_quantities=test_quantities
+    )
+    assert out["values"].shape[0] == len(test_quantities) + num_variables(random_estimates)
+
 
 def test_classifier_two_sample_test(random_samples_a, random_samples_b):
     metric = bf.diagnostics.metrics.classifier_two_sample_test(estimates=random_samples_a, targets=random_samples_a)
@@ -95,6 +122,16 @@ def test_calibration_log_gamma(random_estimates, random_targets):
     assert out["metric_name"] == "Log Gamma"
     assert out["variable_names"] == ["beta_0", "beta_1", "sigma"]
 
+    # test quantities
+    test_quantities = {
+        r"$\beta_1 + \beta_2$": lambda data: np.sum(data["beta"], axis=-1),
+        r"$\beta_1 \cdot \beta_2$": lambda data: np.prod(data["beta"], axis=-1),
+    }
+    out = bf.diagnostics.metrics.calibration_log_gamma(
+        random_estimates, random_targets, test_quantities=test_quantities
+    )
+    assert out["values"].shape[0] == len(test_quantities) + num_variables(random_estimates)
+
 
 def test_calibration_log_gamma_end_to_end():
     # This is a function test for simulation-based calibration.

tests/test_diagnostics/test_diagnostics_plots.py

@@ -85,6 +85,16 @@ def test_calibration_histogram(random_estimates, random_targets):
     assert len(out.axes) == num_variables(random_estimates)
     assert out.axes[0].title._text == "beta_0"
 
+    # test quantities
+    test_quantities = {
+        r"$\beta_1 + \beta_2$": lambda data: np.sum(data["beta"], axis=-1),
+        r"$\beta_1 \cdot \beta_2$": lambda data: np.prod(data["beta"], axis=-1),
+    }
+    out = bf.diagnostics.plots.calibration_histogram(random_estimates, random_targets, test_quantities=test_quantities)
+    assert len(out.axes) == len(test_quantities) + num_variables(random_estimates)
+    assert out.axes[1].title._text == r"$\beta_1 \cdot \beta_2$"
+    assert out.axes[-1].title._text == r"sigma"
+
 
 def test_loss(history):
     out = bf.diagnostics.loss(history)
@@ -102,6 +112,16 @@ def test_recovery_bounds(random_estimates, random_targets):
     assert len(out.axes) == num_variables(random_estimates)
     assert out.axes[2].title._text == "sigma"
 
+    # test quantities
+    test_quantities = {
+        r"$\beta_1 + \beta_2$": lambda data: np.sum(data["beta"], axis=-1),
+        r"$\beta_1 \cdot \beta_2$": lambda data: np.prod(data["beta"], axis=-1),
+    }
+    out = bf.diagnostics.plots.calibration_histogram(random_estimates, random_targets, test_quantities=test_quantities)
+    assert len(out.axes) == len(test_quantities) + num_variables(random_estimates)
+    assert out.axes[1].title._text == r"$\beta_1 \cdot \beta_2$"
+    assert out.axes[-1].title._text == r"sigma"
+
 
 def test_recovery_symmetric(random_estimates, random_targets):
     # basic functionality: automatic variable names
@@ -127,6 +147,16 @@ def test_z_score_contraction(random_estimates, random_targets):
     assert len(out.axes) == num_variables(random_estimates)
     assert out.axes[1].title._text == "beta_1"
 
+    # test quantities
+    test_quantities = {
+        r"$\beta_1 + \beta_2$": lambda data: np.sum(data["beta"], axis=-1),
+        r"$\beta_1 \cdot \beta_2$": lambda data: np.prod(data["beta"], axis=-1),
+    }
+    out = bf.diagnostics.plots.z_score_contraction(random_estimates, random_targets, test_quantities=test_quantities)
+    assert len(out.axes) == len(test_quantities) + num_variables(random_estimates)
+    assert out.axes[1].title._text == r"$\beta_1 \cdot \beta_2$"
+    assert out.axes[-1].title._text == r"sigma"
+
 
 def test_pairs_samples(random_priors):
     out = bf.diagnostics.plots.pairs_samples(
@@ -291,6 +321,16 @@ def test_coverage(random_estimates, random_targets):
     assert out.axes[0].get_xlabel() == "Central interval width"
     assert out.axes[0].get_ylabel() == "Empirical coverage"
 
+    # test quantities
+    test_quantities = {
+        r"$\beta_1 + \beta_2$": lambda data: np.sum(data["beta"], axis=-1),
+        r"$\beta_1 \cdot \beta_2$": lambda data: np.prod(data["beta"], axis=-1),
+    }
+    out = bf.diagnostics.plots.coverage(random_estimates, random_targets, test_quantities=test_quantities)
+    assert len(out.axes) == len(test_quantities) + num_variables(random_estimates)
+    assert out.axes[1].title._text == r"$\beta_1 \cdot \beta_2$"
+    assert out.axes[-1].title._text == r"sigma"
+
 
 def test_coverage_diff(random_estimates, random_targets):
     # basic functionality: automatic variable names

tests/test_links/test_links.py

@@ -23,9 +23,9 @@ def check_ordering(output, axis):
     assert np.all(np.diff(output, axis=axis) > 0), f"is not ordered along specified axis: {axis}."
     for i in range(output.ndim):
         if i != axis % output.ndim:
-            assert not np.all(np.diff(output, axis=i) > 0), (
-                f"is ordered along axis which is not meant to be ordered: {i}."
-            )
+            assert not np.all(
+                np.diff(output, axis=i) > 0
+            ), f"is ordered along axis which is not meant to be ordered: {i}."
 
 
 @pytest.mark.parametrize("axis", [0, 1, 2])

tests/test_networks/test_point_inference_network/test_point_inference_network.py

@@ -44,9 +44,9 @@ def test_save_and_load(tmp_path, point_inference_network, random_samples, random
 
     for key_outer in out1.keys():
         for key_inner in out1[key_outer].keys():
-            assert keras.ops.all(keras.ops.isclose(out1[key_outer][key_inner], out2[key_outer][key_inner])), (
-                "Output of original and loaded model differs significantly."
-            )
+            assert keras.ops.all(
+                keras.ops.isclose(out1[key_outer][key_inner], out2[key_outer][key_inner])
+            ), "Output of original and loaded model differs significantly."
 
 
 def test_copy_unequal(point_inference_network, random_samples, random_conditions):