Datasets¶

@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
    )

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

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

@elapsed
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,
    concurrent_compute_: bool = False,
):
    super().__init__()
    self.id = f"{self.__class__.__name__}_{str(id(self))}"
    sources = [input_source, label_source] if label_source else [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: int = steps
    self.stride: int = stride
    self._concurrent_compute = concurrent_compute_

    assert np.array_equal(sources[0].time, sources[-1].time), (
        "Input and label sources have different time slices!"
    )
    time_ = input_source.time
    self.sources = sources
    self.prognostic_var_names = tuple(prognostic_var_names)
    self.boundary_var_names = tuple(boundary_var_names)
    self.num_prognostic_channels: int = self.input_steps * len(prognostic_var_names)
    self.num_output_channels: int = self.output_steps * len(prognostic_var_names)
    self.num_boundary_channels: int = self.input_steps * len(boundary_var_names)

    # This class will be used only for training and validation
    total_steps = self.input_steps + self.output_steps

    # 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
    )

    self.preprocessors = [
        BatchPreprocessor(
            source,
            self.prognostic_var_names,
            self.boundary_var_names,
            normalize_before_mask=normalize_before_mask,
            masked_fill_value=masked_fill_value,
        )
        for source in sources
    ]

    self.ctx = BatchGrid(
        label_mask=self.sources[-1].masks.prognostic_for_steps(self.output_steps),
        input_resolution_cpu=self.sources[0].resolution,
        output_resolution_cpu=self.sources[-1].resolution,
    )

    self.size: int = (
        time_.size
        - self.steps * self.output_steps * self.stride
        - (self.input_steps - 1) * self.stride
    )

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