Strata

Code for Scaling Storm-Resolving Atmospheric AI Simulation to the Entire Planet (NeurIPS 2026, Main Track).

Strata is a deep-learning weather emulator for the SCREAM global atmosphere model. It trains transformer-based neural networks to emulate SCREAM atmospheric physics on the cubed-sphere grid and supports multi-day global forecasting.

The Strata architecture (a two-stage 3D neighborhood-attention transformer with stereographic rotary position embeddings) lives in NVIDIA PhysicsNeMo as physicsnemo.experimental.models.strata. This repository adds the SCREAM-specific pieces around it: spherical tile geometry, tile-local wind rotation, data pipelines, training configs, rollout/inference tooling (see screamcast/strata_wrappers.py), and a pipeline that builds SCREAM initial conditions from ERA5 or the ECMWF HRES analysis (see Initial conditions).

The public release is named Strata, but the Python package, command-line scripts, and environment variables keep the project's original name: you import screamcast, run the screamcast-prefixed entry points, and configure SCREAM_* variables. Strata and screamcast refer to the same project.

This code is provided for research and development purposes only.

Setup

The models need a recent NVIDIA GPU software stack that cannot come from PyPI alone: a CUDA-tuned PyTorch build, NVIDIA DALI (data pipelines), and a NATTEN build that matches the installed torch. The supported base environment is the NGC PyTorch container nvcr.io/nvidia/pytorch:26.09-py3.

Docker (recommended)

Build the training/inference image from the repository root:

docker build -f docker/Dockerfile -t strata .

The image contains the complete environment, including NATTEN built from source for Ada and Hopper GPUs (pass --build-arg CUDA_ARCH=8.0 for A100; see docker/build_natten.sh) and PhysicsNeMo pinned to a commit that includes the Strata models. It does not contain this repository; mount it at run time.

pip install (plain PyPI environment)

Outside NGC containers the package also installs as a normal Python project. NATTEN's wheel index lets pip pick the CUDA wheel matching your torch build:

pip install -e . -f https://whl.natten.org

Optional extras: .[s3] (S3 data access), .[analysis] (plotting), .[cu12]/.[cu13] (DALI for your CUDA major version), .[dev] (tests), .[notebooks] (marimo). This path installs the newest compatible versions from PyPI; use the Docker image for the tested stack.

Environment variables

The code reads its paths from environment variables. Export them, or put them in a .env file at the repository root (cp envs/example.env .env), which is loaded when present. Which ones you need depends on the task:

Task	Variables
Forecast with the released checkpoint	AUX_DATA_ROOT (auxiliary files, below) and SCREAM_MAIN_ZARR_PATH (the sdecadal store, for initial conditions)
Training	also PROJECT_ROOT (run directories, and the checkpoints that chained configs resume from), ZARR_ROOT (directory or s3:// prefix holding the stores), and SCREAM_MAIN_ZARR_PATH and SCREAM_AUX_ZARR_PATH, both set to the local sdecadal store, which the validation loader reads
Weights & Biases logging	WANDB_API_KEY, optionally WANDB_PROJECT and WANDB_ENTITY; without the key, training skips wandb
s3:// store paths	SCREAM_ZARR_PROFILE, the name of a remote in ~/.config/rclone/rclone.conf that holds the S3 endpoint and keys

Data

The raw SCREAM output used for training and evaluation is at https://portal.nersc.gov/archive/home/n/ndk/www/STRATA2026 (CC BY 4.0, see License). A training-ready version in Zarr format is on Hugging Face, one repository per simulation:

Repository	Simulation	Time steps (10 min)	Size
nvidia/STRATA-SCREAM-sdecadal	sdecadal, 14 days (transient historical, F20TR)	2017	36.8 TB
nvidia/STRATA-SCREAM-sdy1	DYAMOND1, 7 days (F2010)	1009	19.6 TB
nvidia/STRATA-SCREAM-sdy2	DYAMOND2, 7 days (F2010)	1009	19.7 TB

Each repository is a single Zarr v3 store on the ne1024pg2 grid: arrays are (time, lev, ncol) with 32 model levels. The stores have no coordinate arrays; the code attaches them. A store can be read directly over HTTP, which is convenient for exploration:

import xarray as xr

ds = xr.open_zarr("hf://datasets/nvidia/STRATA-SCREAM-sdy1")
ds["T_2m"].isel(time=500, ncol=slice(0, 100)).values  # streams just the chunks it needs

For training, download the stores (hf download <repo> --repo-type dataset --local-dir <dir>); HTTP streaming is too slow for the dataloader. Place each store under ZARR_ROOT with the directory name the training configs expect (_ZARR_SDECADAL / _ZARR_SDY1 / _ZARR_SDY2 in train_configs.py):

Repository	Directory under ZARR_ROOT
nvidia/STRATA-SCREAM-sdecadal	sdecadal.ne1024pg2_ne1024pg2.F20TR-SCREAMv1.c10-sep11.out10min.cubesphere.zarr
nvidia/STRATA-SCREAM-sdy1	sdy1.ne1024pg2_ne1024pg2.F2010-SCREAMv1.c10-sep11-f602da2b98.out10min-4lev.cubesphere.zarr
nvidia/STRATA-SCREAM-sdy2	sdy2.ne1024pg2_ne1024pg2.F2010-SCREAMv1.c10-sep11-f602da2b98.out10min-4lev.cubesphere.zarr

ZARR_ROOT cannot be an hf:// prefix. Download whole stores: a chunk that is missing from a partial download reads back as zeros, not as an error.

Auxiliary files (AUX_DATA_ROOT)

Inference and training look up auxiliary files by name in one directory, AUX_DATA_ROOT. Some come from the model repository; generate the rest with data_prep/scrip_generation/:

File	Source	Used for
checkpoints/strata_v1.pth	nvidia/strata	weights, normalization statistics, training config
phis.npy, sgh30.npy, landfrac.npy	nvidia/strata	static forcing fields (inference)
latlon_ne1024pg2.nc	generated by generate_all.sh	model lat/lon conditioning (inference and training)
ne1024pg2_scrip.nc (2.3 GB)	generated by generate_all.sh	cell areas for regional averages, regridding
ne1024halo256pg2_scrip.nc (5.2 GB)	generated by generate_all.sh	halo exchange between cubed-sphere tiles
scream_vertical_coordinate.nc	copied from data/ in this repository	hybrid level coefficients

From the repository root:

hf download nvidia/strata --local-dir /path/to/strata_v1
export AUX_DATA_ROOT=/path/to/strata_v1
bash data_prep/scrip_generation/generate_all.sh --output-dir $AUX_DATA_ROOT
cp data/scream_vertical_coordinate.nc $AUX_DATA_ROOT/

generate_all.sh takes about 20 minutes on a 64-core node and leaves intermediate meshes (*.g, about 7 GB) that can be deleted.

Running the STRATA v1 checkpoint

The forecast reads only its initial condition from the sdecadal store, so the store can stay remote (SCREAM_MAIN_ZARR_PATH=hf://datasets/nvidia/STRATA-SCREAM-sdecadal). To read it locally instead, download it (36.8 TB) and point SCREAM_MAIN_ZARR_PATH at the copy:

hf download nvidia/STRATA-SCREAM-sdecadal --repo-type dataset \
  --local-dir $ZARR_ROOT/sdecadal.ne1024pg2_ne1024pg2.F20TR-SCREAMv1.c10-sep11.out10min.cubesphere.zarr
export SCREAM_MAIN_ZARR_PATH=$ZARR_ROOT/sdecadal.ne1024pg2_ne1024pg2.F20TR-SCREAMv1.c10-sep11.out10min.cubesphere.zarr

The code labels sdecadal time index 0 as 2020-10-01T00:00, so valid initial times run from 2020-10-01T00:00 to 2020-10-15T00:00 in 10-minute steps. The simulation used October 2020 sea-surface temperatures but started from a 1994 atmospheric initial condition, so its weather does not follow October 2020 observations: evaluate forecasts against the SCREAM output, not against observations or reanalysis. The model was trained on 2020-10-02 to 2020-10-10 of this store (plus sdy1 and sdy2), so evaluate from 2020-10-11 onwards. Initial conditions come from sdecadal only.

Forecast

The forecast runs under torchrun with one process per GPU; the commands below use 4 × H100 80 GB. Run them from the repository root in an environment from Setup. The Docker image does not contain this repository, so mount it, AUX_DATA_ROOT and the output directory into the container; on a cluster, use its own container setup for interactive or batch jobs (for example Pyxis/Enroot, Apptainer or Shifter):

export PYTHONPATH=$PWD   # the script imports data_catalog, which is not in the installed package
export AUX_DATA_ROOT=/path/to/strata_v1
export SCREAM_MAIN_ZARR_PATH=hf://datasets/nvidia/STRATA-SCREAM-sdecadal
export OUTPUT=/path/to/output/strata_v1_24h.zarr
torchrun --nproc_per_node=4 scripts/ace/run_screamcast_nudged.py \
  --run-name strata_v1_24h \
  --checkpoint $AUX_DATA_ROOT/checkpoints/strata_v1.pth \
  --forecast-model screamcast \
  --correction none \
  --n-steps 144 \
  --tile-size 128 \
  --halo-width 16 \
  --halo-adjoint \
  --no-halo-adjoint-cross-face \
  --tile-jitter 4 \
  --tile-jitter-mode cycle16 \
  --sht-omega-lmax 3 \
  --inference-batch-size 1 \
  --output-variables T_2m ps U V PotentialTemperature qv omega precip_liq_surf_mass_flux \
  --output-levels 13,18,24,25,29,30 \
  --initial-time 2020-10-13T00:00:00 \
  $OUTPUT

These are the recommended settings for this checkpoint: tile jitter on (4 px, cycle16), no halo adjoint across cube-face edges, and omega filtered with SHT lmax 3. For a quick check, run --n-steps 2 first; a finished run logs Forecast done: 2 steps. Loading the model and the initial condition takes a few minutes, then each 10-minute step takes about 1.5–2.5 minutes on 4 × H100. The network is compiled with torch.compile, which is about a third faster than eager after the first step (TORCH_COMPILE_DISABLE=1 runs it eagerly); Dynamo's recompile_limit warning during the first step is harmless.

On 4 GPUs keep --inference-batch-size 1: batch 4 runs out of memory on 4 × 80 GB. Each written step holds every selected channel on the full grid (about 100 MB per channel before compression), so choose --output-variables and --output-levels deliberately.

Each step logs the global-mean liquid and ice precipitation in mm/day. Ice precipitation reads 0.000 with this checkpoint (a known limitation); it is a diagnostic output and does not feed back into the forecast.

Forecast output

The forecast is written to the Zarr store $OUTPUT, and area-weighted regional means (global, land, ocean, tropics, extratropics) of every output channel to ${OUTPUT}_regional_averages.nc (--no-zarr-output writes only the latter). Arrays in the store are (time, step, face, x, y) for 2-D fields and (time, step, level, face, x, y) for 3-D fields, on the six 2048 × 2048 cube faces:

• time: initial time; step: lead time in seconds (600 is the first 10-minute step);
• level: SCREAM model-level index (8–31, from the top), not pressure; hyam / hybm give the hybrid coefficients, and --output-levels selects levels by these indices;
• lat, lon: (face, x, y) cell centres in degrees.

Variable units are SCREAM's: precipitation fluxes are in m/s. --output-steps N writes every N-th step, at leads N, 2N, … × 10 minutes (--output-steps 6 gives hourly output); --n-steps must be a multiple of N.

Running under SLURM

slurm/submit_inference.sh is an example sbatch launcher for the same forecast on 2 nodes × 8 GPUs, where --inference-batch-size 4 fits: edit the checkpoint and output paths at the top, and submit it from the repository root (it sets PYTHONPATH from there). The path variables can be exported before sbatch or come from .env. On Perlmutter, perlmutter_scripts/strata_v1_rollout.sh runs a 2-step check with the same model settings (eager, regional means only; setup in docs/nersc_setup.md); for a full forecast, raise --n-steps and drop --no-zarr-output. python3 scripts/ace/run_screamcast_nudged.py --help lists all options.

Initial conditions from ERA5 or ECMWF HRES

screamcast.initial_conditions builds STRATA initial conditions (ICs) on the SCREAM ne1024pg2 physics grid with the E3SM L128 vertical grid, for any hour of ERA5 or of the ECMWF HRES analysis (2016-2022). It follows the E3SM HICCUP steps: conservative horizontal remap, surface-pressure adjustment to the SCREAM topography, vertical remap to L128, supersaturation removal, then potential temperature and hydrostatic heights.

• fast_ic.py is a torch port of those steps (HICCUP is BSD-3-licensed; its notice is kept in the file). It reads ERA5 from the public ARCO store and builds one IC in about 80 s on one GPU. Its output matched HICCUP itself bit for bit on ten cases, apart from last-bit differences in a few values.
• hres_ic.py does the same from the WeatherBench2 HRES hres_t0 store.

Each IC is a one-frame Zarr store with flat ncol and the 32 model levels STRATA uses. They are STRATA inputs, not ICs EAMxx can start from (EAMxx reads its ICs on the np4 GLL grid).

One-time inputs: the SCREAM phis, sgh30 and landfrac arrays from the nvidia/STRATA-SCREAM-sdy1 dataset, the L128 coordinate file from HICCUP, and the ERA5 to ne1024pg2 remap map. build_map.py builds the map (3.2 GB) in about 30 min on 72 cores; it needs conda-forge nco and tempest-remap. Then:

python -m screamcast.initial_conditions.fast_ic \
  --timestamp 2017-09-14T00 \
  --out-dir "$WORK/ics" \
  --map "$WORK/map/map_721x1440_to_ne1024pg2.nc" \
  --phis "$WORK/scream_phis.zarr" \
  --latlon "$AUX_DATA_ROOT/latlon_ne1024pg2.nc" --scream-ref "$WORK/scream_static" \
  --vert "$WORK/vert_coord_E3SM_L128.nc"

To forecast from it, run the v1 forecast above with --ic-store in place of SCREAM_MAIN_ZARR_PATH and --initial-time. The forecast starts at the store's initial_time. A one-frame store has no later truth, so --correction none is required:

torchrun --nproc_per_node=4 scripts/ace/run_screamcast_nudged.py \
  --run-name maria_2017091400 \
  --checkpoint $AUX_DATA_ROOT/checkpoints/strata_v1.pth \
  --ic-store "$WORK/ics/scream_hiccup_ne1024pg2_2017091400.zarr" \
  --forecast-model screamcast --correction none --n-steps 2 \
  --tile-size 128 --halo-width 16 --halo-adjoint --no-halo-adjoint-cross-face \
  --tile-jitter 4 --tile-jitter-mode cycle16 --sht-omega-lmax 3 \
  --inference-batch-size 1 --no-zarr-output \
  $OUTPUT

docs/initial_conditions.md has the downloads, the map build, the Docker command, the HRES path and the validation.

Training

Training settings are named Python configs in train_configs.py; see screamcast/config.py for the full config reference. Add an experiment (optionally branching off an existing one with dataclasses.replace()), then launch it on a GPU node, not a login node:

python train.py <config_name>                                  # single GPU
torchrun --nproc_per_node=<num_gpus> train.py <config_name>    # one node, several GPUs

For multi-node jobs, run the same entry point under SLURM srun. For queued SLURM jobs, experiment_manager.py submit --name <run> train.py <config_name> copies the code to $PROJECT_ROOT/<run> and submits it with sbatch through the script named by SUBMIT_SCRIPT; see docs/nersc_setup.md for a Perlmutter example. Checkpoints (best.pth / latest.pth) and logs are written to the config's rundir (default output/).

The released checkpoint was trained with the config pixeldit_sem1024d24l_pix128d4l_bilineardwgeluproject_unfreeze_3src_const1em5_t128_4stepft, which reads all three stores under the directory names in the table above.

Recommended schedule (production checkpoint)

The 5 km production recipe trains in four stages; each resumes the previous stage's best.pth (experiment.resume_from). "Images" counts training tiles (batch size is 1 per GPU, so one step on 2 nodes / 16 GPUs is 16 images).

We recommend training with the precipitation ReLU off (pipeline.do_precip_relu=False) and clipping precipitation at inference instead: with the ReLU inside the forward pass a precipitation head can saturate below zero during training (zero output, zero gradient) and never recover.

1. Pretrain on 64×64 tiles at a learning rate of 1e-4 (5k-step warm-up, then constant) for about 10M images, or until convergence. The production run reached 10M images in total (started on 1 node, moved to 2 nodes at the same learning rate); starting on 2 nodes should need about the same number of images.
2. Anneal on 64×64 tiles: one cosine from 1e-4 down to 1e-6 over the next ~3M images (200k steps on 2 nodes), keeping the optimizer state. In train.py that is scheduler_type="dampened_cosine_with_hard_restarts" with cosine_n_cycles=1, warmup_steps=0 (otherwise the stage-1 warm-up runs again) and cosine_t_max=10 validations × validate_steps=20000; its floor is min_ratio=0.01 × lr = 1e-6. The LR drops to 0 after those 200k steps, so stop there. (scheduler_type="cosine" only decays to 0.1 × lr = 1e-5.)
3. Fine-tune on 128×128 tiles at a constant 1e-5 (2k-step warm-up) with a fresh optimizer. Better: fine-tune directly on tile + 2×halo, the padded tile the rollout feeds the model. This converges within about 0.5-1M images.
4. Multi-step fine-tuning once stage 3 has converged: 2-step, then 4-step, both at a constant 1e-5. add_multistep_finetune_configs("<stage-3 config>") in train_configs.py defines <stage-3 config>_2stepft and _4stepft; run the 2-step config, then the 4-step one (it resumes the 2-step run's best.pth).

To fine-tune from the released strata_v1.pth (via experiment.resume_from), turn off torch.compile (compute.do_torch_compile=False) for now: with compilation on, its weights are not loaded, and the log does not flag this.

Local evaluation

For a quick rollout on a single tile, open notebooks/run_screamcast.py as a marimo notebook (marimo edit notebooks/run_screamcast.py). It loads a checkpoint via ScreamcastModel and rolls out on one tile so you can inspect predictions interactively. Set checkpoint_path in the notebook to $AUX_DATA_ROOT/checkpoints/strata_v1.pth, and set the same environment variables as for the forecast above.

ACE → SCREAM finetuning

The optional ACE→SCREAM forecast-residual workflow is documented in docs/ace2scream_finetuning.md; its scripts live under scripts/ace/.

Developing

GitHub Actions runs make lint on pushes to main and pull requests. To run the same checks locally, install Python 3.12 and run make install-linters first.

make lint     # SPDX license headers, black, and ruff
pytest

License

• Code in this repository: Apache License 2.0 (LICENSE.txt).
• Model weights (nvidia/strata): OpenMDW License Agreement 1.1; its static fields phis.npy, sgh30.npy and landfrac.npy are derived from the SCREAM data and carry CC BY 4.0 attribution (see ATTRIBUTION.md there).
• SCREAM data (raw output and the nvidia/STRATA-SCREAM-* Zarr stores): CC BY 4.0, as E3SM model output (see the E3SM data documentation).

Disclaimer

This project will download and install additional third-party open-source software. Review the license terms of those projects before use.