Harden tests, improve documentation, and optimize accessors

src/xregrid/core.py

@@ -85,29 +85,37 @@ def _apply_weights_core(
     """
     Apply regridding weights to a data block (NumPy array).
 
+    Handles both standard regridding and NaN-aware regridding (skipna=True)
+    using efficient sparse matrix multiplication.
+
     Parameters
     ----------
     data_block : np.ndarray
         The input data block. Core dimensions must be at the end.
-    weights_matrix : scipy.sparse.csr_matrix or str
-        The sparse weight matrix or a string key for worker-local cache.
-    dims_source : tuple of str
+    weights_matrix : Any
+        The sparse weight matrix (scipy.sparse.csr_matrix) or a string key
+        referencing the matrix in the worker-local cache.
+    dims_source : Tuple[str, ...]
         The names of the source spatial dimensions.
-    shape_target : tuple of int
+    shape_target : Tuple[int, ...]
         The shape of the target spatial grid.
     skipna : bool, default False
-        Whether to handle NaNs by re-normalizing weights.
+        Whether to handle NaNs by re-normalizing weights based on the presence
+        of valid data in each cell.
     total_weights : np.ndarray, optional
-        Pre-computed sum of weights for each destination cell.
+        Pre-computed sum of weights for each destination cell. Used for
+        normalization when skipna=True or for masking low-confidence points.
     na_thres : float, default 1.0
-        Threshold for NaN handling.
+        Threshold for NaN handling. Points with less than this fraction of
+        valid input contribution will be masked.
     weights_key : str, optional
-        Explicit key for the weights in the worker cache.
+        Explicit key for the weights in the worker cache, used for stationary
+        mask optimization.
 
     Returns
     -------
     np.ndarray
-        The regridded data block.
+        The regridded data block reshaped to match the target grid.
     """
     # Worker-local cache retrieval
     weights_matrix_key = weights_key
@@ -163,17 +171,14 @@ def _apply_weights_core(
             is_mask_stationary = True
             if n_other > 1:
                 # Optimized stationary mask detection using heuristic early exit
-                #
                 mask0 = mask[0]
                 sample_size = min(1000, n_spatial)
                 # Check first sample points across all time steps first
                 if not np.all(mask[:, :sample_size] == mask0[:sample_size]):
                     is_mask_stationary = False
                 else:
-                    for i in range(1, n_other):
-                        if not np.array_equal(mask[i], mask0):
-                            is_mask_stationary = False
-                            break
+                    # Fallback to full comparison if sample matches
+                    is_mask_stationary = np.all(mask == mask[0:1])
 
             zero = flat_data.dtype.type(0)
             if is_mask_stationary:

src/xregrid/grid.py

@@ -14,6 +14,9 @@ def _get_non_spatial_dims(ds: Union[xr.Dataset, xr.DataArray]) -> set[str]:
     """
     Identify dimensions that are likely not spatial (Time, Vertical).
 
+    Utilizes cf-xarray axes and standard names, followed by common name heuristics
+    and dtype checks.
+
     Parameters
     ----------
     ds : xr.Dataset or xr.DataArray
@@ -22,7 +25,7 @@ def _get_non_spatial_dims(ds: Union[xr.Dataset, xr.DataArray]) -> set[str]:
     Returns
     -------
     set of str
-        Names of non-spatial dimensions.
+        Names of non-spatial dimensions detected in the object.
     """
     non_spatial_dims = set()
 
@@ -38,8 +41,18 @@ def _get_non_spatial_dims(ds: Union[xr.Dataset, xr.DataArray]) -> set[str]:
         pass
 
     # 2. Heuristics based on dimension names
-    time_names = ["time", "t", "tden", "time_counter", "t_step"]
-    vert_names = [
+    time_names = {
+        "time",
+        "t",
+        "tden",
+        "time_counter",
+        "t_step",
+        "Time",
+        "T",
+        "date",
+        "dtime",
+    }
+    vert_names = {
         "lev",
         "level",
         "depth",
@@ -49,14 +62,42 @@ def _get_non_spatial_dims(ds: Union[xr.Dataset, xr.DataArray]) -> set[str]:
         "height",
         "altitude",
         "z",
-    ]
+        "Level",
+        "Layer",
+        "p",
+        "plev",
+        "bottom_top",
+        "top_bottom",
+    }
+
+    # Pre-compute lower-case set for efficient lookup
+    heuristic_names = {n.lower() for n in time_names | vert_names}
 
     for dim in ds.dims:
-        dim_lower = str(dim).lower()
-        if dim_lower in time_names or dim_lower in vert_names:
+        if str(dim).lower() in heuristic_names:
             non_spatial_dims.add(str(dim))
 
-        # 3. Dtype check for time if it's a coordinate
+    # 3. Use cf-xarray standard names
+    try:
+        std_time = ["time"]
+        std_vert = [
+            "air_pressure",
+            "height",
+            "depth",
+            "altitude",
+            "geopotential_height",
+            "height_above_msl",
+        ]
+        for std_name in std_time + std_vert:
+            if std_name in ds.cf.standard_names:
+                for var_name in ds.cf.standard_names[std_name]:
+                    if var_name in ds.dims:
+                        non_spatial_dims.add(str(var_name))
+    except (KeyError, AttributeError):
+        pass
+
+    # 4. Dtype check for time if it's a coordinate
+    for dim in ds.dims:
         if hasattr(ds, "coords") and dim in ds.coords:
             dtype = ds.coords[dim].dtype
             if np.issubdtype(dtype, np.datetime64) or np.issubdtype(

src/xregrid/regridder.py

@@ -1494,6 +1494,7 @@ def __call__(
         obj: Union[xr.DataArray, xr.Dataset, Any],
         skipna: Optional[bool] = None,
         na_thres: Optional[float] = None,
+        keep_attrs: bool = True,
     ) -> Union[xr.DataArray, xr.Dataset]:
         """
         Apply regridding to an input DataArray or Dataset.
@@ -1509,6 +1510,11 @@ def __call__(
             Threshold for NaN handling.
             If None, uses the value set during initialization.
 
+        keep_attrs : bool, default True
+            If True, merge input-object attributes onto the output.
+            Attributes set by the regridder (e.g. ``history``) take priority
+            so provenance is never lost.
+
         Returns
         -------
         xarray.DataArray or xarray.Dataset
@@ -1534,6 +1540,8 @@ def __call__(
             if self._tgt_was_sorted:
                 # Use sel to restore order from the original target grid
                 res = res.sel({d: self._orig_target_grid[d] for d in self._dims_target})
+            if keep_attrs:
+                res.attrs = {**obj.attrs, **res.attrs}
             return res
         elif isinstance(obj, xr.DataArray):
             # Check if DataArray is regriddable
@@ -1559,13 +1567,21 @@ def __call__(
             if self._tgt_was_sorted:
                 # Use sel to restore order from the original target grid
                 res = res.sel({d: self._orig_target_grid[d] for d in self._dims_target})
+            if keep_attrs:
+                res.attrs = {**obj.attrs, **res.attrs}
             return res
         # Handle uxarray objects if they don't pass isinstance(xr.Dataset)
         elif hasattr(obj, "uxgrid"):
             if hasattr(obj, "data_vars"):
-                return self._regrid_dataset(obj, skipna=skipna, na_thres=na_thres)
+                res = self._regrid_dataset(obj, skipna=skipna, na_thres=na_thres)
+                if keep_attrs:
+                    res.attrs = {**obj.attrs, **res.attrs}
+                return res
             else:
-                return self._regrid_dataarray(obj, skipna=skipna, na_thres=na_thres)
+                res = self._regrid_dataarray(obj, skipna=skipna, na_thres=na_thres)
+                if keep_attrs:
+                    res.attrs = {**obj.attrs, **res.attrs}
+                return res
         else:
             raise TypeError("Input must be an xarray.DataArray or xarray.Dataset.")
 

src/xregrid/utils.py

@@ -1025,28 +1025,29 @@ def create_grid_like(
 
     Automatically detects the CRS and spatial extent of the input object.
     Supports both geographic (lat-lon) and projected coordinate systems.
+    Efficiently handles Dask-backed objects by batching metadata discovery.
 
     Parameters
     ----------
     obj : xr.DataArray or xr.Dataset
         The input object to use as a template.
-    res : float or tuple of float
+    res : Union[float, Tuple[float, float]]
         New grid resolution in the coordinate system units.
         If tuple, (res_x, res_y) or (res_lon, res_lat).
     add_bounds : bool, default True
         Whether to add cell boundary coordinates.
-    chunks : int or dict, optional
+    chunks : Union[int, Dict[str, int]], optional
         Chunk sizes for the resulting dask-backed dataset.
-    extent : tuple of float, optional
+    extent : Tuple[float, float, float, float], optional
         Override the detected extent (min_x, max_x, min_y, max_y).
         Use this to avoid hidden dask.compute() if you already know the extent.
-    crs : str, int, or pyproj.CRS, optional
+    crs : Union[str, int, pyproj.CRS], optional
         Override the detected CRS.
 
     Returns
     -------
     xr.Dataset
-        The new grid dataset.
+        The new grid dataset with consistent metadata.
     """
     if crs is not None:
         if pyproj is not None:
@@ -1101,8 +1102,19 @@ def create_grid_like(
             if dask is not None and (
                 hasattr(x_b.data, "dask") or hasattr(y_b.data, "dask")
             ):
-                vals = dask.compute(x_b.min(), x_b.max(), y_b.min(), y_b.max())
-                extent = tuple(map(float, vals))
+                tasks_dict = {
+                    "xmin": x_b.min(),
+                    "xmax": x_b.max(),
+                    "ymin": y_b.min(),
+                    "ymax": y_b.max(),
+                }
+                results = dask.compute(tasks_dict)[0]
+                extent = (
+                    float(results["xmin"]),
+                    float(results["xmax"]),
+                    float(results["ymin"]),
+                    float(results["ymax"]),
+                )
             elif hasattr(x_b.data, "dask") or hasattr(y_b.data, "dask"):
                 extent = (
                     float(x_b.min()),
@@ -1211,9 +1223,15 @@ def create_grid_like(
             if dask is not None and (
                 hasattr(lat_b.data, "dask") or hasattr(lon_b.data, "dask")
             ):
-                vals = dask.compute(lat_b.min(), lat_b.max(), lon_b.min(), lon_b.max())
-                lat_range = (float(vals[0]), float(vals[1]))
-                lon_range = (float(vals[2]), float(vals[3]))
+                tasks_dict = {
+                    "lat_min": lat_b.min(),
+                    "lat_max": lat_b.max(),
+                    "lon_min": lon_b.min(),
+                    "lon_max": lon_b.max(),
+                }
+                results = dask.compute(tasks_dict)[0]
+                lat_range = (float(results["lat_min"]), float(results["lat_max"]))
+                lon_range = (float(results["lon_min"]), float(results["lon_max"]))
             elif hasattr(lat_b.data, "dask") or hasattr(lon_b.data, "dask"):
                 lat_range = (float(lat_b.min()), float(lat_b.max()))
                 lon_range = (float(lon_b.min()), float(lon_b.max()))

tests/test_aero_protocol.py

@@ -0,0 +1,100 @@
+from __future__ import annotations
+
+import numpy as np
+import pytest
+import xarray as xr
+from xregrid import Regridder
+from xregrid.utils import create_global_grid
+
+
+def test_aero_protocol_equivalence():
+    """
+    Verify that regridding produces identical results for NumPy and Dask backends.
+    Following the Aero Protocol: Flexibility rule.
+    """
+    # 1. Setup grids
+    ds_src = create_global_grid(1.0, 1.0)
+    ds_tgt = create_global_grid(2.0, 2.0)
+
+    # Create source data with some pattern and NaNs
+    data = np.sin(np.deg2rad(ds_src.lat)) * np.cos(np.deg2rad(ds_src.lon))
+    # Coordinates for the DataArray should only include relevant dimensions
+    coords = {
+        k: v for k, v in ds_src.coords.items() if set(v.dims).issubset({"lat", "lon"})
+    }
+    da_numpy = xr.DataArray(data, coords=coords, dims=("lat", "lon"), name="test_data")
+
+    # Add some NaNs to test skipna
+    da_numpy.values[10:20, 10:20] = np.nan
+
+    # 2. Setup Regridder
+    regridder = Regridder(ds_src, ds_tgt, method="bilinear", skipna=True)
+
+    # 3. Eager (NumPy) regridding
+    res_numpy = regridder(da_numpy)
+
+    # 4. Lazy (Dask) regridding
+    da_dask = da_numpy.chunk({"lat": 45, "lon": 90})
+    res_dask = regridder(da_dask)
+
+    # Verify result is still lazy
+    assert hasattr(res_dask.data, "dask")
+
+    # Compute and compare
+    res_dask_computed = res_dask.compute()
+
+    xr.testing.assert_allclose(res_numpy, res_dask_computed)
+
+    # 5. Scientific Hygiene: Check history attribute
+    assert "history" in res_numpy.attrs
+    assert "Regridded using xregrid.Regridder" in res_numpy.attrs["history"]
+    assert "backend=Eager" in res_numpy.attrs["history"]
+
+    assert "history" in res_dask_computed.attrs
+    assert "backend=Distributed (Dask)" in res_dask_computed.attrs["history"]
+
+
+def test_non_spatial_preservation():
+    """
+    Verify that non-spatial dimensions are correctly identified and preserved.
+    """
+    from xregrid.grid import _get_non_spatial_dims
+
+    # Create dataset with various dimension names
+    ds = xr.Dataset(
+        data_vars={
+            "temp": (("time", "lev", "lat", "lon"), np.random.rand(2, 5, 10, 20))
+        },
+        coords={
+            "time": np.arange(2),
+            "lev": np.arange(5),
+            "lat": np.arange(10),
+            "lon": np.arange(20),
+        },
+    )
+
+    non_spatial = _get_non_spatial_dims(ds)
+    assert "time" in non_spatial
+    assert "lev" in non_spatial
+    assert "lat" not in non_spatial
+    assert "lon" not in non_spatial
+
+    # Test with standard names
+    ds_std = xr.Dataset(
+        data_vars={
+            "temp": (
+                ("custom_time", "custom_p", "lat", "lon"),
+                np.random.rand(2, 5, 10, 20),
+            )
+        }
+    )
+    ds_std["custom_time"] = (("custom_time"), np.arange(2), {"standard_name": "time"})
+    ds_std["custom_p"] = (("custom_p"), np.arange(5), {"standard_name": "air_pressure"})
+
+    non_spatial_std = _get_non_spatial_dims(ds_std)
+    assert "custom_time" in non_spatial_std
+    assert "custom_p" in non_spatial_std
+
+
+if __name__ == "__main__":
+    pytest.main([__file__])