Repository navigation
Conversation
Builds EAMxx/SCREAM atmospheric initial conditions on ne1024pg2 with the E3SM L128 vertical grid, in the 32-level Zarr v3 layout of the SCREAM stores: - fast_ic.py: a torch port of the E3SM HICCUP steps (horizontal fv2fv_flx remap, surface-pressure adjustment, L128 vertical remap, remove_supersaturation, PotentialTemperature and hydrostatic z_mid). One GPU builds an IC from ARCO ERA5 in about 80 s; it matches the HICCUP pipeline bit for bit apart from last-bit differences in a few qv values. - hres_ic.py: the same steps on the ECMWF HRES analysis (WeatherBench2 hres_t0, 2016-2022). - run_hiccup_eamxx.py + hiccup_l128_to_zarr.py: the reference pipeline that calls HICCUP directly (2.5-3 h per IC), with hiccup_time_calendar.patch for current xarray; run_tc_batch.py builds the ten tropical-cyclone cases. - docs/initial_conditions.md: inputs, commands, validation and timing. The SCREAM phis/sgh30/landfrac come from the nvidia/STRATA-SCREAM-sdy1 Hugging Face dataset; run_hiccup_eamxx.py now requires the local phis array instead of copying it from internal storage. Checked: tests/test_era5_fast_ic.py and tests/test_era5_hiccup_ic.py pass, and fast_ic run from this tree with the Hugging Face static arrays rebuilt the 2021-07-15T12 IC bit-identical to the reference store in every variable. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
The surface geopotential is checked against the closed form for an isothermal atmosphere, inside the pressure profile and below its lowest level. fetch_hres_t0 is run on a 3-level 5x8 Zarr v2 store in the hres_t0 layout: it must pick the requested time, flip the grid to north-to-south and reject off-hour timestamps, missing times and omega in hPa. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
ncremap -a fv2fv_flx --lrg2sml runs GenerateOverlapMesh and GenerateOfflineMap serially, 13-18 h at this size, and no copy of the map is published. The map is first order, so a destination column's weights depend only on its own overlaps. build_map.py runs the same two TempestRemap commands on contiguous ne1024pg2 column ranges in a process pool and stacks the parts; it also writes the ERA5 SCRIP file with HICCUP's ncremap -G command and keeps the reuse/adopt rules of build_ne1024_map.sh, which it replaces. On 72 Grace cores the full map took 32 min (144 parts; the two polar parts finish last). Against the serial map built on x86 with NCO 5.3.6: row and col identical (38,299,240 links), 31,590 weights differ in float32 (max rel 1.3e-5 on sliver overlaps), and a smooth field remapped with both maps differs by at most 1.8e-5 K. Two different splits agree to 7e-9 with identical float32 weights. tests/test_build_map.py builds a 10-degree lat-lon -> ne4pg2 map in three uneven parts and checks it against serial ncremap (skipped when nco/tempest-remap are not on PATH). Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
build_map.py: - Recompute source-side area_a (summed overlap area) and frac_a from all links after the merge; per-part values only see that part's overlaps. On the full map frac_a now matches the serial map to 1e-15 (was 1.2e-5). - Build parts in "<map>.parts" under --tmp-dir and remove only that, not a user's --tmp-dir. - Run TempestRemap inside the part dir with short relative names; the conda-forge 2.2.0 build aborts on output paths of 98+ characters. - Name parts by column range and record the input grids, so parts from a run with other --parts or other grids are never merged. - Stop all parts on the first failure; adopt an existing map only if its links run from the first to the last column; history and sidecar record the module command with absolute paths. Readers and converter: - fast_ic/hres_ic check the store's time epoch and dimension order (both real stores pass); fast_ic rejects omega in hPa like hres_ic. - hiccup_l128_to_zarr checks the NetCDF time against --timestamp, and load_latlon stops when lat units are missing or do not match the values (degree values labelled "rad" scaled coszr's latitudes by 57.3). - run_hiccup_eamxx no longer stamps the map command with the IC's time. Tests: new end-to-end checks of the ps adjustment (standard atmosphere, including the extrapolated surface T), subsaturated qv kept and virtual temperature in z_mid, supersaturated qv clipped to within 2.5% of Murphy-Koop saturation, 1-based map links, the ARCO reader, the HRES epoch/dimension checks, the IC time check and lat/lon units. Eight mutations that the previous tests missed (hPa/Pa, dry Tv, top-level T, no T*, mean remap, Bolton constant, raw qv, map col offset) now fail. Docs: map and output paths under $WORK in every example, a download for the L128 file, fsspec/aiohttp in the HICCUP env, batch skip rule, and corrected numbers (3.2 GB map, 175k-column parts, 2.5-3 h). Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Mention IC generation in the intro and give it a top-level section: what the pipeline does, the three entry points, the one-time inputs, the map build and an example fast_ic command. Say plainly that run_screamcast_nudged.py does not take a one-frame IC store as its initial state yet; the docs no longer say the rollout loaders read it. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
run_screamcast_nudged.py --ic-store <store> reads the initial state from a one-frame store built by initial_conditions/ instead of SCREAM_MAIN_ZARR_PATH. The forecast starts at the store's initial_time; --initial-time must match it if given, and --correction must be none because the store has no later truth. ScreamData.to_data_source now passes its reference_time to ScreamDataSource instead of relying on the reader's 2020-10-01 default (the same date the SCREAM archive entry uses, so that path is unchanged). data_catalog.ic_store() builds the catalog entry for an IC store. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
nbren12
requested changes
Nov 4, 2026
| @@ -0,0 +1,627 @@ | |||
| #!/usr/bin/env python | |||
Contributor
There was a problem hiding this comment.
Would we ever want to use this instead of fast_ic.py? I'm a bit worried to keep reference implementation code around that won't be used.
Maybe worth removing unless you envision using it regularly for testing. Ditto for the corresponding test module.
|
|
||
| netCDF4 = pytest.importorskip("netCDF4") | ||
|
|
||
| from initial_conditions.build_map import build, looks_complete # noqa: E402 |
Contributor
There was a problem hiding this comment.
should we put this in the source tree so the import doesn't require noqa?
| @@ -0,0 +1,123 @@ | |||
| # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. | |||
Contributor
There was a problem hiding this comment.
how long do these tests take to run?
| @@ -0,0 +1,15 @@ | |||
| # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. | |||
Contributor
There was a problem hiding this comment.
let's not add a new top level package to the repo. scripts can go in scripts/initial_conditions, and any necessary tested modules into screamcast/
Merged
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.
Initial conditions from ERA5 and ECMWF HRES
Adds
initial_conditions/, which builds EAMxx/SCREAM atmospheric initial conditions on ne1024pg2 with the E3SM L128 vertical grid, as one-frame Zarr stores with flatncoland the 32 levels STRATA uses.fast_ic.py: a torch port of the E3SM HICCUP steps (fv2fv_flx horizontal remap, surface-pressure adjustment, L128 vertical remap, remove_supersaturation, potential temperature and hydrostatic z_mid). It builds one IC from ARCO ERA5 in about 80 s on one GPU. Its output matches the HICCUP pipeline bit for bit on ten tropical-cyclone cases, apart from last-bit differences in a few values.hres_ic.py: the same steps on the ECMWF HRES analysis (WeatherBench2hres_t0, 2016-2022). HRES has no data above 50 hPa; STRATA reads only the levels from 100 hPa down.build_map.py: builds the ERA5 -> ne1024pg2 first-order conservative map (ncremap -a fv2fv_flx --lrg2sml) in parallel pieces of destination columns, in about 32 min instead of HICCUP's 13-18 h serial build. The merged map matches a serial ncremap map in every variable;area_aandfrac_aare recomputed after the merge. An existing map is adopted only if its sizes match and it is complete.run_hiccup_eamxx.py+hiccup_l128_to_zarr.py: the reference pipeline that calls HICCUP directly (2.5-3 h per IC);run_tc_batch.pybuilds the ten TC cases.--timestampthat differs from the IC's own time, and lat/lon files whoseunitsare missing or don't match the values.run_screamcast_nudged.py --ic-store <store> --correction nonestarts a forecast from such a store at itsinitial_time.ScreamData.to_data_sourcenow passes the catalog'sreference_timeto the reader instead of relying on the reader's 2020-10-01 default, which is also the SCREAM archive's, so existing runs are unchanged.README.md: a new top-level section with what the pipeline does, its one-time inputs and an example command.docs/initial_conditions.md: inputs, commands, validation and timing. The SCREAM static fields come from thenvidia/STRATA-SCREAM-sdy1Hugging Face dataset; the L128 vertical-coordinate file is downloaded from the HICCUP repository at a pinned commit.Tests:
tests/test_era5_fast_ic.py,tests/test_era5_hiccup_ic.pyandtests/test_hres_ic.py(26 tests) use ground truth the code never sees: an isothermal atmosphere at rest, standard-atmosphere surface pressure, Murphy-Koop saturation, hypsometric heights with virtual temperature, a hand-built map with 1-based unordered links, and the readers on small synthetic stores in the ARCO and WeatherBench2 layouts. Eight deliberately injected bugs (hPa/Pa mix-up, dry instead of virtual temperature, wrong extrapolation, wrong saturation formula, skipped supersaturation removal, ignored map weights, off-by-one map indexing) each fail at least one test.tests/test_ic_store_source.py(3 tests): a one-frame store is read at itsinitial_time, has no frame 10 min later, and a catalog entry's framenisreference_time + n * 10 min.tests/test_build_map.py(3 tests, skipped unlessncoandtempest-remapare on PATH) compares the parallel build with a serial ncremap map on a small grid, and checks that parts from other grids and incomplete maps are not reused.All pass;
make lintpasses. A v1 forecast from a HICCUP IC store (Maria, 2017-09-14T00) on one GB300 loaded all 145 input channels and finished the first model forward. On one rank, tile-jitter post-processing then hits torch's 32-bit index limit (all 1536 tiles in one tensor), so the full 2-step--ic-storeforecast still needs a check on 4 GPUs.🤖 Generated with Claude Code