Upload data_utils.py

data_utils.py

@@ -0,0 +1,993 @@
+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
+
+
+
+
+  
+
+
+
+