neural_lam.datastore.npyfilesmeps.compute_standardization_stats
===============================================================

.. py:module:: neural_lam.datastore.npyfilesmeps.compute_standardization_stats

.. autoapi-nested-parse::

   Utilities for computing MEPS datastore standardization statistics.



Classes
-------

.. autoapisummary::

   neural_lam.datastore.npyfilesmeps.compute_standardization_stats.PaddedWeatherDataset


Functions
---------

.. autoapisummary::

   neural_lam.datastore.npyfilesmeps.compute_standardization_stats.cli
   neural_lam.datastore.npyfilesmeps.compute_standardization_stats.get_rank
   neural_lam.datastore.npyfilesmeps.compute_standardization_stats.get_world_size
   neural_lam.datastore.npyfilesmeps.compute_standardization_stats.main
   neural_lam.datastore.npyfilesmeps.compute_standardization_stats.save_stats
   neural_lam.datastore.npyfilesmeps.compute_standardization_stats.setup


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

.. py:class:: PaddedWeatherDataset(base_dataset: neural_lam.WeatherDataset, world_size: int, batch_size: int)

   Bases: :py:obj:`torch.utils.data.Dataset`


   Wrap :class:`WeatherDataset` to pad samples for distributed runners.

   :param base_dataset: Dataset to pad.
   :type base_dataset: WeatherDataset
   :param world_size: Total number of distributed ranks participating.
   :type world_size: int
   :param batch_size: Per-rank batch size.
   :type batch_size: int


   .. py:method:: real_sample_masks(sampler: torch.utils.data.distributed.DistributedSampler) -> list[torch.Tensor]

      Return one boolean mask per minibatch, ``True`` for non-padded rows.

      ``DistributedSampler`` stripes indices across ranks, so the padded
      tail of this dataset lands on several ranks and cannot be sliced off
      the rank-major gathered tensor. Each rank drops its padded rows with
      these masks instead; the sampler order chunked by ``batch_size`` is
      the order a ``DataLoader(drop_last=False)`` yields batches in.

      :param sampler: Sampler of this rank's ``DataLoader`` over this dataset.
      :type sampler: DistributedSampler

      :returns: Boolean mask per minibatch.
      :rtype: list[torch.Tensor]



   .. py:attribute:: base_dataset


   .. py:attribute:: batch_size


   .. py:attribute:: padded_samples


   .. py:attribute:: total_samples


   .. py:attribute:: world_size


.. py:function:: cli() -> None

   Parse CLI arguments and trigger :func:`main`.


.. py:function:: get_rank() -> int

   Return the rank inferred from SLURM or default to 0.

   :returns: The current process rank.
   :rtype: int


.. py:function:: get_world_size() -> int

   Return the world size inferred from SLURM or default to 1.

   :returns: The total number of processes in the distributed group.
   :rtype: int


.. py:function:: main(datastore_config_path: str | pathlib.Path, batch_size: int, step_length: datetime.timedelta, n_workers: int, distributed: bool) -> None

   Pre-compute and persist standardization statistics from the datastore.

   :param datastore_config_path: Path to the MEPS datastore configuration file.
   :type datastore_config_path: str or pathlib.Path
   :param batch_size: Batch size used while iterating through the dataset.
   :type batch_size: int
   :param step_length: Temporal sampling interval for the difference statistics.
   :type step_length: datetime.timedelta
   :param n_workers: Number of dataloader workers.
   :type n_workers: int
   :param distributed: If ``True``, run using torch.distributed with SLURM settings.
   :type distributed: bool


.. py:function:: save_stats(static_dir_path: str | pathlib.Path, means: list[torch.Tensor], squares: list[torch.Tensor], flux_means: list[torch.Tensor], flux_squares: list[torch.Tensor], filename_prefix: str) -> None

   Aggregate running statistics and persist them to ``static_dir_path``.

   :param static_dir_path: Directory where ``*.pt`` files should be written.
   :type static_dir_path: str or pathlib.Path
   :param means: Shape ``(B, num_state_vars)``. Batch-wise means. Sequence of one or
                 more tensors each of shape ``(B_i, num_state_vars)``; concatenated /
                 stacked along the batch dim before reduction. Typical callers pass a
                 single already-gathered tensor. The list can have length 0 to skip
                 saving.
   :type means: Sequence[torch.Tensor]
   :param squares: Shape ``(B, num_state_vars)``. Batch-wise second moments. Sequence of
                   one or more tensors each of shape ``(B_i, num_state_vars)``;
                   concatenated / stacked along the batch dim before reduction. Typical
                   callers pass a single already-gathered tensor. The list can have
                   length 0 to skip saving.
   :type squares: Sequence[torch.Tensor]
   :param flux_means: Shape ``(B,)``. Flux means. Sequence of one or more tensors each of
                      shape ``(B_i,)``; concatenated / stacked along the batch dim before
                      reduction. Typical callers pass a single already-gathered tensor.
                      The list can have length 0 to skip saving.
   :type flux_means: Sequence[torch.Tensor]
   :param flux_squares: Shape ``(B,)``. Flux second moments. Sequence of one or more tensors
                        each of shape ``(B_i,)``; concatenated / stacked along the batch dim
                        before reduction. Typical callers pass a single already-gathered
                        tensor. The list can have length 0 to skip saving.
   :type flux_squares: Sequence[torch.Tensor]
   :param filename_prefix: Prefix (e.g., ``"parameter"`` or ``"diff"``) for saved tensors.
   :type filename_prefix: str


.. py:function:: setup(rank: int, world_size: int) -> None

   Initialize the distributed group.

   :param rank: The rank of the current process.
   :type rank: int
   :param world_size: The total number of processes.
   :type world_size: int

   :raises RuntimeError: If ``SLURM_JOB_NODELIST`` is set but no hostnames can be retrieved.