Add reusable remap weight support

docs/api.rst

@@ -22,6 +22,7 @@ Top Level Functions
    open_multigrid
    open_mfdataset
    concat
+   load_remap_weights
 
 Tutorial
 --------
@@ -464,6 +465,15 @@ Remapping
 .. seealso::
 
     `Remapping User Guide Section <https://uxarray.readthedocs.io/en/latest/user-guide/remapping.html>`_
+    `Applying External Remap Weights <https://uxarray.readthedocs.io/en/latest/user-guide/remap-weights.html>`_
+
+Helpers
+~~~~~~~
+
+.. autosummary::
+   :toctree: generated/
+
+   RemapWeights
 
 
 UxDataArray
@@ -474,6 +484,7 @@ UxDataArray
    :template: autosummary/accessor_method.rst
 
    UxDataArray.remap
+   UxDataArray.remap.apply_weights
    UxDataArray.remap.nearest_neighbor
    UxDataArray.remap.inverse_distance_weighted
    UxDataArray.remap.bilinear
@@ -487,6 +498,7 @@ UxDataset
    :template: autosummary/accessor_method.rst
 
    UxDataset.remap
+   UxDataset.remap.apply_weights
    UxDataset.remap.nearest_neighbor
    UxDataset.remap.inverse_distance_weighted
    UxDataset.remap.bilinear

docs/user-guide/remap-weights.rst

@@ -0,0 +1,119 @@
+.. currentmodule:: uxarray
+
+Remap Weights
+=============
+
+UXarray can apply precomputed offline remapping weights produced outside of UXarray.
+This is useful when weights are generated once with tools such as ESMF or
+TempestRemap and then reused many times across multiple ensemble members, time
+slices, or variables.
+
+The core workflow is:
+
+1. Generate a weight file for a specific source grid and destination grid.
+2. Load the weight file once with :func:`load_remap_weights`.
+3. Reuse the loaded :class:`RemapWeights` object with :meth:`UxDataArray.remap.apply_weights`
+   or :meth:`UxDataset.remap.apply_weights`.
+
+Basic Usage
+-----------
+
+.. code-block:: python
+
+   import uxarray as ux
+
+   src = ux.open_dataset("source_grid.nc", "source_data.nc")
+   dst = ux.open_grid("destination_grid.nc")
+
+   weights = ux.load_remap_weights("map.nc")
+
+   remapped_temperature = src["temperature"].remap.apply_weights(
+       weights, destination_grid=dst
+   )
+
+   remapped_dataset = src.remap.apply_weights(weights, destination_grid=dst)
+
+Repeated calls with the same path reuse a cached sparse operator, so loading the
+same file again in one Python session avoids rebuilding the matrix.
+
+What A Weight File Represents
+-----------------------------
+
+A remap weight file represents a linear operator from one grid to another:
+
+.. code-block:: text
+
+   target_values = W @ source_values
+
+If the source grid has ``4800`` elements and the destination grid has ``11000``
+elements, then:
+
+- ``source_values.shape = (4800,)``
+- ``W.shape = (11000, 4800)``
+- ``target_values.shape = (11000,)``
+
+So the weight file necessarily encodes both the source grid and the destination
+grid. It is specific to that grid pair and to the ordering of the source and
+destination degrees of freedom.
+
+Supported File Structure
+------------------------
+
+UXarray currently supports the standard sparse offline-map structure used by
+ESMF-style and TempestRemap-style map files. The essential pieces are:
+
+- ``n_a``: source size
+- ``n_b``: destination size
+- ``n_s``: number of nonzero entries
+- ``row``: destination indices
+- ``col``: source indices
+- ``S``: sparse weight values
+
+Common aliases are also accepted:
+
+- ``src_grid_size`` and ``dst_grid_size``
+- ``src_address`` and ``dst_address``
+- ``weights`` instead of ``S``
+
+In full offline map files, these sparse arrays are typically accompanied by
+source and destination metadata such as center coordinates, corner coordinates,
+areas, masks, and grid-dimension metadata.
+
+Tool Compatibility
+------------------
+
+This implementation was verified against real files from both families:
+
+- ESMF-generated offline map files created with ``ESMF_RegridWeightGen``
+- TempestRemap-generated offline map files created with ``GenerateOfflineMap``
+
+In practice, UXarray supports the standard full offline map format used by both
+tools.
+
+Current caveats:
+
+- The source data ordering must match the source ordering encoded in the weight file.
+- Not every possible file variant is guaranteed yet.
+- ESMF ``weight_only`` outputs may require additional handling if they omit
+  source and destination size metadata.
+
+How It Applies Data
+-------------------
+
+When remapping a :class:`UxDataArray` or :class:`UxDataset`, UXarray identifies a
+single spatial dimension whose size matches the source size in the loaded
+weights. That dimension is remapped to the requested destination dimension
+(``faces``, ``edges``, or ``nodes``).
+
+Non-spatial dimensions are preserved, which makes this workflow suitable for
+reusing one operator across many time steps, ensemble members, or variables.
+
+Why Use This Workflow
+---------------------
+
+This path is useful when:
+
+- weight generation is expensive and should be done once
+- remapping needs to be repeated many times
+- external tools already produce trusted offline maps
+- you want to stay in Python for applying the map and preserving array metadata

docs/userguide.rst

@@ -64,6 +64,9 @@ These user guides provide detailed explanations of the core functionality in UXa
 `Remapping <user-guide/remapping.ipynb>`_
  Remap (a.k.a Regrid) between unstructured grids
 
+`Applying External Remap Weights <user-guide/remap-weights.rst>`_
+ Apply precomputed ESMF or TempestRemap offline map files
+
 `Topological Aggregations <user-guide/topological-aggregations.ipynb>`_
  Aggregate data across grid dimensions
 
@@ -119,6 +122,7 @@ These user guides provide additional details about specific features in UXarray.
    user-guide/zonal-average.ipynb
    user-guide/azimuthal-average.ipynb
    user-guide/remapping.ipynb
+   user-guide/remap-weights.rst
    user-guide/topological-aggregations.ipynb
    user-guide/weighted_mean.ipynb
    user-guide/vector_calculus.ipynb

test/precomputed_weights_test.py

@@ -0,0 +1,196 @@
+from pathlib import Path
+
+import numpy as np
+import numpy.testing as nt
+import pytest
+import uxarray as ux
+import xarray as xr
+
+from uxarray.remap.weights import (
+    _WEIGHTS_CACHE,
+    _WEIGHTS_CACHE_MAXSIZE,
+    _normalize_indices,
+)
+
+
+def _write_sparse_map(path: Path, source_size: int, destination_size: int) -> Path:
+    rows = np.arange(1, destination_size + 1, dtype=np.int32)
+    cols = np.arange(source_size, 0, -1, dtype=np.int32)
+    values = np.ones(destination_size, dtype=np.float64)
+
+    ds = xr.Dataset(
+        data_vars={
+            "row": (("n_s",), rows),
+            "col": (("n_s",), cols),
+            "S": (("n_s",), values),
+        },
+        coords={"n_s": np.arange(destination_size, dtype=np.int32)},
+    )
+    ds = ds.assign_coords(
+        n_a=np.arange(source_size, dtype=np.int32),
+        n_b=np.arange(destination_size, dtype=np.int32),
+    )
+    ds.to_netcdf(path)
+    return path
+
+
+def test_load_remap_weights_and_apply_vector(tmp_path, gridpath):
+    grid = ux.open_grid(gridpath("ugrid", "quad-hexagon", "grid.nc"))
+    weight_file = _write_sparse_map(
+        tmp_path / "reverse_map.nc", grid.n_face, grid.n_face
+    )
+
+    weights = ux.load_remap_weights(weight_file)
+    result = weights.apply(np.arange(grid.n_face, dtype=np.float64))
+
+    nt.assert_equal(weights.source_size, grid.n_face)
+    nt.assert_equal(weights.destination_size, grid.n_face)
+    nt.assert_array_equal(result, np.arange(grid.n_face, dtype=np.float64)[::-1])
+    assert isinstance(weights, ux.RemapWeights)
+
+
+def test_apply_weights_to_uxdataarray(tmp_path, gridpath):
+    grid = ux.open_grid(gridpath("ugrid", "quad-hexagon", "grid.nc"))
+    weight_file = _write_sparse_map(
+        tmp_path / "reverse_map.nc", grid.n_face, grid.n_face
+    )
+
+    source = ux.UxDataArray(
+        xr.DataArray(
+            np.arange(grid.n_face, dtype=np.float64),
+            dims=["n_face"],
+            name="temperature",
+            attrs={"units": "K"},
+        ),
+        uxgrid=grid,
+    )
+
+    remapped = source.remap.apply_weights(weight_file, grid)
+
+    nt.assert_array_equal(remapped.values, source.values[::-1])
+    nt.assert_equal(remapped.attrs["units"], "K")
+    nt.assert_equal(remapped.uxgrid, grid)
+
+
+def test_apply_weights_reuses_loaded_operator(tmp_path, gridpath):
+    grid = ux.open_grid(gridpath("ugrid", "quad-hexagon", "grid.nc"))
+    weight_file = _write_sparse_map(
+        tmp_path / "reverse_map.nc", grid.n_face, grid.n_face
+    )
+    weights = ux.load_remap_weights(weight_file)
+    cached_weights = ux.load_remap_weights(weight_file)
+
+    source = ux.UxDataset(
+        xr.Dataset(
+            data_vars={
+                "a": (
+                    ("time", "n_face"),
+                    np.arange(2 * grid.n_face).reshape(2, grid.n_face),
+                ),
+                "flag": (("time",), np.array([1, 0], dtype=np.int32)),
+            },
+            coords={"time": np.array([0, 1], dtype=np.int32)},
+        ),
+        uxgrid=grid,
+    )
+
+    remapped = source.remap.apply_weights(weights, grid)
+    remapped_again = source["a"].remap.apply_weights(weights, grid)
+
+    assert cached_weights is weights
+    nt.assert_array_equal(remapped["a"].values, source["a"].values[:, ::-1])
+    nt.assert_array_equal(remapped["flag"].values, source["flag"].values)
+    nt.assert_array_equal(remapped_again.values, source["a"].values[:, ::-1])
+
+
+def test_normalize_indices_respects_start_index_attr():
+    # 0-based array with an explicit start_index=0 attr — must not shift.
+    arr = xr.DataArray(np.array([0, 1, 2], dtype=np.int32), attrs={"start_index": 0})
+    nt.assert_array_equal(_normalize_indices(arr, 4, "Row"), np.array([0, 1, 2]))
+
+    # 1-based array with explicit start_index=1 attr.
+    arr1 = xr.DataArray(np.array([1, 2, 3], dtype=np.int32), attrs={"start_index": 1})
+    nt.assert_array_equal(_normalize_indices(arr1, 3, "Row"), np.array([0, 1, 2]))
+
+
+def test_normalize_indices_partial_zero_based_not_shifted():
+    # 0-based partial coverage: min=1, max < size. Previous heuristic
+    # would have wrongly shifted to -1; new heuristic keeps as 0-based.
+    arr = np.array([1, 2, 3], dtype=np.int32)
+    nt.assert_array_equal(_normalize_indices(arr, 10, "Row"), arr)
+
+
+def test_normalize_indices_one_based_detected_by_max():
+    arr = np.array([1, 2, 3, 4], dtype=np.int32)
+    nt.assert_array_equal(
+        _normalize_indices(arr, 4, "Row"), np.array([0, 1, 2, 3])
+    )
+
+
+def test_normalize_indices_out_of_bounds_raises():
+    with pytest.raises(ValueError, match="out of bounds"):
+        _normalize_indices(np.array([-1, 0, 1]), 4, "Row")
+
+
+def test_apply_weights_rejects_non_spatial_source_dim(tmp_path, gridpath):
+    grid = ux.open_grid(gridpath("ugrid", "quad-hexagon", "grid.nc"))
+    weight_file = _write_sparse_map(
+        tmp_path / "reverse_map.nc", grid.n_face, grid.n_face
+    )
+
+    source = ux.UxDataArray(
+        xr.DataArray(
+            np.arange(grid.n_face, dtype=np.float64),
+            dims=["n_face"],
+            name="t",
+        ),
+        uxgrid=grid,
+    )
+
+    with pytest.raises(ValueError, match="not a spatial dimension"):
+        source.remap.apply_weights(weight_file, grid, source_dim="time")
+
+
+def test_apply_weights_preserves_aux_coords(tmp_path, gridpath):
+    grid = ux.open_grid(gridpath("ugrid", "quad-hexagon", "grid.nc"))
+    weight_file = _write_sparse_map(
+        tmp_path / "reverse_map.nc", grid.n_face, grid.n_face
+    )
+
+    nt_steps = 3
+    da = xr.DataArray(
+        np.arange(nt_steps * grid.n_face, dtype=np.float64).reshape(
+            nt_steps, grid.n_face
+        ),
+        dims=("time", "n_face"),
+        coords={
+            "time": np.array([10, 20, 30], dtype=np.int64),
+            "time_label": ("time", np.array(["a", "b", "c"])),
+        },
+        name="t",
+    )
+    source = ux.UxDataArray(da, uxgrid=grid)
+    remapped = source.remap.apply_weights(weight_file, grid)
+    assert "time_label" in remapped.coords
+    nt.assert_array_equal(remapped["time_label"].values, np.array(["a", "b", "c"]))
+
+
+def test_clear_remap_weights_cache(tmp_path, gridpath):
+    grid = ux.open_grid(gridpath("ugrid", "quad-hexagon", "grid.nc"))
+    weight_file = _write_sparse_map(
+        tmp_path / "reverse_map.nc", grid.n_face, grid.n_face
+    )
+    ux.load_remap_weights(weight_file)
+    assert len(_WEIGHTS_CACHE) > 0
+    ux.clear_remap_weights_cache()
+    assert len(_WEIGHTS_CACHE) == 0
+
+
+def test_remap_weights_cache_is_lru_bounded(tmp_path, gridpath):
+    grid = ux.open_grid(gridpath("ugrid", "quad-hexagon", "grid.nc"))
+    ux.clear_remap_weights_cache()
+    for i in range(_WEIGHTS_CACHE_MAXSIZE + 5):
+        path = tmp_path / f"map_{i}.nc"
+        _write_sparse_map(path, grid.n_face, grid.n_face)
+        ux.load_remap_weights(path)
+    assert len(_WEIGHTS_CACHE) == _WEIGHTS_CACHE_MAXSIZE

uxarray/__init__.py

@@ -11,6 +11,7 @@
 from .core.dataarray import UxDataArray
 from .core.dataset import UxDataset
 from .grid import Grid
+from .remap import RemapWeights, clear_remap_weights_cache, load_remap_weights
 
 try:
     from importlib.metadata import version as _version
@@ -37,4 +38,7 @@
     "INT_DTYPE",
     "INT_FILL_VALUE",
     "Grid",
+    "RemapWeights",
+    "load_remap_weights",
+    "clear_remap_weights_cache",
 )