Repository: ecmwf/earthkit-data
Pull request #185 (2023-09-15 16:31:43)

Make save work on concatenated numpy fieldlists


Patch:
diff --git a/earthkit/data/sources/numpy_list.py b/earthkit/data/sources/numpy_list.py
index 5212224..77b0d18 100644
--- a/earthkit/data/sources/numpy_list.py
+++ b/earthkit/data/sources/numpy_list.py
@@ -38,7 +38,7 @@ class NumpyField(Field):
         write(f, self.values, self._metadata, check_nans=True)
 
 
-class NumpyFieldList(FieldList):
+class NumpyFieldListCore(FieldList):
     def __init__(self, array, metadata, *args, **kwargs):
         self._array = array
         self._metadata = metadata
@@ -68,12 +68,6 @@ class NumpyFieldList(FieldList):
 
         super().__init__(*args, **kwargs)
 
-    def __getitem__(self, n):
-        return NumpyField(self._array[n], self._metadata[n])
-
-    def __len__(self):
-        return self._array.shape[0]
-
     def _shape_match(self, shape1, shape2):
         if shape1 == shape2:
             return True
@@ -81,12 +75,29 @@ class NumpyFieldList(FieldList):
             return True
         return False
 
+    @classmethod
+    def new_mask_index(self, *args, **kwargs):
+        return NumpyMaskFieldList(*args, **kwargs)
+
+    @classmethod
+    def merge(cls, sources):
+        assert all(isinstance(_, NumpyFieldListCore) for _ in sources)
+        return NumpyMultiFieldList(sources)
+
+
+class NumpyFieldList(NumpyFieldListCore):
+    def __getitem__(self, n):
+        return NumpyField(self._array[n], self._metadata[n])
+
+    def __len__(self):
+        return self._array.shape[0]
+
 
-class NumpyMaskFieldList(NumpyFieldList, MaskIndex):
+class NumpyMaskFieldList(NumpyFieldListCore, MaskIndex):
     def __init__(self, *args, **kwargs):
         MaskIndex.__init__(self, *args, **kwargs)
 
 
-class NumpyMultiFieldList(NumpyFieldList, MultiIndex):
+class NumpyMultiFieldList(NumpyFieldListCore, MultiIndex):
     def __init__(self, *args, **kwargs):
         MultiIndex.__init__(self, *args, **kwargs)


Test patch:
diff --git a/tests/grib/test_grib_concat.py b/tests/grib/test_grib_concat.py
index 2c1630b..2111855 100644
--- a/tests/grib/test_grib_concat.py
+++ b/tests/grib/test_grib_concat.py
@@ -9,7 +9,10 @@
 # nor does it submit to any jurisdiction.
 #
 
+import os
+
 from earthkit.data import from_source
+from earthkit.data.core.temporary import temp_file
 from earthkit.data.testing import earthkit_examples_file
 
 
@@ -22,6 +25,15 @@ def test_grib_concat():
     md = ds1.metadata("param") + ds2.metadata("param")
     assert ds.metadata("param") == md
 
+    # save to disk
+    tmp = temp_file()
+    ds.save(tmp.path)
+    assert os.path.exists(tmp.path)
+    r_tmp = from_source("file", tmp.path)
+    assert len(r_tmp) == 8
+    assert r_tmp.metadata("shortName") == ["2t", "msl", "t", "u", "v", "t", "u", "v"]
+    r_tmp = None
+
 
 if __name__ == "__main__":
     from earthkit.data.testing import main
diff --git a/tests/sources/test_numpy_list.py b/tests/sources/test_numpy_list.py
index 2fd7be6..1fd60c5 100644
--- a/tests/sources/test_numpy_list.py
+++ b/tests/sources/test_numpy_list.py
@@ -148,6 +148,42 @@ def test_numpy_list_grib_write_append():
     assert r_tmp.metadata("shortName") == ["msl", "2d"]
 
 
+def test_numpy_list_grib_concat():
+    ds = from_source("file", earthkit_examples_file("test.grib"))
+
+    assert ds[0].metadata("shortName") == "2t"
+
+    v = ds[0].values
+    v1 = v + 1
+    v2 = v + 2
+
+    md = ds[0].metadata()
+    md1 = md.override(shortName="msl")
+    md2 = md.override(shortName="2d")
+
+    r1 = FieldList.from_numpy(v1, md1)
+    r2 = FieldList.from_numpy(v2, md2)
+
+    r = r1 + r2
+    assert len(r) == 2
+    assert r.metadata("shortName") == ["msl", "2d"]
+
+    f1 = r.sel(shortName="msl")
+    assert len(f1) == 1
+    assert f1.metadata("shortName") == ["msl"]
+    assert np.allclose(f1.values, v1)
+
+    # save to disk
+    tmp = temp_file()
+    r.save(tmp.path)
+    assert os.path.exists(tmp.path)
+    r_tmp = from_source("file", tmp.path)
+    assert len(r_tmp) == 2
+    assert r_tmp.metadata("shortName") == ["msl", "2d"]
+    assert np.allclose(v1, r[0].values)
+    assert np.allclose(v2, r[1].values)
+
+
 if __name__ == "__main__":
     from earthkit.data.testing import main