official-src/aardvark/loader.py

import time as timelib
from time import time

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

from loader_utils_new import *
from data_shapes import *


class WeatherDataset(Dataset):
    """
    Base weather dataset class
    """

    def __init__(
        self,
        device,
        hadisd_mode,
        start_date,
        end_date,
        lead_time,
        era5_mode="train",
        res=1,
        filter_dates=None,
        diff=None,
    ):

        super().__init__()

        # Setup
        self.device = device
        self.mode = hadisd_mode
        self.data_path = "path_to_data/"
        self.aux_data_path = "path_to_auxiliary_data/"
        self.start_date = start_date
        self.end_date = end_date
        self.lead_time = lead_time
        self.era5_mode = era5_mode
        self.res = res
        self.filter_dates = filter_dates
        self.diff = diff

        # Date indexing
        self.dates = pd.date_range(start_date, end_date, freq="6H")
        if self.filter_dates == "start":
            self.index = np.array([i for i, d in enumerate(self.dates) if d.month < 7])
        elif self.filter_dates == "end":
            self.index = np.array([i for i, d in enumerate(self.dates) if d.month >= 7])
        else:
            self.index = np.array(range(len(self.dates)))

        # Load the input modalities
        print("Loading IGRA")
        self.load_igra()

        print("Loading AMSU-A")
        self.load_amsua()

        print("Loading AMSU-B")
        self.load_amsub()

        print("Loading ICOADS")
        self.load_icoads()

        print("Loading IASI")
        self.load_iasi()

        print("Loading GEO")
        self.load_sat_data()

        print("Loading HADISD")
        self.load_hadisd(self.mode)

        print("Loading ASCAT")
        self.load_ascat_data()
        self.load_hirs_data()

        # Load the ground truth data for training
        print("Loading ERA5")
        self.era5_sfc = [
            self.load_era5(year)
            for year in range(int(start_date[:4]), int(end_date[:4]) + 1)
        ]

        # Internal grid to longitude latitude correspondence
        self.era5_x = [
            self.to_tensor(
                np.load(self.data_path + "era5/era5_x_{}.npy".format(self.res))
            )
            / LATLON_SCALE_FACTOR,
            self.to_tensor(
                np.load(self.data_path + "era5/era5_y_{}.npy".format(self.res))
            )
            / LATLON_SCALE_FACTOR,
        ]

        # Orography
        self.era5_elev = self.to_tensor(
            np.load(self.data_path + "era5/elev_vars_{}.npy".format(self.res))
        )
        self.era5_elev = torch.flip(self.era5_elev.permute(0, 2, 1), [-1])
        xx, yy = torch.meshgrid(self.era5_x[0], self.era5_x[1])
        self.era5_lonlat = torch.stack([xx, yy])

        # Climatology
        self.climatology = np.memmap(
            self.data_path + "climatology_data.mmap",
            dtype="float32",
            mode="r",
            shape=CLIMATOLOGY_SHAPE,
        )

        # Setup normalisation factors
        if self.diff:
            self.era5_mean_spatial = np.load(
                self.aux_data_path + "era5_spatial_means.npy"
            )[0, ...]
            self.means = np.load(self.aux_data_path + "era5_avdiff_means.npy")[
                :, np.newaxis, np.newaxis, ...
            ]
            self.stds = np.load(self.aux_data_path + "era5_avdiff_stds.npy")[
                :, np.newaxis, np.newaxis, ...
            ]
        else:
            self.means = np.load(
                self.aux_data_path
                + "norm_factors/mean_{}_{}.npy".format(self.era5_mode, self.res)
            )[:, np.newaxis, np.newaxis, ...]
            self.stds = np.load(
                self.aux_data_path
                + "norm_factors/std_{}_{}.npy".format(self.era5_mode, self.res)
            )[:, np.newaxis, np.newaxis, ...]

    def load_icoads(self):
        """
        Load the ICOADS data
        """

        self.icoads_y = np.memmap(
            self.data_path + "icoads/1999_2021_icoads_y.mmap",
            dtype="float32",
            mode="r",
            shape=ICOADS_Y_SHAPE,
        )

        self.icoads_x = (
            np.memmap(
                self.data_path + "icoads/1999_2021_icoads_x.mmap",
                dtype="float32",
                mode="r",
                shape=ICOADS_X_SHAPE,
            )
            / LATLON_SCALE_FACTOR
        )
        self.icoads_means = self.to_tensor(
            np.load(self.aux_data_path + "norm_factors/mean_icoads.npy")
        )
        self.icoads_stds = self.to_tensor(
            np.load(self.aux_data_path + "norm_factors/std_icoads.npy")
        )
        self.icoads_means = self.to_tensor(
            np.nanmean(self.icoads_y[-365 * 4 :, ...], axis=(0, 2))[:, np.newaxis]
        )
        self.icoads_stds = self.to_tensor(
            np.nanstd(self.icoads_y[-365 * 4 :, ...], axis=(0, 2))[:, np.newaxis]
        )
        self.icoads_index_offset = ICOADS_OFFSETS[self.start_date]
        return

    def load_igra(self):
        """
        Load the IGRA data
        """

        self.igra_y = np.memmap(
            self.data_path + "igra/1999_2021_igra_y.mmap",
            dtype="float32",
            mode="r",
            shape=IGRA_Y_SHAPE,
        )

        self.igra_x = np.copy(
            np.memmap(
                self.data_path + "igra/1999_2021_igra_x.mmap",
                dtype="float32",
                mode="r",
                shape=IGRA_X_SHAPE,
            )
        )
        self.igra_x = self.igra_x / LATLON_SCALE_FACTOR

        self.igra_means = self.to_tensor(
            np.load(self.aux_data_path + "norm_factors/mean_igra.npy")
        )
        self.igra_stds = self.to_tensor(
            np.load(self.aux_data_path + "norm_factors/std_igra.npy")
        )

        self.igra_index_offset = IGRA_OFFSETS[self.start_date]

        return

    def load_amsua(self):
        """
        Load the AMSU-A data
        """

        self.amsua_y = np.memmap(
            self.data_path + "amsua/2007_2021_amsua.mmap",
            dtype="float32",
            mode="r",
            shape=AMSUA_Y_SHAPE,
        )
        self.amsua_index_offset = AMSUA_OFFSETS[self.start_date]

        xx = np.linspace(-180, 179, 360, dtype=np.float32)
        xx = ((xx + 360) % 360) / LATLON_SCALE_FACTOR
        yy = np.linspace(90, -90, 180, dtype=np.float32) / LATLON_SCALE_FACTOR
        self.amsua_x = [xx, yy]

        self.amsua_means = self.to_tensor(
            np.load(self.aux_data_path + "norm_factors/mean_amsua.npy")
        )
        self.amsua_stds = self.to_tensor(
            np.load(self.aux_data_path + "norm_factors/std_amsua.npy")
        )

        return

    def load_amsub(self):
        """
        Load the AMSU-B data
        """

        self.amsub_y = np.memmap(
            self.data_path + "amsub_mhs/2007_2021_amsub.mmap",
            dtype="float32",
            mode="r",
            shape=AMSUB_Y_SHAPE,
        )
        self.amsub_index_offset = AMSUB_OFFSETS[self.start_date]

        xx = np.linspace(0, 359, 360, dtype=np.float32)
        xx = ((xx + 360) % 360) / LATLON_SCALE_FACTOR
        yy = np.linspace(90, -90, 181, dtype=np.float32) / LATLON_SCALE_FACTOR
        self.amsub_x = [xx, yy]

        self.amsub_means = self.to_tensor(
            np.load(self.aux_data_path + "norm_factors/mean_amsub.npy")
        )
        self.amsub_stds = self.to_tensor(
            np.load(self.aux_data_path + "norm_factors/std_amsub.npy")
        )

        return

    def load_ascat_data(self):
        """
        Load the ASCAT data
        """

        self.ascat_y = np.memmap(
            self.data_path + "ascat/2007_2021_ascat.mmap",
            dtype="float32",
            mode="r",
            shape=ASCAT_Y_SHAPE,
        )
        self.ascat_index_offset = ASCAT_OFFSETS[self.start_date]

        xx = np.linspace(0, 359, 360, dtype=np.float32)
        xx = ((xx + 360) % 360) / LATLON_SCALE_FACTOR
        yy = np.linspace(-90, 90, 181, dtype=np.float32) / LATLON_SCALE_FACTOR
        self.ascat_x = [xx, np.copy(yy[::-1])]

        self.ascat_means = self.to_tensor(
            np.load(self.aux_data_path + "norm_factors/mean_ascat.npy")
        )
        self.ascat_stds = self.to_tensor(
            np.load(self.aux_data_path + "norm_factors/std_ascat.npy")
        )

        return

    def load_hirs_data(self):
        """
        Load the HIRS data
        """

        self.hirs_y = np.memmap(
            self.data_path + "hirs/2007_2021_hirs.mmap",
            dtype="float32",
            mode="r",
            shape=HIRS_Y_SHAPE,
        )
        self.hirs_index_offset = ASCAT_OFFSETS[self.start_date]

        xx = np.linspace(0, 359, 360, dtype=np.float32)
        xx = ((xx + 360) % 360) / LATLON_SCALE_FACTOR
        yy = np.linspace(-90, 90, 181, dtype=np.float32) / LATLON_SCALE_FACTOR
        self.hirs_x = [xx, np.copy(yy[::-1])]

        self.hirs_means = self.to_tensor(
            np.load(self.aux_data_path + "norm_factors/hirs_means.npy")
        )
        self.hirs_stds = self.to_tensor(
            np.load(self.aux_data_path + "norm_factors/hirs_stds.npy")
        )

        return

    def load_sat_data(self):
        """
        Load the GRIDSAT data
        """

        self.sat_y = np.memmap(
            self.data_path + "gridsat/gridsat_data.mmap",
            dtype="float32",
            mode="r",
            shape=GRIDSAT_Y_SHAPE,
        )

        xx = np.load(self.data_path + "gridsat/sat_x.npy") / LATLON_SCALE_FACTOR
        yy = np.load(self.data_path + "gridsat/sat_y.npy") / LATLON_SCALE_FACTOR
        self.sat_x = [xx, yy]
        self.sat_index_offset = SAT_OFFSETS[self.start_date]

        self.sat_means = self.to_tensor(
            np.load(self.aux_data_path + "norm_factors/mean_sat.npy")
        )
        self.sat_stds = self.to_tensor(
            np.load(self.aux_data_path + "norm_factors/std_sat.npy")
        )

        return

    def load_iasi(self):
        """
        Load the IASI data
        """

        self.iasi = np.memmap(
            self.data_path + "2007_2021_iasi_subset.mmap",
            dtype="float32",
            mode="r",
            shape=IASI_Y_SHAPE,
        )
        self.iasi_index_offset = ASCAT_OFFSETS[self.start_date]

        xx = np.linspace(0, 359, 360, dtype=np.float32)
        xx = ((xx + 360) % 360) / LATLON_SCALE_FACTOR
        yy = np.linspace(-90, 90, 181, dtype=np.float32) / LATLON_SCALE_FACTOR
        self.iasi_x = [xx, np.copy(yy[::-1])]

        self.iasi_means = self.to_tensor(
            np.load(self.aux_data_path + "norm_factors/mean_iasi.npy")
        )
        self.iasi_stds = self.to_tensor(
            np.load(self.aux_data_path + "norm_factors/std_iasi.npy")
        )

        return

    def load_hadisd(self, mode):
        """
        Load the HADISD data
        """

        self.hadisd_x = []
        self.hadisd_alt = []
        self.hadisd_y = []
        hadisd_vars = ["tas", "tds", "psl", "u", "v"]
        for var in hadisd_vars:
            lon = lon_to_0_360(
                np.load(
                    self.data_path + "hadisd_processed/{}_lon_{}.npy".format(var, mode)
                )
            )
            lat = np.load(
                self.data_path + "hadisd_processed/{}_lat_{}.npy".format(var, mode)
            )
            alt = np.load(
                self.data_path + "hadisd_processed/{}_alt_{}.npy".format(var, mode)
            )

            vals = np.memmap(
                self.data_path
                + "hadisd_processed/{}_vals_{}.memmap".format(var, self.mode),
                dtype="float32",
                mode="r",
                shape=get_hadisd_shape(mode),
            )

            self.hadisd_x.append(np.stack([lon, lat], axis=-1) / LATLON_SCALE_FACTOR)
            self.hadisd_alt.append(alt)
            self.hadisd_y.append(vals)

        self.hadisd_index_offset = HADISD_OFFSETS[self.start_date]

        self.hadisd_means = [
            self.to_tensor(
                np.load(
                    self.aux_data_path + "norm_factors/mean_hadisd_{}.npy".format(var)
                )
            )
            for var in hadisd_vars
        ]
        self.hadisd_stds = [
            self.to_tensor(
                np.load(
                    self.aux_data_path + "norm_factors/std_hadisd_{}.npy".format(var)
                )
            )
            for var in hadisd_vars
        ]

        return

    def load_era5(self, year):
        """
        Load the ERA5 training data
        """

        if year % 4 == 0:
            d = 366 * 4
        else:
            d = 365 * 4

        if self.era5_mode == "sfc":
            levels = 4
        elif self.era5_mode == "13u":
            levels = 69
        else:
            levels = 24

        if self.res == 1:
            x = 240
            y = 121
        elif self.res == 5:
            x = 64
            y = 32
        mmap = np.memmap(
            self.data_path
            + "/era5/era5_{}_{}_6_{}.memmap".format(self.era5_mode, self.res, year),
            dtype="float32",
            mode="r",
            shape=(d, levels, x, y),
        )
        return mmap

    def norm_era5(self, x):

        x = (x - self.means) / self.stds
        return x

    def unnorm_era5(self, x):

        x = x * self.stds + self.means
        return x

    def norm_data(self, x, means, stds):
        return (x - means) / stds

    def norm_hadisd(self, x):
        for i in range(5):
            x[i] = (x[i] - self.hadisd_means[i]) / self.hadisd_stds[i]
        return x

    def __len__(self):
        return self.index.shape[0] - 1 - 1

    def to_tensor(self, arr):
        return torch.from_numpy(arr).float().to(self.device)

    def get_time_aux(self, current_date):
        """
        Return the auxiliary temporal channels given a date
        """

        doy = current_date.dayofyear
        year = (current_date.year - 2007) / 15
        time_of_day = current_date.hour
        return np.array(
            [
                np.cos(np.pi * 2 * doy / DAYS_IN_YEAR),
                np.sin(np.pi * 2 * doy / DAYS_IN_YEAR),
                np.cos(np.pi * 2 * time_of_day / 24),
                np.sin(np.pi * 2 * time_of_day / 24),
                year,
            ]
        )


class WeatherDatasetAssimilation(WeatherDataset):
    """
    Encoder training loader
    """

    def __init__(
        self,
        device,
        hadisd_mode,
        start_date,
        end_date,
        lead_time,
        era5_mode="sfc",
        res=1,
        filter_dates=None,
        var_start=0,
        var_end=24,
        diff=False,
        two_frames=False,
    ):

        super().__init__(
            device,
            hadisd_mode,
            start_date,
            end_date,
            lead_time,
            era5_mode,
            res=res,
            filter_dates=filter_dates,
            diff=diff,
        )

        # Setup

        self.var_start = var_start
        self.var_end = var_end
        self.diff = diff
        self.two_frames = two_frames

    def load_era5_time(self, index):
        """
        ERA5 ground truth data loading
        """

        date = self.dates[index]
        year = date.year
        hour = date.hour
        doy = (date.dayofyear - 1) * 4 + (hour // 6)

        era5 = self.era5_sfc[year - int(self.start_date[:4])][doy, ...]
        era5 = np.copy(era5)
        if self.diff:
            era5 = era5 - self.era5_mean_spatial
        era5 = self.norm_era5(era5[np.newaxis, ...])[0, ...]
        return era5

    def load_year_end(self, year, doy):
        data_1 = self.era5_sfc[year - int(self.start_date[:4])][doy : doy + 1, ...]
        missing = self.lead_time - data_1.shape[0] + 1
        data_2 = self.era5_sfc[year - int(self.start_date[:4]) + 1][:missing, ...]
        data = np.concatenate([data_1, data_2])
        return data

    def load_era5_slice(self, index):
        """
        ERA5 ground truth data loading
        """

        date = self.dates[index]
        year = date.year
        doy = (date.dayofyear - 1) * 4

        next_date = self.dates[index + 1]
        next_year = next_date.year

        if next_year != year:
            era5 = self.load_year_end(year, doy)
        else:
            era5 = self.era5_sfc[year - int(self.start_date[:4])][doy : doy + 1, ...]

        era5 = self.norm_era5(np.copy(era5))
        return era5

    def __getitem__(self, index):

        if self.two_frames:
            # Case 1: loading t=0 and t=-1
            index = index + 1
            current = self.get_index(index, "current")
            prev = self.get_index(index - 1, "prev")
            current["y_target"] = current["y_target_current"]

            return {**current, **prev}
        else:
            # Case 2: loading t=0
            current = self.get_index(index, "current")
            current["y_target"] = current["y_target_current"]

            return {**current}

    def unnorm_pred(self, x):
        dev = x.device
        x = x.detach().cpu().numpy()

        x = (
            x
            * self.stds[np.newaxis, ...].transpose(0, 2, 3, 1)[
                ..., self.var_start : self.var_end
            ]
            + self.means[np.newaxis, ...].transpose(0, 2, 3, 1)[
                ..., self.var_start : self.var_end
            ]
        )
        if bool(self.diff):
            x = (
                x
                + self.era5_mean_spatial[np.newaxis, ...].transpose(0, 3, 2, 1)[
                    ..., self.var_start : self.var_end
                ]
            )
        return torch.from_numpy(x).float().to(dev)

    def get_index(self, index, prefix):
        """
        Load data for the relevant index respecting different offsets depending on the modality
        """

        index = self.index[index]
        date = self.dates[index]

        # ICOADS
        icoads_x = self.icoads_x[index + self.icoads_index_offset, ...]
        icoads_y = self.icoads_y[index + self.icoads_index_offset, ...]
        icoads_x = [icoads_x[0, :], icoads_x[1, :]]
        icoads_x = [self.to_tensor(i) for i in icoads_x]
        icoads_y = self.to_tensor(icoads_y)
        icoads_y = self.norm_data(icoads_y, self.icoads_means, self.icoads_stds)

        # GRIDSAT
        sat_y = self.sat_y[index + self.sat_index_offset, ...]
        sat_x = [self.to_tensor(i) for i in self.sat_x]
        sat_y = self.to_tensor(sat_y)
        sat_y = self.norm_data(sat_y, self.sat_means, self.sat_stds)

        # AMSU-A
        amsua_y = self.to_tensor(self.amsua_y[index + self.amsua_index_offset, ...])
        amsua_y[amsua_y < -998] = torch.nan
        amsua_x = [self.to_tensor(i) for i in self.amsua_x]
        amsua_y[amsua_y < -998] = np.nan
        amsua_y = self.norm_data(amsua_y, self.amsua_means, self.amsua_stds)

        # AMSU-B
        amsub_y = self.to_tensor(self.amsub_y[index + self.amsub_index_offset, ...])
        amsub_y[amsub_y < -998] = torch.nan
        amsub_x = [self.to_tensor(i) for i in self.amsub_x]
        amsub_y[amsub_y < -998] = np.nan
        amsub_y = self.norm_data(amsub_y, self.amsub_means, self.amsub_stds)

        # IASI
        iasi_y = self.to_tensor(self.iasi[index + self.iasi_index_offset, ...])
        iasi_x = [self.to_tensor(i) for i in self.iasi_x]
        iasi_y = self.norm_data(iasi_y, self.iasi_means, self.iasi_stds)

        # IGRA
        igra_y = self.to_tensor(self.igra_y[index + self.igra_index_offset, ...])
        igra_x = [self.igra_x[:, 0], self.igra_x[:, 1]]
        igra_x = [self.to_tensor(i) for i in igra_x]
        igra_y = self.norm_data(igra_y, self.igra_means, self.igra_stds)

        # ASCAT
        ascat_y = self.to_tensor(self.ascat_y[index + self.ascat_index_offset, ...])
        ascat_x = [self.to_tensor(i) for i in self.ascat_x]
        ascat_y[..., 4][ascat_y[..., 4] < -990] = np.nan
        ascat_y = self.norm_data(ascat_y, self.ascat_means, self.ascat_stds)

        # HIRS
        hirs_y = self.to_tensor(self.hirs_y[index + self.hirs_index_offset, ...])
        hirs_y[hirs_y < -998] = np.nan
        hirs_x = [self.to_tensor(i) for i in self.hirs_x]
        hirs_y = self.norm_data(hirs_y, self.hirs_means, self.hirs_stds)

        # HadISD
        x_context_hadisd = self.hadisd_x
        y_context_hadisd = [
            i[index + self.hadisd_index_offset, :] for i in self.hadisd_y
        ]
        x_context_hadisd = [self.to_tensor(i).permute(1, 0) for i in x_context_hadisd]
        y_context_hadisd = [self.to_tensor(i) for i in y_context_hadisd]
        y_context_hadisd = self.norm_hadisd(y_context_hadisd)

        # ERA5
        era5 = self.to_tensor(self.load_era5_time(index))
        era5_target = era5.permute(2, 1, 0)
        era5_x = self.era5_x

        # AUxiliary variables
        aux_time = self.to_tensor(self.get_time_aux(date))
        climatology = self.climatology[date.hour // 6, date.dayofyear - 1, ...]

        task = {
            "x_context_hadisd_{}".format(prefix): x_context_hadisd,
            "y_context_hadisd_{}".format(prefix): y_context_hadisd,
            "climatology_{}".format(prefix): self.to_tensor(climatology),
            "sat_x_{}".format(prefix): sat_x,
            "sat_{}".format(prefix): sat_y,
            "icoads_x_{}".format(prefix): icoads_x,
            "icoads_{}".format(prefix): icoads_y,
            "igra_x_{}".format(prefix): igra_x,
            "igra_{}".format(prefix): igra_y,
            "amsua_{}".format(prefix): amsua_y,
            "amsua_x_{}".format(prefix): amsua_x,
            "amsub_{}".format(prefix): amsub_y,
            "amsub_x_{}".format(prefix): amsub_x,
            "iasi_{}".format(prefix): iasi_y,
            "iasi_x_{}".format(prefix): iasi_x,
            "ascat_{}".format(prefix): ascat_y,
            "ascat_x_{}".format(prefix): ascat_x,
            "hirs_{}".format(prefix): hirs_y,
            "hirs_x_{}".format(prefix): hirs_x,
            "y_target_{}".format(prefix): era5_target[
                ..., self.var_start : self.var_end
            ],
            "era5_x_{}".format(prefix): era5_x,
            "era5_elev_{}".format(prefix): self.era5_elev,
            "era5_lonlat_{}".format(prefix): self.era5_lonlat,
            "aux_time_{}".format(prefix): aux_time,
            "lt": torch.Tensor([self.var_start]),
        }

        return task


class HadISDDataset(Dataset):
    """
    HadISD dataset for decoder training
    """

    def __init__(self, var, mode, device, start_date, end_date):
        super().__init__()

        # Setup
        if not mode in ["train", "val", "test"]:
            raise Exception(f"mode is {mode}. Must be train, val, or test.")

        self.var = var
        self.mode = mode
        self.start_date = start_date
        self.device = device
        dates = pd.date_range(start_date, end_date, freq="6H")
        self.index = np.array(range(len(dates)))

        # Load the hadISD data
        self.load_hadisd()

    def load_hadisd(self):
        """
        Load the raw HadISD data
        """

        data_path = "path_to_data/"
        aux_data_path = "path_to_auxiliary_data/"
        var = self.var
        mode = self.mode

        vals = np.memmap(
            data_path + f"hadisd_processed/{var}_vals_{mode}.memmap",
            dtype="float32",
            mode="r",
            shape=get_hadisd_shape(mode),
        )

        lon = lon_to_0_360(
            np.load(data_path + f"hadisd_processed/{var}_lon_{mode}.npy")
        )
        lat = np.load(data_path + f"hadisd_processed/{var}_lat_{mode}.npy")
        self.hadisd_x = np.stack([lon, lat], axis=-1) / LATLON_SCALE_FACTOR
        self.hadisd_alt = np.load(
            data_path + f"hadisd_processed/{var}_alt_{mode}_final.npy"
        )
        self.hadisd_y = vals

        self.hadisd_index_offset = HADISD_OFFSETS[self.start_date]
        self.hadisd_means = self.to_tensor(
            np.load(aux_data_path + f"norm_factors/mean_hadisd_{var}.npy")
        )
        self.hadisd_stds = self.to_tensor(
            np.load(aux_data_path + f"norm_factors/std_hadisd_{var}.npy")
        )
        return

    def norm_hadisd(self, x):
        return (x - self.hadisd_means) / self.hadisd_stds

    def unnorm_pred(self, x):
        return self.hadisd_means + self.hadisd_stds * x

    def __len__(self):
        return self.index.shape[0] - 2

    def to_tensor(self, arr):
        return torch.from_numpy(np.array(arr)).float().to(self.device)

    def __getitem__(self, index):
        index = self.index[index]

        # Get longitude-latitude locations
        x_target = self.to_tensor(self.hadisd_x).permute(1, 0)

        # Get altitude and normalise
        m_alt = np.expand_dims(np.load("path_to_mean_alt.npy"), 1)
        s_alt = np.expand_dims(np.load("path_to_std_alt.npy"), 1)
        alt_target = self.to_tensor((self.hadisd_alt - m_alt) / s_alt)[:, :]

        # Get observations
        y_target = self.norm_hadisd(
            self.to_tensor(self.hadisd_y[index + self.hadisd_index_offset, :])
        )

        assert x_target.shape[0] == 2
        n_stations = x_target.shape[1]
        assert alt_target.shape[1] == n_stations
        assert y_target.shape[0] == n_stations

        return {"x": x_target, "altitude": alt_target, "y": y_target}


class AardvarkICDataset(Dataset):
    """
    Helper dataset to handle initial condition loading for decoder training
    """

    def __init__(self, device, start_date, end_date, lead_time=0):
        super().__init__()

        # Setup

        if lead_time == 0:
            # If leadtime is 0 load the output of the encoder
            if start_date == "2007-01-02" and end_date == "2017-12-31":
                ic_fname = "ic_train.mmap"
            elif start_date == "2019-01-01" and end_date == "2019-12-01":
                ic_fname = "ic_val.mmap"
            elif start_date == "2018-01-01" and end_date == "2018-12-31":
                ic_fname = "ic_test.mmap"
            else:
                print((start_date, end_date))
                raise Exception("Invalid start and end date")

            dates = pd.date_range(start_date, end_date, freq="6H")

            self.data = np.memmap(
                "path_to_encoder_predictions/" + ic_fname,
                dtype="float32",
                mode="r",
                shape=(len(dates), 121, 240, 24),  # shape of the output
            )
        else:
            # if leadtime >0 load the forecast generated from the encoder prediction
            if start_date == "2007-01-02" and end_date == "2017-12-31":
                ic_fname = f"ic_train_{lead_time}.mmap"

            elif start_date == "2019-01-01" and end_date == "2019-12-01":
                ic_fname = f"ic_val_{lead_time}.mmap"
            elif start_date == "2018-01-01" and end_date == "2018-12-31":
                ic_fname = f"ic_test_{lead_time}.mmap"
            else:
                print((start_date, end_date))
                raise Exception("Invalid start and end date.")

            dates = pd.date_range(start_date, end_date, freq="6H")[(lead_time) * 4 :]
            ic_shape = (len(dates), 121, 240, 24)

            self.data = np.memmap(
                self.data_path + "forecast_finetune/" + ic_fname,
                dtype="float32",
                mode="r",
                shape=ic_shape,
            )

        self.device = device

        # Normalisation
        aux_data_path = "path_to_auxiliary_data/"
        mean_factors_path = aux_data_path + f"norm_factors/mean_4u_1.npy"
        std_factors_path = aux_data_path + f"norm_factors/std_4u_1.npy"
        self.means = np.load(mean_factors_path)[:, np.newaxis, np.newaxis, ...]
        self.stds = np.load(std_factors_path)[:, np.newaxis, np.newaxis, ...]

    def __getitem__(self, index):
        # Load Aardvark prediction and normalise
        data_raw = np.transpose(np.copy(self.data[index, :, :, :]), (2, 1, 0))
        data = (data_raw - self.means) / self.stds
        return torch.from_numpy(data).to(self.device)


class WeatherDatasetDownscaling(Dataset):
    """
    Main decoder training dataset. Uses AardvarkICDataset and HadISDDataset to
    handle processor output and station data
    """

    def __init__(
        self,
        device,
        hadisd_mode,
        start_date,
        end_date,
        context_mode,
        era5_mode="sfc",
        res=1,
        hadisd_var="tas",
        lead_time=1,
    ):
        # The context mode determines whether we make use of ERA5 or our own ICs.
        if not context_mode in ["era5", "aardvark"]:
            raise Exception(
                f"context_mode must be era5 or aardvark, got {context_mode}"
            )

        super().__init__()

        # Setup
        self.lead_time = lead_time

        self.device = device
        self.data_path = "path_to_data/"
        self.aux_data_path = "path_to_auxiliary_data/"
        self.start_date = start_date
        self.end_date = end_date
        self.era5_mode = era5_mode
        self.res = res
        self.context_mode = context_mode

        self.dates = pd.date_range(start_date, end_date, freq="6H")
        self.index = np.array(range(len(self.dates)))

        # Load ERA5 data for pre-training
        self.era5_sfc = [
            self.load_era5(year)
            for year in range(int(start_date[:4]), int(end_date[:4]) + 1)
        ]

        raw_era5_lon = np.load(self.data_path + f"era5/era5_x_{res}.npy")
        raw_era5_lat = np.load(self.data_path + f"era5/era5_y_{res}.npy")
        self.era5_x = [
            self.to_tensor(raw_era5_lon) / LATLON_SCALE_FACTOR,
            self.to_tensor(raw_era5_lat) / LATLON_SCALE_FACTOR,
        ]

        # Load orography
        elev_path = self.data_path + f"era5/elev_vars_{res}.npy"
        self.era5_elev = self.to_tensor(np.load(elev_path)).permute(0, 2, 1)

        # Normalisation
        mean_factors_path = (
            self.aux_data_path + f"norm_factors/mean_{era5_mode}_{res}.npy"
        )
        std_factors_path = (
            self.aux_data_path + f"norm_factors/std_{era5_mode}_{res}.npy"
        )
        self.means = np.load(mean_factors_path)[:, np.newaxis, np.newaxis, ...]
        self.stds = np.load(std_factors_path)[:, np.newaxis, np.newaxis, ...]

        # HadISD data
        self.hadisd_data = HadISDDataset(
            var=hadisd_var,
            mode=hadisd_mode,
            device=device,
            start_date=start_date,
            end_date=end_date,
        )

        if context_mode == "aardvark":
            # Load the Aardvark encoder predictions
            self.aardvark_data = AardvarkICDataset(
                device, start_date, end_date, lead_time
            )

    def load_era5(self, year):
        """
        Load the raw ERA5 data
        """

        if year % 4 == 0:
            d = 366 * 4
        else:
            d = 365 * 4

        if self.era5_mode == "sfc":
            levels = 4
        elif self.era5_mode == "13u":
            levels = 69
        else:
            levels = 24

        if self.res == 1:
            x = 240
            y = 121
        elif self.res == 5:
            x = 64
            y = 32
        mmap = np.memmap(
            self.data_path
            + "era5/era5_{}_{}_6_{}.memmap".format(self.era5_mode, self.res, year),
            dtype="float32",
            mode="r",
            shape=(d, levels, x, y),
        )
        return mmap

    def norm_era5(self, x):
        x = (x - self.means) / self.stds
        return x

    def unnorm_era5(self, x):
        x = x * self.stds + self.means
        return x

    def unnorm_pred(self, x):
        return self.hadisd_data.unnorm_pred(x)

    def norm_data(self, x, means, stds):
        return (x - means) / stds

    def __len__(self):
        return self.index.shape[0] - (self.lead_time) * 4

    def to_tensor(self, arr):
        return torch.from_numpy(np.array(arr)).float().to(self.device)

    def get_time_aux(self, current_date):
        """
        Get auxiliary time variables for a given date
        """

        doy = current_date.dayofyear
        year = (current_date.year - 2007) / 15
        time_of_day = current_date.hour
        return np.array(
            [
                np.cos(np.pi * 2 * doy / DAYS_IN_YEAR),
                np.sin(np.pi * 2 * doy / DAYS_IN_YEAR),
                np.cos(np.pi * 2 * time_of_day / 24),
                np.sin(np.pi * 2 * time_of_day / 24),
                year,
            ]
        )

    def load_era5_time(self, index):
        """
        Load ERA5 training data
        """

        date = self.dates[index]
        year = date.year
        hour = date.hour
        doy = (date.dayofyear - 1) * 4 + (hour // 6)

        era5 = self.era5_sfc[year - int(self.start_date[:4])][doy, ...]
        era5 = np.copy(era5)
        era5 = self.norm_era5(era5[np.newaxis, ...])[0, ...]
        return era5

    def load_year_end(self, year, doy):
        data_1 = self.era5_sfc[year - int(self.start_date[:4])][doy : doy + 1, ...]
        missing = self.lead_time - data_1.shape[0] + 1
        data_2 = self.era5_sfc[year - int(self.start_date[:4]) + 1][:missing, ...]
        data = np.concatenate([data_1, data_2])
        return data

    def __getitem__(self, index):

        index = self.index[index]
        date = self.dates[index + 4 * self.lead_time]

        # Get HadISD data
        hadisd_slice = self.hadisd_data[index + 4 * self.lead_time]

        # Get lon-lat
        x_context = self.era5_x
        n_lon = x_context[0].shape[0]
        n_lat = x_context[1].shape[0]

        # Get auxiliary time
        aux_time = torch.reshape(self.to_tensor(self.get_time_aux(date)), (-1, 1, 1))

        # Load the context (either aardvark or ERA5 for use in pre-training)
        if self.context_mode == "era5":
            y_context_obs = self.to_tensor(
                self.load_era5_time(index + 4 * self.lead_time)
            )

        elif self.context_mode == "aardvark":
            y_context_obs = self.aardvark_data[index]

        else:
            raise Exception

        y_context = torch.cat(
            [
                y_context_obs,
                self.era5_elev.permute(0, 2, 1),
                aux_time.repeat(1, n_lon, n_lat),
            ]
        )

        assert y_context.shape[1] == n_lon
        assert y_context.shape[2] == n_lat

        x = hadisd_slice["x"]
        alt = hadisd_slice["altitude"]
        y = hadisd_slice["y"]

        return {
            "x_target": x,
            "alt_target": alt,
            "y_target": y,
            "y_context": y_context,
            "x_context": x_context,
            "aux_time": aux_time,
            "lt": torch.Tensor([0]),
        }


class ForecasterDatasetDownscaling(Dataset):
    """
    Dataset to generate decoder predictions from pre-saved Aardvark forecasts
    """

    def __init__(
        self,
        start_date,
        end_date,
        lead_time,
        hadisd_var,
        mode,
        device,
        forecast_path,
        region="global",
    ):
        super().__init__()

        # Setup

        if not mode in ["train", "val", "test"]:
            raise Exception(f"Mode is {mode}. Must be either train, val, or test")

        self.device = device
        self.start_date = start_date
        self.end_date = end_date
        self.lead_time = lead_time
        self.mode = mode
        self.offset = np.timedelta64(lead_time, "D").astype("timedelta64[ns]")

        self.dates = pd.date_range(start_date, end_date, freq="6H")[:-30]

        # Normalisation
        aux_data_path = "auxiliary_data_path/"
        self.means = np.load(aux_data_path + "norm_factors/mean_4u_1.npy")
        self.stds = np.load(aux_data_path + "norm_factors/std_4u_1.npy")

        # Load auxiliary data
        self.load_npy_file()
        data_path = "data_path/"
        res = "1"
        raw_era5_lon = np.load(data_path + f"era5/era5_x_{res}.npy")
        raw_era5_lat = np.load(data_path + f"era5/era5_y_{res}.npy")
        self.era5_x = [
            self.to_tensor(raw_era5_lon) / LATLON_SCALE_FACTOR,
            self.to_tensor(raw_era5_lat) / LATLON_SCALE_FACTOR,
        ]
        elev_path = data_path + f"era5/elev_vars_{res}.npy"
        self.era5_elev = self.to_tensor(np.load(elev_path)).permute(0, 2, 1)

        # Load hadISD
        self.hadisd_data = HadISDDataset(
            var=hadisd_var,
            mode="train",
            device=device,
            start_date=start_date,
            end_date=end_date,
        )

        # Subset to region
        self.region = region
        if self.region != "global":
            self.mask = np.load(
                self.data_path + f"hadisd_processed/tas_mask_train_{region}.npy"
            )

    def date_range(self):
        return np.arange(
            start=np.datetime64(self.start_date).astype("datetime64[ns]"),
            stop=np.datetime64(self.end_date).astype("datetime64[ns]"),
            step=np.timedelta64(1, "D").astype("timedelta64[ns]"),
        )

    def load_npy_file(self):
        """
        Load the pre-saved Aardvark forecasts
        """

        dates = pd.date_range(self.start_date, self.end_date, freq="6H")

        if self.mode == "train":
            dates = dates[:-40]  # Need 10 day offset at end of year

        self.Y_context = np.memmap(
            "path_to_forecasts/forecast_{}.mmap".format(self.mode),
            dtype="float32",
            mode="r",
            shape=(len(dates), 121, 240, 24, 11),
        )

        return

    def norm_era5(self, x):
        return (x - self.means) / self.stds

    def norm_hadisd(self, x):
        return self.hadisd_data.norm_hadisd(x)

    def unnorm_pred(self, x):
        return self.hadisd_data.unnorm_pred(x)

    def __len__(self):
        return len(self.dates) - 40  # Need 10 day offset at end of year

    def to_tensor(self, arr):
        return torch.from_numpy(np.array(arr)).float().to(self.device)

    def get_time_aux(self, index):
        """
        Get the auxiliary time variables
        """

        current_date = (self.dates + self.offset)[index]
        doy = current_date.dayofyear
        year = (current_date.year - 2007) / 15
        time_of_day = current_date.hour
        return np.array(
            [
                np.cos(np.pi * 2 * doy / DAYS_IN_YEAR),
                np.sin(np.pi * 2 * doy / DAYS_IN_YEAR),
                np.cos(np.pi * 2 * time_of_day / 24),
                np.sin(np.pi * 2 * time_of_day / 24),
                year,
            ]
        )

    def __getitem__(self, index):

        # Load target data
        hadisd_slice = self.hadisd_data[index + 4 * self.lead_time]

        x_context = self.era5_x
        n_lon = x_context[0].shape[0]
        n_lat = x_context[1].shape[0]

        # Load auxiliary time
        aux_time = torch.reshape(self.to_tensor(self.get_time_aux(index)), (-1, 1, 1))

        # Load input
        y_context = self.norm_era5(self.Y_context[index, ..., self.lead_time])
        y_context = torch.cat(
            [
                self.to_tensor(y_context).permute(2, 1, 0),
                self.era5_elev.permute(0, 2, 1),
                aux_time.repeat(1, n_lon, n_lat),
            ]
        )

        assert y_context.shape[1] == n_lon
        assert y_context.shape[2] == n_lat

        # Handle region masking
        if self.region != "global":
            hadisd_slice["y"][self.mask] = np.nan

        return {
            "x_target": hadisd_slice["x"],
            "alt_target": hadisd_slice["altitude"],
            "y_target": hadisd_slice["y"],
            "y_context": y_context,
            "x_context": x_context,
            "aux_time": aux_time,
            "lt": torch.Tensor([0]),
        }


class ForecastLoader(Dataset):
    """
    Loader for finetuning the processor module
    """

    def __init__(
        self,
        device,
        mode,
        lead_time,
        era5_mode="sfc",
        res=5,
        frequency=24,
        norm=True,
        diff=False,
        rollout=False,
        random_lt=False,
        u_only=False,
        ic_path=None,
        finetune_step=None,
        finetune_eval_every=100,
        eval_steps=False,
    ):

        super().__init__()

        # Setup
        self.device = device
        self.mode = mode
        self.data_path = "data_path/"

        self.lead_time = lead_time
        self.era5_mode = era5_mode
        self.res = res
        self.frequency = frequency
        self.norm = norm
        self.diff = diff
        self.rollout = rollout
        self.random_lt = random_lt
        self.u_only = u_only
        self.ic_path = ic_path

        self.finetune_step = finetune_step
        self.finetune_eval_every = finetune_eval_every
        self.eval_steps = eval_steps

        if self.frequency == 6:
            self.lead_time = self.lead_time * 4
            freq = "6H"

        else:
            freq = "1D"

        if self.mode == "train":
            self.dates = pd.date_range("1979-01-01", "2017-12-31", freq=freq)
        elif self.mode == "tune":
            self.dates = pd.date_range("2018-01-01", "2018-12-31", freq=freq)
        elif self.mode == "test":
            self.dates = pd.date_range("2018-01-01", "2018-12-31", freq=freq)
        elif self.mode == "val":
            self.dates = pd.date_range("2019-01-01", "2019-12-31", freq=freq)

        # Load the predictions from the previous leadtime to be the new context set
        if self.finetune_step is not None:

            if self.mode == "train":
                self.dates = pd.date_range("2007-01-02", "2017-12-31", freq=freq)
                ic_shape = (
                    len(self.dates) - max(0, (self.finetune_step - 1) * 4),
                    121,
                    240,
                    24,
                )
            elif self.mode == "val":
                self.dates = pd.date_range("2019-01-01", "2019-12-31", freq=freq)
                ic_shape = (
                    len(self.dates) - max(0, (self.finetune_step - 1) * 4),
                    121,
                    240,
                    24,
                )
            elif self.mode == "test":
                self.dates = pd.date_range("2018-01-01", "2018-12-31", freq=freq)
                ic_shape = (
                    len(self.dates) - max(0, (self.finetune_step - 1) * 4),
                    121,
                    240,
                    24,
                )

            if self.finetune_step > 1:
                print(ic_shape)
                self.ic = np.memmap(
                    self.ic_path
                    + "ic_{}_{}.mmap".format(self.mode, self.finetune_step - 1),
                    dtype="float32",
                    mode="r",
                    shape=ic_shape,
                )
            elif self.ic_path is not None:

                self.ic = np.memmap(
                    self.ic_path + "ic_{}.mmap".format(self.mode),
                    dtype="float32",
                    mode="r",
                    shape=ic_shape,
                )

        elif self.ic_path is not None:
            if self.mode == "train":
                self.dates = pd.date_range("2007-01-02", "2017-12-31", freq=freq)
            ic_shape = (len(self.dates), 121, 240, 24)

            self.ic = np.memmap(
                self.ic_path + "/ic_{}.mmap".format(self.mode),
                dtype="float32",
                mode="r",
                shape=ic_shape,
            )

        # Orography
        self.era5_elev = np.float32(
            np.load(self.data_path + "era5/elev_vars_{}.npy".format(res))
        )
        elev_mean = self.era5_elev.mean(axis=(1, 2))[:, np.newaxis, np.newaxis]
        elev_std = self.era5_elev.std(axis=(1, 2))[:, np.newaxis, np.newaxis]
        self.era5_elev = (self.era5_elev - elev_mean) / elev_std

        # ERA5 ground truth data for training
        self.era5_sfc = [
            self.load_era5(year)
            for year in range(int(self.dates[0].year), int(self.dates[-1].year) + 1)
        ]

        # Noramalisation factors
        self.means = (
            self.to_tensor(
                np.load(
                    self.data_path
                    + "norm_factors/mean_{}_{}.npy".format(self.era5_mode, self.res)
                )
            )
            .unsqueeze(1)
            .unsqueeze(1)
        )
        self.stds = (
            self.to_tensor(
                np.load(
                    self.data_path
                    + "norm_factors/std_{}_{}.npy".format(self.era5_mode, self.res)
                )
            )
            .unsqueeze(1)
            .unsqueeze(1)
        )
        self.diff_means = (
            self.to_tensor(
                np.load(
                    self.data_path
                    + "norm_factors/mean_diff_{}_{}.npy".format(
                        self.era5_mode, self.res
                    )
                )
            )
            .unsqueeze(0)
            .unsqueeze(0)
        )
        self.diff_stds = (
            self.to_tensor(
                np.load(
                    self.data_path
                    + "norm_factors/std_diff_{}_{}.npy".format(self.era5_mode, self.res)
                )
            )
            .unsqueeze(0)
            .unsqueeze(0)
        )

        self.diff_means_1 = (
            self.to_tensor(
                np.load(
                    self.data_path
                    + "norm_factors/mean_diff_{}_{}_6h.npy".format(
                        self.era5_mode, self.res
                    )
                )
            )
            .unsqueeze(0)
            .unsqueeze(0)
        )
        self.diff_stds_1 = (
            self.to_tensor(
                np.load(
                    self.data_path
                    + "norm_factors/std_diff_{}_{}_6h.npy".format(
                        self.era5_mode, self.res
                    )
                )
            )
            .unsqueeze(0)
            .unsqueeze(0)
        )

        self.diff_means_2 = (
            self.to_tensor(
                np.load(
                    self.data_path
                    + "norm_factors/mean_diff_{}_{}_12h.npy".format(
                        self.era5_mode, self.res
                    )
                )
            )
            .unsqueeze(0)
            .unsqueeze(0)
        )
        self.diff_stds_2 = (
            self.to_tensor(
                np.load(
                    self.data_path
                    + "norm_factors/std_diff_{}_{}_12h.npy".format(
                        self.era5_mode, self.res
                    )
                )
            )
            .unsqueeze(0)
            .unsqueeze(0)
        )

        self.means_dict = {
            0: self.diff_means,
            2: self.diff_means_2,
            3: self.diff_means_1,
        }

        self.stds_dict = {0: self.diff_stds, 2: self.diff_stds_2, 3: self.diff_stds_1}

    def __len__(self):
        if np.logical_and(self.eval_steps, self.mode == "train"):
            return self.finetune_eval_every * 12 * 4

        return self.dates.shape[0] - self.lead_time

    def to_tensor(self, arr):

        return torch.from_numpy(arr).float().to(self.device)

    def norm_era5(self, x):
        x = (x - self.means) / self.stds
        return x

    def norm_era5_tendency(self, x, lt_offset):

        x = (x - self.means_dict[lt_offset]) / self.stds_dict[lt_offset]
        return x

    def unnorm_pred(self, x):
        x = x * self.diff_stds.unsqueeze(0) + self.diff_means.unsqueeze(0)
        return x

    def unnorm_base_context(self, x):
        x = x * self.stds.unsqueeze(0) + self.means.unsqueeze(0)
        return x

    def load_era5(self, year):
        """
        Load ERA5 data for training
        """

        if year % 4 == 0:
            d = 366
        else:
            d = 365

        if self.frequency == 6:
            d = d * 4

        if self.era5_mode == "sfc":
            levels = 4
        elif self.era5_mode == "13u":
            levels = 69
        else:
            levels = 24

        if self.res == 1:
            x = 240
            y = 121
        elif self.res == 5:
            x = 64
            y = 32

        mmap = np.memmap(
            self.data_path
            + "era5/era5_{}_{}_{}_{}.memmap".format(
                self.era5_mode, self.res, self.frequency, year
            ),
            dtype="float32",
            mode="r",
            shape=(d, levels, x, y),
        )
        return mmap

    def load_era5_time(self, index):
        """
        Load ERA5 data for training
        """
        date = self.dates[index]
        year = date.year
        doy = date.dayofyear - 1
        hour = date.hour
        if self.frequency == 6:
            era5 = self.era5_sfc[year - int(self.dates[0].year)][
                doy * 4 + hour // 6, ...
            ]
        else:
            era5 = self.era5_sfc[year - int(self.dates[0].year)][doy, ...]

        return np.copy(era5)

    def make_time_channels(self, index, x, y):
        """
        Make auxiliary time channels
        """

        date = self.dates[index]
        hour = date.hour
        doy = date.dayofyear - 1
        if date.year % 4 == 0:
            n_days = 366
        else:
            n_days = 365
        hour_sin = np.sin(hour * np.pi / 12) * np.float32(np.ones((1, x, y)))
        hour_cos = np.cos(hour * np.pi / 12) * np.float32(np.ones((1, x, y)))
        doy_sin = np.sin(doy * 2 * np.pi / n_days) * np.float32(np.ones((1, x, y)))
        doy_cos = np.cos(doy * 2 * np.pi / n_days) * np.float32(np.ones((1, x, y)))

        return np.concatenate([hour_sin, hour_cos, doy_sin, doy_cos])

    def __getitem__(self, index):

        # Option to offset to random leadtime
        lt_offset = 0
        if self.random_lt:
            lt_offset = np.random.choice([0, 2, 3])

        # Load ground truth data
        y_target = self.to_tensor(
            self.load_era5_time(index + self.lead_time - lt_offset)
        )

        # Load either initial condition or ERA5 depending on task
        if self.ic_path is not None:
            era5_ts0 = self.ic[index].copy().transpose(2, 1, 0)

        else:
            era5_ts0 = self.load_era5_time(index)

        # Auxiliary time
        time = self.make_time_channels(index, era5_ts0.shape[1], era5_ts0.shape[2])
        era5_ts0 = self.to_tensor(
            np.concatenate([era5_ts0, self.era5_elev, time], axis=0)
        )
        y_context = era5_ts0.permute(0, 2, 1)[:, ...]

        # Normalisation
        if self.diff:
            y_target = (y_target - era5_ts0[:24, ...]).permute(2, 1, 0)
            y_target = self.norm_era5_tendency(y_target, lt_offset)
            y_context[:24, ...] = self.norm_era5(y_context[:24, ...])

        else:
            if self.norm:
                y_context[:24, ...] = self.norm_era5(y_context[:24, ...], lt_offset)
                y_target = self.norm_era5(y_target, lt_offset)
            y_target = y_target.permute(2, 1, 0)

        if self.rollout:
            # Option to return entire timeseries of target data
            targets = []
            for t in range(self.lead_time + 1):
                t = self.to_tensor(self.load_era5_time(index + t))
                targets.append(t.permute(2, 1, 0))
            targets = torch.stack(targets, dim=-1)[..., ::4]

            return {
                "y_context": y_context.permute(0, 2, 1),
                "y_target": y_target,
                "targets": targets,
                "lt": self.to_tensor(np.array([lt_offset])),
            }

        else:
            return {
                "y_context": y_context.permute(0, 2, 1),
                "y_target": y_target[..., :],
                "lt": self.to_tensor(np.array([lt_offset])),
                "target_index": self.to_tensor(np.array([index])),
            }


class WeatherDatasetE2E(WeatherDataset):
    """
    Dataset for running Aardvark end-to-end
    """

    def __init__(
        self,
        device,
        hadisd_mode,
        start_date,
        end_date,
        lead_time,
        mode,
        hadisd_var,
        max_steps_per_epoch=None,
        era5_mode="sfc",
        res=1,
        filter_dates=None,
        var_start=0,
        var_end=24,
        diff=False,
        two_frames=False,
        region="global",
    ):

        super().__init__(
            device,
            hadisd_mode,
            start_date,
            end_date,
            lead_time,
            era5_mode,
            res=res,
            filter_dates=filter_dates,
            diff=diff,
        )

        # Setup
        self.var_start = var_start
        self.var_end = var_end
        self.diff = diff
        self.two_frames = two_frames
        self.region = region
        self.lead_time = lead_time
        self.mode = mode
        self.max_steps_per_epoch = max_steps_per_epoch

        # Initialise encoder dataset
        self.assimilation_dataset = WeatherDatasetAssimilation(
            device="cuda",
            hadisd_mode="train",
            start_date=start_date,
            end_date=end_date,
            lead_time=0,
            era5_mode="4u",
            res=1,
            var_start=0,
            var_end=24,
            diff=False,
            two_frames=False,
        )

        # Initialise forecast dataset
        self.forecast_dataset = ForecastLoader(
            device="cuda",
            mode=mode,
            lead_time=lead_time,
            era5_mode=era5_mode,
            res=1,
            frequency=6,
            diff=True,
            u_only=False,
            random_lt=False,
        )

        # Initialise downscaling dataset
        self.downscaling_dataset = ForecasterDatasetDownscaling(
            start_date=start_date,
            end_date=end_date,
            lead_time=lead_time,
            hadisd_var=hadisd_var,
            mode=mode,
            device=device,
            forecast_path=None,
            region=region,
        )

    def __len__(self):
        if self.max_steps_per_epoch:
            return self.max_steps_per_epoch
        return len(self.downscaling_dataset) - 40  # Need 10 day offset at end of year

    def __getitem__(self, index):

        if self.max_steps_per_epoch:
            index = np.random.choice(
                np.arange(len(self.downscaling_dataset) - 40)
            )  # Need 10 day offset at end of year

        # Get data for each of the three datasets
        assimilation = self.assimilation_dataset.__getitem__(index)
        forecast = self.forecast_dataset.__getitem__(index)
        downscaling = self.downscaling_dataset.__getitem__(index)

        # Create task
        task = {
            "assimilation": assimilation,
            "forecast": forecast,
            "downscaling": downscaling,
            "index": torch.tensor(index),
        }

        # Add y target to allow for end to end finetuning if needed
        task["y_target"] = task["downscaling"]["y_target"]

        return task

    def unnorm_pred(self, x):

        dev = x.device
        x = x.detach().cpu().numpy()

        x = (
            x
            * self.stds[np.newaxis, ...].transpose(0, 2, 3, 1)[
                ..., self.var_start : self.var_end
            ]
            + self.means[np.newaxis, ...].transpose(0, 2, 3, 1)[
                ..., self.var_start : self.var_end
            ]
        )
        if bool(self.diff):
            x = (
                x
                + self.era5_mean_spatial[np.newaxis, ...].transpose(0, 3, 2, 1)[
                    ..., self.var_start : self.var_end
                ]
            )
        return torch.from_numpy(x).float().to(dev)