neural_lam.datastore.npyfilesmeps.compute_standardization_stats

neural_lam.datastore.npyfilesmeps.compute_standardization_stats#

Utilities for computing MEPS datastore standardization statistics.

Module Contents#

class neural_lam.datastore.npyfilesmeps.compute_standardization_stats.PaddedWeatherDataset(base_dataset, world_size, batch_size)#

Bases: torch.utils.data.Dataset

Wrap WeatherDataset to pad samples for distributed runners.

Parameters:

  • base_dataset (WeatherDataset) – Dataset to pad.

  • world_size (int) – Total number of distributed ranks participating.

  • batch_size (int) – Per-rank batch size.

get_original_indices()#

Return indices of the non-padded samples.

get_original_window_indices(step_length)#

Return index mapping for sub-sampled windows at step_length.

base_dataset#
batch_size#
original_indices#
padded_indices#
padded_samples#
total_samples#
world_size#
neural_lam.datastore.npyfilesmeps.compute_standardization_stats.cli()#

Parse CLI arguments and trigger main().

neural_lam.datastore.npyfilesmeps.compute_standardization_stats.get_rank()#

Return the rank inferred from SLURM or default to 0.

neural_lam.datastore.npyfilesmeps.compute_standardization_stats.get_world_size()#

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

neural_lam.datastore.npyfilesmeps.compute_standardization_stats.main(datastore_config_path, batch_size, step_length, n_workers, distributed)#

Pre-compute and persist standardization statistics from the datastore.

Parameters:

  • datastore_config_path (str or pathlib.Path) – Path to the MEPS datastore configuration file.

  • batch_size (int) – Batch size used while iterating through the dataset.

  • step_length (datetime.timedelta) – Temporal sampling interval for the difference statistics.

  • n_workers (int) – Number of dataloader workers.

  • distributed (bool) – If True, run using torch.distributed with SLURM settings.

neural_lam.datastore.npyfilesmeps.compute_standardization_stats.save_stats(static_dir_path, means, squares, flux_means, flux_squares, filename_prefix)#

Aggregate running statistics and persist them to static_dir_path.

Parameters:

  • static_dir_path (str or pathlib.Path) – Directory where *.pt files should be written.

  • means (Sequence[torch.Tensor]) – Batch-wise means with shape (N_batch, d_features).

  • squares (Sequence[torch.Tensor]) – Batch-wise second moments with shape (N_batch, d_features).

  • flux_means (Sequence[torch.Tensor]) – Optional flux means of shape (N_batch,).

  • flux_squares (Sequence[torch.Tensor]) – Optional flux second moments of shape (N_batch,).

  • filename_prefix (str) – Prefix (e.g., "parameter" or "diff") for saved tensors.

neural_lam.datastore.npyfilesmeps.compute_standardization_stats.setup(rank, world_size)#

Initialize the distributed group.

Contents