Fix issue with masked timedelta values

Author: mx-moth

src/emsarray/conventions/ugrid.py

@@ -491,7 +491,11 @@ def _to_index_array(
             # If a data array has a fill value, xarray will convert that data array
             # to a floating point data type, and replace masked values with np.nan.
             # Here we convert a floating point array to a masked integer array.
-            values = np.ma.masked_invalid(values).astype(np.int_)
+            masked_values = np.ma.masked_invalid(values)
+            # numpy will emit a warning when converting an array with np.nan to int,
+            # even if the nans are masked out.
+            masked_values.data[masked_values.mask] = self.sensible_fill_value
+            values = masked_values.astype(self.sensible_dtype)
         elif '_FillValue' in data_array.attrs:
             # The data array has a fill value, but xarray has not applied it.
             # This implied the dataset was opened with mask_and_scale=False,
@@ -1195,13 +1199,17 @@ def apply_clip_mask(self, clip_mask: xr.Dataset, work_dir: Pathish) -> xr.Datase
         # any changes.
         topology_variables: List[xr.DataArray] = [topology.mesh_variable]
 
+        # This is the fill value used in the mask.
+        new_fill_value = clip_mask.data_vars['new_node_index'].encoding['_FillValue']
+
         def integer_indices(data_array: xr.DataArray) -> np.ndarray:
-            masked = np.ma.masked_invalid(data_array.values)
-            masked_integers: np.ndarray = masked.astype(np.int_)
+            masked_values = np.ma.masked_invalid(data_array.values)
+            # numpy will emit a warning when converting an array with np.nan to int,
+            # even if the nans are masked out.
+            masked_values.data[masked_values.mask] = new_fill_value
+            masked_integers: np.ndarray = masked_values.astype(np.int_)
             return masked_integers
 
-        # This is the fill value used in the mask.
-        new_fill_value = clip_mask.data_vars['new_node_index'].encoding['_FillValue']
         new_node_indices = integer_indices(clip_mask.data_vars['new_node_index'])
         new_face_indices = integer_indices(clip_mask.data_vars['new_face_index'])
         has_edges = 'new_edge_index' in clip_mask.data_vars

tests/conventions/test_ugrid.py

@@ -23,7 +23,7 @@
     ConventionViolationError, ConventionViolationWarning
 )
 from emsarray.operations import geometry
-from tests.utils import assert_property_not_cached
+from tests.utils import assert_property_not_cached, filter_warning
 
 
 def make_faces(width: int, height, fill_value: int) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
@@ -239,10 +239,28 @@ def make_dataset(
         },
     )
 
+    one_day = np.timedelta64(1, 'D').astype('timedelta64[ns]')
+    period = xr.DataArray(
+        data=np.concatenate([
+            np.arange(cell_size - 2, dtype=int) * one_day,
+            [np.timedelta64('nat', 'ns')] * 2,
+        ]),
+        dims=[face_dimension],
+        name="period",
+        attrs={
+            "long_name": "Some variable counting days",
+        },
+    )
+    period.encoding.update({
+        "units": "days",
+        "_FillValue": np.int16(-1),
+        "dtype": np.dtype('int16'),
+    })
+
     dataset = xr.Dataset(
         data_vars={var.name: var for var in [
             mesh, face_node_connectivity, node_x, node_y,
-            t, z, botz, eta, temp
+            t, z, botz, eta, temp, period,
         ]},
         attrs={
             'title': "COMPAS defalt version",
@@ -713,7 +731,22 @@ def test_apply_clip_mask(tmp_path):
 
 def test_make_and_apply_clip_mask(tmp_path):
     dataset = make_dataset(width=5)
-    dataset.ems.to_netcdf(tmp_path / "original.nc")
+
+    # When saving a dataset to disk, xarray.coding.times.cast_to_int_if_safe
+    # will check if it is possible to encode a timedelta64 using integer values
+    # by casting the values and checking for equality.
+    # Recent versions of numpy will emit warnings
+    # when casting a data array with dtype timedelta64 to int
+    # if it contains NaT (not a time) values.
+    # xarray will fix this eventually, but for now...
+    # See https://github.com/pydata/xarray/issues/7942
+    with filter_warning(
+        'ignore', category=RuntimeWarning,
+        message='invalid value encountered in cast',
+        module=r'xarray\.coding\.times',
+    ):
+        dataset.ems.to_netcdf(tmp_path / "original.nc")
+
     geometry.write_geojson(dataset, tmp_path / 'original.geojson')
 
     polygon = Polygon([[3.4, 1], [3.4, -1], [6, -1], [6, 1], [3.4, 1]])

tests/masking/test_utils.py

@@ -10,7 +10,7 @@
 
 from emsarray import masking
 from emsarray.utils import to_netcdf_with_fixes
-from tests.utils import mask_from_strings
+from tests.utils import filter_warning, mask_from_strings
 
 
 def assert_raw_values(
@@ -120,7 +120,14 @@ def test_find_fill_value_timedelta_with_missing_value(
         fill_value = masking.find_fill_value(data_array)
         assert np.isnat(fill_value)
 
-        to_netcdf_with_fixes(dataset, tmp_path / 'dataset.nc')
+        # See https://github.com/pydata/xarray/issues/7942
+        with filter_warning(
+            'ignore', category=RuntimeWarning,
+            message='invalid value encountered in cast',
+            module=r'xarray\.coding\.times',
+        ):
+            # Write this out for easier debugging purposes
+            to_netcdf_with_fixes(dataset, tmp_path / 'dataset.nc')
 
 
 def test_calculate_mask_bounds():

tests/test_utils.py

@@ -1,4 +1,5 @@
 import datetime
+import logging
 import pathlib
 
 import netCDF4
@@ -10,6 +11,9 @@
 import xarray.testing
 
 from emsarray import utils
+from tests.utils import filter_warning
+
+logger = logging.getLogger(__name__)
 
 
 @pytest.mark.parametrize(
@@ -76,8 +80,9 @@ def test_disable_default_fill_value(tmp_path: pathlib.Path):
 
     td_data = np.where(
         np.tri(5, 7, dtype=bool),
-        np.arange(35).reshape(5, 7) * np.timedelta64(1, 'D'),
-        np.timedelta64('nat'))
+        np.arange(35).reshape(5, 7) * np.timedelta64(1, 'D').astype('timedelta64[ns]'),
+        np.timedelta64('nat', 'ns'))
+    logger.info('%r', td_data)
     timedelta_with_missing_value_var = xarray.DataArray(
         data=td_data, dims=['j', 'i'])
     timedelta_with_missing_value_var.encoding['missing_value'] = np.float64('1e35')
@@ -91,7 +96,17 @@ def test_disable_default_fill_value(tmp_path: pathlib.Path):
     })
 
     # Save to a netCDF4 and then prove that it is bad
-    dataset.to_netcdf(tmp_path / "bad.nc")
+    # This emits warnings in current versions of xarray / numpy.
+    # See https://github.com/pydata/xarray/issues/7942
+    with filter_warning(
+        'default', category=RuntimeWarning,
+        message='invalid value encountered in cast',
+        module=r'xarray\.coding\.times',
+        record=True,
+    ) as ws:
+        dataset.to_netcdf(tmp_path / "bad.nc")
+        assert len(ws) == 1
+
     with netCDF4.Dataset(tmp_path / "bad.nc", "r") as nc_dataset:
         # This one shouldn't be here because it is an integer datatype. xarray
         # does the right thing already in this case.
@@ -105,7 +120,17 @@ def test_disable_default_fill_value(tmp_path: pathlib.Path):
         assert np.isnan(nc_dataset.variables["timedelta_with_missing_value_var"].getncattr("_FillValue"))
 
     utils.disable_default_fill_value(dataset)
-    dataset.to_netcdf(tmp_path / "good.nc")
+
+    # See https://github.com/pydata/xarray/issues/7942
+    with filter_warning(
+        'default', category=RuntimeWarning,
+        message='invalid value encountered in cast',
+        module=r'xarray\.coding\.times',
+        record=True,
+    ) as ws:
+        dataset.to_netcdf(tmp_path / "good.nc")
+        assert len(ws) == 1
+
     with netCDF4.Dataset(tmp_path / "good.nc", "r") as nc_dataset:
         # This one should still be unset
         assert '_FillValue' not in nc_dataset.variables["int_var"].ncattrs()

tests/utils.py

@@ -1,7 +1,9 @@
 from __future__ import annotations
 
 import abc
+import contextlib
 import itertools
+import warnings
 from functools import cached_property
 from typing import Any, Dict, Hashable, List, Optional, Tuple
 
@@ -14,6 +16,17 @@
 )
 
 
+@contextlib.contextmanager
+def filter_warning(*args, record: bool = False, **kwargs):
+    """
+    A shortcut wrapper around warnings.catch_warning()
+    and warnings.filterwarnings()
+    """
+    with warnings.catch_warnings(record=record) as context:
+        warnings.filterwarnings(*args, **kwargs)
+        yield context
+
+
 def box(minx, miny, maxx, maxy) -> shapely.Polygon:
     """
     Make a box, with coordinates going counterclockwise