Harden code and tests for robustness and memory management

README.md

@@ -1,6 +1,6 @@
 # XRegrid
 
-An optimized ESMF-based regridder for xarray that provides significant performance improvements over xESMF.
+An optimized ESMF-based regridder for xarray that provides significant performance improvements.
 
 ## Overview
 

src/xregrid/cli.py

@@ -9,7 +9,15 @@
 from xregrid.utils import get_rdhpcs_cluster
 
 
-def parse_args():
+def parse_args() -> argparse.Namespace:
+    """
+    Parse command-line arguments for the xregrid CLI.
+
+    Returns
+    -------
+    argparse.Namespace
+        The parsed command-line arguments.
+    """
     parser = argparse.ArgumentParser(description="xregrid CLI: Regrid NetCDF files.")
     parser.add_argument("src", help="Path to the source NetCDF file.")
     parser.add_argument(
@@ -69,7 +77,13 @@ def parse_args():
     return parser.parse_args()
 
 
-def main():
+def main() -> None:
+    """
+    Main entry point for the xregrid CLI.
+
+    Performs regridding of a source NetCDF file to a target grid and saves the result.
+    Supports local and distributed Dask clusters for parallel processing.
+    """
     args = parse_args()
 
     # 1. Setup Dask Client if requested

src/xregrid/core.py

@@ -28,6 +28,26 @@ def _setup_worker_cache(key: str, value: Any) -> None:
     _WORKER_CACHE[key] = value
 
 
+def _remove_from_worker_cache(key_pattern: str) -> int:
+    """
+    Remove all keys matching a pattern from the worker-local cache.
+
+    Parameters
+    ----------
+    key_pattern : str
+        The pattern (substring) to match in keys.
+
+    Returns
+    -------
+    int
+        Number of keys removed.
+    """
+    keys_to_remove = [k for k in _WORKER_CACHE.keys() if key_pattern in k]
+    for k in keys_to_remove:
+        del _WORKER_CACHE[k]
+    return len(keys_to_remove)
+
+
 def _matmul(matrix: Any, data: np.ndarray) -> np.ndarray:
     """
     Backend-agnostic matrix multiplication (matrix @ data.T).T.
@@ -122,6 +142,11 @@ def _apply_weights_core(
     else:
         flat_data = data_block.reshape(n_other, n_spatial)
 
+    # Robustness: Handle empty or all-NaN input arrays
+    if n_spatial == 0 or n_other == 0:
+        new_shape = other_dims_shape + shape_target
+        return np.full(new_shape, np.nan, dtype=data_block.dtype)
+
     if skipna:
         # Use a more memory-efficient NaN detection
         mask = np.isnan(flat_data)

src/xregrid/grid.py

@@ -877,7 +877,13 @@ def _create_esmf_grid(
                     mask_arg = None
                     if mask_var and mask_var in ds:
                         if method == "conservative":
-                            mask_val = ds[mask_var].values
+                            v_mask = ds[mask_var]
+                            mask_isel = {
+                                d: 0 for d in non_spatial_dims if d in v_mask.dims
+                            }
+                            if mask_isel:
+                                v_mask = v_mask.isel(mask_isel, drop=True)
+                            mask_val = v_mask.values
                             element_mask = mask_val[orig_idx].astype(np.int32)
                             mask_arg = element_mask
 

src/xregrid/parallel.py

@@ -264,7 +264,17 @@ def _compute_chunk_weights(
             regrid_kwargs["norm_type"] = esmpy.NormType.FRACAREA
 
         # 4. Generate weights
-        regrid = esmpy.Regrid(src_field, dst_field, **regrid_kwargs)
+        try:
+            regrid = esmpy.Regrid(src_field, dst_field, **regrid_kwargs)
+        except Exception as e:
+            # For workers, we return the error string in the result tuple
+            return (
+                np.array([]),
+                np.array([]),
+                np.array([]),
+                f"Regrid initialization error: {str(e)}",
+            )
+
         weights = regrid.get_weights_dict(deep_copy=True)
 
         # 5. Dask Resource Hygiene: Destroy temporary ESMF objects

src/xregrid/regridder.py

@@ -580,7 +580,22 @@ def _generate_weights(self) -> None:
             regrid_kwargs["norm_type"] = esmpy.NormType.FRACAREA
 
         # Build Regrid object
-        regrid = esmpy.Regrid(src_field, dst_field, **regrid_kwargs)
+        try:
+            regrid = esmpy.Regrid(src_field, dst_field, **regrid_kwargs)
+        except Exception as e:
+            msg = str(e)
+            if "ESMC_RC_ARG_OUTOFRANGE" in msg:
+                raise ValueError(
+                    "ESMF Argument out of range. This often happens if latitudes "
+                    "are not in [-90, 90] or if periodic grids have an extent of exactly 360 degrees."
+                ) from e
+            elif "ESMC_RC_GRID_PARTITION" in msg:
+                raise RuntimeError(
+                    "ESMF Grid partition error. Check for extremely small or degenerate grid cells."
+                ) from e
+            raise RuntimeError(
+                f"ESMPy failed to initialize Regrid object: {msg}"
+            ) from e
 
         # Explicit check for overlaps
         fl, fil = regrid.get_factors()
@@ -998,11 +1013,68 @@ def weights(self) -> csr_matrix:
 
         if hasattr(self._weights_matrix, "key"):
             self._weights_matrix = self._dask_client.gather(self._weights_matrix)
-            if hasattr(self._total_weights, "key"):
+            if self._total_weights is not None and hasattr(self._total_weights, "key"):
                 self._total_weights = self._dask_client.gather(self._total_weights)
 
         return self._weights_matrix
 
+    @classmethod
+    def clear_cache(cls) -> None:
+        """
+        Clear the global driver and worker caches for all Regridder instances.
+
+        This can be used to manually free memory when multiple Regridder
+        instances have been created in a long-running session.
+        """
+        global _DRIVER_CACHE
+        _DRIVER_CACHE.clear()
+
+        # Clear worker-local cache on all connected workers
+        try:
+            import dask.distributed
+
+            client = dask.distributed.get_client()
+            if client is not None:
+
+                def _clear_worker_cache():
+                    """Internal helper to clear the worker-local cache."""
+                    from xregrid.core import _WORKER_CACHE
+
+                    _WORKER_CACHE.clear()
+
+                client.run(_clear_worker_cache)
+        except (ImportError, ValueError):
+            pass
+
+    def clear_instance_cache(self) -> None:
+        """
+        Clear the worker-local cache entries specific to this Regridder instance.
+        """
+        if not self.parallel:
+            return
+
+        try:
+            import dask.distributed
+
+            client = self._dask_client or dask.distributed.get_client()
+            if client is not None:
+                from xregrid.core import _remove_from_worker_cache
+
+                # Remove by UID pattern
+                client.run(_remove_from_worker_cache, self._uid)
+
+                # Also clear driver cache for this instance
+                client_id = getattr(client, "id", id(client))
+                keys_to_remove = [
+                    k
+                    for k in _DRIVER_CACHE.keys()
+                    if k[0] == client_id and self._uid in k[1]
+                ]
+                for k in keys_to_remove:
+                    del _DRIVER_CACHE[k]
+        except (ImportError, ValueError):
+            pass
+
     def diagnostics(self) -> xr.Dataset:
         """
         Generate spatial diagnostics of the regridding weights.
@@ -1268,6 +1340,16 @@ def weights_to_xarray(self) -> xr.Dataset:
         )
         return ds
 
+    def __del__(self) -> None:
+        """
+        Cleanup instance-specific cache on deletion.
+        """
+        try:
+            if hasattr(self, "parallel") and self.parallel:
+                self.clear_instance_cache()
+        except Exception:
+            pass
+
     def __repr__(self) -> str:
         """
         String representation of the Regridder.

tests/test_aero_hardening.py

@@ -0,0 +1,99 @@
+import numpy as np
+import pytest
+import xarray as xr
+from xregrid import Regridder, create_global_grid
+
+
+def test_protocol_code_smells():
+    """
+    Verify that calling Regridder on a lazy DataArray does not trigger
+    immediate computation of the data.
+    """
+    try:
+        import dask.array as da
+    except ImportError:
+        pytest.skip("Dask not installed")
+
+    src = create_global_grid(10, 10)
+    tgt = create_global_grid(5, 5)
+
+    # Check laziness: use a dask array with a delayed function that increments a counter
+    from dask.delayed import delayed
+
+    counter = [0]
+
+    @delayed
+    def count_calls(x):
+        counter[0] += 1
+        return x
+
+    data = da.from_delayed(
+        count_calls(np.random.rand(18, 36)), shape=(18, 36), dtype=float
+    )
+    da_lazy = xr.DataArray(
+        data, dims=("lat", "lon"), coords={"lat": src.lat, "lon": src.lon}
+    )
+
+    regridder = Regridder(src, tgt)
+
+    # This should NOT trigger count_calls
+    res = regridder(da_lazy)
+
+    assert counter[0] == 0, "Regridding triggered immediate computation!"
+
+    # Computing the result SHOULD trigger it
+    _ = res.compute()
+    assert counter[0] > 0, "Computation did not trigger the delayed function!"
+
+
+def test_extreme_coordinate_values():
+    """Verify handling of coordinates slightly outside standard ranges."""
+    # ESMF often fails if lat is exactly 90.000000000001
+    # Our _clip_latitudes should handle this.
+    src_lat = np.array([-90.000001, 0, 90.000001])
+    src_lon = np.array([-0.000001, 180, 360.000001])
+
+    src = xr.Dataset(coords={"lat": (["lat"], src_lat), "lon": (["lon"], src_lon)})
+    src.lat.attrs["units"] = "degrees_north"
+    src.lon.attrs["units"] = "degrees_east"
+
+    tgt = create_global_grid(30, 30)
+
+    # Should not raise ESMC_RC_ARG_OUTOFRANGE
+    regridder = Regridder(src, tgt, method="bilinear")
+    assert regridder is not None
+
+
+def test_multiple_regridders_cache_isolation():
+    """Verify that multiple regridder instances don't interfere via cache."""
+    from conftest import setup_esmpy_mock
+    from distributed import Client, LocalCluster
+
+    with LocalCluster(n_workers=2, threads_per_worker=1) as cluster:
+        with Client(cluster) as client:
+            client.run(setup_esmpy_mock)
+
+            src = create_global_grid(10, 10)
+            tgt1 = create_global_grid(5, 5)
+            tgt2 = create_global_grid(2, 2)
+
+            r1 = Regridder(src, tgt1, parallel=True)
+            r2 = Regridder(src, tgt2, parallel=True)
+
+            assert r1._uid != r2._uid
+
+            da = xr.DataArray(
+                np.random.rand(18, 36),
+                dims=("lat", "lon"),
+                coords={"lat": src.lat, "lon": src.lon},
+            )
+
+            res1 = r1(da)
+            res2 = r2(da)
+
+            assert res1.shape == (36, 72)
+            assert res2.shape == (90, 180)
+
+
+if __name__ == "__main__":
+    pytest.main([__file__])

tests/test_aero_hardening_verification.py

@@ -0,0 +1,60 @@
+import numpy as np
+import pytest
+import xarray as xr
+from xregrid import Regridder, create_global_grid
+
+
+def test_empty_input_robustness():
+    """Verify that Regridder handles zero-sized input dimensions."""
+    src = create_global_grid(10, 10)
+    tgt = create_global_grid(5, 5)
+
+    # Create an empty DataArray along a non-spatial dimension
+    data = np.zeros((0, 18, 36))
+    da = xr.DataArray(
+        data,
+        dims=("time", "lat", "lon"),
+        coords={"lat": src.lat, "lon": src.lon, "time": []},
+    )
+
+    regridder = Regridder(src, tgt)
+    res = regridder(da)
+
+    assert res.shape == (0, 36, 72)
+    assert res.dtype == da.dtype
+
+
+def test_all_nan_input_robustness():
+    """Verify that Regridder handles all-NaN input arrays."""
+    src = create_global_grid(10, 10)
+    tgt = create_global_grid(5, 5)
+
+    data = np.full((18, 36), np.nan)
+    da = xr.DataArray(
+        data, dims=("lat", "lon"), coords={"lat": src.lat, "lon": src.lon}
+    )
+
+    regridder = Regridder(src, tgt, skipna=True)
+    res = regridder(da)
+
+    assert np.isnan(res.values).all()
+
+
+def test_cache_clearing():
+    """Verify Regridder cache clearing methods exist and don't crash."""
+    src = create_global_grid(10, 10)
+    tgt = create_global_grid(5, 5)
+    regridder = Regridder(src, tgt)
+
+    # Test class method
+    Regridder.clear_cache()
+
+    # Test instance method
+    regridder.clear_instance_cache()
+
+    # Check that deletion doesn't crash (triggers __del__)
+    del regridder
+
+
+if __name__ == "__main__":
+    pytest.main([__file__])