rollout_to_netcdf_gen2
======================

.. py:module:: rollout_to_netcdf_gen2

.. autoapi-nested-parse::

   rollout_to_netcdf_gen2.py
   -----------------------
   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 LocalDataset + 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
----------

.. autoapisummary::

   rollout_to_netcdf_gen2.logger


Functions
---------

.. autoapisummary::

   rollout_to_netcdf_gen2._inject_flat_schema
   rollout_to_netcdf_gen2._inject_tracer_inds
   rollout_to_netcdf_gen2._build_output_denorm
   rollout_to_netcdf_gen2._sample_to_batch
   rollout_to_netcdf_gen2._save_worker
   rollout_to_netcdf_gen2.predict
   rollout_to_netcdf_gen2._load_model
   rollout_to_netcdf_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 (same as train_gen2.py).


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


.. py:function:: _sample_to_batch(sample)

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


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


.. py:function:: predict(rank, world_size, conf, p)

.. py:function:: _load_model(conf, device)

.. py:function:: main()