neural_lam.datastore.npyfilesmeps.compute_standardization_stats

Utilities for computing MEPS datastore standardization statistics.

Classes

PaddedWeatherDataset

Wrap WeatherDataset to pad samples for distributed runners.

Functions

cli(→ None)	Parse CLI arguments and trigger main().
get_rank(→ int)	Return the rank inferred from SLURM or default to 0.
get_world_size(→ int)	Return the world size inferred from SLURM or default to 1.
main(→ None)	Pre-compute and persist standardization statistics from the datastore.
save_stats(→ None)	Aggregate running statistics and persist them to static_dir_path.
setup(→ None)	Initialize the distributed group.

Module Contents

class neural_lam.datastore.npyfilesmeps.compute_standardization_stats.PaddedWeatherDataset(base_dataset: neural_lam.WeatherDataset, world_size: int, batch_size: int)

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.

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.

Parameters:

sampler (DistributedSampler) – Sampler of this rank’s DataLoader over this dataset.

Returns:

Boolean mask per minibatch.

Return type:

list[torch.Tensor]

base_dataset

batch_size

padded_samples

total_samples

world_size

neural_lam.datastore.npyfilesmeps.compute_standardization_stats.cli() → None

Parse CLI arguments and trigger main().

neural_lam.datastore.npyfilesmeps.compute_standardization_stats.get_rank() → int

Return the rank inferred from SLURM or default to 0.

Returns:

The current process rank.

Return type:

int

neural_lam.datastore.npyfilesmeps.compute_standardization_stats.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.

Return type:

int

neural_lam.datastore.npyfilesmeps.compute_standardization_stats.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.

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

Parameters:

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

• means (Sequence[torch.Tensor]) – 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.

• squares (Sequence[torch.Tensor]) – 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.

• flux_means (Sequence[torch.Tensor]) – 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.

• flux_squares (Sequence[torch.Tensor]) – 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.

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

neural_lam.datastore.npyfilesmeps.compute_standardization_stats.setup(rank: int, world_size: int) → None

Initialize the distributed group.

Parameters:

• rank (int) – The rank of the current process.

• world_size (int) – The total number of processes.

Raises:

RuntimeError – If SLURM_JOB_NODELIST is set but no hostnames can be retrieved.