Adapter: Minor fixes and more thorough testing (#317)

* FilterTransform: deseralize kwargs and transform_constructor

* The deserialize call for `kwargs` was missing.
* The constructor can be (de)serialized using get_registered_name and
  get_registered_object

Added meaningful error message, if functions are not registered and not
passed as custom objects.

* transforms: ensure serialization cycle consistency

Keras deserializes tuples as lists. This can cause problems, and has to
be undone manually.

* adapter: expand tests to include all transforms

* also, expand test to ensure consistency of outputs before and after
  serialization

Author: vpratz

bayesflow/adapters/transforms/broadcast.py

@@ -79,12 +79,19 @@ def __init__(
 
     @classmethod
     def from_config(cls, config: dict, custom_objects=None) -> "Broadcast":
+        # Deserialize turns tuples to lists, undo it if necessary
+        exclude = deserialize(config["exclude"], custom_objects)
+        exclude = tuple(exclude) if isinstance(exclude, list) else exclude
+        expand = deserialize(config["expand"], custom_objects)
+        expand = tuple(expand) if isinstance(expand, list) else expand
+        squeeze = deserialize(config["squeeze"], custom_objects)
+        squeeze = tuple(squeeze) if isinstance(squeeze, list) else squeeze
         return cls(
             keys=deserialize(config["keys"], custom_objects),
             to=deserialize(config["to"], custom_objects),
-            expand=deserialize(config["expand"], custom_objects),
-            exclude=deserialize(config["exclude"], custom_objects),
-            squeeze=deserialize(config["squeeze"], custom_objects),
+            expand=expand,
+            exclude=exclude,
+            squeeze=squeeze,
         )
 
     def get_config(self) -> dict:

bayesflow/adapters/transforms/filter_transform.py

@@ -4,6 +4,8 @@
 import numpy as np
 from keras.saving import (
     deserialize_keras_object as deserialize,
+    get_registered_name,
+    get_registered_object,
     register_keras_serializable as serializable,
     serialize_keras_object as serialize,
 )
@@ -79,21 +81,33 @@ def extra_repr(self) -> str:
 
     @classmethod
     def from_config(cls, config: dict, custom_objects=None) -> "Transform":
-        def transform_constructor(*args, **kwargs):
-            raise RuntimeError(
-                "Instantiating new elementwise transforms on a deserialized FilterTransform is not yet supported (and"
-                "may never be). As a work-around, you can manually register the elementwise transform constructor after"
-                "deserialization:\n"
-                "obj = deserialize(config)\n"
-                "obj.transform_constructor = MyElementwiseTransform"
-            )
-
+        transform_constructor = get_registered_object(config["transform_constructor"])
+        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 "
+                "are not known during deserialization. Please pass them as `custom_objects`."
+            ) from e
         instance = cls(
             transform_constructor=transform_constructor,
             predicate=deserialize(config["predicate"], custom_objects),
             include=deserialize(config["include"], custom_objects),
             exclude=deserialize(config["exclude"], custom_objects),
-            **config["kwargs"],
+            **kwargs,
         )
 
         instance.transform_map = deserialize(config["transform_map"])
@@ -102,6 +116,7 @@ def transform_constructor(*args, **kwargs):
 
     def get_config(self) -> dict:
         return {
+            "transform_constructor": get_registered_name(self.transform_constructor),
             "predicate": serialize(self.predicate),
             "include": serialize(self.include),
             "exclude": serialize(self.exclude),

bayesflow/adapters/transforms/standardize.py

@@ -41,10 +41,14 @@ def __init__(
 
     @classmethod
     def from_config(cls, config: dict, custom_objects=None) -> "Standardize":
+        # Deserialize turns tuples to lists, undo it if necessary
+        deserialized_axis = deserialize(config["axis"], custom_objects)
+        if isinstance(deserialized_axis, list):
+            deserialized_axis = tuple(deserialized_axis)
         return cls(
             mean=deserialize(config["mean"], custom_objects),
             std=deserialize(config["std"], custom_objects),
-            axis=deserialize(config["axis"], custom_objects),
+            axis=deserialized_axis,
             momentum=deserialize(config["momentum"], custom_objects),
         )
 

tests/test_adapters/conftest.py

@@ -12,7 +12,7 @@ def inverse_transform(x):
 
 @pytest.fixture()
 def custom_objects():
-    return globals() | np.__dict__
+    return dict(forward_transform=forward_transform, inverse_transform=inverse_transform)
 
 
 @pytest.fixture()
@@ -22,15 +22,22 @@ def adapter():
     d = (
         Adapter()
         .to_array()
-        .convert_dtype("float64", "float32")
+        .as_set(["s1", "s2"])
+        .broadcast("t1", to="t2")
+        .as_time_series(["t1", "t2"])
+        .convert_dtype("float64", "float32", exclude="o1")
         .concatenate(["x1", "x2"], into="x")
         .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)
         .constrain("p2", lower=0)
-        .standardize()
+        .standardize(exclude=["t1", "t2", "o1"])
+        .drop("d1")
+        .one_hot("o1", 10)
+        .keep(["x", "y", "z1", "p1", "p2", "s1", "s2", "t1", "t2", "o1"])
+        .rename("o1", "o2")
     )
 
     return d
@@ -46,4 +53,11 @@ def random_data():
         "z1": np.random.standard_normal(size=(32, 2)),
         "p1": np.random.lognormal(size=(32, 2)),
         "p2": np.random.lognormal(size=(32, 2)),
+        "s1": np.random.standard_normal(size=(32, 3, 2)),
+        "s2": np.random.standard_normal(size=(32, 3, 2)),
+        "t1": np.zeros((3, 2)),
+        "t2": np.ones((32, 3, 2)),
+        "d1": np.random.standard_normal(size=(32, 2)),
+        "d2": np.random.standard_normal(size=(32, 2)),
+        "o1": np.random.randint(0, 9, size=(32, 2)),
     }

tests/test_adapters/test_adapters.py

@@ -10,15 +10,24 @@ def test_cycle_consistency(adapter, random_data):
     deprocessed = adapter(processed, inverse=True)
 
     for key, value in random_data.items():
+        if key in ["d1", "d2"]:
+            # dropped
+            continue
         assert key in deprocessed
         assert np.allclose(value, deprocessed[key])
 
 
-def test_serialize_deserialize(adapter, custom_objects):
+def test_serialize_deserialize(adapter, custom_objects, random_data):
+    processed = adapter(random_data)
     serialized = serialize(adapter)
     deserialized = deserialize(serialized, custom_objects)
     reserialized = serialize(deserialized)
 
     assert reserialized.keys() == serialized.keys()
     for key in reserialized:
         assert reserialized[key] == serialized[key]
+
+    random_data["foo"] = random_data["x1"]
+    deserialized_processed = deserialized(random_data)
+    for key, value in processed.items():
+        assert np.allclose(value, deserialized_processed[key])