PrithviWxC/dataloaders/merra2.py

import functools as ft
import os
import random
import re
from collections import defaultdict
from datetime import datetime, timedelta
from pathlib import Path

import h5py
import numpy as np
import pandas as pd
import torch
from torch import Tensor
from torch.utils.data import Dataset


def preproc(batch: list[dict], padding: dict[tuple[int]]) -> dict[str, Tensor]:
    """Prepressing function for MERRA2 Dataset

    Args:
        batch (dict): List of training samples, each sample should be a
            dictionary with the following keys::

            'sur_static': Numpy array of shape (3, lat, lon). For each pixel (lat, lon), the first dimension indexes sin(lat), cos(lon), sin(lon).
            'sur_vals': Torch tensor of shape (parameter, time, lat, lon).
            'sur_tars': Torch tensor of shape (parameter, time, lat, lon).
            'ulv_vals': Torch tensor of shape (parameter, level, time, lat, lon).
            'ulv_tars': Torch tensor of shape (parameter, level, time, lat, lon).
            'sur_climate': Torch tensor of shape (parameter, lat, lon)
            'ulv_climate': Torch tensor of shape (parameter, level, lat, lon)
            'lead_time': Integer.
            'input_time': Integer.

        padding: Dictionary with keys 'level', 'lat', 'lon', each of dim 2.

    Returns:
        Dictionary with the following keys::

            'x': [batch, time, parameter, lat, lon]
            'y': [batch, parameter, lat, lon]
            'static': [batch, parameter, lat, lon]
            'lead_time': [batch]
            'input_time': [batch]
            'climate (Optional)': [batch, parameter, lat, lon]

    Note:
        Here, for x and y, 'parameter' is [surface parameter, upper level,
        parameter x level]. Similarly for the static information we have
        [sin(lat), cos(lon), sin(lon), cos(doy), sin(doy), cos(hod), sin(hod),
        ...].
    """  # noqa: E501
    b0 = batch[0]
    nbatch = len(batch)
    data_keys = set(b0.keys())

    essential_keys = {
        "sur_static",
        "sur_vals",
        "sur_tars",
        "ulv_vals",
        "ulv_tars",
        "input_time",
        "lead_time",
    }

    climate_keys = {
        "sur_climate",
        "ulv_climate",
    }

    all_keys = essential_keys | climate_keys

    if not essential_keys.issubset(data_keys):
        raise ValueError("Missing essential keys.")

    if not data_keys.issubset(all_keys):
        raise ValueError("Unexpected keys in batch.")

    # Bring all tensors from the batch into a single tensor
    upl_x = torch.empty((nbatch, *b0["ulv_vals"].shape))
    upl_y = torch.empty((nbatch, *b0["ulv_tars"].shape))

    sur_x = torch.empty((nbatch, *b0["sur_vals"].shape))
    sur_y = torch.empty((nbatch, *b0["sur_tars"].shape))

    sur_sta = torch.empty((nbatch, *b0["sur_static"].shape))

    lead_time = torch.empty((nbatch,), dtype=torch.float32)
    input_time = torch.empty((nbatch,), dtype=torch.float32)

    for i, rec in enumerate(batch):
        sur_x[i] = rec["sur_vals"]
        sur_y[i] = rec["sur_tars"]

        upl_x[i] = rec["ulv_vals"]
        upl_y[i] = rec["ulv_tars"]

        sur_sta[i] = rec["sur_static"]

        lead_time[i] = rec["lead_time"]
        input_time[i] = rec["input_time"]

    return_value = {
        "lead_time": lead_time,
        "input_time": input_time,
    }

    # Reshape (batch, parameter, level, time, lat, lon) ->
    #  (batch, time, parameter, level, lat, lon)
    upl_x = upl_x.permute((0, 3, 1, 2, 4, 5))
    upl_y = upl_y.permute((0, 3, 1, 2, 4, 5))
    # Reshape (batch, parameter, time, lat, lon) ->
    #  (batch, time, parameter, lat, lon)
    sur_x = sur_x.permute((0, 2, 1, 3, 4))
    sur_y = sur_y.permute((0, 2, 1, 3, 4))

    # Pad
    padding_2d = (*padding["lon"], *padding["lat"])

    def pad2d(x):
        return torch.nn.functional.pad(x, padding_2d, mode="constant", value=0)

    padding_3d = (*padding["lon"], *padding["lat"], *padding["level"])

    def pad3d(x):
        return torch.nn.functional.pad(x, padding_3d, mode="constant", value=0)

    sur_x = pad2d(sur_x).contiguous()
    upl_x = pad3d(upl_x).contiguous()
    sur_y = pad2d(sur_y).contiguous()
    upl_y = pad3d(upl_y).contiguous()
    return_value["static"] = pad2d(sur_sta).contiguous()

    # Remove time for targets
    upl_y = torch.squeeze(upl_y, 1)
    sur_y = torch.squeeze(sur_y, 1)

    # We stack along the combined parameter x level dimension
    return_value["x"] = torch.cat(
        (sur_x, upl_x.view(*upl_x.shape[:2], -1, *upl_x.shape[4:])), dim=2
    )
    return_value["y"] = torch.cat(
        (sur_y, upl_y.view(upl_y.shape[0], -1, *upl_y.shape[3:])), dim=1
    )

    if climate_keys.issubset(data_keys):
        sur_climate = torch.empty((nbatch, *b0["sur_climate"].shape))
        ulv_climate = torch.empty((nbatch, *b0["ulv_climate"].shape))
        for i, rec in enumerate(batch):
            sur_climate[i] = rec["sur_climate"]
            ulv_climate[i] = rec["ulv_climate"]
        sur_climate = pad2d(sur_climate)
        ulv_climate = pad3d(ulv_climate)

        return_value["climate"] = torch.cat(
            (
                sur_climate,
                ulv_climate.view(nbatch, -1, *ulv_climate.shape[3:]),
            ),
            dim=1,
        )

    return return_value


def input_scalers(
    surf_vars: list[str],
    vert_vars: list[str],
    levels: list[float],
    surf_path: str | Path,
    vert_path: str | Path,
) -> tuple[Tensor, Tensor]:
    """Reads the input scalers

    Args:
        surf_vars: surface variables to be used.
        vert_vars: vertical variables to be used.
        levels: MERRA2 levels to use.
        surf_path: path to surface scalers file.
        vert_path: path to vertical level scalers file.

    Returns:
        mu (Tensor): mean values
        var (Tensor): varience values
    """
    with h5py.File(Path(surf_path), "r", libver="latest") as surf_file:
        stats = [x.decode().lower() for x in surf_file["statistic"][()]]
        mu_idx = stats.index("mu")
        sig_idx = stats.index("sigma")

        s_mu = torch.tensor([surf_file[k][()][mu_idx] for k in surf_vars])
        s_sig = torch.tensor([surf_file[k][()][sig_idx] for k in surf_vars])

    with h5py.File(Path(vert_path), "r", libver="latest") as vert_file:
        stats = [x.decode().lower() for x in vert_file["statistic"][()]]
        mu_idx = stats.index("mu")
        sig_idx = stats.index("sigma")

        lvl = vert_file["lev"][()]
        l_idx = [np.where(lvl == v)[0].item() for v in levels]

        v_mu = np.array([vert_file[k][()][mu_idx, l_idx] for k in vert_vars])
        v_sig = np.array([vert_file[k][()][sig_idx, l_idx] for k in vert_vars])

    v_mu = torch.from_numpy(v_mu).view(-1)
    v_sig = torch.from_numpy(v_sig).view(-1)

    mu = torch.cat((s_mu, v_mu), dim=0).to(torch.float32)
    sig = torch.cat((s_sig, v_sig), dim=0).to(torch.float32).clamp(1e-4, 1e4)
    return mu, sig


def static_input_scalers(
    scalar_path: str | Path, stat_vars: list[str], unscaled_params: int = 7
) -> tuple[Tensor, Tensor]:
    scalar_path = Path(scalar_path)

    with h5py.File(scalar_path, "r", libver="latest") as scaler_file:
        stats = [x.decode().lower() for x in scaler_file["statistic"][()]]
        mu_idx = stats.index("mu")
        sig_idx = stats.index("sigma")

        mu = torch.tensor([scaler_file[k][()][mu_idx] for k in stat_vars])
        sig = torch.tensor([scaler_file[k][()][sig_idx] for k in stat_vars])

    z = torch.zeros(unscaled_params, dtype=mu.dtype, device=mu.device)
    o = torch.ones(unscaled_params, dtype=sig.dtype, device=sig.device)
    mu = torch.cat((z, mu), dim=0).to(torch.float32)
    sig = torch.cat((o, sig), dim=0).to(torch.float32)

    return mu, sig.clamp(1e-4, 1e4)


def output_scalers(
    surf_vars: list[str],
    vert_vars: list[str],
    levels: list[float],
    surf_path: str | Path,
    vert_path: str | Path,
) -> Tensor:
    surf_path = Path(surf_path)
    vert_path = Path(vert_path)

    with h5py.File(surf_path, "r", libver="latest") as surf_file:
        svars = torch.tensor([surf_file[k][()] for k in surf_vars])

    with h5py.File(vert_path, "r", libver="latest") as vert_file:
        lvl = vert_file["lev"][()]
        l_idx = [np.where(lvl == v)[0].item() for v in levels]
        vvars = np.array([vert_file[k][()][l_idx] for k in vert_vars])
    vvars = torch.from_numpy(vvars).view(-1)

    var = torch.cat((svars, vvars), dim=0).to(torch.float32).clamp(1e-7, 1e7)

    return var


class SampleSpec:
    """
    A data class to collect the information used to define a sample.
    """

    def __init__(
        self,
        inputs: tuple[pd.Timestamp, pd.Timestamp],
        lead_time: int,
        target: pd.Timestamp | list[pd.Timestamp],
    ):
        """
        Args:
            inputs: Tuple of timestamps. In ascending order.
            lead_time: Lead time. In hours.
            target: Timestamp of the target. Can be before or after the inputs.
        """
        if not inputs[0] < inputs[1]:
            raise ValueError(
                "Timestamps in `inputs` should be in strictly ascending order."
            )

        self.inputs = inputs
        self.input_time = (inputs[1] - inputs[0]).total_seconds() / 3600
        self.lead_time = lead_time
        self.target = target

        self.times = [*inputs, target]
        self.stat_times = [inputs[-1]]

    @property
    def climatology_info(self) -> tuple[int, int]:
        """Get the required climatology info.

        :return: information required to obtain climatology data. Essentially
            this is the day of the year and hour of the day of the target
            timestamp, with the former restricted to the interval [1, 365].
        :rtype: tuple
        """
        return (min(self.target.dayofyear, 365), self.target.hour)

    @property
    def year(self) -> int:
        return self.inputs[1].year

    @property
    def dayofyear(self) -> int:
        return self.inputs[1].dayofyear

    @property
    def hourofday(self) -> int:
        return self.inputs[1].hour

    def _info_str(self) -> str:
        iso_8601 = "%Y-%m-%dT%H:%M:%S"

        return (
            f"Issue time: {self.inputs[1].strftime(iso_8601)}\n"
            f"Lead time: {self.lead_time} hours ahead\n"
            f"Input delta: {self.input_time} hours\n"
            f"Target time: {self.target.strftime(iso_8601)}"
        )

    @classmethod
    def get(cls, timestamp: pd.Timestamp, dt: int, lead_time: int):
        """Given a timestamp and lead time, generates a SampleSpec object
        describing the sample further.

        Args:
            timestamp: Timstamp of the sample, Ie this is the larger of the two
                input timstamps.
            dt: Time between input samples, in hours.
            lead_time: Lead time. In hours.

        Returns:
            SampleSpec
        """  # noqa: E501
        assert dt > 0, "dt should be possitive"
        lt = pd.to_timedelta(lead_time, unit="h")
        dt = pd.to_timedelta(dt, unit="h")

        if lead_time >= 0:
            timestamp_target = timestamp + lt
        else:
            timestamp_target = timestamp - dt + lt

        spec = cls(
            inputs=(timestamp - dt, timestamp),
            lead_time=lead_time,
            target=timestamp_target,
        )

        return spec

    def __repr__(self) -> str:
        return self._info_str()

    def __str__(self) -> str:
        return self._info_str()


class Merra2Dataset(Dataset):
    """MERRA2 dataset. The dataset unifies surface and vertical data as well as
    optional climatology.

    Samples come in the form of a dictionary. Not all keys support all
    variables, yet the general ordering of dimensions is
    parameter, level, time, lat, lon

    Note:
        Data is assumed to be in NetCDF files containing daily data at 3-hourly
        intervals. These follow the naming patterns
        MERRA2_sfc_YYYYMMHH.nc and MERRA_pres_YYYYMMHH.nc and can be located in
        two different locations. Optional climatology data comes from files
        climate_surface_doyDOY_hourHOD.nc and
        climate_vertical_doyDOY_hourHOD.nc.


    Note:
        `_get_valid_timestamps` assembles a set of all timestamps for which
        there is data (with hourly resolutions). The result is stored in
        `_valid_timestamps`. `_get_valid_climate_timestamps` does the same with
        climatology data and stores it in `_valid_climate_timestamps`.

        Based on this information, `samples` generates a list of valid samples,
        stored in `samples`. Here the format is::

            [
                [
                    (timestamp 1, lead time A),
                    (timestamp 1, lead time B),
                    (timestamp 1, lead time C),
                ],
                [
                    (timestamp 2, lead time D),
                    (timestamp 2, lead time E),
                ]
            ]

        That is, the outer list iterates over timestamps (init times), the
        inner over lead times. Only valid entries are stored.
    """

    valid_vertical_vars = [
        "CLOUD",
        "H",
        "OMEGA",
        "PL",
        "QI",
        "QL",
        "QV",
        "T",
        "U",
        "V",
    ]
    valid_surface_vars = [
        "EFLUX",
        "GWETROOT",
        "HFLUX",
        "LAI",
        "LWGAB",
        "LWGEM",
        "LWTUP",
        "PRECTOT",
        "PS",
        "QV2M",
        "SLP",
        "SWGNT",
        "SWTNT",
        "T2M",
        "TQI",
        "TQL",
        "TQV",
        "TS",
        "U10M",
        "V10M",
        "Z0M",
    ]
    valid_static_surface_vars = ["FRACI", "FRLAND", "FROCEAN", "PHIS"]

    valid_levels = [
        34.0,
        39.0,
        41.0,
        43.0,
        44.0,
        45.0,
        48.0,
        51.0,
        53.0,
        56.0,
        63.0,
        68.0,
        71.0,
        72.0,
    ]

    timedelta_input = pd.to_timedelta(3, unit="h")

    def __init__(
        self,
        time_range: tuple[str | pd.Timestamp, str | pd.Timestamp],
        lead_times: list[int],
        input_times: list[int],
        data_path_surface: str | Path,
        data_path_vertical: str | Path,
        climatology_path_surface: str | Path | None = None,
        climatology_path_vertical: str | Path | None = None,
        surface_vars: list[str] | None = None,
        static_surface_vars: list[str] | None = None,
        vertical_vars: list[str] | None = None,
        levels: list[float] | None = None,
        roll_longitudes: int = 0,
        positional_encoding: str = "absolute",
        rtype: type = np.float32,
        dtype: torch.dtype = torch.float32,
    ) -> None:
        """
        Args:
            data_path_surface: Location of surface data.
            data_path_vertical: Location of vertical data.
            climatology_path_surface: Location of (optional) surface
                climatology.
            climatology_path_vertical: Location of (optional) vertical
                climatology.
            surface_vars: Surface variables.
            static_surface_vars: Static surface variables.
            vertical_vars: Vertical variables.
            levels: Levels.
            time_range: Used to subset data.
            lead_times: Lead times for generalized forecasting.
            roll_longitudes: Set to non-zero value to data by random amount
                along longitude dimension.
            position_encoding: possible values are
              ['absolute' (default), 'fourier'].
                'absolute' returns lat lon encoded in 3 dimensions using sine
                  and cosine
                'fourier' returns lat/lon to be encoded by model
                <any other key> returns lat/lon to be encoded by model
            rtype: numpy data type used during read
            dtype: torch data type of data output
        """

        self.time_range = (
            pd.to_datetime(time_range[0]),
            pd.to_datetime(time_range[1]),
        )
        self.lead_times = lead_times
        self.input_times = input_times
        self._roll_longitudes = list(range(roll_longitudes + 1))

        self._uvars = vertical_vars or self.valid_vertical_vars
        self._level = levels or self.valid_levels
        self._svars = surface_vars or self.valid_surface_vars
        self._sstat = static_surface_vars or self.valid_static_surface_vars
        self._nuvars = len(self._uvars)
        self._nlevel = len(self._level)
        self._nsvars = len(self._svars)
        self._nsstat = len(self._sstat)

        self.rtype = rtype
        self.dtype = dtype

        self.positional_encoding = positional_encoding

        self._data_path_surface = Path(data_path_surface)
        self._data_path_vertical = Path(data_path_vertical)

        self.dir_exists(self._data_path_surface)
        self.dir_exists(self._data_path_vertical)

        self._get_coordinates()

        self._climatology_path_surface = Path(climatology_path_surface) or None
        self._climatology_path_vertical = (
            Path(climatology_path_vertical) or None
        )
        self._require_clim = (
            self._climatology_path_surface is not None
            and self._climatology_path_vertical is not None
        )

        if self._require_clim:
            self.dir_exists(self._climatology_path_surface)
            self.dir_exists(self._climatology_path_vertical)
        elif (
            climatology_path_surface is None
            and climatology_path_vertical is None
        ):
            self._climatology_path_surface = None
            self._climatology_path_vertical = None
        else:
            raise ValueError(
                "Either both or neither of"
                "`climatology_path_surface` and"
                "`climatology_path_vertical` should be None."
            )

        if not set(self._svars).issubset(set(self.valid_surface_vars)):
            raise ValueError("Invalid surface variable.")

        if not set(self._sstat).issubset(set(self.valid_static_surface_vars)):
            raise ValueError("Invalid static surface variable.")

        if not set(self._uvars).issubset(set(self.valid_vertical_vars)):
            raise ValueError("Inalid vertical variable.")

        if not set(self._level).issubset(set(self.valid_levels)):
            raise ValueError("Invalid level.")

    @staticmethod
    def dir_exists(path: Path) -> None:
        if not path.is_dir():
            raise ValueError(f"Directory {path} does not exist.")

    @property
    def upper_shape(self) -> tuple:
        """Returns the vertical variables shape
        Returns:
            tuple: vertical variable shape in the following order::

                [VAR, LEV, TIME, LAT, LON]
        """
        return self._nuvars, self._nlevel, 2, 361, 576

    @property
    def surface_shape(self) -> tuple:
        """Returns the surface variables shape

        Returns:
            tuple: surafce shape in the following order::

                [VAR, LEV, TIME, LAT, LON]
        """
        return self._nsvars, 2, 361, 576

    def data_file_surface(self, timestamp: pd.Timestamp) -> Path:
        """Build the surfcae data file name based on timestamp

        Args:
            timestamp: a timestamp

        Returns:
            Path: constructed path
        """
        pattern = "MERRA2_sfc_%Y%m%d.nc"
        data_file = self._data_path_surface / timestamp.strftime(pattern)
        return data_file

    def data_file_vertical(self, timestamp: pd.Timestamp) -> Path:
        """Build the vertical data file name based on timestamp

        Args:
            timestamp: a timestamp

        Returns:
            Path: constructed path
        """
        pattern = "MERRA_pres_%Y%m%d.nc"
        data_file = self._data_path_vertical / timestamp.strftime(pattern)
        return data_file

    def data_file_surface_climate(
        self,
        timestamp: pd.Timestamp | None = None,
        dayofyear: int | None = None,
        hourofday: int | None = None,
    ) -> Path:
        """
        Returns the path to a climatology file based either on a timestamp or
        the dayofyear / hourofday combination.
        Args:
            timestamp: A timestamp.
            dayofyear: Day of the year. 1 to 366.
            hourofday: Hour of the day. 0 to 23.
        Returns:
            Path: Path to climatology file.
        """
        if timestamp is not None and (
            (dayofyear is not None) or (hourofday is not None)
        ):
            raise ValueError(
                "Provide either timestamp or both dayofyear and hourofday."
            )

        if timestamp is not None:
            dayofyear = min(timestamp.dayofyear, 365)
            hourofday = timestamp.hour

        file_name = f"climate_surface_doy{dayofyear:03}_hour{hourofday:02}.nc"
        data_file = self._climatology_path_surface / file_name
        return data_file

    def data_file_vertical_climate(
        self,
        timestamp: pd.Timestamp | None = None,
        dayofyear: int | None = None,
        hourofday: int | None = None,
    ) -> Path:
        """Returns the path to a climatology file based either on a timestamp
        or the dayofyear / hourofday combination.

        Args:
            timestamp: A timestamp. dayofyear: Day of the year. 1 to 366.
            hourofday: Hour of the day. 0 to 23.
        Returns:
            Path: Path to climatology file.
        """
        if timestamp is not None and (
            (dayofyear is not None) or (hourofday is not None)
        ):
            raise ValueError(
                "Provide either timestamp or both dayofyear and hourofday."
            )

        if timestamp is not None:
            dayofyear = min(timestamp.dayofyear, 365)
            hourofday = timestamp.hour

        file_name = f"climate_vertical_doy{dayofyear:03}_hour{hourofday:02}.nc"
        data_file = self._climatology_path_vertical / file_name
        return data_file

    def _get_coordinates(self) -> None:
        """
        Obtains the coordiantes (latitudes and longitudes) from a single data
        file.
        """
        timestamp = next(iter(self.valid_timestamps))

        file = self.data_file_surface(timestamp)
        with h5py.File(file, "r", libver="latest") as handle:
            self.lats = lats = handle["lat"][()].astype(self.rtype)
            self.lons = lons = handle["lon"][()].astype(self.rtype)

        deg_to_rad = np.pi / 180
        self._embed_lat = np.sin(lats * deg_to_rad).reshape(-1, 1)

        self._embed_lon = np.empty((2, 1, len(lons)), dtype=self.rtype)
        self._embed_lon[0, 0] = np.cos(lons * deg_to_rad)
        self._embed_lon[1, 0] = np.sin(lons * deg_to_rad)

    @ft.cached_property
    def lats(self) -> np.ndarray:
        timestamp = next(iter(self.valid_timestamps))

        file = self.data_file_surface(timestamp)
        with h5py.File(file, "r", libver="latest") as handle:
            return handle["lat"][()].astype(self.rtype)

    @ft.cached_property
    def lons(self) -> np.ndarray:
        timestamp = next(iter(self.valid_timestamps))

        file = self.data_file_surface(timestamp)
        with h5py.File(file, "r", libver="latest") as handle:
            return handle["lon"][()].astype(self.rtype)

    @ft.cached_property
    def position_signal(self) -> np.ndarray:
        """Generates the "position signal" that is part of the static
        features.

        Returns:
            Tensor: Torch tensor of dimension (parameter, lat, lon) containing
            sin(lat), cos(lon), sin(lon).
        """

        latitudes, longitudes = np.meshgrid(
            self.lats, self.lons, indexing="ij"
        )

        if self.positional_encoding == "absolute":
            latitudes = latitudes / 360 * 2.0 * np.pi
            longitudes = longitudes / 360 * 2.0 * np.pi
            sur_static = np.stack(
                [np.sin(latitudes), np.cos(longitudes), np.sin(longitudes)],
                axis=0,
            )
        else:
            sur_static = np.stack([latitudes, longitudes], axis=0)

        sur_static = sur_static.astype(self.rtype)

        return sur_static

    @ft.cached_property
    def valid_timestamps(self) -> set[pd.Timestamp]:
        """Generates list of valid timestamps based on available files. Only
        timestamps for which both surface and vertical information is available
        are considered valid.
        Returns:
            list: list of timestamps
        """

        s_glob = self._data_path_surface.glob("MERRA2_sfc_????????.nc")
        s_files = [os.path.basename(f) for f in s_glob]
        v_glob = self._data_path_surface.glob("MERRA_pres_????????.nc")
        v_files = [os.path.basename(f) for f in v_glob]

        s_re = re.compile(r"MERRA2_sfc_(\d{8}).nc\Z")
        v_re = re.compile(r"MERRA_pres_(\d{8}).nc\Z")
        fmt = "%Y%m%d"

        s_times = {
            (datetime.strptime(m[1], fmt))
            for f in s_files
            if (m := s_re.match(f))
        }
        v_times = {
            (datetime.strptime(m[1], fmt))
            for f in v_files
            if (m := v_re.match(f))
        }

        times = s_times.intersection(v_times)

        # Each file contains a day at 3 hour intervals
        times = {
            t + timedelta(hours=i) for i in range(0, 24, 3) for t in times
        }

        start_time, end_time = self.time_range
        times = {pd.Timestamp(t) for t in times if start_time <= t <= end_time}

        return times

    @ft.cached_property
    def valid_climate_timestamps(self) -> set[tuple[int, int]]:
        """Generates list of "timestamps" (dayofyear, hourofday) for which
        climatology data is present. Only instances for which surface and
        vertical data is available are considered valid.
        Returns:
            list: List of tuples describing valid climatology instances.
        """
        if not self._require_clim:
            return set()

        s_glob = self._climatology_path_surface.glob(
            "climate_surface_doy???_hour??.nc"
        )
        s_files = [os.path.basename(f) for f in s_glob]

        v_glob = self._climatology_path_vertical.glob(
            "climate_vertical_doy???_hour??.nc"
        )
        v_files = [os.path.basename(f) for f in v_glob]

        s_re = re.compile(r"climate_surface_doy(\d{3})_hour(\d{2}).nc\Z")
        v_re = re.compile(r"climate_vertical_doy(\d{3})_hour(\d{2}).nc\Z")

        s_times = {
            (int(m[1]), int(m[2])) for f in s_files if (m := s_re.match(f))
        }
        v_times = {
            (int(m[1]), int(m[2])) for f in v_files if (m := v_re.match(f))
        }

        times = s_times.intersection(v_times)

        return times

    def _data_available(self, spec: SampleSpec) -> bool:
        """
        Checks whether data is available for a given SampleSpec object. Does so
        using the internal sets with available data previously constructed. Not
        by checking the file system.
        Args:
            spec: SampleSpec object as returned by SampleSpec.get
        Returns:
            bool: if data is availability.
        """
        valid = set(spec.times).issubset(self.valid_timestamps)

        if self._require_clim:
            sci = spec.climatology_info
            ci = set(sci) if isinstance(sci, list) else set([sci])  # noqa: C405
            valid &= ci.issubset(self.valid_climate_timestamps)

        return valid

    @ft.cached_property
    def samples(self) -> list[tuple[pd.Timestamp, int, int]]:
        """
        Generates list of all valid samlpes.
        Returns:
            list: List of tuples (timestamp, input time, lead time).
        """
        valid_samples = []
        dts = [(it, lt) for it in self.input_times for lt in self.lead_times]

        for timestamp in sorted(self.valid_timestamps):
            timestamp_samples = []
            for it, lt in dts:
                spec = SampleSpec.get(timestamp, -it, lt)

                if self._data_available(spec):
                    timestamp_samples.append((timestamp, it, lt))

            if timestamp_samples:
                valid_samples.append(timestamp_samples)

        return valid_samples

    def _to_torch(
        self,
        data: dict[str, Tensor | list[Tensor]],
        dtype: torch.dtype = torch.float32,
    ) -> dict[str, Tensor | list[Tensor]]:
        out = {}
        for k, v in data.items():
            if isinstance(v, list):
                out[k] = [torch.from_numpy(x).to(dtype) for x in v]
            else:
                out[k] = torch.from_numpy(v).to(dtype)

        return out

    def _lat_roll(
        self, data: dict[str, Tensor | list[Tensor]], n: int
    ) -> dict[str, Tensor | list[Tensor]]:
        out = {}
        for k, v in data.items():
            if isinstance(v, list):
                out[k] = [torch.roll(x, shifts=n, dims=-1) for x in v]
            else:
                out[k] = torch.roll(v, shifts=n, dims=-1)

        return out

    def _read_static_data(
        self, file: str | Path, doy: int, hod: int
    ) -> np.ndarray:
        with h5py.File(file, "r", libver="latest") as handle:
            lats_surf = handle["lat"]
            lons_surf = handle["lon"]

            nll = (len(lats_surf), len(lons_surf))

            npos = len(self.position_signal)
            ntime = 4

            nstat = npos + ntime + self._nsstat
            data = np.empty((nstat, *nll), dtype=self.rtype)

            for i, key in enumerate(self._sstat, start=npos + ntime):
                data[i] = handle[key][()].astype(dtype=self.rtype)

        # [possition signal], cos(doy), sin(doy), cos(hod), sin(hod)
        data[0:npos] = self.position_signal
        data[npos + 0] = np.cos(2 * np.pi * doy / 366)
        data[npos + 1] = np.sin(2 * np.pi * doy / 366)
        data[npos + 2] = np.cos(2 * np.pi * hod / 24)
        data[npos + 3] = np.sin(2 * np.pi * hod / 24)

        return data

    def _read_surface(
        self, tidx: int, nll: tuple[int, int], handle: h5py.File
    ) -> np.ndarray:
        data = np.empty((self._nsvars, *nll), dtype=self.rtype)

        for i, key in enumerate(self._svars):
            data[i] = handle[key][tidx][()].astype(dtype=self.rtype)

        return data

    def _read_levels(
        self, tidx: int, nll: tuple[int, int], handle: h5py.File
    ) -> np.ndarray:
        lvls = handle["lev"][()]
        lidx = self._level_idxs(lvls)

        data = np.empty((self._nuvars, self._nlevel, *nll), dtype=self.rtype)

        for i, key in enumerate(self._uvars):
            data[i] = handle[key][tidx, lidx][()].astype(dtype=self.rtype)

        return np.ascontiguousarray(np.flip(data, axis=1))

    def _level_idxs(self, lvls):
        lidx = [np.argwhere(lvls == int(lvl)).item() for lvl in self._level]
        return sorted(lidx)

    @staticmethod
    def _date_to_tidx(date: datetime | pd.Timestamp, handle: h5py.File) -> int:
        if isinstance(date, pd.Timestamp):
            date = date.to_pydatetime()

        time = handle["time"]

        t0 = time.attrs["begin_time"][()].item()
        d0 = f"{time.attrs['begin_date'][()].item()}"

        offset = datetime.strptime(d0, "%Y%m%d")

        times = [offset + timedelta(minutes=int(t + t0)) for t in time[()]]
        return times.index(date)

    def _read_data(
        self, file_pair: tuple[str, str], date: datetime
    ) -> dict[str, np.ndarray]:
        s_file, v_file = file_pair

        with h5py.File(s_file, "r", libver="latest") as shandle:
            lats_surf = shandle["lat"]
            lons_surf = shandle["lon"]

            nll = (len(lats_surf), len(lons_surf))

            tidx = self._date_to_tidx(date, shandle)

            sdata = self._read_surface(tidx, nll, shandle)

        with h5py.File(v_file, "r", libver="latest") as vhandle:
            lats_vert = vhandle["lat"]
            lons_vert = vhandle["lon"]

            nll = (len(lats_vert), len(lons_vert))

            tidx = self._date_to_tidx(date, vhandle)

            vdata = self._read_levels(tidx, nll, vhandle)

        data = {"vert": vdata, "surf": sdata}

        return data

    def _read_climate(
        self, file_pair: tuple[str, str]
    ) -> dict[str, np.ndarray]:
        s_file, v_file = file_pair

        with h5py.File(s_file, "r", libver="latest") as shandle:
            lats_surf = shandle["lat"]
            lons_surf = shandle["lon"]

            nll = (len(lats_surf), len(lons_surf))

            sdata = np.empty((self._nsvars, *nll), dtype=self.rtype)

            for i, key in enumerate(self._svars):
                sdata[i] = shandle[key][()].astype(dtype=self.rtype)

        with h5py.File(v_file, "r", libver="latest") as vhandle:
            lats_vert = vhandle["lat"]
            lons_vert = vhandle["lon"]

            nll = (len(lats_vert), len(lons_vert))

            lvls = vhandle["lev"][()]
            lidx = self._level_idxs(lvls)

            vdata = np.empty(
                (self._nuvars, self._nlevel, *nll), dtype=self.rtype
            )

            for i, key in enumerate(self._uvars):
                vdata[i] = vhandle[key][lidx][()].astype(dtype=self.rtype)

        data = {
            "vert": np.ascontiguousarray(np.flip(vdata, axis=1)),
            "surf": sdata,
        }

        return data

    def get_data_from_sample_spec(
        self, spec: SampleSpec
    ) -> dict[str, Tensor | int | float]:
        """Loads and assembles sample data given a SampleSpec object.

        Args:
            spec (SampleSpec): Full details regarding the data to be loaded
        Returns:
            dict: Dictionary with the following keys::

                'sur_static': Torch tensor of shape [parameter, lat, lon]. For
                each pixel (lat, lon), the first 7 dimensions index sin(lat),
                cos(lon), sin(lon), cos(doy), sin(doy), cos(hod), sin(hod).
                Where doy is the day of the year [1, 366] and hod the hour of
                the day [0, 23].
                'sur_vals': Torch tensor of shape [parameter, time, lat, lon].
                'sur_tars': Torch tensor of shape [parameter, time, lat, lon].
                'ulv_vals': Torch tensor of shape [parameter, level, time, lat, lon].
                'ulv_tars': Torch tensor of shape [parameter, level, time, lat, lon].
                'sur_climate': Torch tensor of shape [parameter, lat, lon].
                'ulv_climate': Torch tensor of shape [paramter, level, lat, lon].
                'lead_time': Float.
                'input_time': Float.

        """  # noqa: E501

        # We assemble the unique timestamps for which we need data.
        vals_required = {*spec.times}
        stat_required = {*spec.stat_times}

        # We assemble the unique data files from which we need value data
        vals_file_map = defaultdict(list)
        for t in vals_required:
            data_files = (
                self.data_file_surface(t),
                self.data_file_vertical(t),
            )
            vals_file_map[data_files].append(t)

        # We assemble the unique data files from which we need static data
        stat_file_map = defaultdict(list)
        for t in stat_required:
            data_files = (
                self.data_file_surface(t),
                self.data_file_vertical(t),
            )
            stat_file_map[data_files].append(t)

        # Load the value data
        data = {}
        for data_files, times in vals_file_map.items():
            for time in times:
                data[time] = self._read_data(data_files, time)

        # Combine times
        sample_data = {}

        input_upl = np.stack([data[t]["vert"] for t in spec.inputs], axis=2)
        sample_data["ulv_vals"] = input_upl

        target_upl = data[spec.target]["vert"]
        sample_data["ulv_tars"] = target_upl[:, :, None]

        input_sur = np.stack([data[t]["surf"] for t in spec.inputs], axis=1)
        sample_data["sur_vals"] = input_sur

        target_sur = data[spec.target]["surf"]
        sample_data["sur_tars"] = target_sur[:, None]

        # Load the static data
        data_files, times = stat_file_map.popitem()
        time = times[0].dayofyear, times[0].hour
        sample_data["sur_static"] = self._read_static_data(
            data_files[0], *time
        )

        # If required load the surface data
        if self._require_clim:
            ci_year, ci_hour = spec.climatology_info

            surf_file = self.data_file_surface_climate(
                dayofyear=ci_year,
                hourofday=ci_hour,
            )

            vert_file = self.data_file_vertical_climate(
                dayofyear=ci_year,
                hourofday=ci_hour,
            )

            clim_data = self._read_climate((surf_file, vert_file))

            sample_data["sur_climate"] = clim_data["surf"]
            sample_data["ulv_climate"] = clim_data["vert"]

        # Move the data from numpy to torch
        sample_data = self._to_torch(sample_data, dtype=self.dtype)

        # Optionally roll
        if len(self._roll_longitudes) > 0:
            roll_by = random.choice(self._roll_longitudes)
            sample_data = self._lat_roll(sample_data, roll_by)

        # Now that we have rolled, we can add the static data
        sample_data["lead_time"] = spec.lead_time
        sample_data["input_time"] = spec.input_time

        return sample_data

    def get_data(
        self, timestamp: pd.Timestamp, input_time: int, lead_time: int
    ) -> dict[str, Tensor | int]:
        """
        Loads data based on timestamp and lead time.
        Args:
            timestamp: Timestamp.
            input_time: time between input samples.
            lead_time: lead time.
         Returns:
            Dictionary with keys 'sur_static', 'sur_vals', 'sur_tars',
                'ulv_vals', 'ulv_tars', 'sur_climate', 'ulv_climate',
                'lead_time'.
        """
        spec = SampleSpec.get(timestamp, -input_time, lead_time)
        sample_data = self.get_data_from_sample_spec(spec)
        return sample_data

    def __getitem__(self, idx: int) -> dict[str, Tensor | int]:
        """
        Loads data based on sample index and random choice of sample.
        Args:
            idx: Sample index.
         Returns:
            Dictionary with keys 'sur_static', 'sur_vals', 'sur_tars',
                'ulv_vals', 'ulv_tars', 'sur_climate', 'ulv_climate',
                'lead_time', 'input_time'.
        """
        sample_set = self.samples[idx]
        timestamp, input_time, lead_time, *nsteps = random.choice(sample_set)
        sample_data = self.get_data(timestamp, input_time, lead_time)
        return sample_data

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