diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 5375967..c97bc1c 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -2,9 +2,11 @@ name: CI on: push: - branches: [ main ] + branches: ["*"] pull_request: - branches: [ main ] + branches: ["*"] + types: [opened, synchronize, reopened] + workflow_dispatch: jobs: test: @@ -13,6 +15,8 @@ jobs: FORCE_JAVASCRIPT_ACTIONS_TO_NODE24: true steps: - uses: actions/checkout@v4 + with: + persist-credentials: false - name: Setup Micromamba uses: mamba-org/setup-micromamba@v2 diff --git a/src/xregrid/utils.py b/src/xregrid/utils.py index 2241835..c11fc42 100644 --- a/src/xregrid/utils.py +++ b/src/xregrid/utils.py @@ -1364,8 +1364,12 @@ def create_mesh_from_coords( # Map any single-dimension chunking to n_pts n_pts_chunks = next(iter(chunks.values())) chunks = {"n_pts": n_pts_chunks} + else: + n_pts_chunks = chunks x_da = x_da.chunk(chunks) y_da = y_da.chunk(chunks) + else: + n_pts_chunks = None # Backend detection for provenance is_lazy = chunks is not None or hasattr(x_da.data, "dask") @@ -1396,7 +1400,10 @@ def create_mesh_from_coords( ds = xr.Dataset( coords={ - "n_pts": (["n_pts"], np.arange(x_da.size)), + "n_pts": ( + ["n_pts"], + _lazy_arange(0, x_da.size, 1, chunks=n_pts_chunks), + ), "x": ( ["n_pts"], x_da.data, diff --git a/tests/test_aero_grid_gen.py b/tests/test_aero_grid_gen.py index 03cb643..6cd5e16 100644 --- a/tests/test_aero_grid_gen.py +++ b/tests/test_aero_grid_gen.py @@ -1,6 +1,12 @@ +import numpy as np import xarray as xr import pytest -from xregrid.utils import create_global_grid, create_regional_grid, create_grid_from_crs +from xregrid.utils import ( + create_global_grid, + create_regional_grid, + create_grid_from_crs, + create_mesh_from_coords, +) def test_global_grid_backend_consistency(): @@ -63,3 +69,57 @@ def test_laziness_large_grid(): # Check shape to ensure it's correct assert ds_lazy.lat.size == 180000 assert ds_lazy.lon.size == 360000 + + +def test_mesh_laziness_backend_consistency(): + """ + Aero Protocol: Double-Check test for mesh laziness. + Verifies that coordinates are lazy when chunks are provided and + results are identical between Eager and Lazy backends. + """ + pytest.importorskip("pyproj") + n_pts = 1000 + x = np.linspace(-10, 10, n_pts) + y = np.linspace(-10, 10, n_pts) + crs = "EPSG:3857" # Web Mercator + + # 1. Eager Mesh + ds_eager = create_mesh_from_coords(x, y, crs=crs) + + assert not hasattr(ds_eager.n_pts.data, "dask") + assert "(Eager)" in ds_eager.attrs["history"] + + # 2. Lazy Mesh + chunks = 100 + ds_lazy = create_mesh_from_coords(x, y, crs=crs, chunks=chunks) + + # n_pts might be eager in xarray as it's an index, + # but we should check if we can verify laziness another way or if we should check data variables + # For now, let's check x and y which are non-index coords in the Dataset output + assert hasattr(ds_lazy.x.data, "dask") + assert ds_lazy.x.chunks is not None + assert "(Lazy)" in ds_lazy.attrs["history"] + + # lat/lon should also be lazy (via apply_ufunc) + assert hasattr(ds_lazy.lat.data, "dask") + assert hasattr(ds_lazy.lon.data, "dask") + + # 3. Numerical Identity + xr.testing.assert_identical(ds_eager.n_pts, ds_lazy.n_pts.compute()) + xr.testing.assert_allclose(ds_eager.lat, ds_lazy.lat.compute()) + xr.testing.assert_allclose(ds_eager.lon, ds_lazy.lon.compute()) + + +def test_mesh_laziness_dict_chunks(): + """Verify that dict-based chunks are also handled correctly for n_pts.""" + pytest.importorskip("pyproj") + n_pts = 500 + x = np.linspace(-10, 10, n_pts) + y = np.linspace(-10, 10, n_pts) + crs = "EPSG:4326" + + chunks = {"n_pts": 250} + ds_lazy = create_mesh_from_coords(x, y, crs=crs, chunks=chunks) + + assert hasattr(ds_lazy.x.data, "dask") + assert ds_lazy.x.chunks[0][0] == 250