Repository: ecmwf/earthkit-data
Pull request #186 (2023-09-16 11:04:22)

Allow concatenation for empty fieldlist


Patch:
diff --git a/docs/examples/numpy_fieldlist.ipynb b/docs/examples/numpy_fieldlist.ipynb
index cbb4938..6074946 100644
--- a/docs/examples/numpy_fieldlist.ipynb
+++ b/docs/examples/numpy_fieldlist.ipynb
@@ -196,7 +196,7 @@
     {
      "data": {
       "text/plain": [
-       "<earthkit.data.readers.grib.metadata.GribMetadata at 0x10fd9df00>"
+       "<earthkit.data.readers.grib.metadata.GribMetadata at 0x1063b9510>"
       ]
      },
      "execution_count": 7,
@@ -913,6 +913,265 @@
     "ds1 = earthkit.data.from_source(\"file\", path)\n",
     "ds1.ls()"
    ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "3f4d3b46-e905-4ac6-88c2-bd514101d61b",
+   "metadata": {},
+   "source": [
+    "#### Performing the computations in a loop"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "0963540f-3a6d-4f0d-a046-3806c5f65e30",
+   "metadata": {},
+   "source": [
+    "In this example we create an **empty fieldlist** and add the results of the computations to it in a loop."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 21,
+   "id": "6c6e8651-f5a3-4534-9227-812e8e69407d",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "fs = ds.sel(shortName=\"t\", level=[850, 700, 500])\n",
+    "\n",
+    "# create an empty fieldlist\n",
+    "ds_r = FieldList()\n",
+    "\n",
+    "for f in fs:\n",
+    "    p = f.metadata(\"level\")*100. # hPa -> Pa\n",
+    "    t_new = potential_temperature(f.values, p)\n",
+    "    md_new = f.metadata().override(shortName=\"pt\")\n",
+    "    \n",
+    "    # create new numpy fieldlist with a single field\n",
+    "    ds_new = FieldList.from_numpy(t_new, md_new)\n",
+    "\n",
+    "    # add it to the resulting fieldlist\n",
+    "    ds_r += ds_new"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 22,
+   "id": "32c53494-e8f4-4a22-aa93-e16b5188483a",
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/html": [
+       "<div>\n",
+       "<style scoped>\n",
+       "    .dataframe tbody tr th:only-of-type {\n",
+       "        vertical-align: middle;\n",
+       "    }\n",
+       "\n",
+       "    .dataframe tbody tr th {\n",
+       "        vertical-align: top;\n",
+       "    }\n",
+       "\n",
+       "    .dataframe thead th {\n",
+       "        text-align: right;\n",
+       "    }\n",
+       "</style>\n",
+       "<table border=\"1\" class=\"dataframe\">\n",
+       "  <thead>\n",
+       "    <tr style=\"text-align: right;\">\n",
+       "      <th></th>\n",
+       "      <th>centre</th>\n",
+       "      <th>shortName</th>\n",
+       "      <th>typeOfLevel</th>\n",
+       "      <th>level</th>\n",
+       "      <th>dataDate</th>\n",
+       "      <th>dataTime</th>\n",
+       "      <th>stepRange</th>\n",
+       "      <th>dataType</th>\n",
+       "      <th>number</th>\n",
+       "      <th>gridType</th>\n",
+       "    </tr>\n",
+       "  </thead>\n",
+       "  <tbody>\n",
+       "    <tr>\n",
+       "      <th>0</th>\n",
+       "      <td>ecmf</td>\n",
+       "      <td>pt</td>\n",
+       "      <td>isobaricInhPa</td>\n",
+       "      <td>850</td>\n",
+       "      <td>20180801</td>\n",
+       "      <td>1200</td>\n",
+       "      <td>0</td>\n",
+       "      <td>an</td>\n",
+       "      <td>0</td>\n",
+       "      <td>regular_ll</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>1</th>\n",
+       "      <td>ecmf</td>\n",
+       "      <td>pt</td>\n",
+       "      <td>isobaricInhPa</td>\n",
+       "      <td>700</td>\n",
+       "      <td>20180801</td>\n",
+       "      <td>1200</td>\n",
+       "      <td>0</td>\n",
+       "      <td>an</td>\n",
+       "      <td>0</td>\n",
+       "      <td>regular_ll</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>2</th>\n",
+       "      <td>ecmf</td>\n",
+       "      <td>pt</td>\n",
+       "      <td>isobaricInhPa</td>\n",
+       "      <td>500</td>\n",
+       "      <td>20180801</td>\n",
+       "      <td>1200</td>\n",
+       "      <td>0</td>\n",
+       "      <td>an</td>\n",
+       "      <td>0</td>\n",
+       "      <td>regular_ll</td>\n",
+       "    </tr>\n",
+       "  </tbody>\n",
+       "</table>\n",
+       "</div>"
+      ],
+      "text/plain": [
+       "  centre shortName    typeOfLevel  level  dataDate  dataTime stepRange  \\\n",
+       "0   ecmf        pt  isobaricInhPa    850  20180801      1200         0   \n",
+       "1   ecmf        pt  isobaricInhPa    700  20180801      1200         0   \n",
+       "2   ecmf        pt  isobaricInhPa    500  20180801      1200         0   \n",
+       "\n",
+       "  dataType  number    gridType  \n",
+       "0       an       0  regular_ll  \n",
+       "1       an       0  regular_ll  \n",
+       "2       an       0  regular_ll  "
+      ]
+     },
+     "execution_count": 22,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "ds_r.ls()"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "b76eaf59-868c-4a59-8823-8d4bc53376f3",
+   "metadata": {},
+   "source": [
+    "We can save the NumpyFieldList into a GRIB file:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 23,
+   "id": "a149caa0-b1ba-4b8e-b2d0-6aad7ee224a6",
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/html": [
+       "<div>\n",
+       "<style scoped>\n",
+       "    .dataframe tbody tr th:only-of-type {\n",
+       "        vertical-align: middle;\n",
+       "    }\n",
+       "\n",
+       "    .dataframe tbody tr th {\n",
+       "        vertical-align: top;\n",
+       "    }\n",
+       "\n",
+       "    .dataframe thead th {\n",
+       "        text-align: right;\n",
+       "    }\n",
+       "</style>\n",
+       "<table border=\"1\" class=\"dataframe\">\n",
+       "  <thead>\n",
+       "    <tr style=\"text-align: right;\">\n",
+       "      <th></th>\n",
+       "      <th>centre</th>\n",
+       "      <th>shortName</th>\n",
+       "      <th>typeOfLevel</th>\n",
+       "      <th>level</th>\n",
+       "      <th>dataDate</th>\n",
+       "      <th>dataTime</th>\n",
+       "      <th>stepRange</th>\n",
+       "      <th>dataType</th>\n",
+       "      <th>number</th>\n",
+       "      <th>gridType</th>\n",
+       "    </tr>\n",
+       "  </thead>\n",
+       "  <tbody>\n",
+       "    <tr>\n",
+       "      <th>0</th>\n",
+       "      <td>ecmf</td>\n",
+       "      <td>pt</td>\n",
+       "      <td>isobaricInhPa</td>\n",
+       "      <td>850</td>\n",
+       "      <td>20180801</td>\n",
+       "      <td>1200</td>\n",
+       "      <td>0</td>\n",
+       "      <td>an</td>\n",
+       "      <td>0</td>\n",
+       "      <td>regular_ll</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>1</th>\n",
+       "      <td>ecmf</td>\n",
+       "      <td>pt</td>\n",
+       "      <td>isobaricInhPa</td>\n",
+       "      <td>700</td>\n",
+       "      <td>20180801</td>\n",
+       "      <td>1200</td>\n",
+       "      <td>0</td>\n",
+       "      <td>an</td>\n",
+       "      <td>0</td>\n",
+       "      <td>regular_ll</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>2</th>\n",
+       "      <td>ecmf</td>\n",
+       "      <td>pt</td>\n",
+       "      <td>isobaricInhPa</td>\n",
+       "      <td>500</td>\n",
+       "      <td>20180801</td>\n",
+       "      <td>1200</td>\n",
+       "      <td>0</td>\n",
+       "      <td>an</td>\n",
+       "      <td>0</td>\n",
+       "      <td>regular_ll</td>\n",
+       "    </tr>\n",
+       "  </tbody>\n",
+       "</table>\n",
+       "</div>"
+      ],
+      "text/plain": [
+       "  centre shortName    typeOfLevel  level  dataDate  dataTime stepRange  \\\n",
+       "0   ecmf        pt  isobaricInhPa    850  20180801      1200         0   \n",
+       "1   ecmf        pt  isobaricInhPa    700  20180801      1200         0   \n",
+       "2   ecmf        pt  isobaricInhPa    500  20180801      1200         0   \n",
+       "\n",
+       "  dataType  number    gridType  \n",
+       "0       an       0  regular_ll  \n",
+       "1       an       0  regular_ll  \n",
+       "2       an       0  regular_ll  "
+      ]
+     },
+     "execution_count": 23,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "path = \"_pt_from_loop.grib\"\n",
+    "ds_r.save(path)\n",
+    "ds1 = earthkit.data.from_source(\"file\", path)\n",
+    "ds1.ls()"
+   ]
   }
  ],
  "metadata": {
diff --git a/earthkit/data/core/fieldlist.py b/earthkit/data/core/fieldlist.py
index c91d67f..2fd9620 100644
--- a/earthkit/data/core/fieldlist.py
+++ b/earthkit/data/core/fieldlist.py
@@ -537,6 +537,15 @@ class FieldList(Index):
 
         return NumpyFieldList(array, metadata)
 
+    def ignore(self):
+        # When the concrete type is Fieldlist we assume the object was
+        # created with Fieldlist() i.e. it is empty. We ignore it from
+        # all the merge operations.
+        if type(self) is FieldList:
+            return True
+        else:
+            return False
+
     @cached_method
     def _default_index_keys(self):
         if len(self) > 0:
diff --git a/earthkit/data/sources/__init__.py b/earthkit/data/sources/__init__.py
index d67f57d..55f5982 100644
--- a/earthkit/data/sources/__init__.py
+++ b/earthkit/data/sources/__init__.py
@@ -44,6 +44,12 @@ class Source(Base):
         return self
 
     def ignore(self):
+        """Indicates to ignore this source in concatenation/merging.
+
+        Returns
+        -------
+        bool
+        """
         # Used by multi-source
         return False
 


Test patch:
diff --git a/tests/grib/test_grib_concat.py b/tests/grib/test_grib_concat.py
index 2111855..2f895e1 100644
--- a/tests/grib/test_grib_concat.py
+++ b/tests/grib/test_grib_concat.py
@@ -11,28 +11,108 @@
 
 import os
 
+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 test_grib_concat():
+def _check_save_to_disk(ds, len_ref, meta_ref):
+    # 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) == len_ref
+    assert r_tmp.metadata("shortName") == meta_ref
+    r_tmp = None
+
+
+@pytest.mark.parametrize("mode", ["oper", "multi"])
+def test_grib_concat(mode):
     ds1 = from_source("file", earthkit_examples_file("test.grib"))
     ds2 = from_source("file", earthkit_examples_file("test6.grib"))
-    ds = ds1 + ds2
+
+    if mode == "oper":
+        ds = ds1 + ds2
+    else:
+        ds = from_source("multi", ds1, ds2)
 
     assert len(ds) == 8
     md = ds1.metadata("param") + ds2.metadata("param")
     assert ds.metadata("param") == md
+    _check_save_to_disk(ds, 8, 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
+
+@pytest.mark.parametrize("mode", ["oper", "multi"])
+def test_grib_concat_3a(mode):
+    ds1 = from_source("file", earthkit_examples_file("test.grib"))
+    ds2 = from_source("file", earthkit_examples_file("test6.grib"))
+    ds3 = from_source("file", earthkit_examples_file("tuv_pl.grib"))
+    md = ds1.metadata("param") + ds2.metadata("param") + ds3.metadata("param")
+
+    if mode == "oper":
+        ds = ds1 + ds2
+        ds = ds + ds3
+    else:
+        ds = from_source("multi", ds1, ds2)
+        ds = from_source("multi", ds, ds3)
+
+    assert len(ds) == 26
+    assert ds.metadata("param") == md
+    _check_save_to_disk(ds, 26, md)
+
+
+@pytest.mark.parametrize("mode", ["oper", "multi"])
+def test_grib_concat_3b(mode):
+    ds1 = from_source("file", earthkit_examples_file("test.grib"))
+    ds2 = from_source("file", earthkit_examples_file("test6.grib"))
+    ds3 = from_source("file", earthkit_examples_file("tuv_pl.grib"))
+    md = ds1.metadata("param") + ds2.metadata("param") + ds3.metadata("param")
+
+    if mode == "oper":
+        ds = ds1 + ds2 + ds3
+    else:
+        ds = from_source("multi", ds1, ds2, ds3)
+
+    assert len(ds) == 26
+    assert ds.metadata("param") == md
+    _check_save_to_disk(ds, 26, md)
+
+
+def test_grib_from_empty_1():
+    ds_e = FieldList()
+    ds = from_source("file", earthkit_examples_file("test.grib"))
+    md = ds.metadata("param")
+
+    ds1 = ds_e + ds
+    assert id(ds1) == id(ds)
+    assert len(ds1) == 2
+    _check_save_to_disk(ds1, 2, md)
+
+
+def test_grib_from_empty_2():
+    ds_e = FieldList()
+    ds = from_source("file", earthkit_examples_file("test.grib"))
+    md = ds.metadata("param")
+
+    ds1 = ds + ds_e
+    assert id(ds1) == id(ds)
+    assert len(ds1) == 2
+    _check_save_to_disk(ds1, 2, md)
+
+
+def test_grib_from_empty_3():
+    ds_e = FieldList()
+    ds1 = from_source("file", earthkit_examples_file("test.grib"))
+    ds2 = from_source("file", earthkit_examples_file("test6.grib"))
+    md = ds1.metadata("param") + ds2.metadata("param")
+
+    ds3 = ds_e + ds1 + ds2
+    assert len(ds3) == 8
+    _check_save_to_disk(ds3, 8, md)
 
 
 if __name__ == "__main__":
diff --git a/tests/sources/test_numpy_list.py b/tests/sources/test_numpy_list.py
index 1fd60c5..033e742 100644
--- a/tests/sources/test_numpy_list.py
+++ b/tests/sources/test_numpy_list.py
@@ -13,6 +13,7 @@ import logging
 import os
 
 import numpy as np
+import pytest
 
 from earthkit.data import from_source
 from earthkit.data.core.fieldlist import FieldList
@@ -22,6 +23,38 @@ from earthkit.data.testing import earthkit_examples_file
 LOG = logging.getLogger(__name__)
 
 
+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 _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)
+
+
 def test_numpy_list_grib_single_field():
     ds = from_source("file", earthkit_examples_file("test.grib"))
 
@@ -91,10 +124,9 @@ def test_numpy_list_grib_write_missing():
     assert ds[0].metadata("shortName") == "2t"
 
     v = ds[0].values
-    v1 = np.array(v)
+    v1 = v + 1
     assert not np.isnan(v1[0])
     assert not np.isnan(v1[1])
-
     v1[0] = np.nan
     assert np.isnan(v1[0])
     assert not np.isnan(v1[1])
@@ -111,9 +143,9 @@ def test_numpy_list_grib_write_missing():
     r.save(tmp.path)
     assert os.path.exists(tmp.path)
     r_tmp = from_source("file", tmp.path)
-
-    assert np.isnan(r_tmp[0].values[0])
-    assert not np.isnan(r_tmp[0].values[1])
+    v_tmp = r_tmp[0].values
+    assert np.isnan(v_tmp[0])
+    assert not np.isnan(v_tmp[1])
 
 
 def test_numpy_list_grib_write_append():
@@ -141,6 +173,7 @@ def test_numpy_list_grib_write_append():
     assert r_tmp.metadata("shortName") == ["msl"]
     r_tmp = None
 
+    # append
     r2.save(tmp.path, append=True)
     assert os.path.exists(tmp.path)
     r_tmp = from_source("file", tmp.path)
@@ -148,40 +181,117 @@ 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"))
+@pytest.mark.parametrize("mode", ["oper", "multi"])
+def test_numpy_list_grib_concat_2a(mode):
+    ds1, ds2, md = _prepare_ds(2)
 
-    assert ds[0].metadata("shortName") == "2t"
+    if mode == "oper":
+        ds = ds1 + ds2
+    else:
+        ds = from_source("multi", ds1, ds2)
 
-    v = ds[0].values
-    v1 = v + 1
-    v2 = v + 2
+    assert len(ds) == 8
+    assert ds.metadata("param") == md
 
-    md = ds[0].metadata()
-    md1 = md.override(shortName="msl")
-    md2 = md.override(shortName="2d")
+    f1 = ds.sel(shortName="msl")
+    assert len(f1) == 1
+    assert f1.metadata("shortName") == ["msl"]
+    assert np.allclose(f1[0].values, ds1[1].values)
 
-    r1 = FieldList.from_numpy(v1, md1)
-    r2 = FieldList.from_numpy(v2, md2)
+    _check_save_to_disk(ds, 8, md)
 
-    r = r1 + r2
-    assert len(r) == 2
-    assert r.metadata("shortName") == ["msl", "2d"]
 
-    f1 = r.sel(shortName="msl")
+def test_numpy_list_grib_concat_2b():
+    ds1, ds2, md = _prepare_ds(2)
+    ds1 += ds2
+
+    assert len(ds1) == 8
+    assert ds1.metadata("param") == md
+
+    f1 = ds1.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)
+    assert np.allclose(f1[0].values, ds1[1].values)
+
+    _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)
+
+    if mode == "oper":
+        ds = ds1 + ds2
+        ds = ds + ds3
+    else:
+        ds = from_source("multi", ds1, ds2)
+        ds = from_source("multi", ds, ds3)
+
+    assert len(ds) == 26
+    assert ds.metadata("param") == 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)
+
+    if mode == "oper":
+        ds = ds1 + ds2 + ds3
+    else:
+        ds = from_source("multi", ds1, ds2, ds3)
+
+    assert len(ds) == 26
+    assert ds.metadata("param") == md
+    _check_save_to_disk(ds, 26, md)
+
+
+def test_numpy_list_grib_from_empty_1():
+    ds_e = FieldList()
+    ds, md = _prepare_ds(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)
+
+
+def test_numpy_list_grib_from_empty_2():
+    ds_e = FieldList()
+    ds, md = _prepare_ds(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)
+
+
+def test_numpy_list_grib_from_empty_3():
+    ds_e = FieldList()
+    ds1, ds2, md = _prepare_ds(2)
+    ds = ds_e + ds1 + ds2
+    assert len(ds) == 8
+    assert ds.metadata("param") == md
+    _check_save_to_disk(ds, 8, md)
+
+
+def test_numpy_list_grib_from_empty_4():
+    ds = FieldList()
+    ds1, md = _prepare_ds(1)
+    ds += ds1
+    assert id(ds) == id(ds1)
+    assert len(ds) == 2
+    assert ds.metadata("param") == md
+    _check_save_to_disk(ds, 2, md)
+
+
+def test_numpy_list_grib_from_empty_5():
+    ds = FieldList()
+    ds1, ds2, md = _prepare_ds(2)
+    ds += ds1 + ds2
+    assert len(ds) == 8
+    assert ds.metadata("param") == md
+    _check_save_to_disk(ds, 8, md)
 
 
 if __name__ == "__main__":