Repository navigation
Feat/ic generation - #13
Merged
Merged
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> Signed-off-by: Zeyuan Hu <zeyuanh@nvidia.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> Signed-off-by: Zeyuan Hu <zeyuanh@nvidia.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> Signed-off-by: Zeyuan Hu <zeyuanh@nvidia.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> Signed-off-by: Zeyuan Hu <zeyuanh@nvidia.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> Signed-off-by: Zeyuan Hu <zeyuanh@nvidia.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> Signed-off-by: Zeyuan Hu <zeyuanh@nvidia.com>
- Move initial_conditions/ to scripts/initial_conditions/ instead of a new top-level package; run commands are python -m scripts.initial_conditions.*. - Remove the HICCUP reference pipeline (run_hiccup_eamxx.py, run_tc_batch.py, the HICCUP calendar patch, the NetCDF-to-Zarr converter and tests/test_era5_hiccup_ic.py) and fast_ic's --era5-atm/--era5-sfc input. The store helpers the converter shared move to ic_store.py, and the lat/lon units test moves to test_era5_fast_ic.py. - test_build_map.py imports build_map at the top (netCDF4 is a core dependency), so it no longer needs importorskip or noqa. - fast_ic.py carries HICCUP's BSD-3-Clause notice (Copyright 2026 Lawrence Livermore National Security, LLC; added in E3SM-Project/HICCUP@0dec8bb). The docs gain a License section, pin the L128 coordinate download to 0dec8bb, and say these are STRATA ICs on the pg2 physics grid, not ICs EAMxx can start from. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Zeyuan Hu <zeyuanh@nvidia.com>
Collaborator
Author
|
@nbren12, this replaces #12 and addresses your review there: No new top-level package: everything is under scripts/initial_conditions/, and tests import scripts.initial_conditions.* the way scripts/ace/run_pipeline.py imports scripts.ace. Reference implementation: removed run_hiccup_eamxx.py, run_tc_batch.py, the HICCUP patch, the NetCDF-to-Zarr converter and tests/test_era5_hiccup_ic.py. The one test there that covered kept code (lat/lon units) moved to test_era5_fast_ic.py. The HICCUP comparison is described in the docs as a past validation. noqa in test_build_map.py: it was only needed because importorskip("netCDF4") came first. netCDF4 is a core dependency, so both are gone. Test time: the three build_map tests take about 5 s with nco and tempest-remap installed, and are skipped without them. License: HICCUP now has a BSD-3 LICENSE (E3SM-Project/HICCUP@0dec8bb, © 2026 LLNS). fast_ic.py carries the full notice, and the docs have a License section. |
nbren12
reviewed
Nov 7, 2026
Contributor
|
Can you put all the python modules code that is imported in the main source folder strata/ please? I think this will make it easier to avoid PYTHONPATH issues since it will be isntalled alongside the package. Currenly the code will only run when the working directory is the repo root. scripts/ should only contain thin wrapper scripts or code that are not imported anywhere. |
Address review on PR 13: the modules that are imported (fast_ic, hres_ic, build_map, ic_store, layout) now live in the installed package as screamcast.initial_conditions, so they import without PYTHONPATH or running from the repository root. scripts/initial_conditions/ keeps only thin command-line wrappers around each module's main(). Also update the README's IC forecast example to the recommended v1 settings (no cross-face halo adjoint, SHT omega lmax 3). Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Zeyuan Hu <zeyuanh@nvidia.com>
The tools run as `python -m screamcast.initial_conditions.<name>`, which works whenever the wrappers did (installed package or repository root), and every example already uses that form. scripts/ keeps code that is not imported anywhere; a second entry point per tool adds nothing. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Zeyuan Hu <zeyuanh@nvidia.com>
Collaborator
Author
|
@nbren12 the IC modules now live in |
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
scripts/initial_conditions/, which builds STRATA initial conditions on the SCREAM ne1024pg2 physics grid with the E3SM L128 vertical grid, as one-frame Zarr stores with flatncoland the 32 levels STRATA uses. These are STRATA inputs, not ICs EAMxx can start from (EAMxx reads its ICs on the np4 GLL grid).fast_ic.py: a torch port of the E3SM HICCUP steps (fv2fv_flx horizontal remap, surface-pressure adjustment, L128 vertical remap, remove_supersaturation), then potential temperature and hydrostatic z_mid. It builds one IC from ARCO ERA5 in about 80 s on one GPU. On ten tropical-cyclone cases its output matched ICs built by HICCUP itself bit for bit, 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, same TempestRemap commands) in parallel pieces of destination columns, in about 32 min instead of 13-18 h serially. The merged map matches a serial ncremap map;area_aandfrac_aare recomputed after the merge.ic_store.py,layout.py: static inputs,coszr, store creation and checks, constants.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.mdanddocs/initial_conditions.md: inputs, commands, validation, timing.License.
fast_ic.pycontains code ported from HICCUP, released under the 3-Clause BSD license (Copyright 2026, Lawrence Livermore National Security, LLC; LICENSE added in E3SM-Project/HICCUP@0dec8bb). The full HICCUP notice is kept at the top offast_ic.py, and the docs have a License section. The L128 coordinate file is downloaded from the HICCUP repository at 0dec8bb. The SCREAM static fields come from thenvidia/STRATA-SCREAM-sdy1Hugging Face dataset.Tests:
tests/test_era5_fast_ic.pyandtests/test_hres_ic.pyuse 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, lat/lon units checks, and the readers on small synthetic stores in the ARCO and WeatherBench2 layouts.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, about 5 s; 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.In the STRATA GB300 image: 24 passed, 3 skipped (build_map, no nco there); the 3 build_map tests pass in a conda env with nco and tempest-remap.
python -m scripts.initial_conditions.{fast_ic,hres_ic,build_map} --helprun.make lintpasses. A v1 forecast from a HICCUP-built IC store (Maria, 2017-09-14T00) on one GB300 loaded all 145 input channels and finished the first model forward; the full 2-step--ic-storeforecast on 4 GPUs has not been run yet.🤖 Generated with Claude Code