import xarray as xr import numpy as np import pandas as pd import matplotlib.pyplot as plt import pickle import glob, os import re import tensorflow as tf import netCDF4 import copy import string import h5py from tqdm import tqdm class data_utils: def __init__(self, grid_info, input_mean, input_max, input_min, output_scale): self.data_path = None self.input_vars = [] self.target_vars = [] self.input_feature_len = None self.target_feature_len = None self.grid_info = grid_info self.level_name = 'lev' self.sample_name = 'sample' self.latlonnum = len(self.grid_info['ncol']) # number of unique lat/lon grid points # make area-weights self.grid_info['area_wgt'] = self.grid_info['area']/self.grid_info['area'].mean(dim = 'ncol') self.area_wgt = self.grid_info['area_wgt'].values # map ncol to nsamples dimension # to_xarray = {'area_wgt':(self.sample_name,np.tile(self.grid_info['area_wgt'], int(n_samples/len(self.grid_info['ncol']))))} # to_xarray = xr.Dataset(to_xarray) self.input_mean = input_mean self.input_max = input_max self.input_min = input_min self.output_scale = output_scale self.lats, self.lats_indices = np.unique(self.grid_info['lat'].values, return_index=True) self.lons, self.lons_indices = np.unique(self.grid_info['lon'].values, return_index=True) self.sort_lat_key = np.argsort(self.grid_info['lat'].values[np.sort(self.lats_indices)]) self.sort_lon_key = np.argsort(self.grid_info['lon'].values[np.sort(self.lons_indices)]) self.indextolatlon = {i: (self.grid_info['lat'].values[i%self.latlonnum], self.grid_info['lon'].values[i%self.latlonnum]) for i in range(self.latlonnum)} def find_keys(dictionary, value): keys = [] for key, val in dictionary.items(): if val[0] == value: keys.append(key) return keys indices_list = [] for lat in self.lats: indices = find_keys(self.indextolatlon, lat) indices_list.append(indices) indices_list.sort(key = lambda x: x[0]) self.lat_indices_list = indices_list self.hyam = self.grid_info['hyam'].values self.hybm = self.grid_info['hybm'].values self.p0 = 1e5 # code assumes this will always be a scalar self.pressure_grid_train = None self.pressure_grid_val = None self.pressure_grid_scoring = None self.pressure_grid_test = None self.dp_train = None self.dp_val = None self.dp_scoring = None self.dp_test = None self.train_regexps = None self.train_stride_sample = None self.train_filelist = None self.val_regexps = None self.val_stride_sample = None self.val_filelist = None self.scoring_regexps = None self.scoring_stride_sample = None self.scoring_filelist = None self.test_regexps = None self.test_stride_sample = None self.test_filelist = None # physical constants from E3SM_ROOT/share/util/shr_const_mod.F90 self.grav = 9.80616 # acceleration of gravity ~ m/s^2 self.cp = 1.00464e3 # specific heat of dry air ~ J/kg/K self.lv = 2.501e6 # latent heat of evaporation ~ J/kg self.lf = 3.337e5 # latent heat of fusion ~ J/kg self.lsub = self.lv + self.lf # latent heat of sublimation ~ J/kg self.rho_air = 101325/(6.02214e26*1.38065e-23/28.966)/273.15 # density of dry air at STP ~ kg/m^3 # ~ 1.2923182846924677 # SHR_CONST_PSTD/(SHR_CONST_RDAIR*SHR_CONST_TKFRZ) # SHR_CONST_RDAIR = SHR_CONST_RGAS/SHR_CONST_MWDAIR # SHR_CONST_RGAS = SHR_CONST_AVOGAD*SHR_CONST_BOLTZ self.rho_h20 = 1.e3 # density of fresh water ~ kg/m^ 3 self.v1_inputs = ['state_t', 'state_q0001', 'state_ps', 'pbuf_SOLIN', 'pbuf_LHFLX', 'pbuf_SHFLX'] self.v1_outputs = ['ptend_t', 'ptend_q0001', 'cam_out_NETSW', 'cam_out_FLWDS', 'cam_out_PRECSC', 'cam_out_PRECC', 'cam_out_SOLS', 'cam_out_SOLL', 'cam_out_SOLSD', 'cam_out_SOLLD'] self.var_lens = {#inputs 'state_t':60, 'state_q0001':60, 'state_ps':1, 'pbuf_SOLIN':1, 'pbuf_LHFLX':1, 'pbuf_SHFLX':1, #outputs 'ptend_t':60, 'ptend_q0001':60, 'cam_out_NETSW':1, 'cam_out_FLWDS':1, 'cam_out_PRECSC':1, 'cam_out_PRECC':1, 'cam_out_SOLS':1, 'cam_out_SOLL':1, 'cam_out_SOLSD':1, 'cam_out_SOLLD':1 } self.var_short_names = {'ptend_t':'$dT/dt$', 'ptend_q0001':'$dq/dt$', 'cam_out_NETSW':'NETSW', 'cam_out_FLWDS':'FLWDS', 'cam_out_PRECSC':'PRECSC', 'cam_out_PRECC':'PRECC', 'cam_out_SOLS':'SOLS', 'cam_out_SOLL':'SOLL', 'cam_out_SOLSD':'SOLSD', 'cam_out_SOLLD':'SOLLD'} self.target_energy_conv = {'ptend_t':self.cp, 'ptend_q0001':self.lv, 'cam_out_NETSW':1., 'cam_out_FLWDS':1., 'cam_out_PRECSC':self.lv*self.rho_h20, 'cam_out_PRECC':self.lv*self.rho_h20, 'cam_out_SOLS':1., 'cam_out_SOLL':1., 'cam_out_SOLSD':1., 'cam_out_SOLLD':1. } # for metrics self.input_train = None self.target_train = None self.preds_train = None self.samples_train = None self.target_weighted_train = {} self.preds_weighted_train = {} self.samples_weighted_train = {} self.metrics_train = [] self.metrics_idx_train = {} self.metrics_var_train = {} self.input_val = None self.target_val = None self.preds_val = None self.samples_val = None self.target_weighted_val = {} self.preds_weighted_val = {} self.samples_weighted_val = {} self.metrics_val = [] self.metrics_idx_val = {} self.metrics_var_val = {} self.input_scoring = None self.target_scoring = None self.preds_scoring = None self.samples_scoring = None self.target_weighted_scoring = {} self.preds_weighted_scoring = {} self.samples_weighted_scoring = {} self.metrics_scoring = [] self.metrics_idx_scoring = {} self.metrics_var_scoring = {} self.input_test = None self.target_test = None self.preds_test = None self.samples_test = None self.target_weighted_test = {} self.preds_weighted_test = {} self.samples_weighted_test = {} self.metrics_test = [] self.metrics_idx_test = {} self.metrics_var_test = {} self.model_names = [] self.metrics_names = [] self.metrics_dict = {'MAE': self.calc_MAE, 'RMSE': self.calc_RMSE, 'R2': self.calc_R2, 'CRPS': self.calc_CRPS, 'bias': self.calc_bias } self.linecolors = ['#0072B2', '#E69F00', '#882255', '#009E73', '#D55E00' ] def set_to_v1_vars(self): ''' This function sets the inputs and outputs to the V1 subset. ''' self.input_vars = self.v1_inputs self.target_vars = self.v1_outputs self.input_feature_len = 124 self.target_feature_len = 128 def get_xrdata(self, file, file_vars = None): ''' This function reads in a file and returns an xarray dataset with the variables specified. file_vars must be a list of strings. ''' ds = xr.open_dataset(file, engine = 'netcdf4') if file_vars is not None: ds = ds[file_vars] ds = ds.merge(self.grid_info[['lat','lon']]) ds = ds.where((ds['lat']>-999)*(ds['lat']<999), drop=True) ds = ds.where((ds['lon']>-999)*(ds['lon']<999), drop=True) return ds def get_input(self, input_file): ''' This function reads in a file and returns an xarray dataset with the input variables for the emulator. ''' # read inputs return self.get_xrdata(input_file, self.input_vars) def get_target(self, input_file): ''' This function reads in a file and returns an xarray dataset with the target variables for the emulator. ''' # read inputs ds_input = self.get_input(input_file) ds_target = self.get_xrdata(input_file.replace('.mli.','.mlo.')) # each timestep is 20 minutes which corresponds to 1200 seconds ds_target['ptend_t'] = (ds_target['state_t'] - ds_input['state_t'])/1200 # T tendency [K/s] ds_target['ptend_q0001'] = (ds_target['state_q0001'] - ds_input['state_q0001'])/1200 # Q tendency [kg/kg/s] ds_target = ds_target[self.target_vars] return ds_target def set_regexps(self, data_split, regexps): ''' This function sets the regular expressions used for getting the filelist for train, val, scoring, and test. ''' assert data_split in ['train', 'val', 'scoring', 'test'], 'Provided data_split is not valid. Available options are train, val, scoring, and test.' if data_split == 'train': self.train_regexps = regexps elif data_split == 'val': self.val_regexps = regexps elif data_split == 'scoring': self.scoring_regexps = regexps elif data_split == 'test': self.test_regexps = regexps def set_stride_sample(self, data_split, stride_sample): ''' This function sets the stride_sample for train, val, scoring, and test. ''' assert data_split in ['train', 'val', 'scoring', 'test'], 'Provided data_split is not valid. Available options are train, val, scoring, and test.' if data_split == 'train': self.train_stride_sample = stride_sample elif data_split == 'val': self.val_stride_sample = stride_sample elif data_split == 'scoring': self.scoring_stride_sample = stride_sample elif data_split == 'test': self.test_stride_sample = stride_sample def set_filelist(self, data_split): ''' This function sets the filelists corresponding to data splits for train, val, scoring, and test. ''' filelist = [] assert data_split in ['train', 'val', 'scoring', 'test'], 'Provided data_split is not valid. Available options are train, val, scoring, and test.' if data_split == 'train': assert self.train_regexps is not None, 'regexps for train is not set.' assert self.train_stride_sample is not None, 'stride_sample for train is not set.' for regexp in self.train_regexps: filelist = filelist + glob.glob(self.data_path + "*/" + regexp) self.train_filelist = sorted(filelist)[::self.train_stride_sample] elif data_split == 'val': assert self.val_regexps is not None, 'regexps for val is not set.' assert self.val_stride_sample is not None, 'stride_sample for val is not set.' for regexp in self.val_regexps: filelist = filelist + glob.glob(self.data_path + "*/" + regexp) self.val_filelist = sorted(filelist)[::self.val_stride_sample] elif data_split == 'scoring': assert self.scoring_regexps is not None, 'regexps for scoring is not set.' assert self.scoring_stride_sample is not None, 'stride_sample for scoring is not set.' for regexp in self.scoring_regexps: filelist = filelist + glob.glob(self.data_path + "*/" + regexp) self.scoring_filelist = sorted(filelist)[::self.scoring_stride_sample] elif data_split == 'test': assert self.test_regexps is not None, 'regexps for test is not set.' assert self.test_stride_sample is not None, 'stride_sample for test is not set.' for regexp in self.test_regexps: filelist = filelist + glob.glob(self.data_path + "*/" + regexp) self.test_filelist = sorted(filelist)[::self.test_stride_sample] def get_filelist(self, data_split): ''' This function returns the filelist corresponding to data splits for train, val, scoring, and test. ''' assert data_split in ['train', 'val', 'scoring', 'test'], 'Provided data_split is not valid. Available options are train, val, scoring, and test.' if data_split == 'train': assert self.train_filelist is not None, 'filelist for train is not set.' return self.train_filelist elif data_split == 'val': assert self.val_filelist is not None, 'filelist for val is not set.' return self.val_filelist elif data_split == 'scoring': assert self.scoring_filelist is not None, 'filelist for scoring is not set.' return self.scoring_filelist elif data_split == 'test': assert self.test_filelist is not None, 'filelist for test is not set.' return self.test_filelist def load_ncdata_with_generator(self, data_split): ''' This function works as a dataloader when training the emulator with raw netCDF files. This can be used as a dataloader during training or it can be used to create entire datasets. When used as a dataloader for training, I/O can slow down training considerably. This function also normalizes the data. mli corresponds to input mlo corresponds to target ''' filelist = self.get_filelist(data_split) def gen(): for file in filelist: # read inputs ds_input = self.get_input(file) # read targets ds_target = self.get_target(file) # normalization, scaling ds_input = (ds_input - self.input_mean)/(self.input_max - self.input_min) ds_target = ds_target*self.output_scale # stack # ds = ds.stack({'batch':{'sample','ncol'}}) ds_input = ds_input.stack({'batch':{'ncol'}}) ds_input = ds_input.to_stacked_array('mlvar', sample_dims=['batch'], name='mli') # dso = dso.stack({'batch':{'sample','ncol'}}) ds_target = ds_target.stack({'batch':{'ncol'}}) ds_target = ds_target.to_stacked_array('mlvar', sample_dims=['batch'], name='mlo') yield (ds_input.values, ds_target.values) return tf.data.Dataset.from_generator( gen, output_types = (tf.float64, tf.float64), output_shapes = ((None,124),(None,128)) ) def save_as_npy(self, data_split, save_path = '', save_latlontime_dict = False): ''' This function saves the training data as a .npy file. Prefix should be train or val. ''' prefix = save_path + data_split data_loader = self.load_ncdata_with_generator(data_split) npy_iterator = list(data_loader.as_numpy_iterator()) npy_input = np.concatenate([npy_iterator[x][0] for x in range(len(npy_iterator))]) npy_target = np.concatenate([npy_iterator[x][1] for x in range(len(npy_iterator))]) with open(save_path + prefix + '_input.npy', 'wb') as f: np.save(f, np.float32(npy_input)) with open(save_path + prefix + '_target.npy', 'wb') as f: np.save(f, np.float32(npy_target)) if data_split == 'train': data_files = self.train_filelist elif data_split == 'val': data_files = self.val_filelist elif data_split == 'scoring': data_files = self.scoring_filelist elif data_split == 'test': data_files = self.test_filelist if save_latlontime_dict: dates = [re.sub('^.*mli\.', '', x) for x in data_files] dates = [re.sub('\.nc$', '', x) for x in dates] repeat_dates = [] for date in dates: for i in range(self.latlonnum): repeat_dates.append(date) latlontime = {i: [(self.grid_info['lat'].values[i%self.latlonnum], self.grid_info['lon'].values[i%self.latlonnum]), repeat_dates[i]] for i in range(npy_input.shape[0])} with open(save_path + prefix + '_indextolatlontime.pkl', 'wb') as f: pickle.dump(latlontime, f) def reshape_npy(self, var_arr, var_arr_dim): ''' This function reshapes the a variable in numpy such that time gets its own axis (instead of being num_samples x num_levels). Shape of target would be (timestep, lat/lon combo, num_levels) ''' var_arr = var_arr.reshape((int(var_arr.shape[0]/self.latlonnum), self.latlonnum, var_arr_dim)) return var_arr @staticmethod def ls(dir_path = ''): ''' You can treat this as a Python wrapper for the bash command "ls". ''' return os.popen(' '.join(['ls', dir_path])).read().splitlines() @staticmethod def set_plot_params(): ''' This function sets the plot parameters for matplotlib. ''' plt.close('all') plt.rcParams.update(plt.rcParamsDefault) plt.rc('font', family='sans') plt.rcParams.update({'font.size': 32, 'lines.linewidth': 2, 'axes.labelsize': 32, 'axes.titlesize': 32, 'xtick.labelsize': 32, 'ytick.labelsize': 32, 'legend.fontsize': 32, 'axes.linewidth': 2, "pgf.texsystem": "pdflatex" }) # %config InlineBackend.figure_format = 'retina' # use the above line when working in a jupyter notebook @staticmethod def load_npy_file(load_path = ''): ''' This function loads the prediction .npy file. ''' with open(load_path, 'rb') as f: pred = np.load(f) return pred @staticmethod def load_h5_file(load_path = ''): ''' This function loads the prediction .h5 file. ''' hf = h5py.File(load_path, 'r') pred = np.array(hf.get('pred')) return pred def set_pressure_grid(self, data_split): ''' This function sets the pressure weighting for metrics. ''' assert data_split in ['train', 'val', 'scoring', 'test'], 'Provided data_split is not valid. Available options are train, val, scoring, and test.' if data_split == 'train': assert self.input_train is not None state_ps = self.input_train[:,120]*(self.input_max['state_ps'].values - self.input_min['state_ps'].values) + self.input_mean['state_ps'].values state_ps = np.reshape(state_ps, (-1, self.latlonnum)) pressure_grid_p1 = np.array(self.grid_info['P0']*self.grid_info['hyai'])[:,np.newaxis,np.newaxis] pressure_grid_p2 = self.grid_info['hybi'].values[:, np.newaxis, np.newaxis] * state_ps[np.newaxis, :, :] self.pressure_grid_train = pressure_grid_p1 + pressure_grid_p2 self.dp_train = self.pressure_grid_train[1:61,:,:] - self.pressure_grid_train[0:60,:,:] self.dp_train = self.dp_train.transpose((1,2,0)) elif data_split == 'val': assert self.input_val is not None state_ps = self.input_val[:,120]*(self.input_max['state_ps'].values - self.input_min['state_ps'].values) + self.input_mean['state_ps'].values state_ps = np.reshape(state_ps, (-1, self.latlonnum)) pressure_grid_p1 = np.array(self.grid_info['P0']*self.grid_info['hyai'])[:,np.newaxis,np.newaxis] pressure_grid_p2 = self.grid_info['hybi'].values[:, np.newaxis, np.newaxis] * state_ps[np.newaxis, :, :] self.pressure_grid_val = pressure_grid_p1 + pressure_grid_p2 self.dp_val = self.pressure_grid_val[1:61,:,:] - self.pressure_grid_val[0:60,:,:] self.dp_val = self.dp_val.transpose((1,2,0)) elif data_split == 'scoring': assert self.input_scoring is not None state_ps = self.input_scoring[:,120]*(self.input_max['state_ps'].values - self.input_min['state_ps'].values) + self.input_mean['state_ps'].values state_ps = np.reshape(state_ps, (-1, self.latlonnum)) pressure_grid_p1 = np.array(self.grid_info['P0']*self.grid_info['hyai'])[:,np.newaxis,np.newaxis] pressure_grid_p2 = self.grid_info['hybi'].values[:, np.newaxis, np.newaxis] * state_ps[np.newaxis, :, :] self.pressure_grid_scoring = pressure_grid_p1 + pressure_grid_p2 self.dp_scoring = self.pressure_grid_scoring[1:61,:,:] - self.pressure_grid_scoring[0:60,:,:] self.dp_scoring = self.dp_scoring.transpose((1,2,0)) elif data_split == 'test': assert self.input_test is not None state_ps = self.input_test[:,120]*(self.input_max['state_ps'].values - self.input_min['state_ps'].values) + self.input_mean['state_ps'].values state_ps = np.reshape(state_ps, (-1, self.latlonnum)) pressure_grid_p1 = np.array(self.grid_info['P0']*self.grid_info['hyai'])[:,np.newaxis,np.newaxis] pressure_grid_p2 = self.grid_info['hybi'].values[:, np.newaxis, np.newaxis] * state_ps[np.newaxis, :, :] self.pressure_grid_test = pressure_grid_p1 + pressure_grid_p2 self.dp_test = self.pressure_grid_test[1:61,:,:] - self.pressure_grid_test[0:60,:,:] self.dp_test = self.dp_test.transpose((1,2,0)) def get_pressure_grid_plotting(self, data_split): ''' This function creates the temporally and zonally averaged pressure grid corresponding to a given data split. ''' filelist = self.get_filelist(data_split) ps = np.concatenate([self.get_xrdata(file, ['state_ps'])['state_ps'].values[np.newaxis, :] for file in tqdm(filelist)], axis = 0)[:, :, np.newaxis] hyam_component = self.hyam[np.newaxis, np.newaxis, :]*self.p0 hybm_component = self.hybm[np.newaxis, np.newaxis, :]*ps pressures = np.mean(hyam_component + hybm_component, axis = 0) pg_lats = [] def find_keys(dictionary, value): keys = [] for key, val in dictionary.items(): if val[0] == value: keys.append(key) return keys for lat in self.lats: indices = find_keys(self.indextolatlon, lat) pg_lats.append(np.mean(pressures[indices, :], axis = 0)[:, np.newaxis]) pressure_grid_plotting = np.concatenate(pg_lats, axis = 1) return pressure_grid_plotting def output_weighting(self, output, data_split): ''' This function does four transformations, and assumes we are using V1 variables: [0] Undos the output scaling [1] Weight vertical levels by dp/g [2] Weight horizontal area of each grid cell by a[x]/mean(a[x]) [3] Unit conversion to a common energy unit ''' assert data_split in ['train', 'val', 'scoring', 'test'], 'Provided data_split is not valid. Available options are train, val, scoring, and test.' num_samples = output.shape[0] heating = output[:,:60].reshape((int(num_samples/self.latlonnum), self.latlonnum, 60)) moistening = output[:,60:120].reshape((int(num_samples/self.latlonnum), self.latlonnum, 60)) netsw = output[:,120].reshape((int(num_samples/self.latlonnum), self.latlonnum)) flwds = output[:,121].reshape((int(num_samples/self.latlonnum), self.latlonnum)) precsc = output[:,122].reshape((int(num_samples/self.latlonnum), self.latlonnum)) precc = output[:,123].reshape((int(num_samples/self.latlonnum), self.latlonnum)) sols = output[:,124].reshape((int(num_samples/self.latlonnum), self.latlonnum)) soll = output[:,125].reshape((int(num_samples/self.latlonnum), self.latlonnum)) solsd = output[:,126].reshape((int(num_samples/self.latlonnum), self.latlonnum)) solld = output[:,127].reshape((int(num_samples/self.latlonnum), self.latlonnum)) # heating = heating.transpose((2,0,1)) # moistening = moistening.transpose((2,0,1)) # scalar_outputs = scalar_outputs.transpose((2,0,1)) # [0] Undo output scaling heating = heating/self.output_scale['ptend_t'].values[np.newaxis, np.newaxis, :] moistening = moistening/self.output_scale['ptend_q0001'].values[np.newaxis, np.newaxis, :] netsw = netsw/self.output_scale['cam_out_NETSW'].values flwds = flwds/self.output_scale['cam_out_FLWDS'].values precsc = precsc/self.output_scale['cam_out_PRECSC'].values precc = precc/self.output_scale['cam_out_PRECC'].values sols = sols/self.output_scale['cam_out_SOLS'].values soll = soll/self.output_scale['cam_out_SOLL'].values solsd = solsd/self.output_scale['cam_out_SOLSD'].values solld = solld/self.output_scale['cam_out_SOLLD'].values # [1] Weight vertical levels by dp/g # only for vertically-resolved variables, e.g. ptend_{t,q0001} # dp/g = -\rho * dz if data_split == 'train': dp = self.dp_train elif data_split == 'val': dp = self.dp_val elif data_split == 'scoring': dp = self.dp_scoring elif data_split == 'test': dp = self.dp_test heating = heating * dp/self.grav moistening = moistening * dp/self.grav # [2] weight by area heating = heating * self.area_wgt[np.newaxis, :, np.newaxis] moistening = moistening * self.area_wgt[np.newaxis, :, np.newaxis] netsw = netsw * self.area_wgt[np.newaxis, :] flwds = flwds * self.area_wgt[np.newaxis, :] precsc = precsc * self.area_wgt[np.newaxis, :] precc = precc * self.area_wgt[np.newaxis, :] sols = sols * self.area_wgt[np.newaxis, :] soll = soll * self.area_wgt[np.newaxis, :] solsd = solsd * self.area_wgt[np.newaxis, :] solld = solld * self.area_wgt[np.newaxis, :] # [3] unit conversion heating = heating * self.target_energy_conv['ptend_t'] moistening = moistening * self.target_energy_conv['ptend_q0001'] netsw = netsw * self.target_energy_conv['cam_out_NETSW'] flwds = flwds * self.target_energy_conv['cam_out_FLWDS'] precsc = precsc * self.target_energy_conv['cam_out_PRECSC'] precc = precc * self.target_energy_conv['cam_out_PRECC'] sols = sols * self.target_energy_conv['cam_out_SOLS'] soll = soll * self.target_energy_conv['cam_out_SOLL'] solsd = solsd * self.target_energy_conv['cam_out_SOLSD'] solld = solld * self.target_energy_conv['cam_out_SOLLD'] return {'ptend_t':heating, 'ptend_q0001':moistening, 'cam_out_NETSW':netsw, 'cam_out_FLWDS':flwds, 'cam_out_PRECSC':precsc, 'cam_out_PRECC':precc, 'cam_out_SOLS':sols, 'cam_out_SOLL':soll, 'cam_out_SOLSD':solsd, 'cam_out_SOLLD':solld} def reweight_target(self, data_split): ''' data_split should be train, val, scoring, or test weights target variables assuming V1 outputs using the output_weighting function ''' assert data_split in ['train', 'val', 'scoring', 'test'], 'Provided data_split is not valid. Available options are train, val, scoring, and test.' if data_split == 'train': assert self.target_train is not None self.target_weighted_train = self.output_weighting(self.target_train, data_split) elif data_split == 'val': assert self.target_val is not None self.target_weighted_val = self.output_weighting(self.target_val, data_split) elif data_split == 'scoring': assert self.target_scoring is not None self.target_weighted_scoring = self.output_weighting(self.target_scoring, data_split) elif data_split == 'test': assert self.target_test is not None self.target_weighted_test = self.output_weighting(self.target_test, data_split) def reweight_preds(self, data_split): ''' weights predictions assuming V1 outputs using the output_weighting function ''' assert data_split in ['train', 'val', 'scoring', 'test'], 'Provided data_split is not valid. Available options are train, val, scoring, and test.' assert self.model_names is not None if data_split == 'train': assert self.preds_train is not None for model_name in self.model_names: self.preds_weighted_train[model_name] = self.output_weighting(self.preds_train[model_name], data_split) elif data_split == 'val': assert self.preds_val is not None for model_name in self.model_names: self.preds_weighted_val[model_name] = self.output_weighting(self.preds_val[model_name], data_split) elif data_split == 'scoring': assert self.preds_scoring is not None for model_name in self.model_names: self.preds_weighted_scoring[model_name] = self.output_weighting(self.preds_scoring[model_name], data_split) elif data_split == 'test': assert self.preds_test is not None for model_name in self.model_names: self.preds_weighted_test[model_name] = self.output_weighting(self.preds_test[model_name], data_split) def calc_MAE(self, pred, target, avg_grid = True): ''' calculate 'globally averaged' mean absolute error for vertically-resolved variables, shape should be time x grid x level for scalars, shape should be time x grid returns vector of length level or 1 ''' assert pred.shape[1] == self.latlonnum assert pred.shape == target.shape mae = np.abs(pred - target).mean(axis = 0) if avg_grid: return mae.mean(axis = 0) # we decided to average globally at end else: return mae def calc_RMSE(self, pred, target, avg_grid = True): ''' calculate 'globally averaged' root mean squared error for vertically-resolved variables, shape should be time x grid x level for scalars, shape should be time x grid returns vector of length level or 1 ''' assert pred.shape[1] == self.latlonnum assert pred.shape == target.shape sq_diff = (pred - target)**2 rmse = np.sqrt(sq_diff.mean(axis = 0)) # mean over time if avg_grid: return rmse.mean(axis = 0) # we decided to separately average globally at end else: return rmse def calc_R2(self, pred, target, avg_grid = True): ''' calculate 'globally averaged' R-squared for vertically-resolved variables, input shape should be time x grid x level for scalars, input shape should be time x grid returns vector of length level or 1 ''' assert pred.shape[1] == self.latlonnum assert pred.shape == target.shape sq_diff = (pred - target)**2 tss_time = (target - target.mean(axis = 0)[np.newaxis, ...])**2 # mean over time r_squared = 1 - sq_diff.sum(axis = 0)/tss_time.sum(axis = 0) # sum over time if avg_grid: return r_squared.mean(axis = 0) # we decided to separately average globally at end else: return r_squared def calc_bias(self, pred, target, avg_grid = True): ''' calculate bias for vertically-resolved variables, input shape should be time x grid x level for scalars, input shape should be time x grid returns vector of length level or 1 ''' assert pred.shape[1] == self.latlonnum assert pred.shape == target.shape bias = pred.mean(axis = 0) - target.mean(axis = 0) if avg_grid: return bias.mean(axis = 0) # we decided to separately average globally at end else: return bias def calc_CRPS(self, preds, target, avg_grid = True): ''' calculate 'globally averaged' continuous ranked probability score for vertically-resolved variables, input shape should be time x grid x level x num_crps_samples for scalars, input shape should be time x grid x num_crps_samples returns vector of length level or 1 ''' assert preds.shape[1] == self.latlonnum num_crps = preds.shape[-1] mae = np.mean(np.abs(preds - target[..., np.newaxis]), axis = (0, -1)) # mean over time and crps samples diff = preds[..., 1:] - preds[..., :-1] count = np.arange(1, num_crps) * np.arange(num_crps - 1, 0, -1) spread = (diff * count[np.newaxis, np.newaxis, np.newaxis, :]).mean(axis = (0, -1)) # mean over time and crps samples crps = mae - spread/(num_crps*(num_crps-1)) # already divided by two in spread by exploiting symmetry if avg_grid: return crps.mean(axis = 0) # we decided to separately average globally at end else: return crps def create_metrics_df(self, data_split): ''' creates a dataframe of metrics for each model ''' assert data_split in ['train', 'val', 'scoring', 'test'], \ 'Provided data_split is not valid. Available options are train, val, scoring, and test.' assert len(self.model_names) != 0 assert len(self.metrics_names) != 0 assert len(self.target_vars) != 0 assert self.target_feature_len is not None if data_split == 'train': assert len(self.preds_weighted_train) != 0 assert len(self.target_weighted_train) != 0 for model_name in self.model_names: df_var = pd.DataFrame(columns = self.metrics_names, index = self.target_vars) df_var.index.name = 'variable' df_idx = pd.DataFrame(columns = self.metrics_names, index = range(self.target_feature_len)) df_idx.index.name = 'output_idx' for metric_name in self.metrics_names: current_idx = 0 for target_var in self.target_vars: metric = self.metrics_dict[metric_name](self.preds_weighted_train[model_name][target_var], self.target_weighted_train[target_var]) df_var.loc[target_var, metric_name] = np.mean(metric) df_idx.loc[current_idx:current_idx + self.var_lens[target_var] - 1, metric_name] = np.atleast_1d(metric) current_idx += self.var_lens[target_var] self.metrics_var_train[model_name] = df_var self.metrics_idx_train[model_name] = df_idx elif data_split == 'val': assert len(self.preds_weighted_val) != 0 assert len(self.target_weighted_val) != 0 for model_name in self.model_names: df_var = pd.DataFrame(columns = self.metrics_names, index = self.target_vars) df_var.index.name = 'variable' df_idx = pd.DataFrame(columns = self.metrics_names, index = range(self.target_feature_len)) df_idx.index.name = 'output_idx' for metric_name in self.metrics_names: current_idx = 0 for target_var in self.target_vars: metric = self.metrics_dict[metric_name](self.preds_weighted_val[model_name][target_var], self.target_weighted_val[target_var]) df_var.loc[target_var, metric_name] = np.mean(metric) df_idx.loc[current_idx:current_idx + self.var_lens[target_var] - 1, metric_name] = np.atleast_1d(metric) current_idx += self.var_lens[target_var] self.metrics_var_val[model_name] = df_var self.metrics_idx_val[model_name] = df_idx elif data_split == 'scoring': assert len(self.preds_weighted_scoring) != 0 assert len(self.target_weighted_scoring) != 0 for model_name in self.model_names: df_var = pd.DataFrame(columns = self.metrics_names, index = self.target_vars) df_var.index.name = 'variable' df_idx = pd.DataFrame(columns = self.metrics_names, index = range(self.target_feature_len)) df_idx.index.name = 'output_idx' for metric_name in self.metrics_names: current_idx = 0 for target_var in self.target_vars: metric = self.metrics_dict[metric_name](self.preds_weighted_scoring[model_name][target_var], self.target_weighted_scoring[target_var]) df_var.loc[target_var, metric_name] = np.mean(metric) df_idx.loc[current_idx:current_idx + self.var_lens[target_var] - 1, metric_name] = np.atleast_1d(metric) current_idx += self.var_lens[target_var] self.metrics_var_scoring[model_name] = df_var self.metrics_idx_scoring[model_name] = df_idx elif data_split == 'test': assert len(self.preds_weighted_test) != 0 assert len(self.target_weighted_test) != 0 for model_name in self.model_names: df_var = pd.DataFrame(columns = self.metrics_names, index = self.target_vars) df_var.index.name = 'variable' df_idx = pd.DataFrame(columns = self.metrics_names, index = range(self.target_feature_len)) df_idx.index.name = 'output_idx' for metric_name in self.metrics_names: current_idx = 0 for target_var in self.target_vars: metric = self.metrics_dict[metric_name](self.preds_weighted_test[model_name][target_var], self.target_weighted_test[target_var]) df_var.loc[target_var, metric_name] = np.mean(metric) df_idx.loc[current_idx:current_idx + self.var_lens[target_var] - 1, metric_name] = np.atleast_1d(metric) current_idx += self.var_lens[target_var] self.metrics_var_test[model_name] = df_var self.metrics_idx_test[model_name] = df_idx def reshape_daily(self, output): ''' This function returns two numpy arrays, one for each vertically resolved variable (heating and moistening). Dimensions of expected input are num_samples by 128 (number of target features). Output argument is espected to be have dimensions of num_samples by features. Heating is expected to be the first feature, and moistening is expected to be the second feature. Data is expected to use a stride_sample of 6. (12 samples per day, 20 min timestep). ''' num_samples = output.shape[0] heating = output[:,:60].reshape((int(num_samples/self.latlonnum), self.latlonnum, 60)) moistening = output[:,60:120].reshape((int(num_samples/self.latlonnum), self.latlonnum, 60)) heating_daily = np.mean(heating.reshape((heating.shape[0]//12, 12, self.latlonnum, 60)), axis = 1) # Nday x lotlonnum x 60 moistening_daily = np.mean(moistening.reshape((moistening.shape[0]//12, 12, self.latlonnum, 60)), axis = 1) # Nday x lotlonnum x 60 heating_daily_long = [] moistening_daily_long = [] for i in range(len(self.lats)): heating_daily_long.append(np.mean(heating_daily[:,self.lat_indices_list[i],:],axis=1)) moistening_daily_long.append(np.mean(moistening_daily[:,self.lat_indices_list[i],:],axis=1)) heating_daily_long = np.array(heating_daily_long) # lat x Nday x 60 moistening_daily_long = np.array(moistening_daily_long) # lat x Nday x 60 return heating_daily_long, moistening_daily_long def plot_r2_analysis(self, pressure_grid_plotting, save_path = ''): ''' This function plots the R2 pressure latitude figure shown in the SI. ''' self.set_plot_params() n_model = len(self.model_names) fig, ax = plt.subplots(2,n_model, figsize=(n_model*12,18)) y = np.array(range(60)) X, Y = np.meshgrid(np.sin(self.lats*np.pi/180), y) Y = pressure_grid_plotting/100 test_heat_daily_long, test_moist_daily_long = self.reshape_daily(self.target_scoring) for i, model_name in enumerate(self.model_names): pred_heat_daily_long, pred_moist_daily_long = self.reshape_daily(self.preds_scoring[model_name]) coeff = 1 - np.sum( (pred_heat_daily_long-test_heat_daily_long)**2, axis=1)/np.sum( (test_heat_daily_long-np.mean(test_heat_daily_long, axis=1)[:,None,:])**2, axis=1) coeff = coeff[self.sort_lat_key,:] coeff = coeff.T contour_plot = ax[0,i].pcolor(X, Y, coeff,cmap='Blues', vmin = 0, vmax = 1) # pcolormesh ax[0,i].contour(X, Y, coeff, [0.7], colors='orange', linewidths=[4]) ax[0,i].contour(X, Y, coeff, [0.9], colors='yellow', linewidths=[4]) ax[0,i].set_ylim(ax[0,i].get_ylim()[::-1]) ax[0,i].set_title(self.model_names[i] + " - Heating") ax[0,i].set_xticks([]) coeff = 1 - np.sum( (pred_moist_daily_long-test_moist_daily_long)**2, axis=1)/np.sum( (test_moist_daily_long-np.mean(test_moist_daily_long, axis=1)[:,None,:])**2, axis=1) coeff = coeff[self.sort_lat_key,:] coeff = coeff.T contour_plot = ax[1,i].pcolor(X, Y, coeff,cmap='Blues', vmin = 0, vmax = 1) # pcolormesh ax[1,i].contour(X, Y, coeff, [0.7], colors='orange', linewidths=[4]) ax[1,i].contour(X, Y, coeff, [0.9], colors='yellow', linewidths=[4]) ax[1,i].set_ylim(ax[1,i].get_ylim()[::-1]) ax[1,i].set_title(self.model_names[i] + " - Moistening") ax[1,i].xaxis.set_ticks([np.sin(-50/180*np.pi), 0, np.sin(50/180*np.pi)]) ax[1,i].xaxis.set_ticklabels(['50$^\circ$S', '0$^\circ$', '50$^\circ$N']) ax[1,i].xaxis.set_tick_params(width = 2) if i != 0: ax[0,i].set_yticks([]) ax[1,i].set_yticks([]) # lines below for x and y label axes are valid if 3 models are considered # we want to put only one label for each axis # if nbr of models is different from 3 please adjust label location to center it #ax[1,1].xaxis.set_label_coords(-0.10,-0.10) ax[0,0].set_ylabel("Pressure [hPa]") ax[0,0].yaxis.set_label_coords(-0.2,-0.09) # (-1.38,-0.09) ax[0,0].yaxis.set_ticks([1000,800,600,400,200,0]) ax[1,0].yaxis.set_ticks([1000,800,600,400,200,0]) fig.subplots_adjust(right=0.8) cbar_ax = fig.add_axes([0.82, 0.12, 0.02, 0.76]) cb = fig.colorbar(contour_plot, cax=cbar_ax) cb.set_label("Skill Score "+r'$\left(\mathrm{R^{2}}\right)$',labelpad=50.1) plt.suptitle("Baseline Models Skill for Vertically Resolved Tendencies", y = 0.97) plt.subplots_adjust(hspace=0.13) plt.show() plt.savefig(save_path + 'press_lat_diff_models.png', bbox_inches='tight', pad_inches=0.1 , dpi = 300) @staticmethod def reshape_input_for_cnn(npy_input, save_path = ''): ''' This function reshapes a numpy input array to be compatible with CNN training. Each variable becomes its own channel. For the input there are 6 channels, each with 60 vertical levels. The last 4 channels correspond to scalars repeated across all 60 levels. This is for V1 data only! (V2 data has more variables) ''' npy_input_cnn = np.stack([ npy_input[:, 0:60], npy_input[:, 60:120], np.repeat(npy_input[:, 120][:, np.newaxis], 60, axis = 1), np.repeat(npy_input[:, 121][:, np.newaxis], 60, axis = 1), np.repeat(npy_input[:, 122][:, np.newaxis], 60, axis = 1), np.repeat(npy_input[:, 123][:, np.newaxis], 60, axis = 1)], axis = 2) if save_path != '': with open(save_path + 'train_input_cnn.npy', 'wb') as f: np.save(f, np.float32(npy_input_cnn)) return npy_input_cnn @staticmethod def reshape_target_for_cnn(npy_target, save_path = ''): ''' This function reshapes a numpy target array to be compatible with CNN training. Each variable becomes its own channel. For the input there are 6 channels, each with 60 vertical levels. The last 4 channels correspond to scalars repeated across all 60 levels. This is for V1 data only! (V2 data has more variables) ''' npy_target_cnn = np.stack([ npy_target[:, 0:60], npy_target[:, 60:120], np.repeat(npy_target[:, 120][:, np.newaxis], 60, axis = 1), np.repeat(npy_target[:, 121][:, np.newaxis], 60, axis = 1), np.repeat(npy_target[:, 122][:, np.newaxis], 60, axis = 1), np.repeat(npy_target[:, 123][:, np.newaxis], 60, axis = 1), np.repeat(npy_target[:, 124][:, np.newaxis], 60, axis = 1), np.repeat(npy_target[:, 125][:, np.newaxis], 60, axis = 1), np.repeat(npy_target[:, 126][:, np.newaxis], 60, axis = 1), np.repeat(npy_target[:, 127][:, np.newaxis], 60, axis = 1)], axis = 2) if save_path != '': with open(save_path + 'train_target_cnn.npy', 'wb') as f: np.save(f, np.float32(npy_target_cnn)) return npy_target_cnn @staticmethod def reshape_target_from_cnn(npy_predict_cnn, save_path = ''): ''' This function reshapes CNN target to (num_samples, 128) for standardized metrics. This is for V1 data only! (V2 data has more variables) ''' npy_predict_cnn_reshaped = np.concatenate([ npy_predict_cnn[:,:,0], npy_predict_cnn[:,:,1], np.mean(npy_predict_cnn[:,:,2], axis = 1)[:, np.newaxis], np.mean(npy_predict_cnn[:,:,3], axis = 1)[:, np.newaxis], np.mean(npy_predict_cnn[:,:,4], axis = 1)[:, np.newaxis], np.mean(npy_predict_cnn[:,:,5], axis = 1)[:, np.newaxis], np.mean(npy_predict_cnn[:,:,6], axis = 1)[:, np.newaxis], np.mean(npy_predict_cnn[:,:,7], axis = 1)[:, np.newaxis], np.mean(npy_predict_cnn[:,:,8], axis = 1)[:, np.newaxis], np.mean(npy_predict_cnn[:,:,9], axis = 1)[:, np.newaxis]], axis = 1) if save_path != '': with open(save_path + 'cnn_predict_reshaped.npy', 'wb') as f: np.save(f, np.float32(npy_predict_cnn_reshaped)) return npy_predict_cnn_reshaped