Updated tests for calc_xr

pyicon/pyicon_calc_xr.py

@@ -22,13 +22,18 @@ def convert_tgrid_data(ds_tg, check_previous_conversion=True, set_dim_order=True
         Transpose the dataset so dimensions appear in the standard pyicon
         order, or the order listed
 
+    Returns
+    -------
+    ds_icd : xr.Dataset
+        A tgrid dataset compatible with pyicon functions
+
+    Notes
+    -----
     Open classical ICON grid file by:
     ds_tg = xr.open_dataset(fpath_tg, chunks=dict())
 
     Then convert by:
     ds_IcD = pyic.convert_tgrid_data(ds_tg)
-
-    ds_tg and ds_IcD are both lazy xarray data sets containing dask arrays.
     """
     if check_previous_conversion:
         if "converted_tgrid" in ds_tg.attrs:
@@ -129,7 +134,7 @@ def convert_tgrid_data(ds_tg, check_previous_conversion=True, set_dim_order=True
     ds_IcD.attrs["converted_tgrid"] = True
 
     if set_dim_order == True:
-        standard_order = ["cell", "vertex", "edge", "nc", "nv", "ne", "cart"]
+        standard_order = ["cell", "vertex", "edge", "nc", "nv", "ne", "cart", ...]
         ds_IcD = ds_IcD.transpose(*standard_order, missing_dims="ignore")
     elif set_dim_order != False:
         ds_IcD = ds_IcD.transpose(*set_dim_order, missing_dims="ignore")
@@ -293,6 +298,23 @@ def xr_crop_tgrid(ds_tg, ireg_c, verbose=1):
         "edge_vertices"
                     ]
 
+    reindex_dict = {
+        "vertex_of_cell": {
+            "from": "c", "to": "v"},
+        "vertices_of_vertex": {
+            "from": "v", "to": "v"},
+        "edge_of_cell": {
+            "from": "c", "to": "e"},
+        "edges_of_vertex": {
+            "from": "v", "to": "e"},
+        "adjacent_cell_of_edge": {
+            "from": "e", "to": "c"},
+        "cells_of_vertex": {
+            "from": "v", "to": "c"},
+        "edge_vertices": {
+            "from": "e", "to": "v"}
+    }
+
     # --- cut all variables remaining
     for var in ds_tg.keys():
         if var not in reindex_vars:

pyicon/tests/test_pyicon_calc_xr.py

-from .conftest import fruit, raw_grid
-# content of test_sample.py
+from itertools import product
+import pytest
+import numpy as np
+import pyicon as pyic
+from .conftest import raw_grid, processed_tgrid
 
-def test_fixture(fruit):
-    assert fruit == "banana"
 
+def test_convert_tgrid_data(raw_grid):
+    converted_tgrid = pyic.convert_tgrid_data(raw_grid)
+
+    # Check conversion to pythonic indexing of neighbour info has worked
+    neighbour_information = [
+        "vertex_of_cell",
+        "edge_of_cell",
+        "vertices_of_vertex",
+        "edges_of_vertex",
+        "edge_vertices",
+        "adjacent_cell_of_edge",
+        "cells_of_vertex",
+    ]
+    for info in neighbour_information:
+        assert converted_tgrid[info].min().values == 0 or -1
+
+        if info.startswith("v") or info.startswith("edge_vertices"):
+            assert converted_tgrid[info].max().values == \
+                converted_tgrid.dims["vertex"] - 1
+        elif info.startswith("e"):
+            assert converted_tgrid[info].max().values == \
+                converted_tgrid.dims["edge"] - 1
+        elif info.startswith("c") or info.startswith("a"):
+            assert converted_tgrid[info].max().values == \
+                converted_tgrid.dims["cell"] - 1
+
+    # Conversion of ecv lat and lon to degrees
+    for point, dim in product("ecv", ("lat", "lon")):
+        coord = point + dim
+        assert converted_tgrid[coord].attrs["units"] == "degrees"
+
+    # Converted tgrid attribute is there
+    assert converted_tgrid.attrs["converted_tgrid"]
+
+    # Check we can't convert a converted grid
+    with pytest.raises(ValueError):
+        pyic.convert_tgrid_data(converted_tgrid)
+
+    # Check the transposition?
+    # something with xr.DataArray.get_axis_num(dim)
+
+    # Dimension ncells is not present and cell is
+    assert "ncells" not in converted_tgrid.dims
+    assert "cell" in converted_tgrid.dims
+
+
+@pytest.mark.parametrize("tgrid", ["raw_grid", "processed_tgrid"])
+def test_xr_crop_tgrid(tgrid, request):
+    # Set ireg_c and crop the grid
+    tgrid = request.getfixturevalue(tgrid)
+
+    for point, dim in product("cev", ["lon", "lat"]):
+        coord = point + dim
+        if tgrid[coord].units == "radian":
+            tgrid[coord] = np.degrees(tgrid[coord])
+
+    ireg_c = tgrid["cell"].where(
+        (tgrid["clon"] > -5) & (tgrid["clon"] < 5)
+        & (tgrid["clat"] > -5) & (tgrid["clat"] < 5),
+        drop=True).astype("int32")
+
+    # This checks ireg_c is as expected
+    assert ireg_c.sum() == 301614
+    assert ireg_c.prod() == -8253145384319188992
+
+    cropped_tgrid = pyic.xr_crop_tgrid(tgrid, ireg_c)
+
+    # Check ireg_[cev] is present
+    for point in "cev":
+        assert f"ireg_{point}" in cropped_tgrid.keys()
+
+    # Check ncells == len(ireg_c)
+    assert cropped_tgrid.dims["cell"] == ireg_c.sizes["cell"]
+
+    # Check ireg_[cev] is correct
+    # Ideally we would hash the array and compare, but this will probably do
+    assert cropped_tgrid["ireg_c"].sum() == 301614
+    assert cropped_tgrid["ireg_c"].prod() == -8253145384319188992
+
+    assert cropped_tgrid["ireg_e"].sum() == 839941
+    assert cropped_tgrid["ireg_e"].prod() == 0
+
+    assert cropped_tgrid["ireg_v"].sum() == 135385
+    assert cropped_tgrid["ireg_v"].prod() == -1427286351937536000