Enabling target mask for LocStream input (#445)

* Enabling target masks for LocStream input
- Adding function xesmf.smm.post_apply_target_mask_to_weights
  to zero out weights of masked target grid cells
- Adding application of this function to BaseRegridder in the case
  of LocStream input, Grid output incl. mask
- This may lead to unexpected results when not using with nearest_s2d,
  as the masking is applied *post* weight generation

* Restricting automatic application of new masking option to nearest_s2d

* Update new tests

* Updated changelog.

* Updated changelog.

* Pin cf_xarray as latest update seems to have introduced an issue to cfxr.bounds_to_vertices

* Pin cf_xarray for ci as latest update seems to have introduced an issue to cfxr.bounds_to_vertices

* unpin cfxr

---------

Co-authored-by: Pascal Bourgault <[email protected]>

Author: sol1105

CHANGES.rst

@@ -1,9 +1,10 @@
 What's new
 ==========
 
-0.8.11 (unreleased)
--------------------
-* ``xe.util.cf_grid_2d`` returns bounds as coordinates, as ``grid_2d`` does and as usually expected. (:pull:453`). `By `Pascal Bourgault <https://github.com/aulemahal>`_.
+0.9.0 (unreleased)
+------------------
+* Added support for target masks when regridding ``LocStream`` to ``Grid`` with ``nearest_s2d`` (:pull:`445`). By `Martin Schupfner <https://github.com/sol1105>`_.
+* ``xesmf.util.cf_grid_2d`` returns bounds as coordinates, as ``grid_2d`` does and as usually expected. (:pull:`453`). `By `Pascal Bourgault <https://github.com/aulemahal>`_.
 
 0.8.10 (2025-04-29)
 -------------------

xesmf/frontend.py

@@ -18,6 +18,7 @@
     add_nans_to_weights,
     apply_weights,
     check_shapes,
+    post_apply_target_mask_to_weights,
     read_weights,
 )
 from .util import LAT_CF_ATTRS, LON_CF_ATTRS, split_polygons_and_holes
@@ -376,6 +377,22 @@ def __init__(
             # Convert weights, whatever their format, to a sparse coo matrix
             self.weights = read_weights(weights, self.n_in, self.n_out)
 
+            # Optionally post-apply output mask for LocStream input and Grid output
+            # as xesmf.backend.esmf_regrid_build filters the output mask in that case
+            # (ESMF does not support output masks for LocStream input and Grid output)
+            # Only method supported is `nearest_s2d`:
+            #   For other methods the masking approach may lead to unexpected results
+            #   as the weights are applied post weight generation and other methods may have
+            #   the source-target mapping depending on the location of masked cells.
+            if (
+                isinstance(grid_in, LocStream)
+                and isinstance(grid_out, Grid)
+                and grid_out.mask is not None
+                and grid_out.mask[0] is not None
+                and method == 'nearest_s2d'
+            ):
+                self.weights = post_apply_target_mask_to_weights(self.weights, grid_out.mask[0])
+
             # replace zeros by NaN for weight matrix entries of unmapped target cells if specified or a mask is present
             if (
                 (grid_out.mask is not None) and (grid_out.mask[0] is not None)

xesmf/smm.py

@@ -201,6 +201,60 @@ def apply_weights(weights, indata, shape_in, shape_out):
     return outdata
 
 
+def post_apply_target_mask_to_weights(weights, target_mask_2d):
+    """
+    Set all contributions to masked target grid cells to zero.
+
+    Parameters
+    ----------
+    weights : DataArray backed by a sparse.COO array
+      Sparse weights matrix.
+    target_mask_2d : array-like
+        Mask array of shape (nx, ny) for the target grid.
+        False / 0 indicates a masked cell whose weights shall be set to zero.
+
+    Returns
+    -------
+    DataArray backed by a sparse.COO array
+      Modified sparse weights matrix with all rows corresponding to masked target cells set to zero.
+
+    Notes
+    -----
+    This defines a post-processing step applied after ESMF weight generation.
+    It is useful in cases where ESMF/ESMPy masks cannot be used directly,
+    which is the case when source or target are LocStream objects / sequences.
+    """
+    # Ensure mask can be converted to array
+    try:
+        target_mask_2d = np.asarray(target_mask_2d, dtype=weights.data.data.dtype)
+    except Exception as e:
+        raise TypeError(
+            "Argument 'target_mask_2d' must be array-like and convertible to a numeric/boolean array"
+        ) from e
+
+    # Validate dimensionality and shape
+    if target_mask_2d.ndim != 2:
+        raise ValueError(f"Argument 'target_mask_2d' must be 2D, got shape {target_mask_2d.shape}")
+    n_target, n_source = weights.data.shape
+    if target_mask_2d.size != n_target:
+        raise ValueError(
+            f"Mismatch: weight matrix has {n_target} target cells, "
+            f"but mask has {target_mask_2d.size} elements"
+        )
+
+    # Flatten mask array to align with weight matrix target index (Fortran order for ESMF layout)
+    target_mask_flat = target_mask_2d.ravel(order='F')
+
+    # Multiply row-wise by mask to zero out weights of masked target cells
+    W = weights.data * target_mask_flat[:, None]
+
+    # Create weights DataArray backed by sparse.COO
+    weights = xr.DataArray(
+        sps.COO(coords=W.coords, data=W.data, shape=W.shape), dims=('out_dim', 'in_dim')
+    )
+    return weights
+
+
 def add_nans_to_weights(weights):
     """Add nan in empty rows of the regridding weights sparse matrix.
 

xesmf/tests/test_frontend.py

@@ -1018,3 +1018,36 @@ def test_spatial_averager_mask():
     savg = xe.SpatialAverager(dsm, [poly], geom_dim_name='my_geom')
     out = savg(dsm.abc)
     assert_allclose(out, 2, rtol=1e-3)
+
+
+def test_locstream_input_grid_output_with_target_mask_applied():
+    # Create locstream input (6 coordinate points)
+    locstream_in = xr.Dataset(
+        {'var': xr.DataArray(np.ones((6)), dims=['location'])},
+        coords={
+            'lat': ('location', np.linspace(0, 5, 6)),
+            'lon': ('location', np.linspace(0, 10, 6)),
+        },
+    )
+
+    # Create Grid output with target mask (3x3 grid)
+    ds_out = xe.util.cf_grid_2d(0, 10, 10.0 / 3.0, 0, 5, 5.0 / 3.0)
+    target_mask_2d = np.ones((3, 3), dtype=bool)
+    target_mask_2d[-1, :] = False
+    ds_out['target_mask'] = xr.DataArray(target_mask_2d, dims=['lat', 'lon'])
+
+    # Create Grid output with target mask (3x3 grid)
+    ds_out = xe.util.cf_grid_2d(0, 10, 10.0 / 3.0, 0, 5, 5.0 / 3.0)
+    target_mask_2d = np.ones((3, 3), dtype=bool)
+    target_mask_2d[-1, :] = False
+    ds_out['mask'] = xr.DataArray(target_mask_2d, dims=['lat', 'lon'])
+
+    # Generate weights
+    regridder = xe.Regridder(
+        ds_in=locstream_in, ds_out=ds_out, method='nearest_s2d', locstream_in=True
+    )
+
+    # Apply weights and check results - the northmost cells should be masked
+    da_out = regridder(locstream_in)['var']
+    assert np.all(np.isnan(da_out[-1, :]))
+    assert np.all(da_out[:-1, :] == 1)

xesmf/tests/test_smm.py

@@ -1,4 +1,5 @@
 import numpy as np
+import pytest
 import sparse as sps
 import xarray as xr
 
@@ -27,3 +28,61 @@ def test_add_nans_to_weights():
 
     Matout = xe.smm.add_nans_to_weights(xr.DataArray(Matin, dims=('in', 'out')))
     assert np.allclose(Matin.todense(), Matout.data.todense())
+
+
+def test_post_apply_target_mask_to_weights():
+    # Create a small sparse weights matrix with shape (9 target, 4 source)
+    #  coords = [[target_indices], [source_indices]]
+    coords = np.array([[0, 1, 1, 2, 3, 3, 4, 5], [0, 0, 1, 1, 2, 3, 2, 3]])
+    data = np.array([0.1, 0.2, 0.3, 0.4, 0.5, 0.45, 0.7, 0.8])
+    shape = (6, 4)
+    W_sparse = sps.COO(coords, data, shape=shape)
+    weights = xr.DataArray(W_sparse, dims=('out_dim', 'in_dim'))
+
+    # Define a 3x3 mask for target (flattened size = 9):
+    #   If all goes to plan, weights of cells 3 and 4 (i.e. index 2 and 3)
+    #   will be set to 0.
+    target_mask_2d = np.array([[True, False], [True, True], [False, True]])
+
+    # Apply mask
+    masked_weights = xe.smm.post_apply_target_mask_to_weights(weights, target_mask_2d)
+
+    # Check results
+    np.testing.assert_array_equal(masked_weights.data.data, np.array([0.1, 0.2, 0.3, 0.7, 0.8]))
+    np.testing.assert_array_equal(
+        masked_weights.data.coords, np.array([[0, 1, 1, 4, 5], [0, 0, 1, 2, 3]])
+    )
+
+
+def test_post_apply_target_mask_to_weights_exceptions():
+    # Create a weights DataArray & mask
+    coords = np.array([[0, 1], [0, 1]])
+    data = np.array([0.5, 0.5])
+    shape = (2, 2)
+    W_sparse = sps.COO(coords, data, shape=shape)
+    weights = xr.DataArray(W_sparse, dims=('out_dim', 'in_dim'))
+    valid_mask = np.array([[True, False]])
+
+    # Mask not array-like
+    with pytest.raises(
+        TypeError,
+        match="Argument 'target_mask_2d' must be array-like and convertible to a numeric/boolean array",
+    ):
+        xe.smm.post_apply_target_mask_to_weights(weights, 'not_array_like')
+
+    # Shape mismatch
+    wrong_shape_mask = np.array([[True, False, True]])
+    with pytest.raises(
+        ValueError, match='Mismatch: weight matrix has 2 target cells, but mask has 3 elements'
+    ):
+        xe.smm.post_apply_target_mask_to_weights(weights, wrong_shape_mask)
+
+    # Mask not 2D
+    wrong_shape_mask = np.array([[[True]], [[True]]])
+    with pytest.raises(
+        ValueError, match="Argument 'target_mask_2d' must be 2D, got shape \\(2, 1, 1\\)"
+    ):
+        xe.smm.post_apply_target_mask_to_weights(weights, wrong_shape_mask)
+
+    # That should work
+    xe.smm.post_apply_target_mask_to_weights(weights, valid_mask)