handle nesting: ConvertDType, ToArray, relax Concatenate

Concatenate can be equal to rename if only one key is supplied. By not
calling concatenate in that case, we can accept arbitrary inputs in the
transform, as long as only one is supplied. This simplifies things e.g.
in the `BasicWorkflow`, where the user passes the `summary_variables` to
concatenate, which may be a single dict, which does not need to be
concatenated.

Author: vpratz

bayesflow/adapters/transforms/concatenate.py

@@ -49,7 +49,7 @@ def get_config(self) -> dict:
         return serialize(config)
 
     def forward(self, data: dict[str, any], *, strict: bool = True, **kwargs) -> dict[str, any]:
-        if not strict and self.indices is None:
+        if not strict and self.indices is None and len(self.keys) != 1:
             raise ValueError("Cannot call `forward` with `strict=False` before calling `forward` with `strict=True`.")
 
         # copy to avoid side effects
@@ -69,6 +69,10 @@ def forward(self, data: dict[str, any], *, strict: bool = True, **kwargs) -> dic
                 data.pop(key)
 
             return data
+        elif len(required_keys) == 1:
+            # only a rename
+            data[self.into] = data.pop(self.keys[0])
+            return data
 
         if self.indices is None:
             # remember the indices of the parts in the concatenated array
@@ -86,7 +90,7 @@ def forward(self, data: dict[str, any], *, strict: bool = True, **kwargs) -> dic
         return data
 
     def inverse(self, data: dict[str, any], *, strict: bool = False, **kwargs) -> dict[str, any]:
-        if self.indices is None:
+        if self.indices is None and len(self.keys) != 1:
             raise RuntimeError("Cannot call `inverse` before calling `forward` at least once.")
 
         # copy to avoid side effects
@@ -98,6 +102,9 @@ def inverse(self, data: dict[str, any], *, strict: bool = False, **kwargs) -> di
         elif self.into not in data:
             # nothing to do
             return data
+        elif len(self.keys) == 1:
+            data[self.keys[0]] = data.pop(self.into)
+            return data
 
         # split the concatenated array and remove the concatenated key
         keys = self.keys
@@ -141,7 +148,7 @@ def log_det_jac(
         available_keys = set(log_det_jac.keys())
         common_keys = available_keys & required_keys
 
-        if len(common_keys) == 0:
+        if len(common_keys) == 0 or len(self.keys) == 1:
             return log_det_jac
 
         parts = [log_det_jac.pop(key) for key in common_keys]

bayesflow/adapters/transforms/convert_dtype.py

@@ -1,4 +1,5 @@
 import numpy as np
+from keras.tree import map_structure
 
 from bayesflow.utils.serialization import serializable, serialize
 
@@ -32,7 +33,7 @@ def get_config(self) -> dict:
         return serialize(config)
 
     def forward(self, data: np.ndarray, **kwargs) -> np.ndarray:
-        return data.astype(self.to_dtype, copy=False)
+        return map_structure(lambda d: d.astype(self.to_dtype, copy=False), data)
 
     def inverse(self, data: np.ndarray, **kwargs) -> np.ndarray:
-        return data.astype(self.from_dtype, copy=False)
+        return map_structure(lambda d: d.astype(self.from_dtype, copy=False), data)

bayesflow/adapters/transforms/to_array.py

@@ -2,6 +2,7 @@
 
 import numpy as np
 
+from bayesflow.utils.tree import map_dict, get_value_at_path, map_dict_with_path
 from bayesflow.utils.serialization import serializable, serialize
 
 from .elementwise_transform import ElementwiseTransform
@@ -35,13 +36,36 @@ def get_config(self) -> dict:
 
     def forward(self, data: any, **kwargs) -> np.ndarray:
         if self.original_type is None:
-            self.original_type = type(data)
+            if isinstance(data, dict):
+                self.original_type = map_dict(type, data)
+            else:
+                self.original_type = type(data)
 
+        if isinstance(self.original_type, dict):
+            # use self.original_type in check to preserve serializablitiy
+            return map_dict(np.asarray, data)
         return np.asarray(data)
 
-    def inverse(self, data: np.ndarray, **kwargs) -> any:
+    def inverse(self, data: np.ndarray | dict, **kwargs) -> any:
         if self.original_type is None:
             raise RuntimeError("Cannot call `inverse` before calling `forward` at least once.")
+        if isinstance(self.original_type, dict):
+            # use self.original_type in check to preserve serializablitiy
+
+            def restore_original_type(path, value):
+                try:
+                    original_type = get_value_at_path(self.original_type, path)
+                    return original_type(value)
+                except KeyError:
+                    pass
+                except TypeError:
+                    pass
+                except ValueError:
+                    # separate statements, as optree does not allow (KeyError | TypeError | ValueError)
+                    pass
+                return value
+
+            return map_dict_with_path(restore_original_type, data)
 
         if issubclass(self.original_type, Number):
             try:

bayesflow/utils/__init__.py

@@ -7,6 +7,7 @@
     logging,
     numpy_utils,
     serialization,
+    tree,
 )
 
 from .callbacks import detailed_loss_callback
@@ -104,4 +105,4 @@
 
 from ._docs import _add_imports_to_all
 
-_add_imports_to_all(include_modules=["keras_utils", "logging", "numpy_utils", "serialization"])
+_add_imports_to_all(include_modules=["keras_utils", "logging", "numpy_utils", "serialization", "tree"])

bayesflow/utils/tree.py

@@ -0,0 +1,69 @@
+import optree
+
+
+def flatten_shape(structure):
+    def is_shape_tuple(x):
+        return isinstance(x, (list, tuple)) and all(isinstance(e, (int, type(None))) for e in x)
+
+    leaves, _ = optree.tree_flatten(
+        structure,
+        is_leaf=is_shape_tuple,
+        none_is_leaf=True,
+        namespace="keras",
+    )
+    return leaves
+
+
+def map_dict(func, *structures):
+    def is_not_dict(x):
+        return not isinstance(x, dict)
+
+    if not structures:
+        raise ValueError("Must provide at least one structure")
+
+    # Add check for same structures, otherwise optree just maps to shallowest.
+    def func_with_check(*args):
+        if not all(optree.tree_is_leaf(s, is_leaf=is_not_dict, none_is_leaf=True, namespace="keras") for s in args):
+            raise ValueError("Structures don't have the same nested structure.")
+        return func(*args)
+
+    map_func = func_with_check if len(structures) > 1 else func
+
+    return optree.tree_map(
+        map_func,
+        *structures,
+        is_leaf=is_not_dict,
+        none_is_leaf=True,
+        namespace="keras",
+    )
+
+
+def map_dict_with_path(func, *structures):
+    def is_not_dict(x):
+        return not isinstance(x, dict)
+
+    if not structures:
+        raise ValueError("Must provide at least one structure")
+
+    # Add check for same structures, otherwise optree just maps to shallowest.
+    def func_with_check(*args):
+        if not all(optree.tree_is_leaf(s, is_leaf=is_not_dict, none_is_leaf=True, namespace="keras") for s in args):
+            raise ValueError("Structures don't have the same nested structure.")
+        return func(*args)
+
+    map_func = func_with_check if len(structures) > 1 else func
+
+    return optree.tree_map_with_path(
+        map_func,
+        *structures,
+        is_leaf=is_not_dict,
+        none_is_leaf=True,
+        namespace="keras",
+    )
+
+
+def get_value_at_path(structure, path):
+    output = structure
+    for accessor in path:
+        output = output.__getitem__(accessor)
+    return output

tests/test_adapters/conftest.py

@@ -13,7 +13,9 @@ def serializable_fn(x):
 
     return (
         Adapter()
+        .group(["p1", "p2"], into="ps", prefix="p")
         .to_array()
+        .ungroup("ps", prefix="p")
         .as_set(["s1", "s2"])
         .broadcast("t1", to="t2")
         .as_time_series(["t1", "t2"])
@@ -37,8 +39,6 @@ def serializable_fn(x):
         .rename("o1", "o2")
         .random_subsample("s3", sample_size=33, axis=0)
         .take("s3", indices=np.arange(0, 32), axis=0)
-        .group(["p1", "p2"], into="ps", prefix="p")
-        .ungroup("ps", prefix="p")
     )
 
 

tests/test_workflows/conftest.py

@@ -81,13 +81,6 @@ def sample(self, batch_shape: Shape, num_observations: int = 4) -> dict[str, Ten
 
             x = mean[:, None] + noise
 
-            return dict(mean=mean, a=x, b=x)
+            return dict(mean=mean, observables=dict(a=x, b=x))
 
     return FusionSimulator()
-
-
-@pytest.fixture
-def fusion_adapter():
-    from bayesflow import Adapter
-
-    return Adapter.create_default(["mean"]).group(["a", "b"], "summary_variables")

tests/test_workflows/test_basic_workflow.py

@@ -34,14 +34,13 @@ def test_basic_workflow(tmp_path, inference_network, summary_network):
     assert samples["parameters"].shape == (5, 3, 2)
 
 
-def test_basic_workflow_fusion(
-    tmp_path, fusion_inference_network, fusion_summary_network, fusion_simulator, fusion_adapter
-):
+def test_basic_workflow_fusion(tmp_path, fusion_inference_network, fusion_summary_network, fusion_simulator):
     workflow = bf.BasicWorkflow(
-        adapter=fusion_adapter,
         inference_network=fusion_inference_network,
         summary_network=fusion_summary_network,
         simulator=fusion_simulator,
+        inference_variables=["mean"],
+        summary_variables=["observables"],
         checkpoint_filepath=str(tmp_path),
     )