Repository: ecmwf/earthkit-data
Pull request #153 (2023-08-14 07:51:18)

Add data method to fieldlist
The `data()` method only existed on a field. This PR implements it on a fieldlist too.

Example:
<img width="771" alt="image" src="https://github.com/ecmwf/earthkit-data/assets/9033020/bf6ac6be-5923-4863-912d-1052c39c4059">



Patch:
diff --git a/docs/examples/grib_lat_lon_value.ipynb b/docs/examples/grib_lat_lon_value.ipynb
index c0a9782..d54c924 100644
--- a/docs/examples/grib_lat_lon_value.ipynb
+++ b/docs/examples/grib_lat_lon_value.ipynb
@@ -198,7 +198,23 @@
    "id": "c7ec94f3-4e68-4380-a88d-128b2fb46cd8",
    "metadata": {},
    "source": [
-    "This *data()* method is only available on fields. It returns the latitude, longitude and values arrays from a field as a tuple. By default the field's shape is kept."
+    "The simplest way to access the latitudes, longitudes and values is to use the *data()* method. It works both on fields and fieldlists."
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "c8c67310-78c7-4dea-be92-2bf51860b2a4",
+   "metadata": {},
+   "source": [
+    "#### Fields"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "443922d8-a8b4-45b4-8139-2c5fd27370e0",
+   "metadata": {},
+   "source": [
+    "*data()* returns the latitude, longitude and values arrays from a field as a tuple. By default the field's shape is kept."
    ]
   },
   {
@@ -302,6 +318,104 @@
     "[x.shape for x in llv]"
    ]
   },
+  {
+   "cell_type": "markdown",
+   "id": "c754b09b-50fe-4b06-8ba4-2cef8aed1cd3",
+   "metadata": {},
+   "source": [
+    "#### FieldLists"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "a21f46b4-6c79-4306-a005-390ef2a02173",
+   "metadata": {},
+   "source": [
+    "*data()* only works on a fieldlist if all the fields has the same grid. The first two elements of the resulting tuple are the latitude and longitude arrays (shared between fields), while the third element is the array of the field value arrays (the same as returned by *to_numpy()*):"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 6,
+   "id": "03c4a09d-ec99-442c-91e3-292d0541c1dd",
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "[(7, 12), (7, 12), (6, 7, 12)]"
+      ]
+     },
+     "execution_count": 6,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "llv = ds.data()\n",
+    "[x.shape for x in llv]"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "e4e9c3d3-6905-4aa9-9857-eb2a58406ac2",
+   "metadata": {},
+   "source": [
+    "E.g. the first value from each field can be extracted as:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 7,
+   "id": "9c461ac5-5eea-460d-af41-c93af51e6f45",
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "array([296.56417847, 288.53916931, 279.26531982, 259.84306335,\n",
+       "       248.00323486, 230.65315247])"
+      ]
+     },
+     "execution_count": 7,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "llv[2][:,1,1]"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "5b0cd71f-86da-4aab-b456-b50b382a6650",
+   "metadata": {},
+   "source": [
+    "When using the *flatten* keyword 1D arrays are returned for latitude and longitude, and the values array will be flattened per field:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 8,
+   "id": "dc0b8aa4-088b-4573-9688-46a0f0a3fdb1",
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "[(84,), (84,), (6, 84)]"
+      ]
+     },
+     "execution_count": 8,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "llv = ds.data(flatten=True)\n",
+    "[x.shape for x in llv]"
+   ]
+  },
   {
    "cell_type": "markdown",
    "id": "04812891-8465-458b-b3f2-9f818de4939d",
@@ -315,7 +429,7 @@
    "id": "fc4f152e-460e-410c-bbbe-d14f6615325a",
    "metadata": {},
    "source": [
-    "These methods are available both for fields and fieldlists."
+    "We can also get the latitudes, longitudes and values by using the *to_latlon()* and *to_numpy()* methods."
    ]
   },
   {
@@ -336,7 +450,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 6,
+   "execution_count": 9,
    "id": "b9af7f2c-f120-4437-b59f-b804ce15e84b",
    "metadata": {},
    "outputs": [
@@ -365,7 +479,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 7,
+   "execution_count": 10,
    "id": "c0926aa8-7fa4-4bb6-adba-11711bc7e463",
    "metadata": {},
    "outputs": [
@@ -375,7 +489,7 @@
        "(7, 12)"
       ]
      },
-     "execution_count": 7,
+     "execution_count": 10,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -395,7 +509,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 8,
+   "execution_count": 11,
    "id": "abfb402a-743e-44e5-a01b-bcc98f0e6ed4",
    "metadata": {},
    "outputs": [
@@ -435,7 +549,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 9,
+   "execution_count": 12,
    "id": "dda20d91-1f1b-4ec5-991e-e50b841312ea",
    "metadata": {},
    "outputs": [
@@ -464,7 +578,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 10,
+   "execution_count": 13,
    "id": "e21d501a-22ef-4e82-a996-261baa2594de",
    "metadata": {},
    "outputs": [
@@ -474,7 +588,7 @@
        "(6, 7, 12)"
       ]
      },
-     "execution_count": 10,
+     "execution_count": 13,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -494,7 +608,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 11,
+   "execution_count": 14,
    "id": "1edbf41c-437f-41b4-aadb-c840249519b8",
    "metadata": {},
    "outputs": [
@@ -505,7 +619,7 @@
        "       248.00323486, 230.65315247])"
       ]
      },
-     "execution_count": 11,
+     "execution_count": 14,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -532,7 +646,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 12,
+   "execution_count": 15,
    "id": "bf898af0-7382-40e8-89fd-8a07bc250378",
    "metadata": {},
    "outputs": [
@@ -543,7 +657,7 @@
        "      dtype=float32)"
       ]
      },
-     "execution_count": 12,
+     "execution_count": 15,
      "metadata": {},
      "output_type": "execute_result"
     }
diff --git a/docs/examples/grib_metadata.ipynb b/docs/examples/grib_metadata.ipynb
index 1a60310..0fd912c 100644
--- a/docs/examples/grib_metadata.ipynb
+++ b/docs/examples/grib_metadata.ipynb
@@ -64,7 +64,7 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "head() displays the first 5 messages by default:"
+    "*head()* displays the first 5 messages by default:"
    ]
   },
   {
@@ -826,7 +826,7 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "tail() displays the last 5 messages by default:"
+    "*tail()* displays the last 5 messages by default:"
    ]
   },
   {
@@ -973,7 +973,7 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "By default all the fields are listed:"
+    "By default *ls()* lists all the fields:"
    ]
   },
   {
@@ -1600,7 +1600,7 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "metadata() with positional arguments works for both single and multiple fields. key type qualifiers are allowed to be used."
+    "*metadata()* with positional arguments works for both single and multiple fields. Key type qualifiers are allowed to be used."
    ]
   },
   {
@@ -1717,7 +1717,7 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "dump() gives a tabbed interface to inspect the various namespaces:"
+    "*dump()* gives a tabbed interface to inspect the various namespaces:"
    ]
   },
   {
diff --git a/earthkit/data/readers/grib/fieldlist.py b/earthkit/data/readers/grib/fieldlist.py
index a93924b..f388a37 100644
--- a/earthkit/data/readers/grib/fieldlist.py
+++ b/earthkit/data/readers/grib/fieldlist.py
@@ -426,6 +426,77 @@ class FieldListMixin(PandasMixIn, XarrayMixIn):
             return True
         return False
 
+    def data(self, keys=("lat", "lon", "value"), flatten=False):
+        r"""Return the values and/or the geographical coordinates.
+
+        Only works when all the fields have the same grid geometry.
+
+        Parameters
+        ----------
+        keys: :obj:`str`, :obj:`list` or :obj:`tuple`
+            Specifies the type of data to be returned. Any combination of "lat", "lon" and "value"
+            is allowed here.
+        flatten: bool
+            When it is True a flat ndarray for "lat" and "lon" is returned, while the
+            resulting ndarray for "value" will be flattened per field. Otherwise an ndarray
+            with one field's :obj:`shape` is returned for "lat" and "lon", while for
+            "value" each array per field will keep the field's :obj:`shape`.
+
+        Returns
+        -------
+        ndarray or tuple of ndarrays
+            When ``keys`` is a single value an ndarray is returned. Otherwise a tuple containing one ndarray
+            per key is returned (following the order in ``keys``). The ndarray for a
+            given key is constructed as follows:
+
+            * "lat": the latitudes array from the first field returned
+            * "lon": the longitudes array from the first field returned
+            * "values": the array of the field values array returned
+              (same as with :obj:`to_numpy`)
+
+        Raises
+        ------
+        ValueError
+            When not all the fields have the same grid geometry.
+
+        Examples
+        --------
+        - :ref:`/examples/grib_lat_lon_value.ipynb`
+
+
+        See Also
+        --------
+        to_latlon
+        to_points
+        to_numpy
+        values
+
+        """
+        if self._is_shared_grid():
+            if isinstance(keys, str):
+                keys = [keys]
+
+            if "lat" in keys or "lon" in keys:
+                latlon = self[0].to_latlon(flatten=flatten)
+
+            r = []
+            for k in keys:
+                if k == "lat":
+                    r.append(latlon["lat"])
+                elif k == "lon":
+                    r.append(latlon["lon"])
+                elif k == "value":
+                    r.append(self.to_numpy(flatten=flatten))
+                else:
+                    raise ValueError(f"data: invalid argument: {k}")
+
+            return r[0] if len(r) == 1 else tuple(r)
+
+        elif len(self) == 0:
+            return tuple()
+        else:
+            raise ValueError("Fields do not have the same grid geometry")
+
     def to_points(self, **kwargs):
         r"""Returns the geographical coordinates shared by all the fields in
         the data's original Coordinate Reference System (CRS).


Test patch:
diff --git a/tests/grib/test_grib_contents.py b/tests/grib/test_grib_contents.py
index 2cb6ba0..4bf4e4e 100644
--- a/tests/grib/test_grib_contents.py
+++ b/tests/grib/test_grib_contents.py
@@ -651,6 +651,87 @@ def test_grib_values_with_missing():
     assert np.count_nonzero(np.isnan(m)) == 38
 
 
+@pytest.mark.parametrize(
+    "kwarg,expected_shape",
+    [({}, (11, 19)), ({"flatten": True}, (209,)), ({"flatten": False}, (11, 19))],
+)
+def test_grib_field_data(kwarg, expected_shape):
+    ds = from_source("file", earthkit_examples_file("test.grib"))
+
+    latlon = ds[0].to_latlon(**kwarg)
+    v = ds[0].to_numpy(**kwarg)
+
+    d = ds[0].data(**kwarg)
+    assert isinstance(d, tuple)
+    assert len(d) == 3
+    assert d[0].shape == expected_shape
+    assert np.allclose(d[0], latlon["lat"])
+    assert np.allclose(d[1], latlon["lon"])
+    assert np.allclose(d[2], v)
+
+    d = ds[0].data(keys="lat", **kwarg)
+    assert np.allclose(d, latlon["lat"])
+
+    d = ds[0].data(keys="lon", **kwarg)
+    assert np.allclose(d, latlon["lon"])
+
+    d = ds[0].data(keys="value", **kwarg)
+    assert np.allclose(d, v)
+
+    d = ds[0].data(keys=("value", "lon"), **kwarg)
+    assert isinstance(d, tuple)
+    assert len(d) == 2
+    assert np.allclose(d[0], v)
+    assert np.allclose(d[1], latlon["lon"])
+
+
+@pytest.mark.parametrize(
+    "kwarg,expected_latlon_shape,expected_values_shape",
+    [
+        ({}, (11, 19), (2, 11, 19)),
+        (
+            {"flatten": True},
+            (209,),
+            (
+                2,
+                209,
+            ),
+        ),
+        ({"flatten": False}, (11, 19), (2, 11, 19)),
+    ],
+)
+def test_grib_fieldlist_data(kwarg, expected_latlon_shape, expected_values_shape):
+    ds = from_source("file", earthkit_examples_file("test.grib"))
+
+    latlon = ds.to_latlon(**kwarg)
+    v = ds.to_numpy(**kwarg)
+
+    d = ds.data(**kwarg)
+    assert isinstance(d, tuple)
+    assert len(d) == 3
+    assert d[0].shape == expected_latlon_shape
+    assert d[1].shape == expected_latlon_shape
+    assert d[2].shape == expected_values_shape
+    assert np.allclose(d[0], latlon["lat"])
+    assert np.allclose(d[1], latlon["lon"])
+    assert np.allclose(d[2], v)
+
+    d = ds.data(keys="lat", **kwarg)
+    assert np.allclose(d, latlon["lat"])
+
+    d = ds.data(keys="lon", **kwarg)
+    assert np.allclose(d, latlon["lon"])
+
+    d = ds.data(keys="value", **kwarg)
+    assert np.allclose(d, v)
+
+    d = ds.data(keys=("value", "lon"), **kwarg)
+    assert isinstance(d, tuple)
+    assert len(d) == 2
+    assert np.allclose(d[0], v)
+    assert np.allclose(d[1], latlon["lon"])
+
+
 @pytest.mark.parametrize("index", [0, None])
 def test_grib_to_latlon_single(index):
     f = from_source("file", earthkit_test_data_file("test_single.grib"))