Add open_multigrid list_grid_names functions (#1404)

* o Add open_multigrid/list_grid_names API with initial OASIS multi-grid SCRIP support (https://cerfacs.fr/oa4web/oasis3-mct_5.0/oasis3mct_UserGuide/node50.html)

* o Use pytest style

* Fix assumption that center_lon is a NumPy array

* o Handle stacking/masking by flattening multi-dimensional cell axes before building grids

* o Try to fix mask errors reported by Nils for OASIS format.

* o Add test for open_multigrid_mask_zero_faces

* Add OASIS mask fixture for multigrid mask test

* o customizable mask values in the open_multigrid function and tests

* o edge case for when a mask zeros out all faces

* o Keep SCRIP grid center calc lazy, use where construct instead of asarray

---------

Co-authored-by: Orhan Eroglu <[email protected]>

Author: rajeeja

docs/api.rst

@@ -18,6 +18,7 @@ Top Level Functions
 
    open_grid
    open_dataset
+   open_multigrid
    open_mfdataset
    concat
 

test/core/test_api.py

@@ -4,10 +4,13 @@
 import pytest
 import tempfile
 import xarray as xr
+from pathlib import Path
 from unittest.mock import patch
 from uxarray.core.utils import _open_dataset_with_fallback
 import os
 
+TEST_MESHFILES = Path(__file__).resolve().parent.parent / "meshfiles"
+
 def test_open_geoflow_dataset(gridpath, datasetpath):
     """Loads a single dataset with its grid topology file using uxarray's
     open_dataset call."""
@@ -163,3 +166,71 @@ def mock_open_dataset(*args, **kwargs):
             ds_fallback.close()
         if os.path.exists(tmp_path):
             os.unlink(tmp_path)
+
+
+def test_list_grid_names_multigrid(gridpath):
+    """List grids from an OASIS-style multi-grid file."""
+    grid_file = gridpath("scrip", "oasis", "grids.nc")
+    grid_names = ux.list_grid_names(grid_file)
+
+    assert isinstance(grid_names, list)
+    assert set(grid_names) == {"ocn", "atm"}
+
+
+def test_list_grid_names_single_scrip():
+    """List grids from a standard single-grid SCRIP file."""
+    grid_path = TEST_MESHFILES / "scrip" / "outCSne8" / "outCSne8.nc"
+    grid_names = ux.list_grid_names(grid_path)
+
+    assert isinstance(grid_names, list)
+    assert grid_names == ["grid"]
+
+
+def test_open_multigrid_all_grids(gridpath):
+    """Open all grids from a multi-grid file."""
+    grid_file = gridpath("scrip", "oasis", "grids.nc")
+    grids = ux.open_multigrid(grid_file)
+
+    assert isinstance(grids, dict)
+    assert set(grids.keys()) == {"ocn", "atm"}
+    assert grids["ocn"].n_face == 12
+    assert grids["atm"].n_face == 20
+
+
+def test_open_multigrid_specific_grids(gridpath):
+    """Open a subset of grids from a multi-grid file."""
+    grid_file = gridpath("scrip", "oasis", "grids.nc")
+    grids = ux.open_multigrid(grid_file, gridnames=["ocn"])
+
+    assert set(grids.keys()) == {"ocn"}
+    assert grids["ocn"].n_face == 12
+
+
+def test_open_multigrid_with_masks(gridpath):
+    """Open grids with a companion mask file."""
+    grid_file = gridpath("scrip", "oasis", "grids.nc")
+    mask_file = gridpath("scrip", "oasis", "masks.nc")
+
+    grids = ux.open_multigrid(grid_file, mask_filename=mask_file)
+
+    assert grids["ocn"].n_face == 8
+    assert grids["atm"].n_face == 20
+
+
+def test_open_multigrid_mask_zero_faces(gridpath):
+    """Applying masks that deactivate an entire grid should not fail."""
+    grid_file = gridpath("scrip", "oasis", "grids.nc")
+    mask_file = gridpath("scrip", "oasis", "masks_no_atm.nc")
+
+    grids = ux.open_multigrid(grid_file, mask_filename=mask_file)
+
+    assert grids["ocn"].n_face == 8
+    assert grids["atm"].n_face == 0
+
+
+def test_open_multigrid_missing_grid_error(gridpath):
+    """Requesting a missing grid should raise."""
+    grid_file = gridpath("scrip", "oasis", "grids.nc")
+
+    with pytest.raises(ValueError, match="Grid 'land' not found"):
+        ux.open_multigrid(grid_file, gridnames=["land"])

test/grid/grid/test_initialization.py

@@ -1,6 +1,7 @@
 import numpy as np
 import numpy.testing as nt
 import pytest
+import xarray as xr
 
 import uxarray as ux
 from uxarray.constants import INT_FILL_VALUE, ERROR_TOLERANCE
@@ -77,3 +78,23 @@ def test_from_topology():
         face_node_connectivity=face_node_connectivity,
         fill_value=-1,
     )
+
+
+def test_grid_init_handles_empty_longitude_fields():
+    """Ensure grids with empty longitude arrays don't error during initialization."""
+    empty_lon = np.array([], dtype=np.float64)
+    ds = xr.Dataset(
+        {
+            "node_lon": (("n_node",), empty_lon),
+            "node_lat": (("n_node",), empty_lon),
+            "face_node_connectivity": (
+                ("n_face", "n_max_face_nodes"),
+                np.empty((0, 0), dtype=np.int64),
+            ),
+            "face_lon": (("n_face",), empty_lon),
+        }
+    )
+
+    uxgrid = ux.Grid(ds, source_grid_spec="UGRID")
+
+    assert uxgrid.n_face == 0

test/io/test_scrip.py

@@ -1,11 +1,13 @@
 import os
 import xarray as xr
 import warnings
+import numpy as np
 import numpy.testing as nt
 import pytest
 
 import uxarray as ux
 from uxarray.constants import INT_DTYPE, INT_FILL_VALUE
+from uxarray.io._scrip import _detect_multigrid
 
 
 def test_read_ugrid(gridpath, mesh_constants):
@@ -50,3 +52,69 @@ def test_to_xarray_ugrid(gridpath):
     reloaded_grid._ds.close()
     del reloaded_grid
     os.remove("scrip_ugrid_csne8.nc")
+
+
+def test_oasis_multigrid_format_detection():
+    """Detect OASIS-style multi-grid naming."""
+    ds = xr.Dataset()
+    ds["ocn.cla"] = xr.DataArray(np.random.rand(100, 4), dims=["nc_ocn", "nv_ocn"])
+    ds["ocn.clo"] = xr.DataArray(np.random.rand(100, 4), dims=["nc_ocn", "nv_ocn"])
+    ds["atm.cla"] = xr.DataArray(np.random.rand(200, 4), dims=["nc_atm", "nv_atm"])
+    ds["atm.clo"] = xr.DataArray(np.random.rand(200, 4), dims=["nc_atm", "nv_atm"])
+
+    format_type, grids = _detect_multigrid(ds)
+    assert format_type == "multi_scrip"
+    assert set(grids.keys()) == {"ocn", "atm"}
+
+
+def test_open_multigrid_with_masks(gridpath):
+    """Load OASIS multi-grids with masks applied."""
+    grid_file = gridpath("scrip", "oasis", "grids.nc")
+    mask_file = gridpath("scrip", "oasis", "masks.nc")
+
+    grids = ux.open_multigrid(grid_file, mask_filename=mask_file)
+    assert grids["ocn"].n_face == 8
+    assert grids["atm"].n_face == 20
+
+    ocean_only = ux.open_multigrid(
+        grid_file, gridnames=["ocn"], mask_filename=mask_file
+    )
+    assert set(ocean_only.keys()) == {"ocn"}
+    assert ocean_only["ocn"].n_face == 8
+
+    grid_names = ux.list_grid_names(grid_file)
+    assert set(grid_names) == {"ocn", "atm"}
+
+
+def test_open_multigrid_mask_active_value_default(gridpath):
+    """Default mask semantics keep value==1 active for both grids."""
+    grid_file = gridpath("scrip", "oasis", "grids.nc")
+    mask_file = gridpath("scrip", "oasis", "masks_no_atm.nc")
+
+    grids = ux.open_multigrid(grid_file, mask_filename=mask_file)
+
+    with xr.open_dataset(mask_file) as mask_ds:
+        expected_ocn = int(mask_ds["ocn.msk"].values.sum())
+        expected_atm = int(mask_ds["atm.msk"].values.sum())
+
+    assert grids["ocn"].n_face == expected_ocn
+    assert grids["atm"].n_face == expected_atm
+
+
+def test_open_multigrid_mask_active_value_per_grid_override(gridpath):
+    """Per-grid override supports masks with different active values."""
+    grid_file = gridpath("scrip", "oasis", "grids.nc")
+    mask_file = gridpath("scrip", "oasis", "masks_no_atm.nc")
+
+    grids = ux.open_multigrid(
+        grid_file,
+        mask_filename=mask_file,
+        mask_active_value={"atm": 0, "ocn": 1},
+    )
+
+    with xr.open_dataset(mask_file) as mask_ds:
+        expected_ocn = int(mask_ds["ocn.msk"].values.sum())
+        expected_atm = int((mask_ds["atm.msk"].values == 0).sum())
+
+    assert grids["ocn"].n_face == expected_ocn
+    assert grids["atm"].n_face == expected_atm

test/meshfiles/scrip/oasis/README.md

@@ -0,0 +1,44 @@
+# OASIS Multi-Grid SCRIP Test Files
+
+This directory contains small test files for OASIS/YAC multi-grid SCRIP format support in UXarray.
+
+## Files
+
+- `grids.nc`: Multi-grid file containing two grids
+  - `ocn`: Ocean grid with 12 cells (3x4 regular grid)
+  - `atm`: Atmosphere grid with 20 cells (4x5 regular grid)
+
+- `masks.nc`: Domain masks for the grids
+  - `ocn.msk`: Ocean mask (8 ocean cells, 4 land cells)
+  - `atm.msk`: Atmosphere mask (all 20 cells active)
+
+## OASIS Format
+
+OASIS uses a specific naming convention for multi-grid SCRIP files:
+- Grid variables are prefixed with grid name: `<gridname>.<varname>`
+- Corner latitudes: `<gridname>.cla`
+- Corner longitudes: `<gridname>.clo`
+- Dimensions: `nc_<gridname>` (cells), `nv_<gridname>` (corners)
+
+## Usage in Tests
+
+```python
+import uxarray as ux
+
+# List available grids
+grid_names = ux.list_grid_names("grids.nc")
+# ['ocn', 'atm']
+
+# Load all grids
+grids = ux.open_multigrid("grids.nc")
+
+# Load with masks
+masked_grids = ux.open_multigrid("grids.nc", mask_filename="masks.nc")
+# Ocean grid will have 8 cells, atmosphere grid will have 20 cells
+```
+
+## File Sizes
+
+These files are intentionally small for fast testing:
+- `grids.nc`: ~3 KB
+- `masks.nc`: ~1 KB