Skip to content

Update xarray usage in RAiDER to use lazy loading - #808

Merged
jlmaurer merged 6 commits into
dbekaert:devfrom
jlmaurer:pr/xarray_memory
Aug 11, 2026
Merged

Update xarray usage in RAiDER to use lazy loading#808
jlmaurer merged 6 commits into
dbekaert:devfrom
jlmaurer:pr/xarray_memory

Conversation

@jlmaurer

@jlmaurer jlmaurer commented Aug 5, 2026

Copy link
Copy Markdown
Collaborator

xarray makes available two methods for opening datasets, open_dataset and load_dataset. Using open_dataset allows for lazy loading and memory efficiency. This PR change all the locations were load_dataset was used to `open_dataset', resulting in much better memory performance for RAiDER. This PR also wraps netcdf dataset opening with try/except/finally logic that ensures the datasets are closed properly.

The PR also fixes a small issue where ECMWF requests were being handled in two API requests instead of one.

All unit tests pass locally.

Type of change

  • Bug fix (non-breaking change which fixes an issue)
  • New feature (non-breaking change which adds functionality)
  • Breaking change (fix or feature that would cause existing functionality to not work as expected)

Checklist:

  • I have added an explanation of what your changes do and why you'd like us to include them.
  • I have written new tests for your core changes, as applicable.
  • I have successfully ran tests with your changes locally.
  • My change requires a change to the documentation.
  • I have updated the documentation accordingly.

@sssangha

sssangha commented Aug 10, 2026

Copy link
Copy Markdown
Collaborator

Hello @jlmaurer

I reviewed and tested the changes locally, and added two small safety enhancements:

  1. Robust Resource Cleanup with ExitStack:
    To guarantee that all opened datasets are cleanly closed—even if an exception occurs mid-loop or during .to_netcdf()—I wrapped the file loading using contextlib.ExitStack:

    from contextlib import ExitStack
    
    with ExitStack() as stack:
        datasets = [stack.enter_context(xr.open_dataset(f)) for f in wfiles]
        # ... perform calculations and save to netcdf ...
  2. Preventing In-Place Variable Reassignment
    Assigning ds_out = datasets[0] mutates datasets[0] directly while its disk handle is open. Updated this to explicit shallow copying:

    ds_out = datasets[0].copy(deep=False)
    

Future Consideration on Dask Chunking
Standard lazy loading via open_dataset handles the current memory performance improvements well. If processing ultra-large weather model grids ever hits RAM limits down the road, we can revisit introducing explicit chunking (e.g., chunks={}) to leverage Dask streaming directly to disk.

LGTM overall! I'm ready to approve if these changes look good to you.

jlmaurer and others added 5 commits August 11, 2026 08:52
ERA-5 model-level data was fetched in two separate CDS calls (lnsp+z, then
t+q) and recombined on disk via a temp-file copy. Request all four params in
one retrieval instead: one queue wait instead of two, one temp file instead
of three, and no shutil.copy of the full cube.
…tests

The single combined request for lnsp/z/t/q could not work. ERA-5 archives lnsp
and z at model level 1 only, so a mixed-level request is not something CDS will
return as one netCDF -- the multi-date path added in dbekaert#810 splits the request for
exactly this reason. Where a single file is returned at all, the surface fields
come back padded onto the full 137-level axis, so .squeeze() no longer collapses
the level dimension and calcgeoh raises:

  ValueError: could not broadcast input array from shape (137,11,11) into (11,11)

reproduced against test/weather_files/ERA-5_2020_01_30_T13_52_45.nc, which was
downloaded by the pre-split single-request path and still carries that layout.

Restore the two requests and keep the rest of the cleanup: the full-cube
shutil.copy is dead (to_netcdf overwrites out_path) and is dropped.

Nothing in CI performs a CDS retrieval, which is why _get_from_cds had no cover
at all. test_ecmwf_fetch.py stubs cdsapi.Client and drives the real
post-processing into the real reader, asserting the layout contract the two
share: t/q/z as (time, level, lat, lon) cubes and lnsp recovered as lnsp[0, 0].
The stub refuses a mixed-level request, so re-merging the two retrievals fails
the suite instead of failing at download time. Its synthetic responses use the
valid_time and model_level coordinate names the current CDS returns.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
jlmaurer added a commit to jlmaurer/RAiDER-1 that referenced this pull request Aug 11, 2026
The rebase onto dbekaert#808 left _get_from_cds still issuing one merged request for
lnsp/z/t/q, which cannot work: CDS returns a mixed-level request as separate
files, and where a single file does come back the surface fields are padded onto
the full 137-level axis so .squeeze() no longer yields the (lat, lon) arrays
calcgeoh and _makeDataCubes expect.

_batch_get_from_cds already splits the request correctly and documents why, so
rather than repeat that logic, treat a single date as a one-element batch. Both
level types now share one download path, which is what stops the two from
drifting apart again -- the pl path was already delegating.

Also restores the latitude fix from dbekaert#811 in _load_model_level. It was reverted
by resolving the rebase conflict in favour of this branch's copy of ecmwf.py,
which predates that fix; test_ecmwf_levels.py caught it.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
@jlmaurer

Copy link
Copy Markdown
Collaborator Author

@sssangha thanks, looks good! I checked the PR again and found a couple small issues related to when I pulled it out from #810, so I fixed those (this is in ecmwf.py and related tests). If you want to just take a quick look at those, we should be good here.

@sssangha sssangha left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants