Harden code, improve tests, and consolidate test suite

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
 

docs/api/accessors.md

@@ -0,0 +1,47 @@
+# Accessors
+
+XRegrid provides xarray accessors for both `DataArray` and `Dataset` objects, allowing you to perform regridding using a convenient `.regrid.to()` syntax.
+
+## DataArray Accessor
+
+### to
+
+::: xregrid.accessors.RegridDataArrayAccessor.to
+
+Regrid the DataArray to a target grid.
+
+```python
+import xarray as xr
+from xregrid import create_global_grid
+
+# Load some data
+da = xr.tutorial.open_dataset("air_temperature").air
+
+# Define target grid
+target_grid = create_global_grid(res_lat=1.0, res_lon=1.0)
+
+# Regrid using the accessor
+regridded_da = da.regrid.to(target_grid, method='bilinear')
+```
+
+## Dataset Accessor
+
+### to
+
+::: xregrid.accessors.RegridDatasetAccessor.to
+
+Regrid the Dataset to a target grid.
+
+```python
+import xarray as xr
+from xregrid import create_global_grid
+
+# Load some data
+ds = xr.tutorial.open_dataset("air_temperature")
+
+# Define target grid
+target_grid = create_global_grid(res_lat=1.0, res_lon=1.0)
+
+# Regrid using the accessor
+regridded_ds = ds.regrid.to(target_grid, method='bilinear')
+```

docs/api/utils.md

@@ -91,6 +91,24 @@ metadata = {
 ds = create_grid_from_ioapi(metadata)
 ```
 
+### create_lcc_grid
+
+::: xregrid.utils.create_lcc_grid
+
+Create a structured grid dataset with a Lambert Conformal Conic (LCC) projection.
+
+### create_sinusoidal_grid
+
+::: xregrid.utils.create_sinusoidal_grid
+
+Create a structured grid dataset with a Sinusoidal projection.
+
+### create_rotated_latlon_grid
+
+::: xregrid.utils.create_rotated_latlon_grid
+
+Create a structured grid dataset with a Rotated Pole (Rotated Lat-Lon) projection.
+
 ### create_mesh_from_coords
 
 ::: xregrid.utils.create_mesh_from_coords

docs/examples/scripts/README.md

@@ -19,5 +19,14 @@ Regridding station-like point data to a regular 2D grid using nearest-neighbor m
 ### [Performance Optimization](plot_performance_optimization.py)
 Efficient workflows using weight reuse to speed up repeated regridding operations.
 
+### [Accessor Showcase](plot_accessor_showcase.py)
+Simplified workflows using the `.regrid.to()` xarray accessor.
+
+### [Unstructured Grids](plot_unstructured_grids.py)
+Regridding for MPAS and ICON style unstructured meshes.
+
+### [Larger-than-Memory Data](plot_larger_than_memory.py)
+Handling massive datasets using Dask and parallel weight generation.
+
 ### [ESMPy vs. XRegrid](plot_esmpy_comparison.py)
 A comparison of code complexity between raw ESMPy and the XRegrid API.

docs/index.md

@@ -1 +1,27 @@
 --8<-- "README.md"
+
+## Architecture Overview
+
+XRegrid acts as a high-performance bridge between `xarray` and the `ESMF` (Earth System Modeling Framework). It leverages `esmpy` for robust weight generation and `scipy` for optimized weight application.
+
+```mermaid
+graph TD
+    A[xarray.Dataset / DataArray] --> B{XRegrid}
+    B --> C[ESMF / esmpy]
+    C --> D[Weight Generation]
+    D --> B
+    B --> E[scipy.sparse]
+    E --> F[Optimized Weight Application]
+    F --> G[Regridded xarray Object]
+
+    subgraph "Distributed Backend"
+    H[Dask / Cubed]
+    B --- H
+    end
+```
+
+XRegrid's architecture is designed for:
+1. **Performance**: Optimized sparse matrix operations.
+2. **Scalability**: Seamless integration with Dask for large-scale parallel processing.
+3. **Correctness**: Leveraging the industry-standard ESMF engine.
+4. **Usability**: High-level API that feels natural to xarray users.

docs/user-guide/hygiene.md

@@ -0,0 +1,76 @@
+# Scientific Hygiene
+
+XRegrid is built on the **Aero Protocol**, a set of principles designed to ensure that Earth Science data processing remains flexible, maintainable, and scientifically robust. This guide details how XRegrid helps you maintain high standards of scientific hygiene.
+
+## 1. Provenance Tracking
+
+Automatically tracking the lineage of your data is critical for reproducibility. XRegrid automatically updates the `history` attribute of your xarray objects whenever a transformation occurs.
+
+- **Weight Generation**: When a `Regridder` is initialized, it records the ESMF version, the method used, and any specific parameters (like periodicity or extrapolation).
+- **Data Application**: Every time you call a regridder on a `DataArray` or `Dataset`, a timestamped message is prepended to the `history` attribute, detailing the backend used (Eager, Dask, or Cubed) and the specific regridding parameters.
+
+```python
+# View the history of a regridded object
+print(regridded_da.attrs['history'])
+```
+
+## 2. Handling Missing Data (NaNs)
+
+In many Earth Science datasets, "missing data" (represented by NaNs) must be handled carefully to avoid biasing results, especially during conservative regridding.
+
+### Weight Re-normalization (`skipna=True`)
+
+When `skipna=True` is set in the `Regridder`, XRegrid handles NaNs by re-normalizing the interpolation weights based only on the "valid" (non-NaN) source points.
+
+- **Mechanism**: For every destination cell, XRegrid sums the weights of all contributing non-NaN source points. The interpolated value is then divided by this sum.
+- **Stationary Mask Caching**: XRegrid is optimized for datasets where the mask is stationary over time (e.g., a fixed land-sea mask). It detects if the NaN locations are identical across time steps and caches the normalization factors to provide a ~2x speedup.
+
+### Validation Threshold (`na_thres`)
+
+Even with re-normalization, you may want to mask destination cells that don't have enough valid source data. The `na_thres` parameter (default 1.0) controls this:
+
+- `na_thres=1.0`: Only mask destination cells that have **zero** valid source points.
+- `na_thres=0.5`: Mask destination cells where less than 50% of the original weight sum is represented by valid source points.
+
+```python
+regridder = Regridder(ds_src, ds_tgt, skipna=True, na_thres=0.7)
+```
+
+## 3. Weight Diagnostics
+
+You should always verify the quality of your regridding weights before trusting the results. XRegrid provides built-in tools for this.
+
+### Spatial Diagnostics
+
+The `.diagnostics()` method returns an xarray Dataset on the target grid containing:
+
+- **`weight_sum`**: The sum of weights for each destination cell. For methods like bilinear or conservative, this should ideally be 1.0.
+- **`unmapped_mask`**: A binary mask where 1 indicates a destination cell that does not overlap with any source cells.
+
+```python
+diag = regridder.diagnostics()
+diag.weight_sum.plot()
+```
+
+### Quality Reports
+
+The `.quality_report()` method provides a summary of the regridding quality:
+
+```python
+report = regridder.quality_report()
+print(f"Unmapped fraction: {report['unmapped_fraction']:.2%}")
+```
+
+## 4. Coordinate Reference Systems (CRS)
+
+XRegrid ensures that your data remains "geospatially aware" by propagating CRS metadata.
+
+- **Automated Propagation**: If the target grid has a `crs` (WKT) or `grid_mapping` attribute, XRegrid automatically attaches it to the regridded output.
+- **CF-Compliance**: XRegrid uses `cf-xarray` and `pyproj` to robustly identify and manage coordinate systems, ensuring compatibility with other geospatial tools.
+
+## 5. Backend Agnosticism
+
+Following the **"Optional Dask" Rule**, XRegrid functions are designed to work regardless of whether your data is backed by NumPy (Eager) or Dask/Cubed (Lazy).
+
+- **No Hidden Computes**: XRegrid never calls `.compute()` or `.values` inside a processing function, ensuring that laziness is preserved for large-scale workflows.
+- **Vectorized Logic**: Computations are written using `xarray.apply_ufunc` with `dask='parallelized'`, allowing the same code to run efficiently on single machines or distributed clusters.

original_files.txt

@@ -0,0 +1,87 @@
+conftest.py
+test_aero_apply_weights.py
+test_aero_aux_coord_optimization.py
+test_aero_bounds.py
+test_aero_cf_awareness.py
+test_aero_coord_preservation.py
+test_aero_coord_robustness.py
+test_aero_crs_propagation.py
+test_aero_cubed_support.py
+test_aero_diagnostics.py
+test_aero_diagnostics_backend.py
+test_aero_diagnostics_extended.py
+test_aero_distributed_opt.py
+test_aero_enhancements.py
+test_aero_enhancements_v2.py
+test_aero_enhancements_v3.py
+test_aero_grid_gen.py
+test_aero_grid_lcc.py
+test_aero_grid_like_enhanced.py
+test_aero_hardening.py
+test_aero_hardening_verification.py
+test_aero_hygiene.py
+test_aero_hygiene_mapping.py
+test_aero_ioapi_support.py
+test_aero_lazy_diagnostics.py
+test_aero_lazy_grids.py
+test_aero_memory_opt.py
+test_aero_mesh_provenance.py
+test_aero_mixed_backend.py
+test_aero_native_formats.py
+test_aero_new_features.py
+test_aero_optimization.py
+test_aero_optimization_v2.py
+test_aero_optimization_v3.py
+test_aero_plot_weights.py
+test_aero_protocol.py
+test_aero_protocol_refactor.py
+test_aero_quality_lazy.py
+test_aero_quality_report_opt.py
+test_aero_regrid_robustness.py
+test_aero_rotated_pole.py
+test_aero_sinusoidal.py
+test_aero_smart_features.py
+test_aero_spatial_slice.py
+test_aero_total_weights_opt.py
+test_aero_ufs_names.py
+test_aero_ugrid_full_support.py
+test_aero_unstructured_enhanced.py
+test_aero_utils_lazy.py
+test_aero_utils_new.py
+test_aero_viz.py
+test_aero_viz_interactive_smart.py
+test_aero_viz_unstructured.py
+test_aero_weight_loading.py
+test_backends.py
+test_cf_xarray.py
+test_cli.py
+test_dask_verification.py
+test_diagnostics.py
+test_dimension_robustness.py
+test_grids.py
+test_integration.py
+test_large_grid_optim.py
+test_misc.py
+test_model_formats.py
+test_monotonicity.py
+test_mpi.py
+test_optimization.py
+test_performance_optim.py
+test_persistence.py
+test_protocol.py
+test_rdhpcs_utils.py
+test_real_esmpy_dask.py
+test_regridder.py
+test_regridder_coverage.py
+test_robustness.py
+test_toy_regrid.py
+test_unstructured.py
+test_unstructured_dask.py
+test_unstructured_dask_advanced.py
+test_utils.py
+test_utils_lazy.py
+test_uxarray.py
+test_validation.py
+test_viz.py
+test_viz_coverage.py
+test_xregrid.py

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

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)

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
 

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