From 3938df085e6f34d8003163dbb734fe94a8c42bbc Mon Sep 17 00:00:00 2001
From: Andreas Kloeckner <inform@tiker.net>
Date: Mon, 14 Jun 2021 15:35:27 -0500
Subject: [PATCH] Special-case serialization of 1/2D object arrays
---
arraycontext/container/__init__.py | 11 ++++++++++-
1 file changed, 10 insertions(+), 1 deletion(-)
diff --git a/arraycontext/container/__init__.py b/arraycontext/container/__init__.py
index 397af49..44b7e9e 100644
--- a/arraycontext/container/__init__.py
+++ b/arraycontext/container/__init__.py
@@ -172,7 +172,16 @@ def _serialize_ndarray_container(ary: np.ndarray) -> Iterable[Tuple[Any, Any]]:
raise ValueError(
f"only object arrays are supported, given dtype '{ary.dtype}'")
- return np.ndenumerate(ary)
+ # special-cased for speed
+ if ary.ndim == 1:
+ return [(i, ary[i]) for i in range(ary.shape[0])]
+ elif ary.ndim == 2:
+ return [((i, j), ary[i, j])
+ for i in range(ary.shape[0])
+ for j in range(ary.shape[1])
+ ]
+ else:
+ return np.ndenumerate(ary)
@deserialize_container.register(np.ndarray)
--
GitLab