Numpy transforms (#350)

* add numpy transform

* remove lambda transform

* repurpose `adapter.apply()` for `NumpyTransform` instead of `LambdaTransform`

* fix usages of `adapter.apply()` in example notebooks

* fix tests

* only allow strings as arguments (subject to be fixed by #323)

* unify serialization pattern

* remove old lambda transform error message

* add fail fast to CI tests

* cannot suppport filtering kwargs for numpy ufunc in python 3.10

Author: LarsKue

.github/workflows/tests.yaml

@@ -95,7 +95,7 @@ jobs:
 
       - name: Run Tests
         run: |
-          pytest
+          pytest -x
 
       - name: Create Coverage Report
         run: |

bayesflow/adapters/adapter.py

@@ -1,4 +1,4 @@
-from collections.abc import Callable, MutableSequence, Sequence
+from collections.abc import MutableSequence, Sequence
 
 import numpy as np
 from keras.saving import (
@@ -18,9 +18,9 @@
     ExpandDims,
     FilterTransform,
     Keep,
-    LambdaTransform,
     Log,
     MapTransform,
+    NumpyTransform,
     OneHot,
     Rename,
     Sqrt,
@@ -234,8 +234,8 @@ def __len__(self):
     def apply(
         self,
         *,
-        forward: Callable[[np.ndarray, ...], np.ndarray],
-        inverse: Callable[[np.ndarray, ...], np.ndarray],
+        forward: np.ufunc | str,
+        inverse: np.ufunc | str = None,
         predicate: Predicate = None,
         include: str | Sequence[str] = None,
         exclude: str | Sequence[str] = None,
@@ -271,7 +271,7 @@ def apply(
         to the `custom_objects` argument of the `deserialize` function when deserializing this class.
         """
         transform = FilterTransform(
-            transform_constructor=LambdaTransform,
+            transform_constructor=NumpyTransform,
             predicate=predicate,
             include=include,
             exclude=exclude,

bayesflow/adapters/transforms/__init__.py

@@ -9,9 +9,9 @@
 from .expand_dims import ExpandDims
 from .filter_transform import FilterTransform
 from .keep import Keep
-from .lambda_transform import LambdaTransform
 from .log import Log
 from .map_transform import MapTransform
+from .numpy_transform import NumpyTransform
 from .one_hot import OneHot
 from .rename import Rename
 from .sqrt import Sqrt

bayesflow/adapters/transforms/as_set.py

@@ -1,8 +1,10 @@
 import numpy as np
+from keras.saving import register_keras_serializable as serializable
 
 from .elementwise_transform import ElementwiseTransform
 
 
+@serializable(package="bayesflow.adapters")
 class AsSet(ElementwiseTransform):
     """The `.as_set(["x", "y"])` transform indicates that both `x` and `y` are treated as sets.
 

bayesflow/adapters/transforms/as_time_series.py

@@ -1,8 +1,10 @@
 import numpy as np
+from keras.saving import register_keras_serializable as serializable
 
 from .elementwise_transform import ElementwiseTransform
 
 
+@serializable(package="bayesflow.adapters")
 class AsTimeSeries(ElementwiseTransform):
     """The `.as_time_series` transform can be used to indicate that variables shall be treated as time series.
 

bayesflow/adapters/transforms/expand_dims.py

@@ -1,13 +1,14 @@
 import numpy as np
-
 from keras.saving import (
     deserialize_keras_object as deserialize,
+    register_keras_serializable as serializable,
     serialize_keras_object as serialize,
 )
 
 from .elementwise_transform import ElementwiseTransform
 
 
+@serializable(package="bayesflow.adapters")
 class ExpandDims(ElementwiseTransform):
     """
     Expand the shape of an array.

bayesflow/adapters/transforms/filter_transform.py

@@ -86,18 +86,6 @@ def from_config(cls, config: dict, custom_objects=None) -> "Transform":
         try:
             kwargs = deserialize(config["kwargs"])
         except TypeError as e:
-            if transform_constructor.__name__ == "LambdaTransform":
-                raise TypeError(
-                    "LambdaTransform (created by Adapter.apply) could not be deserialized.\n"
-                    "This is probably because the custom transform functions `forward` and "
-                    "`backward` from `Adapter.apply` were not passed as `custom_objects`.\n"
-                    "For example, if your adapter uses\n"
-                    "`Adapter.apply(forward=forward_transform, inverse=inverse_transform)`,\n"
-                    "you have to pass\n"
-                    '`custom_objects={"forward_transform": forward_transform, '
-                    '"inverse_transform": inverse_transform}`\n'
-                    "to the function you use to load the serialized object."
-                ) from e
             raise TypeError(
                 "The transform could not be deserialized properly. "
                 "The most likely reason is that some classes or functions "

bayesflow/adapters/transforms/lambda_transform.py

@@ -1,13 +1,14 @@
 import numpy as np
-
 from keras.saving import (
     deserialize_keras_object as deserialize,
+    register_keras_serializable as serializable,
     serialize_keras_object as serialize,
 )
 
 from .elementwise_transform import ElementwiseTransform
 
 
+@serializable(package="bayesflow.adapters")
 class Log(ElementwiseTransform):
     """Log transforms a variable.
 

bayesflow/adapters/transforms/log.py

@@ -0,0 +1,79 @@
+import numpy as np
+from keras.saving import register_keras_serializable as serializable
+
+from .elementwise_transform import ElementwiseTransform
+
+
+@serializable(package="bayesflow.adapters")
+class NumpyTransform(ElementwiseTransform):
+    """
+    A class to apply element-wise transformations using plain NumPy functions.
+
+    Attributes:
+    ----------
+    _forward : str
+        The name of the NumPy function to apply in the forward transformation.
+    _inverse : str
+        The name of the NumPy function to apply in the inverse transformation.
+    """
+
+    INVERSE_METHODS = {
+        np.arctan: np.tan,
+        np.exp: np.log,
+        np.expm1: np.log1p,
+        np.square: np.sqrt,
+        np.reciprocal: np.reciprocal,
+    }
+    # ensure the map is symmetric
+    INVERSE_METHODS |= {v: k for k, v in INVERSE_METHODS.items()}
+
+    def __init__(self, forward: str, inverse: str = None):
+        """
+        Initializes the NumpyTransform with specified forward and inverse functions.
+
+        Parameters:
+        ----------
+        forward: str
+            The name of the NumPy function to use for the forward transformation.
+        inverse: str, optional
+            The name of the NumPy function to use for the inverse transformation.
+            By default, the inverse is inferred from the forward argument for supported methods.
+        """
+        super().__init__()
+
+        if isinstance(forward, str):
+            forward = getattr(np, forward)
+
+        if not isinstance(forward, np.ufunc):
+            raise ValueError("Forward transformation must be a NumPy Universal Function (ufunc).")
+
+        if inverse is None:
+            if forward not in self.INVERSE_METHODS:
+                raise ValueError(f"Cannot infer inverse for method {forward!r}")
+
+            inverse = self.INVERSE_METHODS[forward]
+
+        if isinstance(inverse, str):
+            inverse = getattr(np, inverse)
+
+        if not isinstance(inverse, np.ufunc):
+            raise ValueError("Inverse transformation must be a NumPy Universal Function (ufunc).")
+
+        self._forward = forward
+        self._inverse = inverse
+
+    @classmethod
+    def from_config(cls, config: dict, custom_objects=None) -> "ElementwiseTransform":
+        return cls(
+            forward=config["forward"],
+            inverse=config["inverse"],
+        )
+
+    def get_config(self) -> dict:
+        return {"forward": self._forward.__name__, "inverse": self._inverse.__name__}
+
+    def forward(self, data: dict[str, any], **kwargs) -> dict[str, any]:
+        return self._forward(data)
+
+    def inverse(self, data: np.ndarray, **kwargs) -> np.ndarray:
+        return self._inverse(data)