Update NaN mask handling for FSS

PiperOrigin-RevId: 869895126

Author: WeatherBenchX authors

weatherbenchX/metrics/metrics_test.py

@@ -227,6 +227,41 @@ def test_fss(self):
     ])
     np.testing.assert_allclose(out, correct_result)
 
+    # Test SPF and STF with NaNs, n=1
+    prediction1 = xr.DataArray(
+        [[1, 0, np.nan, 1]],
+        dims=['latitude', 'longitude'],
+        name='precipitation',
+    ).to_dataset()
+    target1 = xr.DataArray(
+        [[0, np.nan, 1, 0]],
+        dims=['latitude', 'longitude'],
+        name='precipitation',
+    ).to_dataset()
+    metrics1 = {
+        'fss': spatial.FSS(neighborhood_size_in_pixels=1),
+    }
+    stats1 = metrics_base.compute_unique_statistics_for_all_metrics(
+        metrics1, prediction1, target1
+    )
+    spf_stat = metrics1['fss'].statistics['SquaredPredictionFraction']
+    stf_stat = metrics1['fss'].statistics['SquaredTargetFraction']
+
+    # SPF: predictions are masked where targets are NaN.
+    # prediction becomes [1, 0, nan, 1].where(~[F,T,F,F]) = [1, nan, nan, 1]
+    # Neighborhood size 1: SPF = prediction^2 = [1, nan, nan, 1]
+    np.testing.assert_allclose(
+        stats1[spf_stat.unique_name]['precipitation'].values,
+        [[1.0, np.nan, np.nan, 1.0]],
+    )
+    # STF: targets are masked where predictions are NaN.
+    # target becomes [0, nan, 1, 0].where(~[F,F,T,F]) = [0, nan, nan, 0]
+    # Neighborhood size 1: STF = target^2 = [0, nan, nan, 0]
+    np.testing.assert_allclose(
+        stats1[stf_stat.unique_name]['precipitation'].values,
+        [[0.0, np.nan, np.nan, 0.0]],
+    )
+
   def test_wrapped_metric(self):
     target = (
         test_utils.mock_prediction_data(

weatherbenchX/metrics/spatial.py

@@ -14,7 +14,7 @@
 """Spatial verification metrics."""
 
 import dataclasses
-from typing import Iterable, Mapping, Union
+from typing import Collection, Iterable, Mapping, Union
 import numpy as np
 from scipy import ndimage
 from weatherbenchX.metrics import base
@@ -99,6 +99,41 @@ def neighborhood_averaging(
     )
 
 
+def get_fss_mask(
+    predictions: xr.DataArray,
+    targets: xr.DataArray,
+    neighborhood_size: Union[int, Iterable[int]],
+    wrap_longitude: bool = False,
+) -> xr.DataArray:
+  """Get mask for FSS.
+
+  The mask is True for pixels where FSS is valid, based on neighborhood
+  averaging method used in FSS, which propagates NaNs and applies boundary
+  zeroing for non-wrap-around case.
+
+  If any of predictions or targets is NaN in a neighborhood of a pixel,
+  the neighborhood averaging for that pixel will result in NaN,
+  unless it's on a boundary that gets zeroed out when wrap_longitude=False.
+  This mask is True where neighborhood averaging doesn't produce NaN.
+  This matches masking logic in SquaredPredictionFraction and
+  SquaredTargetFraction.
+
+  Args:
+    predictions: Predictions DataArray.
+    targets: Targets DataArray.
+    neighborhood_size: Neighborhood size for convolution.
+    wrap_longitude: Whether to wrap longitude in convolution.
+
+  Returns:
+    Boolean mask DataArray.
+  """
+  masked_preds = predictions.where(~targets.isnull())
+  neighborhood_preds = neighborhood_averaging(
+      masked_preds, neighborhood_size, wrap_longitude
+  )
+  return ~neighborhood_preds.isnull()
+
+
 def get_suffix(
     neighborhood_size: Union[int, Iterable[int]],
     wrap_longitude: bool = False,
@@ -129,13 +164,21 @@ def _compute_per_variable(
       predictions: xr.DataArray,
       targets: xr.DataArray,
   ) -> xr.DataArray:
+    mask = get_fss_mask(
+        predictions,
+        targets,
+        self.neighborhood_size_in_pixels,
+        self.wrap_longitude,
+    )
     predictions = neighborhood_averaging(
         predictions, self.neighborhood_size_in_pixels, self.wrap_longitude
     )
     targets = neighborhood_averaging(
         targets, self.neighborhood_size_in_pixels, self.wrap_longitude
     )
-    return (predictions - targets) ** 2
+    result = (predictions - targets) ** 2
+    result = result.assign_coords(mask=mask)
+    return result
 
 
 @dataclasses.dataclass
@@ -155,10 +198,18 @@ def _compute_per_variable(
       predictions: xr.DataArray,
       targets: xr.DataArray,
   ) -> xr.DataArray:
+    mask = get_fss_mask(
+        predictions,
+        targets,
+        self.neighborhood_size_in_pixels,
+        self.wrap_longitude,
+    )
     predictions = neighborhood_averaging(
         predictions, self.neighborhood_size_in_pixels, self.wrap_longitude
     )
-    return predictions**2 + xr.zeros_like(targets)
+    result = predictions**2
+    result = result.assign_coords(mask=mask)
+    return result
 
 
 @dataclasses.dataclass
@@ -178,10 +229,18 @@ def _compute_per_variable(
       predictions: xr.DataArray,
       targets: xr.DataArray,
   ) -> xr.DataArray:
+    mask = get_fss_mask(
+        predictions,
+        targets,
+        self.neighborhood_size_in_pixels,
+        self.wrap_longitude,
+    )
     targets = neighborhood_averaging(
         targets, self.neighborhood_size_in_pixels, self.wrap_longitude
     )
-    return targets**2 + xr.zeros_like(predictions)
+    result = targets**2
+    result = result.assign_coords(mask=mask)
+    return result
 
 
 @dataclasses.dataclass