Repository: ecmwf/earthkit-data
Pull request #154 (2023-08-14 12:10:27)

Change return type in data method to ndarray
The `data()` method on a Field or FieldList previously returned a tuple. This PR changes the return type to ndarray.


Patch:
diff --git a/docs/examples/grib_lat_lon_value.ipynb b/docs/examples/grib_lat_lon_value.ipynb
index d54c924..560b573 100644
--- a/docs/examples/grib_lat_lon_value.ipynb
+++ b/docs/examples/grib_lat_lon_value.ipynb
@@ -214,7 +214,7 @@
    "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."
+    "*data()* returns the latitude, longitude and values arrays from a field as an ndarray. By default the field's shape is kept."
    ]
   },
   {
@@ -226,7 +226,7 @@
     {
      "data": {
       "text/plain": [
-       "[(7, 12), (7, 12), (7, 12)]"
+       "(3, 7, 12)"
       ]
      },
      "execution_count": 3,
@@ -236,15 +236,15 @@
    ],
    "source": [
     "llv = ds[0].data()\n",
-    "[x.shape for x in llv]"
+    "llv.shape"
    ]
   },
   {
    "cell_type": "markdown",
-   "id": "0bd409be-44da-4b28-8f65-6c9ae1f1edba",
+   "id": "a95f88f6-4136-472c-9e45-f4c744f97f5f",
    "metadata": {},
    "source": [
-    "We can get e.g. the field values as:"
+    "Here *llv[0]* contains the latitudes, *llv[1]* the longitudes, while *llv[2]* the values. We can get e.g. the first field value as:"
    ]
   },
   {
@@ -256,27 +256,7 @@
     {
      "data": {
       "text/plain": [
-       "array([[272.56417847, 272.56417847, 272.56417847, 272.56417847,\n",
-       "        272.56417847, 272.56417847, 272.56417847, 272.56417847,\n",
-       "        272.56417847, 272.56417847, 272.56417847, 272.56417847],\n",
-       "       [288.56417847, 296.56417847, 288.56417847, 288.56417847,\n",
-       "        288.56417847, 288.56417847, 280.56417847, 288.56417847,\n",
-       "        296.56417847, 288.56417847, 280.56417847, 280.56417847],\n",
-       "       [312.56417847, 304.56417847, 320.56417847, 296.56417847,\n",
-       "        304.56417847, 296.56417847, 296.56417847, 296.56417847,\n",
-       "        288.56417847, 296.56417847, 296.56417847, 296.56417847],\n",
-       "       [296.56417847, 304.56417847, 296.56417847, 304.56417847,\n",
-       "        296.56417847, 296.56417847, 296.56417847, 296.56417847,\n",
-       "        296.56417847, 296.56417847, 296.56417847, 296.56417847],\n",
-       "       [288.56417847, 296.56417847, 288.56417847, 288.56417847,\n",
-       "        280.56417847, 288.56417847, 288.56417847, 288.56417847,\n",
-       "        288.56417847, 288.56417847, 280.56417847, 288.56417847],\n",
-       "       [272.56417847, 264.56417847, 264.56417847, 272.56417847,\n",
-       "        272.56417847, 272.56417847, 280.56417847, 272.56417847,\n",
-       "        264.56417847, 264.56417847, 272.56417847, 272.56417847],\n",
-       "       [240.56417847, 240.56417847, 240.56417847, 240.56417847,\n",
-       "        240.56417847, 240.56417847, 240.56417847, 240.56417847,\n",
-       "        240.56417847, 240.56417847, 240.56417847, 240.56417847]])"
+       "272.5641784667969"
       ]
      },
      "execution_count": 4,
@@ -285,7 +265,7 @@
     }
    ],
    "source": [
-    "llv[2]"
+    "llv[2,0,0]"
    ]
   },
   {
@@ -305,7 +285,7 @@
     {
      "data": {
       "text/plain": [
-       "[(84,), (84,), (84,)]"
+       "(3, 84)"
       ]
      },
      "execution_count": 5,
@@ -315,7 +295,7 @@
    ],
    "source": [
     "llv = ds[0].data(flatten=True)\n",
-    "[x.shape for x in llv]"
+    "llv.shape"
    ]
   },
   {
@@ -331,7 +311,7 @@
    "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()*):"
+    "*data()* only works on a fieldlist if all the fields has the same grid. The first two elements of the resulting ndarray are the latitude and longitude arrays (shared between fields), while the rest of the elements are the value arrays per field. Since we have six fields in our data the size of the first axis of the resulting ndarray is 2+6=8."
    ]
   },
   {
@@ -343,7 +323,7 @@
     {
      "data": {
       "text/plain": [
-       "[(7, 12), (7, 12), (6, 7, 12)]"
+       "(8, 7, 12)"
       ]
      },
      "execution_count": 6,
@@ -353,28 +333,27 @@
    ],
    "source": [
     "llv = ds.data()\n",
-    "[x.shape for x in llv]"
+    "llv.shape"
    ]
   },
   {
    "cell_type": "markdown",
-   "id": "e4e9c3d3-6905-4aa9-9857-eb2a58406ac2",
+   "id": "28912fa6-4471-412c-bcf1-c186212bf092",
    "metadata": {},
    "source": [
-    "E.g. the first value from each field can be extracted as:"
+    "We can get the first latitude as:"
    ]
   },
   {
    "cell_type": "code",
    "execution_count": 7,
-   "id": "9c461ac5-5eea-460d-af41-c93af51e6f45",
+   "id": "c7f453d4-ba7e-4895-8e5b-14aa531129a4",
    "metadata": {},
    "outputs": [
     {
      "data": {
       "text/plain": [
-       "array([296.56417847, 288.53916931, 279.26531982, 259.84306335,\n",
-       "       248.00323486, 230.65315247])"
+       "90.0"
       ]
      },
      "execution_count": 7,
@@ -383,7 +362,66 @@
     }
    ],
    "source": [
-    "llv[2][:,1,1]"
+    "llv[0,0,0]"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "f93bae8f-5906-40a2-a3d9-166358c3a27c",
+   "metadata": {},
+   "source": [
+    "We can get the first longitude as:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 8,
+   "id": "754853f8-e394-445e-9eab-78bbf9eadf30",
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "0.0"
+      ]
+     },
+     "execution_count": 8,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "llv[1,0,0]"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "e4e9c3d3-6905-4aa9-9857-eb2a58406ac2",
+   "metadata": {},
+   "source": [
+    "The first value from each field can be extracted as:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 9,
+   "id": "9c461ac5-5eea-460d-af41-c93af51e6f45",
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "array([272.56417847, 272.53916931, 271.26531982, 255.84306335,\n",
+       "       244.00323486, 226.65315247])"
+      ]
+     },
+     "execution_count": 9,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "llv[2:,0,0]"
    ]
   },
   {
@@ -396,24 +434,24 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 8,
+   "execution_count": 10,
    "id": "dc0b8aa4-088b-4573-9688-46a0f0a3fdb1",
    "metadata": {},
    "outputs": [
     {
      "data": {
       "text/plain": [
-       "[(84,), (84,), (6, 84)]"
+       "(8, 84)"
       ]
      },
-     "execution_count": 8,
+     "execution_count": 10,
      "metadata": {},
      "output_type": "execute_result"
     }
    ],
    "source": [
     "llv = ds.data(flatten=True)\n",
-    "[x.shape for x in llv]"
+    "llv.shape"
    ]
   },
   {
@@ -450,7 +488,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 9,
+   "execution_count": 11,
    "id": "b9af7f2c-f120-4437-b59f-b804ce15e84b",
    "metadata": {},
    "outputs": [
@@ -479,7 +517,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 10,
+   "execution_count": 12,
    "id": "c0926aa8-7fa4-4bb6-adba-11711bc7e463",
    "metadata": {},
    "outputs": [
@@ -489,7 +527,7 @@
        "(7, 12)"
       ]
      },
-     "execution_count": 10,
+     "execution_count": 12,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -509,7 +547,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 11,
+   "execution_count": 13,
    "id": "abfb402a-743e-44e5-a01b-bcc98f0e6ed4",
    "metadata": {},
    "outputs": [
@@ -549,7 +587,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 12,
+   "execution_count": 14,
    "id": "dda20d91-1f1b-4ec5-991e-e50b841312ea",
    "metadata": {},
    "outputs": [
@@ -578,7 +616,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 13,
+   "execution_count": 15,
    "id": "e21d501a-22ef-4e82-a996-261baa2594de",
    "metadata": {},
    "outputs": [
@@ -588,7 +626,7 @@
        "(6, 7, 12)"
       ]
      },
-     "execution_count": 13,
+     "execution_count": 15,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -608,24 +646,24 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 14,
+   "execution_count": 16,
    "id": "1edbf41c-437f-41b4-aadb-c840249519b8",
    "metadata": {},
    "outputs": [
     {
      "data": {
       "text/plain": [
-       "array([296.56417847, 288.53916931, 279.26531982, 259.84306335,\n",
-       "       248.00323486, 230.65315247])"
+       "array([272.56417847, 272.53916931, 271.26531982, 255.84306335,\n",
+       "       244.00323486, 226.65315247])"
       ]
      },
-     "execution_count": 14,
+     "execution_count": 16,
      "metadata": {},
      "output_type": "execute_result"
     }
    ],
    "source": [
-    "v[:,1,1]"
+    "v[:,0,0]"
    ]
   },
   {
@@ -646,18 +684,18 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 15,
+   "execution_count": 17,
    "id": "bf898af0-7382-40e8-89fd-8a07bc250378",
    "metadata": {},
    "outputs": [
     {
      "data": {
       "text/plain": [
-       "array([296.56418, 288.53918, 279.26532, 259.84308, 248.00323, 230.65315],\n",
+       "array([272.56418, 272.53918, 271.26532, 255.84306, 244.00323, 226.65315],\n",
        "      dtype=float32)"
       ]
      },
-     "execution_count": 15,
+     "execution_count": 17,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -665,7 +703,7 @@
    "source": [
     "import numpy as np\n",
     "v = ds.to_numpy(dtype=np.float32)\n",
-    "v[:,1,1]"
+    "v[:,0,0]"
    ]
   },
   {
diff --git a/earthkit/data/core/fieldlist.py b/earthkit/data/core/fieldlist.py
index 92adc3a..f51945a 100644
--- a/earthkit/data/core/fieldlist.py
+++ b/earthkit/data/core/fieldlist.py
@@ -77,12 +77,36 @@ class Field(Base):
 
         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``).
+        ndarray
+            An ndarray containing one ndarray per key is returned
+            (following the order in ``keys``). When ``keys`` is a single value only the
+            ndarray belonging to the key is returned.
+
+
+        Examples
+        --------
+        - :ref:`/examples/grib_lat_lon_value.ipynb`
+
+        >>> import earthkit.data
+        >>> ds = earthkit.data.from_source("file", "docs/examples/test6.grib")
+        >>> d = ds[0].data()
+        >>> d.shape
+        (3, 7, 12)
+        >>> d[0, 0, 0]  # first latitude
+        90.0
+        >>> d[1, 0, 0]  # first longitude
+        0.0
+        >>> d[2, 0, 0]  # first value
+        272.56417847
+        >>> d = ds[0].data(keys="lon")
+        >>> d.shape
+        (7, 12)
+        >>> d[0, 0]  # first longitude
+        0.0
 
         See Also
         --------
+        to_latlon
         to_points
         to_numpy
         values
@@ -105,7 +129,12 @@ class Field(Base):
         if not flatten:
             shape = self.shape
             r = [x.reshape(shape) for x in r]
-        return r[0] if len(r) == 1 else tuple(r)
+        if len(r) == 1:
+            return r[0]
+        else:
+            import numpy as np
+
+            return np.array(r)
 
     def to_points(self, flatten=False):
         r"""Return the geographical coordinates in the data's original
@@ -633,32 +662,48 @@ class FieldList(Index):
             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`.
+            When it is True the "lat", "lon" arrays and the "value" arrays per field
+            will all be flattened. Otherwise they will preserve 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:
+        ndarray
+            The elements of the ndarray (in the order of the ``keys``) are as follows:
+
+            * the latitudes array from the first field  when "lat" is in ``keys``
+            * the longitudes array from the first field when "lon" is in ``keys``
+            * a values array per field when "values" is in ``keys``
 
-            * "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`
 
+        >>> import earthkit.data
+        >>> ds = earthkit.data.from_source("file", "docs/examples/test6.grib")
+        >>> len(ds)
+        6
+        >>> d = ds.data()
+        >>> d.shape
+        (8, 7, 12)
+        >>> d[0, 0, 0]  # first latitude
+        90.0
+        >>> d[1, 0, 0]  # first longitude
+        0.0
+        >>> d[2:, 0, 0]  # first value per field
+        array([272.56417847,  -6.28688049,   7.83348083, 272.53916931,
+                -4.89837646,   8.66096497])
+        >>> d = ds.data(keys="lon")
+        >>> d.shape
+        (1, 7, 12)
+        >>> d[0, 0, 0]  # first longitude
+        0.0
 
         See Also
         --------
@@ -668,6 +713,8 @@ class FieldList(Index):
         values
 
         """
+        import numpy as np
+
         if self._is_shared_grid():
             if isinstance(keys, str):
                 keys = [keys]
@@ -682,14 +729,14 @@ class FieldList(Index):
                 elif k == "lon":
                     r.append(latlon["lon"])
                 elif k == "value":
-                    r.append(self.to_numpy(flatten=flatten))
+                    r.extend([f.to_numpy(flatten=flatten) for f in self])
                 else:
                     raise ValueError(f"data: invalid argument: {k}")
 
-            return r[0] if len(r) == 1 else tuple(r)
+            return np.array(r)
 
         elif len(self) == 0:
-            return tuple()
+            return np.array([])
         else:
             raise ValueError("Fields do not have the same grid geometry")
 


Test patch:
diff --git a/tests/grib/test_grib_values.py b/tests/grib/test_grib_values.py
index 48e3a3a..3da1749 100644
--- a/tests/grib/test_grib_values.py
+++ b/tests/grib/test_grib_values.py
@@ -228,7 +228,7 @@ def test_grib_field_data(kwarg, expected_shape):
     v = ds[0].to_numpy(**kwarg)
 
     d = ds[0].data(**kwarg)
-    assert isinstance(d, tuple)
+    assert isinstance(d, np.ndarray)
     assert len(d) == 3
     assert d[0].shape == expected_shape
     assert np.allclose(d[0], latlon["lat"])
@@ -236,66 +236,67 @@ def test_grib_field_data(kwarg, expected_shape):
     assert np.allclose(d[2], v)
 
     d = ds[0].data(keys="lat", **kwarg)
+    assert d.shape == expected_shape
     assert np.allclose(d, latlon["lat"])
 
     d = ds[0].data(keys="lon", **kwarg)
+    assert d.shape == expected_shape
     assert np.allclose(d, latlon["lon"])
 
     d = ds[0].data(keys="value", **kwarg)
+    assert d.shape == expected_shape
     assert np.allclose(d, v)
 
     d = ds[0].data(keys=("value", "lon"), **kwarg)
-    assert isinstance(d, tuple)
+    assert isinstance(d, np.ndarray)
     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",
+    "kwarg,expected_shape",
     [
-        ({}, (11, 19), (2, 11, 19)),
+        ({}, (11, 19)),
         (
             {"flatten": True},
             (209,),
-            (
-                2,
-                209,
-            ),
         ),
-        ({"flatten": False}, (11, 19), (2, 11, 19)),
+        ({"flatten": False}, (11, 19)),
     ],
 )
-def test_grib_fieldlist_data(kwarg, expected_latlon_shape, expected_values_shape):
+def test_grib_fieldlist_data(kwarg, expected_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 isinstance(d, np.ndarray)
+    assert d.shape == tuple([4, *expected_shape])
     assert np.allclose(d[0], latlon["lat"])
     assert np.allclose(d[1], latlon["lon"])
-    assert np.allclose(d[2], v)
+    assert np.allclose(d[2], v[0])
+    assert np.allclose(d[3], v[1])
 
     d = ds.data(keys="lat", **kwarg)
-    assert np.allclose(d, latlon["lat"])
+    assert d.shape == tuple([1, *expected_shape])
+    assert np.allclose(d[0], latlon["lat"])
 
     d = ds.data(keys="lon", **kwarg)
-    assert np.allclose(d, latlon["lon"])
+    assert d.shape == tuple([1, *expected_shape])
+    assert np.allclose(d[0], latlon["lon"])
 
     d = ds.data(keys="value", **kwarg)
+    assert d.shape == tuple([2, *expected_shape])
     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"])
+    assert isinstance(d, np.ndarray)
+    assert d.shape == tuple([3, *expected_shape])
+    assert np.allclose(d[0], v[0])
+    assert np.allclose(d[1], v[1])
+    assert np.allclose(d[2], latlon["lon"])
 
 
 def test_grib_values_with_missing():