completion and clean up of assign_ugrid_topology

- adding a option to run (or not) the tests that require AWS access.
: add --online to the pytest command to run them

Author: ChrisBarker-NOAA

docs/examples/example_data/SCHISM/horizontalVelX_1.nc

@@ -7485,7 +7485,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.10.13"
+   "version": "3.12.3"
   }
  },
  "nbformat": 4,

docs/examples/example_data/SCHISM/horizontalVelY_1.nc

@@ -33,6 +33,7 @@ dependencies = [
   "numpy",
   "xarray>=2023.10.0",
   "cf_xarray",
+  "cftime",
   "dask[complete]",
   "netcdf4",
 ]
@@ -131,7 +132,7 @@ sphinx-rtd-theme = "*"
 
 [tool.pixi.feature.dev.tasks]
 lint = "ruff check"
-test = "pytest tests/"
+test = "pytest --online tests/"
 
 [tool.pixi.feature.examples.dependencies]
 matplotlib = "*"

docs/examples/example_data/SCHISM/out2d_1.nc

@@ -0,0 +1,38 @@
+# conftest: some configuration for the tests
+
+import pytest
+
+
+def pytest_addoption(parser):
+    # put a @pytest.mark.online decorator on tests that require net access
+    parser.addoption(
+        "--online",
+        action="store_true",  # what is this?
+        default=False,
+        help="run tests that access AWS resources - have to be online",
+    )
+
+
+def pytest_configure(config):
+    # register an additional marker
+    config.addinivalue_line(
+        "markers",
+        "online: mark test to run only when online (using AWS resources)"
+    )
+
+
+def pytest_collection_modifyitems(config, items):
+    if not config.getoption("--online"):
+        # --online not given in cli: skip tests that require online access
+        skip_online = pytest.mark.skip(reason="need --online option to run")
+        for item in items:
+            if "online" in item.keywords:
+                item.add_marker(skip_online)
+
+
+# # example from docs
+# def pytest_runtest_setup(item):
+#     envnames = [mark.args[0] for mark in item.iter_markers(name="env")]
+#     if envnames:
+#         if item.config.getoption("-E") not in envnames:
+#             pytest.skip(f"test requires env in {envnames!r}")

docs/examples/example_data/SCHISM/salinity_1.nc

@@ -2,6 +2,7 @@
 
 import fsspec
 import numpy as np
+import pytest
 import xarray as xr
 
 import xarray_subset_grid.accessor  # noqa: F401
@@ -13,10 +14,14 @@
 test_dir = get_test_file_dir()
 sample_sgrid_file = os.path.join(test_dir, 'arakawa_c_test_grid.nc')
 
+@pytest.mark.online
 def test_polygon_subset():
     '''
     This is a basic integration test for the subsetting of a ROMS sgrid dataset using a polygon.
     '''
+    if fsspec is None:
+        raise ImportError("Must have fsspec installed to run --online tests")
+
     fs = fsspec.filesystem(
         "reference",
         fo="s3://nextgen-dmac-cloud-ingest/nos/wcofs/nos.wcofs.2ds.best.nc.zarr",

docs/examples/example_data/SCHISM/zCoordinates_1.nc

@@ -4,7 +4,10 @@
 
 from pathlib import Path
 
-import fsspec
+try:
+    import fsspec
+except ImportError:
+    fsspec = None
 import numpy as np
 import pytest
 import xarray as xr
@@ -216,17 +219,13 @@
     int cell(nele) ;
         cell:standard_name = "cell number" ;
         cell:long_name = "Mapping to original mesh cell number" ;
-
-
-
-
 """
 
 # topology for TEST_FILE1
 grid_topology = {
-    "node_coordinates": ("lon", "lat"),
+    "node_coordinates": "lon lat",
     "face_node_connectivity": "nv",
-    "face_coordinates": ("lonc", "latc"),
+    "face_coordinates": "lonc latc",
     "face_face_connectivity": "nbe",
 }
 
@@ -394,6 +393,55 @@ def test_data_vars():
         ]
     )
 
+def test_assign_ugrid_topology_all():
+    """
+    it should use all the ones you pass in.
+
+    # The existing one in the file:
+
+    int mesh ;
+        mesh:cf_role = "mesh_topology" ;
+        mesh:long_name = "Topology data of 2D unstructured mesh" ;
+        mesh:topology_dimension = 2LL ;
+        mesh:node_coordinates = "mesh_node_lon mesh_node_lat" ;
+        mesh:edge_node_connectivity = "mesh_edge_nodes" ;
+        mesh:edge_coordinates = "mesh_edge_lon mesh_edge_lat" ;
+        mesh:face_node_connectivity = "mesh_face_nodes" ;
+        mesh:face_coordinates = "mesh_face_lon mesh_face_lat" ;
+        mesh:face_face_connectivity = "mesh_face_links" ;
+        mesh:boundary_node_connectivity = "mesh_boundary_nodes" ;
+    """
+
+    ds = xr.open_dataset(EXAMPLE_DATA / "small_ugrid_zero_based.nc")
+
+    new_attrs = {}
+    # clear out most of the mesh variable -- transfer to new attrs
+    mesh_attrs = ds['mesh'].attrs
+    for key in tuple(mesh_attrs.keys()):
+        if key not in {"cf_role", "long_name", "topology_dimension"}:
+            new_attrs[key] = mesh_attrs.pop(key)
+
+    # new_attrs = {"node_coordinates": "nc",
+    #              "face_coordinates": "fc",
+    #              "edge_coordinates": "ec",
+    #              "face_node_connectivity": "fnc",
+    #              "face_face_connectivity": "ffc",
+    #              "boundary_node_connectivity": "bnc",
+    #              "face_edge_connectivity": "fec",
+    #              "edge_face_connectivity": "efc",
+    #              "edge_node_connectivity": "enc",
+    #              }
+
+    # ds_new = ugrid.assign_ugrid_topology(ds,
+    #                                      face_node_connectivity="mesh_face_nodes",
+    #                                      boundary_node_connectivity="bounds")
+
+    ds_new = ugrid.assign_ugrid_topology(ds, **new_attrs)
+
+    print(ds_new["mesh"].attrs)
+    for key, value in new_attrs.items():
+        assert ds_new["mesh"].attrs[key] == value
+
 
 def test_extra_vars():
     """
@@ -466,8 +514,10 @@ def test_vertical_levels():
         "node",
     )
 
-
+@pytest.mark.online
 def test_3d_selector():
+    if fsspec is None:
+        raise ImportError("Must have fsspec installed to run --online tests")
     bbox = (-70, 40, -60, 55)
     name = "northeastUSA3d"
 
@@ -489,7 +539,10 @@ def test_3d_selector():
     assert bbox_selector == loaded_selector
 
 
+@pytest.mark.online
 def test_2d_selector():
+    if fsspec is None:
+        raise ImportError("Must have fsspec installed to run --online tests")
     bbox = (-70, 40, -60, 50)
     name = "northeastUSA2d"
 

docs/examples/roms-compare.ipynb

@@ -14,7 +14,7 @@
 try:
     import matplotlib.pyplot as plt
 
-    from xarray_subset_grid.visualization.mpl_plotting import plot_ugrid
+    from xarray_subset_grid.visualization.mpl_plotting import plot_ugrid, plot_sgrid
 except ImportError:
     pytestmark = pytest.mark.skip(reason="matplotlib is not installed")
 
@@ -62,3 +62,21 @@ def test_plot_ugrid_start_index_1():
     plot_ugrid(axis, ds)
 
     fig.savefig(OUTPUT_DIR / "ugrid_plot_start_index_1")
+
+
+#############
+# SGRID tests
+#############
+
+def test_plot_sgrid_only_grid():
+    import cftime
+    ds = xr.open_dataset(EXAMPLE_DATA / "wcofs_small_subset.nc")
+
+    fig, axis = plt.subplots()
+
+    plot_sgrid(axis, ds)
+
+    fig.savefig(OUTPUT_DIR / "sgrid_just_plot")
+
+
+

docs/examples/roms.ipynb

@@ -42,16 +42,16 @@ def subset_bottom_level(self, ds: xr.Dataset) -> xr.Dataset:
         if "siglay" in ds.dims:
             siglay = ds["siglay"].isel(node=0)
             if positive_direction == "up":
-                elevation_index = siglay.argmin().values
+                elevation_index = siglay.argmin(axis=0).values
             else:
-                elevation_index = siglay.argmax().values
+                elevation_index = siglay.argmax(axis=0).values
             selections["siglay"] = elevation_index
         if "siglev" in ds.dims:
             siglev = ds["siglev"].isel(node=0)
             if positive_direction == "up":
-                elevation_index = siglev.argmin().values
+                elevation_index = siglev.argmin(axis=0).values
             else:
-                elevation_index = siglev.argmax().values
+                elevation_index = siglev.argmax(axis=0).values
             selections["siglev"] = elevation_index
 
         return ds.isel(selections)
@@ -66,16 +66,16 @@ def subset_top_level(self, ds: xr.Dataset) -> xr.Dataset:
         if "siglay" in ds.dims:
             siglay = ds["siglay"].isel(node=0)
             if positive_direction == "up":
-                elevation_index = siglay.argmax().values
+                elevation_index = siglay.argmax(axis=0).values
             else:
-                elevation_index = siglay.argmin().values
+                elevation_index = siglay.argmin(axis=0).values
             selections["siglay"] = elevation_index
         if "siglev" in ds.dims:
             siglev = ds["siglev"].isel(node=0)
             if positive_direction == "up":
-                elevation_index = siglev.argmax().values
+                elevation_index = siglev.argmax(axis=0).values
             else:
-                elevation_index = siglev.argmin().values
+                elevation_index = siglev.argmin(axis=0).values
             selections["siglev"] = elevation_index
 
         return ds.isel(selections)
@@ -96,11 +96,11 @@ def subset_vertical_level(
         selections = {}
         if "siglay" in ds.dims:
             siglay = ds["siglay"].isel(node=0)
-            elevation_index = int(np.absolute(siglay - level).argmin().values)
+            elevation_index = int(np.absolute(siglay - level).argmin(axis=0).values)
             selections["siglay"] = elevation_index
         if "siglev" in ds.dims:
             siglev = ds["siglev"].isel(node=0)
-            elevation_index = int(np.absolute(siglev - level).argmin().values)
+            elevation_index = int(np.absolute(siglev - level).argmin(axis=0).values)
             selections["siglev"] = elevation_index
 
         return ds.isel(selections)
@@ -127,12 +127,12 @@ def subset_vertical_levels(
         if "siglay" in ds.dims:
             siglay = ds["siglay"].isel(node=0)
             elevation_indexes = [
-                int(np.absolute(siglay - level).argmin().values) for level in levels
+                int(np.absolute(siglay - level).argmin(axis=0).values) for level in levels
             ]
             selections["siglay"] = slice(elevation_indexes[0], elevation_indexes[1])
         if "siglev" in ds.dims:
             siglev = ds["siglev"].isel(node=0)
-            elevation_index = [int(np.absolute(siglev - level).argmin().values) for level in levels]
+            elevation_index = [int(np.absolute(siglev - level).argmin(axis=0).values) for level in levels]
             selections["siglev"] = slice(elevation_index[0], elevation_index[1])
 
         return ds.isel(selections)