added pyicon tests to ensure subsequent changes don't break functionality

pyicon/pyicon_calc_xr.py

@@ -25,7 +25,7 @@ ds_IcD = pyic.convert_tgrid_data(ds_tg)
 ds_tg and ds_IcD are both lazy xarray data sets containing dask arrays.
     """
     ds_IcD = xr.Dataset()
-    
+
     # --- constants (from src/shared/mo_physical_constants.f90)
     ds_IcD['grid_sphere_radius'] = 6.371229e6
     ds_IcD['grav'] = 9.80665

pyicon/tests/conftest.py

+import pytest
+from pathlib import Path
+import xarray as xr
+import pyicon as pyic
+
+
+@pytest.fixture()
+def raw_grid():
+    cur_dir = Path(__file__).parent.resolve()
+    grid_path = cur_dir / "test_data/icon_grid_0014_R02B04_O.nc"
+
+    if not grid_path.exists():
+        import requests
+        grid_download_link = "http://icon-downloads.mpimet.mpg.de/grids/public/mpim/0014/icon_grid_0014_R02B04_O.nc"
+        try:
+            r = requests.get(grid_download_link, allow_redirects=True)
+            with open(grid_path, "wb") as grid_file:
+                grid_file.write(r.content)
+        except:
+            raise FileNotFoundError("{grid_path} does not exist and unable to \
+                download it")
+
+    ds_grid = xr.open_dataset(grid_path)
+    return ds_grid
+
+
+@pytest.fixture()
+def processed_tgrid(raw_grid):
+    return pyic.convert_tgrid_data(raw_grid)

pyicon/tests/test_pyicon_calc_xr.py

+from itertools import product
+import pytest
+import numpy as np
+import pyicon as pyic
+from .conftest import raw_grid, processed_tgrid
+
+
+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
+
+    # 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"])
+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")
+
+    cropped_tgrid = pyic.xr_crop_tgrid(tgrid, ireg_c)
+
+    # Check ireg_[cev] is present
+    for point in "ev":
+        assert f"ireg_{point}" in cropped_tgrid.keys()
+
+    # Check ncells == len(ireg_c)
+    assert cropped_tgrid.dims["cell"] == ireg_c.sizes["cell"]
+
+    # Check ireg_[ev] is correct
+    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