Support AIMIP IC and forcing generation for a different ERA5 build - #1460
Open
brianhenn wants to merge 5 commits into
Open
Support AIMIP IC and forcing generation for a different ERA5 build#1460brianhenn wants to merge 5 commits into
brianhenn wants to merge 5 commits into
Conversation
HGTsfc, DSWRFtoa and land_fraction are read from an ACE ERA5 zarr and differ between builds, so the forcing should be generated from the build the model was trained on. Makes the source selectable via FORCING_EXTRA_ARGS. The insolation extension is now opt-in via --extension-start rather than always appending 2023 onward, since a source that already spans the window needs none. Without it, a source that falls short errors instead of silently truncating. Also makes OUTPUT_ZARR_NAME and MODIFIED_FORCING_NAME overridable so a regeneration does not overwrite the existing datasets.
xr.concat's defaults broadcast variables that have no time dimension across every timestep, inflating the store and slowing the data loader, which handles static variables natively. Skip 0-d variables when setting chunks and shards; zarr rejects empty tuples, which the preserved scalars now hit.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
The AIMIP IC and forcing scripts assumed one ERA5 build and one model's variable set.
Generating inputs for a model trained on a later ERA5 build required editing the scripts,
and the prepend step stored static variables once per timestep.
HGTsfc,DSWRFtoaandland_fractionare read from an ACE ERA5 zarr rather than fromthe AIMIP protocol data, and they differ between ERA5 builds — the March 2026 pipeline
rewrite changed the regridding package and some data sources, moving
HGTsfcby tens ofmetres in mountainous terrain. The source is now selectable so the forcing can match the
build a model was trained on.
Changes:
create_aimip_ic_datasets.py:--include-near-surfaceaddsTMP2m/Q2m/UGRD10m/VGRD10mfor models carrying them as prognostics.
sel(time=[...])keepstimea length-1dimension, which
get_initial_conditionrequires.interpolate_aimip_forcing.py:--extension-start,--repeat-source-startand--repeat-source-endmake the synthetic insolation extension opt-in. Previously 2023onward was always synthesized, which is wrong for a source that already spans the window.
Without the flag, a source ending before
--end-timenow raises instead of silentlytruncating.
prepend_first_timestep_forcing.py: passdata_vars="minimal",coords="minimal",compat="override"toxr.concat, so variables without a time dimension are notbroadcast along it. The data loader handles static variables natively; broadcasting only
inflates the store and slows loading.
encoding.py: skip 0-d variables inset_shards_chunks. zarr rejects empty chunk/shardtuples, which scalars hit once they survive the concat above. Not reachable from the
6-variable AIMIP forcing, but is from the 73-variable combined forcing, which carries 18
ak_*/bk_*scalars.Makefile:IC_EXTRA_ARGSandFORCING_EXTRA_ARGSpassthroughs;OUTPUT_ZARR_NAMEandMODIFIED_FORCING_NAMEbecome?=so a regeneration cannot overwrite existing datasets.README.md: documents the ERA5-build dependency and the new options.Tests added
If dependencies changed, "deps only" image rebuilt and "latest_deps_only_image.txt" file updated