dl4ds.metrics

  0import tensorflow as tf
  1import numpy as np
  2import xarray as xr
  3from joblib import Parallel, delayed
  4from matplotlib import pyplot as plt
  5from sklearn.metrics import mean_squared_error
  6from scipy.stats import spearmanr, pearsonr
  7import os
  8import seaborn as sns
  9import ecubevis as ecv
 10
 11from .utils import checkarray_ndim, Timing
 12
 13
 14def compute_rmse(y, y_hat, over='time', squared=False, n_jobs=40):
 15    """
 16
 17    Parameters
 18    ----------
 19    squared : bool
 20        If True returns MSE value, if False returns RMSE value.
 21
 22    https://scikit-learn.org/stable/modules/generated/sklearn.metrics.mean_squared_error.html
 23    """
 24    def rmse_per_px(tupyx):
 25        y_coord, x_coord = tupyx
 26        return y_coord, x_coord, mean_squared_error(y[:,y_coord,x_coord,0], y_hat[:,y_coord,x_coord,0])
 27
 28    def rmse_gridpair(index):
 29        return mean_squared_error(y[index].flatten(), y_hat[index].flatten(), squared=squared)
 30
 31    #---------------------------------------------------------------------------
 32    if over == 'time':
 33        rmse_map = np.zeros_like(y[0,:,:,0]) 
 34        rmse_map *= np.nan
 35        yy, xx = np.where(y[0,:,:,0])
 36        coords = zip(yy, xx)
 37        out = Parallel(n_jobs=n_jobs, verbose=False)(delayed(rmse_per_px)(i) for i in coords) 
 38
 39        for i in range(len(out)):
 40            y_coord, x_coord, val = out[i]
 41            rmse_map[y_coord, x_coord] = val
 42        return rmse_map
 43
 44    elif over == 'space':
 45        n_timesteps = np.arange(y.shape[0])
 46        out = Parallel(n_jobs=n_jobs, verbose=False)(delayed(rmse_gridpair)(i) for i in n_timesteps)
 47        return out
 48    
 49
 50def compute_correlation(y, y_hat, over='time', mode='spearman', n_jobs=40):
 51    """
 52    https://scipy.github.io/devdocs/generated/scipy.stats.spearmanr.html
 53
 54    https://docs.scipy.org/doc/scipy/reference/generated/scipy.stats.pearsonr.html
 55    """
 56    def corr_per_px(tupyx):
 57        if mode == 'spearman':
 58            f = spearmanr
 59        elif mode == 'pearson':
 60            f = pearsonr
 61        y_coord, x_coord = tupyx
 62        return y_coord, x_coord, f(y[:,y_coord,x_coord,0], y_hat[:,y_coord,x_coord,0])[0]
 63
 64    def corr_per_gridpair(index):
 65        if mode == 'spearman':
 66            f = spearmanr
 67        elif mode == 'pearson':
 68            f = pearsonr
 69        return index, f(y[index].ravel(), y_hat[index].ravel())[0]
 70
 71    #---------------------------------------------------------------------------
 72    if over == 'time':
 73        corrmap = np.zeros_like(y[0,:,:,0]) 
 74        corrmap *= np.nan
 75        yy, xx = np.where(y[0,:,:,0])
 76        coords = zip(yy, xx)
 77        out = Parallel(n_jobs=n_jobs, verbose=False)(delayed(corr_per_px)(i) for i in coords) 
 78
 79        for i in range(len(out)):
 80            y_coord, x_coord, val = out[i]
 81            corrmap[y_coord, x_coord] = val
 82
 83        return corrmap
 84
 85    elif over == 'space':
 86        n_timesteps = np.arange(y.shape[0])
 87        out = Parallel(n_jobs=n_jobs, verbose=False)(delayed(corr_per_gridpair)(i) for i in n_timesteps)
 88
 89        list_corrs = []
 90        list_inds = []
 91        for i in range(len(out)):
 92            index, value = out[i]
 93            list_inds.append(index)
 94            list_corrs.append(value)
 95
 96        return list_corrs
 97
 98
 99def compute_metrics(
100    y_test, 
101    y_test_hat, 
102    dpi=150, 
103    plot_size_px=1000,
104    n_jobs=-1, 
105    scaler=None, 
106    mask=None,
107    save_path=None):
108    """ Compute temporal and spatial-wise metrics, e.g., RMSE and CORRELATION, 
109    based on the groundtruth and prediction ndarrays.
110
111    Parameters
112    ----------
113    y_test : np.ndarray
114        Groundtruth.
115    y_test_hat : np.ndarray
116        Prediction.
117    dpi : int, optional
118        DPI of the plots.  
119    n_jobs : int, optional
120        Number of cores for the computation of metrics (parallelizing over
121        grid points). Passed to joblib.Parallel. If -1 all CPUs are used. If 1 
122        or None is given, no parallel computing code is used at all, which is 
123        useful for debugging.
124    scaler : scaler object
125        Scaler object from preprocessing module. 
126    mask : np.ndarray or None
127        Binary mask with valid (ones) and non-valid (zeroes) grid points.
128    save_path : str or None, optional
129        Path to save results to disk. 
130        
131    """
132    timing = Timing()
133
134    if y_test.ndim == 5:
135        y_test = np.squeeze(y_test, -1)
136        y_test_hat = np.squeeze(y_test_hat, -1)
137
138    y_test = checkarray_ndim(y_test, 4, -1)
139    y_test_hat = checkarray_ndim(y_test_hat, 4, -1)
140
141    # backward transformation with the provided scaler
142    if scaler is not None:
143        if hasattr(scaler, 'inverse_transform'):
144            y_test = scaler.inverse_transform(y_test)
145            y_test_hat = scaler.inverse_transform(y_test_hat)        
146
147    # applying valid grid points mask
148    if mask is not None:
149        if isinstance(mask, xr.DataArray):
150            mask = mask.values.copy()
151        elif isinstance(mask, np.ndarray):
152            mask = mask.copy()
153
154        if mask.ndim == 2:
155            mask = np.expand_dims(mask, -1)
156        y_test = y_test.copy()
157        y_test_hat = y_test_hat.copy()
158        for i in range(y_test.shape[0]):
159            y_test[i] *= mask
160        for i in range(y_test_hat.shape[0]):
161            y_test_hat[i] *= mask
162        mask_nan = mask.astype('float').copy()
163        mask_nan[mask == 0] = np.nan
164        mask = np.squeeze(mask)
165
166    ### Computing metrics
167    drange = max(y_test.max(), y_test_hat.max()) - min(y_test.min(), y_test_hat.min())
168    with tf.device("cpu:0"):
169        psnr = tf.image.psnr(y_test, y_test_hat, drange)
170    mean_psnr = np.mean(psnr)
171    std_psnr = np.std(psnr)
172
173    with tf.device("cpu:0"):
174        ssim = tf.image.ssim(tf.convert_to_tensor(y_test, dtype=tf.float32), 
175                             tf.convert_to_tensor(y_test_hat, dtype=tf.float32), 
176                             drange)
177    mean_ssim = np.mean(ssim)
178    std_ssim = np.std(ssim)
179
180    with tf.device("cpu:0"):
181        maes = tf.keras.metrics.mean_absolute_error(y_test, y_test_hat)
182    maes_pairs = np.mean(maes, axis=(1,2))
183    mean_mae = np.mean(maes_pairs)
184    std_mae = np.std(maes_pairs)
185
186    ### RMSE 
187    temp_rmse_map = compute_rmse(y_test, y_test_hat, n_jobs=n_jobs, over='time')
188    spatial_rmse = compute_rmse(y_test, y_test_hat, n_jobs=n_jobs, over='space')
189    if save_path is not None:
190        np.save(os.path.join(save_path, 'metrics_mse_pergridpair.npy'), spatial_rmse)
191    mean_spatial_rmse = np.mean(spatial_rmse)
192    std_spatial_rmse = np.std(spatial_rmse)
193    mean_temp_rmse = np.nanmean(temp_rmse_map)
194    std_temp_rmse = np.nanstd(temp_rmse_map)
195    if mask is not None:
196        temp_rmse_map[np.where(mask == 0)] = 0
197    subpti = f'RMSE map ($\mu$ = {mean_temp_rmse:.6f})'
198    if save_path is not None:
199        savepath = os.path.join(save_path, 'metrics_pergridpoint_rmse_map.png')
200        np.save(os.path.join(save_path, 'metrics_pergridpoint_rmse_map.npy'), temp_rmse_map)
201    else:
202        savepath = None
203    ecv.plot_ndarray(temp_rmse_map, dpi=dpi, subplot_titles=(subpti), cmap='viridis', 
204                     plot_size_px=plot_size_px, interactive=False, save=savepath)
205
206    ### Normalized per grid point RMSE 
207    norm_temp_rmse_map = temp_rmse_map / (np.mean(y_test) * 100)
208    norm_mean_temp_rmse = np.nanmean(norm_temp_rmse_map)
209    norm_std_temp_rmse = np.nanstd(norm_temp_rmse_map)
210    if mask is not None:
211        norm_temp_rmse_map[np.where(mask == 0)] = 0
212    subpti = f'nRMSE map ($\mu$ = {norm_mean_temp_rmse:.6f})'
213    if save_path is not None:
214        savepath = os.path.join(save_path, 'metrics_pergridpoint_nrmse_map.png')
215        np.save(os.path.join(save_path, 'metrics_pergridpoint_nrmse_map.npy'), norm_temp_rmse_map)
216    else:
217        savepath = None
218    ecv.plot_ndarray(norm_temp_rmse_map, dpi=dpi, subplot_titles=(subpti), cmap='viridis', 
219                     plot_size_px=plot_size_px, interactive=False, save=savepath)
220
221    # Normalized mean bias
222    nmeanbias = np.mean(y_test_hat - y_test, axis=0)
223    nmeanbias /= np.mean(y_test) * 100
224    if mask is not None:
225        nmeanbias *= mask_nan
226    mean_nmeanbias = np.nanmean(nmeanbias)
227    nmeanbias[np.where(mask == 0)] = 0
228    subpti = f'NMBias map ($\mu$ = {mean_nmeanbias:.6f})'
229    if save_path is not None:
230        savepath = os.path.join(save_path, 'metrics_nmeanbias_map.png')
231        np.save(os.path.join(save_path, 'metrics_nmeanbias_map.npy'), nmeanbias)
232    else:
233        savepath = None
234    ecv.plot_ndarray(nmeanbias, dpi=dpi, subplot_titles=(subpti), cmap='viridis', 
235                     plot_size_px=plot_size_px, interactive=False, save=savepath)
236
237    ### Spearman correlation coefficient
238    spatial_spearman_corr = compute_correlation(y_test, y_test_hat, n_jobs=n_jobs, over='space')
239    mean_spatial_spearman_corr = np.mean(spatial_spearman_corr)
240    std_spatial_spearman_corr = np.std(spatial_spearman_corr)
241    if save_path is not None:
242        np.save(os.path.join(save_path, 'metrics_spearcorr_pergridpair.npy'), spatial_spearman_corr)
243
244    ### Pearson correlation coefficient
245    spatial_pearson_corr = compute_correlation(y_test, y_test_hat, mode='pearson', n_jobs=n_jobs, over='space')
246    mean_spatial_pearson_corr = np.mean(spatial_pearson_corr)
247    std_spatial_pearson_corr = np.std(spatial_pearson_corr)
248    if save_path is not None:
249        np.save(os.path.join(save_path, 'metrics_pearcorr_pergridpair.npy'), spatial_pearson_corr)
250    temp_pearson_corrmap = compute_correlation(y_test, y_test_hat, mode='pearson', n_jobs=n_jobs)
251    mean_temp_pearson_corr = np.nanmean(temp_pearson_corrmap)
252    std_temp_pearson_corr = np.nanstd(temp_pearson_corrmap)
253    temp_pearson_corrmap[np.where(mask == 0)] = 0
254    subpti = f'Pearson correlation map ($\mu$ = {mean_temp_pearson_corr:.6f})'
255    if save_path is not None:
256        savepath = os.path.join(save_path, 'metrics_pergridpoint_corrpears_map.png')
257        np.save(os.path.join(save_path, 'metrics_pergridpoint_corrpears_map.npy'), temp_pearson_corrmap)
258    else:
259        savepath = None
260    ecv.plot_ndarray(temp_pearson_corrmap, dpi=dpi, subplot_titles=(subpti), cmap='magma', 
261                     plot_size_px=plot_size_px, interactive=False, save=savepath)
262    
263    ### Plotting violin plots: http://seaborn.pydata.org/tutorial/aesthetics.html
264    sns.set_style("whitegrid") #{"axes.facecolor": ".9"}
265    sns.despine(left=True)
266    sns.set_context("notebook")
267    f, ax = plt.subplots(1, 6, figsize=(15, 5), dpi=dpi)
268    for axis in f.axes:
269        axis.tick_params(labelrotation=40)
270
271    ax_ = sns.violinplot(x=np.array(psnr), ax=ax[0], orient='h', color="skyblue", saturation=1, linewidth=0.8)
272    ax_.set_title('PSNR')
273    ax_.set_xlabel(f'$\mu$ = {mean_psnr:.4f} \n$\sigma$ = {std_psnr:.4f}')
274
275    ax_ = sns.violinplot(x=np.array(ssim), ax=ax[1], orient='h', color="skyblue", saturation=1, linewidth=0.8)
276    ax_.set_title('SSIM')
277    ax_.set_xlabel(f'$\mu$ = {mean_ssim:.4f} \n$\sigma$ = {std_ssim:.4f}')
278
279    ax_ = sns.violinplot(x=maes_pairs, ax=ax[2], orient='h', color="skyblue", saturation=1, linewidth=0.8)
280    ax_.set_title('MAE')
281    ax_.set_xlabel(f'$\mu$ = {mean_mae:.4f} \n$\sigma$ = {std_mae:.4f}')
282
283    ax_ = sns.violinplot(x=spatial_rmse, ax=ax[3], orient='h', color="skyblue", saturation=1, linewidth=0.8)
284    ax_.set_title('RMSE')
285    ax_.set_xlabel(f'$\mu$ = {mean_spatial_rmse:.4f} \n$\sigma$ = {std_spatial_rmse:.4f}')
286
287    ax_ = sns.violinplot(x=spatial_pearson_corr, ax=ax[4], orient='h', color="skyblue", saturation=1, linewidth=0.8)
288    ax_.set_title('Pearson correlation')
289    ax_.set_xlabel(f'$\mu$ = {mean_spatial_pearson_corr:.4f} \n$\sigma$ = {std_spatial_pearson_corr:.4f}')
290
291    ax_ = sns.violinplot(x=spatial_spearman_corr, ax=ax[5], orient='h', color="skyblue", saturation=1, linewidth=0.8)
292    ax_.set_title('Spearman correlation')
293    ax_.set_xlabel(f'$\mu$ = {mean_spatial_spearman_corr:.4f} \n$\sigma$ = {std_spatial_spearman_corr:.4f}')
294
295    f.tight_layout()
296    if save_path is not None: 
297        plt.savefig(os.path.join(save_path, 'metrics_violin_plots.png'))
298        plt.close()
299    else:
300        plt.show()
301    
302    sns.set_style("white")
303
304    if save_path is not None: 
305        f = open(os.path.join(save_path, 'metrics_summary.txt'), "a")
306    else:
307        f = None
308
309    print('Metrics on y_test and y_test_hat:\n', file=f)
310    print(f'PSNR \tmu = {mean_psnr} \tsigma = {std_psnr}', file=f)
311    print(f'SSIM \tmu = {mean_ssim} \tsigma = {std_ssim}', file=f)
312    print(f'MAE \tmu = {mean_mae} \tsigma = {std_mae}', file=f)
313    print(f'Per-grid-point RMSE \tmu = {mean_temp_rmse} \tsigma = {std_temp_rmse}', file=f)
314    print(f'Per-grid-point nRMSE \tmu = {norm_mean_temp_rmse} \tsigma = {norm_std_temp_rmse}', file=f)
315    print(f'Per-grid-point Spearman correlation \tmu = {mean_spatial_spearman_corr} \tsigma = {std_spatial_spearman_corr}', file=f)
316    print(f'Per-grid-point Pearson correlation \tmu = {mean_temp_pearson_corr} \tsigma = {std_temp_pearson_corr}', file=f)
317    print(file=f)
318    print(f'Spatial MSE \tmu = {mean_spatial_rmse} \tsigma = {std_spatial_rmse}', file=f)
319    print(f'Spatial Spearman correlation \tmu = {mean_spatial_spearman_corr} \tsigma = {std_spatial_spearman_corr}', file=f)
320    print(f'Spatial Pearson correlation \tmu = {mean_spatial_pearson_corr} \tsigma = {std_spatial_pearson_corr}', file=f)
321
322    if save_path is not None:
323        f.close()
324
325    timing.runtime()
326    return temp_rmse_map, temp_pearson_corrmap, nmeanbias