| import concurrent
|
| import multiprocessing
|
| import netCDF4
|
| import os
|
|
|
| from matplotlib import gridspec, colors, pyplot as plt
|
| import numpy as np
|
| import torch
|
|
|
| from ..analysis import confmatrix, fss, rank
|
| from ..features import io
|
| from .cm import homeyer_rainbow
|
|
|
|
|
| def reverse_transform_R(R, mean=-0.051, std=0.528):
|
| return 10**(R*std + mean)
|
|
|
|
|
| def plot_precip_image(
|
| ax, R,
|
| Rmin=0, Rmax=25, threshold_mmh=0.1,
|
| transform_R=False,
|
| grid_spacing=64
|
| ):
|
| if isinstance(R, torch.Tensor):
|
| R = R.detach().numpy()
|
| if transform_R:
|
| R = reverse_transform_R(R, mean=mean, std=std)
|
| Rmin = reverse_transform_R(Rmin)
|
| Rmax = reverse_transform_R(Rmax)
|
| if threshold_mmh:
|
| Rmin = max(Rmin, threshold_mmh)
|
| R[R < threshold_mmh] = np.nan
|
| norm = colors.LogNorm(Rmin, Rmax)
|
| ax.set_yticks(np.arange(0, R.shape[0], grid_spacing))
|
| ax.set_xticks(np.arange(0, R.shape[1], grid_spacing))
|
| ax.grid(which='major', alpha=0.35)
|
| ax.tick_params(left=False, bottom=False,
|
| labelleft=False, labelbottom=False)
|
|
|
| return ax.imshow(R, norm=norm, cmap=homeyer_rainbow)
|
|
|
|
|
| def plot_autoencoder_reconstruction(
|
| R, R_hat, samples=8, timesteps=4,
|
| out_file=None
|
| ):
|
| fig = plt.figure(figsize=(2*timesteps*2+0.5, samples*2), dpi=150)
|
|
|
| gs = gridspec.GridSpec(
|
| samples, 2*timesteps+1,
|
| width_ratios=(1,)*(2*timesteps)+(0.2,),
|
| wspace=0.02, hspace=0.02
|
| )
|
| for k in range(samples):
|
| for (i,j) in enumerate(range(-timesteps,0)):
|
| ax = fig.add_subplot(gs[k,i])
|
| im = plot_precip_image(ax, R[k,0,j,:,:])
|
| for (i,j) in enumerate(range(-timesteps,0)):
|
| ax = fig.add_subplot(gs[k,i+timesteps])
|
| im = plot_precip_image(ax, R_hat[k,0,j,:,:])
|
|
|
| cax = fig.add_subplot(gs[:,-1])
|
| plt.colorbar(im, cax=cax)
|
|
|
| if out_file is not None:
|
| out_file = fig.savefig(out_file, bbox_inches='tight')
|
| plt.close(fig)
|
|
|
|
|
| def plot_animation(x, y, out_dir, sample=0, fmt="{}_{:02d}.png"):
|
| def plot_frame(R, label, timestep):
|
| fig = plt.figure()
|
| ax = fig.add_subplot()
|
| im = plot_precip_image(ax, R[sample,0,timestep,:,:])
|
| fn = fmt.format(label, k)
|
| fn = os.path.join(out_dir, fn)
|
| fig.savefig(fn, bbox_inches='tight')
|
| plt.close(fig)
|
|
|
| for k in range(x.shape[2]):
|
| plot_frame(x, "x", k)
|
|
|
| for k in range(y.shape[2]):
|
| plot_frame(y, "y", k)
|
|
|
|
|
| model_colors = {
|
| "mch-dgmr": "#E69F00",
|
| "mch-pysteps": "#009E73",
|
| "mch-iters=50-res=256": "#0072B2",
|
| "mch-persistence": "#888888",
|
| "dwd-dgmr": "#E69F00",
|
| "dwd-pysteps": "#009E73",
|
| "dwd-iters=50-res=256": "#0072B2",
|
| "dwd-persistence": "#888888",
|
| "pm-mch-dgmr": "#E69F00",
|
| "pm-mch-pysteps": "#009E73",
|
| "pm-mch-iters=50-res=256": "#0072B2",
|
| "pm-dwd-dgmr": "#E69F00",
|
| "pm-dwd-pysteps": "#009E73",
|
| "pm-dwd-iters=50-res=256": "#0072B2",
|
| }
|
|
|
| scale_linestyles = {
|
| "1x1": "-",
|
| "8x8": "--",
|
| "64x64": ":"
|
| }
|
|
|
|
|
| def plot_crps(
|
| log=False,
|
| models=("iters=50-res=256", "dgmr", "pysteps"),
|
| scales=("1x1", "8x8", "64x64"),
|
| model_labels=("LDCast", "DGMR", "PySTEPS", "Persist."),
|
| interval_mins=5,
|
| out_fn=None,
|
| ax=None,
|
| add_xlabel=True,
|
| add_ylabel=True,
|
| add_legend=True,
|
| crop_box=None
|
| ):
|
| crps = {}
|
| crps_name = "logcrps" if log else "crps"
|
| for model in models:
|
| crps[model] = {}
|
| fn = f"../results/crps/{crps_name}-{model}.nc"
|
| with netCDF4.Dataset(fn, 'r') as ds:
|
| for scale in scales:
|
| var = f"crps_pool{scale}"
|
| crps_model_scale = np.array(ds[var][:], copy=False)
|
| if crop_box is not None:
|
| scale_int = int(scale.split("x")[0])
|
| crps_model_scale = crps_model_scale[
|
| ...,
|
| crop_box[0][0]//scale_int:crop_box[0][1]//scale_int,
|
| crop_box[1][0]//scale_int:crop_box[1][1]//scale_int
|
| ]
|
| crps[model][scale] = crps_model_scale.mean(axis=(0,1,3,4))
|
| del crps_model_scale
|
|
|
| if ax is None:
|
| fig = plt.figure(figsize=(8,5))
|
| ax = fig.add_subplot()
|
|
|
| max_t = 0
|
| for (model, label) in zip(models, model_labels):
|
| for scale in scales:
|
| score = crps[model][scale]
|
| color = model_colors[model]
|
| linestyle = scale_linestyles[scale]
|
| t = np.arange(
|
| interval_mins, (len(score)+0.1)*interval_mins, interval_mins
|
| )
|
| max_t = max(max_t, t[-1])
|
| ax.plot(t, score, color=color, linestyle=linestyle,
|
| label=label)
|
|
|
| if add_legend:
|
| plt.legend()
|
| if add_xlabel:
|
| plt.xlabel("Lead time [min]", fontsize=12)
|
| if add_ylabel:
|
| plt.ylabel(
|
| "LogCRPS" if log else "CRPS [mm h$^\\mathrm{-1}$]",
|
| fontsize=12
|
| )
|
|
|
| ax.set_xlim((0, max_t))
|
| ylim = ax.get_ylim()
|
| ylim = (0, ylim[1])
|
| ax.set_ylim(ylim)
|
| ax.tick_params(axis='both', which='major', labelsize=12)
|
|
|
| if out_fn is not None:
|
| fig.savefig(out_fn, bbox_inches='tight')
|
| plt.close(fig)
|
|
|
|
|
| def plot_rank_distribution(
|
| models=("iters=50-res=256", "dgmr", "pysteps"),
|
| scales=("1x1", "8x8", "64x64"),
|
| model_labels=("LDCast", "DGMR", "PySTEPS"),
|
| out_fn=None,
|
| num_ensemble_members=32,
|
| ax=None,
|
| add_xlabel=True,
|
| add_ylabel=True,
|
| add_legend=True,
|
| crop_box=None
|
| ):
|
| rank_hist = {}
|
| rank_KL = {}
|
| for model in models:
|
| fn = f"../results/ranks/ranks-{model}.nc"
|
| with netCDF4.Dataset(fn, 'r') as ds:
|
| for scale in scales:
|
| var = f"ranks_pool{scale}"
|
| ranks_model_scale = np.array(ds[var][:], copy=False)
|
| if crop_box is not None:
|
| scale_int = int(scale.split("x")[0])
|
| ranks_model_scale = ranks_model_scale[
|
| ...,
|
| crop_box[0][0]//scale_int:crop_box[0][1]//scale_int,
|
| crop_box[1][0]//scale_int:crop_box[1][1]//scale_int
|
| ]
|
| rank_hist[(model,scale)] = rank.rank_distribution(ranks_model_scale)
|
| rank_KL[(model,scale)] = rank.rank_DKL(rank_hist[(model,scale)])
|
| del ranks_model_scale
|
|
|
| if ax is None:
|
| fig = plt.figure(figsize=(8,5))
|
| ax = fig.add_subplot()
|
|
|
| for scale in scales:
|
| linestyle = scale_linestyles[scale]
|
| for (model, label) in zip(models, model_labels):
|
| h = rank_hist[(model,scale)]
|
| color = model_colors[model]
|
| x = np.linspace(0, 1, num_ensemble_members+1)
|
| label_with_score = f"{label}: {rank_KL[(model,scale)]:.3f}"
|
| ax.plot(x, h, color=color, linestyle=linestyle,
|
| label=label_with_score)
|
| h_ideal = 1/(num_ensemble_members+1)
|
| ax.plot([0, 1], [h_ideal, h_ideal], color=(0.4,0.4,0.4),
|
| linewidth=1.0)
|
|
|
| if add_legend:
|
| ax.legend(loc='upper center')
|
| if add_xlabel:
|
| ax.set_xlabel("Normalized rank", fontsize=12)
|
| if add_ylabel:
|
| ax.set_ylabel("Occurrence", fontsize=12)
|
|
|
| ax.set_xlim((0, 1))
|
| ylim = ax.get_ylim()
|
| ylim = (0, ylim[1])
|
| ax.set_ylim(ylim)
|
| ax.tick_params(axis='both', which='major', labelsize=12)
|
| ax.set_xticks([0, 0.25, 0.5, 0.75, 1])
|
|
|
| ax.set_xticklabels(["0", "0.25", "0.5", "0.75", "1"])
|
|
|
| if out_fn is not None:
|
| fig.savefig(out_fn, bbox_inches='tight')
|
| plt.close(fig)
|
|
|
|
|
| def plot_rank_metric(
|
| models=("iters=50-res=256", "dgmr", "pysteps"),
|
| scales=("1x1", "8x8", "64x64"),
|
| model_labels=("LDCast", "DGMR", "PySTEPS"),
|
| interval_mins=5,
|
| metric_name="KL",
|
| out_fn=None,
|
| ax=None,
|
| add_xlabel=True,
|
| add_ylabel=True,
|
| add_legend=True,
|
| ):
|
| rank_metric = {}
|
| for model in models:
|
| rank_metric[model] = {}
|
| fn = f"../results/ranks/ranks-{model}.nc"
|
| with netCDF4.Dataset(fn, 'r') as ds:
|
| for scale in scales:
|
| var = f"ranks_pool{scale}"
|
| ranks = np.array(ds[var][:], copy=False)
|
| rank_metric[model][scale] = rank.rank_metric_by_leadtime(ranks)
|
| del ranks
|
|
|
| if ax is None:
|
| fig = plt.figure(figsize=(8,5))
|
| ax = fig.add_subplot()
|
|
|
| max_t = 0
|
| for (model, label) in zip(models, model_labels):
|
| for scale in scales:
|
| score = rank_metric[model][scale]
|
| color = model_colors[model]
|
| linestyle = scale_linestyles[scale]
|
| label_with_scale = f"{label} {scale}"
|
| t = np.arange(
|
| interval_mins, (len(score)+0.1)*interval_mins, interval_mins
|
| )
|
| max_t = max(max_t, t[-1])
|
| ax.plot(t, score, color=color, linestyle=linestyle,
|
| label=label_with_scale)
|
|
|
| if add_legend:
|
| plt.legend()
|
| if add_xlabel:
|
| plt.xlabel("Lead time [min]", fontsize=12)
|
| if add_ylabel:
|
| plt.ylabel(f"Rank {metric_name}", fontsize=12)
|
|
|
| ax.set_xlim((0, max_t))
|
| ylim = ax.get_ylim()
|
| ylim = (0, ylim[1])
|
| ax.set_ylim(ylim)
|
| ax.tick_params(axis='both', which='major', labelsize=12)
|
|
|
| if out_fn is not None:
|
| fig.savefig(out_fn, bbox_inches='tight')
|
| plt.close(fig)
|
|
|
|
|
| def load_fss(model, scale, use_timesteps, crop_box):
|
| fn = f"../results/fractions/fractions-{model}.nc"
|
| with netCDF4.Dataset(fn, 'r') as ds:
|
| sn = f"{scale}x{scale}"
|
| obs_frac = np.array(ds[f"obs_frac_scale{sn}"][:], copy=False)
|
| fc_frac = np.array(ds[f"fc_frac_scale{sn}"][:], copy=False)
|
| if crop_box is not None:
|
| obs_frac = obs_frac[
|
| ...,
|
| crop_box[0][0]//scale:crop_box[0][1]//scale,
|
| crop_box[1][0]//scale:crop_box[1][1]//scale
|
| ]
|
| fc_frac = fc_frac[
|
| ...,
|
| crop_box[0][0]//scale:crop_box[0][1]//scale,
|
| crop_box[1][0]//scale:crop_box[1][1]//scale
|
| ]
|
| return fss.fractions_skill_score(
|
| obs_frac, fc_frac, use_timesteps=use_timesteps
|
| )
|
|
|
|
|
| def plot_fss(
|
| log=False,
|
| models=("iters=50-res=256", "dgmr", "pysteps"),
|
| model_labels=("LDCast", "DGMR", "PySTEPS"),
|
| interval_mins=5,
|
| out_fn=None,
|
| ax=None,
|
| add_xlabel=True,
|
| add_ylabel=True,
|
| add_legend=True,
|
| scales=None,
|
| use_timesteps=18,
|
| crop_box=None
|
| ):
|
| if scales is None:
|
| scales = 2**np.arange(9)
|
| fss_scale = {}
|
| N_threads = min(multiprocessing.cpu_count(), len(models)*len(scales))
|
| with concurrent.futures.ProcessPoolExecutor(N_threads) as executor:
|
| for model in models:
|
| fss_scale[model] = {}
|
| for scale in scales:
|
| fss_scale[model][scale] = executor.submit(
|
| load_fss, model, scale, use_timesteps, crop_box
|
| )
|
|
|
| for model in models:
|
| for scale in scales:
|
| fss_scale[model][scale] = fss_scale[model][scale].result()
|
|
|
| if ax is None:
|
| fig = plt.figure(figsize=(8,5))
|
| ax = fig.add_subplot()
|
|
|
| for (model, label) in zip(models, model_labels):
|
| scales = sorted(fss_scale[model])
|
| fss_for_model = [fss_scale[model][s] for s in scales]
|
|
|
| model_parts = model.split("-")
|
| if model.startswith("pm-"):
|
| model_without_threshold = "-".join(model_parts[:1] + model_parts[2:])
|
| else:
|
| model_without_threshold = "-".join(model_parts[1:])
|
| color = model_colors[model_without_threshold]
|
|
|
| ax.plot(scales, fss_for_model, color=color,
|
| label=label)
|
|
|
| if add_legend:
|
| plt.legend()
|
| if add_xlabel:
|
| plt.xlabel("Scale [km]", fontsize=12)
|
| if add_ylabel:
|
| plt.ylabel("FSS")
|
|
|
| ax.set_xlim((scales[0], scales[-1]))
|
| ylim = ax.get_ylim()
|
| ylim = (0, ylim[1])
|
| ax.set_ylim(ylim)
|
| ax.tick_params(axis='both', which='major', labelsize=12)
|
|
|
| if out_fn is not None:
|
| fig.savefig(out_fn, bbox_inches='tight')
|
| plt.close(fig)
|
|
|
|
|
| def plot_csi_threshold(
|
| models=("iters=50-res=256", "dgmr", "pysteps"),
|
| scales=("1x1", "8x8", "64x64"),
|
| prob_thresholds=tuple(np.linspace(0,1,33)),
|
| model_labels=("LDCast", "DGMR", "PySTEPS"),
|
| out_fn=None,
|
| num_ensemble_members=32,
|
| max_timestep=18,
|
| ax=None,
|
| add_xlabel=True,
|
| add_ylabel=True,
|
| add_legend=True,
|
| crop_box=None
|
| ):
|
| csi = {}
|
| for model in models:
|
| fn = f"../results/fractions/fractions-{model}.nc"
|
| with netCDF4.Dataset(fn, 'r') as ds:
|
| for scale in scales:
|
| fc_var = f"fc_frac_scale{scale}"
|
| fc_frac = np.array(ds[fc_var], copy=False)
|
| fc_frac = fc_frac[...,:max_timestep,:,:]
|
| obs_var = f"obs_frac_scale{scale}"
|
| obs_frac = np.array(ds[obs_var], copy=False)
|
| obs_frac = obs_frac[...,:max_timestep,:,:]
|
| conf_matrix = confmatrix.confusion_matrix_thresholds(
|
| fc_frac, obs_frac, prob_thresholds
|
| )
|
| del fc_frac, obs_frac
|
| csi_scale = confmatrix.intersection_over_union(conf_matrix)
|
| csi[(model,scale)] = csi_scale
|
|
|
| if ax is None:
|
| fig = plt.figure(figsize=(8,5))
|
| ax = fig.add_subplot()
|
|
|
| for scale in scales:
|
| linestyle = scale_linestyles[scale]
|
| for (model, label) in zip(models, model_labels):
|
| c = csi[(model,scale)]
|
| model_parts = model.split("-")
|
| if model.startswith("pm-"):
|
| model_without_threshold = "-".join(model_parts[:1] + model_parts[2:])
|
| else:
|
| model_without_threshold = "-".join(model_parts[1:])
|
| color = model_colors[model_without_threshold]
|
| ax.plot(prob_thresholds, c, color=color, linestyle=linestyle, label=label)
|
|
|
| if add_legend:
|
| ax.legend(loc='upper center')
|
| if add_xlabel:
|
| ax.set_xlabel("Prob. threshold", fontsize=12)
|
| if add_ylabel:
|
| ax.set_ylabel("CSI", fontsize=12)
|
|
|
| ax.set_xlim((0, 1))
|
| ylim = ax.get_ylim()
|
| ylim = (0, ylim[1])
|
| ax.set_ylim(ylim)
|
| ax.tick_params(axis='both', which='major', labelsize=12)
|
| ax.set_xticks([0, 0.25, 0.5, 0.75, 1])
|
|
|
| ax.set_xticklabels(["0", "0.25", "0.5", "0.75", "1"])
|
|
|
| if out_fn is not None:
|
| fig.savefig(out_fn, bbox_inches='tight')
|
| plt.close(fig)
|
|
|
|
|
| def plot_csi_leadtime(
|
| models=("iters=50-res=256", "dgmr", "pysteps"),
|
| scales=("1x1", "8x8", "64x64"),
|
| prob_thresholds=tuple(np.linspace(0,1,33)),
|
| model_labels=("LDCast", "DGMR", "PySTEPS"),
|
| out_fn=None,
|
| interval_mins=5,
|
| num_ensemble_members=32,
|
| ax=None,
|
| add_xlabel=True,
|
| add_ylabel=True,
|
| add_legend=True,
|
| crop_box=None
|
| ):
|
| csi = {}
|
| for model in models:
|
| fn = f"../results/fractions/fractions-{model}.nc"
|
| with netCDF4.Dataset(fn, 'r') as ds:
|
| for scale in scales:
|
| fc_var = f"fc_frac_scale{scale}"
|
| fc_frac = np.array(ds[fc_var], copy=False)
|
| obs_var = f"obs_frac_scale{scale}"
|
| obs_frac = np.array(ds[obs_var], copy=False)
|
| conf_matrix = confmatrix.confusion_matrix_thresholds_leadtime(
|
| fc_frac, obs_frac, prob_thresholds
|
| )
|
|
|
| csi_scale = confmatrix.intersection_over_union(conf_matrix)
|
| csi[(model,scale)] = np.nanmax(csi_scale, axis=1)
|
|
|
| max_t = 0
|
| for (model, label) in zip(models, model_labels):
|
| for scale in scales:
|
| score = csi[(model,scale)]
|
|
|
| model_parts = model.split("-")
|
| if model.startswith("pm-"):
|
| model_without_threshold = "-".join(model_parts[:1] + model_parts[2:])
|
| else:
|
| model_without_threshold = "-".join(model_parts[1:])
|
| color = model_colors[model_without_threshold]
|
| linestyle = scale_linestyles[scale]
|
| t = np.arange(
|
| interval_mins, (len(score)+0.1)*interval_mins, interval_mins
|
| )
|
| max_t = max(max_t, t[-1])
|
| ax.plot(t, score, color=color, linestyle=linestyle,
|
| label=label)
|
|
|
| if add_legend:
|
| plt.legend()
|
| if add_xlabel:
|
| plt.xlabel("Lead time [min]", fontsize=12)
|
| if add_ylabel:
|
| plt.ylabel("CSI", fontsize=12)
|
|
|
| ax.set_xlim((0, max_t))
|
| ylim = ax.get_ylim()
|
| ylim = (0, ylim[1])
|
| ax.set_ylim(ylim)
|
| ax.tick_params(axis='both', which='major', labelsize=12)
|
|
|
| if out_fn is not None:
|
| fig.savefig(out_fn, bbox_inches='tight')
|
| plt.close(fig)
|
|
|
|
|
| def plot_cost_loss_value(
|
| models=("iters=50-res=256", "dgmr", "pysteps"),
|
| scales=("1x1", "8x8", "64x64"),
|
| prob_thresholds=tuple(np.linspace(0,1,33)),
|
| model_labels=("LDCast", "DGMR", "PySTEPS"),
|
| out_fn=None,
|
| interval_mins=5,
|
| num_ensemble_members=32,
|
| ax=None,
|
| add_xlabel=True,
|
| add_ylabel=True,
|
| add_legend=True,
|
| crop_box=None
|
| ):
|
| value = {}
|
| loss = 1.0
|
| cost = np.linspace(0.01, 1, 100)
|
| for model in models:
|
| fn = f"../results/fractions/fractions-{model}.nc"
|
| with netCDF4.Dataset(fn, 'r') as ds:
|
| for scale in scales:
|
| fc_var = f"fc_frac_scale{scale}"
|
| fc_frac = np.array(ds[fc_var], copy=False)
|
| obs_var = f"obs_frac_scale{scale}"
|
| obs_frac = np.array(ds[obs_var], copy=False)
|
| conf_matrix = confmatrix.confusion_matrix_thresholds(
|
| fc_frac, obs_frac, prob_thresholds
|
| )
|
|
|
| p_clim = obs_frac.mean()
|
| value_scale = []
|
| for c in cost:
|
| v = confmatrix.cost_loss_value(
|
| conf_matrix, c, loss, p_clim
|
| )
|
| value_scale.append(v[len(v)//2])
|
| value[(model,scale)] = np.array(value_scale)
|
|
|
| max_score = 0
|
| for (model, label) in zip(models, model_labels):
|
| for scale in scales:
|
| score = value[(model,scale)]
|
| max_score = max(max_score, score[np.isfinite(score)].max())
|
|
|
| model_parts = model.split("-")
|
| if model.startswith("pm-"):
|
| model_without_threshold = "-".join(model_parts[:1] + model_parts[2:])
|
| else:
|
| model_without_threshold = "-".join(model_parts[1:])
|
| color = model_colors[model_without_threshold]
|
| linestyle = scale_linestyles[scale]
|
|
|
| ax.plot(cost, score, color=color, linestyle=linestyle,
|
| label=label)
|
|
|
| if add_legend:
|
| plt.legend()
|
| if add_xlabel:
|
| plt.xlabel("Cost/loss ratio", fontsize=12)
|
| if add_ylabel:
|
| plt.ylabel("Value", fontsize=12)
|
|
|
| ax.set_xlim((0, 1))
|
| ylim = (0, max_score*1.05)
|
| ax.set_ylim(ylim)
|
| ax.tick_params(axis='both', which='major', labelsize=12)
|
| ax.set_xticks([0, 0.25, 0.5, 0.75, 1])
|
|
|
| ax.set_xticklabels(["0", "0.25", "0.5", "0.75", "1"])
|
|
|
| if out_fn is not None:
|
| fig.savefig(out_fn, bbox_inches='tight')
|
| plt.close(fig)
|
|
|
