Merge pull request #361 from ecmwf/feature/fix-netcdf-xy-coordinate-getter

Fix netcdf xy coordinate access

Author: sandorkertesz

earthkit/data/readers/netcdf/dataset.py

@@ -43,13 +43,15 @@ def __getattr__(self, name):
     def bbox(self, variable):
         data_array = self[variable]
 
-        coords = self._get_xy_coords(data_array)
-        key = ("bbox", tuple(coords))
+        keys, coords = self._get_xy_coords(data_array)
+        key = ("bbox", tuple(keys), tuple(coords))
         if key in self._cache:
             return self._cache[key]
 
-        lons = self._get_xy(data_array, "x", flatten=False)
-        lats = self._get_xy(data_array, "y", flatten=False)
+        # lons = self._get_xy(data_array, "x", flatten=False)
+        # lats = self._get_xy(data_array, "y", flatten=False)
+
+        lats, lons = self._get_latlon(data_array, flatten=False)
 
         north = np.amax(lats)
         west = np.amin(lons)
@@ -59,143 +61,86 @@ def bbox(self, variable):
         self._cache[key] = (north, west, south, east)
         return self._cache[key]
 
-        # dims = data_array.dims
-
-        # lat = dims[-2]
-        # lon = dims[-1]
-
-        # key = ("bbox", lat, lon)
-        # if key in self._cache:
-        #     return self._cache[key]
-
-        # lats, lons = self.grid_points(variable)
-        # north = np.amax(lats)
-        # west = np.amin(lons)
-        # south = np.amin(lats)
-        # east = np.amax(lons)
-
-        # self._cache[key] = (north, west, south, east)
-        # return self._cache[key]
-
-        # if (lat, lon) not in self._bbox:
-        #     dims = data_array.dims
-
-        #     latitude = data_array[lat]
-        #     longitude = data_array[lon]
-
-        #     self._bbox[(lat, lon)] = (
-        #         np.amax(latitude.data),
-        #         np.amin(longitude.data),
-        #         np.amin(latitude.data),
-        #         np.amax(longitude.data),
-        #     )
-
-        # return self._bbox[(lat, lon)]
-
-    # def grid_points(self, variable):
-    #     data_array = self[variable]
-    #     dims = data_array.dims
-
-    #     lat = dims[-2]
-    #     lon = dims[-1]
-
-    #     key = ("grid_points", lat, lon)
-    #     if key in self._cache:
-    #         return self._cache[key]
-
-    #     if "latitude" in self._ds and "longitude" in self._ds:
-    #         latitude = self._ds["latitude"]
-    #         longitude = self._ds["longitude"]
-
-    #         if latitude.dims == (lat, lon) and longitude.dims == (lat, lon):
-    #             latitude = latitude.data
-    #             longitude = longitude.data
-    #             return latitude.flatten(), longitude.flatten()
-
-    #     latitude = data_array[lat]
-    #     longitude = data_array[lon]
-
-    #     lat, lon = np.meshgrid(latitude.data, longitude.data)
-
-    #     self._cache[key] = lat.flatten(), lon.flatten()
-    #     return self._cache[key]
-
-    # def grid_points_xy(self, variable):
-    #     data_array = self[variable]
-    #     dims = data_array.dims
-
-    #     lat = dims[-2]
-    #     lon = dims[-1]
-
-    #     latitude = data_array[lat].data
-    #     longitude = data_array[lon].data
-
-    #     print(latitude, longitude)
-
-    #     lat, lon = np.meshgrid(latitude, longitude)
-
-    #     return lat.flatten(), lon.flatten()
-    #     # return self._cache[key]
-
-    def _get_xy(self, data_array, axis, flatten=False, dtype=None):
-        if axis not in ("x", "y"):
-            raise ValueError(f"Invalid axis={axis}")
-
-        coords = self._get_xy_coords(data_array)
-        key = ("grid_points", tuple(coords))
-        if key in self._cache:
-            points = self._cache[key]
-        else:
-            points = dict()
-            keys = [x[0] for x in coords]
-            coords = tuple([x[1] for x in coords])
-
-            if "latitude" in self._ds and "longitude" in self._ds:
-                latitude = self._ds["latitude"]
-                longitude = self._ds["longitude"]
-
-                if latitude.dims == coords and longitude.dims == coords:
-                    latitude = latitude.data
-                    longitude = longitude.data
-                    points["x"] = longitude
-                    points["y"] = latitude
-            if not points:
-                v0, v1 = data_array.coords[coords[0]], data_array.coords[coords[1]]
-                points[keys[1]], points[keys[0]] = np.meshgrid(v1, v0)
-                self._cache[key] = points
-
-        if flatten:
-            points[axis] = points[axis].reshape(-1)
-        if dtype is not None:
-            return points[axis].astype(dtype)
-        else:
-            return points[axis]
-
     def _get_xy_coords(self, data_array):
-        c = []
-
         if (
             len(data_array.dims) >= 2
             and data_array.dims[-1] in GEOGRAPHIC_COORDS["x"]
             and data_array.dims[-2] in GEOGRAPHIC_COORDS["y"]
         ):
-            return [("y", data_array.dims[-2]), ("x", data_array.dims[-1])]
+            return ("y", "x"), (data_array.dims[-2], data_array.dims[-1])
 
+        keys = []
+        coords = []
         axes = ("x", "y")
         for dim in data_array.dims:
             for ax in axes:
                 candidates = GEOGRAPHIC_COORDS.get(ax, [])
                 if dim in candidates:
-                    c.append((ax, dim))
+                    keys.append(ax)
+                    coords.append(dim)
                 else:
                     ax = data_array.coords[dim].attrs.get("axis", "").lower()
                     if ax in axes:
-                        c.append([ax, dim])
-            if len(c) == 2:
-                return c
+                        keys.append(ax)
+                        coords.append(dim)
+            if len(keys) == 2:
+                return tuple(keys), tuple(coords)
 
         for ax in axes:
-            if ax not in [x[0] for x in c]:
+            if ax not in keys:
                 raise ValueError(f"No coordinate found with axis '{ax}'")
 
-        return c
+        return keys, coords
+
+    def _get_xy(self, data_array, flatten=False, dtype=None):
+        keys, coords = self._get_xy_coords(data_array)
+        key = ("grid_points", tuple(keys), tuple(coords))
+
+        if key in self._cache:
+            points = self._cache[key]
+        else:
+            points = dict()
+            v0, v1 = data_array.coords[coords[0]], data_array.coords[coords[1]]
+            points[keys[1]], points[keys[0]] = np.meshgrid(v1, v0)
+            self._cache[key] = points
+
+        if flatten:
+            points["x"] = points["x"].reshape(-1)
+            points["y"] = points["y"].reshape(-1)
+
+        if dtype is not None:
+            return points["x"].astype(dtype), points["y"].astype(dtype)
+        else:
+            return points["x"], points["y"]
+
+    def _get_latlon(self, data_array, flatten=False, dtype=None):
+        keys, coords = self._get_xy_coords(data_array)
+
+        points = dict()
+        if "latitude" in self._ds and "longitude" in self._ds:
+            latitude = self._ds["latitude"]
+            longitude = self._ds["longitude"]
+            if latitude.dims == coords and longitude.dims == coords:
+                latitude = latitude.data
+                longitude = longitude.data
+                points["y"] = latitude
+                points["x"] = longitude
+
+        if not points:
+            key = ("grid_points", tuple(keys), tuple(coords))
+
+            if key in self._cache:
+                points = self._cache[key]
+            else:
+                v0, v1 = data_array.coords[coords[0]], data_array.coords[coords[1]]
+                points[keys[1]], points[keys[0]] = np.meshgrid(v1, v0)
+                self._cache[key] = points
+
+        if flatten:
+            points["x"] = points["x"].reshape(-1)
+            points["y"] = points["y"].reshape(-1)
+
+        if dtype is not None:
+            return points["y"].astype(dtype), points["x"].astype(dtype)
+        else:
+            return points["y"], points["x"]

earthkit/data/readers/netcdf/field.py

@@ -30,20 +30,20 @@ def __init__(self, metadata, data_array, ds, variable):
         self.north, self.west, self.south, self.east = self.ds.bbox(variable)
 
     def latitudes(self, dtype=None):
-        return self.y(dtype=dtype)
+        return self.ds._get_latlon(self.data_array, flatten=True, dtype=dtype)[0]
 
     def longitudes(self, dtype=None):
-        return self.x(dtype=dtype)
+        return self.ds._get_latlon(self.data_array, flatten=True, dtype=dtype)[1]
 
     def x(self, dtype=None):
-        return self.ds._get_xy(self.data_array, "x", flatten=True, dtype=dtype)
+        return self.ds._get_xy(self.data_array, flatten=True, dtype=dtype)[0]
 
     def y(self, dtype=None):
-        return self.ds._get_xy(self.data_array, "y", flatten=True, dtype=dtype)
+        return self.ds._get_xy(self.data_array, flatten=True, dtype=dtype)[1]
 
     def shape(self):
-        coords = self.ds._get_xy_coords(self.data_array)
-        return tuple([self.data_array.coords[v[1]].size for v in coords])
+        _, coords = self.ds._get_xy_coords(self.data_array)
+        return tuple([self.data_array.coords[v].size for v in coords])
 
     def _unique_grid_id(self):
         return self.shape

tests/netcdf/test_netcdf_geography.py

@@ -167,6 +167,78 @@ def test_netcdf_proj_string_laea():
     )
 
 
+def test_netcdf_to_points_laea():
+    ds = from_source("url", earthkit_remote_test_data_file("examples", "efas.nc"))
+
+    assert len(ds) == 3
+
+    pos = [(0, 0), (0, -1), (-1, 0), (-1, -1)]
+
+    # we must check multiple fields
+    for idx in range(2):
+        v = ds[idx].to_points()
+        assert isinstance(v, dict)
+
+        # lon
+        assert isinstance(v["x"], np.ndarray)
+        assert v["x"].shape == (950, 1000)
+
+        ref = np.array([2502500.0, 7497500.0, 2502500.0, 7497500.0])
+        for i, x in enumerate(pos):
+            assert np.isclose(v["x"][x], ref[i]), f"{i=}, {x=}"
+
+        # lat
+        assert isinstance(v["y"], np.ndarray)
+        assert v["y"].shape == (950, 1000)
+
+        ref = np.array([5497500.0, 5497500.0, 752500.0, 752500.0])
+        for i, x in enumerate(pos):
+            assert np.isclose(v["y"][x], ref[i]), f"{i=}, {x=}"
+
+
+def test_netcdf_to_latlon_laea():
+    ds = from_source("url", earthkit_remote_test_data_file("examples", "efas.nc"))
+
+    assert len(ds) == 3
+
+    pos = [(0, 0), (0, -1), (-1, 0), (-1, -1)]
+
+    # we must check multiple fields
+    for idx in range(2):
+        v = ds[idx].to_latlon()
+        assert isinstance(v, dict)
+
+        # lon
+        assert isinstance(v["lon"], np.ndarray)
+        assert v["lon"].shape == (950, 1000)
+
+        ref = np.array(
+            [
+                -35.034023999999995,
+                73.93767587613708,
+                -8.229274420493763,
+                41.13970495087975,
+            ]
+        )
+        for i, x in enumerate(pos):
+            assert np.isclose(v["lon"][x], ref[i]), f"{i=}, {x=}"
+
+        # lat
+        assert isinstance(v["lat"], np.ndarray)
+        assert v["lat"].shape == (950, 1000)
+
+        ref = np.array(
+            [
+                66.9821429989222,
+                58.24673887576243,
+                27.802844211251625,
+                23.942342882929605,
+            ]
+        )
+        for i, x in enumerate(pos):
+            assert np.isclose(v["lat"][x], ref[i]), f"{i=}, {x=}"
+
+
 if __name__ == "__main__":
     from earthkit.data.testing import main