Prithvi.py

import streamlit as st
import torch
import random
import numpy as np
import yaml
from pathlib import Path
from io import BytesIO
import random
from pathlib import Path
import matplotlib.pyplot as plt
import numpy as np
import torch
from huggingface_hub import hf_hub_download, snapshot_download
import tempfile
import traceback
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
import logging
from Prithvi import *

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)

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