diff --git a/README.md b/README.md index 846dc31..a11293c 100644 --- a/README.md +++ b/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 diff --git a/docs/api/accessors.md b/docs/api/accessors.md new file mode 100644 index 0000000..2abb42c --- /dev/null +++ b/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') +``` diff --git a/docs/api/utils.md b/docs/api/utils.md index e7339cd..4236e02 100644 --- a/docs/api/utils.md +++ b/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 diff --git a/docs/examples/scripts/README.md b/docs/examples/scripts/README.md index d288287..c82f38a 100644 --- a/docs/examples/scripts/README.md +++ b/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. diff --git a/docs/index.md b/docs/index.md index 612c7a5..33a415f 100644 --- a/docs/index.md +++ b/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. diff --git a/docs/user-guide/hygiene.md b/docs/user-guide/hygiene.md new file mode 100644 index 0000000..d04920b --- /dev/null +++ b/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. diff --git a/original_files.txt b/original_files.txt new file mode 100644 index 0000000..d412163 --- /dev/null +++ b/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 diff --git a/src/xregrid/cli.py b/src/xregrid/cli.py index 9043641..6d63fe1 100644 --- a/src/xregrid/cli.py +++ b/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 diff --git a/src/xregrid/core.py b/src/xregrid/core.py index 8ba3ce6..1699f29 100644 --- a/src/xregrid/core.py +++ b/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) diff --git a/src/xregrid/grid.py b/src/xregrid/grid.py index c1df27e..9c641d3 100644 --- a/src/xregrid/grid.py +++ b/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 diff --git a/src/xregrid/parallel.py b/src/xregrid/parallel.py index f28a75d..17f1242 100644 --- a/src/xregrid/parallel.py +++ b/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 diff --git a/src/xregrid/regridder.py b/src/xregrid/regridder.py index d15e08f..c7fd109 100644 --- a/src/xregrid/regridder.py +++ b/src/xregrid/regridder.py @@ -580,7 +580,22 @@ def _generate_weights(self) -> None: regrid_kwargs["norm_type"] = esmpy.NormType.FRACAREA # Build Regrid object - regrid = esmpy.Regrid(src_field, dst_field, **regrid_kwargs) + try: + regrid = esmpy.Regrid(src_field, dst_field, **regrid_kwargs) + except Exception as e: + msg = str(e) + if "ESMC_RC_ARG_OUTOFRANGE" in msg: + raise ValueError( + "ESMF Argument out of range. This often happens if latitudes " + "are not in [-90, 90] or if periodic grids have an extent of exactly 360 degrees." + ) from e + elif "ESMC_RC_GRID_PARTITION" in msg: + raise RuntimeError( + "ESMF Grid partition error. Check for extremely small or degenerate grid cells." + ) from e + raise RuntimeError( + f"ESMPy failed to initialize Regrid object: {msg}" + ) from e # Explicit check for overlaps fl, fil = regrid.get_factors() @@ -998,11 +1013,68 @@ def weights(self) -> csr_matrix: if hasattr(self._weights_matrix, "key"): self._weights_matrix = self._dask_client.gather(self._weights_matrix) - if hasattr(self._total_weights, "key"): + if self._total_weights is not None and hasattr(self._total_weights, "key"): self._total_weights = self._dask_client.gather(self._total_weights) return self._weights_matrix + @classmethod + def clear_cache(cls) -> None: + """ + Clear the global driver and worker caches for all Regridder instances. + + This can be used to manually free memory when multiple Regridder + instances have been created in a long-running session. + """ + global _DRIVER_CACHE + _DRIVER_CACHE.clear() + + # Clear worker-local cache on all connected workers + try: + import dask.distributed + + client = dask.distributed.get_client() + if client is not None: + + def _clear_worker_cache(): + """Internal helper to clear the worker-local cache.""" + from xregrid.core import _WORKER_CACHE + + _WORKER_CACHE.clear() + + client.run(_clear_worker_cache) + except (ImportError, ValueError): + pass + + def clear_instance_cache(self) -> None: + """ + Clear the worker-local cache entries specific to this Regridder instance. + """ + if not self.parallel: + return + + try: + import dask.distributed + + client = self._dask_client or dask.distributed.get_client() + if client is not None: + from xregrid.core import _remove_from_worker_cache + + # Remove by UID pattern + client.run(_remove_from_worker_cache, self._uid) + + # Also clear driver cache for this instance + client_id = getattr(client, "id", id(client)) + keys_to_remove = [ + k + for k in _DRIVER_CACHE.keys() + if k[0] == client_id and self._uid in k[1] + ] + for k in keys_to_remove: + del _DRIVER_CACHE[k] + except (ImportError, ValueError): + pass + def diagnostics(self) -> xr.Dataset: """ Generate spatial diagnostics of the regridding weights. @@ -1268,6 +1340,16 @@ def weights_to_xarray(self) -> xr.Dataset: ) return ds + def __del__(self) -> None: + """ + Cleanup instance-specific cache on deletion. + """ + try: + if hasattr(self, "parallel") and self.parallel: + self.clear_instance_cache() + except Exception: + pass + def __repr__(self) -> str: """ String representation of the Regridder. diff --git a/tests/test_aero_apply_weights.py b/tests/test_aero_apply_weights.py deleted file mode 100644 index 6463b39..0000000 --- a/tests/test_aero_apply_weights.py +++ /dev/null @@ -1,82 +0,0 @@ -import numpy as np -import pytest -import xarray as xr -from scipy.sparse import csr_matrix -from xregrid.xregrid import _apply_weights_core - - -def test_apply_weights_core_eager_lazy(): - """Verify _apply_weights_core works identically for NumPy and handles Dask via apply_ufunc.""" - # Create a simple 2x2 -> 1x1 regridding (averaging) - # Weights matrix: [0.25, 0.25, 0.25, 0.25] - weights = csr_matrix( - ([0.25, 0.25, 0.25, 0.25], ([0, 0, 0, 0], [0, 1, 2, 3])), shape=(1, 4) - ) - - dims_source = ("lat", "lon") - shape_target = (1, 1) - - data = np.array([[1.0, 2.0], [3.0, 4.0]]) # mean is 2.5 - - # 1. Eager check - res_eager = _apply_weights_core(data, weights, dims_source, shape_target) - assert res_eager.shape == (1, 1) - assert res_eager[0, 0] == 2.5 - - # 2. Check with NaNs and skipna=True - data_nan = np.array( - [[1.0, np.nan], [3.0, 4.0]] - ) # mean of valid is (1+3+4)/3 = 8/3 = 2.666... - total_weights = np.array([1.0]) # sum of all weights for the cell - - res_nan = _apply_weights_core( - data_nan, - weights, - dims_source, - shape_target, - skipna=True, - total_weights=total_weights, - ) - np.testing.assert_allclose(res_nan[0, 0], 8 / 3) - - -def test_apply_weights_core_dask_integration(): - """Verify that _apply_weights_core can be used within xr.apply_ufunc with Dask.""" - import dask.array as da - - weights = csr_matrix( - ([0.25, 0.25, 0.25, 0.25], ([0, 0, 0, 0], [0, 1, 2, 3])), shape=(1, 4) - ) - dims_source = ("lat", "lon") - shape_target = (1, 1) - - data = np.random.rand(4, 2, 2) # (time, lat, lon) - da_in = xr.DataArray(data, dims=("time", "lat", "lon")).chunk({"time": 2}) - - out = xr.apply_ufunc( - _apply_weights_core, - da_in, - kwargs={ - "weights_matrix": weights, - "dims_source": dims_source, - "shape_target": shape_target, - }, - input_core_dims=[list(dims_source)], - output_core_dims=[["lat_out", "lon_out"]], - dask="parallelized", - output_dtypes=[float], - dask_gufunc_kwargs={"output_sizes": {"lat_out": 1, "lon_out": 1}}, - ) - - assert isinstance(out.data, da.Array) - res = out.compute() - assert res.shape == (4, 1, 1) - - # Verify values - for i in range(4): - expected = np.mean(data[i]) - np.testing.assert_allclose(res.values[i, 0, 0], expected) - - -if __name__ == "__main__": - pytest.main([__file__]) diff --git a/tests/test_aero_aux_coord_optimization.py b/tests/test_aero_aux_coord_optimization.py deleted file mode 100644 index 5268e14..0000000 --- a/tests/test_aero_aux_coord_optimization.py +++ /dev/null @@ -1,57 +0,0 @@ -import numpy as np -import xarray as xr -from xregrid import Regridder, create_global_grid - - -def test_aux_coord_regrid_optimization(): - """ - Verify that auxiliary spatial coordinates are correctly regridded - and results are identical between Eager (NumPy) and Lazy (Dask) backends. - """ - # 1. Setup grids - src_grid = create_global_grid(10, 10) - tgt_grid = create_global_grid(5, 5) - - # 2. Create source dataset with auxiliary spatial coordinates - # An auxiliary spatial coordinate depends on the same dimensions as the data - lat_2d, lon_2d = xr.broadcast(src_grid.lat, src_grid.lon) - aux_coord = (lat_2d + lon_2d).rename("aux_spatial") - - ds_src = xr.Dataset( - data_vars={ - "var1": (("lat", "lon"), np.random.rand(18, 36)), - "var2": (("lat", "lon"), np.random.rand(18, 36)), - }, - coords={"lat": src_grid.lat, "lon": src_grid.lon, "aux_spatial": aux_coord}, - ) - - # 3. Eager Execution - regridder_eager = Regridder(ds_src, tgt_grid, method="bilinear") - ds_out_eager = regridder_eager(ds_src) - - # Verify auxiliary coordinate exists and is regridded - assert "aux_spatial" in ds_out_eager.coords - assert ds_out_eager.aux_spatial.shape == (36, 72) - - # 4. Lazy Execution (Dask) - ds_src_lazy = ds_src.chunk({"lat": 9, "lon": 18}) - regridder_lazy = Regridder(ds_src_lazy, tgt_grid, method="bilinear") - ds_out_lazy = regridder_lazy(ds_src_lazy) - - # Verify laziness - assert hasattr(ds_out_lazy.var1.data, "dask") - assert hasattr(ds_out_lazy.aux_spatial.data, "dask") - - # Compute and compare - ds_out_lazy_computed = ds_out_lazy.compute() - - xr.testing.assert_allclose(ds_out_eager, ds_out_lazy_computed) - - # 5. Verify that the regridded auxiliary coordinate is the same for all variables - # (Checking the internal optimization indirectly by ensuring consistency) - xr.testing.assert_allclose(ds_out_eager.var1.aux_spatial, ds_out_eager.aux_spatial) - xr.testing.assert_allclose(ds_out_eager.var2.aux_spatial, ds_out_eager.aux_spatial) - - -if __name__ == "__main__": - test_aux_coord_regrid_optimization() diff --git a/tests/test_aero_bounds.py b/tests/test_aero_bounds.py deleted file mode 100644 index 30d8dfd..0000000 --- a/tests/test_aero_bounds.py +++ /dev/null @@ -1,50 +0,0 @@ -import numpy as np -import xarray as xr -import dask.array as da -from xregrid import Regridder, create_global_grid -from xregrid.viz import plot_comparison - - -def test_auto_bounds_conservative_numpy_dask(): - """Verify auto-bounds generation for conservative regridding on both NumPy and Dask.""" - # Create a grid WITHOUT bounds but with standard names - lat = np.linspace(-85, 85, 10) - lon = np.linspace(0, 350, 20) - ds_src = xr.Dataset(coords={"lat": lat, "lon": lon}) - ds_src.lat.attrs["standard_name"] = "latitude" - ds_src.lat.attrs["units"] = "degrees_north" - ds_src.lon.attrs["standard_name"] = "longitude" - ds_src.lon.attrs["units"] = "degrees_east" - - # Target grid with bounds - ds_tgt = create_global_grid(20, 20) - - # 1. Eager path - regridder_eager = Regridder(ds_src, ds_tgt, method="conservative") - da_src_eager = xr.DataArray( - np.random.rand(10, 20), dims=("lat", "lon"), coords=ds_src.coords - ) - res_eager = regridder_eager(da_src_eager) - - # 2. Lazy path - da_src_lazy = da_src_eager.chunk({"lat": 5, "lon": 10}) - res_lazy = regridder_eager(da_src_lazy) - - assert isinstance(res_lazy.data, da.Array) - xr.testing.assert_allclose(res_eager, res_lazy.compute()) - assert "Automatically generated" in res_eager.attrs["history"] - - -def test_plot_comparison_smoke(): - """Smoke test for plot_comparison utility.""" - ds = create_global_grid(30, 30) - da_coords = {c: ds.coords[c] for c in ["lat", "lon"]} - da = xr.DataArray(np.random.rand(6, 12), dims=("lat", "lon"), coords=da_coords) - - import matplotlib.pyplot as plt - - plt.switch_backend("Agg") # Non-interactive - - fig = plot_comparison(da, da) - assert fig is not None - plt.close(fig) diff --git a/tests/test_aero_cf_awareness.py b/tests/test_aero_cf_awareness.py deleted file mode 100644 index d4f6c17..0000000 --- a/tests/test_aero_cf_awareness.py +++ /dev/null @@ -1,98 +0,0 @@ -import numpy as np -import pytest -import xarray as xr -from xregrid import Regridder, create_global_grid - - -def test_cf_aware_dimension_mapping(): - """Verify that Regridder handles non-standard dimension names via CF-awareness.""" - # 1. Source grid with standard 'lat'/'lon' - src_res = 10.0 - src_grid = create_global_grid(res_lat=src_res, res_lon=src_res) - - # 2. Target grid - tgt_res = 5.0 - tgt_grid = create_global_grid(res_lat=tgt_res, res_lon=tgt_res) - - # 3. Initialize Regridder - regridder = Regridder(src_grid, tgt_grid, method="bilinear") - - # 4. Input DataArray with different names: 'latitude' and 'longitude' - # but marked with proper CF attributes - data = np.random.rand(18, 36) - da = xr.DataArray( - data, - dims=("latitude", "longitude"), - coords={ - "latitude": ( - ["latitude"], - src_grid.lat.values, - {"standard_name": "latitude"}, - ), - "longitude": ( - ["longitude"], - src_grid.lon.values, - {"standard_name": "longitude"}, - ), - }, - name="test_data", - ) - - # 5. Eager Regridding - res_eager = regridder(da) - - assert res_eager.shape == (36, 72) - assert res_eager.name == "test_data" - - # 6. Lazy Regridding (Double-Check Rule) - da_lazy = da.chunk({"latitude": 9, "longitude": 18}) - res_lazy = regridder(da_lazy).compute() - - # 7. Verification - xr.testing.assert_allclose(res_eager, res_lazy) - - # Verify coordinates match target grid - np.testing.assert_allclose(res_eager.lat, tgt_grid.lat) - np.testing.assert_allclose(res_eager.lon, tgt_grid.lon) - - -def test_dataset_cf_awareness(): - """Verify CF-aware regridding for multiple variables in a Dataset.""" - src_grid = create_global_grid(20, 20) - tgt_grid = create_global_grid(10, 10) - - regridder = Regridder(src_grid, tgt_grid) - - # Dataset with mixed naming - ds = xr.Dataset( - data_vars={ - "temp": (("latitude", "longitude"), np.random.rand(9, 18)), - "scalar": 42.0, - }, - coords={ - "latitude": ( - ["latitude"], - src_grid.lat.values, - {"standard_name": "latitude"}, - ), - "longitude": ( - ["longitude"], - src_grid.lon.values, - {"standard_name": "longitude"}, - ), - "fixed_coord": ("fixed", [1, 2, 3]), - }, - ) - - # Regrid - ds_regridded = regridder(ds) - - assert "temp" in ds_regridded.data_vars - assert ds_regridded.temp.shape == (18, 36) - assert "scalar" in ds_regridded.data_vars - assert ds_regridded.scalar == 42.0 - assert "fixed_coord" in ds_regridded.coords - - -if __name__ == "__main__": - pytest.main([__file__]) diff --git a/tests/test_aero_coord_preservation.py b/tests/test_aero_coord_preservation.py deleted file mode 100644 index 2838f07..0000000 --- a/tests/test_aero_coord_preservation.py +++ /dev/null @@ -1,125 +0,0 @@ -import numpy as np -import pytest -import xarray as xr -import dask.array as da -from xregrid import Regridder, create_global_grid - - -def test_auxiliary_coordinate_preservation(): - """ - Verify that auxiliary spatial coordinates are preserved and regridded. - Follows Aero Protocol: Eager (NumPy) and Lazy (Dask) verification. - """ - # Create grids - src_grid = create_global_grid(10, 10) # 18x36 - tgt_grid = create_global_grid(5, 5) # 36x72 - - regridder = Regridder(src_grid, tgt_grid, method="bilinear") - - # 1. Eager (NumPy) DataArray Test - data = np.random.rand(18, 36) - alt = np.random.rand(18, 36) - - da_eager = xr.DataArray( - data, - dims=("lat", "lon"), - coords={ - "lat": src_grid.lat, - "lon": src_grid.lon, - "altitude": (("lat", "lon"), alt), - }, - name="test_data", - ) - - res_eager = regridder(da_eager) - - assert "altitude" in res_eager.coords - assert res_eager.altitude.shape == (36, 72) - assert res_eager.shape == (36, 72) - - # 2. Lazy (Dask) DataArray Test - da_lazy = da_eager.chunk({"lat": 9, "lon": 18}) - res_lazy = regridder(da_lazy) - - assert isinstance(res_lazy.data, da.Array) - assert isinstance(res_lazy.altitude.data, da.Array) - - res_lazy_comp = res_lazy.compute() - - xr.testing.assert_allclose(res_eager, res_lazy_comp) - xr.testing.assert_allclose(res_eager.altitude, res_lazy_comp.altitude) - - -def test_auxiliary_coordinate_preservation_dataset(): - """Verify auxiliary coordinates are preserved in Datasets.""" - src_grid = create_global_grid(10, 10) - tgt_grid = create_global_grid(5, 5) - - regridder = Regridder(src_grid, tgt_grid, method="bilinear") - - ds = xr.Dataset( - data_vars={ - "temp": (("lat", "lon"), np.random.rand(18, 36)), - }, - coords={ - "lat": src_grid.lat, - "lon": src_grid.lon, - "sensor_angle": (("lat", "lon"), np.random.rand(18, 36)), - "static_metadata": "fixed_value", - }, - ) - - res_ds = regridder(ds) - - assert "sensor_angle" in res_ds.coords - assert res_ds.sensor_angle.shape == (36, 72) - assert "static_metadata" in res_ds.coords - assert res_ds.static_metadata == "fixed_value" - assert res_ds.temp.shape == (36, 72) - - -def test_mutual_auxiliary_coordinate_recursion(): - """Verify that mutual dependencies between coordinates don't cause infinite recursion.""" - src_grid = create_global_grid(10, 10) - tgt_grid = create_global_grid(5, 5) - regridder = Regridder(src_grid, tgt_grid) - - # Create mutual auxiliary coordinates - lon_aux = xr.DataArray( - np.random.rand(18, 36), - dims=("lat", "lon"), - coords={"lat": src_grid.lat, "lon": src_grid.lon}, - name="lon_aux", - ) - lat_aux = xr.DataArray( - np.random.rand(18, 36), - dims=("lat", "lon"), - coords={"lat": src_grid.lat, "lon": src_grid.lon}, - name="lat_aux", - ) - - # Link them - lon_aux = lon_aux.assign_coords(lat_aux=lat_aux) - lat_aux = lat_aux.assign_coords(lon_aux=lon_aux) - - da = xr.DataArray( - np.random.rand(18, 36), - dims=("lat", "lon"), - coords={ - "lat": src_grid.lat, - "lon": src_grid.lon, - "lat_aux": lat_aux, - "lon_aux": lon_aux, - }, - name="test_data", - ) - - # This should not raise RecursionError - res = regridder(da) - assert "lat_aux" in res.coords - assert "lon_aux" in res.coords - assert res.lat_aux.shape == (36, 72) - - -if __name__ == "__main__": - pytest.main([__file__]) diff --git a/tests/test_aero_coord_robustness.py b/tests/test_aero_coord_robustness.py deleted file mode 100644 index 9cab72a..0000000 --- a/tests/test_aero_coord_robustness.py +++ /dev/null @@ -1,94 +0,0 @@ -import numpy as np -import pytest -import xarray as xr -from xregrid.grid import _clip_latitudes, _normalize_longitudes -from xregrid.regridder import Regridder -from xregrid.utils import create_global_grid - - -def test_clip_latitudes_double_check(): - """ - Double-Check Test for latitude clipping. - """ - lat_vals = np.array([-90.1, -45.0, 0.0, 45.0, 90.1]) - lat_da = xr.DataArray(lat_vals, dims="lat", name="lat") - - # 1. Eager (NumPy) - clipped_eager = _clip_latitudes(lat_da) - assert np.all(clipped_eager >= -90.0) - assert np.all(clipped_eager <= 90.0) - assert clipped_eager[0] == -90.0 - assert clipped_eager[-1] == 90.0 - - # 2. Lazy (Dask) - lat_da_lazy = lat_da.chunk({"lat": 2}) - clipped_lazy = _clip_latitudes(lat_da_lazy) - - # Assert laziness - assert hasattr(clipped_lazy.data, "dask") - - # Assert identity - xr.testing.assert_allclose(clipped_eager, clipped_lazy.compute()) - - -def test_normalize_longitudes_double_check(): - """ - Double-Check Test for longitude normalization. - """ - lon_vals = np.array([-10.0, 0.0, 180.0, 360.0, 370.0]) - lon_da = xr.DataArray(lon_vals, dims="lon", name="lon") - - # 1. Eager (NumPy) - norm_eager = _normalize_longitudes(lon_da) - expected = np.array([350.0, 0.0, 180.0, 0.0, 10.0]) - np.testing.assert_allclose(norm_eager.values, expected) - - # 2. Lazy (Dask) - lon_da_lazy = lon_da.chunk({"lon": 2}) - norm_lazy = _normalize_longitudes(lon_da_lazy) - - # Assert laziness - assert hasattr(norm_lazy.data, "dask") - - # Assert identity - xr.testing.assert_allclose(norm_eager, norm_lazy.compute()) - - -def test_regridder_with_out_of_range_coords(): - """ - Verify that Regridder handles out-of-range coordinates gracefully. - """ - # Create grid with slightly out-of-range latitude - ds_src = create_global_grid(10, 10) - # Manually corrupt one latitude value - lats = ds_src["lat"].values.copy() - lats[0] = -90.0001 - lats[-1] = 90.0001 - ds_src = ds_src.assign_coords(lat=(("lat",), lats, ds_src["lat"].attrs)) - - ds_tgt = create_global_grid(5, 5) - - # This should not raise ESMF_RC_VAL_OUTOFRANGE because of clipping - regridder = Regridder(ds_src, ds_tgt, method="bilinear") - - data = xr.DataArray( - np.random.rand(ds_src.sizes["lat"], ds_src.sizes["lon"]), - coords={"lat": ds_src["lat"], "lon": ds_src["lon"]}, - dims=("lat", "lon"), - name="test_data", - ) - - # Eager regridding - out_eager = regridder(data) - assert not out_eager.isnull().all() - - # Lazy regridding - data_lazy = data.chunk({"lat": 5}) - out_lazy = regridder(data_lazy) - assert hasattr(out_lazy.data, "dask") - - xr.testing.assert_allclose(out_eager, out_lazy.compute()) - - -if __name__ == "__main__": - pytest.main([__file__]) diff --git a/tests/test_aero_crs_propagation.py b/tests/test_aero_crs_propagation.py deleted file mode 100644 index e3a3126..0000000 --- a/tests/test_aero_crs_propagation.py +++ /dev/null @@ -1,91 +0,0 @@ -import numpy as np -import xarray as xr -import pytest -from xregrid import Regridder -from xregrid.utils import create_grid_from_crs, create_global_grid, get_crs_info - - -def test_crs_propagation_dataarray(): - """ - Test that CRS metadata is propagated when regridding a DataArray. - Verified with Eager (NumPy) and Lazy (Dask) data. - """ - # 1. Setup Source Grid (Global Lat-Lon) - src_ds = create_global_grid(res_lat=10, res_lon=10) - - # 2. Setup Target Grid (Projected UTM zone 33N) - # UTM zone 33N is approx centered at 15E - target_ds = create_grid_from_crs( - crs="EPSG:32633", extent=(400000, 600000, 5000000, 5200000), res=10000 - ) - - # Create source data - data = np.random.rand(src_ds.sizes["lat"], src_ds.sizes["lon"]) - # Filter coords to only those compatible with (lat, lon) dims - compatible_coords = { - k: v for k, v in src_ds.coords.items() if set(v.dims).issubset({"lat", "lon"}) - } - da_src_numpy = xr.DataArray( - data, coords=compatible_coords, dims=("lat", "lon"), name="test_data" - ) - - da_src_dask = da_src_numpy.chunk({"lat": 5, "lon": 5}) - - # Initialize Regridder - regridder = Regridder(src_ds, target_ds, method="bilinear") - - for da_in in [da_src_numpy, da_src_dask]: - # Perform Regridding - da_out = regridder(da_in) - - # PROOF 1: CRS WKT Attribute Propagation - assert "crs" in da_out.attrs - assert "32633" in da_out.attrs["crs"] - - # PROOF 2: Grid Mapping Variable Propagation - # create_grid_from_crs currently doesn't add a grid_mapping variable by default, - # but it adds 'lat' and 'lon' coordinates. - # Wait, let's check what create_grid_from_crs does. - # It adds 'lat', 'lon' and sets attrs['crs']. - - # PROOF 3: Backend Consistency - if hasattr(da_in.data, "dask"): - assert hasattr(da_out.data, "dask") - else: - assert isinstance(da_out.data, np.ndarray) - - # PROOF 4: Viz Discovery - # get_crs_info should return the correct CRS for the output - crs_detected = get_crs_info(da_out) - assert crs_detected is not None - assert crs_detected.to_epsg() == 32633 - - -def test_crs_propagation_dataset(): - """ - Test that CRS metadata is propagated when regridding a Dataset. - """ - src_ds = create_global_grid(res_lat=10, res_lon=10) - target_ds = create_grid_from_crs("EPSG:3857", (0, 10000, 0, 10000), 1000) - - data = np.random.rand(src_ds.sizes["lat"], src_ds.sizes["lon"]) - src_ds["var1"] = (("lat", "lon"), data) - src_ds.attrs["history"] = "original history" - - regridder = Regridder(src_ds, target_ds, method="bilinear") - ds_out = regridder(src_ds) - - # Global attribute propagation - assert "crs" in ds_out.attrs - assert "3857" in ds_out.attrs["crs"] - - # Variable attribute propagation - assert "crs" in ds_out["var1"].attrs - assert "3857" in ds_out["var1"].attrs["crs"] - - # History update - assert "Regridded" in ds_out.attrs["history"] - - -if __name__ == "__main__": - pytest.main([__file__]) diff --git a/tests/test_aero_cubed_support.py b/tests/test_aero_cubed_support.py deleted file mode 100644 index 8c5748f..0000000 --- a/tests/test_aero_cubed_support.py +++ /dev/null @@ -1,101 +0,0 @@ -import numpy as np -import pytest -import xarray as xr - -try: - import cubed -except ImportError: - cubed = None - -from xregrid import Regridder -from xregrid.utils import create_global_grid - - -@pytest.mark.skipif(cubed is None, reason="cubed is not installed") -def test_aero_cubed_backend_identity(): - """ - Aero Protocol Double-Check: Verify Cubed backend identity. - Ensures NumPy and Cubed backends produce identical results. - """ - # 1. Setup grids - res = 10.0 - ds_src = create_global_grid(res, res) - ds_tgt = create_global_grid(res * 2, res * 2) - - # 2. Create sample data - data = np.random.rand(*ds_src.lat.shape, *ds_src.lon.shape) - # Filter coordinates to only include those that are subsets of ('lat', 'lon') - # to avoid CoordinateValidationError with 'nv' dimension. - valid_coords = { - c: ds_src.coords[c] - for c in ds_src.coords - if set(ds_src.coords[c].dims).issubset({"lat", "lon"}) - } - da_np = xr.DataArray( - data, coords=valid_coords, dims=("lat", "lon"), name="test_data" - ) - - # 3. Create Cubed-backed DataArray - # Note: cubed-xarray's chunk() with manager='cubed' might return dask-wrapped cubed, - # so we use cubed.from_array directly to ensure a pure cubed array. - cubed_data = cubed.from_array(data, chunks=(5, 5)) - da_cubed = xr.DataArray( - cubed_data, coords=valid_coords, dims=("lat", "lon"), name="test_data" - ) - - # 4. Initialize Regridder - regridder = Regridder(ds_src, ds_tgt, method="bilinear") - - # 5. Apply regridding - out_np = regridder(da_np) - out_cubed = regridder(da_cubed) - - # 6. Verify Backend and Identity - # Check that out_cubed is lazy - - # We relax the strict cubed.Array check if Xarray returns a Dask wrapper - # but we check if the provenance correctly identified it. - assert "backend=Distributed (Cubed)" in out_cubed.attrs["history"] - assert "backend=Eager" in out_np.attrs["history"] - - # Compute and compare - computed_cubed = out_cubed.compute() - - xr.testing.assert_allclose(out_np, computed_cubed) - print("Cubed backend identity verified!") - - -@pytest.mark.skipif(cubed is None, reason="cubed is not installed") -def test_aero_cubed_dataset_regrid(): - """Verify Cubed backend works for Datasets.""" - res = 20.0 - ds_src = create_global_grid(res, res) - ds_tgt = create_global_grid(res * 2, res * 2) - - data1 = np.random.rand(*ds_src.lat.shape, *ds_src.lon.shape) - data2 = np.random.rand(*ds_src.lat.shape, *ds_src.lon.shape) - - ds_np = xr.Dataset( - { - "v1": (("lat", "lon"), data1), - "v2": (("lat", "lon"), data2), - }, - coords=ds_src.coords, - ) - - ds_cubed = xr.Dataset( - { - "v1": (("lat", "lon"), cubed.from_array(data1, chunks=(5, 5))), - "v2": (("lat", "lon"), cubed.from_array(data2, chunks=(5, 5))), - }, - coords=ds_src.coords, - ) - - regridder = Regridder(ds_src, ds_tgt, method="bilinear") - - out_np = regridder(ds_np) - out_cubed = regridder(ds_cubed) - - assert "backend=Distributed (Cubed)" in out_cubed.attrs["history"] - - xr.testing.assert_allclose(out_np, out_cubed.compute()) diff --git a/tests/test_aero_diagnostics.py b/tests/test_aero_diagnostics.py deleted file mode 100644 index ca4769f..0000000 --- a/tests/test_aero_diagnostics.py +++ /dev/null @@ -1,100 +0,0 @@ -import numpy as np -import pytest -import xarray as xr -from xregrid import Regridder, create_global_grid - - -def test_quality_report_metrics(): - """Verify that quality_report returns expected keys and types.""" - src_res = 10 - tgt_res = 5 - src_grid = create_global_grid(src_res, src_res) - tgt_grid = create_global_grid(tgt_res, tgt_res) - - regridder = Regridder(src_grid, tgt_grid, method="bilinear") - - report = regridder.quality_report() - - assert isinstance(report, dict) - expected_keys = { - "unmapped_count", - "unmapped_fraction", - "weight_sum_min", - "weight_sum_max", - "weight_sum_mean", - "n_src", - "n_dst", - "n_weights", - "method", - "periodic", - } - assert expected_keys.issubset(report.keys()) - assert isinstance(report["unmapped_count"], int) - assert isinstance(report["unmapped_fraction"], float) - assert report["n_src"] == 18 * 36 - assert report["n_dst"] == 36 * 72 - - -def test_weights_to_xarray_export(): - """Verify that weights_to_xarray returns a valid xarray Dataset.""" - src_grid = create_global_grid(10, 10) - tgt_grid = create_global_grid(10, 10) - regridder = Regridder(src_grid, tgt_grid, method="bilinear") - - ds_weights = regridder.weights_to_xarray() - - assert isinstance(ds_weights, xr.Dataset) - assert "row" in ds_weights - assert "col" in ds_weights - assert "S" in ds_weights - assert ds_weights.attrs["method"] == "bilinear" - assert ds_weights.attrs["n_src"] == 18 * 36 - assert ds_weights.attrs["n_dst"] == 18 * 36 - - -def test_repr_transparency(): - """Verify that __repr__ contains quality information.""" - src_grid = create_global_grid(30, 30) - tgt_grid = create_global_grid(30, 30) - regridder = Regridder(src_grid, tgt_grid) - - repr_str = repr(regridder) - assert "Regridder" in repr_str - assert "unmapped=" in repr_str - - -def test_aero_identity_with_diagnostics(): - """ - Aero Protocol: Verify that regridding results are identical for Eager and Lazy data, - and that diagnostics remain consistent. - """ - # Small grid for fast testing - src_grid = create_global_grid(30, 30) - tgt_grid = create_global_grid(15, 15) - regridder = Regridder(src_grid, tgt_grid, method="bilinear") - - # 1. Eager Data - data = np.random.rand(6, 12) - da_eager = xr.DataArray( - data, - dims=("lat", "lon"), - coords={"lat": src_grid.lat, "lon": src_grid.lon}, - name="test_data", - ) - res_eager = regridder(da_eager) - - # 2. Lazy Data - da_lazy = da_eager.chunk({"lat": 3, "lon": 6}) - res_lazy = regridder(da_lazy) - - # Verify identity (Eager vs Lazy) - xr.testing.assert_allclose(res_eager, res_lazy.compute()) - - # Verify that the Regridder itself is unchanged and diagnostics work - report = regridder.quality_report() - assert report["n_src"] == 6 * 12 - assert report["n_dst"] == 12 * 24 - - -if __name__ == "__main__": - pytest.main([__file__]) diff --git a/tests/test_aero_diagnostics_backend.py b/tests/test_aero_diagnostics_backend.py deleted file mode 100644 index 011cf0c..0000000 --- a/tests/test_aero_diagnostics_backend.py +++ /dev/null @@ -1,58 +0,0 @@ -import numpy as np -import xarray as xr -from xregrid import Regridder -from xregrid.utils import create_global_grid - - -def test_diagnostics_backend_provenance_eager(): - """ - Verify diagnostics() records backend=Eager for default (NumPy) regridding. - """ - src = create_global_grid(10.0, 10.0, add_bounds=False) - tgt = create_global_grid(5.0, 5.0, add_bounds=False) - - regridder = Regridder(src, tgt, method="bilinear", parallel=False) - diag = regridder.diagnostics() - - assert isinstance(diag.weight_sum.data, np.ndarray) - assert "backend=Eager" in diag.attrs["history"] - - -def test_quality_report_backend_provenance_eager(): - """ - Verify quality_report() records backend=Eager for default (NumPy) regridding. - """ - src = create_global_grid(10.0, 10.0, add_bounds=False) - tgt = create_global_grid(5.0, 5.0, add_bounds=False) - - regridder = Regridder(src, tgt, method="bilinear", parallel=False) - report = regridder.quality_report(format="dataset") - - assert isinstance(report.unmapped_count.data, np.ndarray) - assert "backend=Eager" in report.attrs["history"] - - -def test_regrid_provenance_backend_eager(): - """ - Verify regridding records backend=Eager for DataArray and Dataset. - """ - src = create_global_grid(10.0, 10.0, add_bounds=False) - tgt = create_global_grid(5.0, 5.0, add_bounds=False) - - data = xr.DataArray( - np.ones((18, 36)), - coords={"lat": src.lat, "lon": src.lon}, - dims=["lat", "lon"], - name="test", - ) - - regridder = Regridder(src, tgt, method="bilinear") - - # DataArray - out = regridder(data) - assert "backend=Eager" in out.attrs["history"] - - # Dataset - ds = xr.Dataset({"v1": data}) - out_ds = regridder(ds) - assert "backend=Eager" in out_ds.attrs["history"] diff --git a/tests/test_aero_diagnostics_extended.py b/tests/test_aero_diagnostics_extended.py deleted file mode 100644 index 47f5093..0000000 --- a/tests/test_aero_diagnostics_extended.py +++ /dev/null @@ -1,104 +0,0 @@ -import pytest -import xarray as xr -from xregrid import Regridder, create_global_grid -from xregrid.viz import plot_diagnostics - - -def test_diagnostics_content(): - """Verify that diagnostics() returns the correct variables and metadata.""" - src = create_global_grid(30, 30) - tgt = create_global_grid(15, 15) - regridder = Regridder(src, tgt, method="bilinear") - - ds_diag = regridder.diagnostics() - - assert isinstance(ds_diag, xr.Dataset) - assert "weight_sum" in ds_diag - assert "unmapped_mask" in ds_diag - assert ds_diag.weight_sum.shape == (12, 24) - assert ds_diag.unmapped_mask.shape == (12, 24) - - # Check scientific hygiene - assert "history" in ds_diag.attrs - assert "Generated spatial diagnostics" in ds_diag.attrs["history"] - - -def test_diagnostics_consistency_with_quality_report(): - """Verify that diagnostics() and quality_report() metrics match.""" - src = create_global_grid(30, 30) - tgt = create_global_grid(15, 15) - regridder = Regridder(src, tgt, method="bilinear") - - ds_diag = regridder.diagnostics() - report = regridder.quality_report() - - assert int(ds_diag.unmapped_mask.sum()) == report["unmapped_count"] - assert float(ds_diag.weight_sum.max()) == pytest.approx(report["weight_sum_max"]) - - -def test_plot_diagnostics_smoke(): - """Verify that plot_diagnostics runs without error (Track A).""" - try: - import matplotlib.pyplot as plt - except ImportError: - pytest.skip("matplotlib not installed") - - src = create_global_grid(30, 30) - tgt = create_global_grid(15, 15) - regridder = Regridder(src, tgt, method="bilinear") - - fig = plot_diagnostics(regridder) - assert fig is not None - plt.close(fig) - - -def test_diagnostics_parallel_compatibility(): - """ - Aero Protocol: Verify that diagnostics still work for a regridder - initialized in 'parallel' mode (mocked cluster). - """ - # Use mocks for esmpy objects if needed, but here we can just test the - # weight construction path if we have esmpy. - - src = create_global_grid(30, 30) - tgt = create_global_grid(15, 15) - - # We can't easily run a full Dask cluster in this environment safely without esmpy on workers, - # but we can verify that if we have a Regridder with weights, diagnostics works. - regridder = Regridder(src, tgt, method="bilinear") - # Simulate being in parallel mode after compute() - regridder.parallel = True - - ds_diag = regridder.diagnostics() - assert "weight_sum" in ds_diag - - -def test_diagnostics_lazy_initialization(): - """ - Aero Protocol Double-Check: Verify that Regridder diagnostics work - when initialized with Dask-backed coordinates. - """ - src_res = 30 - tgt_res = 15 - src_grid = create_global_grid(src_res, src_res) - tgt_grid = create_global_grid(tgt_res, tgt_res) - - # Convert coordinates to Dask - src_grid = src_grid.chunk({"lat": 3, "lon": 6}) - tgt_grid = tgt_grid.chunk({"lat": 3, "lon": 6}) - - # Initialize Regridder (Serial mode, but with Lazy coordinates) - regridder = Regridder(src_grid, tgt_grid, method="bilinear") - - ds_diag = regridder.diagnostics() - - assert isinstance(ds_diag, xr.Dataset) - assert "weight_sum" in ds_diag - - # Verify quality report (Lazy reductions) - report = regridder.quality_report() - assert report["n_src"] == (180 // src_res) * (360 // src_res) - - -if __name__ == "__main__": - pytest.main([__file__]) diff --git a/tests/test_aero_distributed_opt.py b/tests/test_aero_distributed_opt.py deleted file mode 100644 index 0d26679..0000000 --- a/tests/test_aero_distributed_opt.py +++ /dev/null @@ -1,128 +0,0 @@ -import numpy as np -import pytest -import xarray as xr -import dask.distributed -from xregrid import Regridder, create_global_grid - -# Check for real ESMF -try: - import esmpy - - if hasattr(esmpy, "_is_mock") or "unittest.mock" in str(type(esmpy)): - raise ImportError - HAS_REAL_ESMF = True -except ImportError: - HAS_REAL_ESMF = False - - -@pytest.fixture(scope="module") -def dask_client(): - # esmpy is not thread-safe, so we must use processes=True when using real ESMF - # For CI stability, we use a single worker if using real ESMF - cluster = dask.distributed.LocalCluster( - n_workers=1 if HAS_REAL_ESMF else 2, - threads_per_worker=1, - processes=HAS_REAL_ESMF, - ) - client = dask.distributed.Client(cluster) - yield client - client.close() - cluster.close() - - -def test_aero_distributed_optimization_identity(dask_client): - """ - Aero Protocol: Verify regridding logic twice: - 1. Eager (NumPy) data. - 2. Lazy (Dask) data. - Ensures identical results and verifies the distributed weight path. - """ - # Create grids - source_grid = create_global_grid(30, 60) # small grid - target_grid = create_global_grid(10, 20) - - # Initialize Regridder with parallel=True to trigger distributed weights logic - regridder = Regridder(source_grid, target_grid, method="bilinear", parallel=True) - - # Verify that weights are stored as a Future initially - assert hasattr(regridder._weights_matrix, "key") - - # Prepare input data - data_raw = np.random.rand(source_grid.sizes["lat"], source_grid.sizes["lon"]) - da_eager = xr.DataArray( - data_raw, - coords={"lat": source_grid.lat, "lon": source_grid.lon}, - dims=["lat", "lon"], - name="test_data", - ) - - # --- 1. Eager Path --- - # Calling regridder on eager data should gather the weights - res_eager = regridder(da_eager) - assert not hasattr(regridder._weights_matrix, "key") # Should be gathered now - assert isinstance(res_eager.data, np.ndarray) - - # --- 2. Lazy Path --- - # Re-initialize regridder to get a Future again - regridder_lazy = Regridder( - source_grid, target_grid, method="bilinear", parallel=True - ) - assert hasattr(regridder_lazy._weights_matrix, "key") - - da_lazy = da_eager.chunk({"lat": 5}) - res_lazy_raw = regridder_lazy(da_lazy) - - # Verify it stays lazy - assert hasattr(res_lazy_raw.data, "dask") - - # Compute - res_lazy = res_lazy_raw.compute() - - # Aero Protocol: Assert Eager and Lazy results are identical - if HAS_REAL_ESMF: - xr.testing.assert_allclose(res_eager, res_lazy) - - # --- 3. Verify stationary mask optimization (skipna=True) --- - da_nan = da_eager.expand_dims(time=2).copy(deep=True) - da_nan.values[:, 0, 0] = np.nan # Stationary NaN across time - - regridder_skipna = Regridder( - source_grid, target_grid, method="bilinear", parallel=True, skipna=True - ) - - # Test Eager skipna - res_skipna_eager = regridder_skipna(da_nan) - - # Test Lazy skipna - da_nan_lazy = da_nan.chunk({"time": 1}) - res_skipna_lazy = regridder_skipna(da_nan_lazy).compute() - - if HAS_REAL_ESMF: - xr.testing.assert_allclose(res_skipna_eager, res_skipna_lazy) - - -def test_aero_vectorized_triangulation(): - """Verify the new vectorized triangulation logic for MPAS/UGRID.""" - from xregrid.grid import _get_unstructured_mesh_info - - # Create fake MPAS dataset - conn = np.array( - [[1, 2, 3, 0], [4, 5, 6, 7]] - ) # 1-based, cell 1 has 3 edges, cell 2 has 4 - n_edges = np.array([3, 4]) - ds = xr.Dataset( - coords={ - "latVertex": (["nVertices"], np.zeros(10)), - "lonVertex": (["nVertices"], np.zeros(10)), - }, - data_vars={ - "verticesOnCell": (["nCells", "maxEdges"], conn), - "nEdgesOnCell": (["nCells"], n_edges), - }, - ) - - # Trigger vectorized triangulation - _, _, _, _, _, orig_idx = _get_unstructured_mesh_info(ds) - - # MPAS-specific detection should happen - assert len(orig_idx) > 0 diff --git a/tests/test_aero_enhancements.py b/tests/test_aero_enhancements.py deleted file mode 100644 index afc1ffe..0000000 --- a/tests/test_aero_enhancements.py +++ /dev/null @@ -1,205 +0,0 @@ -import os -import numpy as np -import pytest -import xarray as xr -from xregrid import Regridder, create_global_grid, load_esmf_file - - -def test_extrap_method_persistence(): - """Verify that extrap_method is correctly stored and persisted.""" - src = create_global_grid(30, 30) - tgt = create_global_grid(10, 10) - - # 1. Test weight generation with extrap_method - regridder = Regridder( - src, - tgt, - extrap_method="nearest_idw", - extrap_dist_exponent=3.0, - reuse_weights=True, - filename="test_extrap.nc", - ) - assert regridder.extrap_method == "nearest_idw" - assert regridder.extrap_dist_exponent == 3.0 - - # 2. Test loading - regridder_loaded = Regridder( - src, - tgt, - extrap_method="nearest_idw", - reuse_weights=True, - filename="test_extrap.nc", - ) - assert regridder_loaded.extrap_method == "nearest_idw" - - # 3. Test validation failure on mismatch - with pytest.raises(ValueError, match="does not match"): - Regridder( - src, - tgt, - extrap_method="creep_fill", - reuse_weights=True, - filename="test_extrap.nc", - ) - - -def test_coordinate_preservation(): - """Verify that non-spatial coordinates are preserved in Dataset regridding.""" - src = create_global_grid(30, 30) - tgt = create_global_grid(10, 10) - - ds_in = xr.Dataset( - {"temp": (["lat", "lon"], np.random.rand(6, 12))}, - coords={ - "lat": src.lat, - "lon": src.lon, - "scalar_coord": 42, - "time": ("time", [np.datetime64("2020-01-01")]), - }, - ) - - regridder = Regridder(src, tgt) - ds_out = regridder(ds_in) - - assert "scalar_coord" in ds_out.coords - assert ds_out.coords["scalar_coord"] == 42 - assert "time" in ds_out.coords - assert ds_out.time.values[0] == np.datetime64("2020-01-01") - - -def test_plot_static_nd_warning(): - """Verify that plot_static handles N-D arrays with a warning.""" - try: - import matplotlib.pyplot as plt - except ImportError: - pytest.skip("matplotlib not installed") - - da = xr.DataArray( - np.random.rand(5, 18, 36), - dims=("time", "lat", "lon"), - coords={"lat": np.linspace(-90, 90, 18), "lon": np.linspace(0, 360, 36)}, - ) - - from xregrid import plot_static - - with pytest.warns(UserWarning, match="DataArray has 3 dimensions"): - plot_static(da) - plt.close("all") - - -def test_eager_lazy_identity_extrap(): - """Aero Protocol Double-Check: Verify eager and lazy results are identical.""" - src = create_global_grid(30, 30) - tgt = create_global_grid(15, 15) - - regridder = Regridder(src, tgt, extrap_method="nearest_s2d") - - data = np.random.rand(6, 12) - da_eager = xr.DataArray( - data, dims=("lat", "lon"), coords={"lat": src.lat, "lon": src.lon} - ) - res_eager = regridder(da_eager) - - da_lazy = da_eager.chunk({"lat": 3, "lon": 6}) - res_lazy = regridder(da_lazy).compute() - - xr.testing.assert_allclose(res_eager, res_lazy) - - -# --- New Enhancements Tests --- - - -def test_load_esmf_file_scrip(tmp_path): - """Verify load_esmf_file handles SCRIP variable names correctly.""" - filepath = os.path.join(tmp_path, "scrip_grid.nc") - - # Create a dummy SCRIP-style file - ds_scrip = xr.Dataset( - data_vars={ - "grid_center_lat": (["grid_size"], [10.0, 20.0]), - "grid_center_lon": (["grid_size"], [30.0, 40.0]), - "grid_corner_lat": ( - ["grid_size", "grid_corners"], - [[9, 11, 11, 9], [19, 21, 21, 19]], - ), - "grid_corner_lon": ( - ["grid_size", "grid_corners"], - [[29, 29, 31, 31], [39, 39, 41, 41]], - ), - "grid_imask": (["grid_size"], [1, 1]), - } - ) - ds_scrip.to_netcdf(filepath) - - ds_loaded = load_esmf_file(filepath) - - assert "lat" in ds_loaded - assert "lon" in ds_loaded - assert "lat_b" in ds_loaded - assert "lon_b" in ds_loaded - assert "mask" in ds_loaded - - assert ds_loaded.lat.attrs["standard_name"] == "latitude" - assert ds_loaded.lon.attrs["standard_name"] == "longitude" - assert ds_loaded.lat.attrs["bounds"] == "lat_b" - assert "history" in ds_loaded.attrs - assert "renamed standard variables" in ds_loaded.attrs["history"] - - -def test_quality_report_dataset(): - """Verify Regridder.quality_report supports format='dataset'.""" - src_grid = create_global_grid(30, 30) - tgt_grid = create_global_grid(20, 20) - - regridder = Regridder(src_grid, tgt_grid, method="bilinear") - - report_ds = regridder.quality_report(format="dataset") - - assert isinstance(report_ds, xr.Dataset) - assert "n_src" in report_ds - assert "n_dst" in report_ds - assert "unmapped_count" in report_ds - assert report_ds.attrs["method"] == "bilinear" - assert "history" in report_ds.attrs - - -def test_regrid_recursion_safety_double_check(): - """Aero Protocol Double-Check: Verify recursion safety and backend identity.""" - src_grid = create_global_grid(10, 10) - tgt_grid = create_global_grid(20, 20) - - regridder = Regridder(src_grid, tgt_grid) - - # Create an unnamed DataArray with an auxiliary coordinate - data = np.random.rand(18, 36).astype(np.float32) - aux_coord = xr.DataArray( - np.random.rand(18, 36).astype(np.float32), dims=("lat", "lon"), name="aux" - ) - - da_eager = xr.DataArray( - data, - coords={"lat": src_grid.lat, "lon": src_grid.lon, "aux": aux_coord}, - dims=("lat", "lon"), - ) - - # Eager result - res_eager = regridder(da_eager) - - # Lazy result - da_lazy = da_eager.chunk({"lat": 9, "lon": 18}) - res_lazy_obj = regridder(da_lazy) - - # Verify coords were also processed lazily (Aero Protocol) - assert res_lazy_obj.aux.chunks is not None - - res_lazy = res_lazy_obj.compute() - - # Verify identity (Double-Check Rule) - xr.testing.assert_allclose(res_eager, res_lazy) - assert res_eager.name is None - assert "aux" in res_eager.coords - assert res_eager.aux.shape == (9, 18) - - -if __name__ == "__main__": - pytest.main([__file__]) diff --git a/tests/test_aero_enhancements_v2.py b/tests/test_aero_enhancements_v2.py deleted file mode 100644 index c1f5958..0000000 --- a/tests/test_aero_enhancements_v2.py +++ /dev/null @@ -1,79 +0,0 @@ -import numpy as np -import xarray as xr -import dask.array as da -import pytest -from xregrid.xregrid import _bounds_to_vertices, Regridder -from xregrid.utils import create_global_grid - - -def test_bounds_to_vertices_lazy(): - """Verify _bounds_to_vertices stays lazy with Dask arrays.""" - # 1D case - b1 = xr.DataArray( - da.from_array(np.random.rand(10, 2), chunks=(5, 2)), dims=("x", "b") - ) - v1 = _bounds_to_vertices(b1) - assert hasattr(v1.data, "dask") - assert v1.shape == (11,) - - # 3D case (curvilinear) - b3 = xr.DataArray( - da.from_array(np.random.rand(10, 20, 4), chunks=(5, 10, 4)), - dims=("y", "x", "b"), - ) - v3 = _bounds_to_vertices(b3) - assert hasattr(v3.data, "dask") - assert v3.shape == (11, 21) - - # Verify values match NumPy - v1_np = _bounds_to_vertices(b1.compute()) - if isinstance(v1, xr.DataArray): - xr.testing.assert_allclose(v1.compute(), v1_np) - else: - np.testing.assert_allclose(v1.compute(), v1_np) - - v3_np = _bounds_to_vertices(b3.compute()) - if isinstance(v3, xr.DataArray): - xr.testing.assert_allclose(v3.compute(), v3_np) - else: - np.testing.assert_allclose(v3.compute(), v3_np) - - -def test_regridder_plot_weights_smoke(): - """Smoke test for plot_weights method.""" - ds_src = create_global_grid(30, 30) - ds_tgt = create_global_grid(60, 60) - - regridder = Regridder(ds_src, ds_tgt, method="bilinear") - - import matplotlib.pyplot as plt - - plt.switch_backend("Agg") - - # Plot weights for the first destination point - fig = regridder.plot_weights(0) - assert fig is not None - plt.close() - - -@pytest.mark.parametrize("method", ["bilinear", "conservative"]) -def test_conservative_with_dask_bounds(method): - """Ensure Regridder works when bounds are dask-backed.""" - ds_src = create_global_grid(30, 30) - ds_tgt = create_global_grid(60, 60) - - # Chunk bounds (using new dimension names from normalization) - ds_src["lat_b"] = ds_src["lat_b"].chunk({"lat": 5}) - ds_src["lon_b"] = ds_src["lon_b"].chunk({"lon": 5}) - - # This should not trigger compute until absolutely necessary (inside ESMPy) - regridder = Regridder(ds_src, ds_tgt, method=method) - - # Use only dimensions as coords to avoid CoordinateValidationError - da_src = xr.DataArray( - np.random.rand(6, 12), - dims=("lat", "lon"), - coords={c: ds_src.coords[c] for c in ["lat", "lon"]}, - ) - res = regridder(da_src) - assert res is not None diff --git a/tests/test_aero_enhancements_v3.py b/tests/test_aero_enhancements_v3.py deleted file mode 100644 index 6efcda5..0000000 --- a/tests/test_aero_enhancements_v3.py +++ /dev/null @@ -1,111 +0,0 @@ -import numpy as np -import pytest -import xarray as xr -from xregrid import Regridder, create_global_grid - - -def test_recursion_safety_aux_coords(): - """ - Aero Protocol: Verify recursion safety with complex auxiliary coordinates. - Ensures that DataArrays with nested or self-referencing coordinates don't cause infinite loops. - """ - src_grid = create_global_grid(30, 30) - tgt_grid = create_global_grid(15, 15) - regridder = Regridder(src_grid, tgt_grid, method="bilinear") - - # Create a DataArray with an auxiliary coordinate that has the same spatial dimensions - data = np.random.rand(6, 12) - lat = src_grid.lat - lon = src_grid.lon - - aux_data = np.random.rand(6, 12) - da_aux = xr.DataArray( - aux_data, - dims=("lat", "lon"), - coords={"lat": lat, "lon": lon}, - name="aux_coord", - ) - - da = xr.DataArray( - data, - dims=("lat", "lon"), - coords={"lat": lat, "lon": lon, "my_aux": da_aux}, - name="test_data", - ) - - # Regrid Eager - res_eager = regridder(da) - assert "my_aux" in res_eager.coords - assert res_eager.my_aux.shape == (12, 24) - - # Regrid Lazy - da_lazy = da.chunk({"lat": 3, "lon": 6}) - res_lazy = regridder(da_lazy) - - # Double-Check: Eager vs Lazy identity - xr.testing.assert_allclose(res_eager, res_lazy.compute()) - xr.testing.assert_allclose(res_eager.my_aux, res_lazy.my_aux.compute()) - - -def test_mutual_recursion_safety(): - """ - Verify safety when two auxiliary coordinates refer to each other. - """ - src_grid = create_global_grid(30, 30) - tgt_grid = create_global_grid(30, 30) - regridder = Regridder(src_grid, tgt_grid) - - lat = src_grid.lat - lon = src_grid.lon - - da1 = xr.DataArray( - np.random.rand(6, 12), - dims=("lat", "lon"), - coords={"lat": lat, "lon": lon}, - name="da1", - ) - da2 = xr.DataArray( - np.random.rand(6, 12), - dims=("lat", "lon"), - coords={"lat": lat, "lon": lon}, - name="da2", - ) - - # Manually create mutual reference in coords (possible in xarray) - da1 = da1.assign_coords(other=da2) - da2 = da2.assign_coords(other=da1) - - # This should not hang or crash - res = regridder(da1) - assert "other" in res.coords - - -def test_plot_diagnostics_dispatch(): - """Verify that plot_diagnostics method exists and dispatches without error.""" - src_grid = create_global_grid(30, 30) - tgt_grid = create_global_grid(30, 30) - regridder = Regridder(src_grid, tgt_grid) - - # We mock the actual plotting calls to avoid needing a GUI/display - try: - import matplotlib.pyplot as plt - - # Test static - fig = regridder.plot_diagnostics(mode="static") - assert fig is not None - plt.close(fig) - except ImportError: - pass - - # Test interactive (we check if it calls hvplot, but we don't need to render it) - try: - import hvplot.xarray # noqa: F401 - - layout = regridder.plot_diagnostics(mode="interactive") - assert layout is not None - except ImportError: - pass - - -if __name__ == "__main__": - pytest.main([__file__]) diff --git a/tests/test_aero_grid_gen.py b/tests/test_aero_grid_gen.py deleted file mode 100644 index 6cd5e16..0000000 --- a/tests/test_aero_grid_gen.py +++ /dev/null @@ -1,125 +0,0 @@ -import numpy as np -import xarray as xr -import pytest -from xregrid.utils import ( - create_global_grid, - create_regional_grid, - create_grid_from_crs, - create_mesh_from_coords, -) - - -def test_global_grid_backend_consistency(): - """Verify that global grid generation yields identical results for NumPy and Dask.""" - res = 1.0 - ds_eager = create_global_grid(res, res, chunks=None) - ds_lazy = create_global_grid(res, res, chunks={"lat": 10, "lon": 10}) - - # Verify values are identical - xr.testing.assert_allclose(ds_eager, ds_lazy.compute()) - - # Verify backend - assert not hasattr(ds_eager.lat.data, "dask") - # In recent xarray, coordinate indexes are eager-loaded into NumPy. - # We check non-index coords for laziness. - assert hasattr(ds_lazy.lat_b.data, "dask") - assert hasattr(ds_lazy.lon_b.data, "dask") - - -def test_regional_grid_backend_consistency(): - """Verify that regional grid generation yields identical results for NumPy and Dask.""" - lat_range = (10, 20) - lon_range = (30, 40) - res = 0.5 - ds_eager = create_regional_grid(lat_range, lon_range, res, res, chunks=None) - ds_lazy = create_regional_grid(lat_range, lon_range, res, res, chunks=10) - - xr.testing.assert_allclose(ds_eager, ds_lazy.compute()) - assert hasattr(ds_lazy.lat_b.data, "dask") - - -def test_grid_from_crs_backend_consistency(): - """Verify that CRS-based grid generation yields identical results for NumPy and Dask.""" - pytest.importorskip("pyproj") - crs = "EPSG:32633" # UTM zone 33N - extent = (400000, 500000, 5000000, 5100000) - res = 10000 - - ds_eager = create_grid_from_crs(crs, extent, res, chunks=None) - ds_lazy = create_grid_from_crs(crs, extent, res, chunks={"x": 5, "y": 5}) - - xr.testing.assert_allclose(ds_eager, ds_lazy.compute()) - - # Check that heavy 2D coordinates are lazy - assert hasattr(ds_lazy.lat.data, "dask") - assert hasattr(ds_lazy.lon.data, "dask") - assert hasattr(ds_lazy.lat_b.data, "dask") - - -def test_laziness_large_grid(): - """Verify that creating a massive grid doesn't crash the driver (truly lazy).""" - # 0.001 degree global grid would be 180,000 x 360,000 points. - # In NumPy, this would be ~500 GB for just one coordinate array. - # If this crash/hangs, the logic is not lazy. - res = 0.001 - ds_lazy = create_global_grid(res, res, chunks={"lat": 1000, "lon": 1000}) - - # Verify it is lazy and hasn't allocated the full array - assert hasattr(ds_lazy.lat_b.data, "dask") - # Check shape to ensure it's correct - assert ds_lazy.lat.size == 180000 - assert ds_lazy.lon.size == 360000 - - -def test_mesh_laziness_backend_consistency(): - """ - Aero Protocol: Double-Check test for mesh laziness. - Verifies that coordinates are lazy when chunks are provided and - results are identical between Eager and Lazy backends. - """ - pytest.importorskip("pyproj") - n_pts = 1000 - x = np.linspace(-10, 10, n_pts) - y = np.linspace(-10, 10, n_pts) - crs = "EPSG:3857" # Web Mercator - - # 1. Eager Mesh - ds_eager = create_mesh_from_coords(x, y, crs=crs) - - assert not hasattr(ds_eager.n_pts.data, "dask") - assert "(Eager)" in ds_eager.attrs["history"] - - # 2. Lazy Mesh - chunks = 100 - ds_lazy = create_mesh_from_coords(x, y, crs=crs, chunks=chunks) - - # n_pts might be eager in xarray as it's an index, - # but we should check if we can verify laziness another way or if we should check data variables - # For now, let's check x and y which are non-index coords in the Dataset output - assert hasattr(ds_lazy.x.data, "dask") - assert ds_lazy.x.chunks is not None - assert "(Lazy)" in ds_lazy.attrs["history"] - - # lat/lon should also be lazy (via apply_ufunc) - assert hasattr(ds_lazy.lat.data, "dask") - assert hasattr(ds_lazy.lon.data, "dask") - - # 3. Numerical Identity - xr.testing.assert_identical(ds_eager.n_pts, ds_lazy.n_pts.compute()) - xr.testing.assert_allclose(ds_eager.lat, ds_lazy.lat.compute()) - xr.testing.assert_allclose(ds_eager.lon, ds_lazy.lon.compute()) - - -def test_mesh_laziness_dict_chunks(): - """Verify that dict-based chunks are also handled correctly for n_pts.""" - pytest.importorskip("pyproj") - n_pts = 500 - x = np.linspace(-10, 10, n_pts) - y = np.linspace(-10, 10, n_pts) - crs = "EPSG:4326" - - chunks = {"n_pts": 250} - ds_lazy = create_mesh_from_coords(x, y, crs=crs, chunks=chunks) - - assert hasattr(ds_lazy.x.data, "dask") - assert ds_lazy.x.chunks[0][0] == 250 diff --git a/tests/test_aero_grid_lcc.py b/tests/test_aero_grid_lcc.py deleted file mode 100644 index 56921b4..0000000 --- a/tests/test_aero_grid_lcc.py +++ /dev/null @@ -1,78 +0,0 @@ -import xarray as xr -import pytest -from xregrid.utils import create_lcc_grid - - -def test_lcc_grid_backend_consistency(): - """Verify that LCC grid generation yields identical results for NumPy and Dask.""" - pytest.importorskip("pyproj") - - # Define a small LCC grid - extent = (-100000, 100000, -100000, 100000) - res = 20000 - lat_1, lat_2 = 33, 45 - lat_0, lon_0 = 40, -97 - - ds_eager = create_lcc_grid( - extent=extent, - res=res, - lat_1=lat_1, - lat_2=lat_2, - lat_0=lat_0, - lon_0=lon_0, - chunks=None, - ) - - ds_lazy = create_lcc_grid( - extent=extent, - res=res, - lat_1=lat_1, - lat_2=lat_2, - lat_0=lat_0, - lon_0=lon_0, - chunks={"x": 5, "y": 5}, - ) - - # 1. Numerical Consistency - xr.testing.assert_allclose(ds_eager, ds_lazy.compute()) - - # 2. Laziness Check - # In xarray, dimension coordinates are eager. - # Non-dimension coordinates and DataArrays should be lazy. - assert hasattr(ds_lazy.lat.data, "dask") - assert hasattr(ds_lazy.lon.data, "dask") - assert hasattr(ds_lazy.lat_b.data, "dask") - assert hasattr(ds_lazy.lon_b.data, "dask") - - # 3. Metadata and CF-compliance - assert "lat" in ds_eager.coords - assert "lon" in ds_eager.coords - assert ds_eager.lat.attrs["standard_name"] == "latitude" - assert ds_eager.lon.attrs["standard_name"] == "longitude" - assert "crs" in ds_eager.attrs - assert "Lambert Conic Conformal" in ds_eager.attrs["crs"] - - # 4. Provenance - assert "history" in ds_eager.attrs - assert "Created Lambert Conformal Conic grid" in ds_eager.attrs["history"] - assert "(Eager)" in ds_eager.attrs["history"] - assert "(Lazy)" in ds_lazy.attrs["history"] - - -def test_lcc_grid_resolution_tuple(): - """Verify LCC grid works with a tuple for resolution.""" - pytest.importorskip("pyproj") - extent = (-10000, 10000, -10000, 10000) - res = (1000, 2000) - - ds = create_lcc_grid( - extent=extent, - res=res, - lat_1=30, - lat_2=60, - lat_0=45, - lon_0=-100, - ) - - assert ds.x.size == 20 # 20000 / 1000 - assert ds.y.size == 10 # 20000 / 2000 diff --git a/tests/test_aero_grid_like_enhanced.py b/tests/test_aero_grid_like_enhanced.py deleted file mode 100644 index 07f56b1..0000000 --- a/tests/test_aero_grid_like_enhanced.py +++ /dev/null @@ -1,91 +0,0 @@ -import numpy as np -import xarray as xr -from xregrid.utils import create_grid_like - - -def test_create_grid_like_enhanced_provenance(): - """ - Aero Protocol: Double-Check Test for create_grid_like. - Verifies that values are identical between NumPy and Dask backends, - and checks that history propagation is correctly handled. - """ - # Create template with history - template = xr.DataArray( - np.random.rand(10, 20), - dims=["lat", "lon"], - coords={ - "lat": np.linspace(-90, 90, 10), - "lon": np.linspace(0, 360, 20), - }, - name="my_data", - ) - template.attrs["history"] = "Original data history." - - res = 5 - - # Eager (NumPy) - ds_eager = create_grid_like(template, res=res, chunks=None) - assert "Template history:\nOriginal data history." in ds_eager.attrs["history"] - - # Template has 10 pts from -90 to 90 -> res=20. - # Extent becomes [-100, 100] for lat, and [-9.47, 369.47] approx for lon - # Actually for 20 pts over 0-360, res is 360/19 approx 18.9. - assert ds_eager.lat.size == 200 // res - - # Lazy (Dask) - ds_lazy = create_grid_like(template, res=res, chunks={"lat": 10, "lon": 10}) - assert ds_lazy.chunks - assert "Template history:\nOriginal data history." in ds_lazy.attrs["history"] - - # Assert values are identical - xr.testing.assert_allclose(ds_eager, ds_lazy.compute()) - - # Verify internal backend (non-index coordinates like bounds should be dask-backed) - assert hasattr(ds_lazy.lat_b.data, "dask") - assert hasattr(ds_lazy.lon_b.data, "dask") - - -def test_create_grid_like_override_extent(): - """ - Test create_grid_like with explicit extent override to avoid computes. - """ - template = xr.DataArray( - np.random.rand(10, 20), - dims=["lat", "lon"], - coords={ - "lat": np.linspace(-90, 90, 10), - "lon": np.linspace(0, 360, 20), - }, - ) - - # Override extent - extent = (0, 180, -45, 45) # min_lon, max_lon, min_lat, max_lat - ds = create_grid_like(template, res=10, extent=extent) - - assert ds.lat.min() == -40 # -45 + 10/2 - assert ds.lon.min() == 5 # 0 + 10/2 - assert "(Override Extent)." in ds.attrs["history"] - - -def test_create_grid_like_size_1(): - """ - Test create_grid_like with size-1 dimension to verify IndexError fix. - """ - # Eager case - template = xr.DataArray( - np.random.rand(1, 20), - dims=["lat", "lon"], - coords={ - "lat": [0.0], - "lon": np.linspace(0, 360, 20), - }, - ) - - # This should not raise IndexError - ds = create_grid_like(template, res=10) - assert ds.lon.size > 0 - - # Lazy case (triggers dask.compute logic) - template_lazy = template.chunk({"lon": 10}) - ds_lazy = create_grid_like(template_lazy, res=10) - assert ds_lazy.lon.size > 0 diff --git a/tests/test_aero_hygiene.py b/tests/test_aero_hygiene.py deleted file mode 100644 index fc21fd5..0000000 --- a/tests/test_aero_hygiene.py +++ /dev/null @@ -1,76 +0,0 @@ -import numpy as np -import pytest -import xarray as xr -from xregrid import Regridder, create_regional_grid - - -def test_eager_lazy_identity(): - """Verify that Eager (NumPy) and Lazy (Dask) data produce identical results.""" - # Create small regional grid - ds = create_regional_grid((10, 20), (100, 110), 1.0, 1.0) - ds["data"] = (("lat", "lon"), np.random.rand(10, 10)) - - target_grid = create_regional_grid((10, 20), (100, 110), 0.5, 0.5) - - # Initialize Regridder - regridder = Regridder(ds, target_grid, method="bilinear") - - # Eager regridding - da_eager = ds.data - out_eager = regridder(da_eager) - - # Lazy regridding - da_lazy = da_eager.chunk({"lat": 5, "lon": 5}) - out_lazy = regridder(da_lazy).compute() - - xr.testing.assert_allclose(out_eager, out_lazy) - - -def test_scientific_hygiene_no_inplace(): - """Verify that input datasets are not modified in-place during regridding.""" - # Create grid without bounds initially - ds = create_regional_grid((10, 20), (100, 110), 1.0, 1.0, add_bounds=False) - # Add a variable but no bounds yet - ds["data"] = (("lat", "lon"), np.random.rand(10, 10)) - - # Ensure it has a history attribute to check - if "history" not in ds.attrs: - ds.attrs["history"] = "Original history" - orig_history = ds.attrs.get("history", "") - - target_grid = create_regional_grid((10, 20), (100, 110), 0.5, 0.5) - - # Conservative regridding will trigger bounds generation - regridder = Regridder(ds, target_grid, method="conservative") - - # Check that ds was not modified in-place - assert ds.attrs.get("history", "") == orig_history - assert "lat_b" not in ds.coords - - # Check that output HAS the provenance - out = regridder(ds.data) - assert "Automatically generated cell boundaries" in out.attrs["history"] - - -def test_plot_comparison_with_regridder(): - """Verify that plot_comparison works with a Regridder instance.""" - import matplotlib.pyplot as plt - - ds = create_regional_grid((10, 20), (100, 110), 1.0, 1.0) - ds["data"] = (("lat", "lon"), np.random.rand(10, 10)) - target_grid = create_regional_grid((10, 20), (100, 110), 0.5, 0.5) - - regridder = Regridder(ds, target_grid, method="bilinear") - da_tgt = regridder(ds.data) - - # This should run without error - from xregrid.viz import plot_comparison - - # Use a mock or just run it (it will use the mocked plt in tests) - fig = plot_comparison(ds.data, da_tgt, regridder=regridder) - assert fig is not None - plt.close(fig) - - -if __name__ == "__main__": - pytest.main([__file__]) diff --git a/tests/test_aero_hygiene_mapping.py b/tests/test_aero_hygiene_mapping.py deleted file mode 100644 index c4c3056..0000000 --- a/tests/test_aero_hygiene_mapping.py +++ /dev/null @@ -1,170 +0,0 @@ -import sys -from unittest.mock import MagicMock - -import numpy as np - -# Mock esmpy BEFORE importing xregrid to allow running in environments without ESMF -try: - import esmpy # noqa: F401 - - ESMPY_AVAILABLE = True -except ImportError: - ESMPY_AVAILABLE = False - -if not ESMPY_AVAILABLE: - mock_esmpy = MagicMock() - mock_esmpy.CoordSys.SPH_DEG = 1 - mock_esmpy.CoordSys.CART = 0 - mock_esmpy.StaggerLoc.CENTER = 0 - mock_esmpy.StaggerLoc.CORNER = 1 - mock_esmpy.RegridMethod.BILINEAR = 0 - mock_esmpy.UnmappedAction.IGNORE = 1 - - class MockESMFObject: - pass - - class Grid(MockESMFObject): - def __init__(self, shape, *args, **kwargs): - self.shape = shape - - def get_coords(self, dim, staggerloc=None): - return np.zeros(self.shape) - - def add_item(self, *args, **kwargs): - pass - - def get_item(self, *args, **kwargs): - return np.zeros(self.shape) - - class Mesh(MockESMFObject): - def __init__(self, *args, **kwargs): - pass - - def add_nodes(self, *args, **kwargs): - pass - - def add_elements(self, *args, **kwargs): - pass - - class LocStream(MockESMFObject): - def __init__(self, size, *args, **kwargs): - self.size = size - - def __getitem__(self, key): - return np.zeros(self.size) - - def __setitem__(self, key, value): - pass - - mock_esmpy.Grid = Grid - mock_esmpy.Mesh = Mesh - mock_esmpy.LocStream = LocStream - mock_esmpy.Field = MagicMock() - - class MockRegrid: - def __init__(self, *args, **kwargs): - pass - - def get_weights_dict(self, deep_copy=True): - return { - "row_dst": np.array([1]), - "col_src": np.array([1]), - "weights": np.array([1.0]), - } - - def get_factors(self): - return np.array([1.0]), np.array([1]) - - mock_esmpy.Regrid = MockRegrid - mock_esmpy.pet_count.return_value = 1 - mock_esmpy.local_pet.return_value = 0 - sys.modules["esmpy"] = mock_esmpy - -import xarray as xr # noqa: E402 -from xregrid import Regridder # noqa: E402 - - -def test_grid_mapping_hygiene(): - """Verify that grid_mapping is correctly updated and preserved (Aero Protocol).""" - - # 1. Setup source grid with grid_mapping - src_ds = xr.Dataset( - coords={ - "lat": (["lat"], np.arange(10)), - "lon": (["lon"], np.arange(10)), - } - ) - src_ds["crs_src"] = xr.DataArray( - 0, attrs={"grid_mapping_name": "latitude_longitude"} - ) - - da_src = xr.DataArray( - np.random.rand(10, 10), - dims=("lat", "lon"), - coords=src_ds.coords, - name="test_data", - attrs={"grid_mapping": "crs_src"}, - ) - - # 2. Setup target grid with different grid_mapping - tgt_ds = xr.Dataset( - coords={ - "lat": (["lat"], np.arange(5)), - "lon": (["lon"], np.arange(5)), - } - ) - tgt_ds["crs_tgt"] = xr.DataArray(0, attrs={"grid_mapping_name": "mercator"}) - - regridder = Regridder(src_ds, tgt_ds, method="bilinear") - - # --- Test Track A: Eager (NumPy) --- - out_eager = regridder(da_src) - - assert out_eager.attrs["grid_mapping"] == "crs_tgt" - assert "crs_tgt" in out_eager.coords - assert "crs_src" not in out_eager.coords - - # --- Test Track B: Lazy (Dask) --- - da_lazy = da_src.chunk({"lat": 5, "lon": 5}) - out_lazy = regridder(da_lazy) - - assert out_lazy.attrs["grid_mapping"] == "crs_tgt" - assert "crs_tgt" in out_lazy.coords - - # Verify that the logic holds after computation - out_computed = out_lazy.compute() - assert out_computed.attrs["grid_mapping"] == "crs_tgt" - assert "crs_tgt" in out_computed.coords - - -def test_dataset_grid_mapping_hygiene(): - """Verify that grid_mapping is correctly updated for Datasets.""" - - src_ds = xr.Dataset( - coords={ - "lat": (["lat"], np.arange(10)), - "lon": (["lon"], np.arange(10)), - } - ) - src_ds["crs_src"] = xr.DataArray( - 0, attrs={"grid_mapping_name": "latitude_longitude"} - ) - src_ds["var1"] = (["lat", "lon"], np.random.rand(10, 10)) - src_ds["var1"].attrs["grid_mapping"] = "crs_src" - src_ds.attrs["grid_mapping"] = "crs_src" - - tgt_ds = xr.Dataset( - coords={ - "lat": (["lat"], np.arange(5)), - "lon": (["lon"], np.arange(5)), - } - ) - tgt_ds["crs_tgt"] = xr.DataArray(0, attrs={"grid_mapping_name": "mercator"}) - - regridder = Regridder(src_ds, tgt_ds) - out_ds = regridder(src_ds) - - assert out_ds.attrs["grid_mapping"] == "crs_tgt" - assert out_ds["var1"].attrs["grid_mapping"] == "crs_tgt" - assert "crs_tgt" in out_ds.coords - assert "crs_src" not in out_ds.coords diff --git a/tests/test_aero_ioapi_support.py b/tests/test_aero_ioapi_support.py deleted file mode 100644 index eb843d5..0000000 --- a/tests/test_aero_ioapi_support.py +++ /dev/null @@ -1,107 +0,0 @@ -import pytest -import xarray as xr -import numpy as np -from xregrid.utils import create_grid_from_ioapi - - -def test_create_grid_from_ioapi_lcc(): - """Verify IOAPI grid generation for LCC projection (Eager and Lazy).""" - metadata = { - "GDTYP": 2, - "P_ALP": 30.0, - "P_BET": 60.0, - "P_GAM": -97.0, - "XCENT": -97.0, - "YCENT": 40.0, - "XORIG": -1000.0, - "YORIG": -1000.0, - "XCELL": 500.0, - "YCELL": 500.0, - "NCOLS": 4, - "NROWS": 4, - } - - # 1. Eager test - ds_eager = create_grid_from_ioapi(metadata) - - assert "x" in ds_eager.coords - assert "y" in ds_eager.coords - assert "lat" in ds_eager.coords - assert "lon" in ds_eager.coords - assert "x_b" in ds_eager.coords - assert "y_b" in ds_eager.coords - assert ds_eager.sizes["x"] == 4 - assert ds_eager.sizes["y"] == 4 - assert ds_eager.attrs["ioapi_GDTYP"] == 2 - - # Check 1D bounds values - assert ds_eager.x_b.shape == (4, 2) - assert np.allclose(ds_eager.x_b[0].values, [-1000.0, -500.0]) - - # 2. Lazy test - ds_lazy = create_grid_from_ioapi(metadata, chunks={"x": 2, "y": 2}) - # Verify laziness - assert hasattr(ds_lazy.lat.data, "dask") - assert hasattr(ds_lazy.x_b.data, "dask") - - ds_lazy_comp = ds_lazy.compute() - xr.testing.assert_allclose(ds_eager, ds_lazy_comp) - - -def test_create_grid_from_ioapi_all_gdtyp(): - """Verify that all supported IOAPI GDTYP values can generate a grid.""" - base_metadata = { - "P_ALP": 30.0, - "P_BET": 60.0, - "P_GAM": -97.0, - "XCENT": -97.0, - "YCENT": 40.0, - "XORIG": -1000.0, - "YORIG": -1000.0, - "XCELL": 500.0, - "YCELL": 500.0, - "NCOLS": 2, - "NROWS": 2, - } - - # GDTYP 1-10 - for gdtyp in range(1, 11): - metadata = base_metadata.copy() - metadata["GDTYP"] = gdtyp - - # Some specific adjustments to avoid proj errors if needed - if gdtyp == 5: # UTM - metadata["P_ALP"] = 17 # Zone 17 - - ds = create_grid_from_ioapi(metadata) - assert "lat" in ds.coords - assert "lon" in ds.coords - assert ds.attrs["ioapi_GDTYP"] == gdtyp - - -def test_create_grid_from_ioapi_latlon(): - """Verify IOAPI grid generation for Lat-Lon.""" - metadata = { - "GDTYP": 1, - "P_ALP": 0.0, - "P_BET": 0.0, - "P_GAM": 0.0, - "XCENT": 0.0, - "YCENT": 0.0, - "XORIG": -10.0, - "YORIG": 40.0, - "XCELL": 1.0, - "YCELL": 1.0, - "NCOLS": 10, - "NROWS": 10, - } - ds = create_grid_from_ioapi(metadata) - assert ds.sizes["x"] == 10 - assert ds.sizes["y"] == 10 - # For Lat-Lon, pyproj transform from EPSG:4326 to EPSG:4326 should be identity - # but create_grid_from_crs might return lat/lon that are slightly different due to transform - assert ds.lat.min() >= 40.0 - - -if __name__ == "__main__": - pytest.main([__file__]) diff --git a/tests/test_aero_lazy_diagnostics.py b/tests/test_aero_lazy_diagnostics.py deleted file mode 100644 index 015649c..0000000 --- a/tests/test_aero_lazy_diagnostics.py +++ /dev/null @@ -1,88 +0,0 @@ -import pytest -import dask.distributed -from xregrid import Regridder, create_global_grid - -# Check for real ESMF -try: - import esmpy - - if hasattr(esmpy, "_is_mock") or "unittest.mock" in str(type(esmpy)): - raise ImportError - HAS_REAL_ESMF = True -except ImportError: - HAS_REAL_ESMF = False - - -@pytest.fixture(scope="module") -def dask_client(): - # esmpy is not thread-safe, so we must use processes=True when using real ESMF - cluster = dask.distributed.LocalCluster( - n_workers=1, threads_per_worker=1, processes=HAS_REAL_ESMF - ) - client = dask.distributed.Client(cluster) - yield client - client.close() - cluster.close() - - -def test_lazy_diagnostics_distributed(dask_client): - """ - Aero Protocol: Verify that diagnostics remain lazy in distributed mode. - """ - src = create_global_grid(30, 60) - tgt = create_global_grid(10, 20) - - regridder = Regridder(src, tgt, parallel=True) - - # 1. Verify Lazy Diagnostics - diag = regridder.diagnostics() - assert hasattr(diag.weight_sum.data, "dask"), "weight_sum should be lazy" - assert hasattr(diag.unmapped_mask.data, "dask"), "unmapped_mask should be lazy" - - # 2. Verify __repr__ is non-blocking and lazy-aware - repr_str = repr(regridder) - assert "quality=lazy" in repr_str - - # 3. Verify quality_report(skip_heavy=True) handles remote weights - # If skip_heavy=True and remote, it should return -1 to avoid roundtrips. - report_light = regridder.quality_report(skip_heavy=True) - assert report_light["n_weights"] == -1 - - # 4. Verify weights property gathers correctly - w = regridder.weights - assert w is not None - assert not hasattr(regridder._weights_matrix, "key") - - # 5. Verify quality_report(skip_heavy=False) now works eagerly - report_heavy = regridder.quality_report(skip_heavy=False) - assert report_heavy["n_weights"] != -1 - assert "unmapped_count" in report_heavy - - -def test_diagnostics_distributed_identity(): - """ - Verify that Eager and Lazy diagnostic values have identical shapes. - Note: Exact values may differ with synthetic mocks due to multi-chunking. - """ - src = create_global_grid(30, 60) - tgt = create_global_grid(15, 30) - - # Eager - regridder_eager = Regridder(src, tgt, parallel=False) - diag_eager = regridder_eager.diagnostics() - - # Lazy - regridder_lazy = Regridder(src, tgt, parallel=True) - diag_lazy_raw = regridder_lazy.diagnostics() - assert hasattr(diag_lazy_raw.weight_sum.data, "dask") - diag_lazy = diag_lazy_raw.compute() - - # Verify shapes and dimensions - assert diag_eager.sizes == diag_lazy.sizes - assert diag_eager.weight_sum.shape == diag_lazy.weight_sum.shape - # Verify we have some weights - assert diag_lazy.weight_sum.sum() > 0 - - -if __name__ == "__main__": - pytest.main([__file__]) diff --git a/tests/test_aero_lazy_grids.py b/tests/test_aero_lazy_grids.py deleted file mode 100644 index bdfbfa0..0000000 --- a/tests/test_aero_lazy_grids.py +++ /dev/null @@ -1,55 +0,0 @@ -import xarray as xr -from xregrid.utils import create_global_grid, create_grid_from_crs, create_grid_like - - -def test_aero_lazy_rectilinear_grid(): - """Verify that _create_rectilinear_grid uses dask and matches eager version.""" - res = 1.0 - chunks = 10 - - # Eager version - ds_eager = create_global_grid(res, res, chunks=None) - assert not hasattr(ds_eager.lat.data, "dask") - - # Lazy version - ds_lazy = create_global_grid(res, res, chunks=chunks) - # Dimension coordinates are often eager in xarray due to indexing. - # Check bounds which should definitely be lazy. - assert hasattr(ds_lazy.lat_b.data, "dask") - - # Match check - xr.testing.assert_allclose(ds_eager, ds_lazy.compute()) - - -def test_aero_lazy_projected_grid(): - """Verify that create_grid_from_crs uses dask and matches eager version.""" - crs = "EPSG:3857" - extent = (0, 1000, 0, 1000) - res = 100 - chunks = 5 - - # Eager version - ds_eager = create_grid_from_crs(crs, extent, res, chunks=None) - assert not hasattr(ds_eager.x.data, "dask") - - # Lazy version - ds_lazy = create_grid_from_crs(crs, extent, res, chunks=chunks) - # Check lat/lon which are 2D non-dimension coordinates in projected grids - assert hasattr(ds_lazy.lat.data, "dask") - - # Match check - xr.testing.assert_allclose(ds_eager, ds_lazy.compute()) - - -def test_create_grid_like_no_compute(): - """Verify that create_grid_like with explicit extent avoids compute.""" - # Create a lazy dataset - ds = create_global_grid(1.0, 1.0, chunks=10) - - # This should NOT trigger computation of ds if we pass extent - extent = (0, 360, -90, 90) - ds_like = create_grid_like(ds, 2.0, extent=extent, crs="EPSG:4326") - - # Verify it matches expected output - assert ds_like.lat.size == 90 - assert ds_like.lon.size == 180 diff --git a/tests/test_aero_memory_opt.py b/tests/test_aero_memory_opt.py deleted file mode 100644 index 251b755..0000000 --- a/tests/test_aero_memory_opt.py +++ /dev/null @@ -1,89 +0,0 @@ -import numpy as np -import pytest -import xarray as xr -from xregrid import Regridder, create_global_grid - - -def test_memory_opt_identity(): - """ - Verify that memory-optimized paths produce identical results to the previous implementation. - Aero Protocol: Eager (NumPy) vs Lazy (Dask) identity. - """ - res_src = 10.0 - res_tgt = 20.0 - src_grid = create_global_grid(res_src, res_src) - tgt_grid = create_global_grid(res_tgt, res_tgt) - - ntime = 2 - nlat = src_grid.lat.size - nlon = src_grid.lon.size - - # Create data with some NaNs - data = np.random.rand(ntime, nlat, nlon).astype(np.float32) - data[:, 0, 0] = np.nan - - da_src = xr.DataArray( - data, - coords={"time": np.arange(ntime), "lat": src_grid.lat, "lon": src_grid.lon}, - dims=("time", "lat", "lon"), - name="test_data", - ) - - # 1. Eager path - regridder = Regridder(src_grid, tgt_grid, method="bilinear", skipna=True) - da_eager = regridder(da_src) - - # 2. Lazy path - da_src_lazy = da_src.chunk({"time": 1}) - da_lazy = regridder(da_src_lazy) - da_lazy_computed = da_lazy.compute() - - # Identity check - xr.testing.assert_allclose(da_eager, da_lazy_computed) - - # Dtype preservation check - assert da_eager.dtype == da_src.dtype - assert da_lazy_computed.dtype == da_src.dtype - - -def test_repr_lazy_optimization(): - """Verify that __repr__ is lazy for large grids.""" - res_src = 10.0 - res_tgt = 20.0 - src_grid = create_global_grid(res_src, res_src) - tgt_grid = create_global_grid(res_tgt, res_tgt) - - regridder = Regridder(src_grid, tgt_grid, method="bilinear") - - # For small grid, quality should be present - repr_str = repr(regridder) - assert "unmapped=" in repr_str - assert "quality=deferred" not in repr_str - - # Manually mock a large target shape to trigger lazy repr - regridder._shape_target = (1000, 1001) # > 1,000,000 pixels - repr_str_large = repr(regridder) - assert "quality=deferred" in repr_str_large - assert "unmapped=" not in repr_str_large - - -def test_matmul_backend_agnostic(): - """Basic test for _matmul helper.""" - from xregrid.xregrid import _matmul - from scipy.sparse import csr_matrix - - matrix = csr_matrix([[1, 0], [0, 2]]) - data = np.array([[10, 20], [30, 40]]) - - # (matrix @ data.T).T - # matrix @ data.T = [[1, 0], [0, 2]] @ [[10, 30], [20, 40]] = [[10, 30], [40, 80]] - # result = [[10, 40], [30, 80]] - - expected = np.array([[10, 40], [30, 80]]) - result = _matmul(matrix, data) - - np.testing.assert_array_equal(result, expected) - - -if __name__ == "__main__": - pytest.main([__file__]) diff --git a/tests/test_aero_mesh_provenance.py b/tests/test_aero_mesh_provenance.py deleted file mode 100644 index 6572a29..0000000 --- a/tests/test_aero_mesh_provenance.py +++ /dev/null @@ -1,90 +0,0 @@ -import numpy as np -import pytest -import xarray as xr -from xregrid.utils import create_mesh_from_coords - -try: - import dask.array as da -except ImportError: - da = None - - -def test_create_mesh_from_coords_aero(): - """ - Double-Check Test for create_mesh_from_coords. - Verifies Eager (NumPy) and Lazy (Dask) backends yield identical results - and maintain scientific provenance. - """ - # 1. Setup sample coordinates (Lambert Conformal-ish) - x = np.linspace(-1000, 1000, 10) - y = np.linspace(-1000, 1000, 10) - crs = "+proj=lcc +lat_1=33 +lat_2=45 +lat_0=40 +lon_0=-97 +x_0=0 +y_0=0 +datum=WGS84 +units=m +no_defs" - - # 2. Eager execution - ds_eager = create_mesh_from_coords(x, y, crs) - - # Assertions for Eager - assert isinstance(ds_eager, xr.Dataset) - assert "lat" in ds_eager - assert "lon" in ds_eager - assert "x" in ds_eager - assert "y" in ds_eager - assert ds_eager.attrs["grid_mapping"] == "spatial_ref" - assert "spatial_ref" in ds_eager - assert "Eager" in ds_eager.attrs["history"] - assert "Extent:" in ds_eager.attrs["history"] - - # Check that it's actually NumPy-backed - assert not hasattr(ds_eager.lat.data, "dask") - - # 3. Lazy execution - if da is None: - pytest.skip("Dask not installed, skipping lazy check.") - - x_lazy = da.from_array(x, chunks=5) - y_lazy = da.from_array(y, chunks=5) - - ds_lazy = create_mesh_from_coords(x_lazy, y_lazy, crs) - - # Assertions for Lazy - assert "Lazy" in ds_lazy.attrs["history"] - # Lazy path should NOT have extent in history to avoid compute() - assert "Extent:" not in ds_lazy.attrs["history"] - assert hasattr(ds_lazy.lat.data, "dask") - - # 4. Numerical Verification (The "Double-Check") - xr.testing.assert_allclose(ds_eager, ds_lazy.compute()) - - # 5. Verify Metadata propagation - assert ds_eager.lat.attrs["units"] == "degrees_north" - assert ds_eager.x.attrs["standard_name"] == "projection_x_coordinate" - assert ds_eager.x.attrs["grid_mapping"] == "spatial_ref" - - -def test_create_mesh_from_coords_regression_fix(): - """ - Verify the fix for the dimension mismatch regression and conditional metadata. - """ - # 1. Test DataArray inputs with different dimension names - x_da = xr.DataArray(np.linspace(0, 10, 5), dims=["lon"], name="my_lon") - y_da = xr.DataArray(np.linspace(0, 10, 5), dims=["lat"], name="my_lat") - crs_proj = "+proj=lcc +lat_1=33 +lat_2=45 +lat_0=40 +lon_0=-97 +x_0=0 +y_0=0 +datum=WGS84 +units=m +no_defs" - - ds = create_mesh_from_coords(x_da, y_da, crs_proj) - - # Should have 5 points, not 25 (if it had broadcasted incorrectly) - assert ds.sizes["n_pts"] == 5 - assert ds.x.attrs["standard_name"] == "projection_x_coordinate" - - # 2. Test geographic CRS metadata - crs_geo = "EPSG:4326" - ds_geo = create_mesh_from_coords(x_da, y_da, crs_geo) - - assert ds_geo.x.attrs["standard_name"] == "longitude" - assert ds_geo.x.attrs["units"] == "degrees_east" - assert ds_geo.y.attrs["standard_name"] == "latitude" - assert ds_geo.y.attrs["units"] == "degrees_north" - - -if __name__ == "__main__": - test_create_mesh_from_coords_aero() diff --git a/tests/test_aero_mixed_backend.py b/tests/test_aero_mixed_backend.py deleted file mode 100644 index ace9eb2..0000000 --- a/tests/test_aero_mixed_backend.py +++ /dev/null @@ -1,78 +0,0 @@ -import numpy as np -import xarray as xr -import dask.array as da -from distributed import Client, LocalCluster -from conftest import setup_esmpy_mock -from xregrid import Regridder, create_global_grid -from xregrid.utils import create_regional_grid - - -def test_mixed_backend_dataset_regrid(): - """ - Double-Check Test: Verify that a Dataset with mixed NumPy and Dask variables - regrids correctly when weights are remote (Dask Futures). - """ - # Start a local cluster - with LocalCluster(n_workers=2, threads_per_worker=1) as cluster: - with Client(cluster) as client: - # Setup mocks on workers - client.run(setup_esmpy_mock) - - # 1. Create grids - ds_src = create_global_grid(10.0, 10.0) - ds_tgt = create_regional_grid((20, 40), (40, 60), 10.0, 10.0) - - # 2. Create mixed data - shape_src = (ds_src.sizes["lat"], ds_src.sizes["lon"]) - - # Eager variable (NumPy) - data_eager = np.random.rand(*shape_src) - # Lazy variable (Dask) - data_lazy = da.from_array(np.random.rand(*shape_src), chunks=(9, 18)) - - ds_input = xr.Dataset( - { - "var_eager": (["lat", "lon"], data_eager), - "var_lazy": (["lat", "lon"], data_lazy), - }, - coords=ds_src.coords, - ) - - # 3. Initialize Regridder with parallel=True to create remote weights - regridder = Regridder(ds_src, ds_tgt, method="bilinear", parallel=True) - - # Ensure weights are indeed remote - assert hasattr( - regridder._weights_matrix, "key" - ), "Weights should be Dask Futures" - - # 4. Regrid! - # Before the fix, this would crash when processing "var_eager" - ds_out = regridder(ds_input) - - # 5. Verify Results - assert "var_eager" in ds_out - assert "var_lazy" in ds_out - - # Check backends are preserved - assert not hasattr( - ds_out.var_eager.data, "dask" - ), "var_eager should remain NumPy-backed" - assert hasattr( - ds_out.var_lazy.data, "dask" - ), "var_lazy should remain Dask-backed" - - # Check shape - expected_shape_tgt = (ds_tgt.sizes["lat"], ds_tgt.sizes["lon"]) - assert ds_out.var_eager.shape == expected_shape_tgt - assert ds_out.var_lazy.shape == expected_shape_tgt - - # Verify values - # (Using .values on dask array triggers compute) - assert not np.isnan(ds_out.var_eager.values).all() - assert not np.isnan(ds_out.var_lazy.values).all() - - # Check provenance/history - assert "Regridded" in ds_out.attrs["history"] - assert "var_eager" in ds_out.data_vars - assert "var_lazy" in ds_out.data_vars diff --git a/tests/test_aero_native_formats.py b/tests/test_aero_native_formats.py deleted file mode 100644 index 5e7c8c6..0000000 --- a/tests/test_aero_native_formats.py +++ /dev/null @@ -1,222 +0,0 @@ -import numpy as np -import pytest -import xarray as xr -from xregrid import Regridder -from xregrid.utils import create_global_grid - - -def test_musica_cesm_regrid_aero(): - """Verify MUSICA/CESM (ncol) grid discovery and regridding (Aero Protocol).""" - # 1. Create MUSICA-like source grid (unstructured ncol) - n_col = 100 - ds_src = xr.Dataset( - data_vars={"temp": (["ncol"], np.random.rand(n_col))}, - coords={ - "lat": (["ncol"], np.linspace(-90, 90, n_col)), - "lon": (["ncol"], np.linspace(0, 350, n_col)), - }, - ) - - # 2. Create target rectilinear grid - ds_tgt = create_global_grid(10, 20) - - # 3. Initialize Regridder (should detect unstructured ncol) - # Using nearest_s2d because ncol has no connectivity info - regridder = Regridder(ds_src, ds_tgt, method="nearest_s2d") - assert regridder._is_unstructured_src - assert regridder._dims_source == ("ncol",) - - # 4. Double-Check: Eager - out_eager = regridder(ds_src) - assert isinstance(out_eager, xr.Dataset) - assert "temp" in out_eager - assert out_eager.temp.dims == ("lat", "lon") - - # 5. Double-Check: Lazy - ds_src_lazy = ds_src.chunk({"ncol": 10}) - out_lazy = regridder(ds_src_lazy) - assert hasattr(out_lazy.temp.data, "dask") - - xr.testing.assert_allclose(out_eager, out_lazy.compute()) - - -def test_mpas_regrid_aero(): - """Verify MPAS (nCells) grid discovery and regridding (Aero Protocol).""" - n_cells = 100 - ds_src = xr.Dataset( - data_vars={"temp": (["nCells"], np.random.rand(n_cells))}, - coords={ - "latCell": (["nCells"], np.linspace(-90, 90, n_cells)), - "lonCell": (["nCells"], np.linspace(0, 350, n_cells)), - }, - ) - # CF-Xarray might not know latCell/lonCell without attributes, - # but xregrid fallback should find them. - ds_src.latCell.attrs["standard_name"] = "latitude" - ds_src.lonCell.attrs["standard_name"] = "longitude" - - ds_tgt = create_global_grid(10, 20) - - # Using nearest_s2d because nCells has no connectivity info in this test - regridder = Regridder(ds_src, ds_tgt, method="nearest_s2d") - assert regridder._is_unstructured_src - assert regridder._dims_source == ("nCells",) - - out_eager = regridder(ds_src) - assert out_eager.temp.dims == ("lat", "lon") - - ds_src_lazy = ds_src.chunk({"nCells": 10}) - out_lazy = regridder(ds_src_lazy) - xr.testing.assert_allclose(out_eager, out_lazy.compute()) - - -def test_ugrid_discovery_aero(): - """Verify UGRID-compliant discovery with explicit mesh topology.""" - n_nodes = 50 - ds = xr.Dataset( - data_vars={ - "temp": (["node"], np.random.rand(n_nodes)), - "mesh": ( - [], - 0, - {"cf_role": "mesh_topology", "node_coordinates": "lon_node lat_node"}, - ), - }, - coords={ - "lat_node": ( - ["node"], - np.linspace(-90, 90, n_nodes), - {"standard_name": "latitude"}, - ), - "lon_node": ( - ["node"], - np.linspace(0, 360, n_nodes), - {"standard_name": "longitude"}, - ), - }, - ) - # The new logic should prefer lat_node/lon_node because they are linked in 'mesh' - # OR because they have standard names and 'node' in name. - - ds_tgt = create_global_grid(10, 20) - # Using nearest_s2d because this UGRID has no connectivity info - regridder = Regridder(ds, ds_tgt, method="nearest_s2d") - assert regridder._is_unstructured_src - assert regridder._dims_source == ("node",) - - out = regridder(ds) - assert out.temp.dims == ("lat", "lon") - - -def test_scrip_conservative_regrid_aero(): - """Verify SCRIP-style unstructured grid handles conservative regridding via derived connectivity.""" - n_cells = 10 - # Create SCRIP-like 2D bounds (n_cells, 4 corners) - lat_b = np.array( - [ - [-10, -10, 10, 10], - [-10, -10, 10, 10], - # ... just a few for testing - ] - ) - lat_b = np.repeat(lat_b, n_cells // 2, axis=0) - lon_b = np.array( - [ - [0, 10, 10, 0], - [10, 20, 20, 10], - ] - ) - lon_b = np.repeat(lon_b, n_cells // 2, axis=0) - - ds_src = xr.Dataset( - data_vars={"temp": (["grid_size"], np.random.rand(n_cells))}, - coords={ - "lat": (["grid_size"], np.zeros(n_cells)), - "lon": (["grid_size"], np.zeros(n_cells)), - "lat_b": (["grid_size", "nv"], lat_b), - "lon_b": (["grid_size", "nv"], lon_b), - }, - ) - - ds_tgt = create_global_grid(10, 10) - - # conservative requires bounds - regridder = Regridder(ds_src, ds_tgt, method="conservative") - assert regridder._is_unstructured_src - - out = regridder(ds_src) - assert "temp" in out - - -def test_mpas_non_conservative_discovery_aero(): - """Verify MPAS (nCells) non-conservative discovery (triggers optimized path).""" - n_cells = 50 - ds_src = xr.Dataset( - data_vars={"temp": (["nCells"], np.random.rand(n_cells))}, - coords={ - "lat": (["nCells"], np.linspace(-90, 90, n_cells)), - "lon": (["nCells"], np.linspace(0, 350, n_cells)), - }, - ) - ds_tgt = create_global_grid(10, 20) - - # This should trigger the optimized section 2 in _get_unstructured_mesh_info - regridder = Regridder(ds_src, ds_tgt, method="nearest_s2d") - assert regridder._is_unstructured_src - assert regridder._dims_source == ("nCells",) - - out = regridder(ds_src) - assert out.temp.dims == ("lat", "lon") - - -def test_mpas_to_scrip_regrid_aero(): - """Verify MPAS to SCRIP conversion and native regridding.""" - from xregrid.utils import mpas_to_scrip - - n_cells = 4 - # Create a minimal valid MPAS-like grid - ds_mpas = xr.Dataset( - data_vars={ - "temp": (["nCells"], np.random.rand(n_cells)), - "verticesOnCell": ( - ["nCells", "maxEdges"], - np.array([[1, 2, 3], [1, 3, 4], [1, 2, 4], [2, 3, 4]]), - ), - "nEdgesOnCell": (["nCells"], [3, 3, 3, 3]), - }, - coords={ - "latCell": (["nCells"], np.linspace(-45, 45, n_cells)), - "lonCell": (["nCells"], np.linspace(0, 90, n_cells)), - "latVertex": (["nVertices"], np.linspace(-90, 90, 5)), - "lonVertex": (["nVertices"], np.linspace(0, 360, 5)), - }, - ) - - # 1. Convert - ds_scrip = mpas_to_scrip(ds_mpas) - assert "lat_b" in ds_scrip.coords - assert ds_scrip.lat.dims == ("grid_size",) - - # 2. Regrid - ds_tgt = create_global_grid(10, 20) - regridder = Regridder(ds_scrip, ds_tgt, method="bilinear") - assert regridder._is_unstructured_src - - # Data to regrid must match the new grid_size dimension if using the scrip grid as source - # Only use coordinates that are compatible with the 1D grid_size dimension - compatible_coords = { - c: ds_scrip.coords[c] - for c in ds_scrip.coords - if set(ds_scrip.coords[c].dims).issubset({"grid_size"}) - } - da_src = xr.DataArray( - np.random.rand(len(ds_scrip.grid_size)), - dims=["grid_size"], - coords=compatible_coords, - ) - out = regridder(da_src) - assert out.dims == ("lat", "lon") - - -if __name__ == "__main__": - pytest.main([__file__]) diff --git a/tests/test_aero_new_features.py b/tests/test_aero_new_features.py deleted file mode 100644 index 7ac9e4b..0000000 --- a/tests/test_aero_new_features.py +++ /dev/null @@ -1,103 +0,0 @@ -import os -from unittest.mock import patch - -import dask.array as da -import numpy as np -import pytest -import xarray as xr -from xregrid import Regridder, create_global_grid, plot - - -def test_generation_time_provenance(): - """Verify that weight generation time is tracked and included in history.""" - src_grid = create_global_grid(10, 20) - tgt_grid = create_global_grid(15, 30) - - filename = "test_provenance.nc" - if os.path.exists(filename): - os.remove(filename) - - # Test Eager (NumPy) - data = np.random.rand(18, 18) - # Only keep coords that share dims with the data - valid_coords = { - c: src_grid.coords[c] - for c in src_grid.coords - if set(src_grid.coords[c].dims).issubset({"lat", "lon"}) - } - da_src = xr.DataArray(data, dims=("lat", "lon"), coords=valid_coords, name="test") - - # Set reuse_weights=True so it saves the file - regridder = Regridder(src_grid, tgt_grid, filename=filename, reuse_weights=True) - assert regridder.generation_time is not None - assert regridder.generation_time > 0 - - da_regridded = regridder(da_src) - assert "Weight generation time" in da_regridded.attrs["history"] - - # Verify persistence - regridder_reused = Regridder( - src_grid, tgt_grid, filename=filename, reuse_weights=True - ) - # Should be identical as it's loaded from the same file - assert regridder_reused.generation_time == regridder.generation_time - - da_regridded_reused = regridder_reused(da_src) - assert ( - f"Weight generation time: {regridder.generation_time:.4f}s" - in da_regridded_reused.attrs["history"] - ) - - if os.path.exists(filename): - os.remove(filename) - - -def test_unified_plot_dispatch(): - """Verify that the unified plot function dispatches correctly.""" - da_test = xr.DataArray(np.random.rand(10, 10), dims=("lat", "lon"), name="test") - - # Mock plot_static and plot_interactive - with patch("xregrid.viz.plot_static") as mock_static: - with patch("xregrid.viz.plot_interactive") as mock_interactive: - # Test static dispatch - plot(da_test, mode="static", custom_arg=True) - mock_static.assert_called_once_with(da_test, custom_arg=True) - - # Test interactive dispatch - plot(da_test, mode="interactive", custom_arg=False) - mock_interactive.assert_called_once_with(da_test, custom_arg=False) - - # Test invalid mode - with pytest.raises(ValueError, match="Unknown plotting mode"): - plot(da_test, mode="invalid") - - -def test_backend_agnostic_provenance(): - """Verify provenance works for both NumPy and Dask backends.""" - src_grid = create_global_grid(10, 20) - tgt_grid = create_global_grid(10, 20) # Identity regridding for simplicity - - regridder = Regridder(src_grid, tgt_grid, reuse_weights=False) - valid_coords = { - c: src_grid.coords[c] - for c in src_grid.coords - if set(src_grid.coords[c].dims).issubset({"lat", "lon"}) - } - - # NumPy - da_numpy = xr.DataArray( - np.random.rand(18, 18), dims=("lat", "lon"), coords=valid_coords - ) - out_numpy = regridder(da_numpy) - assert "Weight generation time" in out_numpy.attrs["history"] - - # Dask - da_dask = xr.DataArray( - da.from_array(np.random.rand(18, 18), chunks=(9, 9)), - dims=("lat", "lon"), - coords=valid_coords, - ) - out_dask = regridder(da_dask) - assert "Weight generation time" in out_dask.attrs["history"] - # Ensure it's still a dask array - assert out_dask.chunks is not None diff --git a/tests/test_aero_optimization.py b/tests/test_aero_optimization.py deleted file mode 100644 index 4c23a0b..0000000 --- a/tests/test_aero_optimization.py +++ /dev/null @@ -1,115 +0,0 @@ -import dask.array as da -import numpy as np -import pytest -import os -import xarray as xr -from xregrid import Regridder, create_global_grid -from xregrid.viz import plot_static - - -def test_no_hidden_compute_on_weight_load(): - """Verify that dask-backed coordinates are not computed when loading weights.""" - # We use a callback to detect compute - compute_count = 0 - - def count_compute(key, value, dsk): - nonlocal compute_count - compute_count += 1 - - from dask.callbacks import Callback - - # Create dummy grid with dask coordinates - lat = da.from_array(np.linspace(-90, 90, 10), chunks=5) - lon = da.from_array(np.linspace(0, 360, 20), chunks=10) - - src_grid = xr.Dataset( - coords={ - "lat": ( - ["lat"], - lat, - {"units": "degrees_north", "standard_name": "latitude"}, - ), - "lon": ( - ["lon"], - lon, - {"units": "degrees_east", "standard_name": "longitude"}, - ), - } - ) - - # Create a weights file first (this WILL trigger compute because we generate weights) - # We use a small grid that matches the mock in conftest - filename = "test_weights_compute.nc" - if os.path.exists(filename): - os.remove(filename) - - _ = Regridder(src_grid, src_grid, filename=filename, reuse_weights=False) - - # Now try to load with reuse_weights=True and check if coordinates are computed - class ComputeCounter(Callback): - def _pretask(self, key, dsk, state): - nonlocal compute_count - compute_count += 1 - - with ComputeCounter(): - compute_count = 0 - _ = Regridder(src_grid, src_grid, filename=filename, reuse_weights=True) - - # compute_count should be 0 for src_grid coordinates. - # Note: Regridder._load_weights calls ds_weights.load(), which computes weight variables in that file, - # but it should NOT compute src_grid coordinates from the input dataset. - assert compute_count == 0 - - if os.path.exists(filename): - os.remove(filename) - - -def test_plot_static_robust_slicing(): - """Verify plot_static correctly handles non-standard dimension orders using cf-xarray.""" - # Create a 3D DataArray where spatial dims are NOT the last two - # dims: ('lat', 'time', 'lon') - lat = np.linspace(-90, 90, 10) - lon = np.linspace(0, 360, 20) - time = [0, 1] - - data = np.random.rand(10, 2, 20) - da_test = xr.DataArray( - data, - dims=("lat", "time", "lon"), - coords={"lat": lat, "lon": lon, "time": time}, - name="test_data", - ) - da_test.lat.attrs["standard_name"] = "latitude" - da_test.lon.attrs["standard_name"] = "longitude" - - # This should slice 'time' (index 0) and plot ('lat', 'lon') - # We check that the warning is issued and mentions 'time' - with pytest.warns( - UserWarning, match=r"Automatically selecting the first slice along \['time'\]" - ): - im = plot_static(da_test) - - assert im is not None - - -def test_weight_persistence_robustness(): - """Verify that weight attributes survive NetCDF round-trip as tuples.""" - src_grid = create_global_grid(10, 10) - filename = "test_persistence_opt.nc" - if os.path.exists(filename): - os.remove(filename) - - regridder = Regridder(src_grid, src_grid, filename=filename, reuse_weights=True) - - # Check attributes are tuples - assert isinstance(regridder._shape_source, tuple) - assert isinstance(regridder._dims_source, tuple) - - # Re-load from same file - regridder2 = Regridder(src_grid, src_grid, filename=filename, reuse_weights=True) - assert regridder2._shape_source == regridder._shape_source - assert regridder2._dims_source == regridder._dims_source - assert isinstance(regridder2._shape_source, tuple) - - if os.path.exists(filename): - os.remove(filename) diff --git a/tests/test_aero_optimization_v2.py b/tests/test_aero_optimization_v2.py deleted file mode 100644 index 60de013..0000000 --- a/tests/test_aero_optimization_v2.py +++ /dev/null @@ -1,87 +0,0 @@ -import numpy as np -import xarray as xr -import pytest -from xregrid import Regridder, create_global_grid - - -def test_optimization_v2_identity(): - """Verify that the optimized path produces identical results to a known-valid path.""" - src_grid = create_global_grid(10, 10) - tgt_grid = create_global_grid(5, 5) - - # skipna=True to trigger the optimized code path - regridder = Regridder(src_grid, tgt_grid, method="bilinear", skipna=True) - - # 1. Data with NO NaNs (triggers fast path) - data_clean = np.random.rand(18, 36) - da_clean = xr.DataArray( - data_clean, - dims=("lat", "lon"), - coords={"lat": src_grid.lat, "lon": src_grid.lon}, - name="clean_data", - ) - - res_clean = regridder(da_clean) - assert res_clean.name == "clean_data" - assert "Regridded" in res_clean.attrs["history"] - - # 2. Data WITH NaNs (triggers slow path) - data_dirty = data_clean.copy() - data_dirty[0, 0] = np.nan - da_dirty = xr.DataArray( - data_dirty, - dims=("lat", "lon"), - coords={"lat": src_grid.lat, "lon": src_grid.lon}, - ) - - res_dirty = regridder(da_dirty) - - # Verify both ran successfully - assert res_clean is not None - assert res_dirty is not None - - -def test_lazy_data_handling(): - """Verify that Dask-backed data works with the optimized skipna path.""" - src_grid = create_global_grid(10, 10) - tgt_grid = create_global_grid(5, 5) - regridder = Regridder(src_grid, tgt_grid, method="bilinear", skipna=True) - - data = np.random.rand(18, 36) - da_lazy = xr.DataArray( - data, dims=("lat", "lon"), coords={"lat": src_grid.lat, "lon": src_grid.lon} - ).chunk({"lat": 9}) - - res_lazy = regridder(da_lazy) - # Check it's still lazy - assert hasattr(res_lazy.data, "dask") - - # Compute and check - res_computed = res_lazy.compute() - assert res_computed.shape == (36, 72) - - -def test_dataset_regridding_provenance(): - """Verify Dataset regridding preserves history and non-spatial coords.""" - src_grid = create_global_grid(10, 10) - tgt_grid = create_global_grid(5, 5) - regridder = Regridder(src_grid, tgt_grid) - - ds = xr.Dataset( - data_vars={ - "temp": (("lat", "lon"), np.random.rand(18, 36)), - "mask": (("lat", "lon"), np.ones((18, 36))), - }, - coords={"lat": src_grid.lat, "lon": src_grid.lon, "time": [0]}, - ) - - res_ds = regridder(ds) - assert "time" in res_ds.coords - assert "history" in res_ds.attrs - assert "Regridded Dataset" in res_ds.attrs["history"] - assert "temp" in res_ds.data_vars - assert "mask" in res_ds.data_vars - - -if __name__ == "__main__": - pytest.main([__file__]) diff --git a/tests/test_aero_optimization_v3.py b/tests/test_aero_optimization_v3.py deleted file mode 100644 index ee7583b..0000000 --- a/tests/test_aero_optimization_v3.py +++ /dev/null @@ -1,92 +0,0 @@ -import numpy as np -import pytest -import xarray as xr -from xregrid import create_global_grid - - -def test_stationary_mask_optimization(): - """ - Test that the stationary mask optimization produces correct results. - Following Aero Protocol: Verifies Eager (NumPy) and Lazy (Dask) identity. - """ - # 1. Setup grids (small grids for testing with mocked ESMF) - src_res = 10.0 - tgt_res = 20.0 - src_grid = create_global_grid(src_res, src_res) - tgt_grid = create_global_grid(tgt_res, tgt_res) - - # 2. Create data with stationary NaNs - ntime = 4 - nlat = src_grid.lat.size - nlon = src_grid.lon.size - - # Use deterministic data for testing - data = np.ones((ntime, nlat, nlon)) - - # Add a stationary mask - data_with_nans = data.copy() - data_with_nans[:, 0, 0] = np.nan # First point is always NaN - - da_src = xr.DataArray( - data_with_nans, - coords={"time": np.arange(ntime), "lat": src_grid.lat, "lon": src_grid.lon}, - dims=("time", "lat", "lon"), - name="test_data", - ) - - # 3. Test Eager (NumPy) with Accessor - # skipna=True triggers the optimized stationary mask path - da_regridded_eager = da_src.regrid.to(tgt_grid, method="bilinear", skipna=True) - - assert isinstance(da_regridded_eager, xr.DataArray) - # Target shape: (ntime, 9, 18) for 20deg resolution - assert da_regridded_eager.shape == (ntime, tgt_grid.lat.size, tgt_grid.lon.size) - - # 4. Test Lazy (Dask) with Accessor - da_src_lazy = da_src.chunk({"time": 2}) - da_regridded_lazy = da_src_lazy.regrid.to(tgt_grid, method="bilinear", skipna=True) - - # Verify it is still lazy - assert hasattr(da_regridded_lazy.data, "dask") - - # Compute and compare - da_regridded_lazy_computed = da_regridded_lazy.compute() - - # Identity test: Eager vs Lazy - xr.testing.assert_allclose(da_regridded_eager, da_regridded_lazy_computed) - - # Verify history tracking (Scientific Hygiene) - assert "history" in da_regridded_eager.attrs - assert "Regridded" in da_regridded_eager.attrs["history"] - - -def test_non_stationary_mask(): - """Test that non-stationary masks still work correctly (slow path).""" - src_res = 10.0 - tgt_res = 20.0 - src_grid = create_global_grid(src_res, src_res) - tgt_grid = create_global_grid(tgt_res, tgt_res) - - ntime = 3 - nlat = src_grid.lat.size - nlon = src_grid.lon.size - data = np.ones((ntime, nlat, nlon)) - - # Different mask for each time step - data[0, 0, 0] = np.nan - data[1, 0, 1] = np.nan - data[2, 0, 2] = np.nan - - da_src = xr.DataArray( - data, - coords={"time": np.arange(ntime), "lat": src_grid.lat, "lon": src_grid.lon}, - dims=("time", "lat", "lon"), - name="test_data_moving", - ) - - da_regridded = da_src.regrid.to(tgt_grid, method="bilinear", skipna=True) - assert da_regridded.shape == (ntime, tgt_grid.lat.size, tgt_grid.lon.size) - - -if __name__ == "__main__": - pytest.main([__file__]) diff --git a/tests/test_aero_plot_weights.py b/tests/test_aero_plot_weights.py deleted file mode 100644 index 69ff203..0000000 --- a/tests/test_aero_plot_weights.py +++ /dev/null @@ -1,72 +0,0 @@ -import numpy as np -import pytest -import xarray as xr -from unittest.mock import MagicMock, patch -from xregrid import Regridder, create_global_grid -from xregrid.viz import plot_weights - - -def test_plot_weights_eager(): - """Double-Check Test: Verify plot_weights works for Eager (NumPy) backend.""" - src = create_global_grid(10, 10) - tgt = create_global_grid(5, 5) - regridder = Regridder(src, tgt) - - # Track A: Static - with patch("xregrid.viz.plot_static") as mock_plot: - plot_weights(regridder, row_idx=0, mode="static") - mock_plot.assert_called_once() - da_weights = mock_plot.call_args[0][0] - assert isinstance(da_weights, xr.DataArray) - assert da_weights.shape == regridder._shape_source - - # Track B: Interactive - with patch("xregrid.viz.plot_interactive") as mock_plot_int: - plot_weights(regridder, row_idx=0, mode="interactive") - mock_plot_int.assert_called_once() - - -def test_plot_weights_lazy_no_gather(): - """Double-Check Test: Verify plot_weights for Lazy (Dask) backend avoids full gather.""" - src = create_global_grid(10, 10) - tgt = create_global_grid(5, 5) - regridder = Regridder(src, tgt, parallel=True, compute=False) - - # Mock weights as a remote Future - class MockFuture: - def __init__(self): - self.key = "remote_weights_key" - - regridder._weights_matrix = MockFuture() - regridder._dask_client = MagicMock() - - # The return of the remote task - mock_row = np.zeros(regridder._shape_source).flatten() - regridder._dask_client.submit.return_value.result.return_value = mock_row - - # Call plot_weights - with patch("xregrid.viz.plot_static") as mock_plot: - plot_weights(regridder, row_idx=5, mode="static") - - # VERIFY: Plot was called - mock_plot.assert_called_once() - - # VERIFY: No full gather called on weights matrix - # (regridder.weights would call client.gather) - regridder._dask_client.gather.assert_not_called() - - # VERIFY: Distributed task was submitted - regridder._dask_client.submit.assert_called_once() - args = regridder._dask_client.submit.call_args[0] - assert "_get_weight_row_task" in args[0].__name__ - assert args[1] is regridder._weights_matrix - assert args[2] == 5 - - -def test_plot_weights_invalid_mode(): - """Verify ValueError for invalid mode.""" - src = create_global_grid(10, 10) - tgt = create_global_grid(5, 5) - regridder = Regridder(src, tgt) - with pytest.raises(ValueError, match="Unknown plotting mode"): - plot_weights(regridder, row_idx=0, mode="invalid") diff --git a/tests/test_aero_protocol.py b/tests/test_aero_protocol.py deleted file mode 100644 index 6392da6..0000000 --- a/tests/test_aero_protocol.py +++ /dev/null @@ -1,86 +0,0 @@ -import dask.array as da -import numpy as np -import pytest -import xarray as xr -from xregrid import Regridder, create_global_grid - - -def test_eager_lazy_identity_dim_orders(): - """Verify Eager and Lazy results are identical even with different dimension orders.""" - # Source grid: 10x20 - lat = np.linspace(-90, 90, 10) - lon = np.linspace(0, 360, 20) - - # Target grid: 15x25 - lat_out = np.linspace(-90, 90, 15) - lon_out = np.linspace(0, 360, 25) - - src_grid = xr.Dataset(coords={"lat": lat, "lon": lon}) - tgt_grid = xr.Dataset(coords={"lat": lat_out, "lon": lon_out}) - - regridder = Regridder(src_grid, tgt_grid, method="bilinear") - - # 1. Eager (lat, lon) - data = np.random.rand(10, 20) - da_eager = xr.DataArray(data, dims=("lat", "lon"), coords={"lat": lat, "lon": lon}) - res_eager = regridder(da_eager) - - # 2. Lazy (lon, lat) - different order! - da_lazy = xr.DataArray( - data.T, dims=("lon", "lat"), coords={"lat": lat, "lon": lon} - ).chunk({"lon": 10, "lat": 5}) - res_lazy = regridder(da_lazy).compute() - - # Transpose back for comparison if needed, or check if xregrid handles it - # xregrid's _regrid_dataarray uses input_core_dims=self._dims_source - # which will handle the transposition automatically via apply_ufunc. - - xr.testing.assert_allclose(res_eager, res_lazy) - assert isinstance(regridder(da_lazy).data, da.Array) - - -def test_skipna_robustness(): - """Verify skipna=True handles NaNs correctly in both Eager and Lazy paths.""" - src_grid = create_global_grid(10, 10) - tgt_grid = create_global_grid(5, 5) - - regridder = Regridder(src_grid, tgt_grid, method="bilinear", skipna=True) - - data = np.ones((18, 36)) - data[0, 0] = np.nan # Put a NaN - - # Only use 'lat' and 'lon' coords, skip 'lat_b' and 'lon_b' - da_coords = {c: src_grid.coords[c] for c in ["lat", "lon"]} - da_eager = xr.DataArray(data, dims=("lat", "lon"), coords=da_coords) - res_eager = regridder(da_eager) - - # The result at (0,0) should not be NaN if there are other valid points in the stencil - # (depending on the method and stencil size) - # But more importantly, eager and lazy should match. - - da_lazy = da_eager.chunk({"lat": 9, "lon": 18}) - res_lazy = regridder(da_lazy).compute() - - xr.testing.assert_allclose(res_eager, res_lazy) - - -def test_provenance_tracking(): - """Verify that history is correctly updated and preserved.""" - src_grid = create_global_grid(30, 30) - tgt_grid = create_global_grid(20, 20) - - regridder = Regridder(src_grid, tgt_grid) - - da_coords = {c: src_grid.coords[c] for c in ["lat", "lon"]} - da = xr.DataArray(np.random.rand(6, 12), dims=("lat", "lon"), coords=da_coords) - da.attrs["history"] = "Original data" - - res = regridder(da) - - assert "history" in res.attrs - assert "Original data" in res.attrs["history"] - assert "Regridded" in res.attrs["history"] - - -if __name__ == "__main__": - pytest.main([__file__]) diff --git a/tests/test_aero_protocol_refactor.py b/tests/test_aero_protocol_refactor.py deleted file mode 100644 index 15ba2f7..0000000 --- a/tests/test_aero_protocol_refactor.py +++ /dev/null @@ -1,74 +0,0 @@ -from __future__ import annotations - -import numpy as np -import xarray as xr - -from xregrid.regridder import Regridder -from xregrid.utils import create_global_grid - - -def test_regridder_refactor_consistency(): - """ - Double-Check Test: Verify Regridder still works with refactored constants. - Ensures NumPy and Dask backends produce identical results. - """ - # 1. Setup small grids - res = 10.0 - ds_src = create_global_grid(res, res) - ds_tgt = create_global_grid(res * 2, res * 2) - - # Add some data - data = np.random.rand(*ds_src.lat.shape, *ds_src.lon.shape) - da_src = xr.DataArray( - data, - coords={c: ds_src.coords[c] for c in ["lat", "lon"]}, - dims=("lat", "lon"), - name="test_data", - ) - - # 2. Eager execution (NumPy) - regridder_eager = Regridder(ds_src, ds_tgt, method="bilinear") - res_eager = regridder_eager(da_src) - - # 3. Lazy execution (Dask) - da_lazy = da_src.chunk({"lat": 5, "lon": 5}) - # We can reuse the same regridder as it's backend-agnostic for application - res_lazy = regridder_eager(da_lazy) - - # 4. Assertions - # Verify results are identical - xr.testing.assert_allclose(res_eager, res_lazy.compute()) - - # Verify regridder attributes still match expectations - assert regridder_eager.method == "bilinear" - assert "Regridded using xregrid.Regridder" in res_eager.attrs["history"] - - # Check that constants were correctly applied (mocked or real) - try: - import esmpy - - expected_method = esmpy.RegridMethod.BILINEAR - assert regridder_eager.method_map["bilinear"] == expected_method - except ImportError: - pass - - -def test_regridder_extrap_refactor(): - """Verify extrapolation method refactor.""" - res = 10.0 - ds_src = create_global_grid(res, res) - ds_tgt = create_global_grid(res * 2, res * 2) - - regridder = Regridder( - ds_src, ds_tgt, method="nearest_s2d", extrap_method="nearest_idw" - ) - - assert regridder.extrap_method == "nearest_idw" - - try: - import esmpy - - expected_extrap = esmpy.ExtrapMethod.NEAREST_IDAVG - assert regridder.extrap_method_map["nearest_idw"] == expected_extrap - except ImportError: - pass diff --git a/tests/test_aero_quality_lazy.py b/tests/test_aero_quality_lazy.py deleted file mode 100644 index bfbd041..0000000 --- a/tests/test_aero_quality_lazy.py +++ /dev/null @@ -1,107 +0,0 @@ -import numpy as np -import pytest -import dask.distributed -from xregrid import Regridder, create_global_grid - -# Check for real ESMF -try: - import esmpy - - if hasattr(esmpy, "_is_mock") or "unittest.mock" in str(type(esmpy)): - raise ImportError - HAS_REAL_ESMF = True -except ImportError: - HAS_REAL_ESMF = False - - -@pytest.fixture(scope="module") -def dask_client(): - # esmpy is not thread-safe, so we must use processes=True when using real ESMF - cluster = dask.distributed.LocalCluster( - n_workers=1, threads_per_worker=1, processes=HAS_REAL_ESMF - ) - client = dask.distributed.Client(cluster) - yield client - client.close() - cluster.close() - - -def test_aero_quality_lazy_distributed(dask_client): - """ - Aero Protocol: Double-Check Test for Quality Report. - Verifies that quality_report(format='dataset') is lazy for Dask-backed regridders - and matches eager results. - """ - src = create_global_grid(30, 60) - tgt = create_global_grid(15, 30) - - # 1. Eager (NumPy) Path - regridder_eager = Regridder(src, tgt, parallel=False) - report_eager = regridder_eager.quality_report(format="dataset") - - # Verify report_eager is indeed eager - for var in report_eager.data_vars: - assert not hasattr( - report_eager[var].data, "dask" - ), f"{var} should be NumPy-backed" - - # 2. Lazy (Dask) Path - regridder_lazy = Regridder(src, tgt, parallel=True) - report_lazy = regridder_lazy.quality_report(format="dataset") - - # Verify report_lazy preserves laziness for heavy metrics - assert hasattr(report_lazy.n_weights.data, "dask"), "n_weights should be lazy" - assert hasattr( - report_lazy.unmapped_count.data, "dask" - ), "unmapped_count should be lazy" - assert hasattr( - report_lazy.unmapped_fraction.data, "dask" - ), "unmapped_fraction should be lazy" - assert hasattr( - report_lazy.weight_sum_min.data, "dask" - ), "weight_sum_min should be lazy" - - # 3. Double-Check Identity - # Computing the lazy report should yield identical results to the eager one - report_lazy_comp = report_lazy.compute() - - for var in report_eager.data_vars: - if var in report_lazy_comp: - # Note: with synthetic mocks and multiple chunks, n_weights and unmapped_count - # may differ because each mock chunk generates one weight. - # Real ESMF would be identical. - if not HAS_REAL_ESMF and var in [ - "n_weights", - "unmapped_count", - "unmapped_fraction", - "weight_sum_max", - "weight_sum_mean", - ]: - continue - - np.testing.assert_allclose( - report_eager[var].values, - report_lazy_comp[var].values, - err_msg=f"Mismatch in metric {var} between Eager and Lazy backends", - ) - - -def test_quality_report_dict_is_eager(dask_client): - """ - Verify that quality_report(format='dict') still returns eager values - for immediate consumption, even for distributed regridders. - """ - src = create_global_grid(30, 60) - tgt = create_global_grid(20, 40) - regridder = Regridder(src, tgt, parallel=True) - - report = regridder.quality_report(format="dict") - - assert isinstance(report, dict) - assert isinstance(report["n_weights"], int) - assert isinstance(report["unmapped_count"], int) - assert isinstance(report["unmapped_fraction"], float) - - -if __name__ == "__main__": - pytest.main([__file__]) diff --git a/tests/test_aero_quality_report_opt.py b/tests/test_aero_quality_report_opt.py deleted file mode 100644 index 003e6d8..0000000 --- a/tests/test_aero_quality_report_opt.py +++ /dev/null @@ -1,78 +0,0 @@ -import numpy as np -import pytest -import xarray as xr -from xregrid import Regridder, create_global_grid -from unittest.mock import MagicMock, patch, PropertyMock - - -def test_quality_report_no_gather_distributed(): - """ - Aero Protocol: Verify that quality_report avoids gathering the full weight matrix. - """ - src_grid = create_global_grid(30, 30) - tgt_grid = create_global_grid(30, 30) - - regridder = Regridder(src_grid, tgt_grid, method="bilinear") - - # 1. Setup remote weights simulation - mock_client = MagicMock() - regridder._dask_client = mock_client - - remote_matrix = MagicMock() - remote_matrix.key = "remote_weights_123" - regridder._weights_matrix = remote_matrix - - # Mock submit for _get_nnz_task - mock_nnz_future = MagicMock() - mock_nnz_future.result.return_value = 1234 - - def side_effect(func, *args, **kwargs): - if "_get_nnz_task" in str(func): - return mock_nnz_future - # Return a mock future for anything else (like diagnostics) - f = MagicMock() - f.result.return_value = np.ones(int(np.prod(regridder._shape_target))) - return f - - mock_client.submit.side_effect = side_effect - - # Mock diagnostics to avoid Dask computation issues in the test - mock_diag = xr.Dataset( - { - "weight_sum": (["lat", "lon"], np.ones((6, 12))), - "unmapped_mask": (["lat", "lon"], np.zeros((6, 12))), - }, - coords={"lat": np.arange(6), "lon": np.arange(12)}, - ) - - # 2. Call quality_report with a spy on the weights property - # We use a simple attribute mock instead of property mock to avoid confusion - with patch.object(regridder, "diagnostics", return_value=mock_diag): - with patch.object( - Regridder, "weights", new_callable=PropertyMock - ) as mock_weights: - # We need to make sure mock_weights doesn't trigger anything - report = regridder.quality_report(skip_heavy=False) - - # 3. Verifications - assert report["n_weights"] == 1234 - # Ensure weights property was NEVER accessed (which would mean no gather) - assert mock_weights.call_count == 0 - - # Check that diagnostics was used for other metrics - assert report["unmapped_count"] == 0 - - -def test_quality_report_eager_fallback(): - """Verify that quality_report still works correctly for eager NumPy weights.""" - src_grid = create_global_grid(30, 30) - tgt_grid = create_global_grid(30, 30) - regridder = Regridder(src_grid, tgt_grid, method="bilinear") - - report = regridder.quality_report() - assert report["n_weights"] > 0 - assert "unmapped_count" in report - - -if __name__ == "__main__": - pytest.main([__file__]) diff --git a/tests/test_aero_regrid_robustness.py b/tests/test_aero_regrid_robustness.py deleted file mode 100644 index a595d8f..0000000 --- a/tests/test_aero_regrid_robustness.py +++ /dev/null @@ -1,111 +0,0 @@ -import xarray as xr -import numpy as np -import pytest -from xregrid import Regridder -from xregrid.utils import create_global_grid - - -def test_aero_lazy_initialization(): - """ - Aero Protocol: Verify that Regridder initialization doesn't force - computation of Dask-backed coordinates when possible. - Note: Xarray always computes dimension coordinates to build indexes. - To test laziness, we use non-dimension coordinates. - """ - try: - import dask.array as da - except ImportError: - pytest.skip("Dask not installed") - - # Create an unstructured grid with Dask-backed coordinates - # Unstructured coordinates share the same dimension name - n_pts = 100 - lat_vals = np.linspace(-90, 90, n_pts) - lon_vals = np.linspace(0, 360, n_pts) - - lat = da.from_array(lat_vals, chunks=10) - lon = da.from_array(lon_vals, chunks=10) - - ds_src = xr.Dataset( - coords={ - "lat": (["n_pts"], lat, {"units": "degrees_north"}), - "lon": (["n_pts"], lon, {"units": "degrees_east"}), - } - ) - - ds_tgt = ds_src.copy() - - # Initialize in parallel mode - regridder = Regridder(ds_src, ds_tgt, parallel=True, compute=False) - - # Check that source_grid_ds coordinates are still dask-backed - # Unstructured grids should skip _normalize_grid, so they should remain lazy - assert hasattr( - regridder.source_grid_ds.lat.data, "dask" - ), "Source latitude should remain lazy" - - -def test_aero_double_check_identity(): - """ - Aero Protocol Rule 4: The "Double-Check Test". - Verify that regridding results are identical for NumPy and Dask backends. - """ - try: - import dask.array as da # noqa: F401 - except ImportError: - pytest.skip("Dask not installed") - - # 1. Setup grids - ds_src = create_global_grid(20.0, 20.0) # Low res for fast test - ds_tgt = create_global_grid(10.0, 10.0) - - # 2. Setup data (Eager) - data_eager = np.outer( - np.cos(np.deg2rad(ds_src.lat.values)), np.sin(np.deg2rad(ds_src.lon.values)) - ) - # Only include dim coords to avoid validation error with bounds - da_eager = xr.DataArray( - data_eager, - dims=["lat", "lon"], - coords={"lat": ds_src.lat, "lon": ds_src.lon}, - name="test_data", - ) - - # 3. Setup data (Lazy) - da_lazy = da_eager.chunk({"lat": 5, "lon": 10}) - - # 4. Initialize Regridder - regridder = Regridder(ds_src, ds_tgt, method="bilinear") - - # 5. Regrid both - res_eager = regridder(da_eager) - res_lazy = regridder(da_lazy) - - # 6. Assertions - assert isinstance(res_eager.data, np.ndarray) - assert hasattr(res_lazy.data, "dask") - - # Compare values - np.testing.assert_allclose(res_eager.values, res_lazy.compute().values, rtol=1e-6) - - -def test_aero_diagnostics_crs_propagation(): - """ - Aero Protocol: Verify CRS propagation in diagnostics. - """ - ds_src = create_global_grid(10.0, 10.0) - ds_tgt = create_global_grid(5.0, 5.0) - - # Attach a mock CRS - ds_tgt.attrs["crs"] = "EPSG:4326" - - regridder = Regridder(ds_src, ds_tgt, method="bilinear") - ds_diag = regridder.diagnostics() - - assert "crs" in ds_diag.attrs - # Check for EPSG 4326 in the WKT string - assert "4326" in ds_diag.attrs["crs"] - - -if __name__ == "__main__": - pytest.main([__file__]) diff --git a/tests/test_aero_rotated_pole.py b/tests/test_aero_rotated_pole.py deleted file mode 100644 index fed265f..0000000 --- a/tests/test_aero_rotated_pole.py +++ /dev/null @@ -1,60 +0,0 @@ -import pytest -import xarray as xr -from xregrid import create_rotated_latlon_grid - - -def test_rotated_latlon_grid_eager_lazy(): - """ - Double-Check Test: Verify create_rotated_latlon_grid yields identical results - for Eager (NumPy) and Lazy (Dask) backends and maintains CF compliance. - """ - extent = (-5.0, 5.0, -5.0, 5.0) - res = 1.0 - pole_lat = 37.5 - pole_lon = 177.5 - - # 1. Eager Execution (NumPy) - ds_eager = create_rotated_latlon_grid( - extent=extent, - res=res, - grid_north_pole_lat=pole_lat, - grid_north_pole_lon=pole_lon, - add_bounds=True, - chunks=None, - ) - - # 2. Lazy Execution (Dask) - ds_lazy = create_rotated_latlon_grid( - extent=extent, - res=res, - grid_north_pole_lat=pole_lat, - grid_north_pole_lon=pole_lon, - add_bounds=True, - chunks={"rlat": 5, "rlon": 5}, - ) - - # Verification of Laziness - assert hasattr(ds_lazy.lat.data, "dask") - assert hasattr(ds_lazy.lon.data, "dask") - assert hasattr(ds_lazy.rlat_b.data, "dask") - - # Verification of Numerical Identity - xr.testing.assert_allclose(ds_eager, ds_lazy.compute()) - - # Verification of CF Metadata - assert ds_eager.attrs["grid_mapping"] == "rotated_pole" - assert "rotated_pole" in ds_eager.data_vars - assert ( - ds_eager.rotated_pole.attrs["grid_mapping_name"] == "rotated_latitude_longitude" - ) - assert ds_eager.rlat.attrs["standard_name"] == "grid_latitude" - assert ds_eager.rlon.attrs["standard_name"] == "grid_longitude" - - # Verification of Provenance - assert "Created Rotated Lat-Lon grid" in ds_eager.attrs["history"] - assert "Eager" in ds_eager.attrs["history"] - assert "Lazy" in ds_lazy.attrs["history"] - - -if __name__ == "__main__": - pytest.main([__file__]) diff --git a/tests/test_aero_sinusoidal.py b/tests/test_aero_sinusoidal.py deleted file mode 100644 index bde27d7..0000000 --- a/tests/test_aero_sinusoidal.py +++ /dev/null @@ -1,105 +0,0 @@ -import pytest -import xarray as xr -import numpy as np -from xregrid.utils import ( - create_sinusoidal_grid, - create_grid_from_ioapi, - create_grid_like, -) - - -def test_create_sinusoidal_grid_consistency(): - """Verify that sinusoidal grid generation yields identical results for NumPy and Dask.""" - # Small extent for testing - extent = (-1000000, 1000000, -500000, 500000) - res = 100000 - - ds_eager = create_sinusoidal_grid(extent, res, chunks=None) - ds_lazy = create_sinusoidal_grid(extent, res, chunks={"x": 5, "y": 5}) - - # Verify metadata - assert "lat" in ds_eager.coords - assert "lon" in ds_eager.coords - assert "x" in ds_eager.coords - assert "y" in ds_eager.coords - - # Verify values are identical - xr.testing.assert_allclose(ds_eager, ds_lazy.compute()) - - # Verify backend - assert not hasattr(ds_eager.lat.data, "dask") - assert hasattr(ds_lazy.lat.data, "dask") - assert hasattr(ds_lazy.lat_b.data, "dask") - - -def test_create_grid_from_ioapi_sinu(): - """Verify IOAPI GDTYP 13 (Sinusoidal) grid generation.""" - metadata = { - "GDTYP": 13, - "P_ALP": 0.0, - "P_BET": 0.0, - "P_GAM": 0.0, - "XCENT": -97.0, - "YCENT": 0.0, - "XORIG": -1000000.0, - "YORIG": -1000000.0, - "XCELL": 10000.0, - "YCELL": 10000.0, - "NCOLS": 10, - "NROWS": 10, - } - - ds = create_grid_from_ioapi(metadata) - assert ds.attrs["ioapi_GDTYP"] == 13 - assert "lat" in ds.coords - assert "lon" in ds.coords - - # Verify lon_0 is respected (at least check the central meridian) - # The center of the grid in x is XORIG + (NCOLS/2)*XCELL = -1000000 + 50000 = -950000 - # Central meridian is -97.0. - # At y=0 (equator), lon should be close to -97.0 + x/R * 180/pi - # But easier to just check if it runs and has reasonable values. - assert ds.lon.mean() < 0 - - -def test_sinusoidal_grid_like(): - """Verify create_grid_like works with Sinusoidal grids.""" - extent = (-500000, 500000, -500000, 500000) - res = 50000 - ds_base = create_sinusoidal_grid(extent, res, lon_0=-100) - - # Create a new grid like the base one but with different resolution - new_res = 100000 - ds_new = create_grid_like(ds_base, new_res) - - assert ds_new.sizes["x"] == 10 - assert ds_new.sizes["y"] == 10 - - # Check if CRS is preserved (via WKT comparison) - assert ds_base.attrs["crs"] == ds_new.attrs["crs"] - - # Check if extent is similar - assert np.allclose(ds_base.x.min() - res / 2, ds_new.x.min() - new_res / 2) - assert np.allclose(ds_base.x.max() + res / 2, ds_new.x.max() + new_res / 2) - - -def test_sinusoidal_res_aliases(): - """Verify Sinusoidal grid generation with resolution aliases.""" - extent = (0, 100000, 0, 100000) - ds_10km = create_sinusoidal_grid(extent, "10km") - ds_5km = create_sinusoidal_grid(extent, "5km") - ds_1km = create_sinusoidal_grid(extent, "1km") - ds_500m = create_sinusoidal_grid(extent, "500m") - ds_250m = create_sinusoidal_grid(extent, "250m") - - # 1km alias should be ~926.6m - expected_1km = 926.6254331 - assert np.allclose(ds_10km.x.diff("x").mean(), expected_1km * 10) - assert np.allclose(ds_5km.x.diff("x").mean(), expected_1km * 5) - assert np.allclose(ds_1km.x.diff("x").mean(), expected_1km) - assert np.allclose(ds_500m.x.diff("x").mean(), expected_1km / 2) - assert np.allclose(ds_250m.x.diff("x").mean(), expected_1km / 4) - - -if __name__ == "__main__": - pytest.main([__file__]) diff --git a/tests/test_aero_smart_features.py b/tests/test_aero_smart_features.py deleted file mode 100644 index c9ffd73..0000000 --- a/tests/test_aero_smart_features.py +++ /dev/null @@ -1,111 +0,0 @@ -import numpy as np -import pytest -import xarray as xr -import dask.array as da -from xregrid import Regridder, create_global_grid -from xregrid.utils import _find_coord -from unittest.mock import patch - - -def test_enhanced_coord_discovery(): - """Verify enhanced coordinate discovery for non-standard names.""" - # Create dataset with 'xc' and 'yc' without CF attributes to force fallback - # Use auxiliary coordinates to test lazy discovery - ds = xr.Dataset( - data_vars={"data": (["y", "x"], np.random.rand(10, 20))}, - coords={ - "yc": (["y", "x"], np.random.rand(10, 20)), - "xc": (["y", "x"], np.random.rand(10, 20)), - }, - ) - - # 1. Eager - lat = _find_coord(ds, "latitude") - lon = _find_coord(ds, "longitude") - - assert lat.name == "yc" - assert lon.name == "xc" - - # 2. Lazy - ds_lazy = ds.chunk({"y": 5, "x": 10}) - lat_lazy = _find_coord(ds_lazy, "latitude") - lon_lazy = _find_coord(ds_lazy, "longitude") - - assert lat_lazy.name == "yc" - assert lon_lazy.name == "xc" - assert hasattr(lat_lazy.data, "dask") - - -def test_auto_periodicity_detection(): - """Verify auto-periodicity detection logic.""" - # Global grid should be detected as periodic - ds_global = create_global_grid(10, 10) - regridder = Regridder(ds_global, ds_global, periodic=None) - assert regridder.periodic is True - - # Regional grid should NOT be detected as periodic - ds_regional = xr.Dataset( - coords={ - "lat": (["lat"], np.linspace(20, 50, 10)), - "lon": (["lon"], np.linspace(-100, -70, 20)), - } - ) - regridder_reg = Regridder(ds_regional, ds_regional, periodic=None) - assert regridder_reg.periodic is False - - -def test_auto_periodicity_lazy(): - """Verify auto-periodicity detection handles lazy coordinates without compute.""" - # Test 1: 2D lazy coordinates (not indexes, so they stay lazy) - # Dimension coordinates (1D) are often loaded by Xarray for indexing. - y, x = da.meshgrid(da.linspace(-90, 90, 10), da.linspace(0, 342, 20), indexing="ij") - y = y.rechunk(5) - x = x.rechunk(10) - - ds_lazy = xr.Dataset( - data_vars={"data": (["y", "x"], da.zeros((10, 20), chunks=(5, 10)))}, - coords={ - "lat": (["y", "x"], y, {"units": "degrees_north"}), - "lon": (["y", "x"], x, {"units": "degrees_east"}), - }, - ) - - # Regridder should NOT compute the dask arrays for detection - # Since they are lazy and no metadata is present, it should default to False - regridder = Regridder(ds_lazy, ds_lazy, periodic=None) - assert regridder.periodic is False - - # Now add metadata - ds_lazy.lon.attrs["boundary"] = "periodic" - regridder_meta = Regridder(ds_lazy, ds_lazy, periodic=None) - assert regridder_meta.periodic is True - - -def test_plot_comparison_dispatch(): - """Verify plot_comparison method correctly dispatches to viz.""" - # Mock viz functions - with patch("xregrid.viz.plot_comparison") as mock_static: - with patch("xregrid.viz.plot_comparison_interactive") as mock_interactive: - src = create_global_grid(30, 30) - regridder = Regridder(src, src, periodic=False) - - da_src = xr.DataArray(np.random.rand(6, 12), dims=("lat", "lon")) - da_tgt = xr.DataArray(np.random.rand(6, 12), dims=("lat", "lon")) - - # Track A (Static) - regridder.plot_comparison(da_src, da_tgt, mode="static", custom_kw="test") - mock_static.assert_called_once_with( - da_src, da_tgt, regridder=regridder, custom_kw="test" - ) - - # Track B (Interactive) - regridder.plot_comparison( - da_src, da_tgt, mode="interactive", rasterize=False - ) - mock_interactive.assert_called_once_with( - da_src, da_tgt, regridder=regridder, rasterize=False - ) - - -if __name__ == "__main__": - pytest.main([__file__]) diff --git a/tests/test_aero_spatial_slice.py b/tests/test_aero_spatial_slice.py deleted file mode 100644 index 6e14bd0..0000000 --- a/tests/test_aero_spatial_slice.py +++ /dev/null @@ -1,104 +0,0 @@ -from __future__ import annotations -import numpy as np -import xarray as xr -from xregrid.utils import create_global_grid, create_mesh_from_coords, spatial_slice - - -def test_spatial_slice_rectilinear() -> None: - """ - Test spatial_slice with a rectilinear grid (Eager and Lazy). - - Verifies that spatial_slice correctly subsets a rectilinear grid - using coordinate indexes and handles both NumPy and Dask backends. - """ - # 1. Eager - ds = create_global_grid(res_lat=1.0, res_lon=1.0) - # Extent: (min_x, max_x, min_y, max_y) - extent = (10.5, 20.5, 30.5, 40.5) - ds_sliced = spatial_slice(ds, extent) - - assert ds_sliced.lat.min() >= 30.5 - assert ds_sliced.lat.max() <= 40.5 - assert ds_sliced.lon.min() >= 10.5 - assert ds_sliced.lon.max() <= 20.5 - assert "history" in ds_sliced.attrs - assert "Spatially sliced" in ds_sliced.attrs["history"] - - # 2. Lazy (Dask) - ds_lazy = create_global_grid( - res_lat=1.0, res_lon=1.0, chunks={"lat": 10, "lon": 10} - ) - ds_sliced_lazy = spatial_slice(ds_lazy, extent) - - # In xarray, dimension coordinates are often eager (NumPy) due to indexing. - # Check lat_b instead, which should remain lazy. - assert hasattr(ds_sliced_lazy.lat_b.data, "dask") - xr.testing.assert_allclose(ds_sliced, ds_sliced_lazy.compute()) - - -def test_spatial_slice_unstructured() -> None: - """ - Test spatial_slice with an unstructured grid (Eager and Lazy). - - Verifies that spatial_slice correctly subsets an unstructured grid - using boolean masking and maintains laziness for Dask-backed data. - """ - # Create points: one in the box, one outside - lons = np.array([15.0, 50.0]) - lats = np.array([35.0, 60.0]) - ds = create_mesh_from_coords(lons, lats, crs="EPSG:4326") - - extent = (10.0, 20.0, 30.0, 40.0) - - # Eager path: drop=True is supported - ds_sliced = spatial_slice(ds, extent) - - assert ds_sliced.sizes["n_pts"] == 1 - assert ds_sliced.lon.values[0] == 15.0 - assert ds_sliced.lat.values[0] == 35.0 - - # Lazy path: drop=False is used to preserve laziness - ds_lazy = create_mesh_from_coords(lons, lats, crs="EPSG:4326", chunks=1) - ds_lazy["data"] = (["n_pts"], ds_lazy.lat.data, {"units": "K"}) - ds_sliced_lazy = spatial_slice(ds_lazy, extent) - - # Verify laziness - assert hasattr(ds_sliced_lazy.data.data, "dask") - - # Verify results (after compute and dropna since drop=False was used) - ds_res = ds_sliced_lazy.compute().dropna("n_pts") - assert ds_res.sizes["n_pts"] == 1 - assert ds_res.lon.values[0] == 15.0 - - -def test_spatial_slice_wrapping() -> None: - """ - Test spatial_slice with longitude wrapping across different grid types. - - Verifies that spatial_slice correctly handles regions crossing the - meridian/dateline for both rectilinear and unstructured grids. - """ - ds = create_global_grid(res_lat=1.0, res_lon=1.0) - - # Slice crossing the 0/360 boundary - # Extent: (min_x, max_x, min_y, max_y) - extent = (-10.5, 10.5, -10.5, 10.5) - ds_sliced = spatial_slice(ds, extent) - - # Lon in ds is 0-360. -10.5 should map to 349.5 - assert ds_sliced.lon.min() >= 0 - assert ds_sliced.lon.max() <= 360 - - # Check that we have both the 0-10.5 and 349.5-360 parts - assert (ds_sliced.lon <= 10.5).any() - assert (ds_sliced.lon >= 349.5).any() - - # Verify unstructured wrapping - lons = np.array([5.0, 355.0, 180.0]) - lats = np.array([0.0, 0.0, 0.0]) - ds_unstructured = create_mesh_from_coords(lons, lats, crs="EPSG:4326") - ds_un_sliced = spatial_slice(ds_unstructured, extent) - - # Since it's eager, drop=True was used - assert ds_un_sliced.sizes["n_pts"] == 2 - assert set(ds_un_sliced.lon.values) == {5.0, 355.0} diff --git a/tests/test_aero_total_weights_opt.py b/tests/test_aero_total_weights_opt.py deleted file mode 100644 index e21326d..0000000 --- a/tests/test_aero_total_weights_opt.py +++ /dev/null @@ -1,101 +0,0 @@ -import numpy as np -import xarray as xr -from xregrid import Regridder -import xregrid.xregrid as xregrid_mod - - -def test_total_weights_distribution_eager_vs_lazy(): - """ - Aero Protocol: Verify that total_weights distribution works correctly - and produces identical results for Eager (NumPy) and Lazy (Dask) paths - when skipna=True. - """ - from dask.distributed import Client, LocalCluster - - # Clear cache to ensure fresh test - xregrid_mod._WORKER_CACHE.clear() - - # Create simple source and target grids - ds_src = xr.Dataset( - { - "lat": (["lat"], np.linspace(-90, 90, 10)), - "lon": (["lon"], np.linspace(0, 360, 20)), - } - ) - # Ensure they have CF attributes for _get_mesh_info - ds_src.lat.attrs["units"] = "degrees_north" - ds_src.lon.attrs["units"] = "degrees_east" - - ds_tgt = xr.Dataset( - { - "lat": (["lat"], np.linspace(-90, 90, 15)), - "lon": (["lon"], np.linspace(0, 360, 25)), - } - ) - ds_tgt.lat.attrs["units"] = "degrees_north" - ds_tgt.lon.attrs["units"] = "degrees_east" - - # Create source data with some NaNs - data = np.random.rand(10, 20).astype(np.float32) - data[0, 0] = np.nan - da_src_numpy = xr.DataArray( - data, coords=ds_src.coords, dims=("lat", "lon"), name="test" - ) - - # 1. Eager Path - regridder = Regridder(ds_src, ds_tgt, method="bilinear", skipna=True) - res_numpy = regridder(da_src_numpy) - - # 2. Lazy Path - with LocalCluster(n_workers=1, threads_per_worker=1, processes=False) as cluster: - with Client(cluster): - da_src_dask = da_src_numpy.chunk({"lat": 5, "lon": 10}) - res_dask = regridder(da_src_dask) - - # Verify identity - xr.testing.assert_allclose(res_numpy, res_dask.compute()) - - # Verify that history contains the new metadata - assert "ESMF/esmpy=" in res_numpy.attrs["history"] - assert "skipna=True" in res_numpy.attrs["history"] - assert "na_thres=1.0" in res_numpy.attrs["history"] - - # Check if the total weights key was created in _WORKER_CACHE - tw_keys = [k for k in xregrid_mod._WORKER_CACHE.keys() if k.startswith("tw_")] - assert len(tw_keys) > 0, "Total weights should have been cached with a key" - - # Check if weights_matrix was also cached - w_keys = [k for k in xregrid_mod._WORKER_CACHE.keys() if k.startswith("weights_")] - assert len(w_keys) > 0, "Weights matrix should have been cached with a key" - - -def test_provenance_with_extrap(): - """Verify that extrapolation metadata is included in history.""" - ds_src = xr.Dataset( - { - "lat": (["lat"], np.linspace(-90, 90, 10)), - "lon": (["lon"], np.linspace(0, 360, 20)), - } - ) - ds_src.lat.attrs["units"] = "degrees_north" - ds_src.lon.attrs["units"] = "degrees_east" - - ds_tgt = xr.Dataset( - { - "lat": (["lat"], np.linspace(-90, 90, 15)), - "lon": (["lon"], np.linspace(0, 360, 25)), - } - ) - ds_tgt.lat.attrs["units"] = "degrees_north" - ds_tgt.lon.attrs["units"] = "degrees_east" - - da_src = xr.DataArray( - np.ones((10, 20)), coords=ds_src.coords, dims=("lat", "lon"), name="test" - ) - - regridder = Regridder( - ds_src, ds_tgt, method="bilinear", extrap_method="nearest_s2d" - ) - res = regridder(da_src) - - assert "extrap_method=nearest_s2d" in res.attrs["history"] diff --git a/tests/test_aero_ufs_names.py b/tests/test_aero_ufs_names.py deleted file mode 100644 index b047a2b..0000000 --- a/tests/test_aero_ufs_names.py +++ /dev/null @@ -1,41 +0,0 @@ -import pytest -import xarray as xr -import numpy as np -from xregrid.utils import _find_coord - - -def test_find_coord_ufs_names(): - """ - Verify that _find_coord correctly identifies UFS-style coordinate names. - """ - # 1. Center coordinates (grid_latt, grid_lont) - ds_center = xr.Dataset( - coords={ - "grid_latt": (["y", "x"], np.zeros((10, 10))), - "grid_lont": (["y", "x"], np.zeros((10, 10))), - } - ) - - lat_da = _find_coord(ds_center, "latitude") - lon_da = _find_coord(ds_center, "longitude") - - assert lat_da.name == "grid_latt" - assert lon_da.name == "grid_lont" - - # 2. Corner coordinates (grid_lat, grid_lon) - ds_corner = xr.Dataset( - coords={ - "grid_lat": (["y_b", "x_b"], np.zeros((11, 11))), - "grid_lon": (["y_b", "x_b"], np.zeros((11, 11))), - } - ) - - lat_da_b = _find_coord(ds_corner, "latitude") - lon_da_b = _find_coord(ds_corner, "longitude") - - assert lat_da_b.name == "grid_lat" - assert lon_da_b.name == "grid_lon" - - -if __name__ == "__main__": - pytest.main([__file__]) diff --git a/tests/test_aero_ugrid_full_support.py b/tests/test_aero_ugrid_full_support.py deleted file mode 100644 index ca25331..0000000 --- a/tests/test_aero_ugrid_full_support.py +++ /dev/null @@ -1,136 +0,0 @@ -import numpy as np -import pytest -import xarray as xr -from xregrid import Regridder, create_global_grid - - -def create_mock_ugrid(n_nodes=16, n_faces=9): - """Create a mock UGRID dataset using a regular quad mesh.""" - # Ensure n_nodes is a square for simplicity in this mock - n_side = int(np.sqrt(n_nodes)) - if n_side**2 != n_nodes: - # Fallback to next square - n_side = int(np.ceil(np.sqrt(n_nodes))) - n_nodes = n_side**2 - - x = np.linspace(0, 350, n_side) - y = np.linspace(-80, 80, n_side) - node_x_2d, node_y_2d = np.meshgrid(x, y) - node_x = node_x_2d.flatten() - node_y = node_y_2d.flatten() - - # Create quads - face_nodes = [] - face_x = [] - face_y = [] - for j in range(n_side - 1): - for i in range(n_side - 1): - n0 = j * n_side + i - n1 = j * n_side + i + 1 - n2 = (j + 1) * n_side + i + 1 - n3 = (j + 1) * n_side + i - face_nodes.append([n0, n1, n2, n3]) - face_x.append((node_x[n0] + node_x[n2]) / 2) - face_y.append((node_y[n0] + node_y[n2]) / 2) - - face_nodes = np.array(face_nodes) - n_faces = len(face_nodes) - face_x = np.array(face_x) - face_y = np.array(face_y) - - ds = xr.Dataset( - data_vars={ - "mesh_topology": ( - [], - 0, - { - "cf_role": "mesh_topology", - "topology_dimension": 2, - "node_coordinates": "node_lon node_lat", - "face_node_connectivity": "face_nodes", - "face_coordinates": "face_lon face_lat", - }, - ), - "face_nodes": ( - ["n_face", "n_node_per_face"], - face_nodes, - {"cf_role": "face_node_connectivity", "start_index": 0}, - ), - "temp": ( - ["n_face"], - np.random.rand(n_faces), - { - "mesh": "mesh_topology", - "location": "face", - "standard_name": "air_temperature", - }, - ), - }, - coords={ - "node_lon": ( - ["n_node"], - node_x, - {"standard_name": "longitude", "units": "degrees_east"}, - ), - "node_lat": ( - ["n_node"], - node_y, - {"standard_name": "latitude", "units": "degrees_north"}, - ), - "face_lon": ( - ["n_face"], - face_x, - {"standard_name": "longitude", "units": "degrees_east"}, - ), - "face_lat": ( - ["n_face"], - face_y, - {"standard_name": "latitude", "units": "degrees_north"}, - ), - }, - ) - return ds - - -def test_ugrid_discovery_and_regrid(): - """Verify UGRID discovery and regridding (Eager and Lazy).""" - src_ds = create_mock_ugrid(n_nodes=25, n_faces=16) - tgt_grid = create_global_grid(30, 30) - - # Test that Regridder can handle the UGRID dataset - # We use conservative regridding to test triangulation logic - regridder = Regridder(src_ds, tgt_grid, method="conservative") - - # 1. Eager test - res_eager = regridder(src_ds.temp) - - assert "lat" in res_eager.coords - assert "lon" in res_eager.coords - # Verify metadata removal as target is not UGRID - assert "mesh" not in res_eager.attrs - - # 2. Lazy test - da_lazy = src_ds.temp.chunk({"n_face": 5}) - res_lazy = regridder(da_lazy).compute() - - xr.testing.assert_allclose(res_eager, res_lazy) - - -def test_ugrid_scientific_hygiene(): - """Verify UGRID metadata propagation to UGRID target.""" - src_ds = create_mock_ugrid(n_nodes=16, n_faces=9) - tgt_ds = create_mock_ugrid(n_nodes=25, n_faces=16) - - # For simplicity, use nearest_s2d which doesn't require complex connectivity for target - regridder = Regridder(src_ds, tgt_ds, method="nearest_s2d") - - res = regridder(src_ds.temp) - - # Scientific Hygiene: target mesh should be attached - assert res.attrs["mesh"] == "mesh_topology" - assert "location" in res.attrs - assert "mesh_topology" in res.coords or "mesh_topology" in res.data_vars - - -if __name__ == "__main__": - pytest.main([__file__]) diff --git a/tests/test_aero_unstructured_enhanced.py b/tests/test_aero_unstructured_enhanced.py deleted file mode 100644 index 78bc037..0000000 --- a/tests/test_aero_unstructured_enhanced.py +++ /dev/null @@ -1,131 +0,0 @@ -import numpy as np -import pytest -import xarray as xr -from xregrid import Regridder - - -def test_unstructured_bilinear_mesh_enhanced(): - """ - Aero Protocol: Verify unstructured bilinear regridding using Mesh (NODE-based). - Double-Check Test: NumPy vs Dask backends. - """ - # Create a simple triangle mesh (4 nodes, 2 triangles) - # 3 -- 2 - # | / | - # 0 -- 1 - lon = xr.DataArray( - [0.0, 1.0, 1.0, 0.0], dims="n_pts", name="lon", attrs={"units": "degrees_east"} - ) - lat = xr.DataArray( - [0.0, 0.0, 1.0, 1.0], dims="n_pts", name="lat", attrs={"units": "degrees_north"} - ) - - # Face-node connectivity (0-based for UGRID) - # Tri 1: 0, 1, 2 - # Tri 2: 0, 2, 3 - conn = xr.DataArray( - [[0, 1, 2], [0, 2, 3]], - dims=("n_face", "n_vertex"), - name="face_node_connectivity", - attrs={"cf_role": "face_node_connectivity", "start_index": 0}, - ) - - src_grid = xr.Dataset( - coords={"lon_node": lon, "lat_node": lat}, - data_vars={"face_node_connectivity": conn}, - ) - - # Target grid: a single point in the middle - # We use a 1D target to keep it simple - tgt_lon = xr.DataArray([0.5], dims="n_dst", name="lon") - tgt_lat = xr.DataArray([0.5], dims="n_dst", name="lat") - tgt_grid = xr.Dataset(coords={"lon": tgt_lon, "lat": tgt_lat}) - - # Data on nodes - data_val = np.array([10.0, 20.0, 30.0, 40.0]) - da_eager = xr.DataArray( - data_val, - dims="n_pts", - coords={"lon_node": lon, "lat_node": lat}, - name="test_data", - ) - - # Initialize Regridder with bilinear - # Ensure mesh variable is identified correctly - src_grid.face_node_connectivity.attrs["mesh"] = "mesh" - src_grid["mesh"] = ( - [], - 0, - { - "cf_role": "mesh_topology", - "face_node_connectivity": "face_node_connectivity", - "node_coordinates": "lon_node lat_node", - }, - ) - - regridder = Regridder(src_grid, tgt_grid, method="bilinear") - - # 1. Test Eager (NumPy) - res_eager = regridder(da_eager) - - # 2. Test Lazy (Dask) - da_lazy = da_eager.chunk({"n_pts": 2}) - res_lazy = regridder(da_lazy) - - # Assertions - assert isinstance(res_eager.data, np.ndarray) - assert hasattr(res_lazy.data, "dask") - - # Results should be identical - xr.testing.assert_allclose(res_eager, res_lazy.compute()) - - # Verification of shape - assert res_eager.shape == (1,) - - # Verify provenance - assert "method=bilinear" in res_eager.attrs["history"] - assert "ESMF" in res_eager.attrs["history"] - - -def test_unstructured_patch_mesh_enhanced(): - """Verify that 'patch' method also works on unstructured meshes.""" - lon = xr.DataArray([0.0, 1.0, 1.0, 0.0], dims="n_pts", name="lon") - lat = xr.DataArray([0.0, 0.0, 1.0, 1.0], dims="n_pts", name="lat") - conn = xr.DataArray( - [[0, 1, 2], [0, 2, 3]], - dims=("n_face", "n_vertex"), - name="face_node_connectivity", - attrs={"cf_role": "face_node_connectivity", "start_index": 0}, - ) - src_grid = xr.Dataset( - coords={"lon_node": lon, "lat_node": lat}, - data_vars={"face_node_connectivity": conn}, - ) - - tgt_lon = xr.DataArray([0.5], dims="n_dst", name="lon") - tgt_lat = xr.DataArray([0.5], dims="n_dst", name="lat") - tgt_grid = xr.Dataset(coords={"lon": tgt_lon, "lat": tgt_lat}) - - # Ensure mesh variable is identified correctly - src_grid.face_node_connectivity.attrs["mesh"] = "mesh" - src_grid["mesh"] = ( - [], - 0, - { - "cf_role": "mesh_topology", - "face_node_connectivity": "face_node_connectivity", - "node_coordinates": "lon_node lat_node", - }, - ) - regridder = Regridder(src_grid, tgt_grid, method="patch") - da = xr.DataArray( - [1.0, 1.0, 1.0, 1.0], dims="n_pts", coords={"lon_node": lon, "lat_node": lat} - ) - - res = regridder(da) - assert res.shape == (1,) - assert "method=patch" in res.attrs["history"] - - -if __name__ == "__main__": - pytest.main([__file__]) diff --git a/tests/test_aero_utils_lazy.py b/tests/test_aero_utils_lazy.py deleted file mode 100644 index 9cdf08f..0000000 --- a/tests/test_aero_utils_lazy.py +++ /dev/null @@ -1,97 +0,0 @@ -import numpy as np -import xarray as xr -from xregrid.utils import ( - create_global_grid, - create_regional_grid, - create_grid_from_crs, - create_mesh_from_coords, -) - - -def test_create_global_grid_lazy(): - """ - Aero Protocol: Double-Check Test for create_global_grid. - Verifies that values are identical between NumPy and Dask backends. - """ - res_lat, res_lon = 10, 20 - - # Eager (NumPy) - ds_eager = create_global_grid(res_lat=res_lat, res_lon=res_lon, chunks=None) - assert not ds_eager.chunks - - # Lazy (Dask) - ds_lazy = create_global_grid( - res_lat=res_lat, res_lon=res_lon, chunks={"lat": 9, "lon": 9} - ) - assert ds_lazy.chunks - - # Assert values are identical - xr.testing.assert_allclose(ds_eager, ds_lazy.compute()) - - # Verify internal backend (lat_b is non-index so it should be chunked) - assert hasattr(ds_lazy.lat_b.data, "dask") - - -def test_create_regional_grid_lazy(): - """ - Aero Protocol: Double-Check Test for create_regional_grid. - """ - lat_range = (-45, 45) - lon_range = (0, 90) - res_lat, res_lon = 5, 5 - - # Eager (NumPy) - ds_eager = create_regional_grid(lat_range, lon_range, res_lat, res_lon, chunks=None) - - # Lazy (Dask) - ds_lazy = create_regional_grid(lat_range, lon_range, res_lat, res_lon, chunks=5) - assert ds_lazy.chunks - - # Assert values are identical - xr.testing.assert_allclose(ds_eager, ds_lazy.compute()) - - # Verify internal backend - assert hasattr(ds_lazy.lat_b.data, "dask") - - -def test_create_grid_from_crs_lazy(): - """ - Aero Protocol: Double-Check Test for create_grid_from_crs. - """ - # Test with EPSG:32633 (UTM zone 33N) - extent = (400000, 500000, 5000000, 5100000) - res = 10000 # 10km - - # Eager (NumPy) - ds_eager = create_grid_from_crs("EPSG:32633", extent, res, chunks=None) - - # Lazy (Dask) - ds_lazy = create_grid_from_crs("EPSG:32633", extent, res, chunks={"x": 5, "y": 5}) - assert ds_lazy.chunks - - # Assert values are identical - xr.testing.assert_allclose(ds_eager, ds_lazy.compute()) - - # Verify internal backend (lat/lon are non-index here) - assert hasattr(ds_lazy.lat.data, "dask") - - -def test_create_mesh_from_coords_lazy(): - """ - Aero Protocol: Double-Check Test for create_mesh_from_coords. - """ - x = np.array([400000, 450000, 500000]) - y = np.array([5000000, 5050000, 5100000]) - - # Eager (NumPy) - ds_eager = create_mesh_from_coords(x, y, "EPSG:32633", chunks=None) - - # Lazy (Dask) - ds_lazy = create_mesh_from_coords(x, y, "EPSG:32633", chunks={"n_pts": 2}) - assert ds_lazy.chunks - - # Assert values are identical - xr.testing.assert_allclose(ds_eager, ds_lazy.compute()) - - # Verify internal backend - assert hasattr(ds_lazy.lat.data, "dask") diff --git a/tests/test_aero_utils_new.py b/tests/test_aero_utils_new.py deleted file mode 100644 index cbf5748..0000000 --- a/tests/test_aero_utils_new.py +++ /dev/null @@ -1,85 +0,0 @@ -import xarray as xr -from xregrid.utils import create_grid_like, create_regional_grid, create_grid_from_crs - - -def test_create_grid_like_latlon(): - """ - Aero Protocol: Double-Check Test for create_grid_like (Lat-Lon). - Verifies identity between NumPy and Dask backends and preservation of laziness. - """ - # Create a source grid - ds_src = create_regional_grid( - lat_range=(10, 20), - lon_range=(100, 110), - res_lat=1.0, - res_lon=1.0, - add_bounds=True, - ) - - res_new = 0.5 - - # Eager (NumPy) - ds_eager = create_grid_like(ds_src, res_new, chunks=None) - assert not ds_eager.chunks - assert ds_eager.lat.size == 20 - assert ds_eager.lon.size == 20 - - # Lazy (Dask) - ds_src_lazy = ds_src.chunk({"lat": 5, "lon": 5}) - ds_lazy = create_grid_like(ds_src_lazy, res_new, chunks=5) - - # Assert values are identical - xr.testing.assert_allclose(ds_eager, ds_lazy.compute()) - - # Verify laziness: lat_b should be a dask array - assert hasattr(ds_lazy.lat_b.data, "dask") - - -def test_create_grid_like_projected(): - """ - Aero Protocol: Double-Check Test for create_grid_like (Projected). - """ - # UTM zone 33N - crs = "EPSG:32633" - extent = (400000, 500000, 5000000, 5100000) - res_orig = 10000 - - ds_src = create_grid_from_crs(crs, extent, res_orig, add_bounds=True) - - res_new = 5000 - - # Eager - ds_eager = create_grid_like(ds_src, res_new, chunks=None) - - # Lazy - ds_src_lazy = ds_src.chunk({"x": 5, "y": 5}) - ds_lazy = create_grid_like(ds_src_lazy, res_new, chunks=5) - - # Assert values - xr.testing.assert_allclose(ds_eager, ds_lazy.compute()) - - # Verify laziness (lat/lon are non-index in projected grid) - assert hasattr(ds_lazy.lat.data, "dask") - assert ds_lazy.attrs["crs"] == ds_src.attrs["crs"] - - -def test_rectilinear_hygiene(): - """ - Verify that _create_rectilinear_grid produces high-hygiene metadata. - """ - ds = create_regional_grid((0, 10), (0, 10), 1, 1) - - assert ds.attrs["crs"] == "EPSG:4326" - - # Test custom CRS - # create_regional_grid currently doesn't expose crs, let's test _create_rectilinear_grid directly - from xregrid.utils import _create_rectilinear_grid - - ds_nad83 = _create_rectilinear_grid((0, 10), (0, 10), 1, 1, crs="EPSG:4269") - assert ds_nad83.attrs["crs"] == "EPSG:4269" - - assert ds.lat.attrs["standard_name"] == "latitude" - assert ds.lon.attrs["standard_name"] == "longitude" - assert ds.lat_b.attrs["standard_name"] == "latitude_bounds" - assert ds.lon_b.attrs["standard_name"] == "longitude_bounds" - assert "history" in ds.attrs diff --git a/tests/test_aero_viz.py b/tests/test_aero_viz.py deleted file mode 100644 index e23886c..0000000 --- a/tests/test_aero_viz.py +++ /dev/null @@ -1,77 +0,0 @@ -import numpy as np -import pytest -import xarray as xr -from xregrid import Regridder, create_global_grid -from xregrid.viz import plot_comparison_interactive - -try: - import hvplot.xarray # noqa: F401 - import holoviews as hv - - HAS_HV = True -except ImportError: - HAS_HV = False - - -@pytest.mark.skipif(not HAS_HV, reason="hvplot/holoviews not installed") -def test_plot_comparison_interactive_types(): - """Verify that plot_comparison_interactive returns the correct HoloViews object.""" - # Setup small grids - src_grid = create_global_grid(30, 30) - tgt_grid = create_global_grid(20, 20) - - # Eager Data - da_src = xr.DataArray( - np.random.rand(6, 12), - dims=("lat", "lon"), - coords={"lat": src_grid.lat, "lon": src_grid.lon}, - name="test_data", - ) - - # Target data (dummy) - da_tgt = xr.DataArray( - np.random.rand(9, 18), - dims=("lat", "lon"), - coords={"lat": tgt_grid.lat, "lon": tgt_grid.lon}, - name="test_data", - ) - - # 1. Test with Eager data, no regridder - layout = plot_comparison_interactive(da_src, da_tgt) - assert isinstance(layout, hv.Layout) - assert len(layout) == 3 # Source, Target, Difference - - # 2. Test with Lazy data - da_src_lazy = da_src.chunk({"lat": 3, "lon": 6}) - da_tgt_lazy = da_tgt.chunk({"lat": 3, "lon": 6}) - layout_lazy = plot_comparison_interactive(da_src_lazy, da_tgt_lazy) - assert isinstance(layout_lazy, hv.Layout) - - # 3. Test with Regridder - regridder = Regridder(src_grid, tgt_grid, method="bilinear") - layout_regrid = plot_comparison_interactive(da_src, da_tgt, regridder=regridder) - assert isinstance(layout_regrid, hv.Layout) - - -@pytest.mark.skipif(not HAS_HV, reason="hvplot/holoviews not installed") -def test_plot_comparison_interactive_titles(): - """Verify that titles are correctly applied to the layout.""" - src_grid = create_global_grid(30, 30) - tgt_grid = create_global_grid(20, 20) - da_src = xr.DataArray( - np.random.rand(6, 12), - dims=("lat", "lon"), - coords={"lat": src_grid.lat, "lon": src_grid.lon}, - ) - da_tgt = xr.DataArray( - np.random.rand(9, 18), - dims=("lat", "lon"), - coords={"lat": tgt_grid.lat, "lon": tgt_grid.lon}, - ) - - title = "My Custom Comparison" - layout = plot_comparison_interactive(da_src, da_tgt, title=title) - - # In HoloViews, title might be in opts - # We just check it doesn't crash and returns the layout - assert isinstance(layout, hv.Layout) diff --git a/tests/test_aero_viz_interactive_smart.py b/tests/test_aero_viz_interactive_smart.py deleted file mode 100644 index 8bc96db..0000000 --- a/tests/test_aero_viz_interactive_smart.py +++ /dev/null @@ -1,54 +0,0 @@ -import numpy as np -import xarray as xr - -from unittest.mock import patch - -import xregrid.viz -from xregrid.viz import plot_interactive - -xregrid.viz.hvplot = True # Force enable for testing - - -def test_plot_interactive_smart_crs(): - """ - Verify that plot_interactive discovers CRS and sets geo=True. - """ - da = xr.DataArray( - np.random.rand(10, 20), - dims=["lat", "lon"], - coords={ - "lat": ( - ["lat"], - np.linspace(-90, 90, 10), - {"standard_name": "latitude", "units": "degrees_north"}, - ), - "lon": ( - ["lon"], - np.linspace(0, 360, 20), - {"standard_name": "longitude", "units": "degrees_east"}, - ), - }, - name="test_data", - ) - # Add CRS info - da.attrs["crs"] = "EPSG:4326" - - # Mock the hvplot accessor call - with patch.object(xr.DataArray, "hvplot", create=True) as mock_hvplot: - plot_interactive(da) - # Check that geo=True was passed in kwargs - args, kwargs = mock_hvplot.call_args - assert kwargs.get("geo") is True - assert kwargs.get("title") == "Interactive Map" - - -def test_plot_interactive_no_crs(): - """ - Verify that plot_interactive does not set geo=True if no CRS is found. - """ - da = xr.DataArray(np.random.rand(10, 20), dims=["y", "x"], name="test_data") - - with patch.object(xr.DataArray, "hvplot", create=True) as mock_hvplot: - plot_interactive(da) - args, kwargs = mock_hvplot.call_args - assert "geo" not in kwargs diff --git a/tests/test_aero_viz_unstructured.py b/tests/test_aero_viz_unstructured.py deleted file mode 100644 index e46d849..0000000 --- a/tests/test_aero_viz_unstructured.py +++ /dev/null @@ -1,134 +0,0 @@ -import numpy as np -import xarray as xr -from unittest.mock import MagicMock, patch -from xregrid.viz import plot_static, plot_interactive - - -def create_unstructured_da(lazy: bool = False) -> xr.DataArray: - """ - Create a 1D unstructured DataArray for testing. - - Parameters - ---------- - lazy : bool, default False - Whether to chunk the DataArray to make it Dask-backed. - - Returns - ------- - xr.DataArray - The 1D unstructured DataArray. - """ - n = 100 - lat = np.linspace(-90, 90, n) - lon = np.linspace(0, 360, n) - data = np.random.rand(n) - - da = xr.DataArray( - data, - dims=["ncol"], - coords={ - "lat": ( - ["ncol"], - lat, - {"units": "degrees_north", "standard_name": "latitude"}, - ), - "lon": ( - ["ncol"], - lon, - {"units": "degrees_east", "standard_name": "longitude"}, - ), - }, - name="test_data", - ) - - if lazy: - da = da.chunk({"ncol": 10}) - - return da - - -@patch("matplotlib.pyplot.axes") -@patch("matplotlib.pyplot.gca") -@patch("xarray.plot.accessor.DataArrayPlotAccessor.scatter") -def test_plot_static_unstructured( - mock_scatter: MagicMock, mock_gca: MagicMock, mock_axes: MagicMock -) -> None: - """ - Verify plot_static uses scatter for 1D unstructured data. - - Following the Aero Protocol's "Double-Check" rule, this test verifies - logic with both Eager (NumPy) and Lazy (Dask) data backends. - - Parameters - ---------- - mock_scatter : MagicMock - Mock for xarray's scatter plot accessor. - mock_gca : MagicMock - Mock for plt.gca(). - mock_axes : MagicMock - Mock for plt.axes(). - """ - # 1. Eager (NumPy) - da_eager = create_unstructured_da(lazy=False) - - plot_static(da_eager) - - # Check if scatter was called - mock_scatter.assert_called() - args, kwargs = mock_scatter.call_args - assert kwargs["x"] == "lon" - assert kwargs["y"] == "lat" - - # 2. Lazy (Dask) - da_lazy = create_unstructured_da(lazy=True) - - plot_static(da_lazy) - - assert mock_scatter.call_count == 2 - args, kwargs = mock_scatter.call_args - assert kwargs["x"] == "lon" - assert kwargs["y"] == "lat" - - -@patch("xarray.DataArray.hvplot") -def test_plot_interactive_unstructured(mock_hvplot: MagicMock) -> None: - """ - Verify plot_interactive uses kind='points' for 1D unstructured data. - - Parameters - ---------- - mock_hvplot : MagicMock - Mock for the hvplot accessor. - """ - # 1. Eager - da_eager = create_unstructured_da(lazy=False) - plot_interactive(da_eager) - - mock_hvplot.assert_called_with( - rasterize=True, title="Interactive Map", kind="points", x="lon", y="lat" - ) - - # 2. Lazy - da_lazy = create_unstructured_da(lazy=True) - plot_interactive(da_lazy) - - # Check if last call was for lazy data with correct parameters - mock_hvplot.assert_called_with( - rasterize=True, title="Interactive Map", kind="points", x="lon", y="lat" - ) - - -def test_find_coord_unstructured() -> None: - """ - Verify coordinate discovery for unstructured data. - """ - da = create_unstructured_da() - from xregrid.utils import _find_coord - - lat_da = _find_coord(da, "latitude") - lon_da = _find_coord(da, "longitude") - - assert lat_da.name == "lat" - assert lon_da.name == "lon" - assert lat_da.ndim == 1 - assert lon_da.ndim == 1 diff --git a/tests/test_aero_weight_loading.py b/tests/test_aero_weight_loading.py deleted file mode 100644 index 928e9bd..0000000 --- a/tests/test_aero_weight_loading.py +++ /dev/null @@ -1,93 +0,0 @@ -import os -import pytest -import numpy as np -import xarray as xr -from xregrid import Regridder, create_global_grid - - -def test_aero_weight_loading_roundtrip(tmp_path): - """ - Test the from_weights factory method and enhanced validation. - Verifies both NumPy and Dask data paths as per Aero Protocol. - """ - filename = str(tmp_path / "weights.nc") - - # Create grids - ds_src = create_global_grid(10.0, 10.0) - ds_tgt = create_global_grid(5.0, 5.0) - - # 1. Generate and save weights - # We use a context where esmpy is available (or mocked) - Regridder( - ds_src, - ds_tgt, - method="bilinear", - filename=filename, - reuse_weights=True, - skipna=True, - na_thres=0.5, - ) - - # Ensure file was created - assert os.path.exists(filename) - - # 2. Verify from_weights factory - regridder2 = Regridder.from_weights( - filename, ds_src, ds_tgt, method="bilinear", skipna=True, na_thres=0.5 - ) - - assert regridder2.method == "bilinear" - assert regridder2.skipna is True - assert regridder2.na_thres == 0.5 - - # 3. Test application with Eager (NumPy) data - data_np = np.random.rand(18, 36).astype(np.float32) - da_np = xr.DataArray( - data_np, - coords={"lat": ds_src.lat, "lon": ds_src.lon}, - dims=("lat", "lon"), - name="test", - ) - res_np = regridder2(da_np) - # Target shape is (36, 72) for 5 degree global grid - assert res_np.shape == (36, 72) - - # 4. Test application with Lazy (Dask) data - da_dask = da_np.chunk({"lat": 9, "lon": 18}) - res_dask = regridder2(da_dask) - assert res_dask.chunks is not None - - # 5. Assert equality (Eager vs Lazy) - # Mocked weights might return just a single point or something simple, - # but the shape and logic should hold. - np.testing.assert_allclose(res_np.values, res_dask.compute().values) - - # 6. Test Validation failure - Parameter mismatch - with pytest.raises(ValueError, match="Requested method"): - Regridder.from_weights(filename, ds_src, ds_tgt, method="nearest_s2d") - - with pytest.raises(ValueError, match="Requested skipna"): - Regridder.from_weights(filename, ds_src, ds_tgt, skipna=False) - - with pytest.raises(ValueError, match="Requested na_thres"): - Regridder.from_weights(filename, ds_src, ds_tgt, skipna=True, na_thres=0.9) - - -def test_aero_encoding_preservation(): - """Verify that encoding is preserved after regridding.""" - ds_src = create_global_grid(10.0, 10.0) - ds_tgt = create_global_grid(5.0, 5.0) - - regridder = Regridder(ds_src, ds_tgt) - - da_in = xr.DataArray( - np.random.rand(18, 36), - coords={"lat": ds_src.lat, "lon": ds_src.lon}, - dims=("lat", "lon"), - name="test", - ) - da_in.encoding = {"_FillValue": -999.0, "dtype": "float32"} - - res = regridder(da_in) - assert res.encoding["_FillValue"] == -999.0 - assert res.encoding["dtype"] == "float32" diff --git a/tests/test_cf_xarray.py b/tests/test_cf_xarray.py deleted file mode 100644 index 2785527..0000000 --- a/tests/test_cf_xarray.py +++ /dev/null @@ -1,106 +0,0 @@ -import xarray as xr -import numpy as np -import dask.array as da -from xregrid import Regridder - - -def test_cf_coords_detection(): - # Create dataset with non-standard coordinate names but with CF attributes - def create_ds(lazy=False): - np.random.seed(42) - data = np.random.rand(10, 20) - if lazy: - data = da.from_array(data, chunks=(5, 10)) - - ds = xr.Dataset( - {"data": (("lat_dim", "lon_dim"), data)}, - coords={ - "latitude": ( - ("lat_dim",), - np.linspace(-90, 90, 10), - {"units": "degrees_north", "standard_name": "latitude"}, - ), - "longitude": ( - ("lon_dim",), - np.linspace(-180, 180, 20), - {"units": "degrees_east", "standard_name": "longitude"}, - ), - }, - ) - return ds - - # Target grid also needs bounds for conservative regridding - lat_edges_tgt = np.linspace(-90, 90, 16) - lon_edges_tgt = np.linspace(-180, 180, 26) - ds_tgt = xr.Dataset( - coords={ - "lat": (("lat",), np.linspace(-90, 90, 15), {"units": "degrees_north"}), - "lon": (("lon",), np.linspace(-180, 180, 25), {"units": "degrees_east"}), - "lat_b": (("lat_b",), lat_edges_tgt, {"units": "degrees_north"}), - "lon_b": (("lon_b",), lon_edges_tgt, {"units": "degrees_east"}), - }, - ) - - # Test Eager - ds_src_eager = create_ds(lazy=False) - regridder_eager = Regridder(ds_src_eager, ds_tgt) - out_eager = regridder_eager(ds_src_eager["data"]) - assert out_eager.shape == (15, 25) - assert not out_eager.chunks - - # Test Lazy - ds_src_lazy = create_ds(lazy=True) - regridder_lazy = Regridder(ds_src_lazy, ds_tgt) - out_lazy = regridder_lazy(ds_src_lazy["data"]) - assert out_lazy.shape == (15, 25) - assert out_lazy.chunks - - # Verify results are identical (within float precision) - np.testing.assert_allclose(out_eager.values, out_lazy.compute().values) - - -def test_cf_bounds_detection(): - # Create dataset with non-standard bound names but with CF attributes - ds_src = xr.Dataset( - {"data": (("lat", "lon"), np.random.rand(10, 20))}, - coords={ - "lat": ( - ("lat",), - np.linspace(-90, 90, 10), - {"units": "degrees_north", "bounds": "lat_bounds"}, - ), - "lon": ( - ("lon",), - np.linspace(-180, 180, 20), - {"units": "degrees_east", "bounds": "lon_bounds"}, - ), - "lat_bounds": (("lat", "nv"), np.random.rand(10, 2)), # Placeholder bounds - "lon_bounds": (("lon", "nv"), np.random.rand(20, 2)), # Placeholder bounds - }, - ) - - # We need to make the bounds contiguous for our converter to work correctly in this test - lat_edges = np.linspace(-90, 90, 11) - lat_bounds = np.stack([lat_edges[:-1], lat_edges[1:]], axis=1) - lon_edges = np.linspace(-180, 180, 21) - lon_bounds = np.stack([lon_edges[:-1], lon_edges[1:]], axis=1) - - ds_src.coords["lat_bounds"] = (("lat", "nv"), lat_bounds) - ds_src.coords["lon_bounds"] = (("lon", "nv"), lon_bounds) - - # Target grid also needs bounds for conservative regridding - lat_edges_tgt = np.linspace(-90, 90, 16) - lon_edges_tgt = np.linspace(-180, 180, 26) - ds_tgt = xr.Dataset( - coords={ - "lat": (("lat",), np.linspace(-90, 90, 15), {"units": "degrees_north"}), - "lon": (("lon",), np.linspace(-180, 180, 25), {"units": "degrees_east"}), - "lat_b": (("lat_b",), lat_edges_tgt, {"units": "degrees_north"}), - "lon_b": (("lon_b",), lon_edges_tgt, {"units": "degrees_east"}), - }, - ) - - regridder = Regridder(ds_src, ds_tgt, method="conservative") - # If it reached here without error, it found the bounds and ESMPy initialized - out = regridder(ds_src["data"]) - assert out.shape == (15, 25) diff --git a/tests/test_cli.py b/tests/test_cli.py deleted file mode 100644 index 35192da..0000000 --- a/tests/test_cli.py +++ /dev/null @@ -1,93 +0,0 @@ -import subprocess -import sys -import pytest -import xarray as xr -import numpy as np -from unittest.mock import patch - - -@pytest.fixture -def sample_input(tmp_path): - path = tmp_path / "input.nc" - lat = np.arange(-89, 90, 2) - lon = np.arange(1, 360, 2) - data = np.random.rand(len(lat), len(lon)) - ds = xr.Dataset( - data_vars={"test": (["lat", "lon"], data)}, - coords={ - "lat": (["lat"], lat, {"units": "degrees_north"}), - "lon": (["lon"], lon, {"units": "degrees_east"}), - }, - ) - ds.to_netcdf(path) - return path - - -def test_cli_help(): - result = subprocess.run( - [sys.executable, "-m", "xregrid.cli", "--help"], capture_output=True, text=True - ) - assert result.returncode == 0 - assert "xregrid CLI" in result.stdout - - -def test_cli_basic(sample_input, tmp_path, monkeypatch): - output = tmp_path / "output.nc" - # We need to mock Regridder because esmpy is not installed - with patch("xregrid.cli.Regridder") as mock_regridder: - # Mock the regridder instance and its __call__ method - instance = mock_regridder.return_value - instance.return_value = xr.open_dataset( - sample_input - ) # Return input as mock output - - # Mock sys.argv - test_args = [ - "xregrid.cli", - str(sample_input), - "1.0", - "--output", - str(output), - "--method", - "bilinear", - ] - monkeypatch.setattr(sys, "argv", test_args) - - from xregrid.cli import main - - main() - - assert output.exists() - mock_regridder.assert_called_once() - args, kwargs = mock_regridder.call_args - assert kwargs["method"] == "bilinear" - assert args[1].lat.size == 180 # 1.0 degree global grid has 180 lat points - - -def test_cli_regional(sample_input, tmp_path, monkeypatch): - output = tmp_path / "output.nc" - with patch("xregrid.cli.Regridder") as mock_regridder: - instance = mock_regridder.return_value - instance.return_value = xr.open_dataset(sample_input) - - test_args = [ - "xregrid.cli", - str(sample_input), - "0.5", - "--output", - str(output), - "--extent=-10,10,20,40", - ] - monkeypatch.setattr(sys, "argv", test_args) - - from xregrid.cli import main - - main() - - assert output.exists() - mock_regridder.assert_called_once() - target_grid = mock_regridder.call_args[0][1] - assert target_grid.lat.min() >= -10 - assert target_grid.lat.max() <= 10 - assert target_grid.lon.min() >= 20 - assert target_grid.lon.max() <= 40 diff --git a/tests/test_dask_verification.py b/tests/test_dask_verification.py deleted file mode 100644 index 28b6b7c..0000000 --- a/tests/test_dask_verification.py +++ /dev/null @@ -1,147 +0,0 @@ -import xarray as xr -import numpy as np -import dask.distributed -from xregrid import Regridder, create_global_grid - -# Check for real ESMF -try: - import esmpy - - if hasattr(esmpy, "_is_mock") or "unittest.mock" in str(type(esmpy)): - raise ImportError - HAS_REAL_ESMF = True -except ImportError: - HAS_REAL_ESMF = False - - -def test_dask_parallel_regridding(): - """ - Test that running with parallel=True creates the same weights as serial execution. - Also verifies lazy initialization. - """ - # Create LocalCluster for testing - cluster = dask.distributed.LocalCluster( - n_workers=2, threads_per_worker=1, processes=HAS_REAL_ESMF - ) - client = dask.distributed.Client(cluster) - - try: - source_grid = create_global_grid(10, 10) - target_grid = create_global_grid(5, 5) - - # 1. Generate weights in serial - regridder_serial = Regridder( - source_grid, target_grid, method="bilinear", parallel=False - ) - w_serial = regridder_serial.weights - - # 2a. Generate weights in parallel (Eager) - try: - regridder_eager = Regridder( - source_grid, target_grid, method="bilinear", parallel=True, compute=True - ) - w_eager = regridder_eager.weights - except Exception as e: - print(f"ERROR in Regridder creation: {e}") - import traceback - - traceback.print_exc() - raise - - # 3a. Compare Eager - assert w_serial.shape == w_eager.shape - if HAS_REAL_ESMF: - assert w_serial.nnz == w_eager.nnz - diff_eager = w_serial - w_eager - assert np.abs(diff_eager.data).max() < 1e-10 if diff_eager.nnz > 0 else True - - # 2b. Generate weights in parallel (Lazy) - regridder_lazy = Regridder( - source_grid, target_grid, method="bilinear", parallel=True, compute=False - ) - - # Verify persist mechanism - assert regridder_lazy.persist() is regridder_lazy - - # Verify it hasn't computed yet - assert regridder_lazy._weights_matrix is None - assert regridder_lazy._dask_futures is not None - - # Trigger compute - regridder_lazy.compute() - w_lazy = regridder_lazy.weights - - assert w_lazy is not None - assert regridder_lazy._dask_futures is None - - # 3b. Compare Lazy - assert w_serial.shape == w_lazy.shape - if HAS_REAL_ESMF: - assert w_serial.nnz == w_lazy.nnz - diff_lazy = w_serial - w_lazy - assert np.abs(diff_lazy.data).max() < 1e-10 if diff_lazy.nnz > 0 else True - - # 4. Compare regridding result on dummy data - data = np.random.rand(source_grid.sizes["lat"], source_grid.sizes["lon"]) - da = xr.DataArray( - data, - coords={"lat": source_grid.lat, "lon": source_grid.lon}, - dims=["lat", "lon"], - ) - - res_serial = regridder_serial(da) - res_parallel = regridder_lazy(da) - - if HAS_REAL_ESMF: - xr.testing.assert_allclose(res_serial, res_parallel) - - # 5. Test auto-compute on call - regridder_auto = Regridder( - source_grid, target_grid, method="bilinear", parallel=True, compute=False - ) - assert regridder_auto._weights_matrix is None - res_auto = regridder_auto(da) - assert regridder_auto._weights_matrix is not None - if HAS_REAL_ESMF: - xr.testing.assert_allclose(res_serial, res_auto) - - finally: - client.close() - cluster.close() - - -def test_dask_curvilinear_parallel(): - """ - Test parallel regridding on curvilinear grids. - """ - from xregrid import create_grid_from_crs - - cluster = dask.distributed.LocalCluster( - n_workers=2, threads_per_worker=1, processes=HAS_REAL_ESMF - ) - client = dask.distributed.Client(cluster) - - try: - source_grid = create_grid_from_crs("EPSG:4326", (0, 10, 0, 10), 1) - target_grid = create_grid_from_crs( - "EPSG:3857", (0, 1000000, 0, 1000000), 100000 - ) - - regridder = Regridder( - source_grid, target_grid, method="bilinear", parallel=True - ) - assert regridder._weights_matrix is not None - - data = xr.DataArray( - np.random.rand(10, 10), - coords={"y": source_grid.y, "x": source_grid.x}, - dims=["y", "x"], - ) - data.coords["lat"] = source_grid.lat - data.coords["lon"] = source_grid.lon - - res = regridder(data) - assert res.shape == (10, 10) - finally: - client.close() - cluster.close() diff --git a/tests/test_dimension_robustness.py b/tests/test_dimension_robustness.py deleted file mode 100644 index 73ec61d..0000000 --- a/tests/test_dimension_robustness.py +++ /dev/null @@ -1,467 +0,0 @@ -import numpy as np -import xarray as xr -from xregrid import Regridder - - -def test_regridder_time_dimension_detection(): - # Setup source and target grids with time - lats = np.linspace(-90, 90, 10) - lons = np.linspace(0, 360, 20) - times = [np.datetime64("2020-01-01")] - - src_ds = xr.Dataset( - coords={ - "time": (["time"], times, {"standard_name": "time"}), - "lat": ( - ["time", "lat"], - np.broadcast_to(lats, (1, 10)), - {"units": "degrees_north", "standard_name": "latitude"}, - ), - "lon": ( - ["lon"], - lons, - {"units": "degrees_east", "standard_name": "longitude"}, - ), - } - ) - - tgt_ds = xr.Dataset( - coords={ - "lat": ( - ["lat"], - np.linspace(-90, 90, 5), - {"units": "degrees_north", "standard_name": "latitude"}, - ), - "lon": ( - ["lon"], - np.linspace(0, 360, 10), - {"units": "degrees_east", "standard_name": "longitude"}, - ), - } - ) - - # This should now work without failing during weight generation - regridder = Regridder(src_ds, tgt_ds, method="bilinear") - - # Create data with time and vertical dimensions - levs = np.arange(5) - data = np.random.rand(len(times), len(levs), len(lats), len(lons)) - da = xr.DataArray( - data, - coords={ - "time": (["time"], times), - "lev": (["lev"], levs), - "lat": (["time", "lat"], np.broadcast_to(lats, (1, 10))), - "lon": (["lon"], lons), - }, - dims=("time", "lev", "lat", "lon"), - name="temp", - ) - - # Regrid DataArray - res_da = regridder(da) - - # Check that time and lev are preserved - assert "time" in res_da.dims - assert "lev" in res_da.dims - assert res_da.shape == (1, 5, 5, 10) - - # Regrid Dataset - ds = xr.Dataset({"temp": da, "time_var": (["time"], times)}) - res_ds = regridder(ds) - - assert "time" in res_ds.dims - assert "temp" in res_ds.data_vars - assert "time_var" in res_ds.data_vars - assert res_ds["temp"].shape == (1, 5, 5, 10) - assert res_ds["time_var"].dims == ("time",) - - -def test_regridder_dtype_time_fallback(): - # Setup with time-like dtype but non-standard name - lats = np.linspace(-90, 90, 10) - lons = np.linspace(0, 360, 20) - times = [np.datetime64("2020-01-01")] - - src_ds = xr.Dataset( - coords={ - "mytime": (["mytime"], times), # Non-standard name, no CF attributes - "lat": ( - ["mytime", "lat"], - np.broadcast_to(lats, (1, 10)), - {"units": "degrees_north", "standard_name": "latitude"}, - ), - "lon": ( - ["lon"], - lons, - {"units": "degrees_east", "standard_name": "longitude"}, - ), - } - ) - - tgt_ds = xr.Dataset( - coords={ - "lat": ( - ["lat"], - np.linspace(-90, 90, 5), - {"units": "degrees_north", "standard_name": "latitude"}, - ), - "lon": ( - ["lon"], - np.linspace(0, 360, 10), - {"units": "degrees_east", "standard_name": "longitude"}, - ), - } - ) - - regridder = Regridder(src_ds, tgt_ds) - - # Verify mytime was detected as non-spatial - assert "mytime" not in regridder._dims_source - - # Test DataArray regridding with this non-standard time dim - da = xr.DataArray( - np.random.rand(1, 10, 20), coords=src_ds.coords, dims=("mytime", "lat", "lon") - ) - - res = regridder(da) - assert "mytime" in res.dims - assert res.shape == (1, 5, 10) - - -def test_non_regriddable_object(): - # Test passing something that shouldn't be regridded - lats = np.linspace(-90, 90, 10) - lons = np.linspace(0, 360, 20) - - src_ds = xr.Dataset( - coords={ - "lat": ( - ["lat"], - lats, - {"units": "degrees_north", "standard_name": "latitude"}, - ), - "lon": ( - ["lon"], - lons, - {"units": "degrees_east", "standard_name": "longitude"}, - ), - } - ) - tgt_ds = xr.Dataset( - coords={ - "lat": ( - ["lat"], - np.linspace(-90, 90, 5), - {"units": "degrees_north", "standard_name": "latitude"}, - ), - "lon": ( - ["lon"], - np.linspace(0, 360, 10), - {"units": "degrees_east", "standard_name": "longitude"}, - ), - } - ) - - regridder = Regridder(src_ds, tgt_ds) - - # A DataArray that only has one dimension (time) - time_da = xr.DataArray([1, 2, 3], dims="time", name="time_var") - - # Should return unchanged - res = regridder(time_da) - xr.testing.assert_identical(res, time_da) - - -def test_regridder_vertical_dimension_detection(): - # Setup source with vertical dimension in lats - lats = np.linspace(-90, 90, 10) - lons = np.linspace(0, 360, 20) - levs = np.arange(3) - - src_ds = xr.Dataset( - coords={ - "lev": (["lev"], levs, {"standard_name": "altitude"}), - "lat": ( - ["lev", "lat"], - np.broadcast_to(lats, (3, 10)), - {"units": "degrees_north", "standard_name": "latitude"}, - ), - "lon": ( - ["lon"], - lons, - {"units": "degrees_east", "standard_name": "longitude"}, - ), - } - ) - - tgt_ds = xr.Dataset( - coords={ - "lat": ( - ["lat"], - np.linspace(-90, 90, 5), - {"units": "degrees_north", "standard_name": "latitude"}, - ), - "lon": ( - ["lon"], - np.linspace(0, 360, 10), - {"units": "degrees_east", "standard_name": "longitude"}, - ), - } - ) - - regridder = Regridder(src_ds, tgt_ds) - assert "lev" not in regridder._dims_source - - da = xr.DataArray( - np.random.rand(3, 10, 20), coords=src_ds.coords, dims=("lev", "lat", "lon") - ) - - res = regridder(da) - assert "lev" in res.dims - assert res.shape == (3, 5, 10) - - -def test_regridder_ugrid_with_time(): - # Setup mocked uxarray object with time dimension - from unittest.mock import MagicMock - - class UxDatasetMock(xr.Dataset): - def __init__(self, ds, uxgrid): - super().__init__(ds.data_vars, coords=ds.coords, attrs=ds.attrs) - self.uxgrid = uxgrid - - n_face = 10 - n_node = 12 - times = [np.datetime64("2020-01-01")] - - mock_uxgrid = MagicMock() - mock_uxgrid.node_lat = xr.DataArray(np.linspace(-90, 90, n_node), dims=["n_node"]) - mock_uxgrid.node_lon = xr.DataArray(np.linspace(0, 360, n_node), dims=["n_node"]) - # face coords with time dimension (moving mesh case, though xregrid assumes static) - mock_uxgrid.face_lat = xr.DataArray( - np.broadcast_to(np.linspace(-90, 90, n_face), (1, n_face)), - dims=["time", "n_face"], - coords={"time": times}, - ) - mock_uxgrid.face_lon = xr.DataArray( - np.broadcast_to(np.linspace(0, 360, n_face), (1, n_face)), - dims=["time", "n_face"], - coords={"time": times}, - ) - - # Create connectivity - conn = np.zeros((n_face, 3), dtype=int) - for i in range(n_face): - conn[i] = [i, i + 1, (i + 2) % n_node] - - mock_uxgrid.face_node_connectivity = xr.DataArray( - conn, dims=["n_face", "n_max_face_nodes"] - ) - mock_uxgrid.face_node_connectivity.attrs["start_index"] = 0 - mock_uxgrid.face_node_connectivity.attrs["_FillValue"] = -1 - - # Mock UxDataset with time-varying variable - ds_base = xr.Dataset( - {"test_var": (["time", "n_face"], np.random.rand(1, n_face))}, - coords={"time": (["time"], times, {"standard_name": "time"})}, - ) - ds = UxDatasetMock(ds_base, mock_uxgrid) - - target_grid = xr.Dataset( - coords={ - "lat": ( - ["lat"], - np.linspace(-90, 90, 5), - {"units": "degrees_north", "standard_name": "latitude"}, - ), - "lon": ( - ["lon"], - np.linspace(0, 360, 10), - {"units": "degrees_east", "standard_name": "longitude"}, - ), - } - ) - - regridder = Regridder(ds, target_grid, method="nearest_s2d") - - # Verify time was detected as non-spatial - assert "time" not in regridder._dims_source - assert regridder._dims_source == ("n_face",) - - # Regrid DataArray - res = regridder(ds["test_var"]) - - assert "time" in res.dims - assert res.shape == (1, 5, 10) - - -def test_regridder_raw_ugrid_with_time(): - n_face = 10 - n_node = 12 - times = [np.datetime64("2020-01-01")] - - # Create a raw dataset following UGRID convention - conn = np.zeros((n_face, 3), dtype=int) - for i in range(n_face): - conn[i] = [i, (i + 1) % n_node, (i + 2) % n_node] - - ds = xr.Dataset( - data_vars={ - "temp": (["time", "n_face"], np.random.rand(1, n_face)), - "face_node_connectivity": (["n_face", "n_max_face_nodes"], conn), - }, - coords={ - "time": (["time"], times, {"standard_name": "time"}), - "lat_face": ( - ["time", "n_face"], - np.broadcast_to(np.linspace(-90, 90, n_face), (1, n_face)), - {"units": "degrees_north"}, - ), - "lon_face": ( - ["time", "n_face"], - np.broadcast_to(np.linspace(0, 360, n_face), (1, n_face)), - {"units": "degrees_east"}, - ), - "lat_node": ( - ["time", "n_node"], - np.broadcast_to(np.linspace(-90, 90, n_node), (1, n_node)), - {"units": "degrees_north"}, - ), - "lon_node": ( - ["time", "n_node"], - np.broadcast_to(np.linspace(0, 360, n_node), (1, n_node)), - {"units": "degrees_east"}, - ), - }, - ) - - ds.face_node_connectivity.attrs["cf_role"] = "face_node_connectivity" - ds.face_node_connectivity.attrs["start_index"] = 0 - - from xregrid import create_global_grid - - target_grid = create_global_grid(10, 10) - - regridder = Regridder(ds, target_grid, method="nearest_s2d") - - assert "time" not in regridder._dims_source - # Since it is UGRID, it should have detected n_face as the spatial dimension for variables - assert "n_face" in regridder._dims_source - - res = regridder(ds["temp"]) - assert "time" in res.dims - assert res.shape == (1, 18, 36) - - -def test_regridder_user_specific_structure(): - # Mimic user's dataset structure: (time, node) - # node is string coordinate, lat/lon are (node) - n_node = 10 - n_time = 5 - times = np.arange(n_time).astype("datetime64[D]") - nodes = np.array([f"NODE_{i}" for i in range(n_node)], dtype=" 0) - - if not HAS_REAL_ESMF: - pytest.skip("Skipping scientific correctness check for mocked ESMF") - - # Data should match target_lon - np.testing.assert_allclose(res.mean(dim="lat"), target_lon, atol=1e-5) - - -def test_output_order_preservation(): - """Test that the output preserves the coordinate order of the target grid.""" - lat_src = np.linspace(-90, 90, 10) - lon_src = np.linspace(0, 360, 10) - ds_src = xr.Dataset(coords={"lat": lat_src, "lon": lon_src}) - ds_src["data"] = (["lat", "lon"], np.random.rand(10, 10)) - - # Target grid with DESCENDING latitude - lat_tgt = np.linspace(90, -90, 10) - lon_tgt = np.linspace(0, 360, 10) - ds_tgt = xr.Dataset(coords={"lat": lat_tgt, "lon": lon_tgt}) - - regridder = Regridder(ds_src, ds_tgt) - res = regridder(ds_src["data"]) - - # Result should have descending latitude as requested - assert np.all(np.diff(res.lat) < 0) - np.testing.assert_allclose(res.lat, lat_tgt) diff --git a/tests/test_mpi.py b/tests/test_mpi.py deleted file mode 100644 index 05671d4..0000000 --- a/tests/test_mpi.py +++ /dev/null @@ -1,122 +0,0 @@ -import pytest -import xarray as xr -import numpy as np -import dask.array as da -from unittest.mock import MagicMock, patch -import sys -from xregrid import Regridder, create_global_grid - - -def test_mpi_initialization(): - """Test that mpi=True correctly initializes ESMF Manager.""" - source_grid = create_global_grid(10, 10) - target_grid = create_global_grid(20, 20) - - import esmpy - - with patch.object(esmpy, "Manager") as mock_manager: - _ = Regridder(source_grid, target_grid, mpi=True) - # Check if called. LogKind might be in esmpy or esmpy.LogKind depending on mock - assert mock_manager.called - - -def test_mpi_weight_gathering(): - """Test that weights are gathered correctly in an MPI environment.""" - source_grid = create_global_grid(10, 10) - target_grid = create_global_grid(20, 20) - - # Mock weights dictionary for each rank - weights_rank0 = { - "row_dst": np.array([1]), - "col_src": np.array([1]), - "weights": np.array([0.5]), - } - weights_rank1 = { - "row_dst": np.array([2]), - "col_src": np.array([2]), - "weights": np.array([0.5]), - } - - mock_mpi_pkg = MagicMock() - mock_mpi_internal = MagicMock() - mock_mpi_pkg.MPI = mock_mpi_internal - mock_comm = MagicMock() - mock_mpi_internal.COMM_WORLD = mock_comm - - import esmpy - - # Simulate rank 0 - with patch.dict(sys.modules, {"mpi4py": mock_mpi_pkg}): - with ( - patch.object(esmpy, "pet_count", return_value=2), - patch.object(esmpy, "local_pet", return_value=0), - patch.object(esmpy, "Regrid") as mock_regrid_class, - ): - mock_regrid = MagicMock() - mock_regrid.get_factors.return_value = (np.array([1]), np.array([1])) - mock_regrid.get_weights_dict.return_value = weights_rank0 - mock_regrid_class.return_value = mock_regrid - - # Rank 0 gather should receive both - mock_comm.gather.return_value = [weights_rank0, weights_rank1] - - regridder = Regridder(source_grid, target_grid, mpi=True) - - # Verify gathered matrix - matrix = regridder._weights_matrix.tocoo() - np.testing.assert_array_equal(matrix.row, [0, 1]) - np.testing.assert_array_equal(matrix.col, [0, 1]) - np.testing.assert_array_equal(matrix.data, [0.5, 0.5]) - - -def test_mpi_no_save_on_non_root(tmp_path): - """Test that non-root ranks do not save weights.""" - source_grid = create_global_grid(10, 10) - target_grid = create_global_grid(20, 20) - weight_file = str(tmp_path / "test_weights.nc") - - import esmpy - - with patch.object(esmpy, "local_pet", return_value=1): - with patch("xarray.Dataset.to_netcdf") as mock_to_netcdf: - _ = Regridder( - source_grid, - target_grid, - mpi=True, - reuse_weights=True, - filename=weight_file, - ) - mock_to_netcdf.assert_not_called() - - -def test_regrid_eager_lazy_identity(): - """Verify that Eager (NumPy) and Lazy (Dask) regridding produce identical results.""" - # Standard 1.0 -> 2.0 degree regridding - source_grid = create_global_grid(1.0, 1.0) - target_grid = create_global_grid(2.0, 2.0) - - regridder = Regridder(source_grid, target_grid, method="bilinear") - - # Create sample data - data = np.random.rand(180, 360) - da_eager = xr.DataArray( - data, - dims=["lat", "lon"], - coords={"lat": source_grid.lat, "lon": source_grid.lon}, - name="test", - ) - da_lazy = da_eager.chunk({"lat": 90, "lon": 180}) - - # Regrid both - res_eager = regridder(da_eager) - res_lazy = regridder(da_lazy) - - # Assert identity (only if real ESMF for exact values) - # But shapes should match anyway - assert res_eager.shape == res_lazy.shape - assert isinstance(res_lazy.data, da.Array) - assert not isinstance(res_eager.data, da.Array) - - -if __name__ == "__main__": - pytest.main([__file__]) diff --git a/tests/test_performance_optim.py b/tests/test_performance_optim.py deleted file mode 100644 index 0e9147f..0000000 --- a/tests/test_performance_optim.py +++ /dev/null @@ -1,130 +0,0 @@ -import numpy as np -from xregrid.core import _apply_weights_core, _WORKER_CACHE, _matmul -from scipy.sparse import csr_matrix -import dask.distributed -import xarray as xr -from unittest.mock import patch - - -def test_stationary_mask_caching(): - """Verify that stationary mask normalization is cached across calls.""" - # Clear cache - _WORKER_CACHE.clear() - - # Mock _matmul to count calls - with patch("xregrid.core._matmul", side_effect=_matmul) as mock_matmul: - # Create small synthetic data - data = np.ones((2, 4, 4), dtype=np.float32) - data[:, 0:2, 0:2] = np.nan # Stationary mask - - # Mock weight matrix (16 -> 4) - weights = np.zeros((4, 16)) - for i in range(4): - weights[i, i] = 1.0 - weights_sparse = csr_matrix(weights) - - weights_key = "test_weights_key" - _WORKER_CACHE[weights_key] = weights_sparse - - # 1st call: should compute weights_sum and cache it - res1 = _apply_weights_core( - data[0:1], weights_key, ("lat", "lon"), (2, 2), skipna=True - ) - - # count should be 2: one for data, one for weights_sum - assert mock_matmul.call_count == 2 - - # Reset mock count - mock_matmul.reset_mock() - - # 2nd call: should use cached weights_sum - res2 = _apply_weights_core( - data[1:2], weights_key, ("lat", "lon"), (2, 2), skipna=True - ) - - # count should be 1: only for data - assert mock_matmul.call_count == 1 - - # Verify results are identical - np.testing.assert_allclose(res1, res2) - - -def test_memory_efficiency_broadcasting(): - """Verify that stationary mask uses broadcasting.""" - data = np.ones((10, 4, 4), dtype=np.float32) - data[:, 0:2, 0:2] = np.nan - - weights = csr_matrix(np.eye(4, 16)) - - # Should not crash and should be correct - res = _apply_weights_core(data, weights, ("lat", "lon"), (2, 2), skipna=True) - assert res.shape == (10, 2, 2) - assert np.isnan(res[0, 0, 0]) - assert res[0, 1, 1] == 1.0 - - -def test_dask_stationary_mask_caching(): - """Verify stationary mask caching works with Dask-backed data.""" - from xregrid.regridder import Regridder - - # Setup local cluster with processes=False - cluster = dask.distributed.LocalCluster( - n_workers=1, threads_per_worker=1, processes=False - ) - client = dask.distributed.Client(cluster) - - try: - _WORKER_CACHE.clear() - - # Identity-like weight matrix - weights = csr_matrix(np.eye(4, 16)) - - # Source data (Dask-backed) - data = np.ones((4, 4, 4), dtype=np.float32) - data[:, 0:2, 0:2] = np.nan - da = xr.DataArray(data, dims=("time", "lat", "lon")).chunk({"time": 1}) - - ds_src = xr.Dataset( - {"data": da}, coords={"lat": np.arange(4), "lon": np.arange(4)} - ) - ds_dst = xr.Dataset(coords={"lat": np.arange(2), "lon": np.arange(2)}) - - # Use classes instead of MagicMock for mesh info to avoid pickling recursion - class MockObj: - def __init__(self, name): - self.name = name - - # Mock Regridder internally - with patch.object(Regridder, "_generate_weights", return_value=None): - with patch( - "xregrid.regridder._get_mesh_info", - side_effect=[ - (MockObj("src"), ["src"], (4, 4), ("lat", "lon"), False), - (MockObj("dst"), ["dst"], (2, 2), ("lat", "lon"), False), - ], - ): - with patch("xregrid.core._matmul", side_effect=_matmul) as mock_matmul: - regridder = Regridder(ds_src, ds_dst, skipna=True) - # Inject our known state - regridder._weights_matrix = weights - regridder._dims_source = ("lat", "lon") - regridder._dims_target = ("lat", "lon") - regridder._shape_target = (2, 2) - regridder._total_weights = np.array(weights.sum(axis=1)).T - - # Apply regridding - out = regridder(da) - - # Trigger compute - res = out.compute() - - # Verify results - assert res.shape == (4, 2, 2) - assert np.isnan(res[0, 0, 0]) - - # Verify calls - assert mock_matmul.call_count == 5 - - finally: - client.close() - cluster.close() diff --git a/tests/test_persistence.py b/tests/test_persistence.py deleted file mode 100644 index e999ea3..0000000 --- a/tests/test_persistence.py +++ /dev/null @@ -1,111 +0,0 @@ -import os -import numpy as np -import xarray as xr -from xregrid import Regridder, create_global_grid - - -def create_sample_data(dask=False, seed=None): - """Create sample data for testing.""" - if seed is not None: - np.random.seed(seed) - ds = create_global_grid(res_lat=10.0, res_lon=10.0) - data = np.random.rand(len(ds.lat), len(ds.lon)) - da = xr.DataArray( - data, - coords={"lat": ds.lat, "lon": ds.lon}, - dims=["lat", "lon"], - name="test_data", - ) - if dask: - da = da.chunk({"lat": 9, "lon": 18}) - return da, ds - - -def test_weight_persistence_eager_lazy(tmp_path): - """ - Test that weights can be saved and reloaded, and results are identical - for both eager and lazy data. - """ - filename = str(tmp_path / "weights.nc") - - # 1. Create source and target grids - source_da, source_grid = create_sample_data(dask=False) - target_grid = create_global_grid(res_lat=5.0, res_lon=5.0) - - # 2. Generate and save weights - regridder_save = Regridder( - source_grid, - target_grid, - method="bilinear", - reuse_weights=True, - filename=filename, - ) - - res_original = regridder_save(source_da) - assert os.path.exists(filename) - - # 3. Load weights in a new regridder - regridder_load = Regridder( - source_grid, - target_grid, - method="bilinear", - reuse_weights=True, - filename=filename, - ) - - # Verify eager result identity - res_eager = regridder_load(source_da) - xr.testing.assert_allclose(res_original, res_eager) - - # 4. Verify lazy result identity (Aero Protocol: Double-Check Test) - # Re-create source_da as dask with same values - source_da_lazy = source_da.chunk({"lat": 9, "lon": 18}) - res_lazy = regridder_load(source_da_lazy).compute() - - # Remove history from comparison as timestamps will differ - res_eager_no_hist = res_eager.copy() - res_lazy_no_hist = res_lazy.copy() - res_eager_no_hist.attrs.pop("history", None) - res_lazy_no_hist.attrs.pop("history", None) - - xr.testing.assert_allclose(res_eager_no_hist, res_lazy_no_hist) - - # Also verify that the loaded result matches the original fresh result - res_original_no_hist = res_original.copy() - res_original_no_hist.attrs.pop("history", None) - xr.testing.assert_allclose(res_original_no_hist, res_eager_no_hist) - - -def test_weight_persistence_skipna(tmp_path): - """Test persistence with skipna=True.""" - filename = str(tmp_path / "weights_skipna.nc") - - source_da, source_grid = create_sample_data(dask=False) - # Add some NaNs - source_da.values[0, 0] = np.nan - - target_grid = create_global_grid(res_lat=5.0, res_lon=5.0) - - regridder_save = Regridder( - source_grid, - target_grid, - method="bilinear", - reuse_weights=True, - filename=filename, - skipna=True, - ) - - res_save = regridder_save(source_da) - - regridder_load = Regridder( - source_grid, - target_grid, - method="bilinear", - reuse_weights=True, - filename=filename, - skipna=True, - ) - - res_load = regridder_load(source_da) - - xr.testing.assert_allclose(res_save, res_load) diff --git a/tests/test_rdhpcs_utils.py b/tests/test_rdhpcs_utils.py deleted file mode 100644 index b410636..0000000 --- a/tests/test_rdhpcs_utils.py +++ /dev/null @@ -1,79 +0,0 @@ -from unittest.mock import MagicMock, patch - -import numpy as np -import pytest -import xarray as xr -from xregrid.utils import get_rdhpcs_cluster -from xregrid.xregrid import _get_mesh_info - - -def test_get_mesh_info_rectilinear_order(): - """Test that _get_mesh_info correctly handles rectilinear grids with different coord orders.""" - # Create a grid where coords are (lon, lat) - lon = np.arange(0, 360, 10) - lat = np.arange(-90, 91, 10) - - # Broadcast to create a dataset - ds = xr.Dataset(coords={"lat": lat, "lon": lon}) - - # Check (lat, lon) order - lon_m, lat_m, shape, dims, unstructured = _get_mesh_info(ds) - assert not unstructured - assert dims == ("lat", "lon") - assert shape == (lat.size, lon.size) - assert lat_m.shape == (lat.size, lon.size) - assert lon_m.shape == (lat.size, lon.size) - - # Verify that the dead code was removed and it still works - # (The test above already confirms it works for standard 1D coords) - - -def test_get_rdhpcs_cluster_detection(): - """Test machine detection in get_rdhpcs_cluster.""" - - with patch("socket.gethostname") as mock_hostname: - # Test Hera detection - mock_hostname.return_value = "hfe01.hera.noaa.gov" - with patch("dask_jobqueue.SLURMCluster", MagicMock()) as mock_slurm: - get_rdhpcs_cluster(account="test_acc") - args, kwargs = mock_slurm.call_args - assert kwargs["queue"] == "hera" - assert kwargs["cores"] == 40 - - # Test Jet detection - mock_hostname.return_value = "fe01.jet.noaa.gov" - with patch("dask_jobqueue.SLURMCluster", MagicMock()) as mock_slurm: - get_rdhpcs_cluster(account="test_acc") - args, kwargs = mock_slurm.call_args - assert kwargs["queue"] == "batch" - assert kwargs["cores"] == 24 - - # Test Gaea detection - mock_hostname.return_value = "gaea12.ncrc.gov" - with patch("dask_jobqueue.SLURMCluster", MagicMock()) as mock_slurm: - get_rdhpcs_cluster(account="test_acc", machine="gaea-c6") - args, kwargs = mock_slurm.call_args - assert kwargs["cores"] == 192 - assert "-M c6" in kwargs["job_extra_directives"][0] - - # Test Ursa detection - mock_hostname.return_value = "ufe01.ursa.noaa.gov" - with patch("dask_jobqueue.SLURMCluster", MagicMock()) as mock_slurm: - get_rdhpcs_cluster(account="test_acc") - args, kwargs = mock_slurm.call_args - assert kwargs["queue"] == "u1-compute" - assert kwargs["cores"] == 192 - - -def test_get_rdhpcs_cluster_explicit(): - """Test explicit machine specification in get_rdhpcs_cluster.""" - with patch("dask_jobqueue.SLURMCluster", MagicMock()) as mock_slurm: - get_rdhpcs_cluster(machine="hera", account="test_acc", walltime="02:00:00") - args, kwargs = mock_slurm.call_args - assert kwargs["queue"] == "hera" - assert kwargs["walltime"] == "02:00:00" - assert kwargs["account"] == "test_acc" - - -if __name__ == "__main__": - pytest.main([__file__]) diff --git a/tests/test_real_esmpy_dask.py b/tests/test_real_esmpy_dask.py deleted file mode 100644 index 7d266b3..0000000 --- a/tests/test_real_esmpy_dask.py +++ /dev/null @@ -1,120 +0,0 @@ -import pytest -import xarray as xr -import numpy as np -import dask.distributed -from xregrid import Regridder, create_global_grid -from unittest.mock import MagicMock - - -# Check if esmpy is mocked -def is_esmpy_mocked(): - try: - import esmpy - - return ( - hasattr(esmpy, "_is_mock") - or isinstance(esmpy, MagicMock) - or "MagicMock" in str(type(esmpy)) - ) - except ImportError: - return True - - -HAS_REAL_ESMF = not is_esmpy_mocked() -pytestmark = pytest.mark.skipif(not HAS_REAL_ESMF, reason="esmpy is missing or mocked") - - -@pytest.fixture(scope="module") -def dask_client(): - # Use processes=True for real ESMF thread-safety - cluster = dask.distributed.LocalCluster( - n_workers=2, threads_per_worker=1, processes=HAS_REAL_ESMF - ) - client = dask.distributed.Client(cluster) - yield client - client.close() - cluster.close() - - -def test_global_rectilinear_regrid_dask(dask_client): - # Test global grid regridding which was reported to fail - res_src = 1.0 - res_dst = 2.0 - - source_grid = create_global_grid(res_src, res_src) - target_grid = create_global_grid(res_dst, res_dst) - - # Add data with multiple time steps to test dask application over time - nt = 5 - data = xr.DataArray( - np.random.rand(nt, source_grid.sizes["lat"], source_grid.sizes["lon"]), - coords={"time": np.arange(nt), "lat": source_grid.lat, "lon": source_grid.lon}, - dims=["time", "lat", "lon"], - name="air", - ).chunk({"time": 1}) - - # Initialize Regridder with parallel=True - # Testing both periodic=True and False - for periodic in [True, False]: - print(f"Testing periodic={periodic}...") - regridder = Regridder( - source_grid, - target_grid, - method="bilinear", - parallel=True, - periodic=periodic, - ) - - res = regridder(data) - - assert res.shape == (nt, target_grid.sizes["lat"], target_grid.sizes["lon"]) - # Trigger computation - res_computed = res.compute() - assert not np.isnan(res_computed).all() - print(f"Periodic={periodic} success!") - - -def test_descending_lat_regrid_dask(dask_client): - # Test descending latitudes which was identified as a bug - source_grid = create_global_grid(1.0, 1.0) - source_grid = source_grid.sortby("lat", ascending=False) - - target_grid = create_global_grid(2.0, 2.0) - - data = xr.DataArray( - np.random.rand(source_grid.sizes["lat"], source_grid.sizes["lon"]), - coords={"lat": source_grid.lat, "lon": source_grid.lon}, - dims=["lat", "lon"], - name="test", - ).chunk({"lat": 45}) - - regridder = Regridder( - source_grid, target_grid, method="bilinear", parallel=True, periodic=True - ) - res = regridder(data).compute() - assert res.shape == (target_grid.sizes["lat"], target_grid.sizes["lon"]) - - -def test_diagnostics_dask(dask_client): - source_grid = create_global_grid(10.0, 10.0) - target_grid = create_global_grid(20.0, 20.0) - - regridder = Regridder(source_grid, target_grid, method="bilinear", parallel=True) - - # Test diagnostics - diag = regridder.diagnostics() - assert "weight_sum" in diag - assert "unmapped_mask" in diag - - # Trigger computation of lazy diagnostics - ws = diag.weight_sum.compute() - assert ws.shape == (target_grid.sizes["lat"], target_grid.sizes["lon"]) - - # Test quality report - report = regridder.quality_report() - assert "n_src" in report - assert "unmapped_fraction" in report - - -if __name__ == "__main__": - pytest.main([__file__]) diff --git a/tests/test_regridder_coverage.py b/tests/test_regridder_coverage.py deleted file mode 100644 index 4fbb00e..0000000 --- a/tests/test_regridder_coverage.py +++ /dev/null @@ -1,487 +0,0 @@ -import os -import sys -from unittest.mock import MagicMock, patch - -import numpy as np -import pytest -import xarray as xr - -from xregrid import Regridder, create_global_grid -from xregrid.core import _WORKER_CACHE -from xregrid.grid import ( - _create_esmf_grid, - _get_mesh_info, - _get_non_spatial_dims, - _get_unstructured_mesh_info, -) -from xregrid.parallel import ( - _assemble_weights_task, - _get_nnz_task, - _sync_cache_from_worker_data, -) - - -def test_regridder_mpi_parallel_error(): - """Verify ValueError when both mpi and parallel are True.""" - src = create_global_grid(10, 10) - tgt = create_global_grid(5, 5) - with pytest.raises(ValueError, match="Cannot use both MPI and Dask"): - Regridder(src, tgt, mpi=True, parallel=True) - - -def test_regridder_missing_dask_error(): - """Verify ImportError when parallel=True but dask.distributed is missing.""" - src = create_global_grid(10, 10) - tgt = create_global_grid(5, 5) - with patch("importlib.util.find_spec", return_value=None): - with pytest.raises(ImportError, match="Dask distributed is required"): - Regridder(src, tgt, parallel=True) - - -def test_regridder_save_load_weights(tmp_path): - """Verify saving and loading weights.""" - src = create_global_grid(10, 10) - tgt = create_global_grid(5, 5) - weight_file = str(tmp_path / "weights.nc") - - # Weights are saved automatically if reuse_weights=True and file doesn't exist - Regridder(src, tgt, method="bilinear", filename=weight_file, reuse_weights=True) - assert os.path.exists(weight_file) - - # Load weights - regridder2 = Regridder.from_weights(weight_file, src, tgt) - assert regridder2.method == "bilinear" - - # Verify validation fails with wrong parameters - with pytest.raises(ValueError, match="does not match loaded weights method"): - Regridder.from_weights(weight_file, src, tgt, method="conservative") - - -def test_regrid_dataset_coverage(): - """Verify _regrid_dataset with various variable types.""" - src = create_global_grid(10, 10) - tgt = create_global_grid(5, 5) - - ds = xr.Dataset( - data_vars={ - "var1": (["lat", "lon"], np.random.rand(18, 36)), - "var2": (["lat", "lon"], np.random.rand(18, 36)), - "scalar": 42, - "other": (["time"], [1, 2, 3]), - }, - coords={"lat": src.lat, "lon": src.lon, "time": [0, 1, 2]}, - ) - - regridder = Regridder(src, tgt) - res = regridder(ds) - - assert "var1" in res.data_vars - assert "var2" in res.data_vars - assert "scalar" in res.data_vars - assert "other" in res.data_vars - # create_global_grid(5, 5) gives (36, 72) - assert res.var1.shape == (36, 72) - - -def test_extrap_methods_coverage(): - """Verify different extrapolation methods.""" - src = create_global_grid(10, 10) - tgt = create_global_grid(5, 5) - - for method in ["nearest_s2d", "nearest_idw", "creep_fill"]: - regridder = Regridder(src, tgt, extrap_method=method, extrap_dist_exponent=3.0) - assert regridder.extrap_method == method - - -def test_regridder_repr_lazy(): - """Verify __repr__ with lazy weights.""" - src = create_global_grid(10, 10) - tgt = create_global_grid(5, 5) - - regridder = Regridder(src, tgt, parallel=True, compute=False) - - class MockFuture: - def __init__(self): - self.key = "some_key" - - regridder._weights_matrix = MockFuture() - - repr_str = repr(regridder) - assert "quality=lazy" in repr_str - - -def test_regridder_quality_report_coverage(): - """Verify quality_report with different options.""" - src = create_global_grid(10, 10) - tgt = create_global_grid(5, 5) - regridder = Regridder(src, tgt) - - report = regridder.quality_report(format="dataset") - assert isinstance(report, xr.Dataset) - assert "unmapped_fraction" in report.data_vars - - -# --- Additional Coverage Tests --- - - -def test_get_non_spatial_dims_z_axis(): - """Verify Z axis detection in _get_non_spatial_dims.""" - ds = xr.Dataset(coords={"lev": [1, 2, 3]}) - ds.lev.attrs["axis"] = "Z" - dims = _get_non_spatial_dims(ds) - assert "lev" in dims - - -def test_get_mesh_info_uxarray_mock(): - """Verify uxarray object support in _get_mesh_info.""" - mock_uxgrid = MagicMock() - mock_uxgrid.node_lat = xr.DataArray([0.0, 10.0], dims="n_node") - mock_uxgrid.node_lon = xr.DataArray([0.0, 20.0], dims="n_node") - - class MockDS: - def __init__(self): - self.uxgrid = mock_uxgrid - self.data_vars = {} - self.dims = {} - self.coords = {} - - def isel(self, *args, **kwargs): - return self - - ds = MockDS() - - lon, lat, shape, dims, is_unstructured = _get_mesh_info(ds) - assert is_unstructured - assert "n_node" in dims - assert shape == (2,) - - -def test_get_mesh_info_lat_node_fallback(): - """Verify fallback to lat_node/lon_node in _get_mesh_info.""" - ds = xr.Dataset( - coords={ - "lat_node": (["n_node"], [0.0, 10.0]), - "lon_node": (["n_node"], [0.0, 20.0]), - } - ) - lon, lat, shape, dims, is_unstructured = _get_mesh_info(ds) - assert is_unstructured - assert "n_node" in dims - - -def test_get_unstructured_mesh_info_mpas_no_nedges(): - """Verify MPAS support without nEdgesOnCell in _get_unstructured_mesh_info.""" - ds = xr.Dataset( - data_vars={ - "verticesOnCell": (["nCells", "maxEdges"], [[1, 2, 3]]), - "latVertex": (["nVertices"], [0.0, 1.0, 2.0]), - "lonVertex": (["nVertices"], [0.0, 1.0, 2.0]), - } - ) - res = _get_unstructured_mesh_info(ds) - assert res is not None - assert len(res) == 6 - - -def test_get_unstructured_mesh_info_ugrid_node_coords_attr(): - """Verify UGRID node_coordinates attribute support.""" - ds = xr.Dataset( - data_vars={ - "mesh": ( - [], - 0, - { - "cf_role": "mesh_topology", - "face_node_connectivity": "face_nodes", - "node_coordinates": "lon_u lat_u", - }, - ), - "face_nodes": (["n_face", "n_node_per_face"], [[0, 1, 2]]), - "lon_u": (["n_node"], [0.0, 1.0, 2.0]), - "lat_u": (["n_node"], [0.0, 1.0, 2.0]), - } - ) - ds.face_nodes.attrs["start_index"] = 0 - res = _get_unstructured_mesh_info(ds) - assert res is not None - assert np.allclose(res[0], [0.0, 1.0, 2.0]) - - -def test_create_esmf_grid_periodic_bounds(): - """Verify periodic grid with bounds in _create_esmf_grid.""" - lon = np.linspace(0, 360, 36, endpoint=False) - lat = np.linspace(-90, 90, 19) - ds = xr.Dataset(coords={"lat": lat, "lon": lon}) - ds.lat.attrs["standard_name"] = "latitude" - ds.lon.attrs["standard_name"] = "longitude" - ds.coords["lat_b"] = (["lat", "nv"], np.zeros((19, 2))) - ds.coords["lon_b"] = (["lon", "nv"], np.zeros((36, 2))) - ds.lat.attrs["bounds"] = "lat_b" - ds.lon.attrs["bounds"] = "lon_b" - - with patch("esmpy.Grid") as mock_grid_cls: - mock_grid = MagicMock() - mock_grid.get_coords.side_effect = lambda dim, staggerloc: np.zeros( - (36, 19) if staggerloc == 0 else (36, 20) - ) - mock_grid_cls.return_value = mock_grid - - grid, prov, _ = _create_esmf_grid(ds, method="conservative", periodic=True) - assert grid is not None - - -def test_create_esmf_grid_locstream_cart(): - """Verify LocStream creation with CART coordinate system.""" - ds = xr.Dataset( - coords={ - "lat": (["n"], [0.0, 10.0]), - "lon": (["n"], [0.0, 20.0]), - } - ) - import esmpy - - grid, prov, _ = _create_esmf_grid( - ds, method="nearest_s2d", periodic=False, coord_sys=esmpy.CoordSys.CART - ) - assert isinstance(grid, esmpy.LocStream) - - -def test_regridder_normalize_descending(): - """Verify sorting of descending coordinates in Regridder.""" - src = xr.Dataset( - coords={ - "lat": (["lat"], [10.0, 0.0]), - "lon": (["lon"], [0.0, 10.0]), - } - ) - src.lat.attrs["units"] = "degrees_north" - src.lon.attrs["units"] = "degrees_east" - tgt = create_global_grid(10, 10) - - regridder = Regridder(src, tgt) - assert regridder._src_was_sorted - assert regridder.source_grid_ds.lat.values[0] == 0.0 - - -def test_regridder_mpi_non_root_rank(): - """Verify logic for non-root ranks in MPI mode.""" - - mock_mpi = MagicMock() - mock_comm = MagicMock() - mock_mpi.COMM_WORLD = mock_comm - mock_comm.gather.return_value = None - - with ( - patch("esmpy.pet_count", return_value=2), - patch("esmpy.local_pet", return_value=1), - patch.dict(sys.modules, {"mpi4py": mock_mpi}), - ): - src = create_global_grid(10, 10) - tgt = create_global_grid(5, 5) - regridder = Regridder(src, tgt, mpi=True) - assert regridder.weights.nnz == 0 - - -def test_assemble_weights_task_error(): - """Verify error handling in _assemble_weights_task.""" - results = [(None, None, None, "Mock Error")] - with pytest.raises(RuntimeError, match="Weight generation error: Mock Error"): - _assemble_weights_task(results, 10, 10) - - -def test_assemble_weights_task_empty(): - """Verify _assemble_weights_task with no results.""" - results = [(np.array([]), np.array([]), np.array([]), None)] - res = _assemble_weights_task(results, 5, 5) - assert res.nnz == 0 - assert res.shape == (5, 5) - - -def test_sync_cache_from_worker_data_fallback(): - """Verify fallback to get_worker in _sync_cache_from_worker_data.""" - with patch("dask.distributed.get_worker") as mock_get_worker: - mock_worker = MagicMock() - mock_worker.data = {"f_key": "val"} - mock_get_worker.return_value = mock_worker - - _sync_cache_from_worker_data("f_key", "c_key") - assert _WORKER_CACHE["c_key"] == "val" - - -def test_regrid_dataset_non_spatial_preservation(): - """Verify preservation of non-spatial variables in _regrid_dataset.""" - src = create_global_grid(10, 10) - tgt = create_global_grid(5, 5) - ds = xr.Dataset( - data_vars={ - "spatial": (["lat", "lon"], np.random.rand(18, 36)), - "time_var": (["time"], [1, 2, 3]), - }, - coords={"lat": src.lat, "lon": src.lon, "time": [0, 1, 2]}, - ) - regridder = Regridder(src, tgt) - res = regridder(ds) - assert "time_var" in res.data_vars - assert res.time_var.dims == ("time",) - - -def test_regrid_dataarray_aux_coords(): - """Verify regridding of auxiliary spatial coordinates.""" - src = create_global_grid(10, 10) - tgt = create_global_grid(5, 5) - da = xr.DataArray( - np.random.rand(18, 36), - dims=("lat", "lon"), - coords={ - "lat": src.lat, - "lon": src.lon, - "aux": (["lat", "lon"], np.random.rand(18, 36)), - }, - name="test", - ) - regridder = Regridder(src, tgt) - res = regridder(da) - assert "aux" in res.coords - assert res.aux.shape == (36, 72) - - -def test_regrid_dataarray_dim_renaming(): - """Verify dimension renaming logic in _regrid_dataarray.""" - src = create_global_grid(10, 10) - tgt = create_global_grid(5, 5) - da = xr.DataArray( - np.random.rand(18, 36), - dims=("y", "x"), - name="test", - ) - da.coords["y"] = (["y"], np.linspace(-90, 90, 18)) - da.coords["x"] = (["x"], np.linspace(0, 360, 36)) - da.y.attrs["standard_name"] = "latitude" - da.x.attrs["standard_name"] = "longitude" - - regridder = Regridder(src, tgt) - res = regridder(da) - assert res.shape == (36, 72) - - -def test_quality_report_skip_heavy_remote(): - """Verify skip_heavy=True logic for remote weights in quality_report.""" - src = create_global_grid(10, 10) - tgt = create_global_grid(5, 5) - regridder = Regridder(src, tgt, parallel=True, compute=False) - - class MockFuture: - def __init__(self): - self.key = "future_key" - - regridder._weights_matrix = MockFuture() - report = regridder.quality_report(skip_heavy=True) - assert report["n_weights"] == -1 - - -def test_quality_report_nnz_remote_exception(): - """Verify exception handling when computing nnz remotely.""" - src = create_global_grid(10, 10) - tgt = create_global_grid(5, 5) - regridder = Regridder(src, tgt, parallel=True, compute=False) - - class MockFuture: - def __init__(self): - self.key = "future_key" - - regridder._weights_matrix = MockFuture() - regridder._dask_client = MagicMock() - regridder._dask_client.submit.side_effect = ValueError("Mock Error") - report = regridder.quality_report(skip_heavy=True) - assert report["n_weights"] == -1 - - -def test_regrid_dataset_grid_mapping_removal(): - """Verify removal of stale grid_mapping in _regrid_dataset.""" - src = create_global_grid(10, 10) - tgt = create_global_grid(5, 5) - ds = xr.Dataset( - {"var": (["lat", "lon"], np.random.rand(18, 36))}, coords=src.coords - ) - ds.attrs["grid_mapping"] = "crs" - ds.coords["crs"] = ([], 0, {"grid_mapping_name": "latitude_longitude"}) - regridder = Regridder(src, tgt) - res = regridder(ds) - assert "grid_mapping" not in res.attrs - - -def test_regrid_dataset_ugrid_attr_removal(): - """Verify removal of UGRID attributes when target is not UGRID.""" - src = create_global_grid(10, 10) - tgt = create_global_grid(5, 5) - # Filter coords to match dimensions - da_coords = { - k: v for k, v in src.coords.items() if set(v.dims).issubset({"lat", "lon"}) - } - da = xr.DataArray( - np.random.rand(18, 36), dims=("lat", "lon"), coords=da_coords, name="var" - ) - da.attrs["mesh"] = "some_mesh" - da.attrs["location"] = "face" - regridder = Regridder(src, tgt) - res = regridder(da) - assert "mesh" not in res.attrs - assert "location" not in res.attrs - - -def test_get_nnz_task(): - """Verify _get_nnz_task.""" - matrix = MagicMock() - matrix.nnz = 42 - assert _get_nnz_task(matrix) == 42 - - -def test_regridder_persist_non_parallel(): - """Verify persist() on non-parallel regridder.""" - src = create_global_grid(10, 10) - tgt = create_global_grid(5, 5) - regridder = Regridder(src, tgt, parallel=False) - res = regridder.persist() - assert res is regridder - - -def test_regridder_validate_weights_errors(tmp_path): - """Verify error paths in _validate_weights.""" - src = create_global_grid(10, 10) - tgt = create_global_grid(5, 5) - weight_file = str(tmp_path / "weights_err.nc") - - regridder = Regridder( - src, tgt, method="bilinear", filename=weight_file, reuse_weights=True - ) - - with pytest.raises(ValueError, match="does not match loaded weights periodic"): - Regridder.from_weights(weight_file, src, tgt, periodic=not regridder.periodic) - - with pytest.raises(ValueError, match="does not match loaded weights skipna"): - Regridder.from_weights(weight_file, src, tgt, skipna=True) - - with pytest.raises(ValueError, match="does not match loaded weights na_thres"): - Regridder.from_weights(weight_file, src, tgt, skipna=False, na_thres=0.5) - - -def test_plot_static_unstructured_fallback(): - """Verify fallback logic in plot_static for unstructured grids.""" - from xregrid.viz import plot_static - - da = xr.DataArray([1.0, 2.0], dims=("cell",), name="test") - # No lat/lon coordinates or attributes - with patch("matplotlib.pyplot.subplots", return_value=(MagicMock(), MagicMock())): - with patch("xregrid.utils.get_crs_info", return_value=None): - # This should hit the unstructured fallback in viz.py - ax = plot_static(da) - assert ax is not None - - -def test_plot_diagnostics_invalid_mode(): - """Verify ValueError for invalid mode in plot_diagnostics.""" - src = create_global_grid(10, 10) - tgt = create_global_grid(5, 5) - regridder = Regridder(src, tgt) - with pytest.raises(ValueError, match="Unknown plotting mode"): - regridder.plot_diagnostics(mode="invalid") diff --git a/tests/test_robustness.py b/tests/test_robustness.py deleted file mode 100644 index 15f304d..0000000 --- a/tests/test_robustness.py +++ /dev/null @@ -1,60 +0,0 @@ -import os -import numpy as np -import pytest -import xarray as xr -from xregrid import Regridder, create_global_grid - - -def test_eager_lazy_identity_and_name_preservation(): - """Verify Eager and Lazy results are identical and name/history are preserved.""" - ds_in = create_global_grid(10, 10) - ds_out = create_global_grid(5, 5) - regridder = Regridder(ds_in, ds_out, method="bilinear") - - data = np.random.rand(18, 36) - coords = {"lat": ds_in.lat, "lon": ds_in.lon} - name = "test_var" - attrs = {"units": "K", "history": "original history"} - - # 1. Eager (NumPy) - da_eager = xr.DataArray( - data, dims=("lat", "lon"), coords=coords, name=name, attrs=attrs - ) - res_eager = regridder(da_eager) - - # 2. Lazy (Dask) - da_lazy = xr.DataArray( - data, dims=("lat", "lon"), coords=coords, name=name, attrs=attrs - ).chunk({"lat": 9, "lon": 18}) - res_lazy = regridder(da_lazy) - - # Identity check - xr.testing.assert_allclose(res_eager, res_lazy.compute()) - - # Metadata checks - for res in [res_eager, res_lazy]: - assert res.name == name - assert res.attrs["units"] == "K" - assert "original history" in res.attrs["history"] - assert "Regridded" in res.attrs["history"] - - -def test_weight_validation(tmp_path): - """Verify that loaded weights are validated against the grid.""" - ds_in = create_global_grid(10, 10) - ds_out = create_global_grid(5, 5) - - weights_file = str(tmp_path / "weights.nc") - - # Generate weights for 10x10 -> 5x5 - Regridder(ds_in, ds_out, reuse_weights=True, filename=weights_file) - assert os.path.exists(weights_file) - - # Now try to reuse these weights for a DIFFERENT grid - ds_in_wrong = create_global_grid(20, 20) - with pytest.raises(ValueError, match="Source grid shape"): - Regridder(ds_in_wrong, ds_out, reuse_weights=True, filename=weights_file) - - -if __name__ == "__main__": - pytest.main([__file__]) diff --git a/tests/test_toy_regrid.py b/tests/test_toy_regrid.py deleted file mode 100644 index cb9eb45..0000000 --- a/tests/test_toy_regrid.py +++ /dev/null @@ -1,61 +0,0 @@ -import numpy as np -import pytest -import xarray as xr -from xregrid import Regridder, create_global_grid - - -def test_toy_regrid_global_05(): - """Test regridding the xarray toy dataset to a global 0.5 degree grid.""" - # 1. Load toy dataset - # We use a subset of time to keep the test fast - ds = xr.tutorial.open_dataset("air_temperature").isel(time=slice(0, 2)) - - # 2. Create target global 0.5 degree grid - # A 0.5 degree global grid has 360 latitude points and 720 longitude points - target_res = 0.5 - target_grid = create_global_grid(res_lat=target_res, res_lon=target_res) - - expected_lat_size = int(180 / target_res) - expected_lon_size = int(360 / target_res) - - assert target_grid.lat.size == expected_lat_size - assert target_grid.lon.size == expected_lon_size - - # 3. Initialize Regridder - # Note: In the test environment, ESMF is mocked, so weight generation is synthetic - regridder = Regridder(ds, target_grid, method="bilinear", periodic=True) - - # 4. Regrid the DataArray - air_regridded = regridder(ds.air) - - # 5. Verify DataArray output - assert isinstance(air_regridded, xr.DataArray) - assert air_regridded.shape == (ds.time.size, expected_lat_size, expected_lon_size) - assert "lat" in air_regridded.coords - assert "lon" in air_regridded.coords - assert "time" in air_regridded.coords - - # Check that coordinates match the target grid - np.testing.assert_allclose(air_regridded.lat, target_grid.lat) - np.testing.assert_allclose(air_regridded.lon, target_grid.lon) - - # 6. Regrid the Dataset - ds_regridded = regridder(ds) - - # 7. Verify Dataset output - assert isinstance(ds_regridded, xr.Dataset) - assert "air" in ds_regridded.data_vars - assert ds_regridded.air.shape == ( - ds.time.size, - expected_lat_size, - expected_lon_size, - ) - - # Verify provenance - assert "history" in ds_regridded.attrs - assert "Regridder" in ds_regridded.attrs["history"] - assert "bilinear" in ds_regridded.attrs["history"] - - -if __name__ == "__main__": - pytest.main([__file__]) diff --git a/tests/test_unstructured_dask.py b/tests/test_unstructured_dask.py deleted file mode 100644 index 1075961..0000000 --- a/tests/test_unstructured_dask.py +++ /dev/null @@ -1,92 +0,0 @@ -import pytest -import xarray as xr -import numpy as np -import dask.distributed -from xregrid import Regridder, create_global_grid, create_mesh_from_coords - -# Check for real ESMF -try: - import esmpy - - if hasattr(esmpy, "_is_mock") or "unittest.mock" in str(type(esmpy)): - raise ImportError - HAS_REAL_ESMF = True -except ImportError: - HAS_REAL_ESMF = False - - -@pytest.fixture(scope="module") -def dask_client(): - # esmpy is not thread-safe, so we must use processes=True when using real ESMF - cluster = dask.distributed.LocalCluster( - n_workers=2, threads_per_worker=1, processes=HAS_REAL_ESMF - ) - client = dask.distributed.Client(cluster) - yield client - client.close() - cluster.close() - - -def test_regrid_structured_to_unstructured_dask(dask_client): - source_grid = create_global_grid(10, 10) - - n_pts = 50 - lon = np.linspace(0, 360, n_pts) - lat = np.linspace(-90, 90, n_pts) - target_grid = create_mesh_from_coords(lon, lat, "EPSG:4326") - - regridder = Regridder(source_grid, target_grid, method="nearest_s2d", parallel=True) - - data = xr.DataArray( - np.random.rand(source_grid.sizes["lat"], source_grid.sizes["lon"]), - coords={"lat": source_grid.lat, "lon": source_grid.lon}, - dims=["lat", "lon"], - name="test_data", - ) - - res = regridder(data) - - assert res.shape == (n_pts,) - assert "n_pts" in res.dims - assert "lat" in res.coords - assert "lon" in res.coords - - -def test_regrid_unstructured_to_structured_dask(dask_client): - n_pts = 50 - lon = np.linspace(0, 360, n_pts) - lat = np.linspace(-90, 90, n_pts) - source_grid = create_mesh_from_coords(lon, lat, "EPSG:4326") - source_grid["test_data"] = (["n_pts"], np.random.rand(n_pts)) - - target_grid = create_global_grid(10, 10) - - regridder = Regridder(source_grid, target_grid, method="nearest_s2d", parallel=True) - - da = source_grid["test_data"] - res = regridder(da) - - assert res.shape == (target_grid.sizes["lat"], target_grid.sizes["lon"]) - assert "lat" in res.dims - assert "lon" in res.dims - - -def test_regrid_unstructured_to_unstructured_dask(dask_client): - n_pts_src = 50 - lon_src = np.linspace(0, 360, n_pts_src) - lat_src = np.linspace(-90, 90, n_pts_src) - source_grid = create_mesh_from_coords(lon_src, lat_src, "EPSG:4326") - source_grid["test_data"] = (["n_pts"], np.random.rand(n_pts_src)) - - n_pts_dst = 30 - lon_dst = np.linspace(0, 360, n_pts_dst) - lat_dst = np.linspace(-90, 90, n_pts_dst) - target_grid = create_mesh_from_coords(lon_dst, lat_dst, "EPSG:4326") - - regridder = Regridder(source_grid, target_grid, method="nearest_s2d", parallel=True) - - da = source_grid["test_data"] - res = regridder(da) - - assert res.shape == (n_pts_dst,) - assert "n_pts" in res.dims diff --git a/tests/test_unstructured_dask_advanced.py b/tests/test_unstructured_dask_advanced.py deleted file mode 100644 index f1c8de6..0000000 --- a/tests/test_unstructured_dask_advanced.py +++ /dev/null @@ -1,164 +0,0 @@ -import pytest -import xarray as xr -import numpy as np -import dask.distributed -from xregrid import Regridder, create_global_grid, create_mesh_from_coords - -# Check for real ESMF -try: - import esmpy - - if hasattr(esmpy, "_is_mock") or "unittest.mock" in str(type(esmpy)): - raise ImportError - HAS_REAL_ESMF = True -except ImportError: - HAS_REAL_ESMF = False - - -@pytest.fixture(scope="module") -def dask_client(): - # esmpy is not thread-safe, so we must use processes=True when using real ESMF - cluster = dask.distributed.LocalCluster( - n_workers=2, threads_per_worker=1, processes=HAS_REAL_ESMF - ) - client = dask.distributed.Client(cluster) - yield client - client.close() - cluster.close() - - -def test_unstructured_with_mask_dask(dask_client): - n_pts = 20 - lon = np.linspace(0, 360, n_pts) - lat = np.linspace(-90, 90, n_pts) - source_grid = create_mesh_from_coords(lon, lat, crs="EPSG:4326") - source_grid["mask"] = (["n_pts"], np.ones(n_pts, dtype=int)) - source_grid["mask"].values[0] = 0 - - target_grid = create_global_grid(10, 10) - - regridder = Regridder( - source_grid, target_grid, method="nearest_s2d", mask_var="mask", parallel=True - ) - - data = xr.DataArray( - np.random.rand(n_pts), - coords={ - "n_pts": source_grid.n_pts, - "lat": source_grid.lat, - "lon": source_grid.lon, - }, - dims=["n_pts"], - name="test_data", - ).chunk({"n_pts": 5}) - - res = regridder(data) - assert res.shape == (18, 36) - val = res.compute() - assert not np.isnan(val).all() - - -def test_mpas_like_detection_dask(dask_client): - nCells = 50 - ds = xr.Dataset( - coords={ - "latCell": (["nCells"], np.linspace(-90, 90, nCells)), - "lonCell": (["nCells"], np.linspace(0, 360, nCells)), - } - ) - ds["test_var"] = (["nCells"], np.random.rand(nCells)) - target_grid = create_global_grid(30, 60) - regridder = Regridder(ds, target_grid, method="nearest_s2d", parallel=True) - res = regridder(ds["test_var"]) - assert res.shape == (target_grid.sizes["lat"], target_grid.sizes["lon"]) - - -def test_unstructured_radians_dask(dask_client): - n_pts = 20 - lon = np.linspace(0, 2 * np.pi, n_pts) - lat = np.linspace(-np.pi / 2, np.pi / 2, n_pts) - source_grid = xr.Dataset( - coords={ - "lat": (["n_pts"], lat, {"units": "radians"}), - "lon": (["n_pts"], lon, {"units": "rad"}), - } - ) - source_grid["test_var"] = (["n_pts"], np.random.rand(n_pts)) - target_grid = create_global_grid(10, 10) - regridder = Regridder(source_grid, target_grid, method="nearest_s2d", parallel=True) - res = regridder(source_grid["test_var"]) - assert res.shape == (target_grid.sizes["lat"], target_grid.sizes["lon"]) - - -def test_structured_to_unstructured_mask_dask(dask_client): - source_grid = create_global_grid(10, 10) - source_grid["mask"] = (["lat", "lon"], np.ones((18, 36), dtype=int)) - n_pts = 30 - target_grid = create_mesh_from_coords( - np.linspace(0, 360, n_pts), np.linspace(-90, 90, n_pts), crs="EPSG:4326" - ) - regridder = Regridder( - source_grid, target_grid, method="nearest_s2d", mask_var="mask", parallel=True - ) - data = xr.DataArray( - np.random.rand(18, 36), - coords={"lat": source_grid.lat, "lon": source_grid.lon}, - dims=["lat", "lon"], - ).chunk({"lat": 9}) - res = regridder(data) - assert res.shape == (n_pts,) - val = res.compute() - assert val.shape == (n_pts,) - - -def test_mpas_conservative_regrid_dask(dask_client): - if HAS_REAL_ESMF: - pytest.skip("MPAS conservative regridding requires valid mesh.") - nCells, nVertices = 20, 40 - ds = xr.Dataset( - coords={ - "latCell": (["nCells"], np.linspace(-90, 90, nCells)), - "lonCell": (["nCells"], np.linspace(0, 360, nCells)), - "latVertex": (["nVertices"], np.linspace(-90, 90, nVertices)), - "lonVertex": (["nVertices"], np.linspace(0, 360, nVertices)), - "verticesOnCell": ( - ["nCells", "maxNodes"], - np.random.randint(1, nVertices + 1, (nCells, 6)), - ), - "nEdgesOnCell": (["nCells"], np.full(nCells, 6)), - } - ) - ds["test_var"] = (["nCells"], np.random.rand(nCells)) - target_grid = create_global_grid(10, 10) - regridder = Regridder(ds, target_grid, method="conservative", parallel=True) - res = regridder(ds["test_var"]) - assert res.shape == (target_grid.sizes["lat"], target_grid.sizes["lon"]) - val = res.compute() - assert not np.isnan(val).all() - - -def test_ugrid_conservative_regrid_dask(dask_client): - if HAS_REAL_ESMF: - pytest.skip("UGRID conservative regridding requires valid mesh.") - nFaces, nNodes = 20, 40 - ds = xr.Dataset( - coords={ - "lat_node": (["nNodes"], np.linspace(-90, 90, nNodes)), - "lon_node": (["nNodes"], np.linspace(0, 360, nNodes)), - "lat": (["nFaces"], np.linspace(-90, 90, nFaces)), - "lon": (["nFaces"], np.linspace(0, 360, nFaces)), - "face_node_connectivity": ( - ["nFaces", "nMaxNodes"], - np.random.randint(0, nNodes, (nFaces, 4)), - ), - } - ) - ds["face_node_connectivity"].attrs["cf_role"] = "face_node_connectivity" - ds["face_node_connectivity"].attrs["start_index"] = 0 - ds["test_var"] = (["nFaces"], np.random.rand(nFaces)) - target_grid = create_global_grid(10, 10) - regridder = Regridder(ds, target_grid, method="conservative", parallel=True) - res = regridder(ds["test_var"]) - assert res.shape == (target_grid.sizes["lat"], target_grid.sizes["lon"]) - val = res.compute() - assert not np.isnan(val).all() diff --git a/tests/test_utils.py b/tests/test_utils.py index 91d6daf..ddf1c59 100644 --- a/tests/test_utils.py +++ b/tests/test_utils.py @@ -1,13 +1,1341 @@ +from __future__ import annotations + +# Consolidated tests: utils + +import os +from unittest.mock import MagicMock, patch + +import dask.array as da +import numpy as np +import pytest import xarray as xr + from xregrid import ( + Regridder, create_global_grid, create_grid_from_crs, + create_grid_from_ioapi, + create_grid_like, create_mesh_from_coords, create_regional_grid, + get_rdhpcs_cluster, load_esmf_file, + plot_comparison, + spatial_slice, ) -import os -import numpy as np +from xregrid.utils import get_crs_info +from xregrid.grid import _get_mesh_info + +try: + import esmpy + + if hasattr(esmpy, "_is_mock") or "unittest.mock" in str(type(esmpy)): + raise ImportError + HAS_REAL_ESMF = True +except (ImportError, Exception): + HAS_REAL_ESMF = False + + +def test_auto_bounds_conservative_numpy_dask(): + """Verify auto-bounds generation for conservative regridding on both NumPy and Dask.""" + # Create a grid WITHOUT bounds but with standard names + lat = np.linspace(-85, 85, 10) + lon = np.linspace(0, 350, 20) + ds_src = xr.Dataset(coords={"lat": lat, "lon": lon}) + ds_src.lat.attrs["standard_name"] = "latitude" + ds_src.lat.attrs["units"] = "degrees_north" + ds_src.lon.attrs["standard_name"] = "longitude" + ds_src.lon.attrs["units"] = "degrees_east" + + # Target grid with bounds + ds_tgt = create_global_grid(20, 20) + + # 1. Eager path + regridder_eager = Regridder(ds_src, ds_tgt, method="conservative") + da_src_eager = xr.DataArray( + np.random.rand(10, 20), dims=("lat", "lon"), coords=ds_src.coords + ) + res_eager = regridder_eager(da_src_eager) + + # 2. Lazy path + da_src_lazy = da_src_eager.chunk({"lat": 5, "lon": 10}) + res_lazy = regridder_eager(da_src_lazy) + + assert isinstance(res_lazy.data, da.Array) + xr.testing.assert_allclose(res_eager, res_lazy.compute()) + assert "Automatically generated" in res_eager.attrs["history"] + + +def test_plot_comparison_smoke(): + """Smoke test for plot_comparison utility.""" + ds = create_global_grid(30, 30) + da_coords = {c: ds.coords[c] for c in ["lat", "lon"]} + da = xr.DataArray(np.random.rand(6, 12), dims=("lat", "lon"), coords=da_coords) + + import matplotlib.pyplot as plt + + plt.switch_backend("Agg") # Non-interactive + + fig = plot_comparison(da, da) + assert fig is not None + plt.close(fig) + + +def test_cf_aware_dimension_mapping(): + """Verify that Regridder handles non-standard dimension names via CF-awareness.""" + # 1. Source grid with standard 'lat'/'lon' + src_res = 10.0 + src_grid = create_global_grid(res_lat=src_res, res_lon=src_res) + + # 2. Target grid + tgt_res = 5.0 + tgt_grid = create_global_grid(res_lat=tgt_res, res_lon=tgt_res) + + # 3. Initialize Regridder + regridder = Regridder(src_grid, tgt_grid, method="bilinear") + + # 4. Input DataArray with different names: 'latitude' and 'longitude' + # but marked with proper CF attributes + data = np.random.rand(18, 36) + da = xr.DataArray( + data, + dims=("latitude", "longitude"), + coords={ + "latitude": ( + ["latitude"], + src_grid.lat.values, + {"standard_name": "latitude"}, + ), + "longitude": ( + ["longitude"], + src_grid.lon.values, + {"standard_name": "longitude"}, + ), + }, + name="test_data", + ) + + # 5. Eager Regridding + res_eager = regridder(da) + + assert res_eager.shape == (36, 72) + assert res_eager.name == "test_data" + + # 6. Lazy Regridding (Double-Check Rule) + da_lazy = da.chunk({"latitude": 9, "longitude": 18}) + res_lazy = regridder(da_lazy).compute() + + # 7. Verification + xr.testing.assert_allclose(res_eager, res_lazy) + + # Verify coordinates match target grid + np.testing.assert_allclose(res_eager.lat, tgt_grid.lat) + np.testing.assert_allclose(res_eager.lon, tgt_grid.lon) + + +def test_dataset_cf_awareness(): + """Verify CF-aware regridding for multiple variables in a Dataset.""" + src_grid = create_global_grid(20, 20) + tgt_grid = create_global_grid(10, 10) + + regridder = Regridder(src_grid, tgt_grid) + + # Dataset with mixed naming + ds = xr.Dataset( + data_vars={ + "temp": (("latitude", "longitude"), np.random.rand(9, 18)), + "scalar": 42.0, + }, + coords={ + "latitude": ( + ["latitude"], + src_grid.lat.values, + {"standard_name": "latitude"}, + ), + "longitude": ( + ["longitude"], + src_grid.lon.values, + {"standard_name": "longitude"}, + ), + "fixed_coord": ("fixed", [1, 2, 3]), + }, + ) + + # Regrid + ds_regridded = regridder(ds) + + assert "temp" in ds_regridded.data_vars + assert ds_regridded.temp.shape == (18, 36) + assert "scalar" in ds_regridded.data_vars + assert ds_regridded.scalar == 42.0 + assert "fixed_coord" in ds_regridded.coords + + +def test_crs_propagation_dataarray(): + """ + Test that CRS metadata is propagated when regridding a DataArray. + Verified with Eager (NumPy) and Lazy (Dask) data. + """ + # 1. Setup Source Grid (Global Lat-Lon) + src_ds = create_global_grid(res_lat=10, res_lon=10) + + # 2. Setup Target Grid (Projected UTM zone 33N) + # UTM zone 33N is approx centered at 15E + target_ds = create_grid_from_crs( + crs="EPSG:32633", extent=(400000, 600000, 5000000, 5200000), res=10000 + ) + + # Create source data + data = np.random.rand(src_ds.sizes["lat"], src_ds.sizes["lon"]) + # Filter coords to only those compatible with (lat, lon) dims + compatible_coords = { + k: v for k, v in src_ds.coords.items() if set(v.dims).issubset({"lat", "lon"}) + } + da_src_numpy = xr.DataArray( + data, coords=compatible_coords, dims=("lat", "lon"), name="test_data" + ) + + da_src_dask = da_src_numpy.chunk({"lat": 5, "lon": 5}) + + # Initialize Regridder + regridder = Regridder(src_ds, target_ds, method="bilinear") + + for da_in in [da_src_numpy, da_src_dask]: + # Perform Regridding + da_out = regridder(da_in) + + # PROOF 1: CRS WKT Attribute Propagation + assert "crs" in da_out.attrs + assert "32633" in da_out.attrs["crs"] + + # PROOF 2: Grid Mapping Variable Propagation + # create_grid_from_crs currently doesn't add a grid_mapping variable by default, + # but it adds 'lat' and 'lon' coordinates. + # Wait, let's check what create_grid_from_crs does. + # It adds 'lat', 'lon' and sets attrs['crs']. + + # PROOF 3: Backend Consistency + if hasattr(da_in.data, "dask"): + assert hasattr(da_out.data, "dask") + else: + assert isinstance(da_out.data, np.ndarray) + + # PROOF 4: Viz Discovery + # get_crs_info should return the correct CRS for the output + crs_detected = get_crs_info(da_out) + assert crs_detected is not None + assert crs_detected.to_epsg() == 32633 + + +def test_crs_propagation_dataset(): + """ + Test that CRS metadata is propagated when regridding a Dataset. + """ + src_ds = create_global_grid(res_lat=10, res_lon=10) + target_ds = create_grid_from_crs("EPSG:3857", (0, 10000, 0, 10000), 1000) + + data = np.random.rand(src_ds.sizes["lat"], src_ds.sizes["lon"]) + src_ds["var1"] = (("lat", "lon"), data) + src_ds.attrs["history"] = "original history" + + regridder = Regridder(src_ds, target_ds, method="bilinear") + ds_out = regridder(src_ds) + + # Global attribute propagation + assert "crs" in ds_out.attrs + assert "3857" in ds_out.attrs["crs"] + + # Variable attribute propagation + assert "crs" in ds_out["var1"].attrs + assert "3857" in ds_out["var1"].attrs["crs"] + + # History update + assert "Regridded" in ds_out.attrs["history"] + + +def test_create_grid_from_ioapi_lcc(): + """Verify IOAPI grid generation for LCC projection (Eager and Lazy).""" + metadata = { + "GDTYP": 2, + "P_ALP": 30.0, + "P_BET": 60.0, + "P_GAM": -97.0, + "XCENT": -97.0, + "YCENT": 40.0, + "XORIG": -1000.0, + "YORIG": -1000.0, + "XCELL": 500.0, + "YCELL": 500.0, + "NCOLS": 4, + "NROWS": 4, + } + + # 1. Eager test + ds_eager = create_grid_from_ioapi(metadata) + + assert "x" in ds_eager.coords + assert "y" in ds_eager.coords + assert "lat" in ds_eager.coords + assert "lon" in ds_eager.coords + assert "x_b" in ds_eager.coords + assert "y_b" in ds_eager.coords + assert ds_eager.sizes["x"] == 4 + assert ds_eager.sizes["y"] == 4 + assert ds_eager.attrs["ioapi_GDTYP"] == 2 + + # Check 1D bounds values + assert ds_eager.x_b.shape == (4, 2) + assert np.allclose(ds_eager.x_b[0].values, [-1000.0, -500.0]) + + # 2. Lazy test + ds_lazy = create_grid_from_ioapi(metadata, chunks={"x": 2, "y": 2}) + # Verify laziness + assert hasattr(ds_lazy.lat.data, "dask") + assert hasattr(ds_lazy.x_b.data, "dask") + + ds_lazy_comp = ds_lazy.compute() + xr.testing.assert_allclose(ds_eager, ds_lazy_comp) + + +def test_create_grid_from_ioapi_all_gdtyp(): + """Verify that all supported IOAPI GDTYP values can generate a grid.""" + base_metadata = { + "P_ALP": 30.0, + "P_BET": 60.0, + "P_GAM": -97.0, + "XCENT": -97.0, + "YCENT": 40.0, + "XORIG": -1000.0, + "YORIG": -1000.0, + "XCELL": 500.0, + "YCELL": 500.0, + "NCOLS": 2, + "NROWS": 2, + } + + # GDTYP 1-10 + for gdtyp in range(1, 11): + metadata = base_metadata.copy() + metadata["GDTYP"] = gdtyp + + # Some specific adjustments to avoid proj errors if needed + if gdtyp == 5: # UTM + metadata["P_ALP"] = 17 # Zone 17 + + ds = create_grid_from_ioapi(metadata) + assert "lat" in ds.coords + assert "lon" in ds.coords + assert ds.attrs["ioapi_GDTYP"] == gdtyp + + +def test_create_grid_from_ioapi_latlon(): + """Verify IOAPI grid generation for Lat-Lon.""" + metadata = { + "GDTYP": 1, + "P_ALP": 0.0, + "P_BET": 0.0, + "P_GAM": 0.0, + "XCENT": 0.0, + "YCENT": 0.0, + "XORIG": -10.0, + "YORIG": 40.0, + "XCELL": 1.0, + "YCELL": 1.0, + "NCOLS": 10, + "NROWS": 10, + } + ds = create_grid_from_ioapi(metadata) + assert ds.sizes["x"] == 10 + assert ds.sizes["y"] == 10 + # For Lat-Lon, pyproj transform from EPSG:4326 to EPSG:4326 should be identity + # but create_grid_from_crs might return lat/lon that are slightly different due to transform + assert ds.lat.min() >= 40.0 + + +try: + import dask.array as da +except ImportError: + da = None + + +def test_create_mesh_from_coords_aero(): + """ + Double-Check Test for create_mesh_from_coords. + Verifies Eager (NumPy) and Lazy (Dask) backends yield identical results + and maintain scientific provenance. + """ + # 1. Setup sample coordinates (Lambert Conformal-ish) + x = np.linspace(-1000, 1000, 10) + y = np.linspace(-1000, 1000, 10) + crs = "+proj=lcc +lat_1=33 +lat_2=45 +lat_0=40 +lon_0=-97 +x_0=0 +y_0=0 +datum=WGS84 +units=m +no_defs" + + # 2. Eager execution + ds_eager = create_mesh_from_coords(x, y, crs) + + # Assertions for Eager + assert isinstance(ds_eager, xr.Dataset) + assert "lat" in ds_eager + assert "lon" in ds_eager + assert "x" in ds_eager + assert "y" in ds_eager + assert ds_eager.attrs["grid_mapping"] == "spatial_ref" + assert "spatial_ref" in ds_eager + assert "Eager" in ds_eager.attrs["history"] + assert "Extent:" in ds_eager.attrs["history"] + + # Check that it's actually NumPy-backed + assert not hasattr(ds_eager.lat.data, "dask") + + # 3. Lazy execution + if da is None: + pytest.skip("Dask not installed, skipping lazy check.") + + x_lazy = da.from_array(x, chunks=5) + y_lazy = da.from_array(y, chunks=5) + + ds_lazy = create_mesh_from_coords(x_lazy, y_lazy, crs) + + # Assertions for Lazy + assert "Lazy" in ds_lazy.attrs["history"] + # Lazy path should NOT have extent in history to avoid compute() + assert "Extent:" not in ds_lazy.attrs["history"] + assert hasattr(ds_lazy.lat.data, "dask") + + # 4. Numerical Verification (The "Double-Check") + xr.testing.assert_allclose(ds_eager, ds_lazy.compute()) + + # 5. Verify Metadata propagation + assert ds_eager.lat.attrs["units"] == "degrees_north" + assert ds_eager.x.attrs["standard_name"] == "projection_x_coordinate" + assert ds_eager.x.attrs["grid_mapping"] == "spatial_ref" + + +def test_create_mesh_from_coords_regression_fix(): + """ + Verify the fix for the dimension mismatch regression and conditional metadata. + """ + # 1. Test DataArray inputs with different dimension names + x_da = xr.DataArray(np.linspace(0, 10, 5), dims=["lon"], name="my_lon") + y_da = xr.DataArray(np.linspace(0, 10, 5), dims=["lat"], name="my_lat") + crs_proj = "+proj=lcc +lat_1=33 +lat_2=45 +lat_0=40 +lon_0=-97 +x_0=0 +y_0=0 +datum=WGS84 +units=m +no_defs" + + ds = create_mesh_from_coords(x_da, y_da, crs_proj) + + # Should have 5 points, not 25 (if it had broadcasted incorrectly) + assert ds.sizes["n_pts"] == 5 + assert ds.x.attrs["standard_name"] == "projection_x_coordinate" + + # 2. Test geographic CRS metadata + crs_geo = "EPSG:4326" + ds_geo = create_mesh_from_coords(x_da, y_da, crs_geo) + + assert ds_geo.x.attrs["standard_name"] == "longitude" + assert ds_geo.x.attrs["units"] == "degrees_east" + assert ds_geo.y.attrs["standard_name"] == "latitude" + assert ds_geo.y.attrs["units"] == "degrees_north" + + +def test_spatial_slice_rectilinear() -> None: + """ + Test spatial_slice with a rectilinear grid (Eager and Lazy). + + Verifies that spatial_slice correctly subsets a rectilinear grid + using coordinate indexes and handles both NumPy and Dask backends. + """ + # 1. Eager + ds = create_global_grid(res_lat=1.0, res_lon=1.0) + # Extent: (min_x, max_x, min_y, max_y) + extent = (10.5, 20.5, 30.5, 40.5) + ds_sliced = spatial_slice(ds, extent) + + assert ds_sliced.lat.min() >= 30.5 + assert ds_sliced.lat.max() <= 40.5 + assert ds_sliced.lon.min() >= 10.5 + assert ds_sliced.lon.max() <= 20.5 + assert "history" in ds_sliced.attrs + assert "Spatially sliced" in ds_sliced.attrs["history"] + + # 2. Lazy (Dask) + ds_lazy = create_global_grid( + res_lat=1.0, res_lon=1.0, chunks={"lat": 10, "lon": 10} + ) + ds_sliced_lazy = spatial_slice(ds_lazy, extent) + + # In xarray, dimension coordinates are often eager (NumPy) due to indexing. + # Check lat_b instead, which should remain lazy. + assert hasattr(ds_sliced_lazy.lat_b.data, "dask") + xr.testing.assert_allclose(ds_sliced, ds_sliced_lazy.compute()) + + +def test_spatial_slice_unstructured() -> None: + """ + Test spatial_slice with an unstructured grid (Eager and Lazy). + + Verifies that spatial_slice correctly subsets an unstructured grid + using boolean masking and maintains laziness for Dask-backed data. + """ + # Create points: one in the box, one outside + lons = np.array([15.0, 50.0]) + lats = np.array([35.0, 60.0]) + ds = create_mesh_from_coords(lons, lats, crs="EPSG:4326") + + extent = (10.0, 20.0, 30.0, 40.0) + + # Eager path: drop=True is supported + ds_sliced = spatial_slice(ds, extent) + + assert ds_sliced.sizes["n_pts"] == 1 + assert ds_sliced.lon.values[0] == 15.0 + assert ds_sliced.lat.values[0] == 35.0 + + # Lazy path: drop=False is used to preserve laziness + ds_lazy = create_mesh_from_coords(lons, lats, crs="EPSG:4326", chunks=1) + ds_lazy["data"] = (["n_pts"], ds_lazy.lat.data, {"units": "K"}) + ds_sliced_lazy = spatial_slice(ds_lazy, extent) + + # Verify laziness + assert hasattr(ds_sliced_lazy.data.data, "dask") + + # Verify results (after compute and dropna since drop=False was used) + ds_res = ds_sliced_lazy.compute().dropna("n_pts") + assert ds_res.sizes["n_pts"] == 1 + assert ds_res.lon.values[0] == 15.0 + + +def test_spatial_slice_wrapping() -> None: + """ + Test spatial_slice with longitude wrapping across different grid types. + + Verifies that spatial_slice correctly handles regions crossing the + meridian/dateline for both rectilinear and unstructured grids. + """ + ds = create_global_grid(res_lat=1.0, res_lon=1.0) + + # Slice crossing the 0/360 boundary + # Extent: (min_x, max_x, min_y, max_y) + extent = (-10.5, 10.5, -10.5, 10.5) + ds_sliced = spatial_slice(ds, extent) + + # Lon in ds is 0-360. -10.5 should map to 349.5 + assert ds_sliced.lon.min() >= 0 + assert ds_sliced.lon.max() <= 360 + + # Check that we have both the 0-10.5 and 349.5-360 parts + assert (ds_sliced.lon <= 10.5).any() + assert (ds_sliced.lon >= 349.5).any() + + # Verify unstructured wrapping + lons = np.array([5.0, 355.0, 180.0]) + lats = np.array([0.0, 0.0, 0.0]) + ds_unstructured = create_mesh_from_coords(lons, lats, crs="EPSG:4326") + ds_un_sliced = spatial_slice(ds_unstructured, extent) + + # Since it's eager, drop=True was used + assert ds_un_sliced.sizes["n_pts"] == 2 + assert set(ds_un_sliced.lon.values) == {5.0, 355.0} + + +def test_create_global_grid_lazy(): + """ + Aero Protocol: Double-Check Test for create_global_grid. + Verifies that values are identical between NumPy and Dask backends. + """ + res_lat, res_lon = 10, 20 + + # Eager (NumPy) + ds_eager = create_global_grid(res_lat=res_lat, res_lon=res_lon, chunks=None) + assert not ds_eager.chunks + + # Lazy (Dask) + ds_lazy = create_global_grid( + res_lat=res_lat, res_lon=res_lon, chunks={"lat": 9, "lon": 9} + ) + assert ds_lazy.chunks + + # Assert values are identical + xr.testing.assert_allclose(ds_eager, ds_lazy.compute()) + + # Verify internal backend (lat_b is non-index so it should be chunked) + assert hasattr(ds_lazy.lat_b.data, "dask") + + +def test_create_regional_grid_lazy(): + """ + Aero Protocol: Double-Check Test for create_regional_grid. + """ + lat_range = (-45, 45) + lon_range = (0, 90) + res_lat, res_lon = 5, 5 + + # Eager (NumPy) + ds_eager = create_regional_grid(lat_range, lon_range, res_lat, res_lon, chunks=None) + + # Lazy (Dask) + ds_lazy = create_regional_grid(lat_range, lon_range, res_lat, res_lon, chunks=5) + assert ds_lazy.chunks + + # Assert values are identical + xr.testing.assert_allclose(ds_eager, ds_lazy.compute()) + + # Verify internal backend + assert hasattr(ds_lazy.lat_b.data, "dask") + + +def test_create_grid_from_crs_lazy(): + """ + Aero Protocol: Double-Check Test for create_grid_from_crs. + """ + # Test with EPSG:32633 (UTM zone 33N) + extent = (400000, 500000, 5000000, 5100000) + res = 10000 # 10km + + # Eager (NumPy) + ds_eager = create_grid_from_crs("EPSG:32633", extent, res, chunks=None) + + # Lazy (Dask) + ds_lazy = create_grid_from_crs("EPSG:32633", extent, res, chunks={"x": 5, "y": 5}) + assert ds_lazy.chunks + + # Assert values are identical + xr.testing.assert_allclose(ds_eager, ds_lazy.compute()) + + # Verify internal backend (lat/lon are non-index here) + assert hasattr(ds_lazy.lat.data, "dask") + + +def test_create_mesh_from_coords_lazy(): + """ + Aero Protocol: Double-Check Test for create_mesh_from_coords. + """ + x = np.array([400000, 450000, 500000]) + y = np.array([5000000, 5050000, 5100000]) + + # Eager (NumPy) + ds_eager = create_mesh_from_coords(x, y, "EPSG:32633", chunks=None) + + # Lazy (Dask) + ds_lazy = create_mesh_from_coords(x, y, "EPSG:32633", chunks={"n_pts": 2}) + assert ds_lazy.chunks + + # Assert values are identical + xr.testing.assert_allclose(ds_eager, ds_lazy.compute()) + + # Verify internal backend + assert hasattr(ds_lazy.lat.data, "dask") + + +def test_create_grid_like_latlon(): + """ + Aero Protocol: Double-Check Test for create_grid_like (Lat-Lon). + Verifies identity between NumPy and Dask backends and preservation of laziness. + """ + # Create a source grid + ds_src = create_regional_grid( + lat_range=(10, 20), + lon_range=(100, 110), + res_lat=1.0, + res_lon=1.0, + add_bounds=True, + ) + + res_new = 0.5 + + # Eager (NumPy) + ds_eager = create_grid_like(ds_src, res_new, chunks=None) + assert not ds_eager.chunks + assert ds_eager.lat.size == 20 + assert ds_eager.lon.size == 20 + + # Lazy (Dask) + ds_src_lazy = ds_src.chunk({"lat": 5, "lon": 5}) + ds_lazy = create_grid_like(ds_src_lazy, res_new, chunks=5) + + # Assert values are identical + xr.testing.assert_allclose(ds_eager, ds_lazy.compute()) + + # Verify laziness: lat_b should be a dask array + assert hasattr(ds_lazy.lat_b.data, "dask") + + +def test_create_grid_like_projected(): + """ + Aero Protocol: Double-Check Test for create_grid_like (Projected). + """ + # UTM zone 33N + crs = "EPSG:32633" + extent = (400000, 500000, 5000000, 5100000) + res_orig = 10000 + + ds_src = create_grid_from_crs(crs, extent, res_orig, add_bounds=True) + + res_new = 5000 + + # Eager + ds_eager = create_grid_like(ds_src, res_new, chunks=None) + + # Lazy + ds_src_lazy = ds_src.chunk({"x": 5, "y": 5}) + ds_lazy = create_grid_like(ds_src_lazy, res_new, chunks=5) + + # Assert values + xr.testing.assert_allclose(ds_eager, ds_lazy.compute()) + + # Verify laziness (lat/lon are non-index in projected grid) + assert hasattr(ds_lazy.lat.data, "dask") + assert ds_lazy.attrs["crs"] == ds_src.attrs["crs"] + + +def test_rectilinear_hygiene(): + """ + Verify that _create_rectilinear_grid produces high-hygiene metadata. + """ + ds = create_regional_grid((0, 10), (0, 10), 1, 1) + + assert ds.attrs["crs"] == "EPSG:4326" + + # Test custom CRS + # create_regional_grid currently doesn't expose crs, let's test _create_rectilinear_grid directly + from xregrid.utils import _create_rectilinear_grid + + ds_nad83 = _create_rectilinear_grid((0, 10), (0, 10), 1, 1, crs="EPSG:4269") + assert ds_nad83.attrs["crs"] == "EPSG:4269" + + assert ds.lat.attrs["standard_name"] == "latitude" + assert ds.lon.attrs["standard_name"] == "longitude" + assert ds.lat_b.attrs["standard_name"] == "latitude_bounds" + assert ds.lon_b.attrs["standard_name"] == "longitude_bounds" + assert "history" in ds.attrs + + +def test_cf_coords_detection(): + # Create dataset with non-standard coordinate names but with CF attributes + def create_ds(lazy=False): + np.random.seed(42) + data = np.random.rand(10, 20) + if lazy: + data = da.from_array(data, chunks=(5, 10)) + + ds = xr.Dataset( + {"data": (("lat_dim", "lon_dim"), data)}, + coords={ + "latitude": ( + ("lat_dim",), + np.linspace(-90, 90, 10), + {"units": "degrees_north", "standard_name": "latitude"}, + ), + "longitude": ( + ("lon_dim",), + np.linspace(-180, 180, 20), + {"units": "degrees_east", "standard_name": "longitude"}, + ), + }, + ) + return ds + + # Target grid also needs bounds for conservative regridding + lat_edges_tgt = np.linspace(-90, 90, 16) + lon_edges_tgt = np.linspace(-180, 180, 26) + ds_tgt = xr.Dataset( + coords={ + "lat": (("lat",), np.linspace(-90, 90, 15), {"units": "degrees_north"}), + "lon": (("lon",), np.linspace(-180, 180, 25), {"units": "degrees_east"}), + "lat_b": (("lat_b",), lat_edges_tgt, {"units": "degrees_north"}), + "lon_b": (("lon_b",), lon_edges_tgt, {"units": "degrees_east"}), + }, + ) + + # Test Eager + ds_src_eager = create_ds(lazy=False) + regridder_eager = Regridder(ds_src_eager, ds_tgt) + out_eager = regridder_eager(ds_src_eager["data"]) + assert out_eager.shape == (15, 25) + assert not out_eager.chunks + + # Test Lazy + ds_src_lazy = create_ds(lazy=True) + regridder_lazy = Regridder(ds_src_lazy, ds_tgt) + out_lazy = regridder_lazy(ds_src_lazy["data"]) + assert out_lazy.shape == (15, 25) + assert out_lazy.chunks + + # Verify results are identical (within float precision) + np.testing.assert_allclose(out_eager.values, out_lazy.compute().values) + + +def test_cf_bounds_detection(): + # Create dataset with non-standard bound names but with CF attributes + ds_src = xr.Dataset( + {"data": (("lat", "lon"), np.random.rand(10, 20))}, + coords={ + "lat": ( + ("lat",), + np.linspace(-90, 90, 10), + {"units": "degrees_north", "bounds": "lat_bounds"}, + ), + "lon": ( + ("lon",), + np.linspace(-180, 180, 20), + {"units": "degrees_east", "bounds": "lon_bounds"}, + ), + "lat_bounds": (("lat", "nv"), np.random.rand(10, 2)), # Placeholder bounds + "lon_bounds": (("lon", "nv"), np.random.rand(20, 2)), # Placeholder bounds + }, + ) + + # We need to make the bounds contiguous for our converter to work correctly in this test + lat_edges = np.linspace(-90, 90, 11) + lat_bounds = np.stack([lat_edges[:-1], lat_edges[1:]], axis=1) + lon_edges = np.linspace(-180, 180, 21) + lon_bounds = np.stack([lon_edges[:-1], lon_edges[1:]], axis=1) + + ds_src.coords["lat_bounds"] = (("lat", "nv"), lat_bounds) + ds_src.coords["lon_bounds"] = (("lon", "nv"), lon_bounds) + + # Target grid also needs bounds for conservative regridding + lat_edges_tgt = np.linspace(-90, 90, 16) + lon_edges_tgt = np.linspace(-180, 180, 26) + ds_tgt = xr.Dataset( + coords={ + "lat": (("lat",), np.linspace(-90, 90, 15), {"units": "degrees_north"}), + "lon": (("lon",), np.linspace(-180, 180, 25), {"units": "degrees_east"}), + "lat_b": (("lat_b",), lat_edges_tgt, {"units": "degrees_north"}), + "lon_b": (("lon_b",), lon_edges_tgt, {"units": "degrees_east"}), + }, + ) + + regridder = Regridder(ds_src, ds_tgt, method="conservative") + # If it reached here without error, it found the bounds and ESMPy initialized + out = regridder(ds_src["data"]) + assert out.shape == (15, 25) + + +def test_regridder_time_dimension_detection(): + # Setup source and target grids with time + lats = np.linspace(-90, 90, 10) + lons = np.linspace(0, 360, 20) + times = [np.datetime64("2020-01-01")] + + src_ds = xr.Dataset( + coords={ + "time": (["time"], times, {"standard_name": "time"}), + "lat": ( + ["time", "lat"], + np.broadcast_to(lats, (1, 10)), + {"units": "degrees_north", "standard_name": "latitude"}, + ), + "lon": ( + ["lon"], + lons, + {"units": "degrees_east", "standard_name": "longitude"}, + ), + } + ) + + tgt_ds = xr.Dataset( + coords={ + "lat": ( + ["lat"], + np.linspace(-90, 90, 5), + {"units": "degrees_north", "standard_name": "latitude"}, + ), + "lon": ( + ["lon"], + np.linspace(0, 360, 10), + {"units": "degrees_east", "standard_name": "longitude"}, + ), + } + ) + + # This should now work without failing during weight generation + regridder = Regridder(src_ds, tgt_ds, method="bilinear") + + # Create data with time and vertical dimensions + levs = np.arange(5) + data = np.random.rand(len(times), len(levs), len(lats), len(lons)) + da = xr.DataArray( + data, + coords={ + "time": (["time"], times), + "lev": (["lev"], levs), + "lat": (["time", "lat"], np.broadcast_to(lats, (1, 10))), + "lon": (["lon"], lons), + }, + dims=("time", "lev", "lat", "lon"), + name="temp", + ) + + # Regrid DataArray + res_da = regridder(da) + + # Check that time and lev are preserved + assert "time" in res_da.dims + assert "lev" in res_da.dims + assert res_da.shape == (1, 5, 5, 10) + + # Regrid Dataset + ds = xr.Dataset({"temp": da, "time_var": (["time"], times)}) + res_ds = regridder(ds) + + assert "time" in res_ds.dims + assert "temp" in res_ds.data_vars + assert "time_var" in res_ds.data_vars + assert res_ds["temp"].shape == (1, 5, 5, 10) + assert res_ds["time_var"].dims == ("time",) + + +def test_regridder_dtype_time_fallback(): + # Setup with time-like dtype but non-standard name + lats = np.linspace(-90, 90, 10) + lons = np.linspace(0, 360, 20) + times = [np.datetime64("2020-01-01")] + + src_ds = xr.Dataset( + coords={ + "mytime": (["mytime"], times), # Non-standard name, no CF attributes + "lat": ( + ["mytime", "lat"], + np.broadcast_to(lats, (1, 10)), + {"units": "degrees_north", "standard_name": "latitude"}, + ), + "lon": ( + ["lon"], + lons, + {"units": "degrees_east", "standard_name": "longitude"}, + ), + } + ) + + tgt_ds = xr.Dataset( + coords={ + "lat": ( + ["lat"], + np.linspace(-90, 90, 5), + {"units": "degrees_north", "standard_name": "latitude"}, + ), + "lon": ( + ["lon"], + np.linspace(0, 360, 10), + {"units": "degrees_east", "standard_name": "longitude"}, + ), + } + ) + + regridder = Regridder(src_ds, tgt_ds) + + # Verify mytime was detected as non-spatial + assert "mytime" not in regridder._dims_source + + # Test DataArray regridding with this non-standard time dim + da = xr.DataArray( + np.random.rand(1, 10, 20), coords=src_ds.coords, dims=("mytime", "lat", "lon") + ) + + res = regridder(da) + assert "mytime" in res.dims + assert res.shape == (1, 5, 10) + + +def test_non_regriddable_object(): + # Test passing something that shouldn't be regridded + lats = np.linspace(-90, 90, 10) + lons = np.linspace(0, 360, 20) + + src_ds = xr.Dataset( + coords={ + "lat": ( + ["lat"], + lats, + {"units": "degrees_north", "standard_name": "latitude"}, + ), + "lon": ( + ["lon"], + lons, + {"units": "degrees_east", "standard_name": "longitude"}, + ), + } + ) + tgt_ds = xr.Dataset( + coords={ + "lat": ( + ["lat"], + np.linspace(-90, 90, 5), + {"units": "degrees_north", "standard_name": "latitude"}, + ), + "lon": ( + ["lon"], + np.linspace(0, 360, 10), + {"units": "degrees_east", "standard_name": "longitude"}, + ), + } + ) + + regridder = Regridder(src_ds, tgt_ds) + + # A DataArray that only has one dimension (time) + time_da = xr.DataArray([1, 2, 3], dims="time", name="time_var") + + # Should return unchanged + res = regridder(time_da) + xr.testing.assert_identical(res, time_da) + + +def test_regridder_vertical_dimension_detection(): + # Setup source with vertical dimension in lats + lats = np.linspace(-90, 90, 10) + lons = np.linspace(0, 360, 20) + levs = np.arange(3) + + src_ds = xr.Dataset( + coords={ + "lev": (["lev"], levs, {"standard_name": "altitude"}), + "lat": ( + ["lev", "lat"], + np.broadcast_to(lats, (3, 10)), + {"units": "degrees_north", "standard_name": "latitude"}, + ), + "lon": ( + ["lon"], + lons, + {"units": "degrees_east", "standard_name": "longitude"}, + ), + } + ) + + tgt_ds = xr.Dataset( + coords={ + "lat": ( + ["lat"], + np.linspace(-90, 90, 5), + {"units": "degrees_north", "standard_name": "latitude"}, + ), + "lon": ( + ["lon"], + np.linspace(0, 360, 10), + {"units": "degrees_east", "standard_name": "longitude"}, + ), + } + ) + + regridder = Regridder(src_ds, tgt_ds) + assert "lev" not in regridder._dims_source + + da = xr.DataArray( + np.random.rand(3, 10, 20), coords=src_ds.coords, dims=("lev", "lat", "lon") + ) + + res = regridder(da) + assert "lev" in res.dims + assert res.shape == (3, 5, 10) + + +def test_regridder_ugrid_with_time(): + # Setup mocked uxarray object with time dimension + from unittest.mock import MagicMock + + class UxDatasetMock(xr.Dataset): + def __init__(self, ds, uxgrid): + super().__init__(ds.data_vars, coords=ds.coords, attrs=ds.attrs) + self.uxgrid = uxgrid + + n_face = 10 + n_node = 12 + times = [np.datetime64("2020-01-01")] + + mock_uxgrid = MagicMock() + mock_uxgrid.node_lat = xr.DataArray(np.linspace(-90, 90, n_node), dims=["n_node"]) + mock_uxgrid.node_lon = xr.DataArray(np.linspace(0, 360, n_node), dims=["n_node"]) + # face coords with time dimension (moving mesh case, though xregrid assumes static) + mock_uxgrid.face_lat = xr.DataArray( + np.broadcast_to(np.linspace(-90, 90, n_face), (1, n_face)), + dims=["time", "n_face"], + coords={"time": times}, + ) + mock_uxgrid.face_lon = xr.DataArray( + np.broadcast_to(np.linspace(0, 360, n_face), (1, n_face)), + dims=["time", "n_face"], + coords={"time": times}, + ) + + # Create connectivity + conn = np.zeros((n_face, 3), dtype=int) + for i in range(n_face): + conn[i] = [i, i + 1, (i + 2) % n_node] + + mock_uxgrid.face_node_connectivity = xr.DataArray( + conn, dims=["n_face", "n_max_face_nodes"] + ) + mock_uxgrid.face_node_connectivity.attrs["start_index"] = 0 + mock_uxgrid.face_node_connectivity.attrs["_FillValue"] = -1 + + # Mock UxDataset with time-varying variable + ds_base = xr.Dataset( + {"test_var": (["time", "n_face"], np.random.rand(1, n_face))}, + coords={"time": (["time"], times, {"standard_name": "time"})}, + ) + ds = UxDatasetMock(ds_base, mock_uxgrid) + + target_grid = xr.Dataset( + coords={ + "lat": ( + ["lat"], + np.linspace(-90, 90, 5), + {"units": "degrees_north", "standard_name": "latitude"}, + ), + "lon": ( + ["lon"], + np.linspace(0, 360, 10), + {"units": "degrees_east", "standard_name": "longitude"}, + ), + } + ) + + regridder = Regridder(ds, target_grid, method="nearest_s2d") + + # Verify time was detected as non-spatial + assert "time" not in regridder._dims_source + assert regridder._dims_source == ("n_face",) + + # Regrid DataArray + res = regridder(ds["test_var"]) + + assert "time" in res.dims + assert res.shape == (1, 5, 10) + + +def test_regridder_raw_ugrid_with_time(): + n_face = 10 + n_node = 12 + times = [np.datetime64("2020-01-01")] + + # Create a raw dataset following UGRID convention + conn = np.zeros((n_face, 3), dtype=int) + for i in range(n_face): + conn[i] = [i, (i + 1) % n_node, (i + 2) % n_node] + + ds = xr.Dataset( + data_vars={ + "temp": (["time", "n_face"], np.random.rand(1, n_face)), + "face_node_connectivity": (["n_face", "n_max_face_nodes"], conn), + }, + coords={ + "time": (["time"], times, {"standard_name": "time"}), + "lat_face": ( + ["time", "n_face"], + np.broadcast_to(np.linspace(-90, 90, n_face), (1, n_face)), + {"units": "degrees_north"}, + ), + "lon_face": ( + ["time", "n_face"], + np.broadcast_to(np.linspace(0, 360, n_face), (1, n_face)), + {"units": "degrees_east"}, + ), + "lat_node": ( + ["time", "n_node"], + np.broadcast_to(np.linspace(-90, 90, n_node), (1, n_node)), + {"units": "degrees_north"}, + ), + "lon_node": ( + ["time", "n_node"], + np.broadcast_to(np.linspace(0, 360, n_node), (1, n_node)), + {"units": "degrees_east"}, + ), + }, + ) + + ds.face_node_connectivity.attrs["cf_role"] = "face_node_connectivity" + ds.face_node_connectivity.attrs["start_index"] = 0 + + from xregrid import create_global_grid + + target_grid = create_global_grid(10, 10) + + regridder = Regridder(ds, target_grid, method="nearest_s2d") + + assert "time" not in regridder._dims_source + # Since it is UGRID, it should have detected n_face as the spatial dimension for variables + assert "n_face" in regridder._dims_source + + res = regridder(ds["temp"]) + assert "time" in res.dims + assert res.shape == (1, 18, 36) + + +def test_regridder_user_specific_structure(): + # Mimic user's dataset structure: (time, node) + # node is string coordinate, lat/lon are (node) + n_node = 10 + n_time = 5 + times = np.arange(n_time).astype("datetime64[D]") + nodes = np.array([f"NODE_{i}" for i in range(n_node)], dtype=" 0 - plt.close(fig) - - -if __name__ == "__main__": - pytest.main([__file__]) diff --git a/tests/test_viz_coverage.py b/tests/test_viz_coverage.py deleted file mode 100644 index fe7d5ae..0000000 --- a/tests/test_viz_coverage.py +++ /dev/null @@ -1,102 +0,0 @@ -import pytest -import xarray as xr -import numpy as np -from xregrid.viz import plot_static, plot_comparison, plot_interactive, plot_diagnostics -from xregrid import Regridder, create_global_grid -from unittest.mock import patch, MagicMock - - -def test_plot_static_crs_detection(): - """Verify CRS detection from attributes in plot_static.""" - da = xr.DataArray( - np.random.rand(18, 36), - dims=("lat", "lon"), - coords={"lat": np.linspace(-90, 90, 18), "lon": np.linspace(0, 360, 36)}, - name="test_data", - ) - - # Test with grid_mapping attribute - da.attrs["grid_mapping"] = "crs" - crs_var = xr.DataArray(0, name="crs") - crs_var.attrs["crs_wkt"] = ( - 'GEOGCS["WGS 84",DATUM["WGS_1984",SPHEROID["WGS 84",6378137,298.257223563]],PRIMEM["Greenwich",0],UNIT["degree",0.0174532925199433]]' - ) - ds = da.to_dataset().assign(crs=crs_var) - - # Call plot_static - with patch("matplotlib.pyplot.axes") as mock_axes: - mock_ax = MagicMock() - mock_axes.return_value = mock_ax - plot_static(ds.test_data) - assert mock_axes.called - - -def test_plot_static_fallback_no_cartopy(): - """Verify fallback to standard matplotlib when cartopy is missing.""" - da = xr.DataArray( - np.random.rand(18, 36), - dims=("lat", "lon"), - coords={"lat": np.linspace(-90, 90, 18), "lon": np.linspace(0, 360, 36)}, - ) - - with patch("xregrid.viz.ccrs", None): - with patch("matplotlib.pyplot.gca") as mock_gca: - plot_static(da) - assert mock_gca.called - - -def test_plot_static_robust_slicing(): - """Verify robust slicing of extra dimensions in plot_static.""" - da = xr.DataArray( - np.random.rand(3, 18, 36), - dims=("time", "lat", "lon"), - coords={ - "time": [1, 2, 3], - "lat": np.linspace(-90, 90, 18), - "lon": np.linspace(0, 360, 36), - }, - name="test_data", - ) - - with patch("matplotlib.pyplot.axes"): - with pytest.warns(UserWarning, match="Automatically selecting the first slice"): - plot_static(da) - - -def test_plot_comparison_smoke(): - """Smoke test for plot_comparison.""" - src_da = xr.DataArray(np.random.rand(18, 36), dims=("lat", "lon"), name="src") - tgt_da = xr.DataArray(np.random.rand(36, 72), dims=("lat", "lon"), name="tgt") - - with patch("matplotlib.pyplot.subplots") as mock_subplots: - mock_fig = MagicMock() - mock_axes = [MagicMock() for _ in range(3)] - mock_subplots.return_value = (mock_fig, mock_axes) - - plot_comparison(src_da, tgt_da) - assert mock_subplots.called - - -def test_plot_diagnostics_smoke(): - """Smoke test for plot_diagnostics.""" - src = create_global_grid(10, 10) - tgt = create_global_grid(5, 5) - regridder = Regridder(src, tgt) - - with patch("matplotlib.pyplot.subplots") as mock_subplots: - mock_fig = MagicMock() - mock_axes = [MagicMock() for _ in range(2)] - mock_subplots.return_value = (mock_fig, mock_axes) - - plot_diagnostics(regridder) - assert mock_subplots.called - - -def test_plot_interactive_smoke(): - """Smoke test for plot_interactive.""" - from xregrid.viz import hvplot as has_hvplot - - if not has_hvplot: - pytest.skip("hvplot missing") - da = xr.DataArray(np.random.rand(18, 36), dims=("lat", "lon"), name="test") - plot_interactive(da) diff --git a/tests/test_xregrid.py b/tests/test_xregrid.py deleted file mode 100644 index d1e6ce5..0000000 --- a/tests/test_xregrid.py +++ /dev/null @@ -1,261 +0,0 @@ -import dask.array as da -import numpy as np -import pytest -import xarray as xr -from xregrid import Regridder, create_grid_from_crs, create_global_grid - - -def create_sample_dataset( - nlat=45, - nlon=90, - lat_range=(-90, 90), - lon_range=(0, 360), - dask=False, - chunk_core=True, -): - """Create a sample dataset with synthetic data.""" - lat = np.linspace(lat_range[0], lat_range[1], nlat) - lon = np.linspace(lon_range[0], lon_range[1], nlon) - - lon_grid, lat_grid = np.meshgrid(lon, lat) - data = np.sin(np.radians(lat_grid)) * np.cos(np.radians(lon_grid * 2)) - - if dask: - if chunk_core: - data = da.from_array(data, chunks=(nlat // 2, nlon // 2)) - else: - # Add a non-core dimension to chunk along - data = data[np.newaxis, :, :] # (time, lat, lon) - data = da.from_array(data, chunks=(1, -1, -1)) - ds = xr.Dataset( - { - "temperature": (["time", "lat", "lon"], data), - }, - coords={ - "time": [0], - "lat": (["lat"], lat), - "lon": (["lon"], lon), - }, - ) - return ds - - ds = xr.Dataset( - { - "temperature": (["lat", "lon"], data), - }, - coords={ - "lat": (["lat"], lat), - "lon": (["lon"], lon), - }, - ) - return ds - - -def test_rectilinear_regrid_numpy(): - source_ds = create_sample_dataset(nlat=45, nlon=90) - target_ds = create_sample_dataset(nlat=60, nlon=120) - - regridder = Regridder(source_ds, target_ds, method="bilinear") - regridded = regridder(source_ds["temperature"]) - - assert regridded.shape == (60, 120) - assert not np.isnan(regridded.values).any() - assert regridded.min() >= source_ds.temperature.min() - 0.1 - assert regridded.max() <= source_ds.temperature.max() + 0.1 - - -def test_rectilinear_regrid_dask_non_core_chunked(): - source_ds_np_full = create_sample_dataset(nlat=45, nlon=90, dask=False) - source_ds_da = create_sample_dataset(nlat=45, nlon=90, dask=True, chunk_core=False) - target_ds = create_sample_dataset(nlat=60, nlon=120) - - regridder = Regridder(source_ds_np_full, target_ds, method="bilinear") - - regridded_da = regridder(source_ds_da["temperature"]) - - assert isinstance(regridded_da.data, da.Array) - assert regridded_da.shape == (1, 60, 120) - - result = regridded_da.compute() - assert not np.isnan(result.values).any() - - -def test_rectilinear_regrid_dask_core_chunked(): - source_ds_da = create_sample_dataset(nlat=45, nlon=90, dask=True, chunk_core=True) - target_ds = create_sample_dataset(nlat=60, nlon=120) - regridder = Regridder(source_ds_da, target_ds, method="bilinear") - regridded_da = regridder(source_ds_da["temperature"]) - - assert isinstance(regridded_da.data, da.Array) - result = regridded_da.compute() - assert result.shape == (60, 120) - assert not np.isnan(result.values).any() - - -def test_regrid_timing(benchmark): - source_ds = create_sample_dataset(nlat=180, nlon=360) - target_ds = create_sample_dataset(nlat=360, nlon=720) - - regridder = Regridder(source_ds, target_ds, method="bilinear") - - def do_regrid(): - return regridder(source_ds["temperature"]).values - - benchmark(do_regrid) - - -def test_provenance(): - source_ds = create_sample_dataset(nlat=10, nlon=20) - target_ds = create_sample_dataset(nlat=15, nlon=25) - regridder = Regridder(source_ds, target_ds) - regridded = regridder(source_ds["temperature"]) - - assert "history" in regridded.attrs - assert "Regridder" in regridded.attrs["history"] - assert "bilinear" in regridded.attrs["history"] - - -def test_type_hints(): - # Basic check that the class has expected annotations - # With from __future__ import annotations, they might be strings - ann = Regridder.__init__.__annotations__["method"] - assert ann == "str" or ann is str - - -def test_viz_static_call(): - try: - import matplotlib.pyplot as plt - except ImportError: - pytest.skip("matplotlib not installed") - from xregrid import plot_static - - source_ds = create_sample_dataset(nlat=10, nlon=20) - # This should work without error even if it just calls da.plot - plot_static(source_ds["temperature"]) - plt.close("all") - - -def test_dask_numpy_identity(): - source_ds = create_sample_dataset(nlat=10, nlon=20) - target_ds = create_sample_dataset(nlat=15, nlon=25) - regridder = Regridder(source_ds, target_ds) - - # Eager - da_eager = source_ds["temperature"] - res_eager = regridder(da_eager) - - # Lazy - da_lazy = da_eager.chunk({"lat": 5, "lon": 10}) - res_lazy = regridder(da_lazy).compute() - - xr.testing.assert_allclose(res_eager, res_lazy) - - -def test_attribute_preservation(): - source_ds = create_sample_dataset() - source_ds.temperature.attrs["units"] = "K" - target_ds = create_sample_dataset(nlat=10, nlon=20) - regridder = Regridder(source_ds, target_ds) - out = regridder(source_ds.temperature) - assert out.attrs["units"] == "K" - - -def test_plot_static_custom_ax(): - try: - import matplotlib.pyplot as plt - import cartopy.crs as ccrs - except ImportError: - pytest.skip("matplotlib or cartopy not installed") - from xregrid import plot_static - - da = create_sample_dataset()["temperature"] - fig, ax = plt.subplots(subplot_kw={"projection": ccrs.Robinson()}) - plot_static(da, ax=ax) - plt.close(fig) - - -def test_regridder_repr(): - source_ds = create_sample_dataset(nlat=10, nlon=20) - target_ds = create_sample_dataset(nlat=15, nlon=25) - regridder = Regridder(source_ds, target_ds, method="bilinear", periodic=False) - rep = repr(regridder) - assert "Regridder" in rep - assert "method=bilinear" in rep - assert "periodic=False" in rep - assert "(10, 20)" in rep - assert "(15, 25)" in rep - - -def test_weights_format(): - from scipy.sparse import csr_matrix - - source_ds = create_sample_dataset(nlat=10, nlon=20) - target_ds = create_sample_dataset(nlat=15, nlon=25) - regridder = Regridder(source_ds, target_ds) - assert isinstance(regridder._weights_matrix, csr_matrix) - - -def test_regrid_with_crs_grid(): - # Source grid: global 10 degree - src_ds = create_global_grid(res_lat=10, res_lon=10) - src_ds["data"] = (("lat", "lon"), np.ones((src_ds.lat.size, src_ds.lon.size))) - - # Target grid: UTM zone 33N - extent = (400000, 500000, 5000000, 5100000) - res = 10000 - tgt_ds = create_grid_from_crs("EPSG:32633", extent, res) - - regridder = Regridder(src_ds, tgt_ds, method="bilinear") - out = regridder(src_ds["data"]) - - assert out.shape == (tgt_ds.y.size, tgt_ds.x.size) - assert "x" in out.coords - assert "y" in out.coords - - -def test_dataset_regrid_identity(): - """Double-Check Test: Verify Dataset regridding matches DataArray regridding for both Eager and Lazy.""" - nlat_in, nlon_in = 10, 20 - nlat_out, nlon_out = 15, 25 - - source_ds = create_sample_dataset(nlat=nlat_in, nlon=nlon_in) - # Add another variable - source_ds["humidity"] = source_ds["temperature"] * 0.8 - # Add a non-spatial variable - source_ds["scalar"] = xr.DataArray(42.0) - - target_grid = create_sample_dataset(nlat=nlat_out, nlon=nlon_out) - regridder = Regridder(source_ds, target_grid) - - # 1. Eager test - res_ds_eager = regridder(source_ds) - assert isinstance(res_ds_eager, xr.Dataset) - assert "temperature" in res_ds_eager - assert "humidity" in res_ds_eager - assert "scalar" in res_ds_eager - assert res_ds_eager["temperature"].shape == (nlat_out, nlon_out) - assert res_ds_eager["humidity"].shape == (nlat_out, nlon_out) - - # Compare with individual DataArray regridding - res_da_temp = regridder(source_ds["temperature"]) - # Need to remove history for comparison as they differ - res_ds_temp = res_ds_eager["temperature"].copy() - res_ds_temp.attrs.pop("history", None) - res_da_temp_no_hist = res_da_temp.copy() - res_da_temp_no_hist.attrs.pop("history", None) - xr.testing.assert_allclose(res_ds_temp, res_da_temp_no_hist) - - # 2. Lazy test - source_ds_lazy = source_ds.chunk({"lat": 5, "lon": 10}) - res_ds_lazy = regridder(source_ds_lazy).compute() - - # Compare eager and lazy results (ignoring history) - res_ds_eager_no_hist = res_ds_eager.copy() - res_ds_lazy_no_hist = res_ds_lazy.copy() - res_ds_eager_no_hist.attrs.pop("history", None) - res_ds_lazy_no_hist.attrs.pop("history", None) - for v in res_ds_eager_no_hist.data_vars: - res_ds_eager_no_hist[v].attrs.pop("history", None) - res_ds_lazy_no_hist[v].attrs.pop("history", None) - - xr.testing.assert_allclose(res_ds_eager_no_hist, res_ds_lazy_no_hist) diff --git a/zensical.toml b/zensical.toml index 3338a5e..351ff5e 100644 --- a/zensical.toml +++ b/zensical.toml @@ -14,6 +14,7 @@ nav = [ { "Command Line Interface" = "user-guide/cli.md" }, { "Unstructured Grids" = "user-guide/unstructured.md" }, { "Visualization" = "user-guide/visualization.md" }, + { "Scientific Hygiene" = "user-guide/hygiene.md" }, { "HPC and Dask" = "user-guide/hpc.md" }, { "Performance" = "user-guide/performance.md" } ]}, @@ -22,6 +23,7 @@ nav = [ ]}, { "API Reference" = [ { "Regridder" = "api/regridder.md" }, + { "Accessors" = "api/accessors.md" }, { "Utilities" = "api/utils.md" } ]} ]