add scale and shift transforms

Author: LarsKue

bayesflow/adapters/adapter.py

@@ -555,6 +555,18 @@ def rename(self, from_key: str, to_key: str):
         self.transforms.append(Rename(from_key, to_key))
         return self
 
+    def scale(self, keys: str | Sequence[str], by: float | np.ndarray):
+        from .transforms import Scale
+
+        self.transforms.append(MapTransform({key: Scale(scale=by) for key in keys}))
+        return self
+
+    def shift(self, keys: str | Sequence[str], by: float | np.ndarray):
+        from .transforms import Shift
+
+        self.transforms.append(MapTransform({key: Shift(shift=by) for key in keys}))
+        return self
+
     def sqrt(self, keys: str | Sequence[str]):
         """Append an :py:class:`~transforms.Sqrt` transform to the adapter.
 

bayesflow/adapters/transforms/__init__.py

@@ -14,6 +14,8 @@
 from .numpy_transform import NumpyTransform
 from .one_hot import OneHot
 from .rename import Rename
+from .scale import Scale
+from .shift import Shift
 from .sqrt import Sqrt
 from .standardize import Standardize
 from .to_array import ToArray

bayesflow/adapters/transforms/scale.py

@@ -0,0 +1,14 @@
+import numpy as np
+
+from .elementwise_transform import ElementwiseTransform
+
+
+class Scale(ElementwiseTransform):
+    def __init__(self, scale: float | np.ndarray):
+        self.scale = scale
+
+    def forward(self, data: np.ndarray, **kwargs) -> np.ndarray:
+        return data * self.scale
+
+    def inverse(self, data: np.ndarray, **kwargs) -> np.ndarray:
+        return data / self.scale

bayesflow/adapters/transforms/shift.py

@@ -0,0 +1,16 @@
+import numpy as np
+
+from .elementwise_transform import ElementwiseTransform
+
+class Shift(ElementwiseTransform):
+    def __init__(self, shift: float | np.ndarray):
+        self.shift = shift
+
+    def forward(self, data: np.ndarray, **kwargs) -> np.ndarray:
+        return data + self.shift
+
+    def inverse(self, data: np.ndarray, **kwargs) -> np.ndarray:
+        return data - self.shift
+
+
+