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 d1784c9..0000000 Binary files a/src/figure1.pdf and /dev/null differ diff --git a/src/layers.py b/src/layers.py new file mode 100644 index 0000000..67991d7 --- /dev/null +++ b/src/layers.py @@ -0,0 +1,110 @@ +import tensorflow as tf +from tensorflow import math +from tensorflow.keras.initializers import Constant +from tensorflow.keras.layers import Layer +from tensorflow.keras.losses import mean_squared_error + + +class WeightedMultiLossLayer(Layer): + def __init__(self, cfg): + super().__init__() + self.num_out = cfg["num_out"] + + def build(self, input_shape=None): + self.log_vars = [] + for i in range(self.num_out): + self.log_vars += [self.add_weight(name='log_var'+str(i), + shape=(1,), + initializer=Constant(1.), + trainable=True, + dtype='float64')] + super().build(input_shape) + + def call(self, y_true, y_pred): + loss_sum = 0 + # cal seperate RMSE loss + 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] + self.metrics_.append(math.sqrt(mean_squared_error(a,b))) + + for i, loss in enumerate(self.metrics_): + loss_sum += 0.5 / (self.log_vars[i] ** 2) * loss + \ + tf.math.log(1 + self.log_vars[i] ** 2) + #self.add_loss(loss_sum) + return loss_sum + + +class MassConsLayer(Layer): + """Mass conserve layer""" + + def __init__(self, cfg): + super().__init__() + self.idx = cfg["resid_idx"] + self.num_out = cfg["num_out"] + + 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, 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') + depth, zi = make_CoLM_soil_depth() # m/cm + soil_depth = [70, 210, 720, 10*(zi[9]-100)] # 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) + diff --git a/src/loss.py b/src/loss.py index 5bb423f..f6cf788 100644 --- a/src/loss.py +++ b/src/loss.py @@ -1,25 +1,21 @@ import tensorflow as tf from tensorflow import math from tensorflow.keras.losses import Loss, mean_squared_error -from tensorflow.keras.metrics import Metric +from utils import make_CoLM_soil_depth -class RMSELoss(Loss): + +class NaNMSELoss(Loss): def __init__(self, cfg): super().__init__() self.model_name = cfg["model_name"] - self.sub = list(range(cfg["num_out"])) - self.sub.pop(cfg["resid_idx"]) self.idx = cfg["resid_idx"] def call(self, y_true, y_pred): - if self.model_name == 'hard_multi_tasks_v3': - y_true = tf.gather(y_true, self.sub, axis=-1) - y_pred = tf.gather(y_pred, self.sub, axis=-1) mask = y_true == y_true y_true = y_true[mask] y_pred = y_pred[mask] - return math.sqrt(mean_squared_error(y_true, y_pred)) + return mean_squared_error(y_true, y_pred) class MassConsLoss(Loss): @@ -30,57 +26,33 @@ def __init__(self, cfg, mean, std): super().__init__() # NOTE: Loss class must use call rather than __call__ - def call(self, aux, pred): + def call(self, aux, pred): #(b,6), (b,2,6) # init - soil_depth = [70, 210, 720, 1864.6] # mm + depth, zi = make_CoLM_soil_depth() # m/cm + soil_depth = [70, 210, 720, 10*(zi[9]-100)] # mm aux = tf.cast(aux, 'float32') pred = tf.cast(pred, 'float32') std = tf.cast(self.std, 'float32') mean = tf.cast(self.mean, 'float32') + pred_prev, pred_now = pred[:,0], pred[:,1] + mask = aux == aux + aux = aux[mask] + pred_now = pred_now[mask] + pred_prev = pred_prev[mask] + std = std[mask] + mean = mean[mask] + # reverse scaling - pred = math.multiply(pred, std)+mean - # cal mass conserve loss - swvl = math.multiply(pred[:, :4], soil_depth) - # unnormalized all mm/day - pred = tf.concat([swvl, pred[:, 4:]], axis=-1) + pred_prev = math.multiply(pred_prev, std)+mean + pred_now = math.multiply(pred_now, std)+mean + + # 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,) # NOTE: `y[0]+math.reduce_sum(y[1:])` != `math.reduce_sum(y)` - phy_loss = tf.abs(aux-pred[:,0]-math.reduce_sum(pred[:,1:], axis=-1)) - return math.multiply(self.scale, math.reduce_mean(phy_loss)) - - -class RMetrics(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"] - - 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_ = [] - if self.model == 'single_task': - num_out = 1 - else: - num_out = 6 + return math.multiply(self.scale, math.reduce_mean(w_b)) - for i in range(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 {"all": self.metrics_[0]} - else: - return {"SWVL_1": self.metrics_[0], - "SWVL_2": self.metrics_[1], - "SWVL_3": self.metrics_[2], - "SWVL_4": self.metrics_[3], - "ET": self.metrics_[4], - "R": self.metrics_[5]} + \ No newline at end of file diff --git a/src/main.py b/src/main.py index d4cb4e0..d18d553 100644 --- a/src/main.py +++ b/src/main.py @@ -3,20 +3,20 @@ <<<<<<<< HEAD Author: Lu Li -11/1/2022: First edition +11/1/2022: V1.0 """ -import os import json +import os import numpy as np import wandb from config import parse_args from data import Dataset -from utils import init_fold -from train import train from eval import eval_multi, eval_single +from train import train +from utils import init_fold def main(cfg): @@ -49,13 +49,11 @@ def main(cfg): else: f = Dataset(cfg, mode='train') x_train, y_train, aux_train = f.fit() - f = Dataset(cfg, mode='test') - x_test, y_test, aux_test = f.fit() - print(aux_test.shape) - print(aux_test.shape) np.save("x_train.npy", x_train) np.save("y_train.npy", y_train) np.save("z_train.npy", aux_train) + f = Dataset(cfg, mode='test') + x_test, y_test, aux_test = f.fit() np.save("x_test.npy", x_test) np.save("y_test.npy", y_test) np.save("z_test.npy", aux_test) @@ -70,14 +68,14 @@ def main(cfg): # train & inference print("[PHydro] Start train!") if cfg["model_name"] == 'single_task': - for i in range(cfg["num_out"]): - for j in range(cfg["num_repeat"]): - train(x_train, y_train[:,:,i:i+1], aux_train, scaler, cfg, j, i) + #for i in range(cfg["num_out"]): + # for j in range(cfg["num_repeat"]): + # train(x_train, y_train[:,:,i:i+1], aux_train, scaler, cfg, j, i) y_pred = eval_single(x_test, y_test, scaler, cfg) else: for j in range(cfg["num_repeat"]): train(x_train, y_train, aux_train, scaler, cfg, j) - y_pred = eval_multi(x_test, y_test, aux_test, scaler, cfg) + y_pred, y_true = eval_multi(x_test, y_test, aux_test, scaler, cfg) # save print('[PHydro] Saving') @@ -85,6 +83,7 @@ def main(cfg): if not os.path.exists(path): os.mkdir(path) np.save(cfg["model_name"]+'_gd_9km.npy', y_pred) + np.save('obs_gd_9km.npy', y_true) os.system('mv {} {}'.format(cfg["model_name"]+'_gd_9km.npy', path)) diff --git a/src/preprocess/make_input_data.py b/src/make_input_data.py similarity index 54% rename from src/preprocess/make_input_data.py rename to src/make_input_data.py index 8791045..4903209 100644 --- a/src/preprocess/make_input_data.py +++ b/src/make_input_data.py @@ -1,4 +1,3 @@ -import glob import os import netCDF4 as nc @@ -7,30 +6,29 @@ import shapely.geometry as geometry import xarray as xr -from colm2era import h2osoi_colm2era +from utils import h2osoi_colm2era + # ------------------------------------------------------------------------------ # params # ------------------------------------------------------------------------------ -years = np.arange(2019, 2021) # 2000-2016 -months = np.arange(1, 13) # 1-12 -PATH = "/tera05/lilu/PHydro/data/" # path contain raw data +years = np.arange(2017, 2020) # 2000-2016 +months = np.arange(1,12) # 1-12 inputs_path = "/tera05/lilu/PHydro/input/" # path contain inputs -mode = 'test' +mode = 'train' + + # ------------------------------------------------------------------------------ # extract guangdong points # ------------------------------------------------------------------------------ -in_shape_idx = np.load(inputs_path+"guangdong_9km_shapefile.npy") -idx = np.where(in_shape_idx == 1)[0] -# FIXME: Remove grids contains all NaN automatically -idx1 = np.delete(idx, [1220,1544,1550,1563]) - - -# codes for extract guangdong points -with xr.open_dataset(PATH+"CoLM/global_MSWX_9km_hist_2000-01.nc", decode_times=False) as f: +""" +# read lat/lon in CoLM outputs +with xr.open_dataset(PATH+"global_grid_ERA5_pft_GD_hist_2015-01.nc", decode_times=False) as f: lat, lon = np.array(f.lat), np.array(f.lon) + rnof = np.array(f.f_rnof) -gd_shp = shapefile.Reader(PATH+"shapefile/guangdong.shp") +# shapefile input +gd_shp = shapefile.Reader("/tera05/lilu/PHydro/data/shapefile/guangdong.shp") feat = gd_shp.shapeRecords()[0] first = feat.shape grid_lon, grid_lat = np.meshgrid(lon, lat) @@ -41,32 +39,84 @@ in_shape_points = [] count = 0 for i, pt in enumerate(flat_points): + print(i) if geometry.Point(pt).within(geometry.shape(first)): in_shape_points.append(pt) in_shape_idx[i] = 1 count += 1 idx = np.where(in_shape_idx == 1)[0] -np.save("guangdong_9km_shapefile.npy", in_shape_idx) +# remove NaN points in runoff outputs +rnof = rnof.reshape(rnof.shape[0], -1)[:, idx] # time, grid +remove_idx = np.unique(np.where(np.isnan(rnof))[-1]) +idx = np.delete(idx, [remove_idx]) +print(idx.shape) + +np.save("gd_9km_shapefile.npy", idx) +""" +idx = np.load(inputs_path+"gd_9km_shapefile.npy") + + +# ------------------------------------------------------------------------------ +# read ancillary variables +# ------------------------------------------------------------------------------ +PATH = "/tera05/lilu/PHydro/data/" +with xr.open_dataset(PATH+'ancillary/Beck_KG_V1_present_0p1.nc') as f: + lat2, lon2 = np.array(f.lat), np.array(f.lon) + lat_idx = np.where((lat2>18) & (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