Upload 6 files

PrithviWxC/__init__.py

@@ -0,0 +1,10 @@
+"""Prithvi-WxC - Weather and climate foundational model."""
+
+__version__ = "1.0.0"
+
+from . import dataloaders, model
+
+__all__ = [
+    "dataloaders",
+    "model",
+]

PrithviWxC/dataloaders/merra2.py

@@ -0,0 +1,1168 @@
+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)

PrithviWxC/dataloaders/merra2_rollout.py

@@ -0,0 +1,512 @@
+import functools as ft
+import random
+from collections import defaultdict
+from copy import deepcopy
+from pathlib import Path
+
+import numpy as np
+import pandas as pd
+import torch
+from torch import Tensor
+
+from PrithviWxC.dataloaders.merra2 import Merra2Dataset, SampleSpec
+
+
+def preproc(
+    batch: list[dict[str, int | float | Tensor]], 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 (nstep, parameter, lat, lon)
+            'ulv_climate': Torch tensor of shape (nstep parameter, level, lat, lon)
+            'lead_time': Integer.
+            'input_time': Interger
+
+        padding: Dictionary with keys 'level', 'lat', 'lon', each of dim 2.
+
+    Returns:
+        Dictionary with the following keys::
+
+            'x': [batch, time, parameter, lat, lon]
+            'ys': [batch, nsteps, parameter, lat, lon]
+            'static': [batch, nstep, parameter, lat, lon]
+            'lead_time': [batch]
+            'input_time': [batch]
+            'climate (Optional)': [batch, nsteps, parameter, lat, lon]
+
+    Note:
+        Here, for x and ys, '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, *b0["lead_time"].shape),
+        dtype=torch.float32,
+    )
+    input_time = torch.empty((nbatch,), dtype=torch.float32)
+
+    for i, rec in enumerate(batch):
+        sur_x[i] = torch.Tensor(rec["sur_vals"])
+        sur_y[i] = torch.Tensor(rec["sur_tars"])
+
+        upl_x[i] = torch.Tensor(rec["ulv_vals"])
+        upl_y[i] = torch.Tensor(rec["ulv_tars"])
+
+        sur_sta[i] = torch.Tensor(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,
+        "target_time": torch.sum(lead_time).reshape(-1),
+    }
+
+    # 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["statics"] = pad2d(sur_sta).contiguous()
+
+    # We stack along the combined parameter 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["ys"] = torch.cat(
+        (sur_y, upl_y.view(*upl_y.shape[:2], -1, *upl_y.shape[4:])), dim=2
+    )
+
+    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)
+
+        ulv_climate = ulv_climate.view(
+            *ulv_climate.shape[:2], -1, *ulv_climate.shape[4:]
+        )
+        return_value["climates"] = torch.cat((sur_climate, ulv_climate), dim=2)
+
+    return return_value
+
+
+class RolloutSpec(SampleSpec):
+    """
+    A data class to collect the information used to define a rollout sample.
+    """
+
+    def __init__(
+        self,
+        inputs: tuple[pd.Timestamp, pd.Timestamp],
+        lead_time: int,
+        target: 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.
+        """
+        super().__init__(inputs, lead_time, target)
+
+        self.dt = dt = pd.Timedelta(lead_time, unit="h")
+        self.inters = list(pd.date_range(inputs[-1], target, freq=dt))
+
+        self._ctimes = deepcopy(self.inters)
+        self.stat_times = deepcopy(self.inters)
+
+        self.stat_times.pop(-1)
+        self._ctimes.pop(0)
+        self.inters.pop(0)
+        self.inters.pop(-1)
+
+        self.times = [*inputs, *self.inters, target]
+        self.targets = self.times[2:]
+        self.nsteps = len(self.times) - 2
+
+    @property
+    def climatology_info(self) -> dict[pd.Timestamp, tuple[int, int]]:
+        """Returns information required to obtain climatology data.
+        Returns:
+            list: list containing required climatology info.
+        """
+        return [(min(t.dayofyear, 365), t.hour) for t in self._ctimes]
+
+    def _info_str(self) -> str:
+        iso_8601 = "%Y-%m-%dT%H:%M:%S"
+
+        inter_str = "\n".join(t.strftime(iso_8601) for t in self.inters)
+
+        return (
+            f"Issue time: {self.inputs[1].strftime(iso_8601)}\n"
+            f"Lead time: {self.lead_time} hours ahead\n"
+            f"Target time: {self.target.strftime(iso_8601)}\n"
+            f"Intermediate times: {inter_str}"
+        )
+
+    @classmethod
+    def get(cls, timestamp: pd.Timestamp, lead_time: int, nsteps: int):
+        """Given a timestamp and lead time, generates a RolloutSpec object
+        describing the sample further.
+
+        Args:
+            timestamp: Timstamp (issue time) of the sample.
+            lead_time: Lead time. In hours.
+
+        Returns:
+            SampleSpec object.
+        """
+        if lead_time > 0:
+            dt = pd.to_timedelta(lead_time, unit="h")
+            timestamp_target = timestamp + nsteps * dt
+        else:
+            raise ValueError("Rollout is only forwards")
+
+        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 Merra2RolloutDataset(Merra2Dataset):
+    """Dataset class that read MERRA2 data for performing rollout.
+
+    Implementation details::
+
+        Samples stores the list of valid samples. This takes the form
+        ```
+        [
+            [(timestamp 1, -input_time, n_steps)],
+            [(timestamp 2, -input_time, n_steps)],
+        ]
+        ```
+        The nested list is for compatibility reasons with Merra2Dataset. Note
+        that input time and n_steps are always the same value. For some reason
+        the sign of input_time is the opposite to that in Merra2Dataset
+    """
+
+    input_time_len = 2
+
+    def __init__(
+        self,
+        time_range: tuple[str | pd.Timestamp, str | pd.Timestamp],
+        input_time: int | float | pd.Timedelta,
+        lead_time: int | float,
+        data_path_surface: str | Path,
+        data_path_vertical: str | Path,
+        climatology_path_surface: str | Path | None,
+        climatology_path_vertical: str | Path | None,
+        surface_vars: list[str],
+        static_surface_vars: list[str],
+        vertical_vars: list[str],
+        levels: list[float],
+        roll_longitudes: int = 0,
+        positional_encoding: str = "absolute",
+    ):
+        """
+        Args:
+            time_range: time range to consider when building dataset
+            input_time: requested time between inputs
+            lead_time: requested time to predict
+            data_path_surface: path of surface data directory
+            data_path_vertical: path of vertical data directory
+            climatology_path_surface: path of surface climatology data
+            directory
+            climatology_path_vertical: path of vertical climatology data
+            directory
+            surface_vars: surface variables to return
+            static_surface_vars: static surface variables to return
+            vertical_vars: vertical variables to return
+            levels: MERA2 vertical levels to consider
+            roll_longitudes: Whether and now uch to randomly roll latitudes by.
+            Defaults to 0.
+            positional_encoding: The type of possitional encodeing to use.
+            Defaults to "absolute".
+
+        Raises:
+            ValueError: If lead time is not integer multiple of input time
+        """
+
+        self._target_lead = lead_time
+
+        if isinstance(input_time, int) or isinstance(input_time, float):
+            self.timedelta_input = pd.to_timedelta(-input_time, unit="h")
+        else:
+            self.timedelta_input = -input_time
+
+        lead_times = [self.timedelta_input / pd.to_timedelta(1, unit="h")]
+
+        super().__init__(
+            time_range,
+            lead_times,
+            [input_time],
+            data_path_surface,
+            data_path_vertical,
+            climatology_path_surface,
+            climatology_path_vertical,
+            surface_vars,
+            static_surface_vars,
+            vertical_vars,
+            levels,
+            roll_longitudes,
+            positional_encoding,
+        )
+
+        nstep_float = (
+            pd.to_timedelta(self._target_lead, unit="h") / self.timedelta_input
+        )
+
+        if abs(nstep_float % 1) > 1e-5:
+            raise ValueError("Leadtime not multiple of input time")
+
+        self.nsteps = round(nstep_float)
+
+    @ft.cached_property
+    def samples(self) -> list[tuple[pd.Timestamp, int, int]]:
+        """Generates list of all valid samlpes.
+
+        Returns:
+            List of tuples (timestamp, input time, lead time).
+        """
+        valid_samples = []
+
+        for timestamp in sorted(self.valid_timestamps):
+            timestamp_samples = []
+            for lt in self.lead_times:
+                spec = RolloutSpec.get(timestamp, lt, self.nsteps)
+
+                if self._data_available(spec):
+                    timestamp_samples.append(
+                        (timestamp, self.input_times[0], lt, self.nsteps)
+                    )
+
+            if timestamp_samples:
+                valid_samples.append(timestamp_samples)
+
+        return valid_samples
+
+    def get_data_from_rollout_spec(
+        self, spec: RolloutSpec
+    ) -> dict[str, Tensor | int | float]:
+        """Loads and assembles sample data given a RolloutSpec object.
+
+        Args:
+            spec (RolloutSpec): Full details regarding the data to be loaded
+        Returns:
+            dict: Dictionary with keys 'sur_static', 'sur_vals', 'sur_tars',
+            'ulv_vals', 'ulv_tars', 'sur_climate', 'ulv_climate',c'lead_time',
+            'input_time'. For each, the value is as follows::
+
+            {
+                '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
+                [nsteps, parameter, level, time, lat, lon].
+                'sur_climate': Torch tensor of shape
+                [nsteps, parameter, lat, lon].
+                'ulv_climate': Torch tensor of shape
+                [nsteps, paramter, level, lat, lon].
+                'lead_time': Float.
+                'input_time': Float.
+            }
+
+        """
+
+        # 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)
+
+        # Load the static data
+        stat = {}
+        for data_files, times in stat_file_map.items():
+            for time in times:
+                hod, doy = time.hour, time.dayofyear
+                stat[time] = self._read_static_data(data_files[0], hod, doy)
+
+        # 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 = np.stack([data[t]["vert"] for t in spec.targets], axis=2)
+        sample_data["ulv_tars"] = target_upl
+
+        input_sur = np.stack([data[t]["surf"] for t in spec.inputs], axis=1)
+        sample_data["sur_vals"] = input_sur
+
+        target_sur = np.stack([data[t]["surf"] for t in spec.targets], axis=1)
+        sample_data["sur_tars"] = target_sur
+
+        # Load the static data
+        static = np.stack([stat[t] for t in spec.stat_times], axis=0)
+        sample_data["sur_static"] = static
+
+        # If required load the climate data
+        if self._require_clim:
+            clim_data = {}
+            for ci in spec.climatology_info:
+                ci_year, ci_hour = ci
+
+                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[ci] = self._read_climate((surf_file, vert_file))
+
+            clim_surf = [clim_data[ci]["surf"] for ci in spec.climatology_info]
+            sample_data["sur_climate"] = np.stack(clim_surf, axis=0)
+
+            clim_surf = [clim_data[ci]["vert"] for ci in spec.climatology_info]
+            sample_data["ulv_climate"] = np.stack(clim_surf, axis=0)
+
+        # 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
+        lt = torch.tensor([spec.lead_time] * self.nsteps).to(self.dtype)
+        sample_data["lead_time"] = lt
+        sample_data["input_time"] = spec.input_time
+
+        return sample_data
+
+    def get_data(
+        self, timestamp: pd.Timestamp, *args, **kwargs
+    ) -> dict[Tensor | int]:
+        """Loads data based on timestamp and lead time.
+
+        Args:
+            timestamp: Timestamp.
+         Returns:
+            Dictionary with keys 'sur_static', 'sur_vals', 'sur_tars',
+              'ulv_vals', 'ulv_tars', 'sur_climate', 'ulv_climate',
+              'lead_time', 'input_time'
+        """
+        rollout_spec = RolloutSpec.get(
+            timestamp, self.lead_times[0], self.nsteps
+        )
+        sample_data = self.get_data_from_rollout_spec(rollout_spec)
+        return sample_data

PrithviWxC/model.py

@@ -0,0 +1,1321 @@
+from functools import cached_property
+from importlib.metadata import version
+
+from torch import Tensor
+from torch.utils.checkpoint import checkpoint
+
+if version("torch") > "2.3.0":
+    from torch.nn.attention import SDPBackend, sdpa_kernel
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+# DropPath code is straight from timm
+# (https://huggingface.co/spaces/Roll20/pet_score/blame/main/lib/timm/models/layers/drop.py)
+def drop_path(
+    x: Tensor,
+    drop_prob: float = 0.0,
+    training: bool = False,
+    scale_by_keep: bool = True,
+) -> Tensor:
+    """Drop paths (Stochastic Depth) per sample (when applied in main path of
+    residual blocks). Taken form timm.
+
+    Args:
+        x (Tensor): Input tensor.
+        drop_prob (float): Probability of dropping `x`, defaults to 0.
+        training (bool): Whether model is in in traingin of eval mode,
+            defaults to False.
+        scale_by_keep (bool): Whether the output should scaled by
+            (`1 - drop_prob`), defaults to True.
+    Returns:
+        Tensor: Tensor that may have randomly dropped with proability
+            `drop_path`
+    """
+    if drop_prob == 0.0 or not training:
+        return x
+    keep_prob = 1 - drop_prob
+    shape = (x.shape[0],) + (1,) * (x.ndim - 1)
+    random_tensor = x.new_empty(shape).bernoulli_(keep_prob)
+    if keep_prob > 0.0 and scale_by_keep:
+        random_tensor.div_(keep_prob)
+    return x * random_tensor
+
+
+class DropPath(nn.Module):
+    """
+    Drop paths (Stochastic Depth) per sample (when applied in main path of
+    residual blocks).
+    """
+
+    def __init__(
+        self, drop_prob: float | None = None, scale_by_keep: bool = True
+    ) -> None:
+        super(DropPath, self).__init__()
+        self.drop_prob = drop_prob
+        self.scale_by_keep = scale_by_keep
+
+    def forward(self, x: Tensor) -> Tensor:
+        """Runs drop path on input tensor
+
+        Args:
+            x: input
+
+        Returns:
+            tensor: output after drop_path
+        """
+        return drop_path(x, self.drop_prob, self.training, self.scale_by_keep)
+
+
+class Mlp(nn.Module):
+    """
+    Multi layer perceptron.
+    """
+
+    def __init__(
+        self, features: int, hidden_features: int, dropout: float = 0.0
+    ) -> None:
+        """
+        Args:
+            features: Input/output dimension.
+            hidden_features: Hidden dimension.
+            dropout: Dropout.
+        """
+        super().__init__()
+        self.net = nn.Sequential(
+            nn.Linear(features, hidden_features),
+            nn.GELU(),
+            nn.Dropout(dropout),
+            nn.Linear(hidden_features, features),
+            nn.Dropout(dropout),
+        )
+
+    def forward(self, x: Tensor) -> Tensor:
+        """
+        Args:
+            x (Tesnor): Tensor of shape [..., channel]
+        Returns:
+            Tenosr: Tensor of same shape as x.
+        """
+        return self.net(x)
+
+
+class LayerNormPassThrough(nn.LayerNorm):
+    """Normalising layer that allows the attention mask to be passed through"""
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+
+    def forward(
+        self, d: tuple[Tensor, Tensor | None]
+    ) -> tuple[Tensor, Tensor | None]:
+        """Forwards function
+
+        Args:
+            d (tuple): tuple of the data tensor and the attention mask
+        Returns:
+            output (Tensor): normalised output data
+            attn_mask (Tensor): the attention mask that was passed in
+        """
+        input, attn_mask = d
+        output = F.layer_norm(
+            input, self.normalized_shape, self.weight, self.bias, self.eps
+        )
+        return output, attn_mask
+
+
+class MultiheadAttention(nn.Module):
+    """Multihead attention layer for inputs of shape
+    [..., sequence, features].
+    """
+
+    def __init__(self, features: int, n_heads: int, dropout: float) -> None:
+        """
+        Args:
+            features: Number of features for inputs to the layer.
+            n_heads: Number of attention heads. Should be a factor of features.
+                (I.e. the layer uses features // n_heads.)
+            dropout: Dropout.
+        """  # noqa: E501
+        super().__init__()
+
+        if (features % n_heads) != 0:
+            raise ValueError(
+                f"Features '{features}' is not divisible by heads '{n_heads}'."
+            )
+
+        self.features = features
+        self.n_heads = n_heads
+        self.dropout = dropout
+
+        self.qkv_layer = torch.nn.Linear(features, features * 3, bias=False)
+        self.w_layer = torch.nn.Linear(features, features, bias=False)
+
+    def forward(self, d: tuple[Tensor, Tensor | None]) -> Tensor:
+        """
+        Args:
+            d (tuple): tuple containing Tensor of shape [..., sequence, features] and the attention mask
+        Returns:
+            Tensor: Tensor of shape [..., sequence, features]
+        """  # noqa: E501
+        x, attn_mask = d
+
+        if not x.shape[-1] == self.features:
+            raise ValueError(
+                f"Expecting tensor with last dimension size {self.features}."
+            )
+
+        passenger_dims = x.shape[:-2]
+        B = passenger_dims.numel()
+        S = x.shape[-2]
+        C = x.shape[-1]
+        x = x.reshape(B, S, C)
+
+        # x [B, S, C]
+        # q, k, v [B, H, S, C/H]
+        q, k, v = (
+            self.qkv_layer(x)
+            .view(B, S, self.n_heads, 3 * (C // self.n_heads))
+            .transpose(1, 2)
+            .chunk(chunks=3, dim=3)
+        )
+
+        # Let us enforce either flash (A100+) or memory efficient attention.
+        if version("torch") > "2.3.0":
+            with sdpa_kernel(
+                [SDPBackend.FLASH_ATTENTION, SDPBackend.EFFICIENT_ATTENTION]
+            ):
+                # x [B, H, S, C//H]
+                x = F.scaled_dot_product_attention(
+                    q, k, v, attn_mask=attn_mask, dropout_p=self.dropout
+                )
+        else:
+            with torch.backends.cuda.sdp_kernel(
+                enable_flash=True, enable_math=False, enable_mem_efficient=True
+            ):
+                # x [B, H, S, C//H]
+                x = F.scaled_dot_product_attention(
+                    q, k, v, dropout_p=self.dropout
+                )
+
+        # x [B, S, C]
+        x = x.transpose(1, 2).view(B, S, C)
+
+        # x [B, S, C]
+        x = self.w_layer(x)
+
+        # Back to input shape
+        x = x.view(*passenger_dims, S, self.features)
+        return x
+
+
+class Transformer(nn.Module):
+    """
+    Transformer for inputs of shape [..., S, features].
+    """
+
+    def __init__(
+        self,
+        features: int,
+        mlp_multiplier: int,
+        n_heads: int,
+        dropout: float,
+        drop_path: float,
+    ) -> None:
+        """
+        Args:
+            features: Number of features for inputs to the layer.
+            mlp_multiplier: Model uses features*mlp_multiplier hidden units.
+            n_heads: Number of attention heads. Should be a factor of features.
+            (I.e. the layer uses features // n_heads.) dropout: Dropout.
+            drop_path: DropPath.
+        """
+        super().__init__()
+
+        self.features = features
+        self.mlp_multiplier = mlp_multiplier
+        self.n_heads = n_heads
+        self.dropout = dropout
+        self.drop_path = (
+            DropPath(drop_path) if drop_path > 0.0 else nn.Identity()
+        )
+
+        self.attention = nn.Sequential(
+            LayerNormPassThrough(features),
+            MultiheadAttention(features, n_heads, dropout),
+        )
+
+        self.ff = nn.Sequential(
+            nn.LayerNorm(features),
+            Mlp(
+                features=features,
+                hidden_features=features * mlp_multiplier,
+                dropout=dropout,
+            ),
+        )
+
+    def forward(self, d: tuple[Tensor, Tensor | None]) -> Tensor:
+        """
+        Args:
+            x: Tensor of shape [..., sequence, features]
+        Returns:
+            Tensor: Tensor of shape [..., sequence, features]
+        """
+        x, attn_mask = d
+        if not x.shape[-1] == self.features:
+            raise ValueError(
+                f"Expecting tensor with last dimension size {self.features}."
+            )
+
+        attention_x = self.attention(d)
+
+        x = x + self.drop_path(attention_x)
+        x = x + self.drop_path(self.ff(x))
+
+        return x
+
+
+class _Shift(nn.Module):
+    """Private base class for the shifter. This allows some behaviour to be
+    easily handled when the shifter isn't used.
+    """
+
+    def __init__(self):
+        super().__init__()
+
+        self._shifted = False
+
+    @torch.no_grad()
+    def reset(self) -> None:
+        """
+        Resets the bool tracking whether the data is shifted
+        """
+        self._shifted: bool = False
+
+    def forward(self, data: Tensor) -> tuple[Tensor, dict[bool, None]]:
+        return data, {True: None, False: None}
+
+
+class SWINShift(_Shift):
+    """
+    Handles the shifting of patches similar to how SWIN works. However if we
+    shift the latitudes then the poles will wrap and potentially that might be
+    problematic. The possition tokens should handle it but masking is safer.
+    """
+
+    def __init__(
+        self,
+        mu_shape: tuple[int, int],
+        global_shape: tuple[int, int],
+        local_shape: tuple[int, int],
+        patch_shape: tuple[int, int],
+        n_context_tokens: int = 2,
+    ) -> None:
+        """
+        Args:
+            mu_shape: the shape to the masking units
+            global_shape: number of global patches in lat and lon
+            local_shape: size of the local patches
+            patch_shape: patch size
+            n_context_token: number of additional context tokens at start of
+            _each_ local sequence
+        """
+        super().__init__()
+
+        self._mu_shape = ms = mu_shape
+        self._g_shape = gs = global_shape
+        self._l_shape = ls = local_shape
+        self._p_shape = ps = patch_shape
+        self._lat_patch = (gs[0], ls[0], gs[1], ls[1])
+        self._n_context_tokens = n_context_tokens
+
+        self._g_shift_to = tuple(
+            int(0.5 * x / p) for x, p in zip(ms, ps, strict=False)
+        )
+        self._g_shift_from = tuple(
+            -int(0.5 * x / p) for x, p in zip(ms, ps, strict=False)
+        )
+
+        # Define the attention masks for the shifted MaxViT.
+        nglobal = global_shape[0] * global_shape[1]
+        nlocal = (
+            local_shape[0] * local_shape[1] + self._n_context_tokens
+        )  # "+ 1" for leadtime
+
+        lm = torch.ones((nglobal, 1, nlocal, nlocal), dtype=bool)
+        mwidth = int(0.5 * local_shape[1]) * local_shape[0]
+        lm[
+            : gs[1],
+            :,
+            self._n_context_tokens : mwidth + self._n_context_tokens,
+            self._n_context_tokens : mwidth + self._n_context_tokens,
+        ] = False
+        self.register_buffer("local_mask", lm)
+
+        gm = torch.ones((nlocal, 1, nglobal, nglobal), dtype=bool)
+        gm[: int(0.5 * ls[1]) * ls[0], :, : gs[1], : gs[1]] = False
+        self.register_buffer("global_mask", gm)
+
+    def _to_grid_global(self, x: Tensor) -> Tensor:
+        """
+        Shuffle and reshape the data from the global/local setting back to the
+        lat/lon grid setting
+        Args:
+            x: the data tensor to be shuffled.
+        Returns:
+            x: data in the global/local setting
+        """
+        nbatch, *other = x.shape
+
+        y1 = x.view(nbatch, *self._g_shape, *self._l_shape, -1)
+        y2 = y1.permute(0, 5, 1, 3, 2, 4).contiguous()
+
+        s = y2.shape
+        return y2.view((nbatch, -1, s[2] * s[3], s[4] * s[5]))
+
+    def _to_grid_local(self, x: Tensor) -> Tensor:
+        """
+        Shuffle and reshape the data from the local/global setting to the
+        lat/lon grid setting
+        Args:
+            x: the data tensor to be shuffled.
+        Returns:
+            x: data in the lat/lon setting.
+        """
+        x = x.transpose(2, 1).contiguous()
+        return self._to_grid_global(x)
+
+    def _from_grid_global(self, x: Tensor) -> Tensor:
+        """
+        Shuffle and reshape the data from the lat/lon grid to the global/local
+        setting
+        Args:
+            x: the data tensor to be shuffled.
+        Returns:
+            x: data in the global/local setting
+        """
+        nbatch, *other = x.shape
+
+        z1 = x.view(nbatch, -1, *self._lat_patch)
+        z2 = z1.permute(0, 2, 4, 3, 5, 1).contiguous()
+
+        s = z2.shape
+        return z2.view(nbatch, s[1] * s[2], s[3] * s[4], -1)
+
+    def _from_grid_local(self, x: Tensor) -> Tensor:
+        """
+        Shuffle and reshape the data from the lat/lon grid to the local/global
+        setting
+        Args:
+            x: the data tensor to be shuffled.
+        Returns:
+            x: data in the local/global setting
+        """
+        x = self._from_grid_global(x)
+        return x.transpose(2, 1).contiguous()
+
+    def _shift(self, x: Tensor) -> Tensor:
+        """
+        Shifts data in the gridded lat/lon setting by half the mask unit shape
+        Args:
+            x: data to be shifted
+        Returns:
+            x: either the hsifted or unshifted data
+        """
+        shift = self._g_shift_from if self._shifted else self._g_shift_to
+        x_shifted = torch.roll(x, shift, (-2, -1))
+
+        self._shifted = not self._shifted
+        return x_shifted
+
+    def _sep_lt(self, x: Tensor) -> tuple[Tensor, Tensor]:
+        """
+        Seperate off the leadtime from the local patches
+        Args:
+            x: data to have leadtime removed from
+        Returns:
+            lt: leadtime
+            x: data without the lead time in the local patch
+        """
+        lt_it = x[:, : self._n_context_tokens, :, :]
+        x_stripped = x[:, self._n_context_tokens :, :, :]
+
+        return lt_it, x_stripped
+
+    def forward(self, data: Tensor) -> tuple[Tensor, Tensor]:
+        """Shift or unshift the the data depending on whether the data is
+        already shifted, as defined by self._shifte.
+
+        Args:
+            data: data to be shifted
+        Returns:
+            Tensor: shifted data Tensor
+        """
+        lt, x = self._sep_lt(data)
+
+        x_grid = self._to_grid_local(x)
+        x_shifted = self._shift(x_grid)
+        x_patched = self._from_grid_local(x_shifted)
+
+        # Mask has to be repeated based on batch size
+        n_batch = x_grid.shape[0]
+        local_rep = [n_batch] + [1] * (self.local_mask.ndim - 1)
+        global_rep = [n_batch] + [1] * (self.global_mask.ndim - 1)
+
+        if self._shifted:
+            attn_mask = {
+                True: self.local_mask.repeat(local_rep),
+                False: self.global_mask.repeat(global_rep),
+            }
+        else:
+            attn_mask = {True: None, False: None}
+
+        return torch.cat((lt, x_patched), axis=1), attn_mask
+
+
+class LocalGlobalLocalBlock(nn.Module):
+    """
+    Applies alternating block and grid attention. Given a parameter n_blocks,
+    the entire module contains 2*n_blocks+1 transformer blocks. The first,
+    third, ..., last apply local (block) attention. The second, fourth, ...
+    global (grid) attention.
+
+    This is heavily inspired by
+    Tu et al. "MaxViT: Multi-Axis Vision Transformer"
+    (https://arxiv.org/abs/2204.01697).
+    """
+
+    def __init__(
+        self,
+        features: int,
+        mlp_multiplier: int,
+        n_heads: int,
+        dropout: float,
+        n_blocks: int,
+        drop_path: float,
+        shifter: nn.Module | None = None,
+        checkpoint: list[int] | None = None,
+    ) -> None:
+        """
+        Args:
+            features: Number of features for inputs to the layer.
+            mlp_multiplier: Model uses features*mlp_multiplier hidden units.
+            n_heads: Number of attention heads. Should be a factor of features.
+            (I.e. the layer uses features // n_heads.)
+            dropout: Dropout.
+            drop_path: DropPath.
+            n_blocks: Number of local-global transformer pairs.
+        """
+        super().__init__()
+
+        self.features = features
+        self.mlp_multiplier = mlp_multiplier
+        self.n_heads = n_heads
+        self.dropout = dropout
+        self.drop_path = drop_path
+        self.n_blocks = n_blocks
+        self._checkpoint = checkpoint or []
+
+        if not all(0 <= c < 2 * n_blocks + 1 for c in self._checkpoint):
+            raise ValueError(
+                "Checkpoints should be 0 <= i < 2*n_blocks+1. "
+                f"{self._checkpoint=}."
+            )
+
+        self.transformers = nn.ModuleList(
+            [
+                Transformer(
+                    features=features,
+                    mlp_multiplier=mlp_multiplier,
+                    n_heads=n_heads,
+                    dropout=dropout,
+                    drop_path=drop_path,
+                )
+                for _ in range(2 * n_blocks + 1)
+            ]
+        )
+
+        self.evaluator = [
+            self._checkpoint_wrapper
+            if i in self._checkpoint
+            else lambda m, x: m(x)
+            for i, _ in enumerate(self.transformers)
+        ]
+
+        self.shifter = shifter or _Shift()
+
+    @staticmethod
+    def _checkpoint_wrapper(
+        model: nn.Module, data: tuple[Tensor, Tensor | None]
+    ) -> Tensor:
+        return checkpoint(model, data, use_reentrant=False)
+
+    def forward(self, x: Tensor) -> Tensor:
+        """
+        Args:
+            x: Tensor of shape::
+
+                [batch, global_sequence, local_sequence, features]
+
+        Returns:
+            Tensor: Tensor of shape::
+
+                [batch, global_sequence, local_sequence, features]
+        """
+        if x.shape[-1] != self.features:
+            raise ValueError(
+                f"Expecting tensor with last dimension size {self.features}."
+            )
+        if x.ndim != 4:
+            raise ValueError(
+                f"Expecting tensor with exactly four dimensions. {x.shape=}."
+            )
+
+        self.shifter.reset()
+        local: bool = True
+        attn_mask = {True: None, False: None}
+
+        transformer_iter = zip(self.evaluator, self.transformers, strict=False)
+
+        # First local block
+        evaluator, transformer = next(transformer_iter)
+        x = evaluator(transformer, (x, attn_mask[local]))
+
+        for evaluator, transformer in transformer_iter:
+            local = not local
+            # We are making exactly 2*n_blocks transposes.
+            # So the output has the same shape as input.
+            x = x.transpose(1, 2)
+
+            x = evaluator(transformer, (x, attn_mask[local]))
+
+            if not local:
+                x, attn_mask = self.shifter(x)
+
+        return x
+
+
+class PatchEmbed(nn.Module):
+    """
+    Patch embedding via 2D convolution.
+    """
+
+    def __init__(
+        self, patch_size: int | tuple[int, ...], channels: int, embed_dim: int
+    ):
+        super().__init__()
+
+        self.patch_size = patch_size
+        self.channels = channels
+        self.embed_dim = embed_dim
+
+        self.proj = nn.Conv2d(
+            channels,
+            embed_dim,
+            kernel_size=patch_size,
+            stride=patch_size,
+            bias=True,
+        )
+
+    def forward(self, x: Tensor) -> Tensor:
+        """
+        Args:
+            x: Tensor of shape [batch, channels, lat, lon].
+        Returns:
+            Tensor: Tensor with shape
+                [batch, embed_dim, lat//patch_size, lon//patch_size]
+        """
+
+        H, W = x.shape[-2:]
+
+        if W % self.patch_size[1] != 0:
+            raise ValueError(
+                f"Cannot do patch embedding for tensor of shape {x.size()}"
+                " with patch size {self.patch_size}. (Dimensions are BSCHW.)"
+            )
+        if H % self.patch_size[0] != 0:
+            raise ValueError(
+                f"Cannot do patch embedding for tensor of shape {x.size()}"
+                f" with patch size {self.patch_size}. (Dimensions are BSCHW.)"
+            )
+
+        x = self.proj(x)
+
+        return x
+
+
+class PrithviWxCEncoderDecoder(nn.Module):
+    """
+    Hiera-MaxViT encoder/decoder code.
+    """
+
+    def __init__(
+        self,
+        embed_dim: int,
+        n_blocks: int,
+        mlp_multiplier: float,
+        n_heads: int,
+        dropout: float,
+        drop_path: float,
+        shifter: nn.Module | None = None,
+        transformer_cp: list[int] | None = None,
+    ) -> None:
+        """
+        Args:
+            embed_dim: Embedding dimension
+            n_blocks: Number of local-global transformer pairs.
+            mlp_multiplier: MLP multiplier for hidden features in feed forward
+                networks.
+            n_heads: Number of attention heads.
+            dropout: Dropout.
+            drop_path: DropPath.
+        """
+        super().__init__()
+
+        self.embed_dim = embed_dim
+        self.n_blocks = n_blocks
+        self.mlp_multiplier = mlp_multiplier
+        self.n_heads = n_heads
+        self.dropout = dropout
+        self._transformer_cp = transformer_cp
+
+        self.lgl_block = LocalGlobalLocalBlock(
+            features=embed_dim,
+            mlp_multiplier=mlp_multiplier,
+            n_heads=n_heads,
+            dropout=dropout,
+            drop_path=drop_path,
+            n_blocks=n_blocks,
+            shifter=shifter,
+            checkpoint=transformer_cp,
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """
+        Args:
+            x: Tensor of shape
+              [batch, global sequence, local sequence, embed_dim]
+        Returns:
+            Tensor of shape
+                [batch, mask_unit_sequence, local_sequence, embed_dim].
+                Identical in shape to the input x.
+        """
+
+        x = self.lgl_block(x)
+
+        return x
+
+
+class PrithviWxC(nn.Module):
+    """Encoder-decoder fusing Hiera with MaxViT. See
+    - Ryali et al. "Hiera: A Hierarchical Vision Transformer without the
+    Bells-and-Whistles" (https://arxiv.org/abs/2306.00989)
+    - Tu et al. "MaxViT: Multi-Axis Vision Transformer"
+    (https://arxiv.org/abs/2204.01697)
+    """
+
+    def __init__(
+        self,
+        in_channels: int,
+        input_size_time: int,
+        in_channels_static: int,
+        input_scalers_mu: Tensor,
+        input_scalers_sigma: Tensor,
+        input_scalers_epsilon: float,
+        static_input_scalers_mu: Tensor,
+        static_input_scalers_sigma: Tensor,
+        static_input_scalers_epsilon: float,
+        output_scalers: Tensor,
+        n_lats_px: int,
+        n_lons_px: int,
+        patch_size_px: tuple[int],
+        mask_unit_size_px: tuple[int],
+        mask_ratio_inputs: float,
+        embed_dim: int,
+        n_blocks_encoder: int,
+        n_blocks_decoder: int,
+        mlp_multiplier: float,
+        n_heads: int,
+        dropout: float,
+        drop_path: float,
+        parameter_dropout: float,
+        residual: str,
+        masking_mode: str,
+        positional_encoding: str,
+        decoder_shifting: bool = False,
+        checkpoint_encoder: list[int] | None = None,
+        checkpoint_decoder: list[int] | None = None,
+    ) -> None:
+        """
+        Args:
+            in_channels: Number of input channels.
+            input_size_time: Number of timestamps in input.
+            in_channels_static: Number of input channels for static data.
+            input_scalers_mu: Tensor of size (in_channels,). Used to rescale
+                input.
+            input_scalers_sigma: Tensor of size (in_channels,). Used to rescale
+                input.
+            input_scalers_epsilon: Float. Used to rescale input.
+            static_input_scalers_mu: Tensor of size (in_channels_static). Used
+                to rescale static inputs.
+            static_input_scalers_sigma: Tensor of size (in_channels_static).
+                Used to rescale static inputs.
+            static_input_scalers_epsilon: Float. Used to rescale static inputs.
+            output_scalers: Tensor of shape (in_channels,). Used to rescale
+                output.
+            n_lats_px: Total latitudes in data. In pixels.
+            n_lons_px: Total longitudes in data. In pixels.
+            patch_size_px: Patch size for tokenization. In pixels lat/lon.
+            mask_unit_size_px: Size of each mask unit. In pixels lat/lon.
+            mask_ratio_inputs: Masking ratio for inputs. 0 to 1.
+            embed_dim: Embedding dimension
+            n_blocks_encoder: Number of local-global transformer pairs in
+                encoder.
+            n_blocks_decoder: Number of local-global transformer pairs in
+                decoder.
+            mlp_multiplier: MLP multiplier for hidden features in feed forward
+                networks.
+            n_heads: Number of attention heads.
+            dropout: Dropout.
+            drop_path: DropPath.
+            parameter_dropout: Dropout applied to parameters.
+            residual: Indicates whether and how model should work as residual
+                model. Accepted values are 'climate', 'temporal' and 'none'
+            positional_encoding: possible values are
+              ['absolute' (default), 'fourier'].
+                'absolute'  lat lon encoded in 3 dimensions using sine and
+                  cosine
+                'fourier' lat/lon to be encoded using various frequencies
+            masking_mode: String ['local', 'global', 'both'] that controls the
+                type of masking used.
+            checkpoint_encoder: List of integers controlling if gradient
+              checkpointing is used on encoder.
+                Format: [] for no gradient checkpointing. [3, 7] for
+                  checkpointing after 4th and 8th layer etc.
+            checkpoint_decoder: List of integers controlling if gradient
+              checkpointing is used on decoder.
+                Format: See `checkpoint_encoder`.
+            masking_mode: The type of masking to use
+              {'global', 'local', 'both'}
+            decoder_shifting: Whether to use swin shifting in the decoder.
+        """
+        super().__init__()
+
+        self.in_channels = in_channels
+        self.input_size_time = input_size_time
+        self.in_channels_static = in_channels_static
+        self.n_lats_px = n_lats_px
+        self.n_lons_px = n_lons_px
+        self.patch_size_px = patch_size_px
+        self.mask_unit_size_px = mask_unit_size_px
+        self.mask_ratio_inputs = mask_ratio_inputs
+        self.embed_dim = embed_dim
+        self.n_blocks_encoder = n_blocks_encoder
+        self.n_blocks_decoder = n_blocks_decoder
+        self.mlp_multiplier = mlp_multiplier
+        self.n_heads = n_heads
+        self.dropout = dropout
+        self.drop_path = drop_path
+        self.residual = residual
+        self._decoder_shift = decoder_shifting
+        self.positional_encoding = positional_encoding
+        self._checkpoint_encoder = checkpoint_encoder
+        self._checkpoint_decoder = checkpoint_decoder
+
+        assert self.n_lats_px % self.mask_unit_size_px[0] == 0
+        assert self.n_lons_px % self.mask_unit_size_px[1] == 0
+        assert self.mask_unit_size_px[0] % self.patch_size_px[0] == 0
+        assert self.mask_unit_size_px[1] % self.patch_size_px[1] == 0
+
+        if self.patch_size_px[0] != self.patch_size_px[1]:
+            raise NotImplementedError(
+                "Current pixel shuffle symmetric patches."
+            )
+
+        self.local_shape_mu = (
+            self.mask_unit_size_px[0] // self.patch_size_px[0],
+            self.mask_unit_size_px[1] // self.patch_size_px[1],
+        )
+        self.global_shape_mu = (
+            self.n_lats_px // self.mask_unit_size_px[0],
+            self.n_lons_px // self.mask_unit_size_px[1],
+        )
+
+        assert input_scalers_mu.shape == (in_channels,)
+        assert input_scalers_sigma.shape == (in_channels,)
+        assert output_scalers.shape == (in_channels,)
+
+        if self.positional_encoding != "fourier":
+            assert static_input_scalers_mu.shape == (in_channels_static,)
+            assert static_input_scalers_sigma.shape == (in_channels_static,)
+
+        # Input shape [batch, time, parameter, lat, lon]
+        self.input_scalers_epsilon = input_scalers_epsilon
+        self.register_buffer(
+            "input_scalers_mu", input_scalers_mu.reshape(1, 1, -1, 1, 1)
+        )
+        self.register_buffer(
+            "input_scalers_sigma", input_scalers_sigma.reshape(1, 1, -1, 1, 1)
+        )
+
+        # Static inputs shape [batch, parameter, lat, lon]
+        self.static_input_scalers_epsilon = static_input_scalers_epsilon
+        self.register_buffer(
+            "static_input_scalers_mu",
+            static_input_scalers_mu.reshape(1, -1, 1, 1),
+        )
+        self.register_buffer(
+            "static_input_scalers_sigma",
+            static_input_scalers_sigma.reshape(1, -1, 1, 1),
+        )
+
+        # Output shape [batch, parameter, lat, lon]
+        self.register_buffer(
+            "output_scalers", output_scalers.reshape(1, -1, 1, 1)
+        )
+
+        self.parameter_dropout = nn.Dropout2d(p=parameter_dropout)
+
+        self.patch_embedding = PatchEmbed(
+            patch_size=patch_size_px,
+            channels=in_channels * input_size_time,
+            embed_dim=embed_dim,
+        )
+
+        if self.residual == "climate":
+            self.patch_embedding_static = PatchEmbed(
+                patch_size=patch_size_px,
+                channels=in_channels + in_channels_static,
+                embed_dim=embed_dim,
+            )
+        else:
+            self.patch_embedding_static = PatchEmbed(
+                patch_size=patch_size_px,
+                channels=in_channels_static,
+                embed_dim=embed_dim,
+            )
+
+        self.input_time_embedding = nn.Linear(1, embed_dim // 4, bias=True)
+        self.lead_time_embedding = nn.Linear(1, embed_dim // 4, bias=True)
+
+        self.mask_token = nn.Parameter(torch.randn(1, 1, 1, self.embed_dim))
+        self._nglobal_mu = np.prod(self.global_shape_mu)
+        self._global_idx = torch.arange(self._nglobal_mu)
+
+        self._nlocal_mu = np.prod(self.local_shape_mu)
+        self._local_idx = torch.arange(self._nlocal_mu)
+
+        self.encoder = PrithviWxCEncoderDecoder(
+            embed_dim=embed_dim,
+            n_blocks=n_blocks_encoder,
+            mlp_multiplier=mlp_multiplier,
+            n_heads=n_heads,
+            dropout=dropout,
+            drop_path=drop_path,
+            transformer_cp=checkpoint_encoder,
+        )
+
+        if n_blocks_decoder != 0:
+            if self._decoder_shift:
+                self.decoder_shifter = d_shifter = SWINShift(
+                    self.mask_unit_size_px,
+                    self.global_shape_mu,
+                    self.local_shape_mu,
+                    self.patch_size_px,
+                    n_context_tokens=0,
+                )
+            else:
+                self.decoder_shifter = d_shifter = None
+
+            self.decoder = PrithviWxCEncoderDecoder(
+                embed_dim=embed_dim,
+                n_blocks=n_blocks_decoder,
+                mlp_multiplier=mlp_multiplier,
+                n_heads=n_heads,
+                dropout=dropout,
+                drop_path=0.0,
+                shifter=d_shifter,
+                transformer_cp=checkpoint_decoder,
+            )
+
+            self.unembed = nn.Linear(
+                self.embed_dim,
+                self.in_channels
+                * self.patch_size_px[0]
+                * self.patch_size_px[1],
+                bias=True,
+            )
+
+        self.masking_mode = masking_mode.lower()
+        match self.masking_mode:
+            case "local":
+                self.generate_mask = self._gen_mask_local
+            case "global":
+                self.generate_mask = self._gen_mask_global
+            case "both":
+                self._mask_both_local: bool = True
+                self.generate_mask = self._gen_mask_both
+            case _:
+                raise ValueError(
+                    f"Masking mode '{masking_mode}' not supported"
+                )
+
+    def swap_masking(self) -> None:
+        self._mask_both_local = not self._mask_both_local
+
+    @cached_property
+    def n_masked_global(self):
+        return int(self.mask_ratio_inputs * np.prod(self.global_shape_mu))
+
+    @cached_property
+    def n_masked_local(self):
+        return int(self.mask_ratio_inputs * np.prod(self.local_shape_mu))
+
+    @staticmethod
+    def _shuffle_along_axis(a, axis):
+        idx = torch.argsort(input=torch.rand(*a.shape), dim=axis)
+        return torch.gather(a, dim=axis, index=idx)
+
+    def _gen_mask_local(self, sizes: tuple[int]) -> tuple[Tensor]:
+        """
+        Args:
+            batch_size: Number of elements in batch
+        Returns:
+            Tuple of torch tensors. [indices masked, indices unmasked].
+            Each of these is a tensor of shape (batch, global sequene)
+        """
+        # Identify which indices (values) should be masked
+
+        maskable_indices = self._local_idx.view(1, -1).expand(*sizes[:2], -1)
+
+        maskable_indices = self._shuffle_along_axis(maskable_indices, 2)
+
+        indices_masked = maskable_indices[:, :, : self.n_masked_local]
+        indices_unmasked = maskable_indices[:, :, self.n_masked_local :]
+
+        return indices_masked, indices_unmasked
+
+    def _gen_mask_global(self, sizes: tuple[int]) -> tuple[Tensor]:
+        """
+        Args:
+            batch_size: Number of elements in batch
+        Returns:
+            Tuple of torch tensors. [indices masked, indices unmasked].
+            Each of these is a tensor of shape (batch, global sequene)
+        """
+        # Identify which indices (values) should be masked
+
+        maskable_indices = self._global_idx.view(1, -1).expand(*sizes[:1], -1)
+
+        maskable_indices = self._shuffle_along_axis(maskable_indices, 1)
+
+        indices_masked = maskable_indices[:, : self.n_masked_global]
+        indices_unmasked = maskable_indices[:, self.n_masked_global :]
+
+        return indices_masked, indices_unmasked
+
+    def _gen_mask_both(self, sizes: tuple[int]) -> tuple[Tensor]:
+        if self._mask_both_local:
+            return self._gen_mask_local(sizes)
+        else:
+            return self._gen_mask_global(sizes)
+
+    @staticmethod
+    def reconstruct_batch(
+        idx_masked: Tensor,
+        idx_unmasked: Tensor,
+        data_masked: Tensor,
+        data_unmasked: Tensor,
+    ) -> Tensor:
+        """Reconstructs a tensor along the mask unit dimension. Batched
+        version.
+
+        Args:
+            idx_masked: Tensor of shape `batch, mask unit sequence`.
+            idx_unmasked: Tensor of shape `batch, mask unit sequence`.
+            data_masked: Tensor of shape `batch, mask unit sequence, ...`.
+                Should have same size along mask unit sequence dimension as
+                idx_masked. Dimensions beyond the first two, marked here as ...
+                will typically be `local_sequence, channel` or
+                `channel, lat, lon`. These dimensions should agree with
+                data_unmasked.
+            data_unmasked: Tensor of shape `batch, mask unit sequence, ...`.
+                Should have same size along mask unit sequence dimension as
+                idx_unmasked. Dimensions beyond the first two, marked here as
+                ... will typically be `local_sequence, channel` or `channel,
+                lat, lon`. These dimensions should agree with data_masked.
+        Returns:
+            Tensor: Tensor of same shape as inputs data_masked and
+                data_unmasked. I.e. `batch, mask unit sequence, ...`. Index for
+                the total data composed of the masked and the unmasked part.
+        """
+        dim: int = idx_masked.ndim
+
+        idx_total = torch.argsort(
+            torch.cat([idx_masked, idx_unmasked], dim=-1), dim=-1
+        )
+        idx_total = idx_total.view(
+            *idx_total.shape, *[1] * (data_unmasked.ndim - dim)
+        )
+        idx_total = idx_total.expand(
+            *idx_total.shape[:dim], *data_unmasked.shape[dim:]
+        )
+
+        data = torch.cat([data_masked, data_unmasked], dim=dim - 1)
+        data = torch.gather(data, dim=dim - 1, index=idx_total)
+
+        return data, idx_total
+
+    def fourier_pos_encoding(self, x_static: Tensor) -> Tensor:
+        """
+        Args
+            x_static: B x C x H x W. first two channels are lat, and lon
+        Returns
+            Tensor: Tensor of shape B x E x H x W where E is the embedding
+                dimension.
+        """
+
+        # B x C x H x W -> B x 1 x H/P x W/P
+        latitudes_patch = F.avg_pool2d(
+            x_static[:, [0]],
+            kernel_size=self.patch_size_px,
+            stride=self.patch_size_px,
+        )
+        longitudes_patch = F.avg_pool2d(
+            x_static[:, [1]],
+            kernel_size=self.patch_size_px,
+            stride=self.patch_size_px,
+        )
+
+        modes = (
+            torch.arange(self.embed_dim // 4, device=x_static.device).view(
+                1, -1, 1, 1
+            )
+            + 1.0
+        )
+        pos_encoding = torch.cat(
+            (
+                torch.sin(latitudes_patch * modes),
+                torch.sin(longitudes_patch * modes),
+                torch.cos(latitudes_patch * modes),
+                torch.cos(longitudes_patch * modes),
+            ),
+            axis=1,
+        )
+
+        return pos_encoding  # B x E x H/P x W/P
+
+    def time_encoding(self, input_time, lead_time):
+        """
+        Args:
+            input_time: Tensor of shape [batch].
+            lead_time: Tensor of shape [batch].
+        Returns:
+            Tensor: Tensor of shape [batch, embed_dim, 1, 1]
+        """
+        input_time = self.input_time_embedding(input_time.view(-1, 1, 1, 1))
+        lead_time = self.lead_time_embedding(lead_time.view(-1, 1, 1, 1))
+
+        time_encoding = torch.cat(
+            (
+                torch.cos(input_time),
+                torch.cos(lead_time),
+                torch.sin(input_time),
+                torch.sin(lead_time),
+            ),
+            axis=3,
+        )
+        return time_encoding
+
+    def to_patching(self, x: Tensor) -> Tensor:
+        """Transform data from lat/lon space to two axis patching
+
+        Args: ->
+            x: Tesnor in lat/lon space (N, C, Nlat//P_0, Nlon//P_1)
+
+        Returns:
+            Tensor in patch space (N, G, L, C)
+        """
+        n_batch = x.shape[0]
+
+        x = x.view(
+            n_batch,
+            -1,
+            self.global_shape_mu[0],
+            self.local_shape_mu[0],
+            self.global_shape_mu[1],
+            self.local_shape_mu[1],
+        )
+        x = x.permute(0, 2, 4, 3, 5, 1).contiguous()
+
+        s = x.shape
+        return x.view(n_batch, s[1] * s[2], s[3] * s[4], -1)
+
+    def from_patching(self, x: Tensor) -> Tensor:
+        """Transform data from two axis patching to lat/lon space
+
+        Args:
+            x: Tensor in patch space with shape (N, G, L, C*P_0*P_1)
+
+        Returns:
+            Tensor: Tensor in lat/lon space
+                (N, C*P_0*P_1, Nlat//P_0, Nlon // P_1)
+        """
+        n_batch = x.shape[0]
+
+        x = x.view(
+            n_batch,
+            self.global_shape_mu[0],
+            self.global_shape_mu[1],
+            self.local_shape_mu[0],
+            self.local_shape_mu[1],
+            -1,
+        )
+        x = x.permute(0, 5, 1, 3, 2, 4).contiguous()
+
+        s = x.shape
+        return x.view(n_batch, -1, s[2] * s[3], s[4] * s[5])
+
+    def forward(self, batch: dict[str, torch.Tensor]) -> torch.Tensor:
+        """
+        Args:
+            batch: Dictionary the following keys::
+
+                'x': Tensor of shape [batch, time, parameter, lat, lon]
+                'y': Tensor of shape [batch, parameter, lat, lon]
+                'static': Tensor of shape [batch, channel_static, lat, lon]
+                'climate': Optional tensor of shape [batch, parameter, lat, lon]
+                'input_time': Tensor of shape [batch]. Or none.
+                'lead_time': Tensor of shape [batch]. Or none.
+
+        Returns:
+            Tensor: Tensor of shape [batch, parameter, lat, lon].
+        """  # noqa: E501
+        x_rescaled = (batch["x"] - self.input_scalers_mu) / (
+            self.input_scalers_sigma + self.input_scalers_epsilon
+        )
+        batch_size = x_rescaled.shape[0]
+
+        if self.positional_encoding == "fourier":
+            x_static_pos = self.fourier_pos_encoding(batch["static"])
+            x_static = (
+                batch["static"][:, 2:] - self.static_input_scalers_mu[:, 3:]
+            ) / (
+                self.static_input_scalers_sigma[:, 3:]
+                + self.static_input_scalers_epsilon
+            )
+        else:
+            x_static = (batch["static"] - self.static_input_scalers_mu) / (
+                self.static_input_scalers_sigma
+                + self.static_input_scalers_epsilon
+            )
+
+        if self.residual == "temporal":
+            # We create a residual of same shape as y
+            index = torch.where(
+                batch["lead_time"] > 0, batch["x"].shape[1] - 1, 0
+            )
+            index = index.view(-1, 1, 1, 1, 1)
+            index = index.expand(batch_size, 1, *batch["x"].shape[2:])
+            x_hat = torch.gather(batch["x"], dim=1, index=index)
+            x_hat = x_hat.squeeze(1)
+        elif self.residual == "climate":
+            climate_scaled = (
+                batch["climate"] - self.input_scalers_mu.view(1, -1, 1, 1)
+            ) / (
+                self.input_scalers_sigma.view(1, -1, 1, 1)
+                + self.input_scalers_epsilon
+            )
+
+        # [batch, time, parameter, lat, lon]
+        # -> [batch, time x parameter, lat, lon]
+        x_rescaled = x_rescaled.flatten(1, 2)
+        # Parameter dropout
+        x_rescaled = self.parameter_dropout(x_rescaled)
+
+        x_embedded = self.patch_embedding(x_rescaled)
+
+        if self.residual == "climate":
+            static_embedded = self.patch_embedding_static(
+                torch.cat((x_static, climate_scaled), dim=1)
+            )
+        else:
+            static_embedded = self.patch_embedding_static(x_static)
+
+        if self.positional_encoding == "fourier":
+            static_embedded += x_static_pos
+
+        x_embedded = self.to_patching(x_embedded)
+        static_embedded = self.to_patching(static_embedded)
+
+        time_encoding = self.time_encoding(
+            batch["input_time"], batch["lead_time"]
+        )
+
+        tokens = x_embedded + static_embedded + time_encoding
+
+        # Now we generate masks based on masking_mode
+        indices_masked, indices_unmasked = self.generate_mask(
+            (batch_size, self._nglobal_mu)
+        )
+        indices_masked = indices_masked.to(device=tokens.device)
+        indices_unmasked = indices_unmasked.to(device=tokens.device)
+        maskdim: int = indices_masked.ndim
+
+        # Unmasking
+        unmask_view = (*indices_unmasked.shape, *[1] * (tokens.ndim - maskdim))
+        unmasked = torch.gather(
+            tokens,
+            dim=maskdim - 1,
+            index=indices_unmasked.view(*unmask_view).expand(
+                *indices_unmasked.shape, *tokens.shape[maskdim:]
+            ),
+        )
+
+        # Encoder
+        x_encoded = self.encoder(unmasked)
+
+        # Generate and position encode the mask tokens
+        # [1, 1, 1, embed_dim]
+        # -> [batch, global_seq_masked, local seq, embed_dim]
+        mask_view = (*indices_masked.shape, *[1] * (tokens.ndim - maskdim))
+        masking = self.mask_token.repeat(*static_embedded.shape[:3], 1)
+        masked = masking + static_embedded
+        masked = torch.gather(
+            masked,
+            dim=maskdim - 1,
+            index=indices_masked.view(*mask_view).expand(
+                *indices_masked.shape, *tokens.shape[maskdim:]
+            ),
+        )
+
+        recon, _ = self.reconstruct_batch(
+            indices_masked, indices_unmasked, masked, x_encoded
+        )
+
+        x_decoded = self.decoder(recon)
+
+        # Output: [batch, global sequence, local sequence,
+        #          in_channels * patch_size[0] * patch_size[1]]
+        x_unembed = self.unembed(x_decoded)
+
+        # Reshape to [batch, global_lat, global_lon, local_lat, local_lon,
+        #             in_channels * patch_size[0] * patch_size[1]]
+        x_out = self.from_patching(x_unembed)
+
+        # Pixel shuffle to [batch, in_channels, lat, lon]
+        x_out = F.pixel_shuffle(x_out, self.patch_size_px[0])
+
+        if self.residual == "temporal":
+            x_out = self.output_scalers * x_out + x_hat
+        elif self.residual == "climate":
+            x_out = self.output_scalers * x_out + batch["climate"]
+        elif self.residual == "none":
+            x_out = (
+                self.output_scalers * x_out
+                + self.input_scalers_mu.reshape(1, -1, 1, 1)
+            )
+
+        return x_out

PrithviWxC/rollout.py

@@ -0,0 +1,49 @@
+import torch
+from torch import Tensor, nn
+
+
+def rollout_iter(
+    nsteps: int,
+    model: nn.Module,
+    batch: dict[str, Tensor | int | float],
+) -> Tensor:
+    """A helper function for performing autoregressive rollout.
+
+    Args:
+        nsteps (int): The number of rollout steps to take
+        model (nn.Module): A model.
+        batch (dict): A data dictionary common to the Prithvi models.
+
+    Raises:
+        ValueError: If the number of steps isn't positive.
+
+    Returns:
+        Tensor: the output of the model after nsteps autoregressive iterations.
+    """
+    if nsteps < 1:
+        raise ValueError("'nsteps' shouold be a positive int.")
+
+    xlast = batch["x"][:, 1]
+    batch["lead_time"] = batch["lead_time"][..., 0]
+
+    # Save the masking ratio to be restored later
+    mask_ratio_tmp = model.mask_ratio_inputs
+
+    for step in range(nsteps):
+        # After first step, turn off masking
+        if step > 0:
+            model.mask_ratio_inputs = 0.0
+
+        batch["static"] = batch["statics"][:, step]
+        batch["climate"] = batch["climates"][:, step]
+        batch["y"] = batch["ys"][:, step]
+
+        out = model(batch)
+
+        batch["x"] = torch.cat((xlast[:, None], out[:, None]), dim=1)
+        xlast = out
+
+    # Restore the masking ratio
+    model.mask_ratio_inputs = mask_ratio_tmp
+
+    return xlast