Add more netcdf geography tests (#873)

Author: sandorkertesz

tests/netcdf/test_netcdf_geography.py

@@ -12,6 +12,7 @@
 import numpy as np
 import pytest
 
+from earthkit.data import from_object
 from earthkit.data import from_source
 from earthkit.data.testing import earthkit_examples_file
 from earthkit.data.testing import earthkit_remote_examples_file
@@ -260,6 +261,221 @@ def test_netcdf_forecast_reference_time():
     assert ds[5].metadata("valid_datetime") == "2020-01-23T05:00:00"
 
 
+@pytest.mark.parametrize("lat_name,lon_name", [("lat", "lon"), ("latitude", "longitude")])
+def test_netcdf_geography_2d_1(lat_name, lon_name):
+    # Dimensions:  (level: 2, lat: 3, lon: 3)
+    # Coordinates:
+    #   * level    (level) int64 16B 700 500
+    #   * lat      (lat) int64 24B 50 40 30
+    #   * lon      (lon) int64 24B 0 10 20
+    # Data variables:
+    #     a        (level, lat, lon) int64 144B 11 12 13 21 22 23 ... 25 26 34 35 36
+
+    import xarray as xr
+
+    dims = {"level": 2, lat_name: 3, lon_name: 3}
+    coords = {
+        "level": np.array([700, 500]),
+        lat_name: np.array([50, 40, 30]),
+        lon_name: np.array([0, 10, 20]),
+    }
+
+    lats = [[50, 50, 50], [40, 40, 40], [30, 30, 30]]
+    lons = [[0, 10, 20], [0, 10, 20], [0, 10, 20]]
+
+    data = np.array(
+        [
+            [[11, 12, 13], [21, 22, 23], [31, 32, 33]],
+            [[14, 15, 16], [24, 25, 26], [34, 35, 36]],
+        ]
+    )
+
+    a = xr.Variable(dims, data)
+    v = {"a": a}
+    ds_in = xr.Dataset(v, coords=coords)
+
+    ds = from_object(ds_in)
+    assert len(ds) == 2
+    assert np.allclose(ds.metadata("level"), coords["level"])
+
+    for ll in [ds[0].to_latlon(), ds.to_latlon()]:
+        assert ll["lat"].shape == (3, 3)
+        assert ll["lon"].shape == (3, 3)
+        assert np.allclose(ll["lat"], lats)
+        assert np.allclose(ll["lon"], lons)
+
+
+@pytest.mark.parametrize("lat_name,lon_name", [("lat", "lon"), ("latitude", "longitude")])
+def test_netcdf_geography_2d_2(lat_name, lon_name):
+    # Dimensions:  (level: 2, y: 3, x: 2)
+    # Coordinates:
+    #   * level    (level) int64 16B 700 500
+    #     lat      (y, x) int64 48B 50 50 40 40 30 30
+    #     lon      (y, x) int64 48B 0 10 0 10 0 10
+    # Dimensions without coordinates: y, x
+    # Data variables:
+    #     a        (level, y, x) int64 96B 11 12 21 22 31 32 14 15 24 25 34 35
+
+    import xarray as xr
+
+    dims = {"level": 2, "y": 3, "x": 2}
+    coords = {
+        "level": np.array([700, 500]),
+        lat_name: (["y", "x"], np.array([[50, 50], [40, 40], [30, 30]])),
+        lon_name: (["y", "x"], np.array([[0, 10], [0, 10], [0, 10]])),
+    }
+
+    data = np.array(
+        [
+            [[11, 12], [21, 22], [31, 32]],
+            [[14, 15], [24, 25], [34, 35]],
+        ]
+    )
+
+    a = xr.Variable(dims, data)
+    v = {"a": a}
+    ds_in = xr.Dataset(v, coords=coords)
+
+    ds = from_object(ds_in)
+    assert len(ds) == 2
+    assert np.allclose(ds.metadata("level"), coords["level"])
+
+    for ll in [ds[0].to_latlon(), ds.to_latlon()]:
+        assert ll["lat"].shape == (3, 2)
+        assert ll["lon"].shape == (3, 2)
+        assert np.allclose(ll["lat"], coords[lat_name][1])
+        assert np.allclose(ll["lon"], coords[lon_name][1])
+
+
+@pytest.mark.skip(reason="To be seen if lat-lon as variables have to be supported")
+@pytest.mark.parametrize("lat_name,lon_name", [("lat", "lon"), ("latitude", "longitude")])
+def test_netcdf_geography_2d_3(lat_name, lon_name):
+    # Dimensions:  (level: 2, y: 3, x: 2)
+    # Coordinates:
+    #   * level    (level) int64 16B 700 500
+    # Dimensions without coordinates: y, x
+    # Data variables:
+    #     a        (level, y, x) int64 96B 11 12 21 22 31 32 14 15 24 25 34 35
+    #     lat      (y, x) int64 48B 50 50 40 40 30 30
+    #     lon      (y, x) int64 48B 0 10 0 10 0 10
+
+    import xarray as xr
+
+    dims = {"level": 2, "y": 3, "x": 2}
+    coords = {
+        "level": np.array([700, 500]),
+    }
+
+    data = np.array(
+        [
+            [[11, 12], [21, 22], [31, 32]],
+            [[14, 15], [24, 25], [34, 35]],
+        ]
+    )
+
+    a = xr.Variable(dims, data)
+    lat = xr.Variable({"y": 3, "x": 3}, np.array([[50, 50], [40, 40], [30, 30]]))
+    lon = xr.Variable({"y": 3, "x": 3}, np.array([[0, 10], [0, 10], [0, 10]]))
+    v = {"a": a, lat_name: lat, lon_name: lon}
+    ds_in = xr.Dataset(v, coords=coords)
+
+    ds = from_object(ds_in)
+    assert len(ds) == 2
+    assert np.allclose(ds.metadata("level"), coords["level"])
+
+    for ll in [ds[0].to_latlon(), ds.to_latlon()]:
+        assert ll["lat"].shape == (3, 2)
+        assert ll["lon"].shape == (3, 2)
+        assert np.allclose(ll["lat"], lat.data)
+        assert np.allclose(ll["lon"], lon.data)
+
+
+@pytest.mark.parametrize("lat_name,lon_name", [("lat", "lon"), ("latitude", "longitude")])
+def test_netcdf_geography_1d_1(lat_name, lon_name):
+    # Dimensions:  (level: 2, values: 9)
+    # Coordinates:
+    #   * level    (level) int64 16B 700 500
+    #     lat      (values) int64 72B 50 50 50 40 40 40 30 30 30
+    #     lon      (values) int64 72B 0 10 20 0 10 20 0 10 20
+    # Dimensions without coordinates: values
+    # Data variables:
+    #     a        (level, values) int64 144B 11 12 13 21 22 23 ... 24 25 26 34 35 36
+
+    import xarray as xr
+
+    dims = {"level": 2, "values": 9}
+    coords = {
+        "level": np.array([700, 500]),
+        lat_name: ("values", np.array([50, 50, 50, 40, 40, 40, 30, 30, 30])),
+        lon_name: ("values", np.array([0, 10, 20, 0, 10, 20, 0, 10, 20])),
+    }
+
+    data = np.array(
+        [
+            [11, 12, 13, 21, 22, 23, 31, 32, 33],
+            [14, 15, 16, 24, 25, 26, 34, 35, 36],
+        ]
+    )
+
+    a = xr.Variable(dims, data)
+    v = {"a": a}
+    ds_in = xr.Dataset(v, coords=coords)
+
+    ds = from_object(ds_in)
+    assert len(ds) == 2
+    assert np.allclose(ds.metadata("level"), coords["level"])
+
+    for ll in [ds[0].to_latlon(), ds.to_latlon()]:
+        assert ll["lat"].shape == (9,)
+        assert ll["lon"].shape == (9,)
+        assert np.allclose(ll["lat"], coords[lat_name][1])
+        assert np.allclose(ll["lon"], coords[lon_name][1])
+
+
+@pytest.mark.skip(reason="To be seen if lat-lon as variables have to be supported")
+@pytest.mark.parametrize("lat_name,lon_name", [("lat", "lon"), ("latitude", "longitude")])
+def test_netcdf_geography_1d_2(lat_name, lon_name):
+    # Dimensions:  (level: 2, values: 9)
+    # Coordinates:
+    # * level    (level) int64 16B 700 500
+    # Dimensions without coordinates: values
+    # Data variables:
+    #     a        (level, values) int64 144B 11 12 13 21 22 23 ... 24 25 26 34 35 36
+    #     lat      (values) int64 72B 50 50 50 40 40 40 30 30 30
+    #     lon      (values) int64 72B 0 10 20 0 10 20 0 10 20
+
+    import xarray as xr
+
+    dims = {"level": 2, "values": 9}
+    coords = {
+        "level": np.array([700, 500]),
+    }
+
+    data = np.array(
+        [
+            [11, 12, 13, 21, 22, 23, 31, 32, 33],
+            [14, 15, 16, 24, 25, 26, 34, 35, 36],
+        ]
+    )
+
+    a = xr.Variable(dims, data)
+    lat = xr.Variable({"values": 9}, np.array([50, 50, 50, 40, 40, 40, 30, 30, 30]))
+    lon = xr.Variable({"values": 9}, np.array([0, 10, 20, 0, 10, 20, 0, 10, 20]))
+
+    v = {"a": a, lat_name: lat, lon_name: lon}
+    ds_in = xr.Dataset(v, coords=coords)
+
+    ds = from_object(ds_in)
+    assert len(ds) == 2
+    assert np.allclose(ds.metadata("level"), coords["level"])
+
+    for ll in [ds[0].to_latlon(), ds.to_latlon()]:
+        assert ll["lat"].shape == (9,)
+        assert ll["lon"].shape == (9,)
+        assert np.allclose(ll["lat"], lat.data)
+        assert np.allclose(ll["lon"], lon.data)
+
+
 if __name__ == "__main__":
     from earthkit.data.testing import main