add tests for adapter jacobians

Author: Kucharssim

tests/test_adapters/conftest.py

@@ -49,12 +49,52 @@ def random_data():
         "z1": np.random.standard_normal(size=(32, 2)),
         "p1": np.random.lognormal(size=(32, 2)),
         "p2": np.random.lognormal(size=(32, 2)),
+        "p3": np.random.lognormal(size=(32, 2)),
+        "n1": 1 - np.random.lognormal(size=(32, 2)),
         "s1": np.random.standard_normal(size=(32, 3, 2)),
         "s2": np.random.standard_normal(size=(32, 3, 2)),
         "t1": np.zeros((3, 2)),
         "t2": np.ones((32, 3, 2)),
         "d1": np.random.standard_normal(size=(32, 2)),
         "d2": np.random.standard_normal(size=(32, 2)),
         "o1": np.random.randint(0, 9, size=(32, 2)),
+        "u1": np.random.uniform(low=-1, high=2, size=(32, 1)),
         "key_to_split": np.random.standard_normal(size=(32, 10)),
     }
+
+
+@pytest.fixture()
+def adapter_jacobian():
+    from bayesflow.adapters import Adapter
+
+    adapter = (
+        Adapter()
+        .scale("x1", by=2)
+        .log("p1", p1=True)
+        .sqrt("p2")
+        .constrain("p3", lower=0)
+        .constrain("n1", upper=1)
+        .constrain("u1", lower=-1, upper=2)
+        .concatenate(["p1", "p2", "p3"], into="p")
+        .rename("u1", "u")
+    )
+
+    return adapter
+
+
+@pytest.fixture()
+def adapter_jacobian_inverse():
+    from bayesflow.adapters import Adapter
+
+    adapter = (
+        Adapter()
+        .standardize("x1", mean=1, std=2)
+        .log("p1")
+        .sqrt("p2")
+        .constrain("p3", lower=0, method="log")
+        .constrain("n1", upper=1, method="log")
+        .constrain("u1", lower=-1, upper=2)
+        .scale(["p1", "p2", "p3"], by=3.5)
+    )
+
+    return adapter

tests/test_adapters/test_adapters.py

@@ -13,7 +13,7 @@ def test_cycle_consistency(adapter, random_data):
     deprocessed = adapter(processed, inverse=True)
 
     for key, value in random_data.items():
-        if key in ["d1", "d2"]:
+        if key in ["d1", "d2", "p3", "n1", "u1"]:
             # dropped
             continue
         assert key in deprocessed
@@ -230,3 +230,36 @@ def test_to_dict_transform():
 
     # category should have 5 one-hot categories, even though it was only passed 4
     assert processed["category"].shape[-1] == 5
+
+
+def test_jacobian(adapter_jacobian, random_data):
+    d, jacobian = adapter_jacobian(random_data, jacobian=True)
+
+    assert np.allclose(jacobian["x1"], np.log(2))
+
+    p1 = -np.log1p(random_data["p1"])
+    p2 = -0.5 * np.log(random_data["p2"]) + 0.5
+    p3 = random_data["p3"] - np.log(np.exp(random_data["p3"]) - 1)
+    p = np.sum(p1, axis=-1) + np.sum(p2, axis=-1) + np.sum(p3, axis=-1)
+
+    assert np.allclose(jacobian["p"], p)
+
+    n1 = -(random_data["n1"] - 1)
+    n1 = n1 - np.log(np.exp(n1) - 1)
+    n1 = np.sum(n1, axis=-1)
+
+    assert np.allclose(jacobian["n1"], n1)
+
+    u1 = random_data["u1"]
+    u1 = (u1 + 1) / 3
+    u1 = -np.log(u1) - np.log1p(-u1) - np.log(3)
+
+    assert np.allclose(jacobian["u"], u1[:, 0])
+
+
+def test_jacobian_inverse(adapter_jacobian_inverse, random_data):
+    d, forward_jacobian = adapter_jacobian_inverse(random_data, jacobian=True)
+    d, inverse_jacobian = adapter_jacobian_inverse(d, inverse=True, jacobian=True)
+
+    for key in forward_jacobian.keys():
+        assert np.allclose(forward_jacobian[key], -inverse_jacobian[key])