Repository: ecmwf/earthkit-data (python)

Issue:
to_xarray should not crash when filter_by_keys backend_kwargs is used
### What happened?

We have a GRIB file containing model level data. (Note that since param `lnsp` is only defined on level 1 this data does not form a hypercube, but it is irrelevant for this issue).

![image](https://github.com/ecmwf/earthkit-data/assets/9033020/c7beba2a-05e4-41fb-9577-98d27cee3c51)

When we use `to_xarray()` with the`filter_by_keys` option:

```python
import earthkit.data
ds = earthkit.data.from_source("file", "omega_ml.grib")
r = ds.to_xarray(
    xarray_open_dataset_kwargs={'backend_kwargs': 
                               {'filter_by_keys': {'shortName': 't', 'typeOfLevel': 'hybrid', 'level': 1}}})
```

there is a crash:

```
---------------------------------------------------------------------------
AssertionError                            Traceback (most recent call last)
Cell In[8], line 1
----> 1 r = ds.to_xarray(
      2     xarray_open_dataset_kwargs={'backend_kwargs': 
      3                                {'filter_by_keys': {'shortName': 't', 'typeOfLevel': 'hybrid', 'level': 1}}})

File ~/git/earthkit-data/earthkit/data/readers/grib/xarray.py:168, in XarrayMixIn.to_xarray(self, **kwargs)
    162 two_d_fields = sum(number_of_gribs(result[v]) for v in result.data_vars)
    164 # Make sure all the fields are converted
    165 # There may be more 2D xarray fields than GRB fields
    166 # if some missing dimension are filled with NaN values
--> 168 assert two_d_fields >= len(self), (
    169     "Not all GRIB fields were converted to xarray"
    170     f" ({len(self)} GRIBs > {two_d_fields} 2D-field(s) in xarray)"
    171 )
    173 return result

AssertionError: Not all GRIB fields were converted to xarray (9 GRIBs > 1 2D-field(s) in xarray)

```

**Note**: we could easily perform the selection in earthkit-data then convert the results to xarray. E.g. the following code would work:

```python
import earthkit.data
ds = earthkit.data.from_source("file", "omega_ml.grib")
ds.sel(param="t", typeOfLevel="hybrid", level=1).to_xarray()
```


### What are the steps to reproduce the bug?

See above.

### Version

<=0.4.0

### Platform (OS and architecture)

all

### Relevant log output

_No response_

### Accompanying data

_No response_

### Organisation

ECMWF

Pull request description:
Fix crash in to_xarray when filter_by_keys is used
Fixes #250 

Patch:
diff --git a/earthkit/data/readers/grib/xarray.py b/earthkit/data/readers/grib/xarray.py
index f7ae1c6..a73794b 100644
--- a/earthkit/data/readers/grib/xarray.py
+++ b/earthkit/data/readers/grib/xarray.py
@@ -8,7 +8,6 @@
 #
 
 import logging
-import math
 import warnings
 
 from earthkit.data.utils.kwargs import Kwargs
@@ -141,33 +140,4 @@ class XarrayMixIn:
             **xarray_open_dataset_kwargs,
         )
 
-        def math_prod(lst):
-            if not hasattr(math, "prod"):
-                # python 3.7 does not have math.prod
-                n = 1
-                for x in lst:
-                    n = n * x
-                return n
-            return math.prod(lst)
-
-        def number_of_gribs(da):
-            # Assumes last two dimensions are lat/lon coordinates
-            skip = 2
-            if da.dims[-1] == "values":
-                # Assumes last dimension is the one-dimensional
-                # lat/lon coordinate (non-regular grid)
-                skip = 1
-            return math_prod(list(da.shape)[:-skip])
-
-        two_d_fields = sum(number_of_gribs(result[v]) for v in result.data_vars)
-
-        # Make sure all the fields are converted
-        # There may be more 2D xarray fields than GRB fields
-        # if some missing dimension are filled with NaN values
-
-        assert two_d_fields >= len(self), (
-            "Not all GRIB fields were converted to xarray"
-            f" ({len(self)} GRIBs > {two_d_fields} 2D-field(s) in xarray)"
-        )
-
         return result


Test patch:
diff --git a/tests/grib/test_grib_convert.py b/tests/grib/test_grib_convert.py
index 251f385..ee9a27c 100644
--- a/tests/grib/test_grib_convert.py
+++ b/tests/grib/test_grib_convert.py
@@ -33,6 +33,23 @@ def test_icon_to_xarray(mode):
     assert ds["pres"].sizes["values"] == 6
 
 
+@pytest.mark.parametrize("mode", ["file", "numpy_fs"])
+def test_to_xarray_filter_by_keys(mode):
+    g = load_file_or_numpy_fs("tuv_pl.grib", mode)
+    g = g.sel(param="t", level=500) + g.sel(param="u")
+    assert len(g) > 1
+
+    # see github #250
+    r = g.to_xarray(
+        xarray_open_dataset_kwargs={
+            "backend_kwargs": {"filter_by_keys": {"shortName": "t"}}
+        }
+    )
+
+    assert len(r.data_vars) == 1
+    assert r["t"].sizes["isobaricInhPa"] == 1
+
+
 @pytest.mark.parametrize("mode", ["file", "numpy_fs"])
 def test_grib_to_pandas(mode):
     f = load_file_or_numpy_fs("test_single.grib", mode, folder="data")