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

Author: LarsKue

bayesflow/adapters/transforms/numpy_transform.py

@@ -1,9 +1,11 @@
 import numpy as np
+from keras.saving import register_keras_serializable as serializable
 
 from bayesflow.utils import filter_kwargs
 from .elementwise_transform import ElementwiseTransform
 
 
+@serializable(package="bayesflow.adapters")
 class NumpyTransform(ElementwiseTransform):
     """
     A class to apply element-wise transformations using plain NumPy functions.
@@ -17,45 +19,46 @@ class NumpyTransform(ElementwiseTransform):
     """
 
     INVERSE_METHODS = {
-        "arctan": "tan",
-        "exp": "log",
-        "expm1": "log1p",
-        "square": "sqrt",
-        "reciprocal": "reciprocal",
+        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: np.ufunc | str, inverse: np.ufunc | str = None):
+    def __init__(self, forward: str, inverse: str = None):
         """
         Initializes the NumpyTransform with specified forward and inverse functions.
 
         Parameters:
         ----------
-        forward : str
+        forward: str
             The name of the NumPy function to use for the forward transformation.
-        inverse : str
+        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, np.ufunc):
-            forward = forward.__name__
+        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]
-        elif isinstance(inverse, np.ufunc):
-            inverse = inverse.__name__
 
-        if forward not in dir(np):
-            raise ValueError(f"Method {forward!r} not found in numpy version {np.__version__}")
+        if isinstance(inverse, str):
+            inverse = getattr(np, inverse)
 
-        if inverse not in dir(np):
-            raise ValueError(f"Method {inverse!r} not found in numpy version {np.__version__}")
+        if not isinstance(inverse, np.ufunc):
+            raise ValueError("Inverse transformation must be a NumPy Universal Function (ufunc).")
 
         self._forward = forward
         self._inverse = inverse
@@ -68,14 +71,12 @@ def from_config(cls, config: dict, custom_objects=None) -> "ElementwiseTransform
         )
 
     def get_config(self) -> dict:
-        return {"forward": self._forward, "inverse": self._inverse}
+        return {"forward": self._forward.__name__, "inverse": self._inverse.__name__}
 
     def forward(self, data: dict[str, any], **kwargs) -> dict[str, any]:
-        forward = getattr(np, self._forward)
-        kwargs = filter_kwargs(kwargs, forward)
-        return forward(data, **kwargs)
+        kwargs = filter_kwargs(kwargs, self._forward)
+        return self._forward(data, **kwargs)
 
     def inverse(self, data: np.ndarray, **kwargs) -> np.ndarray:
-        inverse = getattr(np, self._inverse)
-        kwargs = filter_kwargs(kwargs, inverse)
-        return inverse(data, **kwargs)
+        kwargs = filter_kwargs(kwargs, self._inverse)
+        return self._inverse(data, **kwargs)

tests/test_adapters/conftest.py

@@ -18,8 +18,8 @@ def adapter():
         .expand_dims(["z1"], axis=2)
         .log("p1")
         .constrain("p2", lower=0)
-        .apply(include="p2", forward=np.exp, inverse=np.log)
-        .apply(include="p2", forward="logp1")
+        .apply(include="p2", forward="exp", inverse="log")
+        .apply(include="p2", forward="log1p")
         .standardize(exclude=["t1", "t2", "o1"])
         .drop("d1")
         .one_hot("o1", 10)