use map transform

Author: paul-buerkner

bayesflow/adapters/adapter.py

@@ -483,7 +483,7 @@ def expand_dims(self, keys: str | Sequence[str], *, axis: int | tuple):
         if isinstance(keys, str):
             keys = [keys]
 
-        transform = ExpandDims(keys, axis=axis)
+        transform = MapTransform({key: ExpandDims(axis=axis) for key in keys})
         self.transforms.append(transform)
         return self
 
@@ -508,14 +508,14 @@ def log(self, keys: str | Sequence[str], *, p1: bool = False):
         Parameters
         ----------
         keys : str or Sequence of str
-            The names of the variables to expand.
+            The names of the variables to transform.
         p1 : boolean
             Add 1 to the input before taking the logarithm?
         """
         if isinstance(keys, str):
             keys = [keys]
 
-        transform = Log(keys, p1=p1)
+        transform = MapTransform({key: Log(p1=p1) for key in keys})
         self.transforms.append(transform)
         return self
 
@@ -555,12 +555,12 @@ def sqrt(self, keys: str | Sequence[str]):
         Parameters
         ----------
         keys : str or Sequence of str
-            The names of the variables to expand.
+            The names of the variables to transform.
         """
         if isinstance(keys, str):
             keys = [keys]
 
-        transform = Sqrt(keys)
+        transform = MapTransform({key: Sqrt() for key in keys})
         self.transforms.append(transform)
         return self
 

bayesflow/adapters/transforms/expand_dims.py

@@ -5,7 +5,6 @@
     serialize_keras_object as serialize,
 )
 
-from collections.abc import Sequence
 from .elementwise_transform import ElementwiseTransform
 
 
@@ -15,8 +14,6 @@ class ExpandDims(ElementwiseTransform):
 
     Parameters
     ----------
-    keys : str or Sequence of str
-        The names of the variables to expand.
     axis : int or tuple
         The axis to expand.
 
@@ -49,29 +46,23 @@ class ExpandDims(ElementwiseTransform):
     It is recommended to precede this transform with a :class:`bayesflow.adapters.transforms.ToArray` transform.
     """
 
-    def __init__(self, keys: Sequence[str], *, axis: int | tuple):
+    def __init__(self, *, axis: int | tuple):
         super().__init__()
-
-        self.keys = keys
         self.axis = axis
 
     @classmethod
     def from_config(cls, config: dict, custom_objects=None) -> "ExpandDims":
         return cls(
-            keys=deserialize(config["keys"], custom_objects),
             axis=deserialize(config["axis"], custom_objects),
         )
 
     def get_config(self) -> dict:
         return {
-            "keys": serialize(self.keys),
             "axis": serialize(self.axis),
         }
 
-    # noinspection PyMethodOverriding
-    def forward(self, data: dict[str, any], **kwargs) -> dict[str, np.ndarray]:
-        return {k: (np.expand_dims(v, axis=self.axis) if k in self.keys else v) for k, v in data.items()}
+    def forward(self, data: np.ndarray, **kwargs) -> np.ndarray:
+        return np.expand_dims(data, axis=self.axis)
 
-    # noinspection PyMethodOverriding
-    def inverse(self, data: dict[str, any], **kwargs) -> dict[str, np.ndarray]:
-        return {k: (np.squeeze(v, axis=self.axis) if k in self.keys else v) for k, v in data.items()}
+    def inverse(self, data: np.ndarray, **kwargs) -> np.ndarray:
+        return np.squeeze(data, axis=self.axis)

bayesflow/adapters/transforms/log.py

@@ -1,6 +1,9 @@
 import numpy as np
 
-from collections.abc import Sequence
+from keras.saving import (
+    deserialize_keras_object as deserialize,
+    serialize_keras_object as serialize,
+)
 
 from .elementwise_transform import ElementwiseTransform
 
@@ -18,26 +21,29 @@ class Log(ElementwiseTransform):
     >>> adapter = bf.Adapter().log(["x"])
     """
 
-    def __init__(self, keys: Sequence[str], *, p1: bool = False):
+    def __init__(self, *, p1: bool = False):
         super().__init__()
-        self.keys = keys
         self.p1 = p1
 
-    def forward(self, data: dict[str, any], **kwargs) -> dict[str, np.ndarray]:
+    def forward(self, data: np.ndarray, **kwargs) -> np.ndarray:
         if self.p1:
-            return {k: (np.log1p(v) if k in self.keys else v) for k, v in data.items()}
+            return np.log1p(data)
         else:
-            return {k: (np.log(v) if k in self.keys else v) for k, v in data.items()}
+            return np.log(data)
 
-    def inverse(self, data: dict[str, any], **kwargs) -> dict[str, np.ndarray]:
+    def inverse(self, data: np.ndarray, **kwargs) -> np.ndarray:
         if self.p1:
-            return {k: (np.expm1(v) if k in self.keys else v) for k, v in data.items()}
+            return np.expm1(data)
         else:
-            return {k: (np.exp(v) if k in self.keys else v) for k, v in data.items()}
+            return np.exp(data)
 
     @classmethod
     def from_config(cls, config: dict, custom_objects=None) -> "Log":
-        return cls()
+        return cls(
+            p1=deserialize(config["p1"], custom_objects),
+        )
 
     def get_config(self) -> dict:
-        return {}
+        return {
+            "p1": serialize(self.p1),
+        }

bayesflow/adapters/transforms/sqrt.py

@@ -1,7 +1,5 @@
 import numpy as np
 
-from collections.abc import Sequence
-
 from .elementwise_transform import ElementwiseTransform
 
 
@@ -13,15 +11,11 @@ class Sqrt(ElementwiseTransform):
     >>> adapter = bf.Adapter().sqrt(["x"])
     """
 
-    def __init__(self, keys: Sequence[str]):
-        super().__init__()
-        self.keys = keys
-
-    def forward(self, data: dict[str, any], **kwargs) -> dict[str, np.ndarray]:
-        return {k: (np.sqrt(v) if k in self.keys else v) for k, v in data.items()}
+    def forward(self, data: np.ndarray, **kwargs) -> np.ndarray:
+        return np.sqrt(data)
 
-    def inverse(self, data: dict[str, any], **kwargs) -> dict[str, np.ndarray]:
-        return {k: (np.square(v) if k in self.keys else v) for k, v in data.items()}
+    def inverse(self, data: np.ndarray, **kwargs) -> np.ndarray:
+        return np.square(data)
 
     @classmethod
     def from_config(cls, config: dict, custom_objects=None) -> "Sqrt":

tests/test_adapters/conftest.py

@@ -30,8 +30,7 @@ def adapter():
         .concatenate(["y1", "y2"], into="y")
         .expand_dims(["z1"], axis=2)
         .apply(forward=forward_transform, inverse=inverse_transform)
-        # TODO: fix this in keras
-        # .apply(include="p1", forward=np.log, inverse=np.exp)
+        .log("p1")
         .constrain("p2", lower=0)
         .standardize(exclude=["t1", "t2", "o1"])
         .drop("d1")

tests/test_adapters/test_adapters.py

@@ -89,8 +89,8 @@ def test_constrain():
     assert np.isinf(result["x_both_disc2"][-1])
 
 
-def test_log_sqrt(random_data):
-    # check if constraint-implied transforms are applied correctly
+def test_simple_transforms(random_data):
+    # check if simple transforms are applied correctly
     from bayesflow.adapters import Adapter
 
     adapter = Adapter().log(["o1", "p2"]).log("t1", p1=True).sqrt("p1")