Repository: ecmwf/earthkit-data (python)

Issue:
Add geographical coordinate access to FieldList
On a GribField we can extract the latitudes/longitudes with [to_latlon](https://earthkit-data.readthedocs.io/en/latest/_api/data/readers/grib/codes/index.html). It returns a dict as:

`{"lat": array, "lon": array}`

This method is not available on a FieldList. The proposal is that on a FieldList `to_latlon` could return the `to_latlon` value from the first field if the grid geometry is the same for all the fields. Otherwise it would throw an exception.

The situation is exactly the same for [to_points](https://earthkit-data.readthedocs.io/en/latest/_api/data/readers/grib/codes/index.html), so it should be implemented in a similar way on a FieldList.

Relates to #116 


Pull request:
Implement to_latlon and to_points methods on fieldlists
fixes #114 

Patch:
diff --git a/docs/examples/grib_overview.ipynb b/docs/examples/grib_overview.ipynb
index 1f7045f..979e4ff 100644
--- a/docs/examples/grib_overview.ipynb
+++ b/docs/examples/grib_overview.ipynb
@@ -1893,7 +1893,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.8.16"
+   "version": "3.8.17"
   }
  },
  "nbformat": 4,
diff --git a/earthkit/data/core/__init__.py b/earthkit/data/core/__init__.py
index d9d215c..5717c9a 100644
--- a/earthkit/data/core/__init__.py
+++ b/earthkit/data/core/__init__.py
@@ -106,6 +106,14 @@ class Base(metaclass=MetaBase):
         """Reorder the elements of the object."""
         self._not_implemented()
 
+    @abstractmethod
+    def to_points(self, *args, **kwargs):
+        self._not_implemented()
+
+    @abstractmethod
+    def to_latlon(self, *args, **kwargs):
+        self._not_implemented()
+
     def __add__(self, other):
         self._not_implemented()
 
diff --git a/earthkit/data/readers/grib/fieldlist.py b/earthkit/data/readers/grib/fieldlist.py
index 58c3d31..96c85da 100644
--- a/earthkit/data/readers/grib/fieldlist.py
+++ b/earthkit/data/readers/grib/fieldlist.py
@@ -13,6 +13,7 @@ from collections import defaultdict
 
 from earthkit.data.core.caching import auxiliary_cache_file
 from earthkit.data.core.index import ScaledIndex
+from earthkit.data.decorators import cached_method
 
 from .pandas import PandasMixIn
 from .xarray import XarrayMixIn
@@ -414,6 +415,94 @@ class FieldListMixin(PandasMixIn, XarrayMixIn):
             valid.add(d["valid_time"])
         return {"base_time": sorted(base), "valid_time": sorted(valid)}
 
+    @cached_method
+    def _is_shared_grid(self):
+        if len(self) > 0:
+            for i, f in enumerate(self):
+                if i == 0:
+                    grid = f.metadata("md5GridSection")
+                elif f.metadata("md5GridSection") != grid:
+                    return False
+            return True
+        return False
+
+    def to_points(self, **kwargs):
+        r"""Returns the geographical coordinates shared by all the fields in
+        the data's original Coordinate Reference System (CRS).
+
+        Parameters
+        ----------
+        **kwargs: dict, optional
+            Keyword arguments passed to
+            :obj:`GribField.to_points() <data.readers.grib.codes.GribField.to_points>`
+
+        Returns
+        -------
+        dict
+            Dictionary with items "x" and "y", containing the ndarrays of the x and
+            y coordinates, respectively.
+
+        Raises
+        ------
+        ValueError
+            When not all the fields have the same grid geometry.
+        """
+        if self._is_shared_grid():
+            return self[0].to_points(**kwargs)
+        elif len(self) == 0:
+            return dict(x=None, y=None)
+        else:
+            raise ValueError("Fields do not have the same grid geometry")
+
+    def to_latlon(self, **kwargs):
+        r"""Returns the latitudes/longitudes shared by all the fields.
+
+        Parameters
+        ----------
+        **kwargs: dict, optional
+            Keyword arguments passed to
+            :obj:`GribField.to_latlon() <data.readers.grib.codes.GribField.to_latlon>`
+
+        Returns
+        -------
+        dict
+            Dictionary with items "lat" and "lon", containing the ndarrays of the latitudes and
+            longitudes, respectively.
+
+        Raises
+        ------
+        ValueError
+            When not all the fields have the same grid geometry
+
+        Examples
+        --------
+        >>> import earthkit.data
+        >>> ds = earthkit.data.from_source("file", "docs/examples/test.grib")
+        >>> for f in ds:
+        ...     print(f.shape)
+        ...
+        (11, 19)
+        (11, 19)
+        >>> r = ds.to_latlon()
+        >>> for k, v in r.items():
+        ...     print(f"{k}: shape={v.shape}")
+        ...
+        lat: shape=(11, 19)
+        lon: shape=(11, 19)
+        >>> r["lon"][:2]
+        array([[-27., -23., -19., -15., -11.,  -7.,  -3.,   1.,   5.,   9.,  13.,
+         17.,  21.,  25.,  29.,  33.,  37.,  41.,  45.],
+        [-27., -23., -19., -15., -11.,  -7.,  -3.,   1.,   5.,   9.,  13.,
+         17.,  21.,  25.,  29.,  33.,  37.,  41.,  45.]])
+
+        """
+        if self._is_shared_grid():
+            return self[0].to_latlon(**kwargs)
+        elif len(self) == 0:
+            return dict(lat=None, lon=None)
+        else:
+            raise ValueError("Fields do not have the same grid geometry")
+
     def bounding_box(self):
         r"""Returns the bounding box for each field.
 


Test patch:
diff --git a/tests/grib/test_grib_contents.py b/tests/grib/test_grib_contents.py
index e09693c..5309b22 100644
--- a/tests/grib/test_grib_contents.py
+++ b/tests/grib/test_grib_contents.py
@@ -603,11 +603,13 @@ def test_grib_values_with_missing():
     assert np.count_nonzero(np.isnan(m)) == 38
 
 
-def test_grib_to_latlon_single():
+@pytest.mark.parametrize("index", [0, None])
+def test_grib_to_latlon_single(index):
     f = from_source("file", earthkit_test_data_file("test_single.grib"))
 
     eps = 1e-5
-    v = f[0].to_latlon(flatten=True)
+    g = f[index] if index is not None else f
+    v = g.to_latlon(flatten=True)
     assert isinstance(v, dict)
     assert isinstance(v["lon"], np.ndarray)
     assert isinstance(v["lat"], np.ndarray)
@@ -629,10 +631,12 @@ def test_grib_to_latlon_single():
     )
 
 
-def test_grib_to_latlon_single_shape():
+@pytest.mark.parametrize("index", [0, None])
+def test_grib_to_latlon_single_shape(index):
     f = from_source("file", earthkit_test_data_file("test_single.grib"))
 
-    v = f[0].to_latlon()
+    g = f[index] if index is not None else f
+    v = g.to_latlon()
     assert isinstance(v, dict)
     assert isinstance(v["lon"], np.ndarray)
     assert isinstance(v["lat"], np.ndarray)
@@ -648,11 +652,35 @@ def test_grib_to_latlon_single_shape():
         assert np.allclose(y, np.ones(12) * (90 - i * 30))
 
 
-def test_grib_to_points_single():
+def test_grib_to_latlon_multi():
+    f = from_source("file", earthkit_examples_file("test.grib"))
+
+    v_ref = f[0].to_latlon(flatten=True)
+    v = f.to_latlon(flatten=True)
+    assert isinstance(v, dict)
+    assert v.keys() == v_ref.keys()
+
+    assert isinstance(v, dict)
+    assert np.allclose(v["lat"], v_ref["lat"])
+    assert np.allclose(v["lon"], v_ref["lon"])
+
+
+def test_grib_to_latlon_multi_non_shared_grid():
+    f1 = from_source("file", earthkit_examples_file("test.grib"))
+    f2 = from_source("file", earthkit_examples_file("test4.grib"))
+    f = f1 + f2
+
+    with pytest.raises(ValueError):
+        f.to_latlon()
+
+
+@pytest.mark.parametrize("index", [0, None])
+def test_grib_to_points_single(index):
     f = from_source("file", earthkit_test_data_file("test_single.grib"))
 
     eps = 1e-5
-    v = f[0].to_points(flatten=True)
+    g = f[index] if index is not None else f
+    v = g.to_points(flatten=True)
     assert isinstance(v, dict)
     assert isinstance(v["x"], np.ndarray)
     assert isinstance(v["y"], np.ndarray)
@@ -680,6 +708,28 @@ def test_grib_to_points_unsupported_grid():
         f[0].to_points()
 
 
+def test_grib_to_points_multi():
+    f = from_source("file", earthkit_examples_file("test.grib"))
+
+    v_ref = f[0].to_points(flatten=True)
+    v = f.to_points(flatten=True)
+    assert isinstance(v, dict)
+    assert v.keys() == v_ref.keys()
+
+    assert isinstance(v, dict)
+    assert np.allclose(v["x"], v_ref["x"])
+    assert np.allclose(v["y"], v_ref["y"])
+
+
+def test_grib_to_points_multi_non_shared_grid():
+    f1 = from_source("file", earthkit_examples_file("test.grib"))
+    f2 = from_source("file", earthkit_examples_file("test4.grib"))
+    f = f1 + f2
+
+    with pytest.raises(ValueError):
+        f.to_points()
+
+
 def test_grib_datetime():
     s = from_source("file", earthkit_examples_file("test.grib"))