rollout_realtime_gen2#

Operational/realtime forecast rollout for CREDIT v2 models.

Designed to be run with a single command — no pre-editing of config files:

python rollout_realtime_gen2.py
-c config/wxformer_025deg_6hr_gen2.yml –init-time 2024-01-15T00 –steps 40 –save-dir /path/to/output

All forecast-specific parameters (init time, steps, save dir) come from CLI args and override whatever is in the config. The config only needs to describe the model, the data source paths, and the normalization files.

The script uses LocalDataset directly — the same dataset class used for training — so there is no separate “predict dataset” to maintain.

Output: one NetCDF file per forecast step saved to: <save_dir>/<YYYY-MM-DDTHH>Z/pred_<YYYY-MM-DDTHH>Z_<FHR:03d>.nc

Attributes#

logger

Functions#

`_inject_flat_schema`(conf)	Inject v1-style flat keys into conf['data'] so output.py utilities work.
`_inject_tracer_inds`(conf)	Compute tracer_inds for TracerFixer from v2 variable layout.
`_build_output_denorm`(conf, device[, dtype])	Return (mean, std) of shape (1, C_out, 1, 1, 1) for inverse-normalizing y_pred.
`_sample_to_batch`(sample)	Add batch dim and wrap LocalDataset sample for preblock input.
`_save_worker`(shm_name, arr_shape, arr_dtype, init_str, ...)
`run_gfs_init`(→ str)	Download GFS analysis for init_time, regrid to CREDIT grid, save as zarr.
`run_forecast`(conf, init_time, n_steps, save_dir, pool)	Run a single autoregressive forecast from init_time for n_steps steps.
`main`()

Module Contents#

rollout_realtime_gen2.logger#

rollout_realtime_gen2._inject_flat_schema(conf)#: Inject v1-style flat keys into conf[‘data’] so output.py utilities work.

rollout_realtime_gen2._inject_tracer_inds(conf)#: Compute tracer_inds for TracerFixer from v2 variable layout.

rollout_realtime_gen2._build_output_denorm(conf, device, dtype=torch.float32)#: Return (mean, std) of shape (1, C_out, 1, 1, 1) for inverse-normalizing y_pred.

rollout_realtime_gen2._sample_to_batch(sample)#: Add batch dim and wrap LocalDataset sample for preblock input.

rollout_realtime_gen2._save_worker(shm_name, arr_shape, arr_dtype, init_str, step, fhr_per_step, lat, lon, meta_data, conf)#

rollout_realtime_gen2.run_gfs_init(conf, init_time: pandas.Timestamp, n_procs: int = 1) → str#

Download GFS analysis for init_time, regrid to CREDIT grid, save as zarr.

Patches conf[‘data’][‘source’][‘ERA5’][‘variables’][‘prognostic’][‘path’] to the generated zarr so LocalDataset loads the GFS IC at step 0. Returns the zarr path.

rollout_realtime_gen2.run_forecast(conf, init_time: pandas.Timestamp, n_steps: int, save_dir: str, pool, rank=0, world_size=1)#

Run a single autoregressive forecast from init_time for n_steps steps.

Parameters:

conf – Full configuration dict (v2 schema, flat keys already injected).
init_time – Forecast initialization timestamp.
n_steps – Number of autoregressive steps to run.
save_dir – Directory for output NetCDF files.
pool – multiprocessing.Pool for async saves.
rank/world_size – For DDP; single-GPU callers use (0, 1).

rollout_realtime_gen2.main()#