Datasets¶

`samudra.datasets` ¶

`InferenceDataset(source, prognostic_var_names, boundary_var_names, input_steps, output_steps, normalize_before_mask, masked_fill_value, long_rollout)` ¶

Bases: Dataset

Dataset of overlapping windows used for autoregressive inference.

Each window contains input_steps historical states followed by output_steps target states. Consecutive windows advance by output_steps so that every model output becomes the newest part of the next input history. For example::

input_steps=1, output_steps=1: [0 | 1], [1 | 2], [2 | 3]
input_steps=2, output_steps=2: [0, 1 | 2, 3], [2, 3 | 4, 5]
input_steps=2, output_steps=1: [0, 1 | 2], [1, 2 | 3], [2, 3 | 4]

The values before | are inputs and those after it are targets.

Source code in src/samudra/datasets.py

@elapsed
def __init__(
    self,
    source: CanonicalSource,
    prognostic_var_names,
    boundary_var_names,
    input_steps,
    output_steps,
    normalize_before_mask,
    masked_fill_value,
    long_rollout,
):
    super().__init__()
    # NOTE: Keep tensors on CPU during initialization. This allows the dataset
    # to be passed between DataLoader worker processes. Call to(device) before
    # using the dataset for inference.

    self.input_steps = input_steps
    self.output_steps = output_steps
    if not 1 <= self.output_steps <= self.input_steps:
        raise ValueError(
            f"output_steps must be between 1 and {self.input_steps}, "
            f"got {self.output_steps}"
        )
    self.prognostic_var_names = tuple(prognostic_var_names)
    self.boundary_var_names = tuple(boundary_var_names)

    self.num_prognostic_channels = self.input_steps * len(prognostic_var_names)
    self.num_output_channels = self.output_steps * len(prognostic_var_names)
    self.input_resolution = source.resolution
    self._device = torch.device("cpu")
    self._source = source
    self._times = source.time
    self.preprocessor = BatchPreprocessor(
        source,
        self.prognostic_var_names,
        self.boundary_var_names,
        normalize_before_mask=normalize_before_mask,
        masked_fill_value=masked_fill_value,
    )

    window_size = self.input_steps + self.output_steps
    num_windows = (source.time.size - window_size) // self.output_steps + 1
    if num_windows <= 0:
        rolling_values = np.empty((0, window_size), dtype=int)
    else:
        starts = np.arange(num_windows) * self.output_steps
        rolling_values = starts[:, None] + np.arange(window_size)[None, :]
    self.rolling_indices = xr.DataArray(
        rolling_values,
        dims=["window_dim", "time"],
    )

    if long_rollout:
        logger.info(
            f"Long rollout will use input at time {source.time.values[0]} and produce"
            f" output at {source.time.values[self.input_steps]}"
        )

    self.wet_label = source.masks.prognostic_for_steps(self.output_steps)
    self.size = len(self.rolling_indices)

    if using_gpu():
        self.wet_label = self.wet_label.pin_memory()

    # Inference only currently supports the same output resolution as the input
    # resolution.
    self.ctx = BatchGrid(
        self.wet_label, self.input_resolution, self.input_resolution
    )

`to(device)` ¶

Move the dataset's context tensors to the specified device.

Call this before using the dataset for inference to ensure tensors are on the correct device (GPU).

Source code in src/samudra/datasets.py

def to(self, device: torch.device) -> "InferenceDataset":
    """Move the dataset's context tensors to the specified device.

    Call this before using the dataset for inference to ensure tensors
    are on the correct device (GPU).
    """
    self.ctx = self.ctx.to(device)
    self._device = device
    self.wet_label = self.wet_label.to(device, non_blocking=True)
    return self

`get_boundary(step)` ¶

Return boundary at the requested step.

Source code in src/samudra/datasets.py

def get_boundary(self, step: int) -> Boundary:
    """Return boundary at the requested step."""
    x_index = self._get_x_index(step)
    boundary = self._get_boundary(x_index)
    return boundary

`HostBatch(dataset_id)` ¶

Source code in src/samudra/datasets.py

def __init__(self, dataset_id: "TorchTrainDataset.Id"):
    self.dataset_id: TorchTrainDataset.Id = dataset_id
    self.steps: list[RolloutStep] = []
    self.load_stats: LoadStats | None = None

`append(input_, boundary, label)` ¶

Add a prognostic input, boundary, and prognostic label as the last step.

Source code in src/samudra/datasets.py

def append(
    self,
    input_: torch.Tensor,
    boundary: torch.Tensor,
    label: torch.Tensor,
):
    """Add a prognostic input, boundary, and prognostic label as the last step."""
    self.steps.append(RolloutStep(input_, boundary, label))

`ModelBatch(ctx)` ¶

A single batch of training data.

A single batch contains multiple steps worth of RolloutStep entries, each of which is a (prognostic_input, boundary_input, label) triple. The prognostic and boundary tensors are carried separately because the samudra-multi model encodes them separately (Samudra just concatenates them later).

Source code in src/samudra/datasets.py

def __init__(self, ctx: BatchGrid):
    self.ctx = ctx
    self.steps: list[RolloutStep] = []
    self.load_stats: LoadStats | None = None

`append(prognostic_input, boundary_input, label)` ¶

Add another RolloutStep as a new step.

Source code in src/samudra/datasets.py

def append(
    self, prognostic_input: Prognostic, boundary_input: Boundary, label: Prognostic
) -> None:
    """Add another RolloutStep as a new step."""
    self.steps.append(RolloutStep(prognostic_input, boundary_input, label))

`getitem(step)` ¶

Converts index (step) into (prognostic, boundary, label) triple.

Source code in src/samudra/datasets.py

def __getitem__(self, step: int) -> RolloutStep:
    """Converts index (step) into (prognostic, boundary, label) triple."""
    return self.steps[step]

`record_stream(stream)` ¶

Keep batch storage alive until work queued on stream completes.

Source code in src/samudra/datasets.py

def record_stream(self, stream: torch.cuda.Stream) -> None:
    """Keep batch storage alive until work queued on ``stream`` completes."""
    self.ctx.label_mask.record_stream(stream)
    for prognostic, boundary, label in self.steps:
        prognostic.record_stream(stream)
        boundary.record_stream(stream)
        label.record_stream(stream)

`TrainBatchLoader` ¶

Bases: Protocol

Training-loop contract shared by Torch and native batch loaders.

`BatchReadUse(source, request, mask, role)` `dataclass` ¶

One canonical source read and its model-facing preparation policy.

`BatchReadPlan(dataset_id, steps)` `dataclass` ¶

Storage-independent reads for a homogeneous batch and full rollout.

`TrainingWindows(input_source, label_source, prognostic_var_names, boundary_var_names, input_steps, output_steps, steps, normalize_before_mask, masked_fill_value, stride=1, windows_id=None)` ¶

Training-window semantics shared by sample and native batch readers.

Source code in src/samudra/datasets.py

def __init__(
    self,
    input_source: CanonicalSource,
    label_source: CanonicalSource | None,
    prognostic_var_names: PrognosticVarNames,
    boundary_var_names: BoundaryVarNames,
    input_steps: int,
    output_steps: int,
    steps: int,
    normalize_before_mask: bool,
    masked_fill_value: float,
    stride: int = 1,
    windows_id: str | None = None,
) -> None:
    self.id = windows_id or f"TrainingWindows_{id(self)}"
    self.input_source = input_source
    self.label_source = label_source or input_source
    self.input_steps = input_steps
    self.output_steps = output_steps
    if not 1 <= self.output_steps <= self.input_steps:
        raise ValueError(
            f"output_steps must be between 1 and {self.input_steps}, "
            f"got {self.output_steps}"
        )
    self.steps = steps
    self.stride = stride
    self.normalize_before_mask = normalize_before_mask
    self.masked_fill_value = masked_fill_value
    self.prognostic_var_names = tuple(prognostic_var_names)
    self.boundary_var_names = tuple(boundary_var_names)
    self.num_prognostic_channels = input_steps * len(prognostic_var_names)
    self.num_output_channels = output_steps * len(prognostic_var_names)
    self.num_boundary_channels = input_steps * len(boundary_var_names)
    if not np.array_equal(self.input_source.time, self.label_source.time):
        raise ValueError("Input and label sources have different time slices")

    total_times = input_steps + output_steps
    num_windows = input_source.time.size - (total_times - 1) * stride
    indices = xr.DataArray(np.arange(num_windows), dims=["window"])
    offsets = xr.DataArray(np.arange(total_times), dims=["time"])
    self.rolling_indices: Int[xr.DataArray, "window time"] = (
        indices + stride * offsets
    )
    self.input_prognostic_mask = input_source.masks.prognostic
    self.label_prognostic_mask = self.label_source.masks.prognostic
    self.boundary_mask = input_source.masks.boundary
    self.ctx = BatchGrid(
        label_mask=self.label_source.masks.prognostic_for_steps(output_steps),
        input_resolution_cpu=input_source.resolution,
        output_resolution_cpu=self.label_source.resolution,
    )
    self.size = (
        input_source.time.size
        - steps * output_steps * stride
        - (input_steps - 1) * stride
    )

`batch_compatibility_key` `property` ¶

Preserve the current homogeneous dataset-ID batching contract.

`window_plan(indices)` ¶

Plan all canonical reads for a batch and full rollout.

Source code in src/samudra/datasets.py

def window_plan(self, indices: list[int]) -> BatchReadPlan:
    """Plan all canonical reads for a batch and full rollout."""
    if not indices:
        raise ValueError("Cannot plan an empty training batch")
    planned_steps = []
    for step in range(self.steps):
        relative = np.stack(
            [self.window_indices(index, step) for index in indices]
        ).astype(np.int64, copy=False)
        current = relative[:, : self.input_steps]
        forecast = relative[:, self.input_steps :]
        planned_steps.append(
            BatchReadStep(
                input=BatchReadUse(
                    self.input_source,
                    CanonicalReadRequest(current, self.prognostic_var_names),
                    self.input_prognostic_mask,
                    "input",
                ),
                boundary=BatchReadUse(
                    self.input_source,
                    CanonicalReadRequest(current, self.boundary_var_names),
                    self.boundary_mask,
                    "boundary",
                ),
                label=BatchReadUse(
                    self.label_source,
                    CanonicalReadRequest(forecast, self.prognostic_var_names),
                    self.label_prognostic_mask,
                    "label",
                ),
            )
        )
    return BatchReadPlan(self.id, tuple(planned_steps))

`BatchPreparer(windows)` ¶

Rank-local normalization, masking, shaping, and device-static caches.

Source code in src/samudra/datasets.py

def __init__(self, windows: TrainingWindows) -> None:
    self.windows = windows
    self._input = BatchPreprocessor(
        windows.input_source,
        windows.prognostic_var_names,
        windows.boundary_var_names,
        normalize_before_mask=windows.normalize_before_mask,
        masked_fill_value=windows.masked_fill_value,
    )
    self._label = (
        self._input
        if windows.label_source is windows.input_source
        else BatchPreprocessor(
            windows.label_source,
            windows.prognostic_var_names,
            windows.boundary_var_names,
            normalize_before_mask=windows.normalize_before_mask,
            masked_fill_value=windows.masked_fill_value,
        )
    )
    self._device_ctx: dict[torch.device, BatchGrid] = {}

`normalize_and_mask_device_planes(use, data, device)` ¶

Prepare unique (time, variable, lat, lon) planes on device.

Source code in src/samudra/datasets.py

def normalize_and_mask_device_planes(
    self,
    use: BatchReadUse,
    data: torch.Tensor,
    device: torch.device,
) -> torch.Tensor:
    """Prepare unique ``(time, variable, lat, lon)`` planes on device."""
    if data.device.type != device.type or (
        device.index is not None and data.device.index != device.index
    ):
        raise ValueError(
            f"Unique ocean planes are on {data.device}; expected {device}"
        )

    preprocessor = self._label if use.role == "label" else self._input
    return preprocessor.prepare_planes(
        data, device, boundary=use.role == "boundary"
    )

`TorchTrainDataset(windows, *, concurrent_compute_=False)` ¶

Bases: Dataset[HostBatch]

Adapt training windows to PyTorch's per-sample Dataset interface.

Source code in src/samudra/datasets.py

@elapsed
def __init__(self, windows: TrainingWindows, *, concurrent_compute_: bool = False):
    super().__init__()
    self.windows = windows
    self.preparer = BatchPreparer(windows)
    self.id = windows.id
    self._concurrent_compute = concurrent_compute_

`to_model_batch(host_batch, device)` ¶

Convert HostBatch to ModelBatch, moving tensors to the specified device.

Parameters:

Name	Type	Description	Default
`host_batch`	`HostBatch`	CPU data from worker process	required
`device`	`device`	Target device (typically GPU) to move tensors to	required

Returns:

Type	Description
`ModelBatch`	ModelBatch with tensors on the target device

Source code in src/samudra/datasets.py

def to_model_batch(self, host_batch: HostBatch, device: torch.device) -> ModelBatch:
    """Convert HostBatch to ModelBatch, moving tensors to the specified device.

    Args:
        host_batch: CPU data from worker process
        device: Target device (typically GPU) to move tensors to

    Returns:
        ModelBatch with tensors on the target device
    """
    return self.preparer.prepare_host_batch(host_batch, device)

`TorchBatchLoader(host_loader, datasets, device)` ¶

Wrapper around a torch DataLoader that handles GPU post-processing.

This class wraps a DataLoader[HostBatch] and converts the raw data to ModelBatch by applying GPU-based normalization and masking. This allows the data loading process to handle I/O while the main process handles GPU operations.

Since the data samples flow from one process to the other, we want to tie them back to the dataset they came from which knows how to do that second half once they're in the main process which has GPU access set up. To do that, each data sample (which could come from a different dataset) has a dataset ID -- datasets maps from those IDs to the original datasets.

Source code in src/samudra/datasets.py

def __init__(
    self,
    host_loader: torch.utils.data.DataLoader[HostBatch],
    datasets: list[TorchTrainDataset],
    device: torch.device,
):
    self._host_loader = host_loader
    self._datasets = {dataset.id: dataset for dataset in datasets}
    self._device = device

`iter()` ¶

Iterate over the dataloader, converting HostBatch to ModelBatch.

Source code in src/samudra/datasets.py

def __iter__(self):
    """Iterate over the dataloader, converting HostBatch to ModelBatch."""
    for host_batch in self._host_loader:
        dataset = self._datasets[host_batch.dataset_id]
        model_batch = dataset.to_model_batch(host_batch, self._device)
        yield model_batch

`close()` ¶

Release persistent PyTorch workers owned by this loader.

Source code in src/samudra/datasets.py

def close(self) -> None:
    """Release persistent PyTorch workers owned by this loader."""
    close_pytorch_dataloader(self._host_loader)

`getitem(index)` ¶

Access a single item by index, converting HostBatch to ModelBatch.

Note: This bypasses the DataLoader's sampling/batching and directly accesses the underlying dataset for test purposes.

Source code in src/samudra/datasets.py

def __getitem__(self, index: int) -> ModelBatch:
    """Access a single item by index, converting HostBatch to ModelBatch.

    Note: This bypasses the DataLoader's sampling/batching and directly accesses
    the underlying dataset for test purposes.
    """
    # Access the underlying dataset directly
    host_batch = self._host_loader.dataset[index]
    # Apply the collate function to add batch dimension (expects a list)
    collate_fn = self._host_loader.collate_fn
    if collate_fn is not None:
        host_batch = collate_fn([host_batch])
    # Get the dataset that created this raw data
    dataset = self._datasets[host_batch.dataset_id]
    # Convert to ModelBatch
    model_batch = dataset.to_model_batch(host_batch, self._device)
    return model_batch

`close_pytorch_dataloader(dataloader)` ¶

Deterministically stop persistent workers created by a PyTorch loader.

Source code in src/samudra/datasets.py

def close_pytorch_dataloader(dataloader: torch.utils.data.DataLoader) -> None:
    """Deterministically stop persistent workers created by a PyTorch loader."""
    iterator = dataloader._iterator
    if iterator is not None:
        shutdown_workers = getattr(iterator, "_shutdown_workers", None)
        if shutdown_workers is not None:
            shutdown_workers()
        dataloader._iterator = None

Datasets¶

samudra.datasets ¶

InferenceDataset(source, prognostic_var_names, boundary_var_names, input_steps, output_steps, normalize_before_mask, masked_fill_value, long_rollout) ¶

to(device) ¶

get_boundary(step) ¶

HostBatch(dataset_id) ¶

append(input_, boundary, label) ¶

ModelBatch(ctx) ¶

append(prognostic_input, boundary_input, label) ¶

__getitem__(step) ¶

record_stream(stream) ¶

TrainBatchLoader ¶

BatchReadUse(source, request, mask, role) dataclass ¶

BatchReadPlan(dataset_id, steps) dataclass ¶

TrainingWindows(input_source, label_source, prognostic_var_names, boundary_var_names, input_steps, output_steps, steps, normalize_before_mask, masked_fill_value, stride=1, windows_id=None) ¶

batch_compatibility_key property ¶

window_plan(indices) ¶

BatchPreparer(windows) ¶

normalize_and_mask_device_planes(use, data, device) ¶

TorchTrainDataset(windows, *, concurrent_compute_=False) ¶

to_model_batch(host_batch, device) ¶

TorchBatchLoader(host_loader, datasets, device) ¶

__iter__() ¶

close() ¶

__getitem__(index) ¶

close_pytorch_dataloader(dataloader) ¶

`samudra.datasets` ¶

`InferenceDataset(source, prognostic_var_names, boundary_var_names, input_steps, output_steps, normalize_before_mask, masked_fill_value, long_rollout)` ¶

`to(device)` ¶

`get_boundary(step)` ¶

`HostBatch(dataset_id)` ¶

`append(input_, boundary, label)` ¶

`ModelBatch(ctx)` ¶

`append(prognostic_input, boundary_input, label)` ¶

`getitem(step)` ¶

`record_stream(stream)` ¶

`TrainBatchLoader` ¶

`BatchReadUse(source, request, mask, role)` `dataclass` ¶

`BatchReadPlan(dataset_id, steps)` `dataclass` ¶

`TrainingWindows(input_source, label_source, prognostic_var_names, boundary_var_names, input_steps, output_steps, steps, normalize_before_mask, masked_fill_value, stride=1, windows_id=None)` ¶

`batch_compatibility_key` `property` ¶

`window_plan(indices)` ¶

`BatchPreparer(windows)` ¶

`normalize_and_mask_device_planes(use, data, device)` ¶

`TorchTrainDataset(windows, *, concurrent_compute_=False)` ¶

`to_model_batch(host_batch, device)` ¶

`TorchBatchLoader(host_loader, datasets, device)` ¶

`iter()` ¶

`close()` ¶

`getitem(index)` ¶

`close_pytorch_dataloader(dataloader)` ¶