datasets.py

# SPDX-FileCopyrightText: 2026 Samudra Authors
#
# SPDX-License-Identifier: Apache-2.0

import logging
import time
from concurrent.futures import wait
from concurrent.futures.thread import ThreadPoolExecutor
from typing import ClassVar, final

import numpy as np
import torch
import xarray as xr
from einops import rearrange
from jaxtyping import Float
from torch.utils.data import Dataset
from xarray_einstats.einops import rearrange as xr_rearrange  # noqa: F401

from samudra.constants import (
    Boundary,
    BoundaryVarNames,
    Example,
    GridMask,
    Input,
    LoaderVersion,
    Prognostic,
    PrognosticMask,
    PrognosticVarNames,
)
from samudra.utils.ctx import GridContext
from samudra.utils.data import DataSource, LoadStats, OceanData, conditional_rearrange
from samudra.utils.device import using_gpu
from samudra.utils.logging import elapsed

logger = logging.getLogger(__name__)


class InferenceDataset(Dataset):
    """This class is used for inference rollouts.

    It creates rolling indices to keep track of histories/past states.
    For example,
    Hist=0 ; 0->[0, 1]; 1->[1, 2]; 2->[2, 3]; 3->[3, 4]
    Hist=1 ; 0->[[0, 1], [2, 3]]; 1->[[2, 3], [4, 5]];
            2->[[4, 5], [6, 7]]; 3->[[6, 7], [8, 9]]
    Hist=2 ; 0->[[0, 1, 2], [3, 4, 5]];
            1->[[3, 4, 5], [6, 7, 8]];
            2->[[6, 7, 8], [9, 10, 11]];
            3->[[9, 10, 11], [12, 13, 14]]
    """

    @elapsed
    def __init__(
        self,
        src: DataSource,
        prognostic_var_names,
        boundary_var_names,
        hist,
        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.hist = hist

        self.num_prognostic_channels = (hist + 1) * len(prognostic_var_names)
        data = src.data
        self.input_res = src.resolution
        self._prognostic_src = src.filter(prognostic_var_names, prefix="prognostic")
        self._boundary_src = src.filter(boundary_var_names, prefix="boundary")
        self._times = data.time
        self.normalize_before_mask = normalize_before_mask
        self.masked_fill_value = masked_fill_value

        time_indices = np.arange(data.time.size)
        indices = xr.DataArray(
            time_indices,
            dims=["time"],
            coords={"time": time_indices},
        )
        total_steps = 2 * self.hist + 1
        rolling_indices = indices.rolling(
            time=len(time_indices) - total_steps, center=False
        ).construct("window_dim")
        rolling_indices = rolling_indices.transpose("window_dim", "time").isel(
            time=slice(len(time_indices) - total_steps - 1, None)
        )  # Remove first few null indices
        self.rolling_indices = rolling_indices.isel(
            window_dim=slice(0, None, self.hist + 1)
        )  # Skip indices based on history
        self.rolling_indices = self.rolling_indices.astype(int)

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

        self.wet: PrognosticMask = src.masks.prognostic
        self.wet_surface: GridMask = src.masks.boundary
        self.wet_label = src.masks.prognostic_with_hist(self.hist)
        self.size = len(self.rolling_indices)

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

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

    def __len__(self):
        return self.size

    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.wet_label = self.wet_label.to(device, non_blocking=True)
        return self

    @property
    def initial_prognostic(self):
        x_index = self._get_x_index(0)
        data_in = self._get_prognostic(x_index)
        return data_in

    def inference_target(self, step: int | slice):
        x_index = self._get_x_index(step)
        label = self._get_label(x_index)
        return label

    def get_initial_input(self) -> tuple[Prognostic, Boundary]:
        prog, boundary, _ = self.__getitem__(0)
        return prog, boundary

    def get_target_time(self, start_step: int, num_steps: int):
        x_index = self._get_x_index(start_step)
        batch_index = x_index.values[0]
        steps_predicted = len(batch_index) // 2
        start_target_index = batch_index[steps_predicted]

        return self._times.isel(
            time=slice(
                start_target_index, start_target_index + num_steps * steps_predicted
            )
        )

    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

    @elapsed(level=logging.DEBUG)
    def __getitem__(self, idx):
        x_index = self._get_x_index(idx)
        data_in_prog = self._get_prognostic(x_index)
        data_in_boundary = self._get_boundary(x_index)
        label = self._get_label(x_index)
        return (data_in_prog, data_in_boundary, label)

    def _get_x_index(self, idx):
        if isinstance(idx, slice):
            if (
                (idx.start is not None and idx.start < 0)
                or (idx.stop is not None and idx.stop < 0)
                or (idx.step is not None and idx.step < 0)
            ):
                raise IndexError("Sorry, negative indexing is not supported!")
            if idx.step is None:
                idx = slice(idx.start, idx.stop, 1)
            if idx.start is None and idx.stop is None:
                idx = slice(0, self.size, idx.step)
            elif idx.start is None:
                idx = slice(0, idx.stop, idx.step)
            elif idx.stop is None:
                idx = slice(idx.start, self.size, idx.step)
        elif isinstance(idx, int):
            if idx < 0:
                raise IndexError("Sorry, negative indexing is not supported!")
            elif idx >= self.size:
                raise IndexError(f"Index {idx} out of range with size {self.size}")
            idx = slice(idx, idx + 1, 1)

        rolling_idx = self.rolling_indices.isel(window_dim=idx)
        x_index = xr.Variable(["window_dim", "time"], rolling_idx)

        return x_index

    def _get_prognostic(self, x_index):
        data_in_src = self._prognostic_src.map_data(
            lambda ds: ds.isel(time=x_index).isel(time=slice(None, self.hist + 1))
        )
        if self.normalize_before_mask:
            data_in_ds = data_in_src.normalize()
        else:
            data_in_ds = data_in_src.data

        if "lev" in data_in_ds.dims:
            data_in_np: np.ndarray = (
                conditional_rearrange(
                    data_in_ds,
                    "window_dim time (variable lev)=var lat lon",
                    concat_dim="var",
                )
                .rename({"var": "variable"})
                .to_numpy()
            )
        else:
            data_in_np = (
                data_in_ds.to_array()
                .transpose("window_dim", "time", "variable", "lat", "lon")
                .to_numpy()
            )
        data_in: torch.Tensor = torch.from_numpy(data_in_np).float()
        data_in = torch.where(self.wet, data_in, self.masked_fill_value)
        if not self.normalize_before_mask:
            data_in = self._prognostic_src.normalize_with(data_in, variable_axis=2)
        data_in = rearrange(
            data_in,
            "window_dim time variable lat lon -> window_dim (time variable) lat lon",
        )
        return data_in

    def _get_boundary(self, x_index):
        """
        This function returns the boundary condition for the current time step.

        With hist > 0, the boundary condition considered is always the last step of
        the input.
        """
        data_in_boundary_src = self._boundary_src.map_data(
            lambda ds: ds.isel(time=x_index).isel(time=slice(None, self.hist + 1))
        )
        if self.normalize_before_mask:
            data_in_boundary_ds = data_in_boundary_src.normalize()
        else:
            data_in_boundary_ds = data_in_boundary_src.data
        data_in_boundary_np: np.ndarray = (
            data_in_boundary_ds.to_array()
            .transpose("window_dim", "time", "variable", "lat", "lon")
            .to_numpy()
        )
        data_in_boundary: torch.Tensor = torch.from_numpy(data_in_boundary_np).float()
        data_in_boundary = torch.where(
            self.wet_surface, data_in_boundary, self.masked_fill_value
        )
        if not self.normalize_before_mask:
            data_in_boundary = self._boundary_src.normalize_with(
                data_in_boundary, variable_axis=2
            )
        data_in_boundary = rearrange(
            data_in_boundary,
            "window_dim time variable lat lon -> window_dim (time variable) lat lon",
        )
        return data_in_boundary

    def _get_label(self, x_index):
        label_src = self._prognostic_src.map_data(
            lambda ds: ds.isel(time=x_index).isel(time=slice(self.hist + 1, None))
        )
        if self.normalize_before_mask:
            label_ds = label_src.normalize()
        else:
            label_ds = label_src.data
        if "lev" in label_ds.dims:
            label_np: np.ndarray = (
                conditional_rearrange(
                    label_ds,
                    "window_dim time (variable lev)=var lat lon",
                    concat_dim="var",
                )
                .rename({"var": "variable"})
                .to_numpy()
            )
        else:
            label_np = (
                label_ds.to_array()
                .transpose("window_dim", "time", "variable", "lat", "lon")
                .to_numpy()
            )
        label: torch.Tensor = torch.from_numpy(label_np).float()
        label = torch.where(self.wet, label, self.masked_fill_value)
        if not self.normalize_before_mask:
            label = self._prognostic_src.normalize_with(label, variable_axis=2)
        label = rearrange(
            label,
            "window_dim time variable lat lon -> window_dim (time variable) lat lon",
        )
        return label

    def get_coords_dict(self):
        return {
            co: self._prognostic_src.data[co] for co in self._prognostic_src.data.coords
        }


class InferenceDatasets(Dataset):
    def __init__(self, datasets: list[InferenceDataset], lengths: list[int]):
        self.datasets = datasets
        self.lengths = lengths

    def __len__(self):
        return len(self.datasets)

    def __getitem__(self, idx):
        return (self.datasets[idx], self.lengths[idx])


class RawTrainData:
    def __init__(self, dataset_id: "TorchTrainDataset.Id"):
        self.dataset_id: TorchTrainDataset.Id = dataset_id
        self.raw_data: list[tuple[torch.Tensor, torch.Tensor, torch.Tensor]] = []
        self.load_stats: LoadStats | None = None

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

    def to(self, device: torch.device):
        self.raw_data = [
            (
                input_.to(device, non_blocking=True),
                boundary.to(device, non_blocking=True),
                label.to(device, non_blocking=True),
            )
            for input_, boundary, label in self.raw_data
        ]

    def pin_memory(self):
        self.raw_data = [
            (
                input_.pin_memory(),
                boundary.pin_memory(),
                label.pin_memory(),
            )
            for input_, boundary, label in self.raw_data
        ]
        return self


class TrainData:
    """A single batch of training data.

    A single batch contains multiple steps worth of ``Example`` 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).
    """

    def __init__(
        self, num_prognostic_channels: int, num_boundary_channels: int, ctx: GridContext
    ):
        self.num_prognostic_channels = num_prognostic_channels
        self.num_boundary_channels = num_boundary_channels
        self.ctx = ctx
        self.example_by_step: list[Example] = []
        self.load_stats: LoadStats | None = None

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

    def get_initial_input(self) -> tuple[Prognostic, Boundary]:
        return self.get_input(0)

    def get_input(self, step: int) -> tuple[Prognostic, Boundary]:
        prog, boundary, _ = self.example_by_step[step]
        return prog, boundary

    def get_label(self, step: int) -> Prognostic:
        return self.example_by_step[step][2]

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

    def __len__(self) -> int:
        return len(self.example_by_step)

    def __iter__(self):
        return iter(range(len(self)))

    def to(self, device: torch.device) -> None:
        for step in self:
            prog, boundary, label = self.example_by_step[step]
            self.example_by_step[step] = (
                prog.to(device, non_blocking=True),
                boundary.to(device, non_blocking=True),
                label.to(device, non_blocking=True),
            )

    def pin_memory(self):
        for step in self:
            prog, boundary, label = self.example_by_step[step]
            self.example_by_step[step] = (
                prog.pin_memory(),
                boundary.pin_memory(),
                label.pin_memory(),
            )
        return self


@final
class TorchTrainDataset(Dataset[RawTrainData]):
    """
    This class is used for training and validation.

    It creates rolling indices to keep track of histories/past states. But different
    from InferenceDataset, as it creates rolling indices based on stride. By default,
    the sliding window / stride is 1.

    We make use of TrainData class to store a single sample.

    For example,
    Hist=0 ; TD: step=0->[0, 1]; step=1->[1, 2]; step=2->[2, 3]; step=3->[3, 4]
    Hist=1 ; TD: step=0->[[0, 1], [2, 3]]; step=1->[[2, 3], [4, 5]];
            step=2->[[4, 5], [6, 7]]; step=3->[[6, 7], [8, 9]]
    Hist=2 ; TD: step=0->[[0, 1, 2], [3, 4, 5]];
            step=1->[[3, 4, 5], [6, 7, 8]];
            step=2->[[6, 7, 8], [9, 10, 11]];
            step=3->[[9, 10, 11], [12, 13, 14]]
    """

    type Id = str

    FLAG = LoaderVersion.OM4_TORCH

    # Shared across all instances within a process. Created lazily on first
    # __getitem__ call so that each forked DataLoader worker gets its own
    # clean executor — avoids inheriting fork-corrupted locks from the parent.
    _shared_executor: ClassVar[ThreadPoolExecutor | None] = None

    @classmethod
    def _get_executor(cls) -> ThreadPoolExecutor:
        if cls._shared_executor is None:
            cls._shared_executor = ThreadPoolExecutor(
                max_workers=None, thread_name_prefix="concurrent_compute"
            )
        return cls._shared_executor

    @elapsed
    def __init__(
        self,
        src: DataSource,
        dst: DataSource | None,
        prognostic_var_names: PrognosticVarNames,
        boundary_var_names: BoundaryVarNames,
        hist: int,
        steps: int,
        normalize_before_mask: bool,
        masked_fill_value: float,
        stride: int = 1,
        concurrent_compute_: bool = False,
    ):
        super().__init__()
        self.id = f"{self.__class__.__name__}_{str(id(self))}"
        # If the src and dst DataSource are the same, we can do a lot less work.
        srcs = [src, dst] if dst else [src]

        self.hist: int = hist
        self.steps: int = steps
        self.stride: int = stride
        self.normalize_before_mask: bool = normalize_before_mask
        self.masked_fill_value: float = masked_fill_value
        self._concurrent_compute = concurrent_compute_

        self.num_prognostic_channels: int = (hist + 1) * len(prognostic_var_names)
        self.num_boundary_channels: int = (hist + 1) * len(boundary_var_names)
        assert np.array_equal(srcs[0].data.time, srcs[-1].data.time), (
            "src and dst DataSource have different time slices!"
        )
        time_ = src.data.time
        self.prognostic_srcs = [
            src.filter(prognostic_var_names, prefix="prog") for src in srcs
        ]
        self.boundary_src = src.filter(boundary_var_names, prefix="boundary")

        # This class will be used only for training and validation
        total_steps: int = 2 * self.hist + 2

        # Calculate the number of windows
        num_windows = time_.size - (total_steps - 1) * self.stride

        # Create base indices
        indices = np.arange(num_windows)
        indices_da = xr.DataArray(indices, dims=["window"])

        # Create window dimension
        window_dim = xr.DataArray(np.arange(total_steps), dims=["time"])

        # Construct rolling indices
        self.rolling_indices: Float[xr.DataArray, "window time"] = (
            indices_da + stride * window_dim
        )

        # NB(alxmrs): Keep masks on CPU - will be moved to GPU in to_train_data()
        self.wet_prognostic: list[PrognosticMask] = [
            src.masks.prognostic for src in srcs
        ]
        self.wet_surface: GridMask = src.masks.boundary

        self.ctx = GridContext(
            label_mask=self.prognostic_srcs[-1].masks.prognostic_with_hist(self.hist),
            input_resolution_cpu=self.prognostic_srcs[0].resolution,
            output_resolution_cpu=self.prognostic_srcs[-1].resolution,
        )

        self.size: int = (
            time_.size
            - self.steps * (self.hist + 1) * self.stride
            - self.hist * self.stride
        )

    def __len__(self) -> int:
        return self.size

    @elapsed(level=logging.DEBUG)
    def __getitem__(self, idx: int):
        start_time = time.perf_counter()
        TD = RawTrainData(self.id)

        for step in range(self.steps):
            x_index = self._get_x_index(idx, step)
            current_x_index = x_index.isel(time=slice(0, self.hist + 1))
            forecast_x_index = x_index.isel(time=slice(self.hist + 1, None))

            # Only materialize the time ranges we actually use to reduce memory.
            input_selected = self.prognostic_srcs[0].data.isel(time=current_x_index)
            boundary_selected = self.boundary_src.data.isel(time=current_x_index)
            label_selected = self.prognostic_srcs[-1].data.isel(
                time=forecast_x_index
            )  # forecasted data
            prognostic_selected = [input_selected, label_selected]

            if self._concurrent_compute:
                datasets = prognostic_selected + [boundary_selected]
                concurrent_compute(
                    *datasets,
                    executor=self._get_executor(),
                )

            if "lev" in prognostic_selected[0].dims:
                prognostics = [
                    torch.from_numpy(
                        conditional_rearrange(
                            selected,
                            "time (variable lev)=var lat lon",
                            concat_dim="var",
                        )
                        .rename({"var": "variable"})
                        .to_numpy()
                        .astype(np.float32, copy=False)
                    )
                    for selected in prognostic_selected
                ]
            else:
                prognostics = [
                    torch.from_numpy(
                        selected.to_array()
                        .transpose("time", "variable", "lat", "lon")
                        .to_numpy()
                        .astype(np.float32, copy=False)
                    )
                    for selected in prognostic_selected
                ]
            boundary = torch.from_numpy(
                boundary_selected.to_array()
                .transpose("time", "variable", "lat", "lon")
                .to_numpy()
                .astype(np.float32, copy=False)
            )
            input_, label = prognostics[0], prognostics[-1]
            TD.insert(input_, boundary, label)
        TD.load_stats = LoadStats(time.perf_counter() - start_time)

        return TD

    def to_train_data(
        self, raw_train_data: RawTrainData, device: torch.device
    ) -> TrainData:
        """Convert RawTrainData to TrainData, moving tensors to the specified device.

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

        Returns:
            TrainData with tensors on the target device
        """
        train_data = TrainData(
            self.num_prognostic_channels,
            self.num_boundary_channels,
            self.ctx.to(device),
        )
        for input_, boundary, label in raw_train_data.raw_data:
            prog_input, boundary_input, label_tensor = self._to_example(
                OceanData.from_data_source(
                    input_,
                    self.wet_prognostic[0],
                    self.prognostic_srcs[0],
                ).to(device=device, non_blocking=True),
                OceanData.from_data_source(
                    boundary,
                    self.wet_surface,
                    self.boundary_src,
                ).to(device=device, non_blocking=True),
                OceanData.from_data_source(
                    label, self.wet_prognostic[-1], self.prognostic_srcs[-1]
                ).to(device=device, non_blocking=True),
            )
            train_data.append(prog_input, boundary_input, label_tensor)
        train_data.load_stats = raw_train_data.load_stats
        return train_data

    def _to_example(
        self, input_: OceanData, boundary: OceanData, label: OceanData
    ) -> Example:
        # Input/boundary only include current steps; label only includes forecasted steps.
        prog_input = self._prep_tensor_steps(input_)
        boundary_input = self._prep_tensor_steps(boundary)
        label_tensor = self._prep_tensor_steps(label)
        return prog_input, boundary_input, label_tensor

    def _prep_tensor_steps(self, ocean_data: OceanData) -> Input:
        """Normalize, mask, and flatten (time, variable) dims into a channel dim."""
        steps = ocean_data.normalize_and_mask(
            self.normalize_before_mask, self.masked_fill_value
        )
        return rearrange(
            steps, "batch time variable lat lon -> batch (time variable) lat lon"
        )

    def _get_x_index(self, idx: int, step: int) -> xr.DataArray:
        assert isinstance(idx, int)
        if idx < 0:
            raise IndexError("Sorry, negative indexing is not supported!")
        if idx >= len(self):
            raise IndexError("Index out of range")

        window_index = idx + step * (self.hist + 1) * self.stride
        return self.rolling_indices.isel(window=window_index, drop=True)


def concurrent_compute(
    *datasets: xr.Dataset,
    executor: ThreadPoolExecutor,
) -> None:
    def load_variable_data(var: xr.Variable) -> None:
        var.load()

    futures = []
    for ds in datasets:
        for var in ds.data_vars.variables.values():
            futures.append(executor.submit(load_variable_data, var))

    wait(futures)


@final
class TrainDataLoader:
    """Wrapper around a torch DataLoader that handles GPU post-processing.

    This class wraps a DataLoader[RawTrainData] and converts the raw data
    to TrainData 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.
    """

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

    def __iter__(self):
        """Iterate over the dataloader, converting RawTrainData to TrainData."""
        for raw_train_data in self._dataloader:
            dataset = self._datasets[raw_train_data.dataset_id]
            train_data = dataset.to_train_data(raw_train_data, self._device)
            yield train_data

    def __len__(self) -> int:
        return len(self._dataloader)

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

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

    @property
    def dataset(self):
        return self._dataloader.dataset

    @property
    def datasets(self) -> list[TorchTrainDataset]:
        return list(self._datasets.values())

    @property
    def sampler(self):
        return self._dataloader.sampler