From 7ba1ac3ef6cf5459bf112dbfae936cc61e2f099b Mon Sep 17 00:00:00 2001 From: leelew Date: Sat, 10 Dec 2022 20:37:31 +0800 Subject: [PATCH] New edition --- src/config.py | 6 +- src/data.py | 211 +++++------------------- src/data_gen.py | 12 +- src/eval.py | 82 ++++++--- src/figure1.pdf | Bin 9185 -> 0 bytes src/layers.py | 110 ++++++++++++ src/loss.py | 78 +++------ src/main.py | 23 ++- src/{preprocess => }/make_input_data.py | 169 +++++++++++-------- src/metric.py | 37 +++++ src/model.py | 204 +++++++++++------------ src/plot.py | 139 +++++++++++----- src/postprocess.py | 2 + src/preprocess/colm2era.py | 14 -- src/train.py | 161 +++++++++--------- src/utils.py | 45 ++++- 16 files changed, 715 insertions(+), 578 deletions(-) delete mode 100644 src/figure1.pdf create mode 100644 src/layers.py rename src/{preprocess => }/make_input_data.py (54%) create mode 100644 src/metric.py delete mode 100644 src/preprocess/colm2era.py diff --git a/src/config.py b/src/config.py index 00e82b0..a7b2084 100644 --- a/src/config.py +++ b/src/config.py @@ -10,9 +10,9 @@ def parse_args(): parser.add_argument('--work_path', type=str, default="/data/lilu/PHydro") # data - parser.add_argument('--reuse_input', type=bool, default=False) + parser.add_argument('--reuse_input', type=bool, default=True) parser.add_argument('--use_ancillary', type=bool, default=True) - parser.add_argument('--seq_len', type=int, default=365) + parser.add_argument('--seq_len', type=int, default=100) parser.add_argument('--interval', type=int, default=365) parser.add_argument('--window_size', type=int, default=0) parser.add_argument('--num_out', type=int, default=6) @@ -23,7 +23,7 @@ def parse_args(): # model parser.add_argument('--model_name', type=str, default="hard_multi_tasks_v1") parser.add_argument('--learning_rate', type=float, default=0.001) - parser.add_argument('--epochs', type=int, default=300) + parser.add_argument('--epochs', type=int, default=200) parser.add_argument('--niter', type=int, default=100) parser.add_argument('--hidden_size', type=int, default=64) parser.add_argument('--alpha', type=float, default=0.5) diff --git a/src/data.py b/src/data.py index 850ce6d..1c64e16 100644 --- a/src/data.py +++ b/src/data.py @@ -3,11 +3,13 @@ <<<<<<<< HEAD Author: Lu Li -11/1/2022 - First edition +11/1/2022 - V1.0 +12/9/2022 - V2.0 """ import json import numpy as np +from utils import make_CoLM_soil_depth class Dataset(): @@ -20,24 +22,26 @@ def __init__(self, cfg, mode): self.window_size = cfg["window_size"] self.num_out = cfg["num_out"] self.mode = mode + # calculate soil depth in CoLM + depth, zi = make_CoLM_soil_depth() # m/cm + self.soil = [70, 210, 720, 10*(zi[9]-100)] # mm def fit(self): - # load input data [nt,ngrid,...] + # load input data shape as [nt,ngrid,.] forcing, hydro, ancillary = self._load_input() + # remove unreasonable runoff + hydro = self._remove_rnof_outlier(hydro, threshold=20) + # Correct water balance [nt-1,ngrid,...] - # NOTE: CoLM soil moisture is available as mean of - # day, however, the delta(swc) in water balance - # should be swc(24)-swc(0). Thus we need to - # correct the water balance by add the residual - # of water balance into precipitation. + # NOTE: CoLM water balance is defined as + # P = ET + R + delta(SWC+ldew+wa+scv) + # In our study, we focus on simulating variables ET, + # R, SWC, thus we treat `wa`... as static variable. + # Thus, to enforce water balance, we need to + # remove the residual of water balance to P. forcing, hydro, aux = self._correct_mass_conserve(forcing, hydro) - # remove unreasonable runoff - # NOTE: 30mm/day may not be the best setting, but we didn't - # found a better way to remove unreasonable runoff. - hydro = self._remove_outlier(hydro, threshold=30) - # get scaler if self.mode == 'train': scaler = self._get_z_scaler(forcing, hydro) @@ -61,62 +65,29 @@ def fit(self): ancillary = np.tile(ancillary[np.newaxis],(forcing.shape[0],1,1)) forcing = np.concatenate([forcing, ancillary], axis=-1) - # trans nt and ngrid [ngrid, nt, nfeat] + # trans shape as [ngrid,nt,.] forcing = np.transpose(forcing, (1, 0, 2)) hydro = np.transpose(hydro, (1, 0, 2)) aux = np.transpose(aux, (1, 0)) # make training/test data - if self.mode == 'train': - return forcing, hydro, aux - elif self.mode == 'test': - forcing, hydro, aux = self._make_inference_data(forcing, hydro, aux, self.seq_len) - return forcing, hydro, aux + return forcing, hydro, aux def _load_input(self): - forcing = np.load(self.inputs_path+"forcing_gd_9km_{}.npy".format(self.mode)) - hydro = np.load(self.inputs_path+"hydro_gd_9km_{}.npy".format(self.mode)) - ancillary = np.load(self.inputs_path+"ancil_gd_9km.npy") + forcing = np.load(self.inputs_path+"gd_9km_forcing_{}.npy".format(self.mode)) + hydro = np.load(self.inputs_path+"gd_9km_hydrology_{}.npy".format(self.mode)) + ancillary = np.load(self.inputs_path+"gd_9km_ancillary.npy") return forcing, hydro, ancillary - def _correct_mass_conserve(self, forcing, hydro): - soil_depth = [70, 210, 720, 1864.6] # mm - hydro_prev, hydro, forcing = hydro[:-1], hydro[1:], forcing[1:] - swvl, swvl_prev = 0, 0 - for i in range(4): - swvl += hydro[:, :, i]*soil_depth[i] - swvl_prev += hydro_prev[:,:,i]*soil_depth[i] - mc_in = np.nansum(forcing[:,:,:2], axis=-1) + swvl_prev - mc_out = swvl + np.nansum(hydro[:,:,4:], axis=-1) - diff = mc_in - mc_out - - # if diff>0, then add to runoff; - # if diff<0, then add to precipitation; - for i in range(diff.shape[0]): - for j in range(diff.shape[1]): - tmp = diff[i,j] - if tmp < 0: - forcing[i,j,0] = forcing[i,j,0]-diff[i,j] - else: - hydro[i,j,-1] = hydro[i,j,-1]+diff[i,j] - # get mass in - aux = np.nansum(forcing[:,:,:2], axis=-1) + swvl_prev - return forcing, hydro, aux - - def _remove_outlier(self, hydro, threshold=30): - """ - std = np.nanstd(input, axis=(0), keepdims=True) # (1, ngrids, nfeat) - mean = np.nanmean(input, axis=(0), keepdims=True) # (1, ngrids, nfeat) - input[np.where(input > (mean+3*std))] = np.nan - input[np.where(input < (mean-3*std))] = np.nan - self.remove_outlier = True - """ - """ - # @(Zhongwang Wei): remove unreasonable runoff > 200mm/day and - # interplote by adjacency two days + def _remove_rnof_outlier(self, hydro, threshold=20): + # NOTE: In CoLM, there are some runoff larger than 1e3 mm/h, this is + # caused by the parameterization of soil moisture, which cannot + # be solved now. Thus I remove runoff larger than 20 mm/h to + # ensure the robustness of results. rnof = hydro[:, :, -1] rnof[rnof > threshold] = np.nan - nt, ngrid, nout = hydro.shape + # @(Zhongwang Wei): interplote by adjacency two days + """ for i in range(ngrid): tmp = rnof[:, i] if np.isnan(tmp).any(): @@ -129,20 +100,24 @@ def _remove_outlier(self, hydro, threshold=30): else: tmp[j] = np.nanmean(tmp) # (tmp[j-1]+tmp[j+1])/2 rnof[:, i] = tmp - hydro[:, :, -1] = rnof """ - rnof = hydro[:, :, -1] - rnof[rnof > threshold] = np.nan hydro[:, :, -1] = rnof return hydro - def _get_minmax_scaler(self, x, y): - scaler = {} - scaler["x_min"] = np.nanmin(x, axis=(0), keepdims=True).tolist() - scaler["x_max"] = np.nanmax(x, axis=(0), keepdims=True).tolist() - scaler["y_min"] = np.nanmin(y, axis=(0), keepdims=True).tolist() - scaler["y_max"] = np.nanmax(y, axis=(0), keepdims=True).tolist() - return scaler + def _correct_mass_conserve(self, forcing, hydro): + hydro_prev, hydro, forcing = hydro[:-1], hydro[1:], forcing[1:] + # cal mass in/out + swvl, swvl_prev = 0, 0 + for i in range(4): + swvl += hydro[:, :, i]*self.soil[i] + swvl_prev += hydro_prev[:,:,i]*self.soil[i] + mc_in = np.sum(forcing[:,:,:2], axis=-1) + swvl_prev + mc_out = swvl + np.sum(hydro[:,:,4:], axis=-1) + # cal diff in mass balance caused by wa, ldew, scv, xerror + diff = mc_in - mc_out + # get mass in after remove diff in balance + aux = np.sum(forcing[:,:,:2], axis=-1) - diff + return forcing, hydro, aux def _get_z_scaler(self, x, y): scaler = {} @@ -162,6 +137,7 @@ def _load_scaler(self, inputs_path): return scaler def _z_normalize(self, input, scaler, is_feat): + """normalize features using pre-computed statistics.""" if is_feat: input = (input - np.array(scaler["x_mean"])) / ( np.array(scaler["x_std"])) @@ -170,107 +146,8 @@ def _z_normalize(self, input, scaler, is_feat): np.array(scaler["y_std"])) return input - def _minmax_normalize(self, input, scaler, is_feat): - """normalize features using pre-computed statistics.""" - if is_feat: - input = (input - np.array(scaler["x_min"])) / ( - np.array(scaler["x_max"])-np.array(scaler["x_min"])) - else: - input = (input - np.array(scaler["y_min"])) / ( - np.array(scaler["y_max"])-np.array(scaler["y_min"])) - return input - def _spatial_normalize(self, static): # (ngrid, nfeat) for static data mean = np.nanmean(static, axis=(0), keepdims=True) std = np.nanstd(static, axis=(0), keepdims=True) - return (static-mean)/std - - def _make_inference_data(self, - x, - y, - aux, - seq_len=365, - interval=1, - window_size=0): - x_, y_, aux_ = [], [], [] - for i in range(x.shape[0]): - tmpx, tmpy, tmp_aux = self._reshape_1d_data( - x[i], y[i], aux[i], seq_len, interval, window_size) - x_.append(tmpx) - y_.append(tmpy) - aux_.append(tmp_aux) - # (ngrids, nsamples, seq_len, nfeat) - return np.stack(x_, axis=0), np.stack(y_, axis=0), np.stack(aux_, axis=0) - - def _reshape_1d_data(self, - x, - y, - aux, - seq_len=365, - interval=1, - window_size=0): - """reshape data into LSTM many-to-one input samples - - Parameters - ---------- - x : np.ndarray - Input features of shape [num_samples, num_features] - y : np.ndarray - Output feature of shape [num_samples, 1] - seq_length : int - Length of the requested input sequences. - interval: int - interval of time length to generate samples. - window_size: int - window size between x and y - - Returns - ------- - x_new: np.ndarray - shape of [num_samples*, seq_length, num_features], where - num_samples* is equal to num_samples - seq_length + 1, due to - the need of a warm start at the beginning - y_new: np.ndarray - The target value for each sample in x_new - """ - num_samples, num_features = x.shape - _, num_out = y.shape - n = (num_samples-seq_len+1) // interval - x_new = np.zeros((n, seq_len, num_features))*np.nan - y_new = np.zeros((n, num_out))*np.nan - aux_new = np.zeros((n, ))*np.nan - - for i in range(n): - x_new[i] = x[i*interval:i*interval+seq_len] - y_new[i] = y[i*interval+seq_len-1] - aux_new[i] = aux[i*interval+seq_len-1] - return x_new, y_new, aux_new - - def __split_into_batch(self, X, y, seq_len=365, offset=1, window_size=0): - """split training data into batches with size of batch_size - - Params - ------ - offset: [float] 0-1, how to offset the batches (e.g., 0.5 means that - the first batch will be 0-365 and the second will be 182-547) - """ - #(nt_, ngrid, nfeat) - x_batchs, y_batchs = [], [] - for i in range(int(1 / offset)): - start = int(i * offset * seq_len) - idx = np.arange(start, y.shape[0]+1, seq_len) - split_x = np.split(X, indices_or_sections=idx, - axis=0) # (seq_len,ngrid,nfeat) - split_y = np.split(y, indices_or_sections=idx, axis=0) - # add all but the first and last batch since they will be smaller - for s in split_x: - if s.shape[0] == seq_len: - x_batchs.append(s) - for s in split_y: - if s.shape[0] == seq_len: - y_batchs.append(s) - x_batchs = np.concatenate(x_batchs, axis=1) - # (seq_len,ngrid*nyears*1/offset,nfeat) - y_batchs = np.concatenate(y_batchs, axis=1) - return np.transpose(x_batchs, (1, 0, 2)), np.transpose(y_batchs, (1, 0, 2)) + return (static-mean)/std \ No newline at end of file diff --git a/src/data_gen.py b/src/data_gen.py index e266548..94cabdf 100644 --- a/src/data_gen.py +++ b/src/data_gen.py @@ -27,19 +27,19 @@ def load_train_data(cfg, x, y, aux, scaler): return x, y, aux, mean, std -def load_test_data(cfg, x, y, aux, scaler): +def load_test_data(cfg, x, y, aux, scaler, stride=1): ngrid, nt, _ = x.shape - n = (nt-cfg["seq_len"]+1) + n = (nt-cfg["seq_len"]+1)//stride x_new = np.zeros((ngrid, n, cfg["seq_len"], cfg["num_feat"]))*np.nan y_new = np.zeros((ngrid, n, y.shape[-1]))*np.nan aux_new = np.zeros((ngrid, n))*np.nan - mean, std = np.array(scaler["y_mean"]), np.array(scaler["y_std"]) #(1, ngrid, nout) + mean, std = np.array(scaler["y_mean"]), np.array(scaler["y_std"]) mean = np.transpose(mean, (1,0,2)) std = np.transpose(std, (1,0,2)) for i in range(n): - x_new[:,i] = x[:,i:i+cfg["seq_len"]] - y_new[:,i] = y[:,i+cfg["seq_len"]-1] - aux_new[:,i] = aux[:,i+cfg["seq_len"]-1] + x_new[:,i] = x[:,i*stride:i*stride+cfg["seq_len"]] + y_new[:,i] = y[:,i*stride+cfg["seq_len"]-1] + aux_new[:,i] = aux[:,i*stride+cfg["seq_len"]-1] return x_new, y_new, aux_new, np.tile(mean, (1, n, 1)), np.tile(std, (1,n,1)) diff --git a/src/eval.py b/src/eval.py index f63a3a7..6cba93b 100644 --- a/src/eval.py +++ b/src/eval.py @@ -1,8 +1,33 @@ import numpy as np + +from model import HardMTLModel_v1, MTLModel_v1, MTLModel_v2, STModel from utils import r2_score -from model import VanillaLSTM, MTLLSTM, MTLHardLSTM_v1, MTLHardLSTM_v2 +def batcher(x_test, y_test, aux_test, seq_len): + n_t, n_feat = x_test.shape + _, n_out = y_test.shape + n = (n_t-seq_len+1) + x_new = np.zeros((n, seq_len, n_feat))*np.nan + y_new = np.zeros((n, n_out))*np.nan + aux_new = np.zeros((n, ))*np.nan + for i in range(n): + x_new[i] = x_test[i:i+seq_len] + y_new[i] = y_test[i+seq_len-1] + aux_new[i] = aux_test[i+seq_len-1] + return x_new, y_new, aux_new + +def single_batcher(x_test, y_test, seq_len): + n_t, n_feat = x_test.shape + _, n_out = y_test.shape + n = (n_t-seq_len+1) + x_new = np.zeros((n, seq_len, n_feat))*np.nan + y_new = np.zeros((n, n_out))*np.nan + for i in range(n): + x_new[i] = x_test[i:i+seq_len] + y_new[i] = y_test[i+seq_len-1] + return x_new, y_new + def eval_single(x, y, scaler, cfg): # get mean, std (1, ngrid, 6) mean, std = np.array(scaler["y_mean"]), np.array(scaler["y_std"]) @@ -17,24 +42,31 @@ def eval_single(x, y, scaler, cfg): # for each random seed for m in range(cfg["num_repeat"]): # load model - model = VanillaLSTM(cfg) + model = STModel(cfg) model.load_weights(save_folder+str(m)+'/') tmp = [] + y_true = [] # for each grid for i in range(x.shape[0]): - pred = model(x[i]) - pred = pred*std[:, i, j]+mean[:, i, j] + x_new, y_new = single_batcher(x[i], y[i], cfg["seq_len"]) + pred = model(x_new) + pred = pred*std[:, i, j]+mean[:, i, j] #(nsample,2,1) tmp.append(pred) - y_pred_seed.append(np.stack(tmp, axis=0)) # (ngrid, nsample, 1) + y_true.append(y_new) + y_pred_seed.append(np.stack(tmp, axis=0)) # (ngrid, nsample, 2, 1) + + # cal ensemble mean of diff exps (ngrid, nsample, 1) y_pred_seed = np.concatenate(y_pred_seed, axis=-1) - y_pred_seed = np.nanmean(y_pred_seed, axis=-1, keepdims=True) + y_pred_seed = np.nanmean(y_pred_seed, axis=-1, keepdims=True) #(ngrid, nsample, 2,1) + y_true = np.stack(y_true, axis=0) #(ngrid, nsample, 1) + print(y_true.shape, y_pred_seed.shape) # log r2_ens = [] for i in range(y_pred_seed.shape[0]): - r2_ens.append(r2_score(y[i, :, j], y_pred_seed[i, :, 0])) + r2_ens.append(r2_score(y_true[i, :, j], y_pred_seed[i, :, 1, 0])) print('\033[1;31m%s\033[0m' % "Var {} Median NSE {:.3f}".format(j, np.nanmedian(r2_ens))) @@ -42,6 +74,9 @@ def eval_single(x, y, scaler, cfg): return np.concatenate(y_pred_ens, axis=-1) + + + def eval_multi(x, y, aux, scaler, cfg): #FIXME: still have some bugs # get mean, std @@ -54,35 +89,44 @@ def eval_multi(x, y, aux, scaler, cfg): # for each random seed for m in range(cfg["num_repeat"]): # load model - if cfg["model_name"] in ['multi_tasks', 'soft_multi_tasks']: - model = MTLLSTM(cfg) - elif cfg["model_name"] in ['hard_multi_tasks_v2', 'hard_multi_tasks_v3']: - model = MTLHardLSTM_v2(cfg) - elif cfg["model_name"] == 'hard_multi_tasks_v1': - model = MTLHardLSTM_v1(cfg) + if cfg["model_name"] in ['multi_tasks_v1', 'soft_multi_tasks']: + model = MTLModel_v1(cfg) + elif cfg["model_name"] in ["multi_tasks_v2"]: + model = MTLModel_v2(cfg) + elif cfg["model_name"] in ['hard_multi_tasks_v1']: + model = HardMTLModel_v1(cfg) model.load_weights(save_folder+str(m)+'/') tmp = [] + y_true = [] # for each grid for i in range(x.shape[0]): - if cfg["model_name"] in ['single_task', 'multi_tasks', 'soft_multi_tasks']: - pred = model(x[i], training=False) + print(i) + # generate batch + x_new, y_new, aux_new = batcher(x[i], y[i], aux[i], cfg["seq_len"]) + + if cfg["model_name"] in ['single_task', 'multi_tasks_v1', 'soft_multi_tasks']: + pred = model(x_new, training=False) else: - pred = model(x[i], aux[i], mean[0,i], std[0,i], training=False) + pred = model(x_new, aux_new, mean[0,i], std[0,i], training=False) pred = pred*std[:, i]+mean[:, i] tmp.append(pred) + y_true.append(y_new) y_pred_ens.append(np.stack(tmp, axis=0)) # (ngrid, nsample, 6) # cal ensemble mean of diff exps (ngrid, nsample, 6) y_pred_ens = np.stack(y_pred_ens, axis=-1) y_pred_ens = np.nanmean(y_pred_ens, axis=-1) + y_true = np.stack(y_true, axis=0) # log - r2_ens = [] + print(y_pred_ens.shape) + print(y_true.shape) for j in range(cfg["num_out"]): + r2_ens = [] for i in range(y_pred_ens.shape[0]): - r2_ens.append(r2_score(y[i, :, j], y_pred_ens[i, :, j])) + r2_ens.append(r2_score(y_true[i, :, j], y_pred_ens[i, :, 1, j])) print('\033[1;31m%s\033[0m' % "Var {} Median NSE {:.3f}".format(j, np.nanmedian(r2_ens))) - return y_pred_ens + return y_pred_ens, y_true diff --git a/src/figure1.pdf b/src/figure1.pdf deleted file mode 100644 index d1784c91eed115c7555e0d4dea25d604067483cd..0000000000000000000000000000000000000000 GIT binary patch literal 0 HcmV?d00001 literal 9185 zcmb_?2{_bG_;`-QxTE`)L$-j9mJXPo>`#Xj^}PX2jqj!A_fj9ShZN1+ zNiH4~HweEt+{(zwlk7}~lo*r$WCyie0Cft5O813O$VCn`0XJj{Xdo_r-$i!uAZbwj z0S_<+7LQaQKo~4e0T0;&21zu4yd=IU*_V2N<_z=!3;aVLu!@!`1D<$EdpaIoKwd~u z#|tP#i|kBw0s1o}Q{3q85Nc7WfNftojZE@J_@BRPyECTyHecniFl!c^nu-9(yN@uPQuQ7ij#ve8=J0k5%B`Sk2cl zqaEBIJ9ApAxUus`(+e+(9lqB(9b3;kGrP?__$=GH5`&px@@(_ALT1%>lZA)5Yg%i{ zcl2$kB6Ww~(~Xd|)|N=*c#!9Aaj#i>?@cq8lNpCHL+svSJ_kGATUil0clV*gCe3jJ zT~SR}q~Y3PgrJ8%Z(SvCRbQ#tLI%oAG7&bg;-?=7Tx4D@B;oif-FTW+~Jbjas zp8p|zo;VjWcVYh9Ow8Cb-Wo?$Zt-vfXDBC8+RV*rtn5?hijIsp+&%ZE`rg{mZCpiW z9bNsmGs`|js7ni8+?dYH_OiCx4ue#viA$-*u-W-nSKrQ%oISHK_Mj5#dOO$2noS~C z1T7gIo~HM7o$nTk`LeLIlU`xuIf6u5LOEjA&twGewqmz`bby8B{oDU!Lf!PK0c^?-`VDrx1u(1p|udS6pNY%R}fAnjJ~*z9~%VMeU+1jn6yXA*9< z)c0v{oQ{tUvvoaqTbw4NV*aUbj~)hJFl0#MqUvplKVR6p_w?)Su>oPF4rGD@*OuV; z3TEDe66fOEA|5vkT}UBZd~@(&!bSTh=lTWdRKdO@EsY5){Mf~(UY!arbCA}5u=zDp00dq*K>Av)cFz~UW%Tr zT{oqUKxHJBRhQ|e+aDIk&(F8@s7r4a*r$Io*TGaxYx)>JN36n~vx6yf+D*kzdprA5 zk7p6HQO_# z$+6&N&b*p3Qr#>=Ok#=ssSdaBYa!AW+)=O-~lmRg_xdW z9<^<&$H~_x$D-UEUlwoAo?N&eVJG%@*vkBfn`p|ej**R-#!db1*XLz2Q#?ez>dQ}F zEO9c9zI*)6&itHF5oZ=-6YQ#NQC0IA*8HOv1cKtL*aNig@r7LC?{Cgn|CTOn61<8k zOKrGl@yY$=oDNU6=I;Hy@tnzNhQTi-O3o6A92L!0Eu#ii@j=Hg80%j|5ohl!cKNo{ zx_L>*>apKrs@im{(@hsuN?o_B?oB+VMA^uT+?^+NB_5+1L{l}2RwC^--e4K! zru_9wS~qsYc#m+b=q*>T=XSf5S_0OGdrt8gXCtCEGe6Vc>}pegsBh=D&DUn+@Ken1 zRj;MpqF&u{Z-4KTTuIR^%V^oPd|l7ZzdARk{Xi7a99S9t_{!09pJlVe?65->LL(-* z1T@Z$oj4VvDPSq7DpR~=@U#C#^dljIS4CsdE|$-Vw_jFaJ@eL?chxIfq01aH!y9>C zL~YD4q$Bh#(zzEN7DRokZJ*t5b5qJ|WaaY^n?y~ekdD>KyvGQxl8MEq+Y$ql4wmEx zKYOP`_BJ&46v=ujRT#q7x}{>(u;(2HSTfZgpU@>*mPfy{N(YUDXsA_R@E?;a?v2g(w-n~T>_{t9(M{HI9#?x%M_hTh$k^2k|5b1E z^QS{>G*#Pp%rpE{yj7pBM_q<+5XXUJiy()V-R>Lkaku)fcwQxXjqfr;XbK=2JA)-J z-+S|RB3zb5&Z}JXqj64Zoo%YJq~+_Klh)+~{m$4^0Z!Im24eXd2Hp@`3o702R`ztC z(Q#4QW^iNQg&kiBp_;g_?73f*?p8CkMMRsPXtnm*=q=7iNKx13u8WyK!wvoe-E8qO0~^Zk zq$T`6WA-KLp4&IaWc$I3bT=gtl|IA2!|)<`Qg|PeZ#q6G zt))fZ01PXSP4)a^IdW<7!!;`MiQvW~yGTNK`!0sI6TfP+{8L z?f}oEY%-y=BldHE*VCrUQoJZTRNCi}_d#u4UC~&Q*8ULg7d)QB#bOW6oDE<#4Os1T z6D1?gimbtZ5O5p}&wFm+ldTZ?p^cs5@HFg}5ZSCYeEe0e*zU_?4ZP>_hkI<5pY_^j z4UP>tCZya+#-}~V`no2tGemA)S^Aq>@Usua^)l#ps%oiz;zPFA-$J(Q71v7kM%yUSmU30TwfPk59DTBocz zZ;DSH`zjs4ar)WR=*HlX&-G8~-|G3yUE1tbamT}vC!Hd2G_RHcaMQCfY) z6k2Rr|NC;F1tJvRd}J_DO2c z_8exniyG1bZPP2xW-PF=BEpD=j4y=u^ekE~?jwf4uU{i}t0tv^_2 zIq*t6;k#uV*t5_)KA&E=`u^k?i7UeTl37J}ux~Y`D0xp9SG?IRzsk$uuVPlphAVQ^ zl?}XU!tF5UOlvICK8r7QKQbCotHjd1zJ4`(^8WLUiEF#?gAZ zjobD|J-R79c%I(K+>JcNF35IA+R)cCGPyBjev4Xw{E$;cWx(e7C1s^U5)l{1xqzE@s4=%s#fkk*J1&4jV|p@m=fK{u%Kr9x|Wc z+s5~!^7nfO;|L9(wb2ADVd*l$IO)TUw%`Q6ga}^|{^d}PA>!bp{Erj4is?>Dof>=B z;2OJ(r=QNX2AJ;YsrX9Wr)ioeR}qP-M}{<=DI3^n&%ggfW2r-dt@7ny1zVfcfQa%_ z(Ne0nx|2Uh+_I-D+#rbRj&q_f$Z6$1y7DA*MkDFOsx!MkWpdsRAbp2>M15dXZh9nNQ1G0xtEO7t?$1y`7k_Qo~!XmifKf z3aya@^J`~LyKNGv66|qv+OT)UQ2FNg;uSq(XKC>lMJ_7)Kkoa z9On|Eo(I!c%H?6l6UHC!DSpgp6&b)kMV!-{8Eo^C2-`+!8;=|i$JC+uTlVfG?#)bo zqx-7M`T$>nXxZ%+`VsNOHXp;1;X@B5jppfHHjfTe(Acxo26m}`D(5NCxlvn-R=p4{ z8<-`3UHEK;4(1+dHDB9w%>x$UZ?uZl*qnWmC03H>ZaQg2pY-xoxy#yS__2=IRwHoV z>*IyC%pwd%PU2|ioiNW~HNWD@&px-(cVLFHy#HeS<+slN>!O9m5`VGjE))HeEJ&{M zLBTI-5)SbdIhv`OgA46VJST{(9EEZ5NhwyGyr|>i@+vd2623~`L=I(jlv@YJHFhkd z3`&}r%$VjXIBD2Nh_q(%eNxHX$8UY(?if-2VybWDD)AGySuvSR&LR5fvfQ_c5(2xq z4s77Dxpe>bKBt;KnWMN?xkphd7PE z*Soe$J`_V-LT}%?DeQTVa5iR6>G-r4IFB<&Ut11!uivqQy@9|Yb z?5qps>Kq6}r;&7g_|!u5YIN}4L7orCS90D8&=FsmDD5C~dUntwR)3Fd_a**-1s*r)Gr0jzoG(msI)B}${3n-a_+I``YZs0Fg-x7E zxpy6sYXrA)9TGYu_2>=QlAPmlhWRg;AKy;a<+r6{*$ zdQ|L0IZ|aQm>NyCG%oot`~62Il8~Hsy2ht3j9etWsC>^FJ$c#m^GDu8+@JIB=BXZKySyns63j zRe^N)mSbDESG+FaPWcOI_;KrBNMnF2FApv_5L{NG*t=R~Z>wd0GHUsJ<;tl_5L^WA z3K}B!zL8b5e!8>rxl~tnO~0-9wA2la-k8y|4^8W)*A}m7-CC&I%2&D-!D=F)VwI{W zDEYi~VrW~gk3sqoOZL(=A{EEdT$brYuJ{{pahkpFNT$9eMat zMcPZ5YnL)7RIF2*U&imuek{6c_YhNEQ=NW>z5Z-QYKnnYT=HNsUpT~uS} z2tLv~Yo+v{Shm!|+BMs*H!xSYo_(=rXzzWoaM>tJ&f;7#uW?&}#7m;r-UxLddz+DF zPMNQ-uYCU2iG1|Kz#ClQ{fJSMc>5Zk4@r*>vOjU37g^i%%2_ZjR5#}aKeq>=I-)ww zfl%l0S%An>9pQZ8tUQ;qyL3+A>B#EGX~o&3`UWwZ$;OB(Tord-#~POXbkl_GHs*fK;U~;*j?iaI5E%{fNWA!%iY_e0xv@Sl&A#7%segat(H%9}Ih+l%w{8>;JP?f$vl85xsa1`x zs~r+fYWZ;8)%H-5M|F(<@$*~Swl-L;G!^L~jGmm%4A@%47I8N<=_EqbB695^^|1`y z=#Z)P=2p;?(GB`OQIW>6p?*g?^|**Z@_r`b6PV2?fuT+j?QZoWQgJE7x;ylVmotxV zWX%5;rBt;|?wzfTwewR|kCDiD6l( zCKmpsn(^F^Pz@5FdD3cxI1Iihuc?ZW5)e&z@REl+^RT1X#i;Z6Ugo!1JG=!~@$?nf zIFua`G+4JUZkj32aOlXpg~G9$mF?}~qNAbpe=!dH=KC+kLI1Wwx0>kFSk&3eTenOc zcg=a}oV+Exf1#bdX%!pyDK)xHUq)0qByU+U5S#B_^AdwP`f_FXp)Kv+R{d(`kiN4q zF4J!;a6;_A{V5y|V}C)Jj*&0kE`Qqdez@3ypz5~@h|;%4z{d) z_9p02;i>DQfr^eJ?8_-h7QJ92}gTV{zPT8}Sz%!Vlm7q(juN$5@;Rk%vWs zYh+q*UTss9qQ6nr=122YkD{Jfh~2XbD@l+@9XE~7l^8Hn4NO%yf2X_sp(l1Qsbs5f z=q~|O5yEszIi11mc(Er zu)mqvCEXyQxWATqW1^VN^27HG4MT^piKd< z^Kf!OSs9_I<>Bf|22cO+TfZ$x3P6fJWSR%n1;P>_#XvHR3Mo=4;3=Pe5Y*`IG%~2W zg4_W3@8JjPz8?M%4vZpuxVeLj2E_v)gLJ}T{^SQG5DOG_JZQdjm>a;kBFwJFVlrYY z3W-KRTd{Zygi-?MFBBSsfm0wbiWb?|ndad`r_vx4Bgq1$0nNKebQ+)w5*Z3eBpj*1 z-;Wv?YDoG{35~-+in~ePfYCoxgHSX(;6!KSd0=iaOG%$aOpU0~0m1@n zg2WU=zz`rj29F@(kdP7q389s65Qd0k01og8_rv%YJZiDTD8cyu7a*4UF#s0^c#b8& zVNf7%1}lJo1pp=@5+Efd@EwDOaDccH3^+7|CgKnPAUyDS91ctZyjBS;6(s^-014`F z3HM_Ffb5(S7!S`$BoINl*anvjVK7PvxP+O2g@OSO91b3f1A;Hk0dN6bMje9%iXbom zr$l6wL?ys4Xvd?_K)^-1SU?vCmK25vkHbPZP!fQ$;F1Ao7!S`-0$jrM8AgK@2zUUB zMZh$1U<^YkOU%I~!_#0&SojMPRuRlLEE-$_VAusFw+ILv5Roww*6+^>U~q8?G>{nX zfw=;lFS)^zxt2;`GQ<)i46d+BC|GO;`>-GU!YV8i!G@AN&doT?4;>ePS6f`~mifC1CiJKqLU0z12Vvj}u zFDl*3!wK@kDBu(@@{l{7?&G^nQPKNHmjadM1_pt1gUbPDkk(&5%E!eOaw0kJ1Ge}+ zgX}WB_|wR)2(V!j5Ma4~|A5^R2YkyF`q2g-Mc@xaS!yF-KoI@C4FyL1(S`@&!=Ghi z(fB{-LxYv~J01oJJoeYINF@Fb8b}OS8NcCS;MMh88xsBJd>Ax{3cumu;RC|&ZA$P$ zfImOw01NUDUXXEK7D@O64~xeAAsZGC-VJ}Ffh8dS#3Q2rkPQc4 z^?w_SQvxTr-`c?D_QzPT87-qrr;$9o$TY@par7y!R1n`5_W)BW618) & (lat2<28))[0] + lon_idx = np.where((lon2>108) & (lon2<119))[0] + climate_zone = np.array(f.Band1) +with xr.open_dataset(PATH+'ancillary/DEM_0p1.nc') as f: + dem = np.array(f.elv) +with xr.open_dataset(PATH+'ancillary/kosugi_0p1.nc') as f: + kosugi = np.array(f.kosugi) +with xr.open_dataset(PATH+'ancillary/LC_0p1.nc') as f: + land_cover = np.array(f["class"]) +with xr.open_dataset(PATH+'ancillary/soilgrid_0p1.nc') as f: + soilgrid = np.array(f.soilgrids) +tmp = np.stack([climate_zone, dem, land_cover, soilgrid[0], soilgrid[1], + soilgrid[7], soilgrid[8], soilgrid[14], soilgrid[15]], axis=0) +static = np.concatenate([tmp, kosugi], axis=0) +static = static[:,lat_idx, :][:,:,lon_idx] +static = np.transpose(static, (1, 2, 0)) # (lat, lon, feats) +static = np.flip(static, axis=0) +Nlat, Nlon, Nfeat = static.shape +static = static.reshape(-1, Nfeat)[idx, :] +print(static.shape) + +for i in range(Nfeat): + tmp = static[:,i] + tmp[np.isnan(tmp)] = np.nanmean(tmp) + static[:,i] = tmp + +print(np.isnan(static).any()) +np.save("gd_9km_ancillary.npy", static) # ------------------------------------------------------------------------------ # read CoLM hydrology & forcing variables # ------------------------------------------------------------------------------ +PATH = "/tera04/zhwei/cases/global_grid_ERA5_pft_GD/history/" # path contain raw data forcing = [] hydrology = [] for year in years: for month in months: print(year, month) - file_name = "global_MSWX_9km_hist_{year}-{month:02}.nc".format( + file_name = "global_grid_ERA5_pft_GD_hist_{year}-{month:02}.nc".format( year=year, month=month) - if os.path.exists(PATH+'CoLM/'+file_name): - with xr.open_dataset(PATH+'CoLM/'+file_name, decode_times=False) as f: + if os.path.exists(PATH+file_name): + with xr.open_dataset(PATH+file_name, decode_times=False) as f: # hydrology evpa = np.array(f.f_fevpa) # mm/s rnof = np.array(f.f_rnof) # mm/s h2osoi = np.array(f.f_h2osoi) # m3/m3 - swvl1, swvl2, swvl3, swvl4 = h2osoi_colm2era(h2osoi) # m3/m3 + wa = np.array(f.f_wa) #mm + ldew = np.array(f.f_ldew) #mm + scv = np.array(f.f_scv) #mm + xerr = np.array(f.f_xerr) # mm/s + wice_soisno = np.array(f.f_wice_soisno)[:,:,:,5:] #kg/m2 + wliq_soisno = np.array(f.f_wliq_soisno)[:,:,:,5:] #kg/m2 + swvl1, swvl2, swvl3, swvl4, zi = h2osoi_colm2era(wice_soisno, wliq_soisno) # m3/m3 # forcing prc = np.array(f.f_xy_prc) # mm/s convective @@ -79,65 +129,48 @@ fr = np.array(f.f_xy_frl) # W/m2 solarin = np.array(f.f_xy_solarin) # W/m2 - # turn all hydrology variables to mm/day - evpa = evpa*24*60*60 # mm - rnof = rnof*24*60*60 # mm - prc = prc*24*60*60 # mm - prl = prl*24*60*60 # mm - tmp = np.stack([prc, prl, t, q, pbot, fr, solarin, us, vs], axis=-1) - Nt, Nlat, Nlon, Nfeat = tmp.shape - tmp = tmp.reshape(Nt, -1, Nfeat)[:, idx, :] - forcing.append(tmp) - tmp = np.stack([swvl1, swvl2, swvl3, swvl4, evpa, rnof], axis=-1) - Nt, Nlat, Nlon, Nfeat = tmp.shape - tmp = tmp.reshape(Nt, -1, Nfeat)[:, idx, :] - hydrology.append(tmp) - + #p = swvl1*70+swvl2*210+swvl3*720+swvl4*(zi[9]-100)*10+wa+scv+ldew + # turn all hydrology variables to mm/hour + evpa = evpa*60*60 # mm + rnof = rnof*60*60 # mm + prc = prc*60*60 # mm + prl = prl*60*60 # mm + xerr = xerr*60*60 + + tmp = np.stack([prc, prl, t, q, pbot, fr, solarin, us, vs], axis=-1) + print(tmp.shape) + forcing.append(tmp) + tmp = np.stack([swvl1, swvl2, swvl3, swvl4, evpa, rnof], axis=-1) + hydrology.append(tmp) + +# check physical consistency +#m = np.diff(p, axis=0) +#d = m-(prc[1:]+prl[1:]-evpa[1:]-rnof[1:]+xerr[1:]) +#mean_phy = np.nanmean(np.abs(d), axis=0) +#mean_phy = mean_phy.reshape(-1,)[idx] +#mask_unphy = np.ones_like(mean_phy) +#mask_unphy[mean_phy>0.1] = 0 +#idx = np.delete(idx, np.where(mask_unphy==0)[0]) +print(len(forcing)) forcing = np.concatenate(forcing, axis=0) -np.save("guangdong_9km_forcing_{}.npy".format(mode), forcing) +Nt, Nlat, Nlon, Nfeat = forcing.shape +forcing = forcing.reshape(Nt, -1, Nfeat)[:, idx, :] +np.save("gd_9km_forcing_{}.npy".format(mode), forcing) print(forcing.shape) del forcing hydrology = np.concatenate(hydrology, axis=0) -np.save("guangdong_9km_hydrology_{}.npy".format(mode), hydrology) +Nt, Nlat, Nlon, Nfeat = hydrology.shape +hydrology = hydrology.reshape(Nt, -1, Nfeat)[:, idx, :] +np.save("gd_9km_hydrology_{}.npy".format(mode), hydrology) print(hydrology.shape) del hydrology -# ------------------------------------------------------------------------------ -# read ancillary variables -# ------------------------------------------------------------------------------ -with xr.open_dataset(PATH+'ancillary/Beck_KG_V1_present_0p1.nc') as f: - climate_zone = np.array(f.Band1) -with xr.open_dataset(PATH+'ancillary/DEM_0p1.nc') as f: - dem = np.array(f.elv) -with xr.open_dataset(PATH+'ancillary/kosugi_0p1.nc') as f: - kosugi = np.array(f.kosugi) -with xr.open_dataset(PATH+'ancillary/LC_0p1.nc') as f: - land_cover = np.array(f["class"]) -with xr.open_dataset(PATH+'ancillary/soilgrid_0p1.nc') as f: - soilgrid = np.array(f.soilgrids) - -tmp = np.stack([climate_zone, dem, land_cover, soilgrid[0], soilgrid[1], - soilgrid[7], soilgrid[8], soilgrid[14], soilgrid[15]], axis=0) -static = np.concatenate([tmp, kosugi], axis=0) -static = np.transpose(static, (1, 2, 0)) # (lat, lon, feats) -static = np.flip(static, axis=0) -Nlat, Nlon, Nfeat = static.shape -static = static.reshape(-1, Nfeat)[idx, :] -print(static.shape) - -for i in range(Nfeat): - tmp = static[:,i] - tmp[np.isnan(tmp)] = np.nanmean(tmp) - static[:,i] = tmp - -np.save("guangdong_9km_ancillary.npy", static) - - # ------------------------------------------------------------------------------ # move to input path # ------------------------------------------------------------------------------ -os.system("mv {} {}".format('guangdong*npy', inputs_path)) +np.save("gd_9km_shapefile.npy", idx) +os.system("mv {} {}".format('gd_9km*npy', inputs_path)) diff --git a/src/metric.py b/src/metric.py new file mode 100644 index 0000000..24c6371 --- /dev/null +++ b/src/metric.py @@ -0,0 +1,37 @@ +import tensorflow as tf +from tensorflow.keras.metrics import Metric + + +class NSEMetrics(Metric): + def __init__(self, cfg): + super().__init__() + self.total = self.add_weight('total', initializer='zeros') + self.count = self.add_weight('count', initializer='zeros') + self.model = cfg["model_name"] + if self.model == 'single_task': self.num_out = 1 + else: self.num_out = 6 + + def update_state(self, y_true, y_pred): + y_true = tf.cast(y_true, dtype=tf.float64) + y_pred = tf.cast(y_pred, dtype=tf.float64) + + self.metrics_ = [] + for i in range(self.num_out): + a,b = y_true[:,i], y_pred[:,i] + mask = a == a + a, b = a[mask], b[mask] + unexplained_error = tf.reduce_sum(tf.square(a-b)) + total_error = tf.reduce_sum(tf.square(a - tf.reduce_mean(a))) + r2 = 1. - tf.divide(unexplained_error, total_error) + self.metrics_.append(r2) + + def result(self): + if self.model == 'single_task': + return {"loss": self.metrics_[0]} + else: + return {"0": self.metrics_[0], + "1": self.metrics_[1], + "2": self.metrics_[2], + "3": self.metrics_[3], + "4": self.metrics_[4], + "5": self.metrics_[5]} diff --git a/src/model.py b/src/model.py index b0f59b8..a6c23b3 100644 --- a/src/model.py +++ b/src/model.py @@ -1,142 +1,145 @@ +""" +All model structure + +<<<<<<<< HEAD +Author: Lu Li +11/1/2022 - V1.0 +12/9/2022 - V2.0 +""" + import tensorflow as tf -from tensorflow.keras import Model -from tensorflow.keras.layers import LSTM, Dense, Dropout, Layer from tensorflow import math +from tensorflow.keras import Model +from tensorflow.keras.layers import LSTM, Dense, Dropout +from utils import make_CoLM_soil_depth +from layers import WeightedMultiLossLayer, MassConsLayer + + + +class STModel(Model): + """single task model""" -class VanillaLSTM(Model): def __init__(self, cfg): super().__init__() - self.lstm = LSTM(cfg["hidden_size"], return_sequences=False) + self.lstm = LSTM(cfg["hidden_size"], return_sequences=True) self.drop = Dropout(cfg["dropout_rate"]) self.dense = Dense(1) def call(self, inputs): x = self.lstm(inputs) x = self.drop(x) - x = self.dense(x) + # we only predict the last two steps + x = self.dense(x[:,-2:]) return x -class MTLLSTM(Model): - """LSTM with multi-tasks""" +class MTLModel_v1(Model): + """multitasks model with average loss""" def __init__(self, cfg): super().__init__() self.num_out = cfg["num_out"] - self.shared_layer = LSTM(cfg["hidden_size"], - return_sequences=False, - name='shared_layer') + self.shared_layer = LSTM(cfg["hidden_size"], return_sequences=True) self.drop = Dropout(cfg["dropout_rate"]) self.head_layers = [] for i in range(cfg["num_out"]): - self.head_layers.append(Dense(1, name='head_layer_'+str(i+1))) + self.head_layers.append(Dense(1, name='head_layer_'+str(i+1))) def call(self, inputs): x = self.shared_layer(inputs) # shared layer x = self.drop(x) pred = [] for i in range(self.num_out): # each heads - pred.append(self.head_layers[i](x)) + pred.append(self.head_layers[i](x[:,-2:])) pred = tf.concat(pred, axis=-1) return pred -class MassConsLayer(Layer): - """Mass conserve layer""" +class MTLModel_v2(Model): + """multitasks model with adaptive loss""" def __init__(self, cfg): super().__init__() - self.idx = cfg["resid_idx"] self.num_out = cfg["num_out"] + self.shared_layer = LSTM(cfg["hidden_size"], return_sequences=True) + self.drop = Dropout(cfg["dropout_rate"]) + self.head_layers = [] + for i in range(cfg["num_out"]): + self.head_layers.append(Dense(1, name='head_layer_'+str(i+1))) + self.loss_layer = WeightedMultiLossLayer(cfg) - def _fill_matrix(self, x, idx, res=None): - if res is not None: - empty = res - else: - empty = x[:, 0:1] - if idx == 0: - x = tf.concat([empty, x], axis=-1) - elif idx == self.num_out: - x = tf.concat([x, empty], axis=-1) - else: - x = tf.concat([x[:, :idx], empty, x[:, idx+1:]], axis=-1) - return x - - def _slice_matrix(self, x, idx): - if idx == 0: - x = x[:, 1:] - elif idx == self.num_out: - x = x[:, :-1] - else: - x = tf.concat([x[:, :idx], x[:, idx+1:]], axis=-1) - return x + def call(self, inputs, y_true=None): + x = self.shared_layer(inputs) # shared layer + x = self.drop(x) + pred = [] + for i in range(self.num_out): # each heads + pred.append(self.head_layers[i](x[:,-2:])) + pred = tf.concat(pred, axis=-1) - def call(self, inputs, aux, mean, std): - """ - Args - ---- - inputs: Directly outputs of multi-task models. Notably, - it's z-score normalized value, and if we want - to predict N vars, it only contains (N-1) vars. - """ - # init - inputs = tf.cast(inputs, 'float32') - aux = tf.cast(aux, 'float32') - mean = tf.cast(mean, 'float32') - std = tf.cast(std, 'float32') - soil_depth = [70, 210, 720, 1864.6] # mm - # Concat empty tensor to inputs based on idx (batch, nout) - inputs = self._fill_matrix(inputs, self.idx) - # reverse normalized forecasts - inputs = math.multiply(inputs, std) + mean - # Transform soil moisture in unit mm - swvl = math.multiply(inputs[:, :4], soil_depth) - # unnormalized all mm/day - inputs = tf.concat([swvl, inputs[:, 4:]], axis=-1) - # Calculate residual outputs - inputs = self._slice_matrix(inputs, self.idx) - mass_resid = aux-math.reduce_sum(inputs, axis=-1) - mass_resid = mass_resid[:, tf.newaxis] - # Output - inputs = self._fill_matrix(inputs, self.idx, mass_resid) - swvl = tf.divide(inputs[:, :4], soil_depth) # mm3/mm3 - inputs = tf.concat([swvl, inputs[:, 4:]], axis=-1) - inputs = math.divide(inputs-mean, std) - return inputs - - def compute_output_shape(self, input_shape): - return (input_shape[0][0], input_shape[0][1] + 1) + if y_true is not None: # train mode + #FIXME: Only cal loss on last step + loss_sum = self.loss_layer(y_true, pred) + return pred, loss_sum + else: # inference mode + return pred -class MTLHardLSTM_v2(Model): - """LSTM with hard physical constrain through residual layer.""" +class HardMTLModel_v1(Model): + """multitasks model with hard physical constrain through redistribute layer.""" def __init__(self, cfg): super().__init__() self.num_out = cfg["num_out"] - self.shared_layer = LSTM(cfg["hidden_size"], - return_sequences=False, - name='shared_layer') + self.shared_layer = LSTM(cfg["hidden_size"],return_sequences=True) self.drop = Dropout(cfg["dropout_rate"]) self.head_layers = [] - for i in range(cfg["num_out"]-1): + for i in range(cfg["num_out"]): self.head_layers.append(Dense(1, name='head_layer_'+str(i+1))) - self.resid_layer = MassConsLayer(cfg) def call(self, inputs, aux, mean, std): - x = self.shared_layer(inputs) + x = self.shared_layer(inputs) # shared layer x = self.drop(x) pred = [] - for i in range(self.num_out-1): - pred.append(self.head_layers[i](x)) - pred = tf.concat(pred, axis=-1) - pred = self.resid_layer(pred, aux, mean, std) - return pred + for i in range(self.num_out): # each heads + pred.append(self.head_layers[i](x[:,-2:])) + pred = tf.concat(pred, axis=-1) #(b, 2, 6) + # ------------------------- + # redistribute water budget + # ------------------------- + depth, zi = make_CoLM_soil_depth() # m/cm + soil_depth = [70, 210, 720, 10*(zi[9]-100)] # mm + pred_prev_save, pred_now = pred[:,0], pred[:,1] #(b,6) + pred_prev = math.multiply(pred_prev_save, std) + mean + pred_now = math.multiply(pred_now, std) + mean + print(tf.shape(pred_prev)) + + # cal water budget + swvl_prev = math.multiply(pred_prev[:,:4], soil_depth) # (b,4) + swvl_now = math.multiply(pred_now[:,:4], soil_depth) # (b,4) + delta_swvl = math.reduce_sum(swvl_now-swvl_prev, axis=-1) #(b,) + w_b = aux-delta_swvl-pred_now[:,-2]-pred_now[:,-1] #(b,) + + # cal ratio and distribute + pred_new = [] + w_a = math.reduce_sum(swvl_now, axis=-1) #(b) + for i in range(4): + ratio = math.divide(swvl_now[:,i], w_a) + water_add = math.multiply(w_b, ratio) + pred_new.append((water_add+swvl_now[:,i])/soil_depth[i]) + pred_new.append(pred_now[:,-2]) + pred_new.append(pred_now[:,-1]) + pred_new = tf.stack(pred_new, axis=-1) #(b,6) + pred_new = math.divide(pred_new-mean, std) + pred = tf.stack([pred_prev_save, pred_new], axis=1) #(b,2,6) + print(tf.shape(pred)) + return pred #(b,2,6) + + +class MTLHardLSTM_v2(Model): + """LSTM with hard physical constrain through residual layer.""" -class MTLHardLSTM_v1(Model): - """LSTM with hard physical constrain through redistribute layer.""" def __init__(self, cfg): super().__init__() self.num_out = cfg["num_out"] @@ -145,35 +148,20 @@ def __init__(self, cfg): name='shared_layer') self.drop = Dropout(cfg["dropout_rate"]) self.head_layers = [] - for i in range(cfg["num_out"]): + for i in range(cfg["num_out"]-1): self.head_layers.append(Dense(1, name='head_layer_'+str(i+1))) + self.resid_layer = MassConsLayer(cfg) def call(self, inputs, aux, mean, std): - x = self.shared_layer(inputs) # shared layer + x = self.shared_layer(inputs) x = self.drop(x) pred = [] - for i in range(self.num_out): # each heads + for i in range(self.num_out-1): pred.append(self.head_layers[i](x)) pred = tf.concat(pred, axis=-1) - - # redistribute water budget - soil_depth = [70, 210, 720, 1864.6] # mm - # cal water budget - pred = math.multiply(pred, std) + mean - swvl = math.multiply(pred[:, :4], soil_depth) - pred = tf.concat([swvl, pred[:, 4:]], axis=-1) - w_b = aux-math.reduce_sum(pred, axis=-1) - # cal ratio - swvl_new = [] - water_all = math.reduce_sum(swvl, axis=-1) - for i in range(4): - ratio = math.divide(swvl[:,i], water_all) - water_add = math.multiply(w_b, ratio) - swvl_new.append((water_add+swvl[:,i])/soil_depth[i]) - swvl_new.append(pred[:,4]) - swvl_new.append(pred[:,5]) - pred = tf.stack(swvl_new, axis=-1) - pred = math.divide(pred-mean, std) + pred = self.resid_layer(pred, aux, mean, std) return pred + + diff --git a/src/plot.py b/src/plot.py index 15a57b3..c29191c 100644 --- a/src/plot.py +++ b/src/plot.py @@ -10,14 +10,17 @@ # Figure 1 # --------------------------------- path = cfg["outputs_path"]+'forecast/' -r2_single_task = np.load(path+"single_task_epoch_200/r2_single_task.npy") -r2_multi_tasks = np.load(path+"multi_tasks_epoch_200/r2_multi_tasks.npy") +r2_single_task = np.load(path+"single_task/r2_single_task.npy") +r2_multi_tasks = np.load(path+"multi_tasks/r2_multi_tasks.npy") +r2_hard_multi_tasks_v1 = np.load(path+"hard_multi_tasks_v1/r2_hard_multi_tasks_v1.npy") -phy_single_task = np.load(path+"single_task_epoch_200/phy_cons_single_task.npy") -phy_multi_tasks = np.load(path+"multi_tasks_epoch_200/phy_cons_multi_tasks.npy") -r2_st, r2_mt = [], [] -phy_st, phy_mt = [], [] +phy_single_task = np.load(path+"single_task/phy_cons_single_task.npy") +phy_multi_tasks = np.load(path+"multi_tasks/phy_cons_multi_tasks.npy") +phy_hard_multi_tasks_v1 = np.load(path+"hard_multi_tasks_v1/phy_cons_hard_multi_tasks_v1.npy") + +r2_st, r2_mt, r2_hard = [], [], [] +phy_st, phy_mt, phy_hard = [], [], [] for i in range(6): temp = r2_multi_tasks[:,:,i].reshape(-1, ) temp = np.delete(temp, np.isnan(temp)) @@ -27,14 +30,21 @@ temp = np.delete(temp, np.isnan(temp)) r2_st.append(temp) + temp = r2_hard_multi_tasks_v1[:,:,i].reshape(-1, ) + temp = np.delete(temp, np.isnan(temp)) + r2_hard.append(temp) + temp = phy_multi_tasks.reshape(-1, ) temp = np.delete(temp, np.isnan(temp)) phy_mt.append(temp) temp = phy_single_task.reshape(-1, ) temp = np.delete(temp, np.isnan(temp)) phy_st.append(temp) +temp = phy_hard_multi_tasks_v1.reshape(-1, ) +temp = np.delete(temp, np.isnan(temp)) +phy_hard.append(temp) -#figure +#figure 1 fig = plt.figure() ax = plt.subplot(211) ax.spines['left'].set_linewidth(2) @@ -48,6 +58,11 @@ positions=[0, 1, 2, 3, 4, 5], boxprops=dict(facecolor='lightblue', color='black')) ax.boxplot(r2_mt,notch=True, + patch_artist=True, + showfliers=False, + positions=[0.25, 1.25, 2.25, 3.25, 4.25, 5.25], + boxprops=dict(facecolor='red', color='black')) +ax.boxplot(r2_hard,notch=True, patch_artist=True, showfliers=False, positions=[0.5, 1.5, 2.5, 3.5, 4.5, 5.5], @@ -67,37 +82,44 @@ ax.boxplot(phy_mt,notch=True, patch_artist=True, showfliers=False, - positions=[0.5], + positions=[0.25], boxprops=dict(facecolor='red', color='black')) +ax.boxplot(phy_hard,notch=True, + patch_artist=True, + showfliers=False, + positions=[0.5], + boxprops=dict(facecolor='yellow', color='black')) plt.savefig('figure1.pdf') print('Figure 1 completed!') +#figure 2 +fig = plt.figure() +ax = plt.subplot(211) +ax.spines['left'].set_linewidth(2) +ax.spines['bottom'].set_linewidth(2) +ax.spines['right'].set_linewidth(2) +ax.spines['top'].set_linewidth(2) -""" -m = Basemap(projection='mill', - llcrnrlat=lat_min-1, urcrnrlat=lat_max+1, - llcrnrlon=lon_min-1, urcrnrlon=lon_max+1, - ax=ax) # mill projection -#m.drawcoastlines(linewidth=1) -x, y = m(grid_lon, grid_lat) -sc = m.pcolormesh(x, y, phy_cons_grid[:,:,0], cmap='jet') -m.readshapefile(cfg["inputs_path"]+"guangdong_shp/guangdong", "guangdong", drawbounds=True) -cbar_ax = fig.add_axes([0.25, 0.05, 0.5, 0.02]) -fig.colorbar(sc, cax=cbar_ax, orientation="horizontal",spacing='proportional') -plt.suptitle("Physical consistency of model") -plt.savefig("figure2.pdf") - +# load site and grid lon/lat +path = cfg["inputs_path"] +attr = np.load(path+"coord_gd_9km.npy") +lon_min, lat_min = np.nanmin(attr, axis=0) +lon_max, lat_max = np.nanmax(attr, axis=0) +ease_9km_grids = np.load(path+"coord_global_9km.npy") #(1800, 3600, 2) +lon, lat = ease_9km_grids[0,:,0], ease_9km_grids[:,0,1] +idx_lat = np.where((lat>=lat_min) & (lat<=lat_max)) +idx_lon = np.where((lon>=lon_min) & (lon<=lon_max)) +lon_gd = lon[idx_lon] +lat_gd = lat[idx_lat] +grid_lon, grid_lat = np.meshgrid(lon_gd, lat_gd) -# preliminary plot -# figure 1 (countour) -print(r2.shape, attr.shape, lat_gd.shape, lon_gd.shape) var_list = ["0-7cm soil moisture", "7-28cm soil moisture", "28-100cm soil moisture", "100-286.46cm soil moisture","evapotranspiration","total runoff"] text = ["(a)","(b)","(c)","(d)","(e)","(f)"] -fig = plt.figure(figsize=(8,10)) +fig = plt.figure() for i in range(cfg["num_out"]): - ax = plt.subplot(3, 2, i+1) + ax = plt.subplot(3, 6, i+1) ax.spines['left'].set_linewidth(1) ax.spines['bottom'].set_linewidth(1) ax.spines['right'].set_linewidth(1) @@ -109,22 +131,59 @@ ax=ax) # mill projection #m.drawcoastlines(linewidth=1) x, y = m(grid_lon, grid_lat) - sc = m.pcolormesh(x, y, r2_grid[:,:,i] ,vmin=0.6, vmax=1, cmap='jet') + sc = m.pcolormesh(x, y, r2_single_task[:,:,i] ,vmin=0.6, vmax=1, cmap='jet') m.readshapefile(cfg["inputs_path"]+"guangdong_shp/guangdong", "guangdong", drawbounds=True) - plt.title(text[i]+' '+var_list[i]) + #plt.title(text[i]+' '+var_list[i]) + + ax = plt.subplot(3, 6, i+1+6) + ax.spines['left'].set_linewidth(1) + ax.spines['bottom'].set_linewidth(1) + ax.spines['right'].set_linewidth(1) + ax.spines['top'].set_linewidth(1) + + m = Basemap(projection='mill', + llcrnrlat=lat_min-1, urcrnrlat=lat_max+1, + llcrnrlon=lon_min-1, urcrnrlon=lon_max+1, + ax=ax) # mill projection + #m.drawcoastlines(linewidth=1) + x, y = m(grid_lon, grid_lat) + sc = m.pcolormesh(x, y, r2_multi_tasks[:,:,i] ,vmin=0.6, vmax=1, cmap='jet') + m.readshapefile(cfg["inputs_path"]+"guangdong_shp/guangdong", "guangdong", drawbounds=True) + #plt.title(text[i]+' '+var_list[i]) + + ax = plt.subplot(3, 6, i+1+6+6) + ax.spines['left'].set_linewidth(1) + ax.spines['bottom'].set_linewidth(1) + ax.spines['right'].set_linewidth(1) + ax.spines['top'].set_linewidth(1) + + m = Basemap(projection='mill', + llcrnrlat=lat_min-1, urcrnrlat=lat_max+1, + llcrnrlon=lon_min-1, urcrnrlon=lon_max+1, + ax=ax) # mill projection + #m.drawcoastlines(linewidth=1) + x, y = m(grid_lon, grid_lat) + sc = m.pcolormesh(x, y, r2_hard_multi_tasks_v1[:,:,i] ,vmin=0.6, vmax=1, cmap='jet') + m.readshapefile(cfg["inputs_path"]+"guangdong_shp/guangdong", "guangdong", drawbounds=True) + #plt.title(text[i]+' '+var_list[i]) + + cbar_ax = fig.add_axes([0.25, 0.05, 0.5, 0.02]) fig.colorbar(sc, cax=cbar_ax, orientation="horizontal",spacing='proportional') -plt.suptitle("R2 of single task model") -plt.savefig("figure1.pdf") +plt.savefig("figure2.pdf") -# figure 3 (spatial mean) -fig = plt.figure(figsize=(10,18)) -for i in range(cfg["num_out"]): - ax = plt.subplot(6, 1, i+1) - plt.plot(np.nanmean(y_pred[:,:,i], axis=0)) - plt.plot(np.nanmean(y_test[:,:,i], axis=0)) - plt.title(text[i]+' '+var_list[i]) - plt.legend(['single_task','CoLM']) -plt.savefig("figure3.pdf") +""" +m = Basemap(projection='mill', + llcrnrlat=lat_min-1, urcrnrlat=lat_max+1, + llcrnrlon=lon_min-1, urcrnrlon=lon_max+1, + ax=ax) # mill projection +#m.drawcoastlines(linewidth=1) +x, y = m(grid_lon, grid_lat) +sc = m.pcolormesh(x, y, phy_cons_grid[:,:,0], cmap='jet') +m.readshapefile(cfg["inputs_path"]+"guangdong_shp/guangdong", "guangdong", drawbounds=True) +cbar_ax = fig.add_axes([0.25, 0.05, 0.5, 0.02]) +fig.colorbar(sc, cax=cbar_ax, orientation="horizontal",spacing='proportional') +plt.suptitle("Physical consistency of model") +plt.savefig("figure2.pdf") """ \ No newline at end of file diff --git a/src/postprocess.py b/src/postprocess.py index 7fb27cb..4810cbe 100644 --- a/src/postprocess.py +++ b/src/postprocess.py @@ -44,6 +44,8 @@ def postprocess(cfg): a = np.delete(a, np.isnan(a)) urmse[i, t] = unbiased_rmse(a, b) r2[i, t] = r2_score(a, b) + + print(np.nanmedian(r2, axis=0)) # cal physical consistency phy_cons = cal_phy_cons(aux_test, y_pred, y_test) diff --git a/src/preprocess/colm2era.py b/src/preprocess/colm2era.py deleted file mode 100644 index e118e87..0000000 --- a/src/preprocess/colm2era.py +++ /dev/null @@ -1,14 +0,0 @@ -def h2osoi_colm2era(h2osoi): - # turn h2osoi to ERA5-Land soil depth - # NOTE: 0.71/2.79/6.23/11.89/21.22/36.61/61.98/103.80/172.76/286.46cm - # to 7/28/100/286.46cm - # h2osoi: (time, lat, lon, soil) - swvl1 (time, lat, lon) - swvl1 = (h2osoi[:, :, :, 0]*0.71+h2osoi[:, :, :, 1]*(2.79-0.71) + - h2osoi[:, :, :, 2]*(6.23-2.79)+h2osoi[:, :, :, 3]*(7-6.23))/7 - swvl2 = (h2osoi[:, :, :, 3]*(11.89-7)+h2osoi[:, :, :, 4]*(21.22-11.89) + - h2osoi[:, :, :, 5]*(28-21.22))/(28-7) - swvl3 = (h2osoi[:, :, :, 5]*(36.61-28)+h2osoi[:, :, :, 6]*(61.98-36.61) + - h2osoi[:, :, :, 7]*(100-61.98))/(100-28) - swvl4 = (h2osoi[:, :, :, 7]*(103.80-100)+h2osoi[:, :, :, 8]*(172.76-103.80) + - h2osoi[:, :, :, 9]*(286.46-172.76))/(286.46-100) - return swvl1, swvl2, swvl3, swvl4 diff --git a/src/train.py b/src/train.py index 20bc7e5..ad1db48 100644 --- a/src/train.py +++ b/src/train.py @@ -1,37 +1,15 @@ import time -from tqdm import trange import numpy as np import tensorflow as tf -from tensorflow.keras.optimizers import Adam import tensorflow_addons as tfa +from tensorflow.keras.optimizers import Adam +from tqdm import trange -from data_gen import load_train_data, load_test_data -from loss import RMSELoss, MassConsLoss, RMetrics -from model import MTLHardLSTM_v1, VanillaLSTM, MTLLSTM, MTLHardLSTM_v2 - - -# NOTE: If we add decorator `tf.function` of `train_step`, and -# we try to trained model twice. It will raise error: -# "with ValueError: tf.function only supports singleton -# tf.Variables created on the first call. Make sure the -# tf.Variable is only created once or created outside -# tf.function". Thus, `train_step` only used for multi-task -# model to speed up trainning. see `train_multi`. -@tf.function -def train_step(x, y, model, loss_fn, optim, metric): - with tf.GradientTape() as tape: - pred = model(x, training=True) - loss = loss_fn(y, pred) - grads = tape.gradient(loss, model.trainable_variables) - optim.apply_gradients(zip(grads, model.trainable_variables)) - metric.update_state(y, pred) - - -@tf.function -def test_step(x, y, model, metric): - pred = model(x, training=False) - metric.update_state(y, pred) +from data_gen import load_test_data, load_train_data +from loss import MassConsLoss, NaNMSELoss +from metric import NSEMetrics +from model import HardMTLModel_v1, MTLModel_v1, MTLModel_v2, STModel def train(x, @@ -47,8 +25,8 @@ def train(x, # learing rate schedule. We found `Adam` perform # much better than `Adagrad`, `Adadelta`. optim = Adam(cfg["learning_rate"]) - metric = RMetrics(cfg)#tfa.metrics.RSquare() - patience = 5 + metric = NSEMetrics(cfg) + patience = 10 wait = 0 best = -9999 @@ -67,43 +45,45 @@ def train(x, x_valid, y_valid, aux_valid = x[:, N:], y[:, N:], aux[:, N:] x, y, aux = x[:, :N], y[:, :N], aux[:, :N] x_valid, y_valid, aux_valid, mean_valid, std_valid = load_test_data( - cfg, x_valid, y_valid, aux_valid, scaler) - valid_metric = RMetrics(cfg)#tfa.metrics.RSquare() - valid_save_metric = tfa.metrics.RSquare() + cfg, x_valid, y_valid, aux_valid, scaler, stride=4) + valid_metric = NSEMetrics(cfg) # train and validate # NOTE: We preprare two callbacks for training: # early stopping and save best model. for _ in range(100): # prepare models - if cfg["model_name"] == 'single_task': - model = VanillaLSTM(cfg) - elif cfg["model_name"] in ['multi_tasks', 'soft_multi_tasks']: - model = MTLLSTM(cfg) - elif cfg["model_name"] in ['hard_multi_tasks_v2', 'hard_multi_tasks_v3']: - model = MTLHardLSTM_v2(cfg) - elif cfg["model_name"] == 'hard_multi_tasks_v1': - model = MTLHardLSTM_v1(cfg) + if cfg["model_name"] in ['single_task']: + model = STModel(cfg) + elif cfg["model_name"] in ['multi_tasks_v1', 'soft_multi_tasks']: + model = MTLModel_v1(cfg) + elif cfg["model_name"] in ['multi_tasks_v2']: + model = MTLModel_v2(cfg) + elif cfg["model_name"] in ['hard_multi_tasks_v1']: + model = HardMTLModel_v1(cfg) with trange(1, cfg["epochs"]+1) as pbar: for epoch in pbar: pbar.set_description( cfg["model_name"]+' '+str(num_task)+' member '+str(num_repeat)) + t_begin = time.time() # train MCLoss, MSELoss = 0, 0 - t0 = time.time() for iter in range(0, cfg["niter"]): # generate batch data x_batch, y_batch, aux_batch, mean_batch, std_batch = \ load_train_data(cfg, x, y, aux, scaler) with tf.GradientTape() as tape: # cal MSE loss - if cfg["model_name"] in ['single_task', 'multi_tasks', 'soft_multi_tasks']: - pred = model(x_batch, training=True) + if cfg["model_name"] in \ + ['single_task','multi_tasks_v1','soft_multi_tasks']: + pred = model(x_batch, training=True) + elif cfg["model_name"] in ['multi_tasks_v2']: + pred, adaptive_loss = model(x_batch, y_batch, training=True) else: pred = model(x_batch, aux_batch, mean_batch, std_batch, training=True) - mse_loss = RMSELoss(cfg)(y_batch, pred) + mse_loss = NaNMSELoss(cfg)(y_batch, pred[:,1]) # only cal on last step MSELoss+=mse_loss # cal physic loss @@ -112,31 +92,28 @@ def train(x, MCLoss += phy_loss # cal all loss - if cfg["model_name"] == "soft_multi_tasks": + if cfg["model_name"] in ["soft_multi_tasks"]: loss = (1-cfg["alpha"])*mse_loss + (cfg["alpha"])*phy_loss + elif cfg["model_name"] in ["multi_tasks_v2"]: + loss = adaptive_loss else: loss = mse_loss # gradient tape grads = tape.gradient(loss, model.trainable_variables) optim.apply_gradients(zip(grads, model.trainable_variables)) - metric.update_state(y_batch, pred) - t1 = time.time() + metric.update_state(y_batch, pred[:,1]) + t_end = time.time() # get loss log - train_acc = metric.result() + t_acc = metric.result() if cfg["model_name"] == 'single_task': - train = train_acc["all"].numpy() - loss_str = "Epoch {} Train NSE {:.3f} MSE Loss {:.3f} MC Loss {:.3f} time {:.2f}".format(epoch, train, MSELoss/cfg["niter"], MCLoss/cfg["niter"], t1-t0) + t0 = t_acc["loss"].numpy() + loss_str = "Epoch {} Train NSE {:.3f} MSE Loss {:.3f} MC Loss {:.3f} time {:.2f}".format(epoch,t0,MSELoss/cfg["niter"],MCLoss/cfg["niter"],t_end-t_begin) else: - train_1 = train_acc["SWVL_1"].numpy() - train_2 = train_acc["SWVL_2"].numpy() - train_3 = train_acc["SWVL_3"].numpy() - train_4 = train_acc["SWVL_4"].numpy() - train_5 = train_acc["ET"].numpy() - train_6 = train_acc["R"].numpy() + t1,t2,t3,t4,t5,t6 = t_acc["0"].numpy(),t_acc["1"].numpy(),t_acc["2"].numpy(), \ + t_acc["3"].numpy(),t_acc["4"].numpy(),t_acc["5"].numpy() loss_str = "Epoch {} Train NSE SWVL1 {:.3f} SWVL2 {:.3f} SWVL3 {:.3f} SWVL4 {:.3f} ET {:.3f} R {:.3f} MSE Loss {:.3f} MC Loss {:.3f} time {:.2f}".format( - epoch, train_1, train_2, train_3, train_4, train_5, train_6, - MSELoss/cfg["niter"], MCLoss/cfg["niter"], t1-t0) + epoch,t1,t2,t3,t4,t5,t6,MSELoss/cfg["niter"],MCLoss/cfg["niter"],t_end-t_begin) print(loss_str) metric.reset_states() @@ -151,53 +128,44 @@ def train(x, if epoch % 20 == 0: wait += 1 - # NOTE: We use larger than 3 years for validate, and - # used grids-mean NSE as valid metrics. - # We cannot use `tf.data.Dataset.from_tensor_slices` - # to transform nd.array to tensor. Because it will - # put all valid data into GPU, which exceed memory. - t0 = time.time() + # NOTE: We used grids-mean NSE as valid metrics. + t_begin = time.time() for i in range(x_valid.shape[0]): if cfg["model_name"] in ['single_task', \ - 'multi_tasks', 'soft_multi_tasks']: + 'multi_tasks_v1', 'multi_tasks_v2', 'soft_multi_tasks']: pred = model(x_valid[i], training=False) else: pred = model(x_valid[i], aux_valid[i], \ mean_valid[i], std_valid[i], training=False) # cal mse loss - mse_valid_loss = RMSELoss(cfg)(y_valid[i], pred) + mse_valid_loss = NaNMSELoss(cfg)(y_valid[i], pred[:,1]) MSE_valid_loss+=mse_valid_loss - valid_metric.update_state(y_valid[i], pred) - valid_save_metric.update_state(y_valid[i], pred) + valid_metric.update_state(y_valid[i], pred[:,1]) # cal phy loss phy_loss = MassConsLoss( cfg, mean_valid[i], std_valid[i])(aux_valid[i], pred) MC_valid_loss+=phy_loss - t1 = time.time() + t_end = time.time() # get loss log - val_acc = valid_metric.result() - val_save_acc = valid_save_metric.result().numpy() + v_acc = valid_metric.result() if cfg["model_name"] == 'single_task': - val = val_acc["all"].numpy() + val = v_acc["loss"].numpy() loss_str = '\033[1;31m%s\033[0m' % \ - "Epoch {} Val NSE {:.3f} MSE Loss {:.3f} MC Loss {:.3f} time {:.2f}".format(epoch, val, MSE_valid_loss/x_valid.shape[0], - MC_valid_loss/x_valid.shape[0], t1-t0) + "Epoch {} Val NSE {:.3f} MSE Loss {:.3f} MC Loss {:.3f} time {:.2f}".format(epoch,val,MSE_valid_loss/x_valid.shape[0], + MC_valid_loss/x_valid.shape[0],t_end-t_begin) + val_save_acc = val else: - val_1 = val_acc["SWVL_1"].numpy() - val_2 = val_acc["SWVL_2"].numpy() - val_3 = val_acc["SWVL_3"].numpy() - val_4 = val_acc["SWVL_4"].numpy() - val_5 = val_acc["ET"].numpy() - val_6 = val_acc["R"].numpy() + t1,t2,t3,t4,t5,t6 = v_acc["0"].numpy(),v_acc["1"].numpy(),\ + v_acc["2"].numpy(), v_acc["3"].numpy(),v_acc["4"].numpy(),\ + v_acc["5"].numpy() loss_str = '\033[1;31m%s\033[0m' % \ - "Epoch {} Val NSE SWVL1 {:.3f} SWVL2 {:.3f} SWVL3 {:.3f} SWVL4 {:.3f} ET {:.3f} R {:.3f} MSE Loss {:.3f} MC Loss {:.3f} time {:.2f}".format(epoch, val_1, val_2, val_3, val_4, val_5, val_6, - MSE_valid_loss/x_valid.shape[0], - MC_valid_loss/x_valid.shape[0], t1-t0) + "Epoch {} Val NSE SWVL1 {:.3f} SWVL2 {:.3f} SWVL3 {:.3f} SWVL4 {:.3f} ET {:.3f} R {:.3f} MSE Loss {:.3f} MC Loss {:.3f} time {:.2f}".format(epoch,t1,t2,t3,t4,t5,t6, MSE_valid_loss/x_valid.shape[0], + MC_valid_loss/x_valid.shape[0],t_end-t_begin) + val_save_acc = (t1+t2+t3+t4+t5+t6)/6 print(loss_str) valid_metric.reset_states() - valid_save_metric.reset_states() # save best model by val loss # NOTE: save best MSE results get `single_task` better than `multi_tasks` @@ -220,3 +188,26 @@ def train(x, if wait >= patience: return return + + +# NOTE: If we add decorator `tf.function` of `train_step`, and +# we try to trained model twice. It will raise error: +# "with ValueError: tf.function only supports singleton +# tf.Variables created on the first call. Make sure the +# tf.Variable is only created once or created outside +# tf.function". Thus, `train_step` only used for multi-task +# model to speed up trainning. see `train_multi`. +@tf.function +def train_step(x, y, model, loss_fn, optim, metric): + with tf.GradientTape() as tape: + pred = model(x, training=True) + loss = loss_fn(y, pred) + grads = tape.gradient(loss, model.trainable_variables) + optim.apply_gradients(zip(grads, model.trainable_variables)) + metric.update_state(y, pred) + + +@tf.function +def test_step(x, y, model, metric): + pred = model(x, training=False) + metric.update_state(y, pred) \ No newline at end of file diff --git a/src/utils.py b/src/utils.py index 3b5f010..3e6289d 100644 --- a/src/utils.py +++ b/src/utils.py @@ -2,6 +2,43 @@ import numpy as np +def make_CoLM_soil_depth(): + # central depth unit (m) + fs = 0.025 + z = [fs*(np.exp(0.5*(i-0.5))-1) for i in range(1, 11)] + # thickness of soil (m) + depth = [] + for i in range(10): + if i == 0: + depth.append(0.5*(z[0]+z[1])) + elif i == 9: + depth.append(z[-1]-z[-2]) + else: + depth.append(0.5*(z[i+1]-z[i-1])) + # depth at layer interface (m) + zi = [] + for i in range(10): + if i == 9: + zi.append(z[-1]+0.5*depth[-1]) + else: + zi.append(0.5*(z[i]+z[i+1])) + zi = np.array(zi)*100 # m->cm + return depth, zi + + +def make_swc_CoLM2EC(wice, wliq): + """turn CoLM to EC soil setting.""" + depth, zi = make_CoLM_soil_depth() + # cal h2osoi + h2osoi = (wliq+wice)/(np.array(depth)*1000) + # 1.75 4.50 9,05 16.55 28.91 49.29 82.89 138.28 229.61 343.30 + swvl1 = (h2osoi[:,:,:,0]*zi[0]+h2osoi[:,:,:,1]*(zi[1]-zi[0])+h2osoi[:,:,:,2]*(7-zi[1]))/7 + swvl2 = (h2osoi[:,:,:,2]*(zi[2]-7)+h2osoi[:,:,:,3]*(zi[3]-zi[2])+h2osoi[:,:,:,4]*(28-zi[3]))/21 + swvl3 = (h2osoi[:,:,:,4]*(zi[4]-28)+h2osoi[:,:,:,5]*(zi[5]-zi[4])+h2osoi[:,:,:,6]*(zi[6]-zi[5])+h2osoi[:,:,:,7]*(100-zi[6]))/72 + swvl4 = (h2osoi[:,:,:,7]*(zi[7]-100)+h2osoi[:,:,:,8]*(zi[8]-zi[7])+h2osoi[:,:,:,9]*(zi[9]-zi[8]))/(zi[9]-100) + return swvl1, swvl2, swvl3, swvl4, zi + + def unbiased_rmse(y_true, y_pred): predmean = np.nanmean(y_pred) targetmean = np.nanmean(y_true) @@ -9,14 +46,14 @@ def unbiased_rmse(y_true, y_pred): targetanom = y_true - targetmean return np.sqrt(np.nanmean((predanom-targetanom)**2)) + def r2_score(y_true, y_pred): - np.square mask = y_true == y_true a, b = y_true[mask], y_pred[mask] unexplained_error = np.nansum(np.square(a-b)) total_error = np.nansum(np.square(a - np.nanmean(a))) - r2 = 1. - unexplained_error/total_error - return r2 + return 1. - unexplained_error/total_error + def init_fold(work_path): print("[PHydro] Construct work path in {}".format(work_path)) @@ -31,6 +68,7 @@ def init_fold(work_path): if not os.path.exists(work_path+"/output/forecast"): os.mkdir(work_path+"/output/forecast") + def site2grid(input, site_lat, site_lon, grid_lat, grid_lon): # postprocess (return sites to grids) if input.ndim == 2: @@ -47,6 +85,7 @@ def site2grid(input, site_lat, site_lon, grid_lat, grid_lon): input_grid[idx_lat, idx_lon] = input[i] return input_grid + def cal_phy_cons(aux, y_pred, y_true): print(aux.shape, y_pred.shape, y_true.shape) # cal physical consistency