jacobian -> log_det_jac

Author: Kucharssim

bayesflow/adapters/adapter.py

@@ -80,7 +80,7 @@ def get_config(self) -> dict:
         return serialize(config)
 
     def forward(
-        self, data: dict[str, any], *, stage: str = "inference", jacobian: bool = False, **kwargs
+        self, data: dict[str, any], *, stage: str = "inference", log_det_jac: bool = False, **kwargs
     ) -> dict[str, np.ndarray] | tuple[dict[str, np.ndarray], dict[str, np.ndarray]]:
         """Apply the transforms in the forward direction.
 
@@ -90,18 +90,18 @@ def forward(
             The data to be transformed.
         stage : str, one of ["training", "validation", "inference"]
             The stage the function is called in.
-        jacobian: bool, optional
+        log_det_jac: bool, optional
             Whether to return the log determinant jacobians of the transforms.
         **kwargs : dict
             Additional keyword arguments passed to each transform.
 
         Returns
         -------
         dict | tuple[dict, dict]
-            The transformed data or tuple of transformed data and jacobians.
+            The transformed data or tuple of transformed data and log determinant jacobians.
         """
         data = data.copy()
-        if not jacobian:
+        if not log_det_jac:
             for transform in self.transforms:
                 data = transform(data, stage=stage, **kwargs)
             return data
@@ -114,7 +114,7 @@ def forward(
         return data, log_det_jac
 
     def inverse(
-        self, data: dict[str, np.ndarray], *, stage: str = "inference", jacobian: bool = False, **kwargs
+        self, data: dict[str, np.ndarray], *, stage: str = "inference", log_det_jac: bool = False, **kwargs
     ) -> dict[str, np.ndarray] | tuple[dict[str, np.ndarray], dict[str, np.ndarray]]:
         """Apply the transforms in the inverse direction.
 
@@ -124,18 +124,18 @@ def inverse(
             The data to be transformed.
         stage : str, one of ["training", "validation", "inference"]
             The stage the function is called in.
-        jacobian: bool, optional
+        log_det_jac: bool, optional
             Whether to return the log determinant jacobians of the transforms.
         **kwargs : dict
             Additional keyword arguments passed to each transform.
 
         Returns
         -------
         dict | tuple[dict, dict]
-            The transformed data or tuple of transformed data and jacobians.
+            The transformed data or tuple of transformed data and log determinant jacobians.
         """
         data = data.copy()
-        if not jacobian:
+        if not log_det_jac:
             for transform in reversed(self.transforms):
                 data = transform(data, stage=stage, inverse=True, **kwargs)
             return data
@@ -166,7 +166,7 @@ def __call__(
         Returns
         -------
         dict | tuple[dict, dict]
-            The transformed data or tuple of transformed data and jacobians.
+            The transformed data or tuple of transformed data and log determinant jacobians.
         """
         if inverse:
             return self.inverse(data, stage=stage, **kwargs)

bayesflow/adapters/transforms/concatenate.py

@@ -131,7 +131,7 @@ def log_det_jac(
         if inverse:
             if log_det_jac.get(self.into) is not None:
                 raise ValueError(
-                    "Cannot obtain an inverse jacobian of concatenation. "
+                    "Cannot obtain an inverse jacobian determinant of concatenation. "
                     "Transform your variables before you concatenate."
                 )
 

bayesflow/adapters/transforms/numpy_transform.py

@@ -74,4 +74,4 @@ def inverse(self, data: np.ndarray, **kwargs) -> np.ndarray:
         return self._inverse(data)
 
     def log_det_jac(self, data, inverse=False, **kwargs):
-        raise NotImplementedError("Jacobian of the numpy transforms are not implemented yet")
+        raise NotImplementedError("Log determinand jacobian of the numpy transforms are not implemented yet")

bayesflow/adapters/transforms/sqrt.py

@@ -24,7 +24,7 @@ def get_config(self) -> dict:
         return {}
 
     def log_det_jac(self, data: np.ndarray, inverse: bool = False, **kwargs) -> np.ndarray:
-        ldj = -0.5 * np.log(data) + 0.5
+        ldj = -0.5 * np.log(data) - np.log(2)
         if inverse:
             ldj = -ldj
         return np.sum(ldj, axis=tuple(range(1, ldj.ndim)))

bayesflow/approximators/continuous_approximator.py

@@ -417,14 +417,15 @@ def log_prob(self, data: Mapping[str, np.ndarray], **kwargs) -> np.ndarray | dic
         np.ndarray
             Log-probabilities of the distribution `p(inference_variables | inference_conditions, h(summary_conditions))`
         """
-        data, jacobian = self.adapter(data, strict=False, stage="inference", jacobian=True, **kwargs)
+        data, log_det_jac = self.adapter(data, strict=False, stage="inference", log_det_jac=True, **kwargs)
         data = keras.tree.map_structure(keras.ops.convert_to_tensor, data)
         log_prob = self._log_prob(**data, **kwargs)
         log_prob = keras.tree.map_structure(keras.ops.convert_to_numpy, log_prob)
 
-        jacobian = jacobian.get("inference_variables")
-        if jacobian is not None:
-            log_prob = log_prob + jacobian
+        # change of variables formula
+        log_det_jac = log_det_jac.get("inference_variables")
+        if log_det_jac is not None:
+            log_prob = log_prob + log_det_jac
 
         return log_prob
 

tests/test_adapters/conftest.py

@@ -64,7 +64,7 @@ def random_data():
 
 
 @pytest.fixture()
-def adapter_jacobian():
+def adapter_log_det_jac():
     from bayesflow.adapters import Adapter
 
     adapter = (
@@ -83,7 +83,7 @@ def adapter_jacobian():
 
 
 @pytest.fixture()
-def adapter_jacobian_inverse():
+def adapter_log_det_jac_inverse():
     from bayesflow.adapters import Adapter
 
     adapter = (

tests/test_adapters/test_adapters.py

@@ -232,34 +232,34 @@ def test_to_dict_transform():
     assert processed["category"].shape[-1] == 5
 
 
-def test_jacobian(adapter_jacobian, random_data):
-    d, jacobian = adapter_jacobian(random_data, jacobian=True)
+def test_log_det_jac(adapter_log_det_jac, random_data):
+    d, log_det_jac = adapter_log_det_jac(random_data, log_det_jac=True)
 
-    assert np.allclose(jacobian["x1"], np.log(2))
+    assert np.allclose(log_det_jac["x1"], np.log(2))
 
     p1 = -np.log1p(random_data["p1"])
-    p2 = -0.5 * np.log(random_data["p2"]) + 0.5
+    p2 = -0.5 * np.log(random_data["p2"]) - np.log(2)
     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)
+    assert np.allclose(log_det_jac["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)
+    assert np.allclose(log_det_jac["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])
+    assert np.allclose(log_det_jac["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)
+def test_log_det_jac_inverse(adapter_log_det_jac_inverse, random_data):
+    d, forward_log_det_jac = adapter_log_det_jac_inverse(random_data, log_det_jac=True)
+    d, inverse_log_det_jac = adapter_log_det_jac_inverse(d, inverse=True, log_det_jac=True)
 
-    for key in forward_jacobian.keys():
-        assert np.allclose(forward_jacobian[key], -inverse_jacobian[key])
+    for key in forward_log_det_jac.keys():
+        assert np.allclose(forward_log_det_jac[key], -inverse_log_det_jac[key])