YAC as an optional remapping backend for uxarray

.github/workflows/ci.yml

@@ -11,6 +11,8 @@ jobs:
     #   github.repository == 'UXARRAY/uxarray'
     name: Python (${{ matrix.python-version }}, ${{ matrix.os }})
     runs-on: ${{ matrix.os }}
+    env:
+      MPLBACKEND: Agg
     defaults:
       run:
         shell: bash -l {0}

.github/workflows/yac-optional.yml

@@ -0,0 +1,142 @@
+name: YAC Optional CI
+
+on:
+  pull_request:
+    paths:
+      - ".github/workflows/yac-optional.yml"
+      - "uxarray/remap/**"
+      - "test/test_remap_yac.py"
+  workflow_dispatch:
+
+jobs:
+  yac-optional:
+    name: YAC core v3.14.0_p1 (Ubuntu)
+    runs-on: ubuntu-latest
+    defaults:
+      run:
+        shell: bash -l {0}
+    env:
+      YAC_VERSION: v3.14.0_p1
+      YAXT_VERSION: v0.11.5.1
+      MPIEXEC: /usr/bin/mpirun
+      MPIRUN: /usr/bin/mpirun
+      MPICC: /usr/bin/mpicc
+      MPIFC: /usr/bin/mpif90
+      MPIF90: /usr/bin/mpif90
+      OMPI_ALLOW_RUN_AS_ROOT: 1
+      OMPI_ALLOW_RUN_AS_ROOT_CONFIRM: 1
+    steps:
+      - name: checkout
+        uses: actions/checkout@v4
+        with:
+          token: ${{ github.token }}
+
+      - name: conda_setup
+        uses: conda-incubator/setup-miniconda@v3
+        with:
+          activate-environment: uxarray_build
+          channel-priority: strict
+          python-version: "3.11"
+          channels: conda-forge
+          environment-file: ci/environment.yml
+          miniforge-variant: Miniforge3
+          miniforge-version: latest
+
+      - name: Install build dependencies (apt)
+        run: |
+          sudo apt-get update
+          sudo apt-get install -y \
+            autoconf \
+            automake \
+            gawk \
+            gfortran \
+            libopenmpi-dev \
+            libtool \
+            make \
+            openmpi-bin \
+            pkg-config
+      - name: Verify MPI tools
+        run: |
+          which mpirun
+          which mpicc
+          which mpif90
+          mpirun --version
+          mpicc --version
+          mpif90 --version
+      - name: Install Python build dependencies
+        run: |
+          python -m pip install --upgrade pip
+          python -m pip install cython wheel
+      - name: Build and install YAXT
+        run: |
+          set -euxo pipefail
+          YAC_PREFIX="${GITHUB_WORKSPACE}/yac_prefix"
+          echo "YAC_PREFIX=${YAC_PREFIX}" >> "${GITHUB_ENV}"
+          git clone --depth 1 --branch "${YAXT_VERSION}" https://gitlab.dkrz.de/dkrz-sw/yaxt.git
+          if [ ! -x yaxt/configure ]; then
+            if [ -x yaxt/autogen.sh ]; then
+              (cd yaxt && ./autogen.sh)
+            else
+              (cd yaxt && autoreconf -i)
+            fi
+          fi
+          mkdir -p yaxt/build
+          cd yaxt/build
+          ../configure \
+            --prefix="${YAC_PREFIX}" \
+            --without-regard-for-quality \
+            CC="${MPICC}" \
+            FC="${MPIF90}"
+          make -j2
+          make install
+      - name: Build and install YAC
+        run: |
+          set -euxo pipefail
+          git clone --depth 1 --branch "${YAC_VERSION}" https://gitlab.dkrz.de/dkrz-sw/yac.git
+          if [ ! -x yac/configure ]; then
+            if [ -x yac/autogen.sh ]; then
+              (cd yac && ./autogen.sh)
+            else
+              (cd yac && autoreconf -i)
+            fi
+          fi
+          mkdir -p yac/build
+          cd yac/build
+          ../configure \
+            --prefix="${YAC_PREFIX}" \
+            --with-yaxt-root="${YAC_PREFIX}" \
+            --disable-mci \
+            --disable-utils \
+            --disable-examples \
+            --disable-tools \
+            --disable-netcdf \
+            --enable-python-bindings \
+            CC="${MPICC}" \
+            FC="${MPIF90}"
+          make -j2
+          make install
+      - name: Configure YAC runtime paths
+        run: |
+          set -euxo pipefail
+          PY_VER="$(python -c 'import sys; print(f"{sys.version_info.major}.{sys.version_info.minor}")')"
+          echo "LD_LIBRARY_PATH=${YAC_PREFIX}/lib:${LD_LIBRARY_PATH:-}" >> "${GITHUB_ENV}"
+          echo "PYTHONPATH=${YAC_PREFIX}/lib/python${PY_VER}/site-packages:${YAC_PREFIX}/lib/python${PY_VER}/dist-packages:${PYTHONPATH:-}" >> "${GITHUB_ENV}"
+      - name: Verify YAC core Python bindings
+        run: |
+          python - <<'PY'
+          from pathlib import Path
+          import sys
+          candidates = []
+          for entry in sys.path:
+              pkg = Path(entry) / "yac"
+              candidates.extend(pkg.glob("core*.so"))
+              candidates.extend(pkg.glob("core*.pyd"))
+          assert candidates, "yac.core extension not found on sys.path"
+          print("Found yac.core extension:", candidates[0])
+          PY
+      - name: Install uxarray
+        run: |
+          python -m pip install . --no-deps
+      - name: Run tests (uxarray with YAC)
+        run: |
+          python -m pytest test/test_remap_yac.py

test/test_remap_yac.py

@@ -0,0 +1,244 @@
+import numpy as np
+import pytest
+
+import uxarray as ux
+from uxarray.remap.yac import YacNotAvailableError, _import_yac
+
+
+try:
+    _import_yac()
+except YacNotAvailableError:
+    pytest.skip("yac.core is not available", allow_module_level=True)
+
+
+def test_yac_nnn_node_remap(gridpath, datasetpath):
+    grid_path = gridpath("ugrid", "geoflow-small", "grid.nc")
+    uxds = ux.open_dataset(grid_path, datasetpath("ugrid", "geoflow-small", "v1.nc"))
+    dest = ux.open_grid(grid_path)
+
+    out = uxds["v1"].remap.nearest_neighbor(
+        destination_grid=dest,
+        remap_to="nodes",
+        backend="yac",
+        yac_method="nnn",
+        yac_options={"n": 1},
+    )
+    assert out.size > 0
+    assert "n_node" in out.dims
+
+
+def test_yac_conservative_face_remap(gridpath):
+    mesh_path = gridpath("mpas", "QU", "mesh.QU.1920km.151026.nc")
+    uxds = ux.open_dataset(mesh_path, mesh_path)
+    dest = ux.open_grid(mesh_path)
+
+    out = uxds["latCell"].remap(
+        destination_grid=dest,
+        remap_to="faces",
+        backend="yac",
+        yac_method="conservative",
+        yac_options={"order": 1},
+    )
+    assert out.size == dest.n_face
+
+
+def test_yac_matches_uxarray_nearest_neighbor():
+    verts = np.array([(0.0, 90.0), (-180.0, 0.0), (0.0, -90.0)])
+    grid = ux.open_grid(verts)
+    da = ux.UxDataArray(
+        np.asarray([1.0, 2.0, 3.0]),
+        dims=["n_node"],
+        coords={"n_node": [0, 1, 2]},
+        uxgrid=grid,
+    )
+
+    ux_out = da.remap.nearest_neighbor(
+        destination_grid=grid,
+        remap_to="nodes",
+        backend="uxarray",
+    )
+    yac_out = da.remap.nearest_neighbor(
+        destination_grid=grid,
+        remap_to="nodes",
+        backend="yac",
+        yac_method="nnn",
+        yac_options={"n": 1},
+    )
+    assert ux_out.shape == yac_out.shape
+    assert (ux_out.values == yac_out.values).all()
+
+
+def test_yac_call_defaults_to_nnn():
+    verts = np.array([(0.0, 90.0), (-180.0, 0.0), (0.0, -90.0)])
+    grid = ux.open_grid(verts)
+    da = ux.UxDataArray(
+        np.asarray([1.0, 2.0, 3.0]),
+        dims=["n_node"],
+        coords={"n_node": [0, 1, 2]},
+        uxgrid=grid,
+    )
+
+    out = da.remap(
+        destination_grid=grid,
+        remap_to="nodes",
+        backend="yac",
+    )
+
+    assert out.shape == da.shape
+    np.testing.assert_array_equal(out.values, da.values)
+
+
+def test_yac_invalid_backend_raises():
+    verts = np.array([(0.0, 90.0), (-180.0, 0.0), (0.0, -90.0)])
+    grid = ux.open_grid(verts)
+    da = ux.UxDataArray(
+        np.asarray([1.0, 2.0, 3.0]),
+        dims=["n_node"],
+        coords={"n_node": [0, 1, 2]},
+        uxgrid=grid,
+    )
+
+    with pytest.raises(ValueError, match="Invalid backend"):
+        da.remap.nearest_neighbor(
+            destination_grid=grid,
+            remap_to="nodes",
+            backend="bogus",
+        )
+
+
+def test_yac_idw_not_implemented():
+    verts = np.array([(0.0, 90.0), (-180.0, 0.0), (0.0, -90.0)])
+    grid = ux.open_grid(verts)
+    da = ux.UxDataArray(
+        np.asarray([1.0, 2.0, 3.0]),
+        dims=["n_node"],
+        coords={"n_node": [0, 1, 2]},
+        uxgrid=grid,
+    )
+
+    with pytest.raises(NotImplementedError, match="inverse_distance_weighted"):
+        da.remap.inverse_distance_weighted(
+            destination_grid=grid,
+            remap_to="nodes",
+            backend="yac",
+            yac_method="nnn",
+            yac_options={"n": 1},
+        )
+
+
+def test_yac_bilinear_face_remap(gridpath):
+    mesh_path = gridpath("mpas", "QU", "mesh.QU.1920km.151026.nc")
+    uxds = ux.open_dataset(mesh_path, mesh_path)
+    dest = ux.open_grid(mesh_path)
+
+    out = uxds["latCell"].remap.bilinear(
+        destination_grid=dest,
+        remap_to="faces",
+        backend="yac",
+    )
+
+    assert out.size == dest.n_face
+
+
+def test_yac_conservative_rejects_non_face_data():
+    verts = np.array([(0.0, 90.0), (-180.0, 0.0), (0.0, -90.0)])
+    grid = ux.open_grid(verts)
+    da = ux.UxDataArray(
+        np.asarray([1.0, 2.0, 3.0]),
+        dims=["n_node"],
+        coords={"n_node": [0, 1, 2]},
+        uxgrid=grid,
+    )
+
+    with pytest.raises(ValueError, match="face-centered"):
+        da.remap.nearest_neighbor(
+            destination_grid=grid,
+            remap_to="nodes",
+            backend="yac",
+            yac_method="conservative",
+            yac_options={"order": 1},
+        )
+
+
+def test_yac_preserves_spatial_coordinate_remap():
+    verts = np.array([(0.0, 90.0), (-180.0, 0.0), (0.0, -90.0)])
+    grid = ux.open_grid(verts)
+    da = ux.UxDataArray(
+        np.asarray([1.0, 2.0, 3.0]),
+        dims=["n_node"],
+        coords={
+            "n_node": [0, 1, 2],
+            "node_lon": (
+                "n_node",
+                np.array([0.0, -180.0, 0.0]),
+                {"standard_name": "longitude", "units": "degrees_east"},
+            ),
+            "node_lat": (
+                "n_node",
+                np.array([90.0, 0.0, -90.0]),
+                {"standard_name": "latitude", "units": "degrees_north"},
+            ),
+        },
+        uxgrid=grid,
+    )
+
+    out = da.remap.nearest_neighbor(
+        destination_grid=grid,
+        remap_to="nodes",
+        backend="yac",
+        yac_method="nnn",
+        yac_options={"n": 1},
+    )
+
+    np.testing.assert_array_equal(out.values, da.values)
+    assert "node_lon" in out.coords
+    assert "node_lat" in out.coords
+
+
+def test_yac_batched_remap_with_extra_dimension():
+    verts = np.array([(0.0, 90.0), (-180.0, 0.0), (0.0, -90.0)])
+    grid = ux.open_grid(verts)
+    da = ux.UxDataArray(
+        np.asarray([[1.0, 2.0, 3.0], [10.0, 20.0, 30.0]]),
+        dims=["time", "n_node"],
+        coords={"time": [0, 1], "n_node": [0, 1, 2]},
+        uxgrid=grid,
+    )
+
+    out = da.remap.nearest_neighbor(
+        destination_grid=grid,
+        remap_to="nodes",
+        backend="yac",
+        yac_method="nnn",
+        yac_options={"n": 1},
+    )
+
+    assert out.shape == da.shape
+    np.testing.assert_array_equal(out.values, da.values)
+
+
+def test_yac_batched_remap_with_fractional_mask():
+    verts = np.array([(0.0, 90.0), (-180.0, 0.0), (0.0, -90.0)])
+    grid = ux.open_grid(verts)
+    da = ux.UxDataArray(
+        np.asarray([[1.0, 2.0, 3.0], [10.0, 20.0, 30.0]]),
+        dims=["time", "n_node"],
+        coords={"time": [0, 1], "n_node": [0, 1, 2]},
+        uxgrid=grid,
+    )
+    frac_mask = np.ones_like(da.values, dtype=np.float64)
+
+    out = da.remap.nearest_neighbor(
+        destination_grid=grid,
+        remap_to="nodes",
+        backend="yac",
+        yac_method="nnn",
+        yac_options={
+            "n": 1,
+            "frac_mask_fallback_value": 0.0,
+            "frac_mask": frac_mask,
+        },
+    )
+
+    assert out.shape == da.shape
+    np.testing.assert_array_equal(out.values, da.values)