Add optional YAC remapping backend with NNN/conservative support and CI coverage

.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,68 @@
+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.nearest_neighbor(
+        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()

uxarray/remap/accessor.py

@@ -27,17 +27,36 @@ def __repr__(self) -> str:
             + "  • inverse_distance_weighted(destination_grid, remap_to='faces', power=2, k=8)\n"
         )
 
-    def __call__(self, *args, **kwargs) -> UxDataArray | UxDataset:
+    def __call__(
+        self,
+        *args,
+        backend: str = "uxarray",
+        yac_method: str | None = None,
+        yac_options: dict | None = None,
+        **kwargs,
+    ) -> UxDataArray | UxDataset:
         """
         Shortcut for nearest-neighbor remapping.
 
         Calling `.remap(...)` with no explicit method will invoke
         `nearest_neighbor(...)`.
         """
-        return self.nearest_neighbor(*args, **kwargs)
+        return self.nearest_neighbor(
+            *args,
+            backend=backend,
+            yac_method=yac_method,
+            yac_options=yac_options,
+            **kwargs,
+        )
 
     def nearest_neighbor(
-        self, destination_grid: Grid, remap_to: str = "faces", **kwargs
+        self,
+        destination_grid: Grid,
+        remap_to: str = "faces",
+        backend: str = "uxarray",
+        yac_method: str | None = "nnn",
+        yac_options: dict | None = None,
+        **kwargs,
     ) -> UxDataArray | UxDataset:
         """
         Perform nearest-neighbor remapping.
@@ -51,16 +70,39 @@ def nearest_neighbor(
         remap_to : {'nodes', 'edges', 'faces'}, default='faces'
             Which grid element receives the remapped values.
 
+        backend : {'uxarray', 'yac'}, default='uxarray'
+            Remapping backend to use. When set to 'yac', requires YAC to be
+            available on PYTHONPATH.
+        yac_method : {'nnn', 'conservative'}, optional
+            YAC interpolation method. Defaults to 'nnn' when backend='yac'.
+        yac_options : dict, optional
+            YAC interpolation configuration options.
+
         Returns
         -------
         UxDataArray or UxDataset
             A new object with data mapped onto `destination_grid`.
         """
 
+        if backend == "yac":
+            from uxarray.remap.yac import _yac_remap
+
+            yac_kwargs = yac_options or {}
+            return _yac_remap(
+                self.ux_obj, destination_grid, remap_to, yac_method, yac_kwargs
+            )
         return _nearest_neighbor_remap(self.ux_obj, destination_grid, remap_to)
 
     def inverse_distance_weighted(
-        self, destination_grid: Grid, remap_to: str = "faces", power=2, k=8, **kwargs
+        self,
+        destination_grid: Grid,
+        remap_to: str = "faces",
+        power=2,
+        k=8,
+        backend: str = "uxarray",
+        yac_method: str | None = None,
+        yac_options: dict | None = None,
+        **kwargs,
     ) -> UxDataArray | UxDataset:
         """
         Perform inverse-distance-weighted (IDW) remapping.
@@ -80,18 +122,39 @@ def inverse_distance_weighted(
         k : int, default=8
             Number of nearest source points to include in the weighted average.
 
+        backend : {'uxarray', 'yac'}, default='uxarray'
+            Remapping backend to use. When set to 'yac', requires YAC to be
+            available on PYTHONPATH.
+        yac_method : {'nnn', 'conservative'}, optional
+            YAC interpolation method. Required when backend='yac'.
+        yac_options : dict, optional
+            YAC interpolation configuration options.
+
         Returns
         -------
         UxDataArray or UxDataset
             A new object with data mapped onto `destination_grid`.
         """
 
+        if backend == "yac":
+            from uxarray.remap.yac import _yac_remap
+
+            yac_kwargs = yac_options or {}
+            return _yac_remap(
+                self.ux_obj, destination_grid, remap_to, yac_method, yac_kwargs
+            )
         return _inverse_distance_weighted_remap(
             self.ux_obj, destination_grid, remap_to, power, k
         )
 
     def bilinear(
-        self, destination_grid: Grid, remap_to: str = "faces", **kwargs
+        self,
+        destination_grid: Grid,
+        remap_to: str = "faces",
+        backend: str = "uxarray",
+        yac_method: str | None = None,
+        yac_options: dict | None = None,
+        **kwargs,
     ) -> UxDataArray | UxDataset:
         """
         Perform bilinear remapping.
@@ -103,10 +166,25 @@ def bilinear(
         remap_to : {'nodes', 'edges', 'faces'}, default='faces'
             Which grid element receives the remapped values.
 
+        backend : {'uxarray', 'yac'}, default='uxarray'
+            Remapping backend to use. When set to 'yac', requires YAC to be
+            available on PYTHONPATH.
+        yac_method : {'nnn', 'conservative'}, optional
+            YAC interpolation method. Required when backend='yac'.
+        yac_options : dict, optional
+            YAC interpolation configuration options.
+
         Returns
         -------
         UxDataArray or UxDataset
             A new object with data mapped onto `destination_grid`.
         """
 
+        if backend == "yac":
+            from uxarray.remap.yac import _yac_remap
+
+            yac_kwargs = yac_options or {}
+            return _yac_remap(
+                self.ux_obj, destination_grid, remap_to, yac_method, yac_kwargs
+            )
         return _bilinear(self.ux_obj, destination_grid, remap_to)