diff --git a/.github/dependabot.yml b/.github/dependabot.yml new file mode 100644 index 0000000..556f611 --- /dev/null +++ b/.github/dependabot.yml @@ -0,0 +1,10 @@ +version: 2 +updates: + - package-ecosystem: "github-actions" + directory: "/" # .github/workflows + schedule: + interval: "monthly" + groups: + gha-workflow-deps: + patterns: + - "*" diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 143c38f..c356b86 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -14,12 +14,12 @@ jobs: env: FORCE_JAVASCRIPT_ACTIONS_TO_NODE24: true steps: - - uses: actions/checkout@v4 + - uses: actions/checkout@v7 with: persist-credentials: false - name: Setup Micromamba - uses: mamba-org/setup-micromamba@v2 + uses: mamba-org/setup-micromamba@v3 with: environment-file: environment.yml init-shell: bash diff --git a/.github/workflows/docs.yml b/.github/workflows/docs.yml index 47f6b32..9a59bce 100644 --- a/.github/workflows/docs.yml +++ b/.github/workflows/docs.yml @@ -25,10 +25,10 @@ jobs: env: FORCE_JAVASCRIPT_ACTIONS_TO_NODE24: true steps: - - uses: actions/checkout@v4 + - uses: actions/checkout@v7 - name: Setup Micromamba - uses: mamba-org/setup-micromamba@v2 + uses: mamba-org/setup-micromamba@v3 with: environment-file: environment.yml init-shell: bash @@ -70,7 +70,7 @@ jobs: import subprocess; subprocess.run(['zensical', 'build', '--clean'], check=True)" - name: Upload artifact - uses: actions/upload-pages-artifact@v3 + uses: actions/upload-pages-artifact@v5 with: path: ./site @@ -84,4 +84,4 @@ jobs: steps: - name: Deploy to GitHub Pages id: deployment - uses: actions/deploy-pages@v4 + uses: actions/deploy-pages@v5 diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 5604d71..44e94b0 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -1,14 +1,30 @@ repos: -- repo: https://github.com/astral-sh/ruff-pre-commit - rev: v0.3.0 + - repo: https://github.com/pre-commit/pre-commit-hooks + rev: v6.0.0 hooks: - - id: ruff - args: [ --fix ] - - id: ruff-format + - id: trailing-whitespace + - id: end-of-file-fixer + - id: check-yaml + args: ['--unsafe'] + - id: check-json + - id: check-toml + - id: check-shebang-scripts-are-executable + exclude: '\.sl$' # Slurm job scripts + - id: check-executables-have-shebangs + - id: check-symlinks + - id: check-added-large-files + args: ['--maxkb=1000'] -- repo: local + - repo: https://github.com/astral-sh/ruff-pre-commit + rev: v0.15.22 hooks: - - id: generate-requirements + - id: ruff-check + args: [--fix] + - id: ruff-format + + - repo: local + hooks: + - id: generate-requirements name: generate-requirements entry: python3 scripts/generate_requirements.py language: system diff --git a/benchmarks/real_benchmark.py b/benchmarks/real_benchmark.py index a5609dc..4058ade 100644 --- a/benchmarks/real_benchmark.py +++ b/benchmarks/real_benchmark.py @@ -1,7 +1,9 @@ import time + import numpy as np import xarray as xr import xesmf as xe + from xregrid import Regridder @@ -21,21 +23,15 @@ def create_sample_dataset(nlat, nlon, ntime=1): return ds -def benchmark_resolution( - name, nlat_in, nlon_in, nlat_out, nlon_out, ntime=10, trials=3 -): +def benchmark_resolution(name, nlat_in, nlon_in, nlat_out, nlon_out, ntime=10, trials=3): print(f"\n--- Benchmarking {name} ({ntime} time steps) ---") source_ds = create_sample_dataset(nlat_in, nlon_in, ntime=ntime) target_ds = create_sample_dataset(nlat_out, nlon_out, ntime=1) # --- Weight Generation --- print("Generating weights...") - regridder_xesmf = xe.Regridder( - source_ds.isel(time=0), target_ds.isel(time=0), method="bilinear", periodic=True - ) - regridder_xregrid = Regridder( - source_ds.isel(time=0), target_ds.isel(time=0), method="bilinear", periodic=True - ) + regridder_xesmf = xe.Regridder(source_ds.isel(time=0), target_ds.isel(time=0), method="bilinear", periodic=True) + regridder_xregrid = Regridder(source_ds.isel(time=0), target_ds.isel(time=0), method="bilinear", periodic=True) # --- Weight Application --- print(f"Applying weights ({trials} trials)...") @@ -64,9 +60,7 @@ def benchmark_resolution( times_xregrid.append(time.perf_counter() - start) avg_xregrid = np.mean(times_xregrid) / ntime - print( - f"App avg per time step - xESMF: {avg_xesmf:.6f}s, XRegrid: {avg_xregrid:.6f}s" - ) + print(f"App avg per time step - xESMF: {avg_xesmf:.6f}s, XRegrid: {avg_xregrid:.6f}s") print(f"Application Speedup: {avg_xesmf / avg_xregrid:.1f}x") return { @@ -85,18 +79,12 @@ def benchmark_resolution( results.append(benchmark_resolution("0.25° Global", 720, 1440, 720, 1440, ntime=3)) # 0.1° Global -results.append( - benchmark_resolution("0.1° Global", 1800, 3600, 1800, 3600, ntime=1, trials=1) -) +results.append(benchmark_resolution("0.1° Global", 1800, 3600, 1800, 3600, ntime=1, trials=1)) print("\n\n" + "=" * 50) print("FINAL RESULTS SUMMARY (SINGLE TIME STEP REGRIDDING)") print("=" * 50) -print( - f"{'Resolution':<15} | {'xESMF App (s)':<15} | {'XRegrid App (s)':<15} | {'Speedup':<10}" -) +print(f"{'Resolution':<15} | {'xESMF App (s)':<15} | {'XRegrid App (s)':<15} | {'Speedup':<10}") print("-" * 65) for r in results: - print( - f"{r['name']:<15} | {r['app_xesmf']:<15.6f} | {r['app_xregrid']:<15.6f} | {r['speedup']:<10.1f}x" - ) + print(f"{r['name']:<15} | {r['app_xesmf']:<15.6f} | {r['app_xregrid']:<15.6f} | {r['speedup']:<10.1f}x") diff --git a/benchmarks/run_benchmarks.py b/benchmarks/run_benchmarks.py index 5bdafc2..de3f409 100644 --- a/benchmarks/run_benchmarks.py +++ b/benchmarks/run_benchmarks.py @@ -36,9 +36,7 @@ def generate_mock_weights(n_src, n_dst, weights_per_row=4): return csr_matrix((data, (row, col)), shape=(n_dst, n_src)) -def benchmark_apply( - res_name, n_lat, n_lon, target_n_lat, target_n_lon, n_time=1, skipna=False -): +def benchmark_apply(res_name, n_lat, n_lon, target_n_lat, target_n_lon, n_time=1, skipna=False): n_src = n_lat * n_lon n_dst = target_n_lat * target_n_lon @@ -99,9 +97,7 @@ def benchmark_stationary_mask(n_lat, n_lon, n_time=10): # Time with skipna=True start = time.perf_counter() - _ = _apply_weights_core( - data, weights_key, ("lat", "lon"), (n_lat, n_lon), skipna=True - ) + _ = _apply_weights_core(data, weights_key, ("lat", "lon"), (n_lat, n_lon), skipna=True) end = time.perf_counter() total_time = end - start @@ -126,7 +122,7 @@ def benchmark_stationary_mask(n_lat, n_lon, n_time=10): print("| Time Steps | Resolution | Avg Time per Step |") print("|------------|------------|-------------------|") for n_t in [1, 10, 100]: - for name, ny, nx, tny, tnx in [ + for name, ny, nx, _tny, _tnx in [ ("1.0°", 180, 360, 180, 360), ("0.25°", 720, 1440, 720, 1440), ]: diff --git a/benchmarks/run_dask_benchmark.py b/benchmarks/run_dask_benchmark.py index ca17cce..c7d7bd6 100644 --- a/benchmarks/run_dask_benchmark.py +++ b/benchmarks/run_dask_benchmark.py @@ -29,9 +29,7 @@ def generate_mock_weights(n_src, n_dst, weights_per_row=4): def benchmark_dask(n_workers, n_chunks, n_lat=360, n_lon=720): - cluster = dask.distributed.LocalCluster( - n_workers=n_workers, threads_per_worker=1, processes=True - ) + cluster = dask.distributed.LocalCluster(n_workers=n_workers, threads_per_worker=1, processes=True) client = dask.distributed.Client(cluster) try: @@ -41,9 +39,7 @@ def benchmark_dask(n_workers, n_chunks, n_lat=360, n_lon=720): # 20 time steps data = np.random.rand(20, n_lat, n_lon).astype(np.float32) - da = xr.DataArray(data, dims=("time", "lat", "lon")).chunk( - {"time": 20 // n_chunks} - ) + da = xr.DataArray(data, dims=("time", "lat", "lon")).chunk({"time": 20 // n_chunks}) # We need to distribute weights to workers weights_key = "bench_weights" diff --git a/docs/examples/scripts/create_sample_data.py b/docs/examples/scripts/create_sample_data.py index 2eb07b4..877b2c4 100644 --- a/docs/examples/scripts/create_sample_data.py +++ b/docs/examples/scripts/create_sample_data.py @@ -8,6 +8,7 @@ """ import numpy as np + from xregrid.utils import create_global_grid @@ -18,9 +19,7 @@ def create_sample_data(): # Add some dummy data lat = ds.lat.values lon = ds.lon.values - data = ( - np.sin(np.deg2rad(lat))[:, np.newaxis] * np.cos(np.deg2rad(lon))[np.newaxis, :] - ) + data = np.sin(np.deg2rad(lat))[:, np.newaxis] * np.cos(np.deg2rad(lon))[np.newaxis, :] ds["sample_var"] = (["lat", "lon"], data) ds["sample_var"].attrs["units"] = "dimensionless" diff --git a/docs/examples/scripts/plot_accessor_showcase.py b/docs/examples/scripts/plot_accessor_showcase.py index 0f7fdd5..2b9f9ff 100644 --- a/docs/examples/scripts/plot_accessor_showcase.py +++ b/docs/examples/scripts/plot_accessor_showcase.py @@ -11,9 +11,9 @@ - Passing Regridder parameters through the accessor """ -import xarray as xr -import numpy as np import matplotlib.pyplot as plt +import numpy as np +import xarray as xr # Load air_temperature tutorial dataset ds = xr.tutorial.open_dataset("air_temperature").isel(time=0) diff --git a/docs/examples/scripts/plot_air_temperature.py b/docs/examples/scripts/plot_air_temperature.py index 836d510..922da60 100644 --- a/docs/examples/scripts/plot_air_temperature.py +++ b/docs/examples/scripts/plot_air_temperature.py @@ -12,9 +12,10 @@ - Regridding 3D datasets (time, lat, lon) """ -import xarray as xr -import numpy as np import matplotlib.pyplot as plt +import numpy as np +import xarray as xr + from xregrid import Regridder # Load air_temperature tutorial dataset (North America, 2.5° resolution) diff --git a/docs/examples/scripts/plot_basic_regridding.py b/docs/examples/scripts/plot_basic_regridding.py index 93bb1ea..7a45778 100644 --- a/docs/examples/scripts/plot_basic_regridding.py +++ b/docs/examples/scripts/plot_basic_regridding.py @@ -13,9 +13,10 @@ - Handling global periodicity """ -import xarray as xr -import numpy as np import matplotlib.pyplot as plt +import numpy as np +import xarray as xr + from xregrid import Regridder # Load tutorial dataset (global, 0.75° resolution) diff --git a/docs/examples/scripts/plot_conservative_regridding.py b/docs/examples/scripts/plot_conservative_regridding.py index fd8dd22..7581d4c 100644 --- a/docs/examples/scripts/plot_conservative_regridding.py +++ b/docs/examples/scripts/plot_conservative_regridding.py @@ -10,8 +10,9 @@ (bounds), which are automatically provided by XRegrid's grid creation utilities. """ -import numpy as np import matplotlib.pyplot as plt +import numpy as np + from xregrid import Regridder, create_global_grid # 1. Create a source grid with boundaries (2.0° resolution) diff --git a/docs/examples/scripts/plot_curvilinear_grids.py b/docs/examples/scripts/plot_curvilinear_grids.py index d333241..f4aebb4 100644 --- a/docs/examples/scripts/plot_curvilinear_grids.py +++ b/docs/examples/scripts/plot_curvilinear_grids.py @@ -7,9 +7,10 @@ We use the 'rasm' tutorial dataset which features a curvilinear Arctic grid. """ -import xarray as xr -import numpy as np import matplotlib.pyplot as plt +import numpy as np +import xarray as xr + from xregrid import Regridder # Load rasm tutorial dataset (curvilinear Arctic grid) diff --git a/docs/examples/scripts/plot_dateline_fix.py b/docs/examples/scripts/plot_dateline_fix.py index 9373ed4..04116c3 100644 --- a/docs/examples/scripts/plot_dateline_fix.py +++ b/docs/examples/scripts/plot_dateline_fix.py @@ -9,6 +9,7 @@ import numpy as np import xarray as xr + from xregrid import Regridder from xregrid.utils import create_global_grid @@ -55,9 +56,7 @@ def run_example(): print("Regridding complete.") print(f"Output Variables: {list(ds_regrid.data_vars)}") - print( - f"Source Longitude Range: {ds_src.lon.min().values:.1f} to {ds_src.lon.max().values:.1f}" - ) + print(f"Source Longitude Range: {ds_src.lon.min().values:.1f} to {ds_src.lon.max().values:.1f}") print("Coordinates crossing the dateline are handled correctly by using SPH_DEG.") # In a real environment with matplotlib: diff --git a/docs/examples/scripts/plot_esmpy_comparison.py b/docs/examples/scripts/plot_esmpy_comparison.py index 0ea25f3..c7e9b63 100644 --- a/docs/examples/scripts/plot_esmpy_comparison.py +++ b/docs/examples/scripts/plot_esmpy_comparison.py @@ -11,9 +11,10 @@ API while delivering better performance than other wrappers. """ -import xarray as xr -import numpy as np import matplotlib.pyplot as plt +import numpy as np +import xarray as xr + from xregrid import Regridder # --- Part 1: Load Sample Data --- diff --git a/docs/examples/scripts/plot_larger_than_memory.py b/docs/examples/scripts/plot_larger_than_memory.py index 3427764..83f87b1 100644 --- a/docs/examples/scripts/plot_larger_than_memory.py +++ b/docs/examples/scripts/plot_larger_than_memory.py @@ -17,11 +17,13 @@ - Memory-efficient processing on a local machine. """ -import xarray as xr -import numpy as np +import time + import dask.array as da +import numpy as np +import xarray as xr from dask.distributed import Client, LocalCluster -import time + from xregrid import Regridder @@ -75,9 +77,7 @@ def run_example(): # for high-resolution grids. print("\nGenerating weights in parallel...") start = time.time() - regridder = Regridder( - ds_src, ds_tgt, method="bilinear", periodic=True, parallel=True - ) + regridder = Regridder(ds_src, ds_tgt, method="bilinear", periodic=True, parallel=True) print(f"Weight generation took: {time.time() - start:.2f}s") # 5. Apply regridding (Lazy) diff --git a/docs/examples/scripts/plot_multidimensional_regridding.py b/docs/examples/scripts/plot_multidimensional_regridding.py index 37a0491..0afb661 100644 --- a/docs/examples/scripts/plot_multidimensional_regridding.py +++ b/docs/examples/scripts/plot_multidimensional_regridding.py @@ -12,9 +12,10 @@ - Global periodicity for 4D atmospheric data """ -import xarray as xr -import numpy as np import matplotlib.pyplot as plt +import numpy as np +import xarray as xr + from xregrid import Regridder # Load multidimensional tutorial dataset (ERA-Interim) diff --git a/docs/examples/scripts/plot_performance_optimization.py b/docs/examples/scripts/plot_performance_optimization.py index b6638ba..5df43cd 100644 --- a/docs/examples/scripts/plot_performance_optimization.py +++ b/docs/examples/scripts/plot_performance_optimization.py @@ -7,10 +7,12 @@ This example demonstrates how to save and load weights. """ -import xarray as xr -import numpy as np -import time import os +import time + +import numpy as np +import xarray as xr + from xregrid import Regridder # Load a larger tutorial dataset @@ -26,18 +28,14 @@ # 1. First time: Generate and save weights start = time.time() -regridder = Regridder( - ds, target_grid, method="bilinear", filename=weights_file, reuse_weights=True -) +regridder = Regridder(ds, target_grid, method="bilinear", filename=weights_file, reuse_weights=True) _ = regridder(ds.air) first_time = time.time() - start print(f"First run (with weight generation): {first_time:.2f}s") # 2. Second time: Load weights from disk start = time.time() -regridder_cached = Regridder( - ds, target_grid, method="bilinear", filename=weights_file, reuse_weights=True -) +regridder_cached = Regridder(ds, target_grid, method="bilinear", filename=weights_file, reuse_weights=True) _ = regridder_cached(ds.air) second_time = time.time() - start print(f"Second run (reusing weights): {second_time:.2f}s") diff --git a/docs/examples/scripts/plot_roms_example.py b/docs/examples/scripts/plot_roms_example.py index 4f6abc5..ebe31e2 100644 --- a/docs/examples/scripts/plot_roms_example.py +++ b/docs/examples/scripts/plot_roms_example.py @@ -13,9 +13,10 @@ - Preserving temporal and vertical dimensions """ -import xarray as xr -import numpy as np import matplotlib.pyplot as plt +import numpy as np +import xarray as xr + from xregrid import Regridder # Load ROMS_example tutorial dataset @@ -46,9 +47,7 @@ fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(15, 6)) # Plot original curvilinear data (first time step, surface level) -ds.salt.isel(ocean_time=0, s_rho=-1).plot( - ax=ax1, x="lon_rho", y="lat_rho", cmap="viridis" -) +ds.salt.isel(ocean_time=0, s_rho=-1).plot(ax=ax1, x="lon_rho", y="lat_rho", cmap="viridis") ax1.set_title("Original ROMS Salt (Surface)") # Plot regridded rectilinear data diff --git a/docs/examples/scripts/plot_unstructured_grids.py b/docs/examples/scripts/plot_unstructured_grids.py index c52ac0e..df0e443 100644 --- a/docs/examples/scripts/plot_unstructured_grids.py +++ b/docs/examples/scripts/plot_unstructured_grids.py @@ -13,9 +13,10 @@ - Nearest neighbor interpolation methods """ -import xarray as xr -import numpy as np import matplotlib.pyplot as plt +import numpy as np +import xarray as xr + from xregrid import Regridder # 1. Create a synthetic unstructured grid diff --git a/docs/examples/scripts/plot_weather_data.py b/docs/examples/scripts/plot_weather_data.py index 990ecb5..5f29101 100644 --- a/docs/examples/scripts/plot_weather_data.py +++ b/docs/examples/scripts/plot_weather_data.py @@ -7,10 +7,11 @@ regular 2D grid. """ -import xarray as xr +import matplotlib.pyplot as plt import numpy as np import pandas as pd -import matplotlib.pyplot as plt +import xarray as xr + from xregrid import Regridder # 1. Create toy weather data (similar to xarray docs) @@ -59,9 +60,7 @@ fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(15, 6)) # Plot stations as points -sc = ax1.scatter( - ds.lon, ds.lat, c=ds.tmax.isel(time=0), s=200, cmap="viridis", edgecolor="k" -) +sc = ax1.scatter(ds.lon, ds.lat, c=ds.tmax.isel(time=0), s=200, cmap="viridis", edgecolor="k") plt.colorbar(sc, ax=ax1, label="Temperature") ax1.set_title("Station Data (tmax, day 0)") ax1.set_xlabel("Longitude") diff --git a/docs/installation.md b/docs/installation.md index bd80b17..44bf42a 100644 --- a/docs/installation.md +++ b/docs/installation.md @@ -196,4 +196,4 @@ For large grids: - Linux (tested on Ubuntu, CentOS, RHEL) - macOS (Intel and Apple Silicon) -- Windows (via WSL or native with conda) \ No newline at end of file +- Windows (via WSL or native with conda) diff --git a/docs/stylesheets/extra.css b/docs/stylesheets/extra.css index a6d2bf5..1c47671 100644 --- a/docs/stylesheets/extra.css +++ b/docs/stylesheets/extra.css @@ -72,4 +72,4 @@ table th { .md-search-result__teaser mark { background-color: #ffeb3b; color: #000; -} \ No newline at end of file +} diff --git a/pyproject.toml b/pyproject.toml index d2ac075..28d96dd 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -60,3 +60,22 @@ where = ["src"] [project.scripts] xregrid = "xregrid.cli:main" + +[tool.ruff] +target-version = "py311" +line-length = 132 + +[tool.ruff.lint] +select = [ + "E", # pycodestyle errors + "W", # pycodestyle warnings + "F", # pyflakes + "I", # isort + "B", # flake8-bugbear + "C4", # flake8-comprehensions + "UP", # pyupgrade +] +ignore = [ + "E501", # line too long + "B028", # explicit `stacklevel` in `warnings.warn` +] diff --git a/scripts/generate_requirements.py b/scripts/generate_requirements.py index 22e458b..b2dbaed 100644 --- a/scripts/generate_requirements.py +++ b/scripts/generate_requirements.py @@ -26,18 +26,14 @@ def main(): # Core requirements with open("requirements.txt", "w") as f: - f.write( - "# This file is auto-generated from pyproject.toml. Do not edit directly.\n" - ) + f.write("# This file is auto-generated from pyproject.toml. Do not edit directly.\n") for dep in deps: f.write(f"{dep}\n") print("Successfully generated requirements.txt") # Custom requirements without esmpy (useful for custom ESMF builds) with open("requirements_no_esmpy.txt", "w") as f: - f.write( - "# This file is auto-generated from pyproject.toml. Do not edit directly.\n" - ) + f.write("# This file is auto-generated from pyproject.toml. Do not edit directly.\n") f.write("# It omits esmpy to support custom ESMF installations.\n") for dep in deps: if dep != "esmpy": @@ -47,9 +43,7 @@ def main(): # Optional requirements (e.g. test, viz) optional_deps = data.get("project", {}).get("optional-dependencies", {}) if not optional_deps: - optional_deps = ( - data.get("tool", {}).get("xregrid", {}).get("optional-dependencies", {}) - ) + optional_deps = data.get("tool", {}).get("xregrid", {}).get("optional-dependencies", {}) for extra, extra_deps in optional_deps.items(): if extra == "full": @@ -57,9 +51,7 @@ def main(): filename = f"requirements-{extra}.txt" with open(filename, "w") as f: - f.write( - f"# This file is auto-generated from pyproject.toml [{extra}]. Do not edit directly.\n" - ) + f.write(f"# This file is auto-generated from pyproject.toml [{extra}]. Do not edit directly.\n") for dep in extra_deps: # If it's a reference to the package itself, like xregrid[test], skip or handle if dep.startswith("xregrid["): diff --git a/setup.py b/setup.py index 0e6c167..183f889 100644 --- a/setup.py +++ b/setup.py @@ -1,5 +1,6 @@ -import os import importlib.util +import os + from setuptools import setup # Use tomllib (Python 3.11+) or tomli @@ -35,11 +36,7 @@ def get_install_requires(): with open("pyproject.toml", "rb") as f: data = tomllib.load(f) # Read from the custom [tool.xregrid] section - deps = ( - data.get("tool", {}) - .get("xregrid", {}) - .get("dependencies", default_deps) - ) + deps = data.get("tool", {}).get("xregrid", {}).get("dependencies", default_deps) except Exception: deps = default_deps @@ -63,9 +60,7 @@ def get_install_requires(): # User has ESMF but maybe not esmpy yet; they likely want to build it. print("\n" + "=" * 80) print("NOTICE: ESMFMKFILE detected but esmpy is not installed.") - print( - "We are omitting the 'esmpy' requirement to allow manual installation." - ) + print("We are omitting the 'esmpy' requirement to allow manual installation.") print("Please install esmpy from the ESMF source tree:") print(" cd $ESMF_DIR/src/addon/esmpy && python setup.py install") print("=" * 80 + "\n") diff --git a/slurm/dask_jobqueue_examples.py b/slurm/dask_jobqueue_examples.py index bbc1474..f9da40f 100644 --- a/slurm/dask_jobqueue_examples.py +++ b/slurm/dask_jobqueue_examples.py @@ -12,6 +12,7 @@ import xarray as xr from dask_jobqueue import SLURMCluster from distributed import Client + from xregrid import Regridder, create_global_grid @@ -105,9 +106,7 @@ def run_regridding(cluster): # 3. Initialize Regridder with parallel=True # This will use the Dask cluster to generate weights in parallel - regridder = Regridder( - ds, target_grid, method="bilinear", periodic=True, parallel=True - ) + regridder = Regridder(ds, target_grid, method="bilinear", periodic=True, parallel=True) # 4. Apply regridding # The application itself will also be parallelized across the cluster diff --git a/src/xregrid/__init__.py b/src/xregrid/__init__.py index 9a98bf9..76fa147 100644 --- a/src/xregrid/__init__.py +++ b/src/xregrid/__init__.py @@ -13,6 +13,7 @@ spatial_slice, unstructured_to_scrip, ) + from .viz import plot, plot_comparison, plot_interactive, plot_static from .xregrid import Regridder diff --git a/src/xregrid/accessors.py b/src/xregrid/accessors.py index 5604733..d24c227 100644 --- a/src/xregrid/accessors.py +++ b/src/xregrid/accessors.py @@ -1,6 +1,6 @@ from __future__ import annotations -from typing import Any, Union +from typing import Any import xarray as xr @@ -24,9 +24,7 @@ def __init__(self, xarray_obj: xr.DataArray): """ self._obj = xarray_obj - def to( - self, target_grid: Union[xr.Dataset, Regridder], **kwargs: Any - ) -> xr.DataArray: + def to(self, target_grid: xr.Dataset | Regridder, **kwargs: Any) -> xr.DataArray: """ Regrid the DataArray to a target grid or using a pre-computed Regridder. @@ -70,9 +68,7 @@ def get_regridder(self, target_grid: xr.Dataset, **kwargs: Any) -> Regridder: source_ds = self._obj.to_dataset(name="_tmp_data") return Regridder(source_ds, target_grid, **kwargs) - def plot_diagnostics( - self, target_grid: xr.Dataset, mode: str = "static", **kwargs: Any - ) -> Any: + def plot_diagnostics(self, target_grid: xr.Dataset, mode: str = "static", **kwargs: Any) -> Any: """ Visualize regridding diagnostics between this DataArray and a target grid. @@ -111,9 +107,7 @@ def __init__(self, xarray_obj: xr.Dataset): """ self._obj = xarray_obj - def to( - self, target_grid: Union[xr.Dataset, Regridder], **kwargs: Any - ) -> xr.Dataset: + def to(self, target_grid: xr.Dataset | Regridder, **kwargs: Any) -> xr.Dataset: """ Regrid the Dataset to a target grid or using a pre-computed Regridder. @@ -153,9 +147,7 @@ def get_regridder(self, target_grid: xr.Dataset, **kwargs: Any) -> Regridder: """ return Regridder(self._obj, target_grid, **kwargs) - def plot_diagnostics( - self, target_grid: xr.Dataset, mode: str = "static", **kwargs: Any - ) -> Any: + def plot_diagnostics(self, target_grid: xr.Dataset, mode: str = "static", **kwargs: Any) -> Any: """ Visualize regridding diagnostics between this Dataset and a target grid. diff --git a/src/xregrid/cli.py b/src/xregrid/cli.py index 6d63fe1..c243279 100644 --- a/src/xregrid/cli.py +++ b/src/xregrid/cli.py @@ -5,6 +5,7 @@ import sys import xarray as xr + from xregrid import Regridder, create_global_grid, create_regional_grid from xregrid.utils import get_rdhpcs_cluster @@ -30,9 +31,7 @@ def parse_args() -> argparse.Namespace: choices=["bilinear", "conservative", "nearest_s2d", "nearest_d2s", "patch"], help="Regridding method (default: bilinear).", ) - parser.add_argument( - "--output", "-o", default="output.nc", help="Path to the output NetCDF file." - ) + parser.add_argument("--output", "-o", default="output.nc", help="Path to the output NetCDF file.") parser.add_argument( "--extent", help="Target grid extent as min_lat,max_lat,min_lon,max_lon (only used if target is a resolution).", @@ -47,12 +46,8 @@ def parse_args() -> argparse.Namespace: action="store_true", help="Reuse weights if the weights file already exists.", ) - parser.add_argument( - "--weights-file", default="weights.nc", help="Path to the weights file." - ) - parser.add_argument( - "--skipna", action="store_true", help="Handle NaNs by re-normalizing weights." - ) + parser.add_argument("--weights-file", default="weights.nc", help="Path to the weights file.") + parser.add_argument("--skipna", action="store_true", help="Handle NaNs by re-normalizing weights.") # Dask options dask_group = parser.add_argument_group("Dask options") @@ -100,9 +95,7 @@ def main() -> None: client = Client(args.dask_scheduler) print(f"Connected to Dask scheduler: {args.dask_scheduler}") elif args.dask_jobqueue: - cluster = get_rdhpcs_cluster( - machine=args.dask_jobqueue, account=args.dask_account - ) + cluster = get_rdhpcs_cluster(machine=args.dask_jobqueue, account=args.dask_account) from dask.distributed import Client client = Client(cluster) @@ -121,28 +114,18 @@ def main() -> None: try: res = float(args.target) if args.extent: - lat_min, lat_max, lon_min, lon_max = map( - float, args.extent.split(",") - ) - print( - f"Creating regional target grid: res={res}, extent=[{lat_min}, {lat_max}, {lon_min}, {lon_max}]" - ) - ds_tgt = create_regional_grid( - (lat_min, lat_max), (lon_min, lon_max), res, res - ) + lat_min, lat_max, lon_min, lon_max = map(float, args.extent.split(",")) + print(f"Creating regional target grid: res={res}, extent=[{lat_min}, {lat_max}, {lon_min}, {lon_max}]") + ds_tgt = create_regional_grid((lat_min, lat_max), (lon_min, lon_max), res, res) else: print(f"Creating global target grid: res={res}") ds_tgt = create_global_grid(res, res) except ValueError: - print( - f"Error: target '{args.target}' is neither a file nor a valid resolution." - ) + print(f"Error: target '{args.target}' is neither a file nor a valid resolution.") sys.exit(1) # 4. Initialize Regridder - print( - f"Initializing Regridder (method={args.method}, periodic={args.periodic})" - ) + print(f"Initializing Regridder (method={args.method}, periodic={args.periodic})") regridder = Regridder( ds_src, ds_tgt, diff --git a/src/xregrid/constants.py b/src/xregrid/constants.py index ca40397..b2c28c1 100644 --- a/src/xregrid/constants.py +++ b/src/xregrid/constants.py @@ -1,6 +1,6 @@ from __future__ import annotations -from typing import Any, Dict, Optional +from typing import Any try: import esmpy @@ -8,7 +8,7 @@ esmpy = None -def get_regrid_method_map() -> Dict[str, Any]: +def get_regrid_method_map() -> dict[str, Any]: """ Get the mapping of string names to ESMF RegridMethod constants. @@ -28,7 +28,7 @@ def get_regrid_method_map() -> Dict[str, Any]: } -def get_extrap_method_map() -> Dict[str, Any]: +def get_extrap_method_map() -> dict[str, Any]: """ Get the mapping of string names to ESMF ExtrapMethod constants. @@ -46,7 +46,7 @@ def get_extrap_method_map() -> Dict[str, Any]: } -def get_coord_sys(name: str = "SPH_DEG") -> Optional[Any]: +def get_coord_sys(name: str = "SPH_DEG") -> Any | None: """ Get an ESMF CoordSys constant by name. diff --git a/src/xregrid/core.py b/src/xregrid/core.py index b234531..0fc0ba0 100644 --- a/src/xregrid/core.py +++ b/src/xregrid/core.py @@ -1,13 +1,11 @@ from __future__ import annotations -from typing import Any, Optional, Tuple - -import numpy as np - - # Global cache for workers to reuse ESMF source objects and weight matrices # We use builtins to ensure the cache survives module re-imports in Dask workers. import builtins +from typing import Any + +import numpy as np if not hasattr(builtins, "_XREGRID_WORKER_CACHE"): builtins._XREGRID_WORKER_CACHE = {} # type: ignore @@ -75,12 +73,12 @@ def _matmul(matrix: Any, data: np.ndarray) -> np.ndarray: def _apply_weights_core( data_block: np.ndarray, weights_matrix: Any, - dims_source: Tuple[str, ...], - shape_target: Tuple[int, ...], + dims_source: tuple[str, ...], + shape_target: tuple[int, ...], skipna: bool = False, - total_weights: Optional[np.ndarray] = None, + total_weights: np.ndarray | None = None, na_thres: float = 1.0, - weights_key: Optional[str] = None, + weights_key: str | None = None, ) -> np.ndarray: """ Apply regridding weights to a data block (NumPy array). @@ -124,17 +122,13 @@ def _apply_weights_core( weights_matrix = _WORKER_CACHE.get(weights_matrix_key) if weights_matrix is None: - raise RuntimeError( - f"Weights key '{weights_matrix_key}' not found in worker cache." - ) + raise RuntimeError(f"Weights key '{weights_matrix_key}' not found in worker cache.") if isinstance(total_weights, str): total_weights_key = total_weights total_weights = _WORKER_CACHE.get(total_weights_key) if total_weights is None: - raise RuntimeError( - f"Total weights key '{total_weights_key}' not found in worker cache." - ) + raise RuntimeError(f"Total weights key '{total_weights_key}' not found in worker cache.") original_shape = data_block.shape # Core dimensions are at the end @@ -226,9 +220,7 @@ def _apply_weights_core( # Masking of low-confidence points # Ensure NaN value doesn't force promotion to float64 nan_val = result.dtype.type(np.nan) - result = np.where( - fraction_valid < (1.0 - na_thres - 1e-6), nan_val, result - ) + result = np.where(fraction_valid < (1.0 - na_thres - 1e-6), nan_val, result) else: # Standard path (skipna=False): Just apply weights result = _matmul(weights_matrix, flat_data) diff --git a/src/xregrid/grid.py b/src/xregrid/grid.py index d4fff0b..1f59669 100644 --- a/src/xregrid/grid.py +++ b/src/xregrid/grid.py @@ -1,16 +1,16 @@ from __future__ import annotations -from typing import Any, Optional, Tuple, Union +from typing import Any import cf_xarray # noqa: F401 import numpy as np import xarray as xr -from .utils import _find_coord, is_lazy from .constants import get_coord_sys +from .utils import _find_coord, is_lazy -def _get_non_spatial_dims(ds: Union[xr.Dataset, xr.DataArray]) -> set[str]: +def _get_non_spatial_dims(ds: xr.Dataset | xr.DataArray) -> set[str]: """ Identify dimensions that are likely not spatial (Time, Vertical). @@ -100,9 +100,7 @@ def _get_non_spatial_dims(ds: Union[xr.Dataset, xr.DataArray]) -> set[str]: for dim in ds.dims: if hasattr(ds, "coords") and dim in ds.coords: dtype = ds.coords[dim].dtype - if np.issubdtype(dtype, np.datetime64) or np.issubdtype( - dtype, np.timedelta64 - ): + if np.issubdtype(dtype, np.datetime64) or np.issubdtype(dtype, np.timedelta64): non_spatial_dims.add(str(dim)) return non_spatial_dims @@ -110,9 +108,9 @@ def _get_non_spatial_dims(ds: Union[xr.Dataset, xr.DataArray]) -> set[str]: def _get_mesh_info( ds: xr.Dataset, - method: Optional[str] = None, + method: str | None = None, is_source: bool = True, -) -> Tuple[xr.DataArray, xr.DataArray, Tuple[int, ...], Tuple[str, ...], bool]: +) -> tuple[xr.DataArray, xr.DataArray, tuple[int, ...], tuple[str, ...], bool]: """ Detect grid type and extract coordinates and shape from a dataset. @@ -143,7 +141,7 @@ def _get_mesh_info( # Handle uxarray objects if hasattr(ds, "uxgrid"): - uxgrid = getattr(ds, "uxgrid") + uxgrid = ds.uxgrid try: # Check if data variable is on faces use_faces = False @@ -163,11 +161,7 @@ def _get_mesh_info( if "n_face" in first_var.dims or "nFaces" in first_var.dims: use_faces = True - if ( - use_faces - and hasattr(uxgrid, "face_lat") - and hasattr(uxgrid, "face_lon") - ): + if use_faces and hasattr(uxgrid, "face_lat") and hasattr(uxgrid, "face_lon"): lat = uxgrid.face_lat lon = uxgrid.face_lon else: @@ -177,16 +171,8 @@ def _get_mesh_info( # If they share same dim, it's unstructured if lat.dims == lon.dims: # Apply filtering before returning - lat_isel = { - d: 0 - for d in non_spatial_dims - if d in lat.dims and len(lat.dims) > 1 - } - lon_isel = { - d: 0 - for d in non_spatial_dims - if d in lon.dims and len(lon.dims) > 1 - } + lat_isel = {d: 0 for d in non_spatial_dims if d in lat.dims and len(lat.dims) > 1} + lon_isel = {d: 0 for d in non_spatial_dims if d in lon.dims and len(lon.dims) > 1} if lat_isel: lat = lat.isel(lat_isel, drop=True) if lon_isel: @@ -212,10 +198,9 @@ def _get_mesh_info( if first_var is not None: for c_name in ds.coords: c = ds[c_name] - if ( - c.attrs.get("standard_name") == "latitude" - or "lat" in c_name.lower() - ) and set(c.dims).issubset(set(first_var.dims)): + if (c.attrs.get("standard_name") == "latitude" or "lat" in c_name.lower()) and set(c.dims).issubset( + set(first_var.dims) + ): lat = c break @@ -312,9 +297,7 @@ def _get_mesh_info( is_unstructured_fmt = True elif "mesh" in lon.attrs and "location" in lon.attrs: is_unstructured_fmt = True - elif any(d in lat.dims for d in unstructured_dims) or any( - d in ds.dims for d in unstructured_dims - ): + elif any(d in lat.dims for d in unstructured_dims) or any(d in ds.dims for d in unstructured_dims): is_unstructured_fmt = True else: for var in ds.variables: @@ -335,19 +318,11 @@ def _get_mesh_info( lon_vals = lon.values # Separable if std-dev along axis of variation is ~0 for the *other* axis. # lat should be constant along dim-1 (columns), lon constant along dim-0 (rows). - lat_col_std = float( - _np.nanstd(lat_vals - lat_vals[:, :1], axis=1).max() - ) - lon_row_std = float( - _np.nanstd(lon_vals - lon_vals[:1, :], axis=0).max() - ) + lat_col_std = float(_np.nanstd(lat_vals - lat_vals[:, :1], axis=1).max()) + lon_row_std = float(_np.nanstd(lon_vals - lon_vals[:1, :], axis=0).max()) if lat_col_std < 1e-6 and lon_row_std < 1e-6: - lat_1d = xr.DataArray( - lat_vals[:, 0], dims=[lat.dims[0]], attrs=lat.attrs - ) - lon_1d = xr.DataArray( - lon_vals[0, :], dims=[lon.dims[1]], attrs=lon.attrs - ) + lat_1d = xr.DataArray(lat_vals[:, 0], dims=[lat.dims[0]], attrs=lat.attrs) + lon_1d = xr.DataArray(lon_vals[0, :], dims=[lon.dims[1]], attrs=lon.attrs) lon_mesh, lat_mesh = xr.broadcast(lon_1d, lat_1d) lon_mesh = lon_mesh.transpose(lat_1d.dims[0], lon_1d.dims[0]) lat_mesh = lat_mesh.transpose(lat_1d.dims[0], lon_1d.dims[0]) @@ -402,7 +377,7 @@ def _get_mesh_info( raise ValueError("Latitude and longitude must be 1D or 2D.") -def _bounds_to_vertices(b: xr.DataArray) -> Union[xr.DataArray, np.ndarray]: +def _bounds_to_vertices(b: xr.DataArray) -> xr.DataArray | np.ndarray: """ Convert cell boundary coordinates (bounds) to vertex coordinates for ESMF. @@ -431,22 +406,14 @@ def _bounds_to_vertices(b: xr.DataArray) -> Union[xr.DataArray, np.ndarray]: elif b.ndim == 3 and b.shape[-1] == 4: # 2D curvilinear bounds (Y, X, 4) -> (Y+1, X+1) vertices v0 = b.isel({b.dims[-1]: 0}) # (y, x) - v1_last_col = b.isel({b.dims[-1]: 1}).isel( - {b.dims[1]: slice(-1, None)} - ) # (y, 1) + v1_last_col = b.isel({b.dims[-1]: 1}).isel({b.dims[1]: slice(-1, None)}) # (y, 1) row_block = xr.concat([v0, v1_last_col], dim=b.dims[1]) # (y, x+1) - v3_last_row = b.isel({b.dims[-1]: 3}).isel( - {b.dims[0]: slice(-1, None)} - ) # (1, x) - v2_last_corner = b.isel({b.dims[-1]: 2}).isel( - {b.dims[0]: slice(-1, None), b.dims[1]: slice(-1, None)} - ) # (1, 1) + v3_last_row = b.isel({b.dims[-1]: 3}).isel({b.dims[0]: slice(-1, None)}) # (1, x) + v2_last_corner = b.isel({b.dims[-1]: 2}).isel({b.dims[0]: slice(-1, None), b.dims[1]: slice(-1, None)}) # (1, 1) - last_row_block = xr.concat( - [v3_last_row, v2_last_corner], dim=b.dims[1] - ) # (1, x+1) + last_row_block = xr.concat([v3_last_row, v2_last_corner], dim=b.dims[1]) # (1, x+1) return xr.concat([row_block, last_row_block], dim=b.dims[0]) # (y+1, x+1) @@ -455,9 +422,7 @@ def _bounds_to_vertices(b: xr.DataArray) -> Union[xr.DataArray, np.ndarray]: def _get_grid_bounds( ds: xr.Dataset, -) -> Tuple[ - Optional[Union[xr.DataArray, np.ndarray]], Optional[Union[xr.DataArray, np.ndarray]] -]: +) -> tuple[xr.DataArray | np.ndarray | None, xr.DataArray | np.ndarray | None]: """ Extract grid cell boundaries from a dataset using cf-xarray or standard names. @@ -554,13 +519,13 @@ def _get_unstructured_mesh_info( ds: xr.Dataset, method: str = "conservative", is_source: bool = True, -) -> Tuple[ +) -> tuple[ np.ndarray, # node_lon np.ndarray, # node_lat np.ndarray, # element_conn np.ndarray, # element_types np.ndarray, # element_ids - Optional[np.ndarray], # orig_cell_index + np.ndarray | None, # orig_cell_index ]: """ Extract unstructured mesh connectivity and vertex info for ESMF Mesh. @@ -597,15 +562,13 @@ def _get_unstructured_mesh_info( # 0. Detect uxarray if hasattr(ds, "uxgrid"): - uxgrid = getattr(ds, "uxgrid") + uxgrid = ds.uxgrid try: node_lat = _clip_latitudes(_to_degrees(uxgrid.node_lat)).values node_lon = _normalize_longitudes(_to_degrees(uxgrid.node_lon)).values conn_raw = uxgrid.face_node_connectivity.values start_index = uxgrid.face_node_connectivity.attrs.get("start_index", 0) - fill_value = uxgrid.face_node_connectivity.attrs.get( - "_FillValue", -9223372036854775808 - ) + fill_value = uxgrid.face_node_connectivity.attrs.get("_FillValue", -9223372036854775808) # Vectorized triangulation n_cells, max_edges = conn_raw.shape @@ -656,11 +619,7 @@ def _get_unstructured_mesh_info( node_lat = _clip_latitudes(_to_degrees(v_lat)).values node_lon = _normalize_longitudes(_to_degrees(v_lon)).values conn_raw = v_conn.values - n_edges = ( - ds["nEdgesOnCell"].values - if "nEdgesOnCell" in ds - else np.full(ds.sizes["nCells"], conn_raw.shape[1]) - ) + n_edges = ds["nEdgesOnCell"].values if "nEdgesOnCell" in ds else np.full(ds.sizes["nCells"], conn_raw.shape[1]) n_cells, max_edges = conn_raw.shape max_tris = max_edges - 2 @@ -809,9 +768,7 @@ def _get_unstructured_mesh_info( coords = np.column_stack([flat_lon, flat_lat]) coords_rounded = np.round(coords, 8) - _, unique_indices, inverse_indices = np.unique( - coords_rounded, axis=0, return_index=True, return_inverse=True - ) + _, unique_indices, inverse_indices = np.unique(coords_rounded, axis=0, return_index=True, return_inverse=True) node_lon = flat_lon[unique_indices] node_lat = flat_lat[unique_indices] @@ -839,9 +796,7 @@ def _get_unstructured_mesh_info( ) # 4. Fallback for LocStreams - if method in ["nearest_s2d", "nearest_d2s"] or ( - method in ["bilinear", "patch"] and not is_source - ): + if method in ["nearest_s2d", "nearest_d2s"] or (method in ["bilinear", "patch"] and not is_source): v_lat = _find_coord(ds, "latitude") v_lon = _find_coord(ds, "longitude") @@ -866,19 +821,17 @@ def _get_unstructured_mesh_info( None, ) - raise ValueError( - f"Could not find unstructured mesh connectivity (MPAS or UGRID) for {method} regridding." - ) + raise ValueError(f"Could not find unstructured mesh connectivity (MPAS or UGRID) for {method} regridding.") def _create_esmf_grid( ds: xr.Dataset, method: str, periodic: bool = False, - mask_var: Optional[str] = None, + mask_var: str | None = None, coord_sys: Any = None, is_source: bool = True, -) -> Tuple[Any, list[str], Optional[np.ndarray]]: +) -> tuple[Any, list[str], np.ndarray | None]: """ Create an ESMF Grid or LocStream from an xarray Dataset. @@ -912,9 +865,7 @@ def _create_esmf_grid( coord_sys = get_coord_sys(coord_sys) non_spatial_dims = _get_non_spatial_dims(ds) - lon, lat, shape, dims, is_unstructured = _get_mesh_info( - ds, method=method, is_source=is_source - ) + lon, lat, shape, dims, is_unstructured = _get_mesh_info(ds, method=method, is_source=is_source) provenance = [] orig_idx = None @@ -941,9 +892,7 @@ def _create_esmf_grid( if len(element_ids) > 0: if "lat_b" in ds and "lon_b" in ds and ds["lat_b"].ndim == 2: - provenance.append( - "Derived unstructured mesh connectivity from SCRIP-style bounds." - ) + provenance.append("Derived unstructured mesh connectivity from SCRIP-style bounds.") mesh = esmpy.Mesh( parametric_dim=2, @@ -967,9 +916,7 @@ def _create_esmf_grid( if mask_var and mask_var in ds: if method == "conservative": v_mask = ds[mask_var] - mask_isel = { - d: 0 for d in non_spatial_dims if d in v_mask.dims - } + 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 @@ -994,24 +941,14 @@ def _create_esmf_grid( raise if method not in ["nearest_s2d", "nearest_d2s"] and is_source: - raise NotImplementedError( - f"Method '{method}' requires connectivity information for unstructured grids. " - ) + raise NotImplementedError(f"Method '{method}' requires connectivity information for unstructured grids. ") locstream = esmpy.LocStream(shape[0], coord_sys=coord_sys) if coord_sys == get_coord_sys("CART"): - locstream["ESMF:X"] = _normalize_longitudes(_to_degrees(lon)).values.astype( - np.float64 - ) - locstream["ESMF:Y"] = _clip_latitudes(_to_degrees(lat)).values.astype( - np.float64 - ) + locstream["ESMF:X"] = _normalize_longitudes(_to_degrees(lon)).values.astype(np.float64) + locstream["ESMF:Y"] = _clip_latitudes(_to_degrees(lat)).values.astype(np.float64) else: - locstream["ESMF:Lon"] = _normalize_longitudes( - _to_degrees(lon) - ).values.astype(np.float64) - locstream["ESMF:Lat"] = _clip_latitudes(_to_degrees(lat)).values.astype( - np.float64 - ) + locstream["ESMF:Lon"] = _normalize_longitudes(_to_degrees(lon)).values.astype(np.float64) + locstream["ESMF:Lat"] = _clip_latitudes(_to_degrees(lat)).values.astype(np.float64) if mask_var and mask_var in ds: v_mask = ds[mask_var] @@ -1031,9 +968,7 @@ def _create_esmf_grid( # periodic and cause GridCreate1PeriDim failures. if periodic and shape_f[0] < 2: periodic = False - provenance.append( - "Disabled periodic handling for degenerate grid (periodic dimension < 2 cells)." - ) + provenance.append("Disabled periodic handling for degenerate grid (periodic dimension < 2 cells).") num_peri_dims = 1 if periodic else None periodic_dim = 0 if periodic else None @@ -1046,17 +981,13 @@ def _create_esmf_grid( ds_with_bounds = ds.cf.add_bounds(["latitude", "longitude"]) lat_b, lon_b = _get_grid_bounds(ds_with_bounds) if lat_b is not None and lon_b is not None: - provenance.append( - f"Automatically generated cell boundaries for {method} regridding." - ) + provenance.append(f"Automatically generated cell boundaries for {method} regridding.") except Exception: pass has_bounds = lat_b is not None and lon_b is not None if method == "conservative" and not has_bounds: - raise ValueError( - "Conservative regridding requires cell boundaries (bounds)." - ) + raise ValueError("Conservative regridding requires cell boundaries (bounds).") staggerlocs = [esmpy.StaggerLoc.CENTER] if has_bounds: @@ -1074,12 +1005,8 @@ def _create_esmf_grid( pole_dim=pole_dim, ) - grid.get_coords(0, staggerloc=esmpy.StaggerLoc.CENTER)[...] = lon_f.astype( - np.float64 - ) - grid.get_coords(1, staggerloc=esmpy.StaggerLoc.CENTER)[...] = lat_f.astype( - np.float64 - ) + grid.get_coords(0, staggerloc=esmpy.StaggerLoc.CENTER)[...] = lon_f.astype(np.float64) + grid.get_coords(1, staggerloc=esmpy.StaggerLoc.CENTER)[...] = lat_f.astype(np.float64) if has_bounds: if lon_b.ndim == 1 and lat_b.ndim == 1: @@ -1103,12 +1030,8 @@ def _create_esmf_grid( lon_b_vals_f = lon_b_vals_f[:-1, :] lat_b_vals_f = lat_b_vals_f[:-1, :] - grid.get_coords(0, staggerloc=esmpy.StaggerLoc.CORNER)[...] = ( - lon_b_vals_f.astype(np.float64) - ) - grid.get_coords(1, staggerloc=esmpy.StaggerLoc.CORNER)[...] = ( - lat_b_vals_f.astype(np.float64) - ) + grid.get_coords(0, staggerloc=esmpy.StaggerLoc.CORNER)[...] = lon_b_vals_f.astype(np.float64) + grid.get_coords(1, staggerloc=esmpy.StaggerLoc.CORNER)[...] = lat_b_vals_f.astype(np.float64) if mask_var and mask_var in ds: v_mask = ds[mask_var] @@ -1117,7 +1040,5 @@ def _create_esmf_grid( v_mask = v_mask.isel(mask_isel, drop=True) grid.add_item(esmpy.GridItem.MASK, staggerloc=esmpy.StaggerLoc.CENTER) - grid.get_item(esmpy.GridItem.MASK, staggerloc=esmpy.StaggerLoc.CENTER)[ - ... - ] = v_mask.values.T.astype(np.int32) + grid.get_item(esmpy.GridItem.MASK, staggerloc=esmpy.StaggerLoc.CENTER)[...] = v_mask.values.T.astype(np.int32) return grid, provenance, None diff --git a/src/xregrid/parallel.py b/src/xregrid/parallel.py index d6c58b4..9a847b8 100644 --- a/src/xregrid/parallel.py +++ b/src/xregrid/parallel.py @@ -1,17 +1,17 @@ from __future__ import annotations -from typing import Any, Optional, Tuple, Union +from typing import Any import numpy as np import xarray as xr +from xregrid.constants import get_extrap_method_map, get_regrid_method_map from xregrid.core import _WORKER_CACHE, _setup_worker_cache -from xregrid.constants import get_regrid_method_map, get_extrap_method_map from xregrid.grid import _create_esmf_grid def _assemble_weights_task( - results: list[Tuple[np.ndarray, np.ndarray, np.ndarray, Optional[str]]], + results: list[tuple[np.ndarray, np.ndarray, np.ndarray, str | None]], n_src: int, n_dst: int, ) -> Any: @@ -123,9 +123,7 @@ def _populate_cache_task(value: Any, key: str) -> None: _setup_worker_cache(key, value) -def _sync_cache_from_worker_data( - future_key: str, cache_key: str, dask_worker: Any = None -) -> None: +def _sync_cache_from_worker_data(future_key: str, cache_key: str, dask_worker: Any = None) -> None: """ Internal worker task to sync worker-local cache from Dask worker data. @@ -159,13 +157,13 @@ def _compute_chunk_weights( source_ds: xr.Dataset, chunk_ds: xr.Dataset, method: str, - dest_slice_info: Union[np.ndarray, Tuple[int, int, int, int, int]], - extrap_method: Optional[str] = None, + dest_slice_info: np.ndarray | tuple[int, int, int, int, int], + extrap_method: str | None = None, extrap_dist_exponent: float = 2.0, - mask_var: Optional[str] = None, + mask_var: str | None = None, periodic: bool = False, coord_sys: Any = None, -) -> Tuple[np.ndarray, np.ndarray, np.ndarray, Optional[str]]: +) -> tuple[np.ndarray, np.ndarray, np.ndarray, str | None]: """ Worker function to compute weights for a specific chunk of the target grid. @@ -216,30 +214,18 @@ def _compute_chunk_weights( if src_cache_key in _WORKER_CACHE: src_field, src_orig_idx = _WORKER_CACHE[src_cache_key] else: - src_obj, _, src_orig_idx = _create_esmf_grid( - source_ds, method, periodic, mask_var, coord_sys=coord_sys - ) + src_obj, _, src_orig_idx = _create_esmf_grid(source_ds, method, periodic, mask_var, coord_sys=coord_sys) if isinstance(src_obj, esmpy.Mesh): - meshloc = ( - esmpy.MeshLoc.ELEMENT - if method == "conservative" - else esmpy.MeshLoc.NODE - ) + meshloc = esmpy.MeshLoc.ELEMENT if method == "conservative" else esmpy.MeshLoc.NODE src_field = esmpy.Field(src_obj, name="src", meshloc=meshloc) else: src_field = esmpy.Field(src_obj, name="src") _WORKER_CACHE[src_cache_key] = (src_field, src_orig_idx) # 2. Create target ESMF object (chunk is small, no need to cache) - dst_obj, _, dst_orig_idx = _create_esmf_grid( - chunk_ds, method, periodic=False, mask_var=None, coord_sys=coord_sys - ) + dst_obj, _, dst_orig_idx = _create_esmf_grid(chunk_ds, method, periodic=False, mask_var=None, coord_sys=coord_sys) if isinstance(dst_obj, esmpy.Mesh): - meshloc = ( - esmpy.MeshLoc.ELEMENT - if method == "conservative" - else esmpy.MeshLoc.NODE - ) + meshloc = esmpy.MeshLoc.ELEMENT if method == "conservative" else esmpy.MeshLoc.NODE dst_field = esmpy.Field(dst_obj, name="dst", meshloc=meshloc) else: dst_field = esmpy.Field(dst_obj, name="dst") @@ -284,16 +270,12 @@ def _compute_chunk_weights( dst_field.get_area() dst_areas = np.asarray(dst_field.data) n_dst = int(np.max(dst_orig_idx) + 1) if len(dst_orig_idx) > 0 else 0 - orig_dst_areas = np.bincount( - dst_orig_idx, weights=dst_areas, minlength=n_dst - ) + orig_dst_areas = np.bincount(dst_orig_idx, weights=dst_areas, minlength=n_dst) # Avoid division by zero scale_factors = np.zeros_like(dst_areas) valid = orig_dst_areas[dst_orig_idx] > 0 - scale_factors[valid] = ( - dst_areas[valid] / orig_dst_areas[dst_orig_idx][valid] - ) + scale_factors[valid] = dst_areas[valid] / orig_dst_areas[dst_orig_idx][valid] row_dst_idx = weights["row_dst"] - 1 weights["weights"] = weights["weights"] * scale_factors[row_dst_idx] @@ -322,10 +304,7 @@ def _compute_chunk_weights( else: # Structured target (2D) n0, n1 = i0_end - i0_start, i1_end - i1_start - global_indices = ( - (np.arange(n0)[:, None] + i0_start) * total_size1 - + (np.arange(n1) + i1_start) - ).flatten() + global_indices = ((np.arange(n0)[:, None] + i0_start) * total_size1 + (np.arange(n1) + i1_start)).flatten() row_dst = weights["row_dst"] - 1 col_src = weights["col_src"] - 1 diff --git a/src/xregrid/regridder.py b/src/xregrid/regridder.py index fae2824..8e5c6dd 100644 --- a/src/xregrid/regridder.py +++ b/src/xregrid/regridder.py @@ -2,38 +2,38 @@ import os import time -from typing import TYPE_CHECKING, Any, Optional, Tuple, Type, Union +from typing import TYPE_CHECKING, Any import cf_xarray # noqa: F401 import numpy as np import xarray as xr from scipy.sparse import coo_matrix -from xregrid.utils import ( - update_history, - get_crs_info, - is_dask, - is_cubed, - _get_min_max_lazy_aware, -) from xregrid.constants import ( - get_regrid_method_map, - get_extrap_method_map, get_coord_sys, + get_extrap_method_map, + get_regrid_method_map, ) +from xregrid.core import _apply_weights_core, _setup_worker_cache from xregrid.grid import ( + _create_esmf_grid, _get_mesh_info, _get_non_spatial_dims, - _create_esmf_grid, ) -from xregrid.core import _apply_weights_core, _setup_worker_cache from xregrid.parallel import ( _assemble_weights_task, - _get_weights_sum_task, - _get_nnz_task, _compute_chunk_weights, + _get_nnz_task, + _get_weights_sum_task, _sync_cache_from_worker_data, ) +from xregrid.utils import ( + _get_min_max_lazy_aware, + get_crs_info, + is_cubed, + is_dask, + update_history, +) if TYPE_CHECKING: import dask.distributed @@ -73,10 +73,10 @@ class Regridder: """ # Internal state default values - source_grid_ds: Optional[xr.Dataset] = None - target_grid_ds: Optional[xr.Dataset] = None + source_grid_ds: xr.Dataset | None = None + target_grid_ds: xr.Dataset | None = None method: str = "bilinear" - mask_var: Optional[str] = None + mask_var: str | None = None filename: str = "weights.nc" skipna: bool = False na_thres: float = 1.0 @@ -84,32 +84,32 @@ class Regridder: provenance: list[str] = [] _uid: str = "" - _shape_source: Optional[Tuple[int, ...]] = None - _shape_target: Optional[Tuple[int, ...]] = None - _dims_source: Optional[Tuple[str, ...]] = None - _dims_target: Optional[Tuple[str, ...]] = None + _shape_source: tuple[int, ...] | None = None + _shape_target: tuple[int, ...] | None = None + _dims_source: tuple[str, ...] | None = None + _dims_target: tuple[str, ...] | None = None _is_unstructured_src: bool = False _is_unstructured_tgt: bool = False - _total_weights: Optional[Union[np.ndarray, dask.distributed.Future]] = None - _weights_matrix: Optional[Union[csr_matrix, dask.distributed.Future]] = None - _dask_client: Optional[dask.distributed.Client] = None - _dask_futures: Optional[list[dask.distributed.Future]] = None + _total_weights: np.ndarray | dask.distributed.Future | None = None + _weights_matrix: csr_matrix | dask.distributed.Future | None = None + _dask_client: dask.distributed.Client | None = None + _dask_futures: list[dask.distributed.Future] | None = None def __init__( self, source_grid_ds: xr.Dataset, target_grid_ds: xr.Dataset, method: str = "bilinear", - mask_var: Optional[str] = None, + mask_var: str | None = None, reuse_weights: bool = False, filename: str = "weights.nc", skipna: bool = False, na_thres: float = 1.0, - periodic: Optional[bool] = None, + periodic: bool | None = None, mpi: bool = False, parallel: bool = False, compute: bool = True, - extrap_method: Optional[str] = None, + extrap_method: str | None = None, extrap_dist_exponent: float = 2.0, ) -> None: """ @@ -152,17 +152,14 @@ def __init__( Exponent for IDW extrapolation. """ if mpi and parallel: - raise ValueError( - "Cannot use both MPI and Dask (parallel=True) simultaneously." - ) + raise ValueError("Cannot use both MPI and Dask (parallel=True) simultaneously.") if parallel: import importlib.util if importlib.util.find_spec("dask.distributed") is None: raise ImportError( - "Dask distributed is required for parallel=True. " - "Please install it via `pip install dask distributed`." + "Dask distributed is required for parallel=True. Please install it via `pip install dask distributed`." ) # Initialize ESMF Manager (required for some environments) @@ -176,9 +173,7 @@ def __init__( # Use MULTI logkind for MPI parallelization # Some versions of esmpy don't support logkind in Manager constructor try: - self._manager = esmpy.Manager( - logkind=esmpy.LogKind.MULTI, debug=False - ) + self._manager = esmpy.Manager(logkind=esmpy.LogKind.MULTI, debug=False) except TypeError: self._manager = esmpy.Manager(debug=False) else: @@ -218,18 +213,12 @@ def __init__( # Default to geographic if no CRS found (common for simple lat-lon) is_geographic = True - if (src_crs and not src_crs.is_geographic) or ( - tgt_crs and not tgt_crs.is_geographic - ): + if (src_crs and not src_crs.is_geographic) or (tgt_crs and not tgt_crs.is_geographic): is_geographic = False # Determine if we have unstructured grids - _, _, _, _, is_unstructured_src = _get_mesh_info( - source_grid_ds, method=method, is_source=True - ) - _, _, _, _, is_unstructured_tgt = _get_mesh_info( - target_grid_ds, method=method, is_source=False - ) + _, _, _, _, is_unstructured_src = _get_mesh_info(source_grid_ds, method=method, is_source=True) + _, _, _, _, is_unstructured_tgt = _get_mesh_info(target_grid_ds, method=method, is_source=False) # Use SPH_DEG for all geographic (lat/lon in degrees) grids. # Using CART for geographic data is incorrect: ESMF's CART coordinate @@ -260,21 +249,21 @@ def __init__( self.target_grid_ds, self._tgt_was_sorted = self._normalize_grid(target_grid_ds) # Internal state - self._shape_source: Optional[Tuple[int, ...]] = None - self._shape_target: Optional[Tuple[int, ...]] = None - self._dims_source: Optional[Tuple[str, ...]] = None - self._dims_target: Optional[Tuple[str, ...]] = None + self._shape_source: tuple[int, ...] | None = None + self._shape_target: tuple[int, ...] | None = None + self._dims_source: tuple[str, ...] | None = None + self._dims_target: tuple[str, ...] | None = None self._is_unstructured_src: bool = False self._is_unstructured_tgt: bool = False - self._total_weights: Optional[np.ndarray] = None - self._weights_matrix: Optional[Any] = None - self._loaded_method: Optional[str] = None - self._loaded_periodic: Optional[bool] = None - self._loaded_extrap: Optional[str] = None - self.generation_time: Optional[float] = None - self._dask_futures: Optional[list] = None - self._dask_client: Optional[Any] = None - self._dask_start_time: Optional[float] = None + self._total_weights: np.ndarray | None = None + self._weights_matrix: Any | None = None + self._loaded_method: str | None = None + self._loaded_periodic: bool | None = None + self._loaded_extrap: str | None = None + self.generation_time: float | None = None + self._dask_futures: list | None = None + self._dask_client: Any | None = None + self._dask_start_time: float | None = None self.provenance: list[str] = [] if reuse_weights and os.path.exists(filename): @@ -288,12 +277,12 @@ def __init__( @classmethod def from_weights( - cls: Type["Regridder"], + cls: type[Regridder], filename: str, source_grid_ds: xr.Dataset, target_grid_ds: xr.Dataset, **kwargs: Any, - ) -> "Regridder": + ) -> Regridder: """ Create a Regridder from a pre-computed weights file. @@ -322,7 +311,7 @@ def from_weights( **kwargs, ) - def _normalize_grid(self, ds: xr.Dataset) -> Tuple[xr.Dataset, bool]: + def _normalize_grid(self, ds: xr.Dataset) -> tuple[xr.Dataset, bool]: """ Normalize coordinate names and ensure they are in a predictable order. @@ -346,18 +335,12 @@ def _normalize_grid(self, ds: xr.Dataset) -> Tuple[xr.Dataset, bool]: # Only for rectilinear 1D coordinates # Must be 1D and not shared (unstructured grids share dimensions) - if ( - lat_da.ndim == 1 - and lon_da.ndim == 1 - and lat_da.dims[0] != lon_da.dims[0] - ): + if lat_da.ndim == 1 and lon_da.ndim == 1 and lat_da.dims[0] != lon_da.dims[0]: lat_dim = lat_da.dims[0] lon_dim = lon_da.dims[0] # Only sort if dimension coordinates are numeric - if np.issubdtype(ds[lat_dim].dtype, np.number) and np.issubdtype( - ds[lon_dim].dtype, np.number - ): + if np.issubdtype(ds[lat_dim].dtype, np.number) and np.issubdtype(ds[lon_dim].dtype, np.number): # Use indexes for monotonicity check to remain lazy. # Indexes are always in memory in xarray, so this doesn't trigger # computation of dask-backed coordinates. @@ -376,10 +359,7 @@ def _normalize_grid(self, ds: xr.Dataset) -> Tuple[xr.Dataset, bool]: if x_da.ndim == 1 and y_da.ndim == 1 and x_da.dims[0] != y_da.dims[0]: x_dim, y_dim = x_da.dims[0], y_da.dims[0] - if not ( - ds.indexes[x_dim].is_monotonic_increasing - and ds.indexes[y_dim].is_monotonic_increasing - ): + if not (ds.indexes[x_dim].is_monotonic_increasing and ds.indexes[y_dim].is_monotonic_increasing): ds = ds.sortby([y_dim, x_dim]) was_sorted = True except (KeyError, AttributeError, ValueError): @@ -399,36 +379,20 @@ def _validate_weights(self) -> None: If the loaded weights do not match the current regridding configuration. """ # Get current grid info - _, _, src_shape, src_dims, _ = _get_mesh_info( - self.source_grid_ds, method=self.method, is_source=True - ) - _, _, dst_shape, dst_dims, _ = _get_mesh_info( - self.target_grid_ds, method=self.method, is_source=False - ) + _, _, src_shape, src_dims, _ = _get_mesh_info(self.source_grid_ds, method=self.method, is_source=True) + _, _, dst_shape, dst_dims, _ = _get_mesh_info(self.target_grid_ds, method=self.method, is_source=False) if src_shape != self._shape_source: - raise ValueError( - f"Source grid shape {src_shape} does not match " - f"loaded weights source shape {self._shape_source}" - ) + raise ValueError(f"Source grid shape {src_shape} does not match loaded weights source shape {self._shape_source}") if dst_shape != self._shape_target: - raise ValueError( - f"Target grid shape {dst_shape} does not match " - f"loaded weights target shape {self._shape_target}" - ) + raise ValueError(f"Target grid shape {dst_shape} does not match loaded weights target shape {self._shape_target}") # Check regridding parameters if self._loaded_method is not None and self._loaded_method != self.method: - raise ValueError( - f"Requested method '{self.method}' does not match " - f"loaded weights method '{self._loaded_method}'" - ) + raise ValueError(f"Requested method '{self.method}' does not match loaded weights method '{self._loaded_method}'") if self._loaded_periodic is not None and self._loaded_periodic != self.periodic: - raise ValueError( - f"Requested periodic={self.periodic} does not match " - f"loaded weights periodic={self._loaded_periodic}" - ) + raise ValueError(f"Requested periodic={self.periodic} does not match loaded weights periodic={self._loaded_periodic}") if self._loaded_extrap is not None: current_extrap = self.extrap_method or "none" @@ -440,21 +404,15 @@ def _validate_weights(self) -> None: if hasattr(self, "_loaded_skipna") and self._loaded_skipna is not None: if self._loaded_skipna != self.skipna: - raise ValueError( - f"Requested skipna={self.skipna} does not match " - f"loaded weights skipna={self._loaded_skipna}" - ) + raise ValueError(f"Requested skipna={self.skipna} does not match loaded weights skipna={self._loaded_skipna}") if hasattr(self, "_loaded_na_thres") and self._loaded_na_thres is not None: if abs(self._loaded_na_thres - self.na_thres) > 1e-6: raise ValueError( - f"Requested na_thres={self.na_thres} does not match " - f"loaded weights na_thres={self._loaded_na_thres}" + f"Requested na_thres={self.na_thres} does not match loaded weights na_thres={self._loaded_na_thres}" ) - def _create_esmf_object( - self, ds: xr.Dataset, is_source: bool = True - ) -> Tuple[Any, list[str], Optional[np.ndarray]]: + def _create_esmf_object(self, ds: xr.Dataset, is_source: bool = True) -> tuple[Any, list[str], np.ndarray | None]: """ Creates an ESMF Grid or LocStream and updates internal metadata. @@ -505,31 +463,19 @@ def _generate_weights(self) -> None: return start_time = time.perf_counter() - src_obj, src_prov, src_orig_idx = self._create_esmf_object( - self.source_grid_ds, is_source=True - ) - dst_obj, dst_prov, dst_orig_idx = self._create_esmf_object( - self.target_grid_ds, is_source=False - ) + src_obj, src_prov, src_orig_idx = self._create_esmf_object(self.source_grid_ds, is_source=True) + dst_obj, dst_prov, dst_orig_idx = self._create_esmf_object(self.target_grid_ds, is_source=False) self.provenance.extend(src_prov) self.provenance.extend(dst_prov) if isinstance(src_obj, esmpy.Mesh): - meshloc = ( - esmpy.MeshLoc.ELEMENT - if self.method == "conservative" - else esmpy.MeshLoc.NODE - ) + meshloc = esmpy.MeshLoc.ELEMENT if self.method == "conservative" else esmpy.MeshLoc.NODE src_field = esmpy.Field(src_obj, name="src", meshloc=meshloc) else: src_field = esmpy.Field(src_obj, name="src") if isinstance(dst_obj, esmpy.Mesh): - meshloc = ( - esmpy.MeshLoc.ELEMENT - if self.method == "conservative" - else esmpy.MeshLoc.NODE - ) + meshloc = esmpy.MeshLoc.ELEMENT if self.method == "conservative" else esmpy.MeshLoc.NODE dst_field = esmpy.Field(dst_obj, name="dst", meshloc=meshloc) else: dst_field = esmpy.Field(dst_obj, name="dst") @@ -538,10 +484,7 @@ def _generate_weights(self) -> None: regrid_method = self.method_map[self.method] except KeyError: available_methods = ", ".join(self.method_map.keys()) - raise ValueError( - f"Method '{self.method}' is not supported. " - f"Available methods are: {available_methods}" - ) + raise ValueError(f"Method '{self.method}' is not supported. Available methods are: {available_methods}") from None regrid_kwargs = { "regrid_method": regrid_method, @@ -570,12 +513,8 @@ def _generate_weights(self) -> None: "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 + 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() @@ -595,16 +534,12 @@ def _generate_weights(self) -> None: dst_field.get_area() dst_areas = np.asarray(dst_field.data) n_dst = int(np.prod(self._shape_target)) - orig_dst_areas = np.bincount( - dst_orig_idx, weights=dst_areas, minlength=n_dst - ) + orig_dst_areas = np.bincount(dst_orig_idx, weights=dst_areas, minlength=n_dst) # Avoid division by zero scale_factors = np.zeros_like(dst_areas) valid = orig_dst_areas[dst_orig_idx] > 0 - scale_factors[valid] = ( - dst_areas[valid] / orig_dst_areas[dst_orig_idx][valid] - ) + scale_factors[valid] = dst_areas[valid] / orig_dst_areas[dst_orig_idx][valid] row_dst_idx = weights["row_dst"] - 1 weights["weights"] = weights["weights"] * scale_factors[row_dst_idx] @@ -633,7 +568,7 @@ def _generate_weights(self) -> None: rows = [] cols = [] data = [] - for i, w in enumerate(all_weights): + for w in all_weights: r = w["row_dst"] - 1 c = w["col_src"] - 1 # Note: we already mapped and scaled the weights locally on each PET, @@ -670,9 +605,7 @@ def _generate_weights(self) -> None: n_src = int(np.prod(self._shape_source)) n_dst = int(np.prod(self._shape_target)) - self._weights_matrix = coo_matrix( - (data, (rows, cols)), shape=(n_dst, n_src) - ).tocsr() + self._weights_matrix = coo_matrix((data, (rows, cols)), shape=(n_dst, n_src)).tocsr() if self.skipna: # Optimization: Use sum(axis=1) instead of memory-intensive ones multiplication @@ -698,17 +631,13 @@ def _generate_weights_dask(self, compute: bool = True) -> None: # Get grid info and populate internal state # Source - _, _, src_shape, src_dims, is_unstructured_src = _get_mesh_info( - self.source_grid_ds, method=self.method, is_source=True - ) + _, _, src_shape, src_dims, is_unstructured_src = _get_mesh_info(self.source_grid_ds, method=self.method, is_source=True) self._shape_source = src_shape self._dims_source = src_dims self._is_unstructured_src = is_unstructured_src # Target - _, _, dst_shape, dst_dims, is_unstructured_dst = _get_mesh_info( - self.target_grid_ds, method=self.method, is_source=False - ) + _, _, dst_shape, dst_dims, is_unstructured_dst = _get_mesh_info(self.target_grid_ds, method=self.method, is_source=False) self._shape_target = dst_shape self._dims_target = dst_dims self._is_unstructured_tgt = is_unstructured_dst @@ -825,7 +754,7 @@ def _generate_weights_dask(self, compute: bool = True) -> None: if compute: self.compute() - def persist(self) -> "Regridder": + def persist(self) -> Regridder: """ Ensure tasks are submitted to the cluster. @@ -866,15 +795,11 @@ def compute(self) -> None: # Perform concatenation on a worker to protect driver memory # We use top-level task functions to avoid capturing 'self' and mocks. - self._weights_matrix = self._dask_client.submit( - _assemble_weights_task, self._dask_futures, n_src, n_dst - ) + self._weights_matrix = self._dask_client.submit(_assemble_weights_task, self._dask_futures, n_src, n_dst) if self.skipna: # Compute total weights sum on worker too - self._total_weights = self._dask_client.submit( - _get_weights_sum_task, self._weights_matrix - ) + self._total_weights = self._dask_client.submit(_get_weights_sum_task, self._weights_matrix) if self._dask_start_time: self.generation_time = time.perf_counter() - self._dask_start_time @@ -924,9 +849,7 @@ def _save_weights(self) -> None: "provenance": "; ".join(self.provenance) if self.provenance else "", "extrap_method": self.extrap_method or "none", "extrap_dist_exponent": self.extrap_dist_exponent, - "generation_time": self.generation_time - if self.generation_time - else 0.0, + "generation_time": self.generation_time if self.generation_time else 0.0, }, ) update_history(ds_weights, "Weights generated by Regridder") @@ -944,7 +867,7 @@ def _load_weights(self) -> None: n_src = ds_weights.attrs["n_src"] n_dst = ds_weights.attrs["n_dst"] - def _to_tuple(attr: Any) -> Tuple[Any, ...]: + def _to_tuple(attr: Any) -> tuple[Any, ...]: """ Convert attribute to tuple. @@ -980,9 +903,7 @@ def _to_tuple(attr: Any) -> Tuple[Any, ...]: if loaded_prov: self.provenance = loaded_prov.split("; ") - self._weights_matrix = coo_matrix( - (data, (rows, cols)), shape=(n_dst, n_src) - ).tocsr() + self._weights_matrix = coo_matrix((data, (rows, cols)), shape=(n_dst, n_src)).tocsr() if self.skipna: # Optimization: Use sum(axis=1) instead of memory-intensive ones multiplication @@ -1058,11 +979,7 @@ def clear_instance_cache(self) -> None: # 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] - ] + 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): @@ -1092,18 +1009,14 @@ def diagnostics(self) -> xr.Dataset: import dask.array as da if self._total_weights is None: - self._total_weights = self._dask_client.submit( - _get_weights_sum_task, self._weights_matrix - ) + self._total_weights = self._dask_client.submit(_get_weights_sum_task, self._weights_matrix) # Convert Future to Dask array to preserve laziness n_dst = int(np.prod(self._shape_target)) # Use dask.array.from_delayed to wrap the Future (or NumPy array) # as a lazy Dask array to avoid driver-side blocking. - weights_sum_da = da.from_delayed( - dask.delayed(self._total_weights), shape=(n_dst,), dtype=np.float64 - ) + weights_sum_da = da.from_delayed(dask.delayed(self._total_weights), shape=(n_dst,), dtype=np.float64) weights_sum_2d = weights_sum_da.reshape(self._shape_target) # Preserve laziness for the mask @@ -1122,10 +1035,7 @@ def diagnostics(self) -> xr.Dataset: coords = { c: self.target_grid_ds.coords[c] for c in self.target_grid_ds.coords - if self._dims_target is not None - and set(self.target_grid_ds.coords[c].dims).issubset( - set(self._dims_target) - ) + if self._dims_target is not None and set(self.target_grid_ds.coords[c].dims).issubset(set(self._dims_target)) } dims_target = self._dims_target @@ -1156,9 +1066,7 @@ def diagnostics(self) -> xr.Dataset: ) return ds - def quality_report( - self, skip_heavy: bool = False, format: str = "dict" - ) -> Union[dict[str, Any], xr.Dataset]: + def quality_report(self, skip_heavy: bool = False, format: str = "dict") -> dict[str, Any] | xr.Dataset: """ Generate a scientific quality report of the regridding weights. @@ -1205,15 +1113,11 @@ def quality_report( import dask.array as da client = self._dask_client or dask.distributed.get_client() - n_weights_future = client.submit( - _get_nnz_task, self._weights_matrix - ) + n_weights_future = client.submit(_get_nnz_task, self._weights_matrix) if format == "dataset": # Preserve laziness for dataset output - n_weights = da.from_delayed( - dask.delayed(n_weights_future), shape=(), dtype=int - ) + n_weights = da.from_delayed(dask.delayed(n_weights_future), shape=(), dtype=int) else: # For dict, we still need to wait to satisfy return type n_weights = int(n_weights_future.result()) @@ -1253,9 +1157,7 @@ def quality_report( { "unmapped_count": int(unmapped_count), "unmapped_fraction": float(unmapped_fraction), - "weight_sum_min": float(weight_sum_min) - if int(unmapped_count) < n_dst - else 0.0, + "weight_sum_min": float(weight_sum_min) if int(unmapped_count) < n_dst else 0.0, "weight_sum_max": float(weight_sum_max), "weight_sum_mean": float(weight_sum_mean), } @@ -1299,9 +1201,7 @@ def quality_report( "provenance": "; ".join(self.provenance), }, ) - update_history( - ds_report, f"Generated scientific quality report (backend={backend})." - ) + update_history(ds_report, f"Generated scientific quality report (backend={backend}).") return ds_report return report @@ -1445,13 +1345,9 @@ def plot_comparison( return _plot_static(da_src, da_tgt, regridder=self, **kwargs) elif mode == "interactive": rasterize = kwargs.pop("rasterize", True) - return _plot_interactive( - da_src, da_tgt, regridder=self, rasterize=rasterize, **kwargs - ) + return _plot_interactive(da_src, da_tgt, regridder=self, rasterize=rasterize, **kwargs) else: - raise ValueError( - f"Unknown plotting mode: '{mode}'. Must be 'static' or 'interactive'." - ) + raise ValueError(f"Unknown plotting mode: '{mode}'. Must be 'static' or 'interactive'.") def plot_diagnostics(self, mode: str = "static", **kwargs: Any) -> Any: """ @@ -1482,17 +1378,15 @@ def plot_diagnostics(self, mode: str = "static", **kwargs: Any) -> Any: rasterize = kwargs.pop("rasterize", True) return _plot_interactive(self, rasterize=rasterize, **kwargs) else: - raise ValueError( - f"Unknown plotting mode: '{mode}'. Must be 'static' or 'interactive'." - ) + raise ValueError(f"Unknown plotting mode: '{mode}'. Must be 'static' or 'interactive'.") def __call__( self, - obj: Union[xr.DataArray, xr.Dataset, Any], - skipna: Optional[bool] = None, - na_thres: Optional[float] = None, + obj: xr.DataArray | xr.Dataset | Any, + skipna: bool | None = None, + na_thres: float | None = None, keep_attrs: bool = True, - ) -> Union[xr.DataArray, xr.Dataset]: + ) -> xr.DataArray | xr.Dataset: """ Apply regridding to an input DataArray or Dataset. @@ -1535,8 +1429,7 @@ def __call__( import warnings warnings.warn( - "Applying serial Regridder to Dask-backed data. " - "For better performance, initialize Regridder with parallel=True." + "Applying serial Regridder to Dask-backed data. For better performance, initialize Regridder with parallel=True." ) if isinstance(obj, xr.Dataset): @@ -1557,9 +1450,7 @@ def __call__( if not is_regriddable: try: # Check for logical spatial dimensions - spatial_dims = set(obj.cf["latitude"].dims) | set( - obj.cf["longitude"].dims - ) + spatial_dims = set(obj.cf["latitude"].dims) | set(obj.cf["longitude"].dims) if spatial_dims.issubset(set(obj.dims)): is_regriddable = True except (KeyError, AttributeError): @@ -1597,10 +1488,10 @@ def _regrid_dataarray( self, da_in: xr.DataArray, update_history_attr: bool = True, - _processed_ids: Optional[set[Union[int, str]]] = None, - skipna: Optional[bool] = None, - na_thres: Optional[float] = None, - _precomputed_aux: Optional[dict[str, xr.DataArray]] = None, + _processed_ids: set[int | str] | None = None, + skipna: bool | None = None, + na_thres: float | None = None, + _precomputed_aux: dict[str, xr.DataArray] | None = None, ) -> xr.DataArray: """ Regrid a single DataArray, including auxiliary spatial coordinates. @@ -1651,14 +1542,10 @@ def _regrid_dataarray( if skipna and self._total_weights is None and self._weights_matrix is not None: if hasattr(self._weights_matrix, "key"): # Distributed path: compute total weights on cluster - self._total_weights = self._dask_client.submit( - _get_weights_sum_task, self._weights_matrix - ) + self._total_weights = self._dask_client.submit(_get_weights_sum_task, self._weights_matrix) else: # Eager path: compute locally and flatten to 1D - self._total_weights = np.array( - self._weights_matrix.sum(axis=1) - ).flatten() + self._total_weights = np.array(self._weights_matrix.sum(axis=1)).flatten() # Identify auxiliary coordinates that need regridding aux_coords_to_regrid = {} @@ -1681,9 +1568,7 @@ def _regrid_dataarray( if c_name in non_spatial_dims: continue - if c_name not in da_in.dims and all( - d in c_da.dims for d in self._dims_source - ): + if c_name not in da_in.dims and all(d in c_da.dims for d in self._dims_source): # This is an auxiliary spatial coordinate if _precomputed_aux and c_name in _precomputed_aux: aux_coords_to_regrid[c_name] = _precomputed_aux[c_name] @@ -1746,9 +1631,7 @@ def _regrid_dataarray( da_in = da_in.rename({found_dim: self._dims_source[0]}) else: # Fallback to cf-xarray discovery - da_in = da_in.cf.rename( - {da_in.cf["latitude"].dims[0]: self._dims_source[0]} - ) + da_in = da_in.cf.rename({da_in.cf["latitude"].dims[0]: self._dims_source[0]}) except (KeyError, AttributeError, ValueError): # Handle uxarray if hasattr(da_in, "uxgrid"): @@ -1802,9 +1685,7 @@ def _regrid_dataarray( ) else: # Eager matrix: run on all workers - client.run( - _setup_worker_cache, weights_key_arg, self._weights_matrix - ) + client.run(_setup_worker_cache, weights_key_arg, self._weights_matrix) _DRIVER_CACHE[(client_id, weights_key_arg)] = True weights_arg = weights_key_arg @@ -1850,7 +1731,7 @@ def _regrid_dataarray( da_in.data = da_in.data.rechunk(tuple(new_chunks)) else: # Fallback to xarray's chunk() if it's already wrapped - da_in = da_in.chunk({d: -1 for d in input_core_dims}) + da_in = da_in.chunk(dict.fromkeys(input_core_dims, -1)) # Use allow_rechunk=True to support chunked core dimensions # and move output_sizes to dask_gufunc_kwargs for future compatibility @@ -1873,9 +1754,7 @@ def _regrid_dataarray( vectorize=False, output_dtypes=[da_in.dtype], dask_gufunc_kwargs={ - "output_sizes": { - d: s for d, s in zip(temp_output_core_dims, self._shape_target) - }, + "output_sizes": dict(zip(temp_output_core_dims, self._shape_target, strict=True)), "allow_rechunk": True, }, ) @@ -1887,7 +1766,7 @@ def _regrid_dataarray( # Determine if we need to rename temp dims to target dims rename_dict = {} - for temp, orig in zip(temp_output_core_dims, self._dims_target): + for temp, orig in zip(temp_output_core_dims, self._dims_target, strict=True): if temp in out.dims: rename_dict[temp] = orig @@ -1918,10 +1797,7 @@ def _regrid_dataarray( # Aero Protocol: Ensure assigned coordinates are dimensionally compatible with the output c_dims = set(self.target_grid_ds.coords[c].dims) out_dims = set(out.dims) - if ( - c_dims.issubset(set(self._dims_target)) - or c in [target_gm_name, target_mesh_name] - ) and c_dims.issubset(out_dims): + if (c_dims.issubset(set(self._dims_target)) or c in [target_gm_name, target_mesh_name]) and c_dims.issubset(out_dims): target_coords_to_assign[c] = self.target_grid_ds.coords[c] # Also check data_vars for topology/mapping that might be needed as coords @@ -1947,10 +1823,7 @@ def _regrid_dataarray( if attr in topology_attrs: ref_vars = topology_attrs[attr].split() for rv in ref_vars: - if ( - rv in self.target_grid_ds - and rv not in target_coords_to_assign - ): + if rv in self.target_grid_ds and rv not in target_coords_to_assign: rv_dims = set(self.target_grid_ds[rv].dims) if rv_dims.issubset(set(out.dims)): target_coords_to_assign[rv] = self.target_grid_ds[rv] @@ -2109,10 +1982,7 @@ def _detect_periodicity(self, ds: xr.Dataset) -> bool: # 3. Last fallback: Check dimension name if "lon" in lon.dims or "longitude" in lon.dims: # If it's a global grid from a known generator, it might have attributes - if ( - ds.attrs.get("history") - and "global grid" in ds.attrs.get("history", "").lower() - ): + if ds.attrs.get("history") and "global grid" in ds.attrs.get("history", "").lower(): return True except Exception: pass @@ -2121,8 +1991,8 @@ def _detect_periodicity(self, ds: xr.Dataset) -> bool: def _regrid_dataset( self, ds_in: xr.Dataset, - skipna: Optional[bool] = None, - na_thres: Optional[float] = None, + skipna: bool | None = None, + na_thres: float | None = None, ) -> xr.Dataset: """ Regrid all data variables and auxiliary coordinates in a Dataset. @@ -2146,7 +2016,7 @@ def _regrid_dataset( if na_thres is None: na_thres = self.na_thres - regridded_items: dict[str, Union[xr.DataArray, Any]] = {} + regridded_items: dict[str, xr.DataArray | Any] = {} # Identify non-spatial variables to exclude from regridding non_spatial_dims = _get_non_spatial_dims(ds_in) @@ -2196,9 +2066,7 @@ def _regrid_dataset( else: try: # Check if variable has logical latitude and longitude - spatial_dims = set(da.cf["latitude"].dims) | set( - da.cf["longitude"].dims - ) + spatial_dims = set(da.cf["latitude"].dims) | set(da.cf["longitude"].dims) if spatial_dims.issubset(set(da.dims)): is_regriddable = True except (KeyError, AttributeError): @@ -2252,9 +2120,10 @@ def _regrid_dataset( for c in self.target_grid_ds.coords: if c not in out.coords: - if set(self.target_grid_ds.coords[c].dims).issubset( - set(self._dims_target) - ) or c in [target_gm_name, target_mesh_name]: + if set(self.target_grid_ds.coords[c].dims).issubset(set(self._dims_target)) or c in [ + target_gm_name, + target_mesh_name, + ]: out = out.assign_coords({c: self.target_grid_ds[c]}) # Ensure mapping/topology vars from data_vars are also attached if needed diff --git a/src/xregrid/utils.py b/src/xregrid/utils.py index 07c96b6..7b5fca9 100644 --- a/src/xregrid/utils.py +++ b/src/xregrid/utils.py @@ -4,7 +4,7 @@ import os import socket import warnings -from typing import Any, Dict, Optional, Tuple, Union +from typing import Any import numpy as np @@ -129,9 +129,7 @@ def _get_array_namespace(*objs: Any) -> Any: return np -def _lazy_arange( - start: float, stop: float, step: float, chunks: Optional[int] = None -) -> Any: +def _lazy_arange(start: float, stop: float, step: float, chunks: int | None = None) -> Any: """ Create a lazy or eager range. @@ -163,12 +161,12 @@ def _lazy_arange( def _create_rectilinear_grid( - lat_range: Tuple[float, float], - lon_range: Tuple[float, float], + lat_range: tuple[float, float], + lon_range: tuple[float, float], res_lat: float, res_lon: float, add_bounds: bool = True, - chunks: Optional[Union[int, Dict[str, int]]] = None, + chunks: int | dict[str, int] | None = None, history_msg: str = "", crs: str = "EPSG:4326", ) -> xr.Dataset: @@ -202,12 +200,8 @@ def _create_rectilinear_grid( lat_chunks = chunks.get("lat", -1) if isinstance(chunks, dict) else chunks lon_chunks = chunks.get("lon", -1) if isinstance(chunks, dict) else chunks - lat_arr = _lazy_arange( - lat_range[0] + res_lat / 2, lat_range[1], res_lat, chunks=lat_chunks - ) - lon_arr = _lazy_arange( - lon_range[0] + res_lon / 2, lon_range[1], res_lon, chunks=lon_chunks - ) + lat_arr = _lazy_arange(lat_range[0] + res_lat / 2, lat_range[1], res_lat, chunks=lat_chunks) + lon_arr = _lazy_arange(lon_range[0] + res_lon / 2, lon_range[1], res_lon, chunks=lon_chunks) ds = xr.Dataset( coords={ @@ -227,12 +221,8 @@ def _create_rectilinear_grid( if add_bounds: # Use CF-compliant (N, 2) bounds. # Ensure identical length handling for lazy/eager - lat_b_1d = _lazy_arange( - lat_range[0], lat_range[1] + res_lat, res_lat, chunks=lat_chunks - )[: lat_arr.size + 1] - lon_b_1d = _lazy_arange( - lon_range[0], lon_range[1] + res_lon, res_lon, chunks=lon_chunks - )[: lon_arr.size + 1] + lat_b_1d = _lazy_arange(lat_range[0], lat_range[1] + res_lat, res_lat, chunks=lat_chunks)[: lat_arr.size + 1] + lon_b_1d = _lazy_arange(lon_range[0], lon_range[1] + res_lon, res_lon, chunks=lon_chunks)[: lon_arr.size + 1] xp = _get_array_namespace(lat_b_1d, lon_b_1d) lat_b_2d = xp.stack([lat_b_1d[:-1], lat_b_1d[1:]], axis=1) @@ -266,7 +256,7 @@ def create_global_grid( res_lat: float, res_lon: float, add_bounds: bool = True, - chunks: Optional[Union[int, Dict[str, int]]] = None, + chunks: int | dict[str, int] | None = None, ) -> xr.Dataset: """ Create a global rectilinear grid dataset. @@ -300,12 +290,12 @@ def create_global_grid( def create_regional_grid( - lat_range: Tuple[float, float], - lon_range: Tuple[float, float], + lat_range: tuple[float, float], + lon_range: tuple[float, float], res_lat: float, res_lon: float, add_bounds: bool = True, - chunks: Optional[Union[int, Dict[str, int]]] = None, + chunks: int | dict[str, int] | None = None, ) -> xr.Dataset: """ Create a regional rectilinear grid dataset. @@ -402,7 +392,7 @@ def load_esmf_file(filepath: str) -> xr.Dataset: return ds -def get_crs_info(obj: Union[xr.DataArray, xr.Dataset]) -> Optional[Any]: +def get_crs_info(obj: xr.DataArray | xr.Dataset) -> Any | None: """ Detect CRS information from an xarray object's attributes or encoding. @@ -423,12 +413,7 @@ def get_crs_info(obj: Union[xr.DataArray, xr.Dataset]) -> Optional[Any]: # Try to detect CRS from attributes and encoding # We prioritize 'grid_mapping' then 'crs' - crs_info = ( - obj.attrs.get("grid_mapping") - or obj.encoding.get("grid_mapping") - or obj.attrs.get("crs") - or obj.encoding.get("crs") - ) + crs_info = obj.attrs.get("grid_mapping") or obj.encoding.get("grid_mapping") or obj.attrs.get("crs") or obj.encoding.get("crs") # Try cf-xarray for robust grid mapping discovery if crs_info is None or isinstance(crs_info, str): @@ -446,11 +431,7 @@ def get_crs_info(obj: Union[xr.DataArray, xr.Dataset]) -> Optional[Any]: gm_var = gms[0].array if hasattr(gms[0], "array") else gms[0] if gm_var is not None: - crs_info = ( - gm_var.attrs.get("crs_wkt") - or gm_var.attrs.get("spatial_ref") - or gm_var.attrs.get("grid_mapping_name") - ) + crs_info = gm_var.attrs.get("crs_wkt") or gm_var.attrs.get("spatial_ref") or gm_var.attrs.get("grid_mapping_name") except (AttributeError, KeyError, ImportError): pass @@ -463,9 +444,7 @@ def get_crs_info(obj: Union[xr.DataArray, xr.Dataset]) -> Optional[Any]: return None -def _find_coord( - obj: Union[xr.DataArray, xr.Dataset], key: str -) -> Optional[xr.DataArray]: +def _find_coord(obj: xr.DataArray | xr.Dataset, key: str) -> xr.DataArray | None: """ Find a coordinate in an xarray object by CF standard name or common name. @@ -502,7 +481,7 @@ def _find_coord( if set(obj[m].dims).issubset(set(obj.dims)): return obj[m] elif isinstance(obj, xr.Dataset) and len(obj.data_vars) > 0: - for name, da in obj.data_vars.items(): + for _name, da in obj.data_vars.items(): if da.attrs.get("cf_role") not in [ "mesh_topology", "face_node_connectivity", @@ -561,9 +540,7 @@ def _find_coord( return None -def update_history( - obj: Union[xr.DataArray, xr.Dataset], message: str -) -> Union[xr.DataArray, xr.Dataset]: +def update_history(obj: xr.DataArray | xr.Dataset, message: str) -> xr.DataArray | xr.Dataset: """ Update the 'history' attribute of an xarray object with a timestamped message. @@ -590,7 +567,7 @@ def update_history( def _transform_coords( x_arr: np.ndarray, y_arr: np.ndarray, crs_in: Any, crs_out: str = "EPSG:4326" -) -> Tuple[np.ndarray, np.ndarray]: +) -> tuple[np.ndarray, np.ndarray]: """ Transform coordinates using pyproj. @@ -621,11 +598,11 @@ def _transform_coords( def create_grid_from_crs( - crs: Union[str, int, Any], - extent: Tuple[float, float, float, float], - res: Union[float, Tuple[float, float]], + crs: str | int | Any, + extent: tuple[float, float, float, float], + res: float | tuple[float, float], add_bounds: bool = True, - chunks: Optional[Union[int, Dict[str, int]]] = None, + chunks: int | dict[str, int] | None = None, ) -> xr.Dataset: """ Create a structured grid dataset from a CRS and extent. @@ -757,12 +734,8 @@ def create_grid_from_crs( y_da_1d = xr.DataArray(y, dims=["y"]) # Create (N, 2) bounds - x_b_1d = xr.concat( - [x_da_1d - res_x / 2, x_da_1d + res_x / 2], dim="nbounds" - ).transpose("x", "nbounds") - y_b_1d = xr.concat( - [y_da_1d - res_y / 2, y_da_1d + res_y / 2], dim="nbounds" - ).transpose("y", "nbounds") + x_b_1d = xr.concat([x_da_1d - res_x / 2, x_da_1d + res_x / 2], dim="nbounds").transpose("x", "nbounds") + y_b_1d = xr.concat([y_da_1d - res_y / 2, y_da_1d + res_y / 2], dim="nbounds").transpose("y", "nbounds") ds.coords["x_b"] = (["x", "nbounds"], x_b_1d.data, {"units": units}) ds.coords["y_b"] = (["y", "nbounds"], y_b_1d.data, {"units": units}) @@ -775,18 +748,16 @@ def create_grid_from_crs( # Add extra metadata about the generated bounds if present bounds_msg = " with cell bounds" if add_bounds else "" - update_history( - ds, f"Created grid from CRS {crs} using xregrid ({backend}){bounds_msg}." - ) + update_history(ds, f"Created grid from CRS {crs} using xregrid ({backend}){bounds_msg}.") if chunks is not None: ds = ds.chunk(chunks) return ds def create_grid_from_ioapi( - metadata: Dict[str, Any], + metadata: dict[str, Any], add_bounds: bool = True, - chunks: Optional[Union[int, Dict[str, int]]] = None, + chunks: int | dict[str, int] | None = None, ) -> xr.Dataset: """ Create a structured grid dataset from IOAPI-compliant metadata. @@ -834,49 +805,25 @@ def create_grid_from_ioapi( if gdtyp == 1: # Lat-Lon crs = "EPSG:4326" elif gdtyp == 2: # Lambert Conformal - crs = ( - f"+proj=lcc +lat_1={p_alp} +lat_2={p_bet} +lat_0={ycent} " - f"+lon_0={xcent} +x_0=0 +y_0=0 +datum=WGS84 +units=m +no_defs" - ) + crs = f"+proj=lcc +lat_1={p_alp} +lat_2={p_bet} +lat_0={ycent} +lon_0={xcent} +x_0=0 +y_0=0 +datum=WGS84 +units=m +no_defs" elif gdtyp == 3: # Mercator - crs = ( - f"+proj=merc +lat_ts={p_alp} +lon_0={xcent} +lat_0={ycent} " - f"+x_0=0 +y_0=0 +datum=WGS84 +units=m +no_defs" - ) + crs = f"+proj=merc +lat_ts={p_alp} +lon_0={xcent} +lat_0={ycent} +x_0=0 +y_0=0 +datum=WGS84 +units=m +no_defs" elif gdtyp == 4: # Stereographic - crs = ( - f"+proj=stere +lat_ts={p_alp} +lat_0={ycent} +lon_0={xcent} " - f"+x_0=0 +y_0=0 +datum=WGS84 +units=m +no_defs" - ) + crs = f"+proj=stere +lat_ts={p_alp} +lat_0={ycent} +lon_0={xcent} +x_0=0 +y_0=0 +datum=WGS84 +units=m +no_defs" elif gdtyp == 5: # UTM crs = f"+proj=utm +zone={int(p_alp)} +datum=WGS84 +units=m +no_defs" elif gdtyp == 6: # Polar Stereographic # lat_0 determined by p_alp (1.0 for North, -1.0 for South) lat_0 = 90.0 if p_alp > 0 else -90.0 - crs = ( - f"+proj=stere +lat_0={lat_0} +lat_ts={p_bet} +lon_0={xcent} " - f"+x_0=0 +y_0=0 +datum=WGS84 +units=m +no_defs" - ) + crs = f"+proj=stere +lat_0={lat_0} +lat_ts={p_bet} +lon_0={xcent} +x_0=0 +y_0=0 +datum=WGS84 +units=m +no_defs" elif gdtyp == 7: # Equatorial Mercator - crs = ( - f"+proj=merc +lat_ts={p_alp} +lon_0={xcent} +lat_0=0 " - f"+x_0=0 +y_0=0 +datum=WGS84 +units=m +no_defs" - ) + crs = f"+proj=merc +lat_ts={p_alp} +lon_0={xcent} +lat_0=0 +x_0=0 +y_0=0 +datum=WGS84 +units=m +no_defs" elif gdtyp == 8: # Transverse Mercator - crs = ( - f"+proj=tmerc +lat_0={ycent} +k={p_bet} +lon_0={xcent} " - f"+x_0=0 +y_0=0 +datum=WGS84 +units=m +no_defs" - ) + crs = f"+proj=tmerc +lat_0={ycent} +k={p_bet} +lon_0={xcent} +x_0=0 +y_0=0 +datum=WGS84 +units=m +no_defs" elif gdtyp == 9: # Albers Equal Area - crs = ( - f"+proj=aea +lat_1={p_alp} +lat_2={p_bet} +lat_0={ycent} " - f"+lon_0={xcent} +x_0=0 +y_0=0 +datum=WGS84 +units=m +no_defs" - ) + crs = f"+proj=aea +lat_1={p_alp} +lat_2={p_bet} +lat_0={ycent} +lon_0={xcent} +x_0=0 +y_0=0 +datum=WGS84 +units=m +no_defs" elif gdtyp == 10: # Lambert Azimuthal Equal Area - crs = ( - f"+proj=laea +lat_0={ycent} +lon_0={xcent} " - f"+x_0=0 +y_0=0 +datum=WGS84 +units=m +no_defs" - ) + crs = f"+proj=laea +lat_0={ycent} +lon_0={xcent} +x_0=0 +y_0=0 +datum=WGS84 +units=m +no_defs" elif gdtyp == 13: # Sinusoidal crs = f"+proj=sinu +lon_0={xcent} +x_0=0 +y_0=0 +datum=WGS84 +units=m +no_defs" else: @@ -897,14 +844,14 @@ def create_grid_from_ioapi( def create_lcc_grid( - extent: Tuple[float, float, float, float], - res: Union[float, Tuple[float, float]], + extent: tuple[float, float, float, float], + res: float | tuple[float, float], lat_1: float, lat_2: float, lat_0: float, lon_0: float, add_bounds: bool = True, - chunks: Optional[Union[int, Dict[str, int]]] = None, + chunks: int | dict[str, int] | None = None, ) -> xr.Dataset: """ Create a structured grid dataset with a Lambert Conformal Conic (LCC) projection. @@ -935,10 +882,7 @@ def create_lcc_grid( xr.Dataset The grid dataset containing 'lat', 'lon' and projected coordinates 'x', 'y'. """ - crs = ( - f"+proj=lcc +lat_1={lat_1} +lat_2={lat_2} +lat_0={lat_0} " - f"+lon_0={lon_0} +x_0=0 +y_0=0 +datum=WGS84 +units=m +no_defs" - ) + crs = f"+proj=lcc +lat_1={lat_1} +lat_2={lat_2} +lat_0={lat_0} +lon_0={lon_0} +x_0=0 +y_0=0 +datum=WGS84 +units=m +no_defs" ds = create_grid_from_crs(crs, extent, res, add_bounds=add_bounds, chunks=chunks) # Backend detection for provenance @@ -953,12 +897,12 @@ def create_lcc_grid( def create_sinusoidal_grid( - extent: Tuple[float, float, float, float], - res: Union[float, Tuple[float, float], str], + extent: tuple[float, float, float, float], + res: float | tuple[float, float] | str, lon_0: float = 0.0, radius: float = 6371007.181, add_bounds: bool = True, - chunks: Optional[Union[int, Dict[str, int]]] = None, + chunks: int | dict[str, int] | None = None, ) -> xr.Dataset: """ Create a structured grid dataset with a Sinusoidal projection. @@ -1002,20 +946,18 @@ def create_sinusoidal_grid( crs = f"+proj=sinu +lon_0={lon_0} +x_0=0 +y_0=0 +R={radius} +units=m +no_defs" ds = create_grid_from_crs(crs, extent, res, add_bounds=add_bounds, chunks=chunks) - update_history( - ds, f"Created Sinusoidal grid (lon_0={lon_0}, R={radius}) using xregrid." - ) + update_history(ds, f"Created Sinusoidal grid (lon_0={lon_0}, R={radius}) using xregrid.") return ds def create_rotated_latlon_grid( - extent: Tuple[float, float, float, float], - res: Union[float, Tuple[float, float]], + extent: tuple[float, float, float, float], + res: float | tuple[float, float], grid_north_pole_lat: float, grid_north_pole_lon: float, north_pole_grid_lon: float = 0.0, add_bounds: bool = True, - chunks: Optional[Union[int, Dict[str, int]]] = None, + chunks: int | dict[str, int] | None = None, ) -> xr.Dataset: """ Create a structured grid dataset with a Rotated Pole (Rotated Lat-Lon) projection. @@ -1068,9 +1010,7 @@ def create_rotated_latlon_grid( core_chunks["y"] = core_chunks.pop("rlat") # Use core implementation - ds = create_grid_from_crs( - crs_obj, extent, res, add_bounds=add_bounds, chunks=core_chunks - ) + ds = create_grid_from_crs(crs_obj, extent, res, add_bounds=add_bounds, chunks=core_chunks) # Rename coordinates to standard rotated pole names (rlat, rlon) ds = ds.rename({"y": "rlat", "x": "rlon"}) @@ -1130,7 +1070,7 @@ def create_rotated_latlon_grid( return ds -def _get_min_max_lazy_aware(da_coord: xr.DataArray) -> Tuple[Any, Any, bool]: +def _get_min_max_lazy_aware(da_coord: xr.DataArray) -> tuple[Any, Any, bool]: """ Helper to get min/max from a coordinate DataArray efficiently. @@ -1157,11 +1097,7 @@ def _get_min_max_lazy_aware(da_coord: xr.DataArray) -> Tuple[Any, Any, bool]: return float(da_coord.min()), float(da_coord.max()), True # 2. Check if it's a dimension coordinate in indexes - if ( - da_coord.ndim == 1 - and da_coord.name in da_coord.dims - and da_coord.name in da_coord.indexes - ): + if da_coord.ndim == 1 and da_coord.name in da_coord.dims and da_coord.name in da_coord.indexes: idx = da_coord.indexes[da_coord.name] return float(idx.min()), float(idx.max()), True @@ -1184,12 +1120,12 @@ def _get_min_max_lazy_aware(da_coord: xr.DataArray) -> Tuple[Any, Any, bool]: def create_grid_like( - obj: Union[xr.DataArray, xr.Dataset], - res: Union[float, Tuple[float, float]], + obj: xr.DataArray | xr.Dataset, + res: float | tuple[float, float], add_bounds: bool = True, - chunks: Optional[Union[int, Dict[str, int]]] = None, - extent: Optional[Tuple[float, float, float, float]] = None, - crs: Optional[Union[str, int, Any]] = None, + chunks: int | dict[str, int] | None = None, + extent: tuple[float, float, float, float] | None = None, + crs: str | int | Any | None = None, ) -> xr.Dataset: """ Create a new grid dataset with the same extent and CRS as an existing object. @@ -1239,9 +1175,7 @@ def create_grid_like( history_msg_base += f"\nTemplate history:\n{obj.attrs['history']}" if extent is not None: - if crs_obj is None or ( - hasattr(crs_obj, "is_geographic") and crs_obj.is_geographic - ): + if crs_obj is None or (hasattr(crs_obj, "is_geographic") and crs_obj.is_geographic): # Lat-Lon return _create_rectilinear_grid( (extent[2], extent[3]), # lat_range @@ -1255,9 +1189,7 @@ def create_grid_like( ) else: # Projected - return create_grid_from_crs( - crs_obj, extent, (res_x, res_y), add_bounds=add_bounds, chunks=chunks - ) + return create_grid_from_crs(crs_obj, extent, (res_x, res_y), add_bounds=add_bounds, chunks=chunks) # Aero Optimization: Try to find extent in metadata before falling back to compute. if extent is None: @@ -1277,9 +1209,7 @@ def create_grid_like( pass if extent is not None: - if crs_obj is None or ( - hasattr(crs_obj, "is_geographic") and crs_obj.is_geographic - ): + if crs_obj is None or (hasattr(crs_obj, "is_geographic") and crs_obj.is_geographic): # Lat-Lon return _create_rectilinear_grid( (extent[2], extent[3]), # lat_range @@ -1293,9 +1223,7 @@ def create_grid_like( ) else: # Projected - return create_grid_from_crs( - crs_obj, extent, (res_x, res_y), add_bounds=add_bounds, chunks=chunks - ) + return create_grid_from_crs(crs_obj, extent, (res_x, res_y), add_bounds=add_bounds, chunks=chunks) # Discovery logic: we need min/max. We use batch compute if lazy to minimize roundtrips. # Aero-Optimization: Use Xarray indexes for 1D dimension coordinates to avoid hidden computes. @@ -1368,9 +1296,7 @@ def create_grid_like( y_max_val = float(results.get("y_max", y_max)) res_x_orig = float(results.get("res_x", 0)) - res_y_orig = float( - results.get("res_y", res_x_orig if res_x_orig else 0) - ) + res_y_orig = float(results.get("res_y", res_x_orig if res_x_orig else 0)) extent = ( x_min_val - res_x_orig / 2, @@ -1382,14 +1308,8 @@ def create_grid_like( x_min_val, x_max_val = float(x_min), float(x_max) y_min_val, y_max_val = float(y_min), float(y_max) - res_x_orig = ( - abs(float(x_da.diff(x_da.dims[0]).mean())) if x_da.size > 1 else 0 - ) - res_y_orig = ( - abs(float(y_da.diff(y_da.dims[0]).mean())) - if y_da.size > 1 - else res_x_orig - ) + res_x_orig = abs(float(x_da.diff(x_da.dims[0]).mean())) if x_da.size > 1 else 0 + res_y_orig = abs(float(y_da.diff(y_da.dims[0]).mean())) if y_da.size > 1 else res_x_orig extent = ( x_min_val - res_x_orig / 2, x_max_val + res_x_orig / 2, @@ -1401,9 +1321,7 @@ def create_grid_like( # Fallback to generic geographic if no CRS found crs_obj = "EPSG:4326" - return create_grid_from_crs( - crs_obj, extent, (res_x, res_y), add_bounds=add_bounds, chunks=chunks - ) + return create_grid_from_crs(crs_obj, extent, (res_x, res_y), add_bounds=add_bounds, chunks=chunks) except (KeyError, AttributeError, ValueError): pass @@ -1479,9 +1397,7 @@ def create_grid_like( lon_max_val = float(results.get("lon_max", lon_max)) res_lat_orig = float(results.get("res_lat", 0)) - res_lon_orig = float( - results.get("res_lon", res_lat_orig if res_lat_orig else 0) - ) + res_lon_orig = float(results.get("res_lon", res_lat_orig if res_lat_orig else 0)) lat_range = ( lat_min_val - res_lat_orig / 2, @@ -1495,16 +1411,8 @@ def create_grid_like( lat_min_val, lat_max_val = float(lat_min), float(lat_max) lon_min_val, lon_max_val = float(lon_min), float(lon_max) - res_lat_orig = ( - abs(float(lat_da.diff(lat_da.dims[0]).mean())) - if lat_da.size > 1 - else 0 - ) - res_lon_orig = ( - abs(float(lon_da.diff(lon_da.dims[-1]).mean())) - if lon_da.size > 1 - else res_lat_orig - ) + res_lat_orig = abs(float(lat_da.diff(lat_da.dims[0]).mean())) if lat_da.size > 1 else 0 + res_lon_orig = abs(float(lon_da.diff(lon_da.dims[-1]).mean())) if lon_da.size > 1 else res_lat_orig lat_range = ( lat_min_val - res_lat_orig / 2, lat_max_val + res_lat_orig / 2, @@ -1524,18 +1432,15 @@ def create_grid_like( crs=crs_obj.to_wkt() if crs_obj else "EPSG:4326", history_msg=history_msg_base, ) - except (KeyError, AttributeError, ValueError): - raise ValueError( - "Could not detect spatial coordinates (latitude/longitude or " - "projection_x/y) in input object." - ) + except (KeyError, AttributeError, ValueError) as err: + raise ValueError("Could not detect spatial coordinates (latitude/longitude or projection_x/y) in input object.") from err def create_mesh_from_coords( - x: Union[np.ndarray, xr.DataArray], - y: Union[np.ndarray, xr.DataArray], - crs: Union[str, int, Any], - chunks: Optional[Union[int, Dict[str, int]]] = None, + x: np.ndarray | xr.DataArray, + y: np.ndarray | xr.DataArray, + crs: str | int | Any, + chunks: int | dict[str, int] | None = None, ) -> xr.Dataset: """ Create an unstructured mesh dataset from coordinates and a CRS. @@ -1558,10 +1463,7 @@ def create_mesh_from_coords( The mesh dataset containing 'lat', 'lon' and 'x', 'y' as 1D arrays. """ if pyproj is None: - raise ImportError( - "pyproj is required for create_mesh_from_coords. " - "Install it with `pip install pyproj`." - ) + raise ImportError("pyproj is required for create_mesh_from_coords. Install it with `pip install pyproj`.") crs_obj = pyproj.CRS(crs) # Force n_pts dimension to avoid alignment/broadcasting issues in apply_ufunc @@ -1676,8 +1578,8 @@ def create_mesh_from_coords( def get_rdhpcs_cluster( - machine: Optional[str] = None, - account: Optional[str] = None, + machine: str | None = None, + account: str | None = None, **kwargs: Any, ) -> Any: """ @@ -1703,11 +1605,8 @@ def get_rdhpcs_cluster( """ try: from dask_jobqueue import SLURMCluster - except ImportError: - raise ImportError( - "dask-jobqueue is required for get_rdhpcs_cluster. " - "Install it with `pip install dask-jobqueue`." - ) + except ImportError as err: + raise ImportError("dask-jobqueue is required for get_rdhpcs_cluster. Install it with `pip install dask-jobqueue`.") from err hostname = socket.gethostname() if machine is None: @@ -1722,8 +1621,7 @@ def get_rdhpcs_cluster( machine = "gaea-c5" else: raise ValueError( - f"Could not detect NOAA RDHPCS machine from hostname '{hostname}'. " - "Please specify 'machine' explicitly." + f"Could not detect NOAA RDHPCS machine from hostname '{hostname}'. Please specify 'machine' explicitly." ) defaults = { @@ -1779,19 +1677,17 @@ def get_rdhpcs_cluster( if defaults["account"] is None: import warnings - warnings.warn( - "No SLURM account specified. Please provide 'account' or set SACCOUNT environment variable." - ) + warnings.warn("No SLURM account specified. Please provide 'account' or set SACCOUNT environment variable.") return SLURMCluster(**defaults) def spatial_slice( - obj: Union[xr.DataArray, xr.Dataset], - extent: Tuple[float, float, float, float], - crs: Optional[Union[str, int, Any]] = None, + obj: xr.DataArray | xr.Dataset, + extent: tuple[float, float, float, float], + crs: str | int | Any | None = None, buffer: float = 0.0, -) -> Union[xr.DataArray, xr.Dataset]: +) -> xr.DataArray | xr.Dataset: """ Slice an xarray object to a spatial extent, handling longitude wrapping. @@ -1829,11 +1725,10 @@ def spatial_slice( x_da = obj.cf["projection_x_coordinate"] y_da = obj.cf["projection_y_coordinate"] is_geographic = False - except (KeyError, AttributeError): + except (KeyError, AttributeError) as err: raise ValueError( - "Could not detect spatial coordinates (lat/lon or x/y) for slicing. " - "Ensure your data has CF-compliant coordinates." - ) + "Could not detect spatial coordinates (lat/lon or x/y) for slicing. Ensure your data has CF-compliant coordinates." + ) from err else: x_da, y_da = lon_da, lat_da is_geographic = True @@ -1841,10 +1736,7 @@ def spatial_slice( # 2. CRS Transformation if crs is not None: if pyproj is None: - raise ImportError( - "pyproj is required for CRS-aware slicing. " - "Install it with `pip install pyproj`." - ) + raise ImportError("pyproj is required for CRS-aware slicing. Install it with `pip install pyproj`.") target_crs = get_crs_info(obj) or pyproj.CRS("EPSG:4326") transformer = pyproj.Transformer.from_crs(crs, target_crs, always_xy=True) @@ -1985,7 +1877,7 @@ def unstructured_to_scrip(ds: xr.Dataset) -> xr.Dataset: orig_cell_index, ) = _get_unstructured_mesh_info(ds, method="conservative") except Exception as e: - raise ValueError(f"Failed to extract unstructured connectivity: {e}") + raise ValueError(f"Failed to extract unstructured connectivity: {e}") from e # 3. Reshape connectivity to SCRIP-style (N, 3 for triangles) # _get_unstructured_mesh_info always triangulates. @@ -2018,16 +1910,12 @@ def unstructured_to_scrip(ds: xr.Dataset) -> xr.Dataset: coords={ "lat": ( ["grid_size"], - lat_c.data[orig_cell_index] - if orig_cell_index is not None - else lat_c.data, + lat_c.data[orig_cell_index] if orig_cell_index is not None else lat_c.data, lat_attrs, ), "lon": ( ["grid_size"], - lon_c.data[orig_cell_index] - if orig_cell_index is not None - else lon_c.data, + lon_c.data[orig_cell_index] if orig_cell_index is not None else lon_c.data, lon_attrs, ), "lat_b": ( diff --git a/src/xregrid/viz.py b/src/xregrid/viz.py index 06758e3..c319fe4 100644 --- a/src/xregrid/viz.py +++ b/src/xregrid/viz.py @@ -1,11 +1,11 @@ from __future__ import annotations import warnings -from typing import TYPE_CHECKING, Any, Optional +from typing import TYPE_CHECKING, Any import xarray as xr -from xregrid.utils import get_crs_info, _find_coord +from xregrid.utils import _find_coord, get_crs_info if TYPE_CHECKING: from xregrid.regridder import Regridder @@ -26,8 +26,8 @@ pyproj = None try: - import hvplot.xarray # noqa: F401 import holoviews as hv + import hvplot.xarray # noqa: F401 except ImportError: hvplot = None hv = None @@ -39,7 +39,7 @@ def plot_static( da: xr.DataArray, projection: Any = None, transform: Any = None, - title: Optional[str] = None, + title: str | None = None, **kwargs: Any, ) -> Any: """ @@ -69,10 +69,7 @@ def plot_static( If matplotlib is not installed. """ if plt is None: - raise ImportError( - "Matplotlib is required for plot_static. " - "Install it with `pip install matplotlib`." - ) + raise ImportError("Matplotlib is required for plot_static. Install it with `pip install matplotlib`.") # Handle axes and faceting early to avoid multiple 'ax' arguments ax = kwargs.pop("ax", None) @@ -108,7 +105,7 @@ def plot_static( extra_dims = [d for d in da.dims if d not in spatial_dims and d not in facet_dims] if extra_dims: - first_slice = {d: 0 for d in extra_dims} + first_slice = dict.fromkeys(extra_dims, 0) warnings.warn( f"DataArray has {da.ndim} dimensions, but only 2 spatial dimensions " f"(plus optional faceting) are supported for static plots. " @@ -161,9 +158,7 @@ def plot_static( if ax is not None: if is_faceted: - warnings.warn( - "Providing an 'ax' with faceting ('col' or 'row') is not supported by xarray and will be ignored." - ) + warnings.warn("Providing an 'ax' with faceting ('col' or 'row') is not supported by xarray and will be ignored.") ax = None else: # Ensure the existing axes is a GeoAxes if we are using cartopy @@ -262,9 +257,7 @@ def plot( elif mode == "interactive": return plot_interactive(da, **kwargs) else: - raise ValueError( - f"Unknown plotting mode: '{mode}'. Must be 'static' or 'interactive'." - ) + raise ValueError(f"Unknown plotting mode: '{mode}'. Must be 'static' or 'interactive'.") def plot_interactive( @@ -298,10 +291,7 @@ def plot_interactive( If HvPlot is not installed. """ if not hvplot: - raise ImportError( - "HvPlot is required for plot_interactive. " - "Install it with `pip install hvplot`." - ) + raise ImportError("HvPlot is required for plot_interactive. Install it with `pip install hvplot`.") # Automated CRS discovery for Track B (Interactive) # This ensures "No Ambiguous Plots" even in exploratory mode. @@ -325,7 +315,7 @@ def plot_interactive( def plot_diagnostics( - regridder: "Regridder", + regridder: Regridder, projection: Any = None, **kwargs: Any, ) -> Any: @@ -407,9 +397,9 @@ def plot_diagnostics( def plot_diagnostics_interactive( - regridder: "Regridder", + regridder: Regridder, rasterize: bool = True, - title: Optional[str] = None, + title: str | None = None, **kwargs: Any, ) -> Any: """ @@ -441,21 +431,16 @@ def plot_diagnostics_interactive( """ if not hvplot or hv is None: raise ImportError( - "HvPlot and HoloViews are required for plot_diagnostics_interactive. " - "Install them with `pip install hvplot holoviews`." + "HvPlot and HoloViews are required for plot_diagnostics_interactive. Install them with `pip install hvplot holoviews`." ) ds_diag = regridder.diagnostics() # 1. Weight Sum Plot - p_sum = ds_diag.weight_sum.hvplot( - rasterize=rasterize, cmap="viridis", title="Weight Sum", **kwargs - ) + p_sum = ds_diag.weight_sum.hvplot(rasterize=rasterize, cmap="viridis", title="Weight Sum", **kwargs) # 2. Unmapped Mask Plot - p_mask = ds_diag.unmapped_mask.hvplot( - rasterize=rasterize, cmap="Reds", title="Unmapped Mask (1=Unmapped)", **kwargs - ) + p_mask = ds_diag.unmapped_mask.hvplot(rasterize=rasterize, cmap="Reds", title="Unmapped Mask (1=Unmapped)", **kwargs) layout = (p_sum + p_mask).cols(2) @@ -470,12 +455,12 @@ def plot_diagnostics_interactive( def plot_comparison( da_src: xr.DataArray, da_tgt: xr.DataArray, - regridder: Optional[Any] = None, + regridder: Any | None = None, projection: Any = None, transform: Any = None, cmap: str = "viridis", diff_cmap: str = "RdBu_r", - title: Optional[str] = None, + title: str | None = None, **kwargs: Any, ) -> Any: """ @@ -588,11 +573,11 @@ def plot_comparison( def plot_comparison_interactive( da_src: xr.DataArray, da_tgt: xr.DataArray, - regridder: Optional[Any] = None, + regridder: Any | None = None, rasterize: bool = True, cmap: str = "viridis", diff_cmap: str = "RdBu_r", - title: Optional[str] = None, + title: str | None = None, **kwargs: Any, ) -> Any: """ @@ -632,8 +617,7 @@ def plot_comparison_interactive( """ if not hvplot or hv is None: raise ImportError( - "HvPlot and HoloViews are required for plot_comparison_interactive. " - "Install them with `pip install hvplot holoviews`." + "HvPlot and HoloViews are required for plot_comparison_interactive. Install them with `pip install hvplot holoviews`." ) # 1. Source Plot @@ -669,7 +653,7 @@ def plot_comparison_interactive( def plot_weights( - regridder: "Regridder", + regridder: Regridder, row_idx: int, mode: str = "static", **kwargs: Any, @@ -701,17 +685,13 @@ def plot_weights( return _plot_weights_static(regridder, row_idx, **kwargs) elif mode == "interactive": rasterize = kwargs.pop("rasterize", True) - return plot_weights_interactive( - regridder, row_idx, rasterize=rasterize, **kwargs - ) + return plot_weights_interactive(regridder, row_idx, rasterize=rasterize, **kwargs) else: - raise ValueError( - f"Unknown plotting mode: '{mode}'. Must be 'static' or 'interactive'." - ) + raise ValueError(f"Unknown plotting mode: '{mode}'. Must be 'static' or 'interactive'.") def _plot_weights_static( - regridder: "Regridder", + regridder: Regridder, row_idx: int, **kwargs: Any, ) -> Any: @@ -733,13 +713,11 @@ def _plot_weights_static( The plot object. """ da_weights = _get_weight_row_da(regridder, row_idx) - return plot_static( - da_weights, title=f"Weights for Destination Point {row_idx}", **kwargs - ) + return plot_static(da_weights, title=f"Weights for Destination Point {row_idx}", **kwargs) def plot_weights_interactive( - regridder: "Regridder", + regridder: Regridder, row_idx: int, rasterize: bool = True, **kwargs: Any, @@ -772,7 +750,7 @@ def plot_weights_interactive( ) -def _get_weight_row_da(regridder: "Regridder", row_idx: int) -> xr.DataArray: +def _get_weight_row_da(regridder: Regridder, row_idx: int) -> xr.DataArray: """ Extract a single weight row as a DataArray, optimized for remote weights. @@ -792,9 +770,7 @@ def _get_weight_row_da(regridder: "Regridder", row_idx: int) -> xr.DataArray: # Optimized Distributed Path: extract row on cluster from .parallel import _get_weight_row_task - row = regridder._dask_client.submit( - _get_weight_row_task, regridder._weights_matrix, row_idx - ).result() + row = regridder._dask_client.submit(_get_weight_row_task, regridder._weights_matrix, row_idx).result() else: # Eager Path matrix = regridder.weights @@ -804,19 +780,13 @@ def _get_weight_row_da(regridder: "Regridder", row_idx: int) -> xr.DataArray: coords = { c: regridder.source_grid_ds.coords[c] for c in regridder.source_grid_ds.coords - if regridder._dims_source is not None - and set(regridder.source_grid_ds.coords[c].dims).issubset( - set(regridder._dims_source) - ) + if regridder._dims_source is not None and set(regridder.source_grid_ds.coords[c].dims).issubset(set(regridder._dims_source)) } # Include topology/mapping from source grid for v in regridder.source_grid_ds.data_vars: var_obj = regridder.source_grid_ds[v] - if ( - var_obj.attrs.get("cf_role") == "mesh_topology" - or "grid_mapping_name" in var_obj.attrs - ): + if var_obj.attrs.get("cf_role") == "mesh_topology" or "grid_mapping_name" in var_obj.attrs: coords[v] = var_obj da_weights = xr.DataArray( diff --git a/tests/conftest.py b/tests/conftest.py index bcb60bf..81d761f 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -1,4 +1,5 @@ import sys + import numpy as np @@ -26,11 +27,7 @@ def get_coords(self, coord_dim, staggerloc=0): if staggerloc == 0: # CENTER shape = list(self.max_index) elif staggerloc == 1: # CORNER - if ( - self.num_peri_dims - and self.num_peri_dims > 0 - and self.periodic_dim is not None - ): + if self.num_peri_dims and self.num_peri_dims > 0 and self.periodic_dim is not None: shape[self.periodic_dim] -= 1 self.coords[key] = np.zeros(tuple(shape)) return self.coords[key] @@ -42,11 +39,7 @@ def get_item(self, item, staggerloc=0): if staggerloc == 0: # CENTER shape = list(self.max_index) elif staggerloc == 1: # CORNER - if ( - self.num_peri_dims - and self.num_peri_dims > 0 - and self.periodic_dim is not None - ): + if self.num_peri_dims and self.num_peri_dims > 0 and self.periodic_dim is not None: shape[self.periodic_dim] -= 1 self.items[key] = np.zeros(tuple(shape)) return self.items[key] diff --git a/tests/test_aero_hardening.py b/tests/test_aero_hardening.py index 38bde69..6289045 100644 --- a/tests/test_aero_hardening.py +++ b/tests/test_aero_hardening.py @@ -3,8 +3,9 @@ import numpy as np import pytest import xarray as xr -from xregrid.utils import is_lazy, is_dask, is_cubed, create_grid_like + from xregrid import Regridder +from xregrid.utils import create_grid_like, is_cubed, is_dask, is_lazy def test_backend_utilities(): @@ -54,9 +55,7 @@ def test_create_grid_like_hardening(): # (The warning is issued only if it falls back to compute) for warning in record: if "Triggering hidden compute" in str(warning.message): - pytest.fail( - "create_grid_like triggered a compute despite metadata presence" - ) + pytest.fail("create_grid_like triggered a compute despite metadata presence") assert grid.lat.size == 90 assert grid.lon.size == 180 diff --git a/tests/test_aero_hardening_verification.py b/tests/test_aero_hardening_verification.py index 7935843..b8afc9c 100644 --- a/tests/test_aero_hardening_verification.py +++ b/tests/test_aero_hardening_verification.py @@ -1,6 +1,7 @@ import numpy as np import pytest import xarray as xr + from xregrid import Regridder, create_global_grid @@ -30,9 +31,7 @@ def test_all_nan_input_robustness(): tgt = create_global_grid(5, 5) data = np.full((18, 36), np.nan) - da = xr.DataArray( - data, dims=("lat", "lon"), coords={"lat": src.lat, "lon": src.lon} - ) + da = xr.DataArray(data, dims=("lat", "lon"), coords={"lat": src.lat, "lon": src.lon}) regridder = Regridder(src, tgt, skipna=True) res = regridder(da) diff --git a/tests/test_aero_optimization.py b/tests/test_aero_optimization.py index d0eea4c..84e3ebb 100644 --- a/tests/test_aero_optimization.py +++ b/tests/test_aero_optimization.py @@ -1,9 +1,10 @@ from __future__ import annotations +import dask import numpy as np import xarray as xr -import dask -from xregrid.utils import create_grid_like, create_global_grid + +from xregrid.utils import create_global_grid, create_grid_like class ComputeCounter(dask.callbacks.Callback): diff --git a/tests/test_aero_periodicity.py b/tests/test_aero_periodicity.py index 03ac97e..a62c4b4 100644 --- a/tests/test_aero_periodicity.py +++ b/tests/test_aero_periodicity.py @@ -1,10 +1,11 @@ from __future__ import annotations +import dask.array as da import numpy as np -import xarray as xr import pytest -import dask.array as da +import xarray as xr from dask.callbacks import Callback + from xregrid import Regridder @@ -48,12 +49,8 @@ def test_regridder_periodicity_lazy(): lon_2d, lat_2d = np.meshgrid(lon_raw, lat_raw) # Chunk it to make it Dask-backed - lon_da = xr.DataArray( - da.from_array(lon_2d, chunks=(50, 50)), dims=("y", "x"), name="lon" - ) - lat_da = xr.DataArray( - da.from_array(lat_2d, chunks=(50, 50)), dims=("y", "x"), name="lat" - ) + lon_da = xr.DataArray(da.from_array(lon_2d, chunks=(50, 50)), dims=("y", "x"), name="lon") + lat_da = xr.DataArray(da.from_array(lat_2d, chunks=(50, 50)), dims=("y", "x"), name="lat") ds_src = xr.Dataset({"lon": lon_da, "lat": lat_da}) ds_src.lon.attrs["units"] = "degrees_east" @@ -68,9 +65,7 @@ def test_regridder_periodicity_lazy(): ) # Case 1: Lazy curvilinear triggers a warning - with pytest.warns( - UserWarning, match="Triggering hidden compute in _detect_periodicity" - ): + with pytest.warns(UserWarning, match="Triggering hidden compute in _detect_periodicity"): regridder = Regridder(ds_src, ds_tgt, method="bilinear") assert regridder.periodic is True @@ -86,9 +81,7 @@ def test_regridder_periodicity_lazy(): with counter: regridder = Regridder(ds_src_meta, ds_tgt, method="bilinear") assert regridder.periodic is True - assert ( - counter.count == 0 - ), f"Metadata-based detection triggered {counter.count} computes" + assert counter.count == 0, f"Metadata-based detection triggered {counter.count} computes" # Case 3: Explicit periodicity avoids compute and warning with patch("xregrid.regridder.Regridder._generate_weights"): @@ -96,9 +89,7 @@ def test_regridder_periodicity_lazy(): with counter: regridder = Regridder(ds_src, ds_tgt, method="bilinear", periodic=True) assert regridder.periodic is True - assert ( - counter.count == 0 - ), f"Explicit periodicity triggered {counter.count} computes" + assert counter.count == 0, f"Explicit periodicity triggered {counter.count} computes" if __name__ == "__main__": diff --git a/tests/test_aero_protocol.py b/tests/test_aero_protocol.py index 0fb5b55..f204f94 100644 --- a/tests/test_aero_protocol.py +++ b/tests/test_aero_protocol.py @@ -3,6 +3,7 @@ import numpy as np import pytest import xarray as xr + from xregrid import Regridder from xregrid.utils import create_global_grid @@ -19,9 +20,7 @@ def test_aero_protocol_equivalence(): # Create source data with some pattern and NaNs data = np.sin(np.deg2rad(ds_src.lat)) * np.cos(np.deg2rad(ds_src.lon)) # Coordinates for the DataArray should only include relevant dimensions - coords = { - k: v for k, v in ds_src.coords.items() if set(v.dims).issubset({"lat", "lon"}) - } + coords = {k: v for k, v in ds_src.coords.items() if set(v.dims).issubset({"lat", "lon"})} da_numpy = xr.DataArray(data, coords=coords, dims=("lat", "lon"), name="test_data") # Add some NaNs to test skipna @@ -62,9 +61,7 @@ def test_non_spatial_preservation(): # Create dataset with various dimension names ds = xr.Dataset( - data_vars={ - "temp": (("time", "lev", "lat", "lon"), np.random.rand(2, 5, 10, 20)) - }, + data_vars={"temp": (("time", "lev", "lat", "lon"), np.random.rand(2, 5, 10, 20))}, coords={ "time": np.arange(2), "lev": np.arange(5), diff --git a/tests/test_backend_agnostic.py b/tests/test_backend_agnostic.py index 8b2a807..3a74160 100644 --- a/tests/test_backend_agnostic.py +++ b/tests/test_backend_agnostic.py @@ -1,6 +1,7 @@ from __future__ import annotations import xarray as xr + from xregrid.utils import create_global_grid, create_grid_from_crs, is_dask diff --git a/tests/test_periodicity_lazy.py b/tests/test_periodicity_lazy.py index 22563c8..25fb779 100644 --- a/tests/test_periodicity_lazy.py +++ b/tests/test_periodicity_lazy.py @@ -1,8 +1,9 @@ from __future__ import annotations +import dask.array as da import numpy as np import xarray as xr -import dask.array as da + from xregrid.regridder import Regridder from xregrid.utils import create_global_grid diff --git a/tests/test_property_based.py b/tests/test_property_based.py index b8020ee..3fe7696 100644 --- a/tests/test_property_based.py +++ b/tests/test_property_based.py @@ -3,27 +3,28 @@ import os import uuid import warnings + import numpy as np import xarray as xr -from typing import Tuple -from hypothesis import given, strategies as st, settings, HealthCheck, assume +from hypothesis import HealthCheck, assume, given, settings +from hypothesis import strategies as st from xregrid import ( Regridder, create_global_grid, - create_regional_grid, create_grid_from_crs, + create_grid_from_ioapi, create_grid_like, + create_regional_grid, create_rotated_latlon_grid, create_sinusoidal_grid, - create_grid_from_ioapi, ) from xregrid.utils import ( _get_min_max_lazy_aware, - is_lazy, + create_lcc_grid, is_dask, + is_lazy, spatial_slice, - create_lcc_grid, ) # Configure hypothesis to have higher deadlines since ESMF is involved @@ -37,9 +38,7 @@ add_bounds=st.booleans(), ) @settings(suppress_health_check=[HealthCheck.filter_too_much], max_examples=15) -def test_create_global_grid_properties( - res_lat: float, res_lon: float, add_bounds: bool -) -> None: +def test_create_global_grid_properties(res_lat: float, res_lon: float, add_bounds: bool) -> None: """ Test properties of create_global_grid using Hypothesis. @@ -58,9 +57,7 @@ def test_create_global_grid_properties( Whether to add cell boundaries. """ # Create eager grid - ds_eager = create_global_grid( - res_lat=res_lat, res_lon=res_lon, add_bounds=add_bounds - ) + ds_eager = create_global_grid(res_lat=res_lat, res_lon=res_lon, add_bounds=add_bounds) assert "lat" in ds_eager.coords assert "lon" in ds_eager.coords @@ -84,9 +81,7 @@ def test_create_global_grid_properties( assert ds_eager.lon_b.shape == (lon.size, 2) # Test eager vs lazy parity - ds_lazy = create_global_grid( - res_lat=res_lat, res_lon=res_lon, add_bounds=add_bounds, chunks=5 - ) + ds_lazy = create_global_grid(res_lat=res_lat, res_lon=res_lon, add_bounds=add_bounds, chunks=5) if add_bounds: # Bounds are 2D and not dimension coordinates, so they remain lazy @@ -253,9 +248,7 @@ def test_get_min_max_lazy_aware_properties(ndim: int, size_y: int, size_x: int) assert np.isclose(min_v, -20.0) assert np.isclose(max_v, 20.0) - da_lazy = xr.DataArray( - da.from_array(x, chunks=(3, 3)), dims=("y", "x"), name="lon" - ) + da_lazy = xr.DataArray(da.from_array(x, chunks=(3, 3)), dims=("y", "x"), name="lon") min_v2, max_v2, is_eager2 = _get_min_max_lazy_aware(da_lazy) assert not is_eager2 assert np.isclose(float(min_v2.compute()), -20.0) @@ -269,9 +262,7 @@ def test_get_min_max_lazy_aware_properties(ndim: int, size_y: int, size_x: int) skipna=st.booleans(), has_time=st.booleans(), ) -@settings( - max_examples=10, suppress_health_check=[HealthCheck.filter_too_much], deadline=None -) +@settings(max_examples=10, suppress_health_check=[HealthCheck.filter_too_much], deadline=None) def test_regridder_real_esmf_properties( res_src: float, res_tgt: float, @@ -354,11 +345,7 @@ def test_regridder_real_esmf_properties( # Eager application res_eager = regridder(da_src) - expected_shape = ( - (2, ds_tgt.lat.size, ds_tgt.lon.size) - if has_time - else (ds_tgt.lat.size, ds_tgt.lon.size) - ) + expected_shape = (2, ds_tgt.lat.size, ds_tgt.lon.size) if has_time else (ds_tgt.lat.size, ds_tgt.lon.size) assert res_eager.shape == expected_shape assert "lat" in res_eager.dims assert "lon" in res_eager.dims @@ -380,9 +367,7 @@ def test_regridder_real_esmf_properties( res_lazy_computed = res_lazy.compute() # Parity check - np.testing.assert_allclose( - res_eager.values, res_lazy_computed.values, equal_nan=True - ) + np.testing.assert_allclose(res_eager.values, res_lazy_computed.values, equal_nan=True) # History check assert "history" in res_eager.attrs @@ -403,9 +388,7 @@ def test_regridder_real_esmf_properties( add_bounds=st.booleans(), ) @settings(max_examples=10, suppress_health_check=[HealthCheck.filter_too_much]) -def test_create_rotated_latlon_grid_properties( - pole_lat: float, pole_lon: float, res: float, add_bounds: bool -) -> None: +def test_create_rotated_latlon_grid_properties(pole_lat: float, pole_lon: float, res: float, add_bounds: bool) -> None: """ Test properties of create_rotated_latlon_grid using Hypothesis. @@ -413,9 +396,7 @@ def test_create_rotated_latlon_grid_properties( geographic coordinates lat/lon are computed and eager/lazy parity is preserved. """ extent = (-20.0, 20.0, -15.0, 15.0) - ds_eager = create_rotated_latlon_grid( - extent, res, pole_lat, pole_lon, add_bounds=add_bounds - ) + ds_eager = create_rotated_latlon_grid(extent, res, pole_lat, pole_lon, add_bounds=add_bounds) assert "rlat" in ds_eager.coords assert "rlon" in ds_eager.coords @@ -433,9 +414,7 @@ def test_create_rotated_latlon_grid_properties( assert "lon_b" in ds_eager.coords # Lazy parity check - ds_lazy = create_rotated_latlon_grid( - extent, res, pole_lat, pole_lon, add_bounds=add_bounds, chunks=5 - ) + ds_lazy = create_rotated_latlon_grid(extent, res, pole_lat, pole_lon, add_bounds=add_bounds, chunks=5) assert is_lazy(ds_lazy.lat) xr.testing.assert_allclose(ds_eager, ds_lazy.compute()) @@ -483,9 +462,7 @@ def test_create_sinusoidal_grid_properties(res_alias: str, add_bounds: bool) -> add_bounds=st.booleans(), ) @settings(max_examples=10, suppress_health_check=[HealthCheck.filter_too_much]) -def test_create_lcc_grid_properties( - lat_1: float, lat_2: float, lat_0: float, lon_0: float, add_bounds: bool -) -> None: +def test_create_lcc_grid_properties(lat_1: float, lat_2: float, lat_0: float, lon_0: float, add_bounds: bool) -> None: """ Test properties of Lambert Conformal Conic grid helper using Hypothesis. """ @@ -495,9 +472,7 @@ def test_create_lcc_grid_properties( extent = (-100000.0, 100000.0, -100000.0, 100000.0) res = 50000.0 - ds_eager = create_lcc_grid( - extent, res, lat_1, lat_2, lat_0, lon_0, add_bounds=add_bounds - ) + ds_eager = create_lcc_grid(extent, res, lat_1, lat_2, lat_0, lon_0, add_bounds=add_bounds) assert "x" in ds_eager.coords assert "y" in ds_eager.coords @@ -509,9 +484,7 @@ def test_create_lcc_grid_properties( assert "y_b" in ds_eager.coords # Lazy parity - ds_lazy = create_lcc_grid( - extent, res, lat_1, lat_2, lat_0, lon_0, add_bounds=add_bounds, chunks=5 - ) + ds_lazy = create_lcc_grid(extent, res, lat_1, lat_2, lat_0, lon_0, add_bounds=add_bounds, chunks=5) assert is_lazy(ds_lazy.lat) xr.testing.assert_allclose(ds_eager, ds_lazy.compute()) @@ -523,9 +496,7 @@ def test_create_lcc_grid_properties( add_bounds=st.booleans(), ) @settings(max_examples=10, suppress_health_check=[HealthCheck.filter_too_much]) -def test_create_grid_from_ioapi_properties( - gdtyp: int, ncols: int, nrows: int, add_bounds: bool -) -> None: +def test_create_grid_from_ioapi_properties(gdtyp: int, ncols: int, nrows: int, add_bounds: bool) -> None: """ Test properties of create_grid_from_ioapi using Hypothesis. @@ -575,7 +546,7 @@ def test_create_grid_from_ioapi_properties( ) @settings(max_examples=10, suppress_health_check=[HealthCheck.filter_too_much]) def test_spatial_slice_wrapping_properties( - extent: Tuple[float, float, float, float], + extent: tuple[float, float, float, float], buffer: float, is_lazy_data: bool, ) -> None: @@ -678,12 +649,8 @@ def test_regrid_unstructured_to_unstructured_properties( Test properties of unstructured to unstructured regridding. """ # Create source and target unstructured meshes - ds_src = generate_unstructured_polygon_mesh( - n_cells_src, n_corners, center_lon_start=0.5, step_lon=1.0 - ) - ds_tgt = generate_unstructured_polygon_mesh( - n_cells_tgt, n_corners, center_lon_start=0.5, step_lon=1.0 - ) + ds_src = generate_unstructured_polygon_mesh(n_cells_src, n_corners, center_lon_start=0.5, step_lon=1.0) + ds_tgt = generate_unstructured_polygon_mesh(n_cells_tgt, n_corners, center_lon_start=0.5, step_lon=1.0) # Input constant field of ones data = np.ones(n_cells_src) @@ -721,9 +688,7 @@ def test_regrid_unstructured_to_unstructured_properties( # All points overlap in identical span, so bilinear/nearest on constant should return 1.0 (or close) # where overlap is successful. - np.testing.assert_allclose( - res_comp.values, np.ones(n_cells_tgt), rtol=1e-3, atol=1e-3 - ) + np.testing.assert_allclose(res_comp.values, np.ones(n_cells_tgt), rtol=1e-3, atol=1e-3) finally: if os.path.exists(weights_file): try: @@ -748,9 +713,7 @@ def test_regrid_unstructured_to_rectilinear_properties( """ Test properties of unstructured to rectilinear regridding. """ - ds_src = generate_unstructured_polygon_mesh( - n_cells, n_corners, center_lon_start=15.0, step_lon=25.0 - ) + ds_src = generate_unstructured_polygon_mesh(n_cells, n_corners, center_lon_start=15.0, step_lon=25.0) ds_tgt = create_global_grid(res_lat=res_tgt, res_lon=res_tgt, add_bounds=True) da_src = xr.DataArray( @@ -785,9 +748,7 @@ def test_regrid_unstructured_to_rectilinear_properties( else: res_comp = res - np.testing.assert_allclose( - res_comp.values, np.full(res_comp.shape, 10.0), rtol=1e-3, atol=1e-3 - ) + np.testing.assert_allclose(res_comp.values, np.full(res_comp.shape, 10.0), rtol=1e-3, atol=1e-3) finally: if os.path.exists(weights_file): try: @@ -811,16 +772,12 @@ def test_regrid_unstructured_to_lcc_properties( Test properties of unstructured to projected (LCC) regridding. """ # Create unstructured mesh in geographic coordinates - ds_src = generate_unstructured_polygon_mesh( - n_cells, n_corners, center_lon_start=-97.0, step_lon=0.5 - ) + ds_src = generate_unstructured_polygon_mesh(n_cells, n_corners, center_lon_start=-97.0, step_lon=0.5) # Create LCC target grid extent = (-100000.0, 100000.0, -100000.0, 100000.0) res = 50000.0 - ds_tgt = create_lcc_grid( - extent, res, lat_1=30.0, lat_2=60.0, lat_0=40.0, lon_0=-97.0, add_bounds=True - ) + ds_tgt = create_lcc_grid(extent, res, lat_1=30.0, lat_2=60.0, lat_0=40.0, lon_0=-97.0, add_bounds=True) da_src = xr.DataArray( np.full(n_cells, 20.0), @@ -854,9 +811,7 @@ def test_regrid_unstructured_to_lcc_properties( else: res_comp = res - np.testing.assert_allclose( - res_comp.values, np.full(res_comp.shape, 20.0), rtol=1e-3, atol=1e-3 - ) + np.testing.assert_allclose(res_comp.values, np.full(res_comp.shape, 20.0), rtol=1e-3, atol=1e-3) finally: if os.path.exists(weights_file): try: @@ -882,9 +837,7 @@ def test_regrid_rectilinear_to_unstructured_properties( Test properties of rectilinear to unstructured regridding. """ ds_src = create_global_grid(res_lat=res_src, res_lon=res_src, add_bounds=True) - ds_tgt = generate_unstructured_polygon_mesh( - n_cells, n_corners, center_lon_start=15.0, step_lon=25.0 - ) + ds_tgt = generate_unstructured_polygon_mesh(n_cells, n_corners, center_lon_start=15.0, step_lon=25.0) da_src = xr.DataArray( np.full((ds_src.lat.size, ds_src.lon.size), 30.0), @@ -917,9 +870,7 @@ def test_regrid_rectilinear_to_unstructured_properties( else: res_comp = res - np.testing.assert_allclose( - res_comp.values, np.full(n_cells, 30.0), rtol=1e-3, atol=1e-3 - ) + np.testing.assert_allclose(res_comp.values, np.full(n_cells, 30.0), rtol=1e-3, atol=1e-3) finally: if os.path.exists(weights_file): try: @@ -944,13 +895,9 @@ def test_regrid_lcc_to_unstructured_properties( """ extent = (-100000.0, 100000.0, -100000.0, 100000.0) res = 50000.0 - ds_src = create_lcc_grid( - extent, res, lat_1=30.0, lat_2=60.0, lat_0=40.0, lon_0=-97.0, add_bounds=True - ) + ds_src = create_lcc_grid(extent, res, lat_1=30.0, lat_2=60.0, lat_0=40.0, lon_0=-97.0, add_bounds=True) - ds_tgt = generate_unstructured_polygon_mesh( - n_cells, n_corners, center_lon_start=-97.0, step_lon=0.5 - ) + ds_tgt = generate_unstructured_polygon_mesh(n_cells, n_corners, center_lon_start=-97.0, step_lon=0.5) da_src = xr.DataArray( np.full((ds_src.y.size, ds_src.x.size), 40.0), @@ -983,9 +930,7 @@ def test_regrid_lcc_to_unstructured_properties( else: res_comp = res - np.testing.assert_allclose( - res_comp.values, np.full(n_cells, 40.0), rtol=1e-3, atol=1e-3 - ) + np.testing.assert_allclose(res_comp.values, np.full(n_cells, 40.0), rtol=1e-3, atol=1e-3) finally: if os.path.exists(weights_file): try: @@ -1002,9 +947,7 @@ def test_regrid_lcc_to_unstructured_properties( buffer=st.floats(min_value=0.0, max_value=2.0), ) @settings(max_examples=10, suppress_health_check=[HealthCheck.filter_too_much]) -def test_spatial_slice_properties( - min_x: float, max_x: float, min_y: float, max_y: float, buffer: float -) -> None: +def test_spatial_slice_properties(min_x: float, max_x: float, min_y: float, max_y: float, buffer: float) -> None: """ Test properties of spatial_slice using Hypothesis. @@ -1059,12 +1002,8 @@ def test_spatial_slice_properties( res=st.floats(min_value=20000, max_value=100000), add_bounds=st.booleans(), ) -@settings( - max_examples=10, suppress_health_check=[HealthCheck.filter_too_much], deadline=None -) -def test_create_grid_from_crs_properties( - crs_code: int, extent_offset_x: float, res: float, add_bounds: bool -) -> None: +@settings(max_examples=10, suppress_health_check=[HealthCheck.filter_too_much], deadline=None) +def test_create_grid_from_crs_properties(crs_code: int, extent_offset_x: float, res: float, add_bounds: bool) -> None: """ Test properties of create_grid_from_crs using Hypothesis. @@ -1104,9 +1043,7 @@ def test_create_grid_from_crs_properties( assert "y" in ds_eager.coords # Parity check with Dask using appropriate chunks - ds_lazy = create_grid_from_crs( - crs, extent, res, add_bounds=add_bounds, chunks=min(num_x, num_y, 5) - ) + ds_lazy = create_grid_from_crs(crs, extent, res, add_bounds=add_bounds, chunks=min(num_x, num_y, 5)) assert is_lazy(ds_lazy.lat) xr.testing.assert_allclose(ds_eager, ds_lazy.compute()) @@ -1118,9 +1055,7 @@ def test_create_grid_from_crs_properties( add_bounds=st.booleans(), ) @settings(max_examples=10, suppress_health_check=[HealthCheck.filter_too_much]) -def test_create_grid_like_properties( - res_lat: float, res_lon: float, new_res: float, add_bounds: bool -) -> None: +def test_create_grid_like_properties(res_lat: float, res_lon: float, new_res: float, add_bounds: bool) -> None: """ Test properties of create_grid_like using Hypothesis. @@ -1163,9 +1098,7 @@ def test_create_grid_like_properties( # or just suppress/ignore warning to verify it triggers and functions correctly. with warnings.catch_warnings(): warnings.simplefilter("ignore", category=UserWarning) - ds_new_lazy = create_grid_like( - ds_template_lazy, new_res, add_bounds=add_bounds, chunks=2 - ) + ds_new_lazy = create_grid_like(ds_template_lazy, new_res, add_bounds=add_bounds, chunks=2) assert is_lazy(ds_new_lazy.lat_b) if add_bounds else True xr.testing.assert_allclose(ds_new_eager, ds_new_lazy.compute()) @@ -1177,9 +1110,7 @@ def test_create_grid_like_properties( method=st.sampled_from(["bilinear", "nearest_s2d", "conservative", "patch"]), constant_val=st.floats(min_value=-100.0, max_value=100.0), ) -@settings( - max_examples=8, suppress_health_check=[HealthCheck.filter_too_much], deadline=None -) +@settings(max_examples=8, suppress_health_check=[HealthCheck.filter_too_much], deadline=None) def test_regridder_constant_preservation( res_src: float, res_tgt: float, @@ -1243,9 +1174,7 @@ def test_regridder_constant_preservation( res_src=st.floats(min_value=15.0, max_value=45.0), res_tgt=st.floats(min_value=15.0, max_value=45.0), ) -@settings( - max_examples=8, suppress_health_check=[HealthCheck.filter_too_much], deadline=None -) +@settings(max_examples=8, suppress_health_check=[HealthCheck.filter_too_much], deadline=None) def test_regridder_conservative_conservation( res_src: float, res_tgt: float, diff --git a/tests/test_unstructured_conservative.py b/tests/test_unstructured_conservative.py index a71266d..9115b1f 100644 --- a/tests/test_unstructured_conservative.py +++ b/tests/test_unstructured_conservative.py @@ -4,6 +4,7 @@ import numpy as np import pytest import xarray as xr + from xregrid.regridder import Regridder @@ -180,9 +181,7 @@ def test_unstructured_conservative_weight_scaling(): res_lazy = regridder(da_src_lazy) assert isinstance(res_lazy.data, da.Array) if is_mock: - np.testing.assert_allclose( - res_lazy.compute().values, [0.5, 0.0, 0.0, 0.0], rtol=1e-5 - ) + np.testing.assert_allclose(res_lazy.compute().values, [0.5, 0.0, 0.0, 0.0], rtol=1e-5) else: np.testing.assert_allclose(res_lazy.compute().values, np.ones(4), rtol=1e-5) @@ -193,28 +192,18 @@ def test_unstructured_conservative_weight_scaling(): if is_mock: import dask.distributed - cluster = dask.distributed.LocalCluster( - n_workers=1, threads_per_worker=1, processes=False - ) + cluster = dask.distributed.LocalCluster(n_workers=1, threads_per_worker=1, processes=False) client = dask.distributed.Client(cluster) try: - regridder_parallel = Regridder( - ds_src, ds_tgt, method="conservative", parallel=True - ) + regridder_parallel = Regridder(ds_src, ds_tgt, method="conservative", parallel=True) res_parallel = regridder_parallel(da_src_lazy) assert isinstance(res_parallel.data, da.Array) # Under mock ESMF with n_workers=1, 2 chunks are created (size 2 each). # Each chunk gets its own mock weight of 1.0, scaled to 0.5, # so cell 0 and cell 2 get weight 0.5. - np.testing.assert_allclose( - res_parallel.compute().values, [0.5, 0.0, 0.5, 0.0], rtol=1e-5 - ) - weights_sum_parallel = np.array( - regridder_parallel.weights.sum(axis=1) - ).flatten() - np.testing.assert_allclose( - weights_sum_parallel, [0.5, 0.0, 0.5, 0.0], rtol=1e-5 - ) + np.testing.assert_allclose(res_parallel.compute().values, [0.5, 0.0, 0.5, 0.0], rtol=1e-5) + weights_sum_parallel = np.array(regridder_parallel.weights.sum(axis=1)).flatten() + np.testing.assert_allclose(weights_sum_parallel, [0.5, 0.0, 0.5, 0.0], rtol=1e-5) finally: client.close() cluster.close() diff --git a/tests/test_utils.py b/tests/test_utils.py index 7e4e7f7..9b08aa9 100644 --- a/tests/test_utils.py +++ b/tests/test_utils.py @@ -1,7 +1,6 @@ from __future__ import annotations # Consolidated tests: utils - import os from unittest.mock import MagicMock, patch @@ -23,8 +22,8 @@ plot_comparison, spatial_slice, ) -from xregrid.utils import get_crs_info from xregrid.grid import _get_mesh_info +from xregrid.utils import get_crs_info try: import esmpy @@ -52,9 +51,7 @@ def test_auto_bounds_conservative_numpy_dask(): # 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 - ) + 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 @@ -78,9 +75,7 @@ def test_regridder_keep_attrs(): ds_tgt = create_global_grid(20, 20) regridder = Regridder(ds_src, ds_tgt, method="conservative") - da_src = xr.DataArray( - np.random.rand(10, 20), dims=("lat", "lon"), coords=ds_src.coords, name="foo" - ) + da_src = xr.DataArray(np.random.rand(10, 20), dims=("lat", "lon"), coords=ds_src.coords, name="foo") da_src.attrs["my_custom_attr"] = "preserve_me" # Eager: custom attr preserved AND regridder history kept @@ -214,19 +209,13 @@ def test_crs_propagation_dataarray(): # 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 - ) + 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" - ) + 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}) @@ -489,9 +478,7 @@ def test_spatial_slice_rectilinear() -> None: 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_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. @@ -580,9 +567,7 @@ def test_create_global_grid_lazy(): 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} - ) + 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 @@ -959,9 +944,7 @@ def test_regridder_dtype_time_fallback(): 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") - ) + 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 @@ -1052,9 +1035,7 @@ def test_regridder_vertical_dimension_detection(): 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") - ) + 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 @@ -1094,9 +1075,7 @@ def __init__(self, ds, uxgrid): 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 = 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 @@ -1247,9 +1226,7 @@ def test_regridder_user_specific_structure(): assert "mesh" in res_ds.data_vars # Non-spatial data var should be preserved -@pytest.mark.skip( - reason="ESMF abort (SIGABRT) in ESMP_MeshGetElemCoordPtr with small synthetic mesh — kills process" -) +@pytest.mark.skip(reason="ESMF abort (SIGABRT) in ESMP_MeshGetElemCoordPtr with small synthetic mesh — kills process") def test_regridder_raw_ugrid_conservative_with_time(): n_face = 10 n_node = 12 @@ -1391,9 +1368,7 @@ def test_create_global_grid(): def test_create_regional_grid(): - ds = create_regional_grid( - lat_range=(-45, 45), lon_range=(0, 90), res_lat=5, res_lon=5 - ) + ds = create_regional_grid(lat_range=(-45, 45), lon_range=(0, 90), res_lat=5, res_lon=5) assert ds.lat.size == 18 # 90 / 5 assert ds.lon.size == 18 # 90 / 5 assert ds.lat.min() == -42.5