rollout_to_netcdf

Rollout entry point for the v2 data schema (conf[“data”][“source”]).

Key differences from rollout_to_netcdf.py (v1):

• No credit_main_parser / predict_data_check (those are v1-only)

• Uses ERA5Dataset + build_preblocks/apply_preblocks for data loading

• Inverse-normalization built from mean/std NC files (per-channel)

• Injects flat data schema keys for output.py compatibility

• forecast_len semantics: N steps at lead_time_periods hours per step

Usage:

torchrun –standalone –nnodes=1 –nproc-per-node=<N_GPUS>

applications/rollout_to_netcdf_gen2.py -c config/wxformer_025deg_6hr_gen2.yml

Or submit via PBS (see scripts/casper_gen2.sh with SCRIPT=applications/rollout_to_netcdf_gen2.py).

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 (same as train_gen2.py).
_build_output_denorm(conf, device[, dtype])	Build (mean, std) tensors of shape (1, C_out, 1, 1, 1) for inverse-normalizing y_pred.
_sample_to_batch(sample)	Wrap a single LocalDataset sample (no batch dim) into preblock-compatible dict.
_save_worker(y_pred_np, init_datetime_str, ...)	Called via pool.apply_async — converts numpy array to xarray and saves.
predict(rank, world_size, conf, p)
_load_model(conf, device)
main()

Module Contents

rollout_to_netcdf.logger

rollout_to_netcdf._inject_flat_schema(conf)

Inject v1-style flat keys into conf[‘data’] so output.py utilities work.

rollout_to_netcdf._inject_tracer_inds(conf)

Compute tracer_inds for TracerFixer from v2 variable layout (same as train_gen2.py).

rollout_to_netcdf._build_output_denorm(conf, device, dtype=torch.float32)

Build (mean, std) tensors of shape (1, C_out, 1, 1, 1) for inverse-normalizing y_pred.

Channel order follows LocalDataset target insertion order:

prognostic/3d (each var × n_levels), prognostic/2d, diagnostic/2d

rollout_to_netcdf._sample_to_batch(sample)

Wrap a single LocalDataset sample (no batch dim) into preblock-compatible dict.

rollout_to_netcdf._save_worker(y_pred_np, init_datetime_str, forecast_step, lead_time_periods, lat, lon, meta_data, conf)

Called via pool.apply_async — converts numpy array to xarray and saves.

rollout_to_netcdf.predict(rank, world_size, conf, p)

rollout_to_netcdf._load_model(conf, device)

rollout_to_netcdf.main()