rollout_realtime_gen2
=====================

.. py:module:: rollout_realtime_gen2

.. autoapi-nested-parse::

   rollout_realtime_gen2.py
   ----------------------
   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
----------

.. autoapisummary::

   rollout_realtime_gen2.logger


Functions
---------

.. autoapisummary::

   rollout_realtime_gen2._inject_flat_schema
   rollout_realtime_gen2._inject_tracer_inds
   rollout_realtime_gen2._build_output_denorm
   rollout_realtime_gen2._sample_to_batch
   rollout_realtime_gen2._save_worker
   rollout_realtime_gen2.run_gfs_init
   rollout_realtime_gen2.run_forecast
   rollout_realtime_gen2.main


Module Contents
---------------

.. py:data:: logger

.. py:function:: _inject_flat_schema(conf)

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


.. py:function:: _inject_tracer_inds(conf)

   Compute tracer_inds for TracerFixer from v2 variable layout.


.. py:function:: _build_output_denorm(conf, device, dtype=torch.float32)

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


.. py:function:: _sample_to_batch(sample)

   Add batch dim and wrap LocalDataset sample for preblock input.


.. py:function:: _save_worker(shm_name, arr_shape, arr_dtype, init_str, step, fhr_per_step, lat, lon, meta_data, conf)

.. py:function:: 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.


.. py:function:: 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.

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


.. py:function:: main()