era5_data_loader.py

import numpy as np
import xarray as xr

# Load ERA5 geopotential levels
era5_ds1 = xr.open_dataset("./datasets/GEOP1000_GAN_2017.nc")
era5_ds2 = xr.open_dataset("./datasets/GEOP800_GAN_2017.nc")
era5_ds3 = xr.open_dataset("./datasets/GEOP500_GAN_2017.nc")
era5_times = era5_ds1.time[:].data

# Load ERA5 total precipitation
prec_ds = xr.open_dataset("./datasets/TP_GAN_2017.nc")
prec_times = prec_ds.time[:].data

# Find common dates and shuffle
times = np.intersect1d(era5_times, prec_times)
np.random.shuffle(times)

# Create geopotential normalised stack
z500 = era5_ds3.Geopotential.sel(time=times[::10])[:].data
z500 = (z500 - z500.min()) / (z500.max() - z500.min())
z800 = era5_ds2.Geopotential.sel(time=times[::10])[:].data
z800 = (z800 - z800.min()) / (z800.max() - z800.min())
z1000 = era5_ds1.Geopotential.sel(time=times[::10])[:].data
z1000 = (z1000 - z1000.min()) / (z1000.max() - z1000.min())
z = np.stack((z1000, z800, z500), axis=3)

# Create precipitation normalised stack
tp = prec_ds.tp.sel(time=times[::10])[:].data * 1000
tp = np.clip(tp, 0, 30)
tp = np.log(1+tp)/np.log(31)
tp = tp[:,:,:,None]

geop_train = z[:600,:,:,:]
prec_train = tp[:600,:,:,:]

geop_val = z[600:750,:,:,:]
prec_val = tp[600:,:,:,:]

print(prec_train.shape, prec_val.shape)
print(geop_train.shape, geop_val.shape)