diff --git a/src/xregrid/grid.py b/src/xregrid/grid.py index 11ce350..8561b61 100644 --- a/src/xregrid/grid.py +++ b/src/xregrid/grid.py @@ -177,8 +177,16 @@ def _get_mesh_info( # If they share same dim, it's unstructured if lat.dims == lon.dims: # Apply filtering before returning - lat_isel = {d: 0 for d in non_spatial_dims if d in lat.dims and len(lat.dims) > 1} - lon_isel = {d: 0 for d in non_spatial_dims if d in lon.dims and len(lon.dims) > 1} + lat_isel = { + d: 0 + for d in non_spatial_dims + if d in lat.dims and len(lat.dims) > 1 + } + lon_isel = { + d: 0 + for d in non_spatial_dims + if d in lon.dims and len(lon.dims) > 1 + } if lat_isel: lat = lat.isel(lat_isel, drop=True) if lon_isel: diff --git a/src/xregrid/regridder.py b/src/xregrid/regridder.py index 77f9b31..2ead3c0 100644 --- a/src/xregrid/regridder.py +++ b/src/xregrid/regridder.py @@ -15,6 +15,7 @@ is_dask, is_cubed, _get_min_max_lazy_aware, + _compute_lazy_aware, ) from xregrid.constants import ( get_regrid_method_map, @@ -2065,19 +2066,12 @@ def _detect_periodicity(self, ds: xr.Dataset) -> bool: "or set the 'boundary' attribute to 'periodic' in your longitude coordinate." ) try: - # Use xarray's compute to remain backend-agnostic - # if the tasks are xarray/dask objects. - if hasattr(lon_min_task, "compute"): - lon_min = lon_min_task.compute() - else: - lon_min = lon_min_task - - if hasattr(lon_max_task, "compute"): - lon_max = lon_max_task.compute() - else: - lon_max = lon_max_task - - lon_min, lon_max = float(lon_min), float(lon_max) + # Aero Protocol: Use centralized backend-agnostic compute + # to resolve sampled longitude edges. + res = _compute_lazy_aware( + {"min": lon_min_task, "max": lon_max_task} + ) + lon_min, lon_max = float(res["min"]), float(res["max"]) except Exception: lon_min = lon_max = None diff --git a/src/xregrid/utils.py b/src/xregrid/utils.py index 07c96b6..e2b433d 100644 --- a/src/xregrid/utils.py +++ b/src/xregrid/utils.py @@ -35,6 +35,8 @@ def is_cubed(obj: Any) -> bool: """ Check if an object is a cubed array or a Dataset/DataArray containing one. + Supports checking inside lists, tuples, and dictionaries. + Parameters ---------- obj : Any @@ -52,10 +54,20 @@ def is_cubed(obj: Any) -> bool: return True if isinstance(obj, xr.DataArray): - return isinstance(obj.data, cubed.Array) + return ( + hasattr(obj.data, "__array_namespace__") + and "cubed" in str(obj.data.__array_namespace__()) + or isinstance(obj.data, cubed.Array) + ) if isinstance(obj, xr.Dataset): - return any(isinstance(v.data, cubed.Array) for v in obj.data_vars.values()) + return any(is_cubed(v) for v in obj.data_vars.values()) + + if isinstance(obj, (list, tuple)): + return any(is_cubed(item) for item in obj) + + if isinstance(obj, dict): + return any(is_cubed(item) for item in obj.values()) return False @@ -108,6 +120,54 @@ def is_lazy(obj: Any) -> bool: return is_dask(obj) or is_cubed(obj) +def _compute_lazy_aware(obj: Any) -> Any: + """ + Compute lazy objects (Dask or Cubed) in a backend-agnostic way. + + Supports single objects, lists, and dictionaries. + + Parameters + ---------- + obj : Any + The object(s) to compute. + + Returns + ------- + Any + The computed object(s). + """ + # 1. Check for Cubed backend + if is_cubed(obj) and cubed is not None: + if isinstance(obj, dict): + keys = list(obj.keys()) + vals = [obj[k] for k in keys] + computed_vals = cubed.compute(*vals) + return dict(zip(keys, computed_vals)) + elif isinstance(obj, (list, tuple)): + res = cubed.compute(*obj) + return list(res) if isinstance(obj, list) else res + else: + return cubed.compute(obj)[0] + + # 2. Check for Dask backend + if is_dask(obj) and dask is not None: + # Dask handles dicts and lists natively + res = dask.compute(obj) + return res[0] if len(res) == 1 else res + + # 3. Fallback for objects with .compute() method (Backend-agnostic) + if hasattr(obj, "compute") and callable(obj.compute): + return obj.compute() + + # 4. Handle non-lazy containers that might contain objects with .compute() + if isinstance(obj, list): + return [_compute_lazy_aware(item) for item in obj] + if isinstance(obj, dict): + return {k: _compute_lazy_aware(v) for k, v in obj.items()} + + return obj + + def _get_array_namespace(*objs: Any) -> Any: """ Get the appropriate array namespace (numpy, dask.array, or cubed) for the given objects. @@ -1327,7 +1387,7 @@ def create_grid_like( tasks_dict["ymin"] = y_min tasks_dict["ymax"] = y_max - results = dask.compute(tasks_dict)[0] + results = _compute_lazy_aware(tasks_dict) extent = ( float(results.get("xmin", x_min)), float(results.get("xmax", x_max)), @@ -1361,7 +1421,7 @@ def create_grid_like( if y_da.size > 1: tasks_dict["res_y"] = abs(y_da.diff(y_da.dims[0]).mean()) - results = dask.compute(tasks_dict)[0] + results = _compute_lazy_aware(tasks_dict) x_min_val = float(results.get("x_min", x_min)) x_max_val = float(results.get("x_max", x_max)) y_min_val = float(results.get("y_min", y_min)) @@ -1436,7 +1496,7 @@ def create_grid_like( tasks_dict["lon_min"] = lon_min tasks_dict["lon_max"] = lon_max - results = dask.compute(tasks_dict)[0] + results = _compute_lazy_aware(tasks_dict) lat_range = ( float(results.get("lat_min", lat_min)), float(results.get("lat_max", lat_max)), @@ -1472,7 +1532,7 @@ def create_grid_like( if lon_da.size > 1: tasks_dict["res_lon"] = abs(lon_da.diff(lon_da.dims[-1]).mean()) - results = dask.compute(tasks_dict)[0] + results = _compute_lazy_aware(tasks_dict) lat_min_val = float(results.get("lat_min", lat_min)) lat_max_val = float(results.get("lat_max", lat_max)) lon_min_val = float(results.get("lon_min", lon_min)) diff --git a/tests/test_aero_compute_agnostic.py b/tests/test_aero_compute_agnostic.py new file mode 100644 index 0000000..881cc8c --- /dev/null +++ b/tests/test_aero_compute_agnostic.py @@ -0,0 +1,102 @@ +from __future__ import annotations + +import numpy as np +import pytest +from xregrid.utils import _compute_lazy_aware, is_lazy, is_dask + +try: + import dask.array as da + + HAS_DASK = True +except ImportError: + HAS_DASK = False + +import importlib.util + +HAS_DASK = importlib.util.find_spec("dask") is not None +HAS_CUBED = importlib.util.find_spec("cubed") is not None + + +def test_compute_lazy_aware_numpy(): + """Verify compute_lazy_aware handles eager NumPy data.""" + data = np.array([1, 2, 3]) + res = _compute_lazy_aware(data) + assert np.array_equal(res, data) + assert not is_lazy(res) + + +def test_compute_lazy_aware_dict_numpy(): + """Verify compute_lazy_aware handles dicts of eager data.""" + data = {"a": np.array([1]), "b": 2} + res = _compute_lazy_aware(data) + assert res == data + assert res["a"] is data["a"] + + +@pytest.mark.skipif(not HAS_DASK, reason="Dask not installed") +def test_compute_lazy_aware_dask(): + """Verify compute_lazy_aware handles Dask objects.""" + dask_arr = da.from_array(np.array([1, 2, 3]), chunks=2) + assert is_dask(dask_arr) + + res = _compute_lazy_aware(dask_arr) + assert np.array_equal(res, np.array([1, 2, 3])) + assert not is_lazy(res) + + +@pytest.mark.skipif(not HAS_DASK, reason="Dask not installed") +def test_compute_lazy_aware_dask_dict(): + """Verify compute_lazy_aware handles dicts of Dask objects.""" + dask_arr = da.from_array(np.array([1]), chunks=1) + data = {"a": dask_arr, "b": 2} + + res = _compute_lazy_aware(data) + assert res["a"] == 1 + assert res["b"] == 2 + assert not is_lazy(res["a"]) + + +@pytest.mark.skipif(not HAS_CUBED, reason="Cubed not installed") +def test_compute_lazy_aware_cubed_mock(monkeypatch): + """Verify compute_lazy_aware handles (mocked) Cubed objects.""" + # We can't easily create a real Cubed array without a spec/plan, + # so we mock the cubed.compute call and is_cubed check. + + data_val = np.array([10]) + + def mock_cubed_compute(*args, **kwargs): + return [data_val for _ in args] + + monkeypatch.setattr("cubed.compute", mock_cubed_compute) + + # Create an object that passes is_cubed + # In our implementation, is_cubed checks isinstance(obj, cubed.Array) + # or DataArray.data is cubed.Array + + class FakeCubed: + pass + + monkeypatch.setattr("cubed.Array", FakeCubed) + + obj = FakeCubed() + # Ensure is_cubed returns True + monkeypatch.setattr("xregrid.utils.is_cubed", lambda x: True) + + res = _compute_lazy_aware(obj) + assert np.array_equal(res, data_val) + + +def test_compute_lazy_aware_generic_compute(): + """Verify compute_lazy_aware handles objects with a .compute() method.""" + + class GenericLazy: + def compute(self): + return "computed" + + obj = GenericLazy() + res = _compute_lazy_aware(obj) + assert res == "computed" + + +if __name__ == "__main__": + pytest.main([__file__]) diff --git a/tests/test_aero_periodicity.py b/tests/test_aero_periodicity.py index 03ac97e..a3b2023 100644 --- a/tests/test_aero_periodicity.py +++ b/tests/test_aero_periodicity.py @@ -86,9 +86,9 @@ def test_regridder_periodicity_lazy(): with counter: regridder = Regridder(ds_src_meta, ds_tgt, method="bilinear") assert regridder.periodic is True - assert ( - counter.count == 0 - ), f"Metadata-based detection triggered {counter.count} computes" + assert counter.count == 0, ( + f"Metadata-based detection triggered {counter.count} computes" + ) # Case 3: Explicit periodicity avoids compute and warning with patch("xregrid.regridder.Regridder._generate_weights"): @@ -96,9 +96,9 @@ def test_regridder_periodicity_lazy(): with counter: regridder = Regridder(ds_src, ds_tgt, method="bilinear", periodic=True) assert regridder.periodic is True - assert ( - counter.count == 0 - ), f"Explicit periodicity triggered {counter.count} computes" + assert counter.count == 0, ( + f"Explicit periodicity triggered {counter.count} computes" + ) if __name__ == "__main__":