diff --git a/src/xregrid/accessors.py b/src/xregrid/accessors.py index 5604733..754d370 100644 --- a/src/xregrid/accessors.py +++ b/src/xregrid/accessors.py @@ -93,6 +93,55 @@ def plot_diagnostics( regridder = self.get_regridder(target_grid, **kwargs) return regridder.plot_diagnostics(mode=mode, **kwargs) + def plot_weights( + self, target_grid: xr.Dataset, row_idx: int, mode: str = "static", **kwargs: Any + ) -> Any: + """ + Visualize source points contributing to a specific destination point. + + Parameters + ---------- + target_grid : xr.Dataset + The target grid dataset. + row_idx : int + The index of the destination point. + mode : str, default 'static' + The plotting mode: 'static' or 'interactive'. + **kwargs : Any + Arguments passed to Regridder.plot_weights. + + Returns + ------- + Any + The plot object. + """ + regridder = self.get_regridder(target_grid, **kwargs) + return regridder.plot_weights(row_idx, mode=mode, **kwargs) + + def plot_comparison( + self, target_grid: xr.Dataset, mode: str = "static", **kwargs: Any + ) -> Any: + """ + Unified comparison plot (Source, Target, Difference). + + Parameters + ---------- + target_grid : xr.Dataset + The target grid dataset. + mode : str, default 'static' + The plotting mode: 'static' or 'interactive'. + **kwargs : Any + Arguments passed to Regridder.plot_comparison. + + Returns + ------- + Any + The plot object. + """ + regridder = self.get_regridder(target_grid, **kwargs) + da_tgt = regridder(self._obj) + return regridder.plot_comparison(self._obj, da_tgt, mode=mode, **kwargs) + @xr.register_dataset_accessor("regrid") class RegridDatasetAccessor: @@ -175,3 +224,59 @@ def plot_diagnostics( """ regridder = self.get_regridder(target_grid, **kwargs) return regridder.plot_diagnostics(mode=mode, **kwargs) + + def plot_weights( + self, target_grid: xr.Dataset, row_idx: int, mode: str = "static", **kwargs: Any + ) -> Any: + """ + Visualize source points contributing to a specific destination point. + + Parameters + ---------- + target_grid : xr.Dataset + The target grid dataset. + row_idx : int + The index of the destination point. + mode : str, default 'static' + The plotting mode: 'static' or 'interactive'. + **kwargs : Any + Arguments passed to Regridder.plot_weights. + + Returns + ------- + Any + The plot object. + """ + regridder = self.get_regridder(target_grid, **kwargs) + return regridder.plot_weights(row_idx, mode=mode, **kwargs) + + def plot_comparison( + self, + target_grid: xr.Dataset, + var_name: str, + mode: str = "static", + **kwargs: Any, + ) -> Any: + """ + Unified comparison plot (Source, Target, Difference) for a specific variable. + + Parameters + ---------- + target_grid : xr.Dataset + The target grid dataset. + var_name : str + The name of the variable to compare. + mode : str, default 'static' + The plotting mode: 'static' or 'interactive'. + **kwargs : Any + Arguments passed to Regridder.plot_comparison. + + Returns + ------- + Any + The plot object. + """ + regridder = self.get_regridder(target_grid, **kwargs) + da_src = self._obj[var_name] + da_tgt = regridder(da_src) + return regridder.plot_comparison(da_src, da_tgt, mode=mode, **kwargs) diff --git a/src/xregrid/regridder.py b/src/xregrid/regridder.py index 670b27f..03759f9 100644 --- a/src/xregrid/regridder.py +++ b/src/xregrid/regridder.py @@ -9,7 +9,7 @@ import xarray as xr from scipy.sparse import coo_matrix -from xregrid.utils import update_history, get_crs_info, is_lazy, is_dask, is_cubed +from xregrid.utils import update_history, get_crs_info, is_dask, is_cubed from xregrid.constants import ( get_regrid_method_map, get_extrap_method_map, @@ -2034,7 +2034,7 @@ def _detect_periodicity(self, ds: xr.Dataset) -> bool: True if the grid is detected as periodic. """ try: - from xregrid.utils import _find_coord + from xregrid.utils import _find_coord, _get_min_max_lazy_aware lon = _find_coord(ds, "longitude") if lon is not None: @@ -2042,20 +2042,13 @@ def _detect_periodicity(self, ds: xr.Dataset) -> bool: if lon.attrs.get("boundary") == "periodic": return True - # 2. Check eager values (dimension coordinates are eager in xarray) - lon_for_check = None - if not is_lazy(lon): - # Already in memory - lon_for_check = lon - elif lon.ndim == 1 and lon.name in lon.dims: - # Xarray dimension coordinates are usually eager even if data is lazy, - # but we check is_lazy(lon) above to be sure. - # If we are here, it means it reported as lazy. - lon_for_check = None - else: + # 2. Check eager values (prioritizing Xarray indexes) + lon_min, lon_max, is_eager = _get_min_max_lazy_aware(lon) + + if not is_eager: # Lazy 2D or non-dimension 1D coordinate. # Aero Protocol: Avoid hidden computes. - # We emit a warning if we are forced to compute to detect periodicity. + # We only trigger a compute if we are on the driver and it's absolutely necessary. import warnings warnings.warn( @@ -2063,25 +2056,18 @@ def _detect_periodicity(self, ds: xr.Dataset) -> bool: "To avoid this, provide 'periodic' explicitly in Regridder constructor " "or set the 'boundary' attribute to 'periodic' in your longitude coordinate." ) - try: - # Sample the first row to check extent cheaply without triggering a full compute. - if lon.ndim == 2: - lon_for_check = lon.isel({lon.dims[0]: 0}).compute() - else: - lon_for_check = lon.compute() - except Exception: - lon_for_check = None - - if lon_for_check is not None: - lon_min = float(lon_for_check.min()) - lon_max = float(lon_for_check.max()) - extent = lon_max - lon_min - # ESMF periodic grids must have extent strictly less than 360 - # because the periodicity is handled by connecting the last point to the first. - # If extent is 360, the last point is a duplicate of the first and ESMF will fail. - # Use a tighter bound (354 degrees) to avoid false positives for regional swaths. - if 354.0 <= extent < 360.0: - return True + import dask + + results = dask.compute({"min": lon_min, "max": lon_max})[0] + lon_min, lon_max = results["min"], results["max"] + + extent = float(lon_max) - float(lon_min) + # ESMF periodic grids must have extent strictly less than 360 + # because the periodicity is handled by connecting the last point to the first. + # If extent is 360, the last point is a duplicate of the first and ESMF will fail. + # Use a tighter bound (354 degrees) to avoid false positives for regional swaths. + if 354.0 <= extent < 360.0: + return True # 3. Last fallback: Check dimension name if "lon" in lon.dims or "longitude" in lon.dims: diff --git a/src/xregrid/utils.py b/src/xregrid/utils.py index d14aee8..287d07c 100644 --- a/src/xregrid/utils.py +++ b/src/xregrid/utils.py @@ -1130,6 +1130,61 @@ def create_rotated_latlon_grid( return ds +def _get_min_max_lazy_aware( + da_coord: xr.DataArray, +) -> Tuple[Union[float, Any], Union[float, Any], bool]: + """ + Helper to get min/max from a coordinate DataArray efficiently. + + Aero-Optimization: Prioritizes Xarray indexes for 1D dimension coordinates + (zero compute) and uses edge/corner sampling for 2D lazy coordinates to + minimize scan size. + + Parameters + ---------- + da_coord : xr.DataArray + The coordinate DataArray. + + Returns + ------- + min_val : float or dask.array.Item + The minimum value. + max_val : float or dask.array.Item + The maximum value. + is_eager : bool + True if the values were retrieved without triggering a Dask compute. + """ + # 1. Check if it's already in memory (dimension coordinates or NumPy-backed) + if not is_lazy(da_coord): + return float(da_coord.min()), float(da_coord.max()), True + + # 2. Check if it's a dimension coordinate in indexes + if ( + da_coord.ndim == 1 + and da_coord.name in da_coord.dims + and da_coord.name in da_coord.indexes + ): + idx = da_coord.indexes[da_coord.name] + return float(idx.min()), float(idx.max()), True + + # 3. Corner/Edge sampling heuristic for 2D lazy coordinates to avoid full scan + if da_coord.ndim == 2: + # Curvilinear grids are typically monotonic along edges. + # Sampling edges is much faster than a full array scan. + edges = [ + da_coord.isel({da_coord.dims[0]: 0}), + da_coord.isel({da_coord.dims[0]: -1}), + da_coord.isel({da_coord.dims[1]: 0}), + da_coord.isel({da_coord.dims[1]: -1}), + ] + # Use a dummy dimension for concatenation + combined_edges = xr.concat(edges, dim="_pts") + return combined_edges.min(), combined_edges.max(), False + + # Fallback: return Dask scalars for later batch compute + return da_coord.min(), da_coord.max(), False + + def create_grid_like( obj: Union[xr.DataArray, xr.Dataset], res: Union[float, Tuple[float, float]], @@ -1247,38 +1302,6 @@ def create_grid_like( # Discovery logic: we need min/max. We use batch compute if lazy to minimize roundtrips. # Aero-Optimization: Use Xarray indexes for 1D dimension coordinates to avoid hidden computes. - def _get_min_max_lazy_aware(da_coord): - """Helper to get min/max from a coordinate DataArray efficiently.""" - # 1. Check if it's already in memory (dimension coordinates or NumPy-backed) - if not is_lazy(da_coord): - return float(da_coord.min()), float(da_coord.max()), True - - # 2. Check if it's a dimension coordinate in indexes - if ( - da_coord.ndim == 1 - and da_coord.name in da_coord.dims - and da_coord.name in da_coord.indexes - ): - idx = da_coord.indexes[da_coord.name] - return float(idx.min()), float(idx.max()), True - - # 3. Corner/Edge sampling heuristic for 2D lazy coordinates to avoid full scan - if da_coord.ndim == 2: - # Curvilinear grids are typically monotonic along edges. - # Sampling edges is much faster than a full array scan. - edges = [ - da_coord.isel({da_coord.dims[0]: 0}), - da_coord.isel({da_coord.dims[0]: -1}), - da_coord.isel({da_coord.dims[1]: 0}), - da_coord.isel({da_coord.dims[1]: -1}), - ] - # Use a dummy dimension for concatenation - combined_edges = xr.concat(edges, dim="_pts") - return combined_edges.min(), combined_edges.max(), False - - # Fallback: return Dask scalars for later batch compute - return da_coord.min(), da_coord.max(), False - # 1. Try to find projected coordinates try: x_da = obj.cf["projection_x_coordinate"] diff --git a/tests/test_aero_refactor.py b/tests/test_aero_refactor.py new file mode 100644 index 0000000..3844386 --- /dev/null +++ b/tests/test_aero_refactor.py @@ -0,0 +1,133 @@ +from __future__ import annotations + +import numpy as np +import xarray as xr +from xregrid import Regridder +from xregrid.utils import create_global_grid, create_grid_like, is_lazy + + +def test_aero_periodicity_lazy(): + """ + Verify that periodicity detection works with lazy (Dask) coordinates. + The Aero Protocol: Refactored logic should avoid hidden computes for 1D dimension coords. + """ + # Create a global grid (eager) + ds_eager = create_global_grid(1.0, 1.0) + + # Create a lazy version + ds_lazy = ds_eager.chunk({"lat": 10, "lon": 10}) + + # Check periodicity of eager + regridder_eager = Regridder(ds_eager, ds_eager, method="bilinear") + assert regridder_eager.periodic is True + + # Check periodicity of lazy + # For 1D dimension coordinates, Xarray indexes are eager, so it should be eager. + regridder_lazy = Regridder(ds_lazy, ds_lazy, method="bilinear") + assert regridder_lazy.periodic is True + + +def test_create_grid_like_lazy(): + """ + Verify that create_grid_like works with lazy (Dask) coordinates and avoids hidden computes. + """ + # Use a predictable global grid as template + ds_template = create_global_grid(2.0, 2.0) + + # Lazy version + ds_lazy = ds_template.chunk({"lat": 10, "lon": 10}) + + # create_grid_like should use the edge-sampling heuristic or Xarray indexes. + # For 1D dimension coordinates, it should use indexes (zero compute). + grid_new = create_grid_like(ds_lazy, res=2.0) + + assert grid_new.lat.size == 90 + assert grid_new.lon.size == 180 + + # Check lineage in history + assert "Created grid like input" in grid_new.attrs["history"] + + +def test_create_grid_like_lazy_2d(): + """ + Verify that create_grid_like works with lazy 2D coordinates. + """ + # Create a template grid with 2D coordinates (curvilinear-like) + # Use an exact number of points to make extent calculation predictable + # res_lat = 180 / 45 = 4.0 + # res_lon = 360 / 90 = 4.0 + lon_2d, lat_2d = np.meshgrid(np.linspace(2, 358, 90), np.linspace(-88, 88, 45)) + ds_template = xr.Dataset( + coords={ + "lat": (["y", "x"], lat_2d, {"units": "degrees_north"}), + "lon": (["y", "x"], lon_2d, {"units": "degrees_east"}), + } + ) + + # Lazy version + ds_lazy = ds_template.chunk({"y": 10, "x": 10}) + + # create_grid_like should use the edge-sampling heuristic for 2D lazy coords. + grid_new = create_grid_like(ds_lazy, res=4.0) + + # Extent should be [-90, 90] for lat, [0, 360] for lon + assert grid_new.lat.size == 45 + assert grid_new.lon.size == 90 + + +def test_accessors_new_viz_methods(): + """ + Verify that the new visualization methods in accessors are present and callable. + (We mock the actual plot calls as they require GUI/heavy deps). + """ + ds = create_global_grid(10.0, 10.0) + + # Test DataArray accessor + da = ds["lat"] # Just a dummy variable + assert hasattr(da.regrid, "plot_weights") + assert hasattr(da.regrid, "plot_comparison") + + # Test Dataset accessor + assert hasattr(ds.regrid, "plot_weights") + assert hasattr(ds.regrid, "plot_comparison") + + # Note: We don't call them here because they would trigger matplotlib/cartopy/hvplot + # which might not be fully configured or might try to open windows. + # But we can verify they exist and have correct signatures. + import inspect + + sig = inspect.signature(da.regrid.plot_weights) + assert "target_grid" in sig.parameters + assert "row_idx" in sig.parameters + + sig = inspect.signature(ds.regrid.plot_comparison) + assert "target_grid" in sig.parameters + assert "var_name" in sig.parameters + + +def test_aero_protocol_numpy_vs_dask(): + """ + Double-Check Test: Run the logic on a NumPy array, then convert to Dask + and assert the result is identical. + """ + # 1. Implementation (Numpy) + ds_src = create_global_grid(10.0, 10.0) + ds_tgt = create_global_grid(5.0, 5.0) + + # Create some dummy data + data = np.random.rand(18, 36) + da_src_eager = xr.DataArray( + data, coords=[ds_src.lat, ds_src.lon], dims=["lat", "lon"] + ) + + regridder_eager = Regridder(ds_src, ds_tgt, method="bilinear") + res_eager = regridder_eager(da_src_eager) + + # 2. Implementation (Dask) + da_src_lazy = da_src_eager.chunk({"lat": 9, "lon": 9}) + regridder_lazy = Regridder(ds_src, ds_tgt, method="bilinear", parallel=False) + res_lazy = regridder_lazy(da_src_lazy) + + # Verify results are identical + xr.testing.assert_allclose(res_eager, res_lazy.compute()) + assert is_lazy(res_lazy)