Backend-Agnostic Computation Refactor

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:

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
 

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))

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__])

tests/test_aero_periodicity.py

@@ -86,19 +86,19 @@ 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"):
         counter = ComputeCounter()
         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__":