Repository: ecmwf/earthkit-data
Pull request #197 (2023-09-19 11:59:53)

Simplify structure of concatenated numpy fieldlist


Patch:
diff --git a/earthkit/data/sources/numpy_list.py b/earthkit/data/sources/numpy_list.py
index 8e8e5a5..aeb1473 100644
--- a/earthkit/data/sources/numpy_list.py
+++ b/earthkit/data/sources/numpy_list.py
@@ -34,6 +34,9 @@ class NumpyField(Field):
         else:
             return self._array.astype(dtype)
 
+    def __repr__(self):
+        return f"{self.__class__.__name__}()"
+
     def write(self, f):
         from earthkit.data.writers import write
 
@@ -52,22 +55,38 @@ class NumpyFieldListCore(PandasMixIn, XarrayMixIn, FieldList):
             if not isinstance(md, Metadata):
                 raise TypeError("metadata must be a subclass of MetaData")
 
-        if self._array.shape[0] != len(self._metadata):
-            import numpy as np
-
-            # we have a single array and a single metadata
-            if len(self._metadata) == 1 and self._shape_match(
-                self._array.shape, self._metadata[0].geography.shape()
-            ):
-                self._array = np.array([self._array])
-            else:
+        if isinstance(self._array, np.ndarray):
+            if self._array.shape[0] != len(self._metadata):
+                # we have a single array and a single metadata
+                if len(self._metadata) == 1 and self._shape_match(
+                    self._array.shape, self._metadata[0].geography.shape()
+                ):
+                    self._array = np.array([self._array])
+                else:
+                    raise ValueError(
+                        (
+                            f"first array dimension ({self._array.shape[0]}) differs "
+                            f"from number of metadata objects ({len(self._metadata)})"
+                        )
+                    )
+        elif isinstance(self._array, list):
+            if len(self._array) != len(self._metadata):
                 raise ValueError(
                     (
-                        f"first array dimension ({self._array.shape[0]}) differs "
+                        f"array len ({len(self._array)}) differs "
                         f"from number of metadata objects ({len(self._metadata)})"
                     )
                 )
 
+            for i, a in enumerate(self._array):
+                if not isinstance(a, np.ndarray):
+                    raise ValueError(
+                        f"All array element must be an ndarray. Type at position={i} is {type(a)}"
+                    )
+
+        else:
+            raise TypeError("array must be an ndarray or a list of ndarrays")
+
         super().__init__(*args, **kwargs)
 
     def _shape_match(self, shape1, shape2):
@@ -84,7 +103,44 @@ class NumpyFieldListCore(PandasMixIn, XarrayMixIn, FieldList):
     @classmethod
     def merge(cls, sources):
         assert all(isinstance(_, NumpyFieldListCore) for _ in sources)
-        return NumpyMultiFieldList(sources)
+        merger = ListMerger(sources)
+        # merger = MultiUnwindMerger(sources)
+        return merger.to_fieldlist()
+
+    def __repr__(self):
+        return f"{self.__class__.__name__}(fields={len(self)})"
+
+
+class MultiUnwindMerger:
+    def __init__(self, sources):
+        self.sources = list(self._flatten(sources))
+
+    def _flatten(self, sources):
+        if isinstance(sources, NumpyMultiFieldList):
+            for s in sources.indexes:
+                yield from self._flatten(s)
+        elif isinstance(sources, list):
+            for s in sources:
+                yield from self._flatten(s)
+        else:
+            yield sources
+
+    def to_fieldlist(self):
+        return NumpyMultiFieldList(self.sources)
+
+
+class ListMerger:
+    def __init__(self, sources):
+        self.sources = sources
+
+    def to_fieldlist(self):
+        array = []
+        metadata = []
+        for s in self.sources:
+            for f in s:
+                array.append(f._array)
+                metadata.append(f._metadata)
+        return NumpyFieldList(array, metadata)
 
 
 class NumpyFieldList(NumpyFieldListCore):
@@ -95,7 +151,9 @@ class NumpyFieldList(NumpyFieldListCore):
             return super().__getitem__(n)
 
     def __len__(self):
-        return self._array.shape[0]
+        return (
+            len(self._array) if isinstance(self._array, list) else self._array.shape[0]
+        )
 
 
 class NumpyMaskFieldList(NumpyFieldListCore, MaskIndex):


Test patch:
diff --git a/tests/numpy_fs/numpy_fs_fixtures.py b/tests/numpy_fs/numpy_fs_fixtures.py
new file mode 100644
index 0000000..37f1e90
--- /dev/null
+++ b/tests/numpy_fs/numpy_fs_fixtures.py
@@ -0,0 +1,87 @@
+#!/usr/bin/env python3
+
+# (C) Copyright 2020 ECMWF.
+#
+# This software is licensed under the terms of the Apache Licence Version 2.0
+# which can be obtained at http://www.apache.org/licenses/LICENSE-2.0.
+# In applying this licence, ECMWF does not waive the privileges and immunities
+# granted to it by virtue of its status as an intergovernmental organisation
+# nor does it submit to any jurisdiction.
+#
+
+import os
+
+import numpy as np
+
+from earthkit.data import from_source
+from earthkit.data.core.fieldlist import FieldList
+from earthkit.data.core.temporary import temp_file
+from earthkit.data.testing import earthkit_examples_file
+
+
+def load_numpy_fs(num):
+    assert num in [1, 2, 3]
+    files = ["test.grib", "test6.grib", "tuv_pl.grib"]
+    files = files[:num]
+
+    ds_in = []
+    md = []
+    for fname in files:
+        ds_in.append(from_source("file", earthkit_examples_file(fname)))
+        md += ds_in[-1].metadata("param")
+
+    ds = []
+    for x in ds_in:
+        ds.append(
+            FieldList.from_numpy(
+                x.values, [m.override(edition=1) for m in x.metadata()]
+            )
+        )
+
+    return (*ds, md)
+
+
+def check_numpy_fs(ds, ds_input, md_full):
+    assert len(ds_input) in [1, 2, 3]
+
+    assert len(ds) == len(md_full)
+    assert ds.metadata("param") == md_full
+    assert np.allclose(ds[0].values, ds_input[0][0].values)
+
+    # check slice
+    r = ds[1]
+    assert r.metadata("param") == "msl"
+
+    if len(ds_input) > 1:
+        r = ds[1:3]
+        assert len(r) == 2
+        assert r.metadata("param") == ["msl", "t"]
+        assert r[0].metadata("param") == "msl"
+        assert r[1].metadata("param") == "t"
+        assert np.allclose(r[0].values, ds_input[0][1].values)
+        assert np.allclose(r[1].values, ds_input[1][0].values)
+
+        # check sel
+        r = ds.sel(shortName="msl")
+        assert len(r) == 1
+        assert r.metadata("shortName") == ["msl"]
+        assert r[0].metadata("param") == "msl"
+        assert np.allclose(r[0].values, ds_input[0][1].values)
+
+    if len(ds_input) == 3:
+        r = ds[1:13:4]
+        assert len(r) == 3
+        assert r.metadata("param") == ["msl", "t", "u"]
+        assert r[0].metadata("param") == "msl"
+        assert r[1].metadata("param") == "t"
+        assert r[2].metadata("param") == "u"
+
+
+def check_save_to_disk(ds, len_ref, meta_ref):
+    tmp = temp_file()
+    ds.save(tmp.path)
+    assert os.path.exists(tmp.path)
+    r_tmp = from_source("file", tmp.path)
+    assert len(r_tmp) == len_ref
+    assert r_tmp.metadata("shortName") == meta_ref
+    r_tmp = None
diff --git a/tests/numpy_fs/test_numpy_fs.py b/tests/numpy_fs/test_numpy_fs.py
index 523949b..c5e4105 100644
--- a/tests/numpy_fs/test_numpy_fs.py
+++ b/tests/numpy_fs/test_numpy_fs.py
@@ -9,20 +9,23 @@
 # nor does it submit to any jurisdiction.
 #
 
-import logging
 import os
+import sys
 
 import numpy as np
+import pytest
 
 from earthkit.data import from_source
 from earthkit.data.core.fieldlist import FieldList
 from earthkit.data.core.temporary import temp_file
 from earthkit.data.testing import earthkit_examples_file
 
-LOG = logging.getLogger(__name__)
+here = os.path.dirname(__file__)
+sys.path.insert(0, here)
+from numpy_fs_fixtures import check_numpy_fs  # noqa: E402
 
 
-def test_numpy_list_grib_single_field():
+def test_numpy_fs_grib_single_field():
     ds = from_source("file", earthkit_examples_file("test.grib"))
 
     assert ds[0].metadata("shortName") == "2t"
@@ -54,7 +57,7 @@ def test_numpy_list_grib_single_field():
     _check_field(r_tmp)
 
 
-def test_numpy_list_grib_multi_field():
+def test_numpy_fs_grib_multi_field():
     ds = from_source("file", earthkit_examples_file("test.grib"))
 
     assert ds[0].metadata("shortName") == "2t"
@@ -85,6 +88,31 @@ def test_numpy_list_grib_multi_field():
         assert f.metadata("name") == "2 metre dewpoint temperature", f"name {i}"
 
 
+def test_numpy_fs_grib_from_list_of_arrays():
+    ds = from_source("file", earthkit_examples_file("test.grib"))
+    md_full = ds.metadata("param")
+    assert len(ds) == 2
+
+    v = [ds[0].values, ds[1].values]
+    md = [f.metadata().override(generatingProcessIdentifier=150) for f in ds]
+    r = FieldList.from_numpy(v, md)
+
+    check_numpy_fs(r, [ds], md_full)
+
+
+def test_numpy_fs_grib_from_list_of_arrays_bad():
+    ds = from_source("file", earthkit_examples_file("test.grib"))
+
+    v = ds[0].values
+    md = [f.metadata().override(generatingProcessIdentifier=150) for f in ds]
+
+    with pytest.raises(ValueError):
+        _ = FieldList.from_numpy(v, md)
+
+    with pytest.raises(ValueError):
+        _ = FieldList.from_numpy([v], md)
+
+
 if __name__ == "__main__":
     from earthkit.data.testing import main
 
diff --git a/tests/numpy_fs/test_numpy_fs_concat.py b/tests/numpy_fs/test_numpy_fs_concat.py
index eae187f..b15b543 100644
--- a/tests/numpy_fs/test_numpy_fs_concat.py
+++ b/tests/numpy_fs/test_numpy_fs_concat.py
@@ -10,109 +10,47 @@
 #
 
 import os
+import sys
 
-import numpy as np
 import pytest
 
 from earthkit.data import from_source
 from earthkit.data.core.fieldlist import FieldList
-from earthkit.data.core.temporary import temp_file
-from earthkit.data.testing import earthkit_examples_file
-
-
-def _check_save_to_disk(ds, len_ref, meta_ref):
-    tmp = temp_file()
-    ds.save(tmp.path)
-    assert os.path.exists(tmp.path)
-    r_tmp = from_source("file", tmp.path)
-    assert len(r_tmp) == len_ref
-    assert r_tmp.metadata("shortName") == meta_ref
-    r_tmp = None
-
-
-def _check_contents(ds, ds_input, md_full):
-    assert len(ds_input) in [2, 3]
-
-    assert len(ds) == len(md_full)
-    assert ds.metadata("param") == md_full
-
-    # check slice
-    r = ds[1]
-    assert r.metadata("param") == "msl"
-
-    r = ds[1:3]
-    assert len(r) == 2
-    assert r.metadata("param") == ["msl", "t"]
-    assert r[0].metadata("param") == "msl"
-    assert r[1].metadata("param") == "t"
-    assert np.allclose(r[0].values, ds_input[0][1].values)
-    assert np.allclose(r[1].values, ds_input[1][0].values)
-    with pytest.raises(ValueError):
-        _ = r.values
-
-    # check sel
-    r = ds.sel(shortName="msl")
-    assert len(r) == 1
-    assert r.metadata("shortName") == ["msl"]
-    assert r[0].metadata("param") == "msl"
-    assert np.allclose(r[0].values, ds_input[0][1].values)
-
-    if len(ds_input) == 3:
-        r = ds[1:13:4]
-        assert len(r) == 3
-        assert r.metadata("param") == ["msl", "t", "u"]
-        assert r[0].metadata("param") == "msl"
-        assert r[1].metadata("param") == "t"
-        assert r[2].metadata("param") == "u"
-
-
-def _prepare_ds(num):
-    assert num in [1, 2, 3]
-    files = ["test.grib", "test6.grib", "tuv_pl.grib"]
-    files = files[:num]
-
-    ds_in = []
-    md = []
-    for fname in files:
-        ds_in.append(from_source("file", earthkit_examples_file(fname)))
-        md += ds_in[-1].metadata("param")
-
-    ds = []
-    for x in ds_in:
-        ds.append(
-            FieldList.from_numpy(
-                x.values, [m.override(edition=1) for m in x.metadata()]
-            )
-        )
-
-    return (*ds, md)
+
+here = os.path.dirname(__file__)
+sys.path.insert(0, here)
+from numpy_fs_fixtures import (  # noqa: E402
+    check_numpy_fs,
+    check_save_to_disk,
+    load_numpy_fs,
+)
 
 
 @pytest.mark.parametrize("mode", ["oper", "multi"])
-def test_numpy_list_grib_concat_2a(mode):
-    ds1, ds2, md = _prepare_ds(2)
+def test_numpy_fs_grib_concat_2a(mode):
+    ds1, ds2, md = load_numpy_fs(2)
 
     if mode == "oper":
         ds = ds1 + ds2
     else:
         ds = from_source("multi", ds1, ds2)
 
-    _check_contents(ds, [ds1, ds2], md)
-    _check_save_to_disk(ds, 8, md)
+    check_numpy_fs(ds, [ds1, ds2], md)
+    check_save_to_disk(ds, 8, md)
 
 
-def test_numpy_list_grib_concat_2b():
-    ds1, ds2, md = _prepare_ds(2)
+def test_numpy_fs_grib_concat_2b():
+    ds1, ds2, md = load_numpy_fs(2)
     ds1_ori = ds1
     ds1 += ds2
 
-    _check_contents(ds1, [ds1_ori, ds2], md)
-    _check_save_to_disk(ds1, 8, md)
+    check_numpy_fs(ds1, [ds1_ori, ds2], md)
+    check_save_to_disk(ds1, 8, md)
 
 
 @pytest.mark.parametrize("mode", ["oper", "multi"])
-def test_numpy_list_grib_concat_3a(mode):
-    ds1, ds2, ds3, md = _prepare_ds(3)
+def test_numpy_fs_grib_concat_3a(mode):
+    ds1, ds2, ds3, md = load_numpy_fs(3)
 
     if mode == "oper":
         ds = ds1 + ds2
@@ -121,69 +59,69 @@ def test_numpy_list_grib_concat_3a(mode):
         ds = from_source("multi", ds1, ds2)
         ds = from_source("multi", ds, ds3)
 
-    _check_contents(ds, [ds1, ds2, ds3], md)
-    _check_save_to_disk(ds, 26, md)
+    check_numpy_fs(ds, [ds1, ds2, ds3], md)
+    check_save_to_disk(ds, 26, md)
 
 
 @pytest.mark.parametrize("mode", ["oper", "multi"])
-def test_numpy_list_grib_concat_3b(mode):
-    ds1, ds2, ds3, md = _prepare_ds(3)
+def test_numpy_fs_grib_concat_3b(mode):
+    ds1, ds2, ds3, md = load_numpy_fs(3)
 
     if mode == "oper":
         ds = ds1 + ds2 + ds3
     else:
         ds = from_source("multi", ds1, ds2, ds3)
 
-    _check_contents(ds, [ds1, ds2, ds3], md)
-    _check_save_to_disk(ds, 26, md)
+    check_numpy_fs(ds, [ds1, ds2, ds3], md)
+    check_save_to_disk(ds, 26, md)
 
 
-def test_numpy_list_grib_from_empty_1():
+def test_numpy_fs_grib_from_empty_1():
     ds_e = FieldList()
-    ds, md = _prepare_ds(1)
+    ds, md = load_numpy_fs(1)
     ds1 = ds_e + ds
     assert id(ds1) == id(ds)
     assert len(ds1) == 2
     assert ds1.metadata("param") == md
-    _check_save_to_disk(ds1, 2, md)
+    check_save_to_disk(ds1, 2, md)
 
 
-def test_numpy_list_grib_from_empty_2():
+def test_numpy_fs_grib_from_empty_2():
     ds_e = FieldList()
-    ds, md = _prepare_ds(1)
+    ds, md = load_numpy_fs(1)
     ds1 = ds + ds_e
     assert id(ds1) == id(ds)
     assert len(ds1) == 2
     assert ds1.metadata("param") == md
-    _check_save_to_disk(ds1, 2, md)
+    check_save_to_disk(ds1, 2, md)
 
 
-def test_numpy_list_grib_from_empty_3():
+def test_numpy_fs_grib_from_empty_3():
     ds_e = FieldList()
-    ds1, ds2, md = _prepare_ds(2)
+    ds1, ds2, md = load_numpy_fs(2)
     ds = ds_e + ds1 + ds2
 
-    _check_contents(ds, [ds1, ds2], md)
-    _check_save_to_disk(ds, 8, md)
+    check_numpy_fs(ds, [ds1, ds2], md)
+    check_save_to_disk(ds, 8, md)
 
 
-def test_numpy_list_grib_from_empty_4():
+def test_numpy_fs_grib_from_empty_4():
     ds = FieldList()
-    ds1, md = _prepare_ds(1)
+    ds1, md = load_numpy_fs(1)
     ds += ds1
     assert id(ds) == id(ds1)
     assert len(ds) == 2
     assert ds.metadata("param") == md
-    _check_save_to_disk(ds, 2, md)
+    check_save_to_disk(ds, 2, md)
 
 
-def test_numpy_list_grib_from_empty_5():
+def test_numpy_fs_grib_from_empty_5():
     ds = FieldList()
-    ds1, ds2, md = _prepare_ds(2)
+    ds1, ds2, md = load_numpy_fs(2)
     ds += ds1 + ds2
 
-    _check_contents(ds, [ds1, ds2], md)
-    _check_save_to_disk(ds, 8, md)
+    check_numpy_fs(ds, [ds1, ds2], md)
+    check_save_to_disk(ds, 8, md)
 
 
 if __name__ == "__main__":