| 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__() |
|
|
| |
| 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 |
|
|
| |
| 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))) |
|
|
| |
| 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() |
|
|
| |
| print("Loading ERA5") |
| self.era5_sfc = [ |
| self.load_era5(year) |
| for year in range(int(start_date[:4]), int(end_date[:4]) + 1) |
| ] |
|
|
| |
| 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, |
| ] |
|
|
| |
| 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]) |
|
|
| |
| self.climatology = np.memmap( |
| self.data_path + "climatology_data.mmap", |
| dtype="float32", |
| mode="r", |
| shape=CLIMATOLOGY_SHAPE, |
| ) |
|
|
| |
| 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, |
| ) |
|
|
| |
|
|
| 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: |
| |
| 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: |
| |
| 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_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) |
|
|
| |
| 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) |
|
|
| |
| 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) |
|
|
| |
| 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_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_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_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_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) |
|
|
| |
| 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 = self.to_tensor(self.load_era5_time(index)) |
| era5_target = era5.permute(2, 1, 0) |
| era5_x = self.era5_x |
|
|
| |
| 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__() |
|
|
| |
| 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))) |
|
|
| |
| 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] |
|
|
| |
| x_target = self.to_tensor(self.hadisd_x).permute(1, 0) |
|
|
| |
| 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)[:, :] |
|
|
| |
| 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__() |
|
|
| |
|
|
| if lead_time == 0: |
| |
| 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), |
| ) |
| else: |
| |
| 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 |
|
|
| |
| 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): |
| |
| 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, |
| ): |
| |
| if not context_mode in ["era5", "aardvark"]: |
| raise Exception( |
| f"context_mode must be era5 or aardvark, got {context_mode}" |
| ) |
|
|
| super().__init__() |
|
|
| |
| 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))) |
|
|
| |
| 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, |
| ] |
|
|
| |
| 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) |
|
|
| |
| 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, ...] |
|
|
| |
| self.hadisd_data = HadISDDataset( |
| var=hadisd_var, |
| mode=hadisd_mode, |
| device=device, |
| start_date=start_date, |
| end_date=end_date, |
| ) |
|
|
| if context_mode == "aardvark": |
| |
| 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] |
|
|
| |
| 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] |
|
|
| |
| aux_time = torch.reshape(self.to_tensor(self.get_time_aux(date)), (-1, 1, 1)) |
|
|
| |
| 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__() |
|
|
| |
|
|
| 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] |
|
|
| |
| 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") |
|
|
| |
| 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) |
|
|
| |
| self.hadisd_data = HadISDDataset( |
| var=hadisd_var, |
| mode="train", |
| device=device, |
| start_date=start_date, |
| end_date=end_date, |
| ) |
|
|
| |
| 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] |
|
|
| 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 |
|
|
| 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): |
|
|
| |
| 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] |
|
|
| |
| aux_time = torch.reshape(self.to_tensor(self.get_time_aux(index)), (-1, 1, 1)) |
|
|
| |
| 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 |
|
|
| |
| 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__() |
|
|
| |
| 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) |
|
|
| |
| 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, |
| ) |
|
|
| |
| 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 |
|
|
| |
| self.era5_sfc = [ |
| self.load_era5(year) |
| for year in range(int(self.dates[0].year), int(self.dates[-1].year) + 1) |
| ] |
|
|
| |
| 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): |
|
|
| |
| lt_offset = 0 |
| if self.random_lt: |
| lt_offset = np.random.choice([0, 2, 3]) |
|
|
| |
| y_target = self.to_tensor( |
| self.load_era5_time(index + self.lead_time - lt_offset) |
| ) |
|
|
| |
| 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) |
|
|
| |
| 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)[:, ...] |
|
|
| |
| 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: |
| |
| 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, |
| ) |
|
|
| |
| 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 |
|
|
| |
| 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, |
| ) |
|
|
| |
| 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, |
| ) |
|
|
| |
| 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 |
|
|
| def __getitem__(self, index): |
|
|
| if self.max_steps_per_epoch: |
| index = np.random.choice( |
| np.arange(len(self.downscaling_dataset) - 40) |
| ) |
|
|
| |
| assimilation = self.assimilation_dataset.__getitem__(index) |
| forecast = self.forecast_dataset.__getitem__(index) |
| downscaling = self.downscaling_dataset.__getitem__(index) |
|
|
| |
| task = { |
| "assimilation": assimilation, |
| "forecast": forecast, |
| "downscaling": downscaling, |
| "index": torch.tensor(index), |
| } |
|
|
| |
| 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) |
|
|
