Merge branch 'dev' into adapater_nan

# Conflicts:
#	tests/test_adapters/test_adapters.py

Author: arrjon

README.md

@@ -245,8 +245,9 @@ Depending on your needs, you might not want to upgrade yet if one of the followi
   with the new version. Loading models from version 1.x in version 2.0+ is not supported.
 - You require a feature that was not ported to BayesFlow 2.0+ yet. To our knowledge,
   this applies to:
-  * Two-level/Hierarchical models: `TwoLevelGenerativeModel`, `TwoLevelPrior`.
-  * Sensitivity analysis: functionality from the `bayesflow.sensitivity` module.
+  * Two-level/Hierarchical models (planned for version 2.1): `TwoLevelGenerativeModel`, `TwoLevelPrior`.
+  * Sensitivity analysis (partially discontinued): functionality from the `bayesflow.sensitivity` module. This is still
+    possible, but we do no longer offer a special module for it. We plan to add a tutorial on this, see [#455](https://github.com/bayesflow-org/bayesflow/issues/455).
   * MCMC (discontinued): The `bayesflow.mcmc` module. We are considering other options
     to enable the use of BayesFlow in an MCMC setting.
   * Networks: `EvidentialNetwork`.

bayesflow/adapters/adapter.py

@@ -80,7 +80,9 @@ def get_config(self) -> dict:
 
         return serialize(config)
 
-    def forward(self, data: dict[str, any], *, stage: str = "inference", **kwargs) -> dict[str, np.ndarray]:
+    def forward(
+        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.
 
         Parameters
@@ -89,22 +91,33 @@ def forward(self, data: dict[str, any], *, stage: str = "inference", **kwargs) -
             The data to be transformed.
         stage : str, one of ["training", "validation", "inference"]
             The stage the function is called in.
+        log_det_jac: bool, optional
+            Whether to return the log determinant of the Jacobian of the transforms.
         **kwargs : dict
             Additional keyword arguments passed to each transform.
 
         Returns
         -------
-        dict
-            The transformed data.
+        dict | tuple[dict, dict]
+            The transformed data or tuple of transformed data and log determinant of the Jacobian.
         """
         data = data.copy()
+        if not log_det_jac:
+            for transform in self.transforms:
+                data = transform(data, stage=stage, **kwargs)
+            return data
 
+        log_det_jac = {}
         for transform in self.transforms:
-            data = transform(data, stage=stage, **kwargs)
+            transformed_data = transform(data, stage=stage, **kwargs)
+            log_det_jac = transform.log_det_jac(data, log_det_jac, **kwargs)
+            data = transformed_data
 
-        return data
+        return data, log_det_jac
 
-    def inverse(self, data: dict[str, np.ndarray], *, stage: str = "inference", **kwargs) -> dict[str, any]:
+    def inverse(
+        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.
 
         Parameters
@@ -113,24 +126,32 @@ def inverse(self, data: dict[str, np.ndarray], *, stage: str = "inference", **kw
             The data to be transformed.
         stage : str, one of ["training", "validation", "inference"]
             The stage the function is called in.
+        log_det_jac: bool, optional
+            Whether to return the log determinant of the Jacobian of the transforms.
         **kwargs : dict
             Additional keyword arguments passed to each transform.
 
         Returns
         -------
-        dict
-            The transformed data.
+        dict | tuple[dict, dict]
+            The transformed data or tuple of transformed data and log determinant of the Jacobian.
         """
         data = data.copy()
+        if not log_det_jac:
+            for transform in reversed(self.transforms):
+                data = transform(data, stage=stage, inverse=True, **kwargs)
+            return data
 
+        log_det_jac = {}
         for transform in reversed(self.transforms):
             data = transform(data, stage=stage, inverse=True, **kwargs)
+            log_det_jac = transform.log_det_jac(data, log_det_jac, inverse=True, **kwargs)
 
-        return data
+        return data, log_det_jac
 
     def __call__(
         self, data: Mapping[str, any], *, inverse: bool = False, stage="inference", **kwargs
-    ) -> dict[str, np.ndarray]:
+    ) -> dict[str, np.ndarray] | tuple[dict[str, np.ndarray], dict[str, np.ndarray]]:
         """Apply the transforms in the given direction.
 
         Parameters
@@ -146,8 +167,8 @@ def __call__(
 
         Returns
         -------
-        dict
-            The transformed data.
+        dict | tuple[dict, dict]
+            The transformed data or tuple of transformed data and log determinant of the Jacobian.
         """
         if inverse:
             return self.inverse(data, stage=stage, **kwargs)

bayesflow/adapters/transforms/concatenate.py

@@ -115,3 +115,37 @@ def extra_repr(self) -> str:
             result += f", axis={self.axis}"
 
         return result
+
+    def log_det_jac(
+        self,
+        data: dict[str, np.ndarray],
+        log_det_jac: dict[str, np.ndarray],
+        *,
+        strict: bool = False,
+        inverse: bool = False,
+        **kwargs,
+    ) -> dict[str, np.ndarray]:
+        # copy to avoid side effects
+        log_det_jac = log_det_jac.copy()
+
+        if inverse:
+            if log_det_jac.get(self.into) is not None:
+                raise ValueError(
+                    "Cannot obtain an inverse Jacobian of concatenation. "
+                    "Transform your variables before you concatenate."
+                )
+
+            return log_det_jac
+
+        required_keys = set(self.keys)
+        available_keys = set(log_det_jac.keys())
+        common_keys = available_keys & required_keys
+
+        if len(common_keys) == 0:
+            return log_det_jac
+
+        parts = [log_det_jac.pop(key) for key in common_keys]
+
+        log_det_jac[self.into] = sum(parts)
+
+        return log_det_jac

bayesflow/adapters/transforms/constrain.py

@@ -87,6 +87,11 @@ def constrain(x):
 
                     def unconstrain(x):
                         return inverse_sigmoid((x - lower) / (upper - lower))
+
+                    def ldj(x):
+                        x = (x - lower) / (upper - lower)
+                        return -np.log(x) - np.log1p(-x) - np.log(upper - lower)
+
                 case str() as name:
                     raise ValueError(f"Unsupported method name for double bounded constraint: '{name}'.")
                 case other:
@@ -101,13 +106,22 @@ def constrain(x):
 
                     def unconstrain(x):
                         return inverse_softplus(x - lower)
+
+                    def ldj(x):
+                        x = x - lower
+                        return x - np.log(np.exp(x) - 1)
+
                 case "exp" | "log":
 
                     def constrain(x):
                         return np.exp(x) + lower
 
                     def unconstrain(x):
                         return np.log(x - lower)
+
+                    def ldj(x):
+                        return -np.log(x - lower)
+
                 case str() as name:
                     raise ValueError(f"Unsupported method name for single bounded constraint: '{name}'.")
                 case other:
@@ -122,13 +136,21 @@ def constrain(x):
 
                     def unconstrain(x):
                         return -inverse_softplus(-(x - upper))
+
+                    def ldj(x):
+                        x = -(x - upper)
+                        return x - np.log(np.exp(x) - 1)
+
                 case "exp" | "log":
 
                     def constrain(x):
                         return -np.exp(-x) + upper
 
                     def unconstrain(x):
                         return -np.log(-x + upper)
+
+                    def ldj(x):
+                        return -np.log(-x + upper)
                 case str() as name:
                     raise ValueError(f"Unsupported method name for single bounded constraint: '{name}'.")
                 case other:
@@ -142,6 +164,7 @@ def unconstrain(x):
 
         self.constrain = constrain
         self.unconstrain = unconstrain
+        self.ldj = ldj
 
         # do this last to avoid serialization issues
         match inclusive:
@@ -178,3 +201,9 @@ def forward(self, data: np.ndarray, **kwargs) -> np.ndarray:
     def inverse(self, data: np.ndarray, **kwargs) -> np.ndarray:
         # inverse means network space -> data space, so constrain the data
         return self.constrain(data)
+
+    def log_det_jac(self, data: np.ndarray, inverse: bool = False, **kwargs) -> np.ndarray:
+        ldj = self.ldj(data)
+        if inverse:
+            ldj = -ldj
+        return np.sum(ldj, axis=tuple(range(1, ldj.ndim)))

bayesflow/adapters/transforms/drop.py

@@ -46,3 +46,6 @@ def inverse(self, data: dict[str, any], **kwargs) -> dict[str, any]:
 
     def extra_repr(self) -> str:
         return "[" + ", ".join(map(repr, self.keys)) + "]"
+
+    def log_det_jac(self, data: dict[str, any], log_det_jac: dict[str, any], inverse: bool = False, **kwargs):
+        return self.inverse(data=log_det_jac) if inverse else self.forward(data=log_det_jac)

bayesflow/adapters/transforms/elementwise_transform.py

@@ -25,3 +25,6 @@ def forward(self, data: np.ndarray, **kwargs) -> np.ndarray:
 
     def inverse(self, data: np.ndarray, **kwargs) -> np.ndarray:
         raise NotImplementedError
+
+    def log_det_jac(self, data: np.ndarray, inverse: bool = False, **kwargs) -> np.ndarray | None:
+        return None

bayesflow/adapters/transforms/filter_transform.py

@@ -150,9 +150,35 @@ def _should_transform(self, key: str, value: np.ndarray, inverse: bool = False)
                 return predicate(key, value, inverse=inverse)
 
     def _apply_transform(self, key: str, value: np.ndarray, inverse: bool = False, **kwargs) -> np.ndarray:
+        transform = self._get_transform(key)
+
+        return transform(value, inverse=inverse, **kwargs)
+
+    def _get_transform(self, key: str) -> ElementwiseTransform:
         if key not in self.transform_map:
             self.transform_map[key] = self.transform_constructor(**self.kwargs)
 
-        transform = self.transform_map[key]
+        return self.transform_map[key]
 
-        return transform(value, inverse=inverse, **kwargs)
+    def log_det_jac(
+        self, data: dict[str, np.ndarray], log_det_jac: dict[str, np.ndarray], *, strict: bool = True, **kwargs
+    ):
+        data = data.copy()
+
+        if strict and self.include is not None:
+            missing_keys = set(self.include) - set(data.keys())
+            if missing_keys:
+                raise KeyError(f"Missing keys from include list: {missing_keys!r}")
+
+        for key, value in data.items():
+            if self._should_transform(key, value, inverse=False):
+                transform = self._get_transform(key)
+                ldj = transform.log_det_jac(value, **kwargs)
+                if ldj is None:
+                    continue
+                elif key in log_det_jac:
+                    log_det_jac[key] += ldj
+                else:
+                    log_det_jac[key] = ldj
+
+        return log_det_jac

bayesflow/adapters/transforms/keep.py

@@ -57,3 +57,6 @@ def inverse(self, data: dict[str, any], **kwargs) -> dict[str, any]:
 
     def extra_repr(self) -> str:
         return "[" + ", ".join(map(repr, self.keys)) + "]"
+
+    def log_det_jac(self, data: dict[str, any], log_det_jac: dict[str, any], inverse: bool = False, **kwargs):
+        return self.inverse(data=log_det_jac) if inverse else self.forward(data=log_det_jac)

bayesflow/adapters/transforms/log.py

@@ -37,3 +37,12 @@ def inverse(self, data: np.ndarray, **kwargs) -> np.ndarray:
 
     def get_config(self) -> dict:
         return serialize({"p1": self.p1})
+
+    def log_det_jac(self, data: np.ndarray, inverse: bool = False, **kwargs) -> np.ndarray:
+        if self.p1:
+            ldj = -np.log1p(data)
+        else:
+            ldj = -np.log(data)
+        if inverse:
+            ldj = -ldj
+        return np.sum(ldj, axis=tuple(range(1, ldj.ndim)))

bayesflow/adapters/transforms/map_transform.py

@@ -41,12 +41,8 @@ def get_config(self) -> dict:
     def forward(self, data: dict[str, np.ndarray], *, strict: bool = True, **kwargs) -> dict[str, np.ndarray]:
         data = data.copy()
 
-        required_keys = set(self.transform_map.keys())
-        available_keys = set(data.keys())
-        missing_keys = required_keys - available_keys
-
-        if strict and missing_keys:
-            raise KeyError(f"Missing keys: {missing_keys!r}")
+        if strict:
+            self._check_keys(data)
 
         for key, transform in self.transform_map.items():
             if key in data:
@@ -57,15 +53,40 @@ def forward(self, data: dict[str, np.ndarray], *, strict: bool = True, **kwargs)
     def inverse(self, data: dict[str, np.ndarray], *, strict: bool = False, **kwargs) -> dict[str, np.ndarray]:
         data = data.copy()
 
-        required_keys = set(self.transform_map.keys())
-        available_keys = set(data.keys())
-        missing_keys = required_keys - available_keys
-
-        if strict and missing_keys:
-            raise KeyError(f"Missing keys: {missing_keys!r}")
+        if strict:
+            self._check_keys(data)
 
         for key, transform in self.transform_map.items():
             if key in data:
                 data[key] = transform.inverse(data[key], **kwargs)
 
         return data
+
+    def log_det_jac(
+        self, data: dict[str, np.ndarray], log_det_jac: dict[str, np.ndarray], *, strict: bool = True, **kwargs
+    ) -> dict[str, np.ndarray]:
+        data = data.copy()
+
+        if strict:
+            self._check_keys(data)
+
+        for key, transform in self.transform_map.items():
+            if key in data:
+                ldj = transform.log_det_jac(data[key], **kwargs)
+
+                if ldj is None:
+                    continue
+                elif key in log_det_jac:
+                    log_det_jac[key] += ldj
+                else:
+                    log_det_jac[key] = ldj
+
+        return log_det_jac
+
+    def _check_keys(self, data: dict[str, np.ndarray]):
+        required_keys = set(self.transform_map.keys())
+        available_keys = set(data.keys())
+        missing_keys = required_keys - available_keys
+
+        if missing_keys:
+            raise KeyError(f"Missing keys: {missing_keys!r}")