add a leaf_class parameter to flatten

Author: alexfikl

arraycontext/container/traversal.py

@@ -583,7 +583,10 @@ def thaw(ary: ArrayOrContainerT, actx: ArrayContext) -> ArrayOrContainerT:
 
 # {{{ flatten / unflatten
 
-def flatten(ary: ArrayOrContainerT, actx: ArrayContext) -> Any:
+def flatten(
+        ary: ArrayOrContainerT, actx: ArrayContext, *,
+        leaf_class: Optional[type] = None,
+        ) -> Any:
     """Convert all arrays in the :class:`~arraycontext.ArrayContainer`
     into single flat array of a type :attr:`arraycontext.ArrayContext.array_types`.
 
@@ -629,12 +632,36 @@ def _flatten(subary: ArrayOrContainerT) -> None:
             for _, isubary in iterable:
                 _flatten(isubary)
 
-    _flatten(ary)
+    def _flatten_without_leaf_class(subary: ArrayOrContainerT) -> Any:
+        nonlocal result, common_dtype
+        result = []
+        common_dtype = None
+
+        _flatten(ary)
+
+        if len(result) == 1:
+            return result[0]
+        else:
+            return actx.np.concatenate(result)
+
+    def _flatten_with_leaf_class(subary: ArrayOrContainerT) -> Any:
+        if type(subary) is leaf_class:
+            return _flatten_without_leaf_class(subary)
+
+        try:
+            iterable = serialize_container(subary)
+        except NotAnArrayContainerError:
+            return subary
+        else:
+            return deserialize_container(subary, [
+                (key, _flatten_without_leaf_class(isubary))
+                for key, isubary in iterable
+                ])
 
-    if len(result) == 1:
-        return result[0]
+    if leaf_class is None:
+        return _flatten_without_leaf_class(ary)
     else:
-        return actx.np.concatenate(result)
+        return _flatten_with_leaf_class(ary)
 
 
 def unflatten(

run-pylint.sh

@@ -20,4 +20,4 @@ if [[ -f .pylintrc-local.yml ]]; then
     PYLINT_RUNNER_ARGS+=" --yaml-rcfile=.pylintrc-local.yml"
 fi
 
-python .run-pylint.py $PYLINT_RUNNER_ARGS $(basename $PWD) test/test_*.py examples "$@"
+PYTHONWARNINGS=ignore python .run-pylint.py $PYLINT_RUNNER_ARGS $(basename $PWD) test/test_*.py examples "$@"

test/test_arraycontext.py

@@ -1064,6 +1064,26 @@ def test_flatten_array_container_failure(actx_factory):
         # cannot unflatten partially
         unflatten(ary, flat_ary[:-1], actx)
 
+
+def test_flatten_with_leaf_class(actx_factory):
+    actx = actx_factory()
+
+    from arraycontext import flatten
+    arys = _get_test_containers(actx, shapes=512)
+
+    flat = flatten(arys[0], actx, leaf_class=DOFArray)
+    assert isinstance(flat, actx.array_types)
+
+    flat = flatten(arys[1], actx, leaf_class=DOFArray)
+    assert isinstance(flat, np.ndarray) and flat.dtype == object
+    assert all(isinstance(entry, actx.array_types) for entry in flat)
+
+    flat = flatten(arys[3], actx, leaf_class=DOFArray)
+    assert isinstance(flat, MyContainer)
+    assert isinstance(flat.mass, actx.array_types)
+    assert isinstance(flat.enthalpy, actx.array_types)
+    assert all(isinstance(entry, actx.array_types) for entry in flat.momentum)
+
 # }}}