added the tgrid docstring example to the tests

pyicon/tests/test_pyicon_calc_xr.py

@@ -2,7 +2,12 @@ from itertools import product
 import pytest
 import numpy as np
 import pyicon as pyic
-from .conftest import eager_raw_grid, eager_processed_tgrid, lazy_raw_grid, lazy_processed_tgrid
+from .conftest import (
+    eager_raw_grid,
+    eager_processed_tgrid,
+    lazy_raw_grid,
+    lazy_processed_tgrid,
+)
 
 
 @pytest.mark.parametrize("raw_grid", ["lazy_raw_grid", "eager_raw_grid"])
@@ -26,14 +31,17 @@ def test_convert_tgrid_data(raw_grid, request):
         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
+            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
+            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
+            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")):
@@ -52,7 +60,15 @@ def test_convert_tgrid_data(raw_grid, request):
     assert "cell" in converted_tgrid.dims
 
 
-@pytest.mark.parametrize("tgrid", ["lazy_raw_grid", "lazy_processed_tgrid", "eager_raw_grid", "eager_processed_tgrid"])
+@pytest.mark.parametrize(
+    "tgrid",
+    [
+        "lazy_raw_grid",
+        "lazy_processed_tgrid",
+        "eager_raw_grid",
+        "eager_processed_tgrid",
+    ],
+)
 def test_xr_crop_tgrid(tgrid, request):
     tgrid = request.getfixturevalue(tgrid)
 
@@ -64,10 +80,17 @@ def test_xr_crop_tgrid(tgrid, request):
     tgrid["clon"] = tgrid["clon"].compute()
     tgrid["clat"] = tgrid["clat"].compute()
 
-    ireg_c = tgrid["cell"].where(
-        (tgrid["clon"] > -5) & (tgrid["clon"] < 5)
-        & (tgrid["clat"] > -5) & (tgrid["clat"] < 5),
-        drop=True).astype("int32")
+    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
@@ -93,8 +116,23 @@ def test_xr_crop_tgrid(tgrid, request):
     assert cropped_tgrid["ireg_v"].sum() == 135385
     assert cropped_tgrid["ireg_v"].prod() == -1427286351937536000
 
-
-@pytest.mark.parametrize("processed_tgrid", ["lazy_processed_tgrid", "eager_processed_tgrid"])
+    # Try running the example code from the docstring
+    clon = tgrid.clon.compute().data * 180.0 / np.pi
+    clat = tgrid.clat.compute().data * 180.0 / np.pi
+    lon_reg = [6, 10]
+    lat_reg = [-32, -30]
+    ireg_c = np.where(
+        (clon > lon_reg[0])
+        & (clon <= lon_reg[1])
+        & (clat > lat_reg[0])
+        & (clat <= lat_reg[1])
+    )[0]
+    pyic.xr_crop_tgrid(tgrid, ireg_c)
+
+
+@pytest.mark.parametrize(
+    "processed_tgrid", ["lazy_processed_tgrid", "eager_processed_tgrid"]
+)
 def test_nabla_funcs(processed_tgrid, request):
     processed_tgrid = request.getfixturevalue(processed_tgrid)