diff --git a/README.md b/README.md index 846dc31..a11293c 100644 --- a/README.md +++ b/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 diff --git a/src/xregrid/cli.py b/src/xregrid/cli.py index 9043641..6d63fe1 100644 --- a/src/xregrid/cli.py +++ b/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 diff --git a/src/xregrid/core.py b/src/xregrid/core.py index 8ba3ce6..1699f29 100644 --- a/src/xregrid/core.py +++ b/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) diff --git a/src/xregrid/grid.py b/src/xregrid/grid.py index c1df27e..9c641d3 100644 --- a/src/xregrid/grid.py +++ b/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 diff --git a/src/xregrid/parallel.py b/src/xregrid/parallel.py index f28a75d..17f1242 100644 --- a/src/xregrid/parallel.py +++ b/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 diff --git a/src/xregrid/regridder.py b/src/xregrid/regridder.py index d15e08f..c7fd109 100644 --- a/src/xregrid/regridder.py +++ b/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. diff --git a/tests/test_aero_hardening.py b/tests/test_aero_hardening.py new file mode 100644 index 0000000..b77dd0e --- /dev/null +++ b/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__]) diff --git a/tests/test_aero_hardening_verification.py b/tests/test_aero_hardening_verification.py new file mode 100644 index 0000000..7935843 --- /dev/null +++ b/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__])