Skip to content

Support AIMIP IC and forcing generation for a different ERA5 build - #1460

Open
brianhenn wants to merge 5 commits into
mainfrom
scripts/aimip-ic-generation-fixes
Open

Support AIMIP IC and forcing generation for a different ERA5 build#1460
brianhenn wants to merge 5 commits into
mainfrom
scripts/aimip-ic-generation-fixes

Conversation

@brianhenn

Copy link
Copy Markdown
Contributor

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, DSWRFtoa and land_fraction are read from an ACE ERA5 zarr rather than from
the AIMIP protocol data, and they differ between ERA5 builds — the March 2026 pipeline
rewrite changed the regridding package and some data sources, moving HGTsfc by tens of
metres 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-surface adds TMP2m/Q2m/UGRD10m/VGRD10m
    for models carrying them as prognostics. sel(time=[...]) keeps time a length-1
    dimension, which get_initial_condition requires.

  • interpolate_aimip_forcing.py: --extension-start, --repeat-source-start and
    --repeat-source-end make the synthetic insolation extension opt-in. Previously 2023
    onward was always synthesized, which is wrong for a source that already spans the window.
    Without the flag, a source ending before --end-time now raises instead of silently
    truncating.

  • prepend_first_timestep_forcing.py: pass data_vars="minimal", coords="minimal",
    compat="override" to xr.concat, so variables without a time dimension are not
    broadcast along it. The data loader handles static variables natively; broadcasting only
    inflates the store and slows loading.

  • encoding.py: skip 0-d variables in set_shards_chunks. zarr rejects empty chunk/shard
    tuples, 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_ARGS and FORCING_EXTRA_ARGS passthroughs; OUTPUT_ZARR_NAME and
    MODIFIED_FORCING_NAME become ?= 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

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.
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.

1 participant