Repository: ecmwf/earthkit-data

Issue:
Add flatten argument to to_numpy in xarray wrappers
### Is your feature request related to a problem? Please describe.

Many earthkit-data objects implement a `to_numpy()` method that accepts a `flatten` keyword argument. In other earthkit libraries, like earthkit-plots, we sometimes pass the `flatten` kwarg when calling `to_numpy`, depending on the use case.

However, the xarray Dataset and DataArray wrappers within earthkit-data do not support a `flatten` argument. This inconsistency can lead to errors when flatten is passed to their `to_numpy` methods. (See ecmwf/earthkit-plots#67).

### Describe the solution you'd like

- Update the xarray Dataset and DataArray wrappers in earthkit-data to accept a `flatten` keyword argument in their `to_numpy()` methods.
- If `flatten=True`, reshape the resulting NumPy array to be 1-dimensional after converting with `to_numpy()`.
- If `flatten=False` (default), maintain the current behavior and return the regular array.

This will ensure that all `to_numpy()` methods across earthkit objects have a consistent API and behavior.

### Describe alternatives you've considered

One alternative considered was handling this issue directly within earthkit-plots by wrapping calls to `to_numpy()` in a try-except block. If an error occurred due to an unexpected `flatten` argument, earthkit-plots could retry the call without the flatten parameter.

However, this approach felt suboptimal for several reasons:
- It would hide inconsistencies between earthkit-data object types rather than addressing the root cause.
- It would make the code more fragile and harder to maintain, as future changes to `to_numpy()` behavior across different objects could introduce subtle bugs.
- It would introduce unnecessary complexity into earthkit-plots, which should ideally be agnostic to these internal differences.

### Additional context

ecmwf/earthkit-plots#67

### Organisation

ECMWF

Discussion:
@sandorkertesz: Fixed by #685 

Solution patch:
diff --git a/src/earthkit/data/wrappers/xarray.py b/src/earthkit/data/wrappers/xarray.py
index bdbfb2b..fd79680 100644
--- a/src/earthkit/data/wrappers/xarray.py
+++ b/src/earthkit/data/wrappers/xarray.py
@@ -58,14 +58,17 @@ class XArrayDataArrayWrapper(Wrapper):
         """
         return self.data
 
-    def to_numpy(self):
+    def to_numpy(self, flatten=False):
         """Return a numpy `ndarray` representation of the data.
 
         Returns
         -------
         numpy.ndarray
         """
-        return self.data.to_numpy()
+        arr = self.data.to_numpy()
+        if flatten:
+            arr = arr.flatten()
+        return arr
 
     def to_pandas(self, *args, **kwargs):
         """Return a pandas `dataframe` representation of the data.
@@ -91,14 +94,17 @@ class XArrayDatasetWrapper(XArrayDataArrayWrapper):
     methods.
     """
 
-    def to_numpy(self):
+    def to_numpy(self, flatten=False):
         """Return a numpy `ndarray` representation of the data.
 
         Returns
         -------
         numpy.ndarray
         """
-        return self.data.to_array().to_numpy()
+        arr = self.data.to_array().to_numpy()
+        if flatten:
+            arr = arr.flatten()
+        return arr
 
     # def component(self, component):
     #     """


Test patch:
diff --git a/tests/wrappers/test_xarray.py b/tests/wrappers/test_xarray.py
index d3e43b5..d7466cd 100644
--- a/tests/wrappers/test_xarray.py
+++ b/tests/wrappers/test_xarray.py
@@ -49,3 +49,21 @@ def test_xarray_lazy_fieldlist_scan():
     assert ds._fields is None
     assert len(ds) == 2
     assert len(ds._fields) == 2
+
+
+@pytest.mark.no_eccodes
+def test_xarray_to_numpy():
+    import xarray as xr
+
+    data_array = xr.DataArray(
+        [[1, 2, 3], [4, 5, 6]],
+        dims=["x", "y"],
+        coords={"x": [1, 2], "y": [3, 4, 5]},
+    )
+    ds = from_object(data_array)
+
+    arr_2d = ds.to_numpy()
+    assert arr_2d.shape == (2, 3)
+
+    arr_1d = ds.to_numpy(flatten=True)
+    assert arr_1d.shape == (6,)