Datasets¶

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

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