Select Git revision
graphics.py 9.83 KiB
#!/usr/bin/env python3
# -*- coding: utf-8 -*-
import numpy as np
from matplotlib import pyplot as p
import matplotlib.colors as colors
def plot_CBL_assimilation(exp,figtitle,which_cycle=0,ax=None):
naturerun_color = 'black'
ensmembers_color = 'lightskyblue'
ensmean_color = 'royalblue'
def ensplot(ax,x,x1,x2,z,colors=[ensmean_color,ensmembers_color],label=None):
nens = x.shape[1]
ax.plot(x[x1:x2].mean(axis=1),z,color=colors[0],zorder=5,label=label)
for i in range(nens): ax.plot(x[x1:x2,i],z,color=colors[1],alpha=0.2,zorder=1)
# Shorthand
o = exp.observations[which_cycle]
b = exp.da.backgrounds[which_cycle,:,:]
a = exp.da.analyses[which_cycle,:,:]
ocoords = exp.obs_coordinates
nz = exp.nr.nz
zt = exp.nr.zt
if exp.nr.is_bwind:
ii=3
else:
ii=1
zmax = 2000
if ax is None:
fig = p.figure(151)
fig.set_size_inches(9,3*ii)
#
# Potential temperature
#
ax1 = fig.add_subplot(ii,4,1)
ax1.scatter(o,ocoords,marker="o",s=20,color=naturerun_color,zorder=10)
ax1.set_title(r'$\theta$ (K), nature run')
ax1.set_ylabel(r'$z$ (m)')
ax1.set_ylim([0,zmax])
#
ax2 = fig.add_subplot(ii,4,2,sharey=ax1)
ensplot(ax2,b,0,nz,zt)
ax2.scatter(o,ocoords,marker="o",s=20,color=naturerun_color,zorder=10)
ax2.set_title(r'$\theta$ (K), prior')
#
ax3 = fig.add_subplot(ii,4,3,sharey=ax1)
ensplot(ax3,a,0,nz,zt)
ax3.scatter(o,ocoords,marker="o",s=20,color=naturerun_color,zorder=10)
ax3.set_title(r'$\theta$ (K), posterior')
#
ax4 = fig.add_subplot(ii,4,4,sharey=ax1)
ax4.plot(a[:nz].mean(axis=1)-b[:nz].mean(axis=1),zt,color=ensmean_color)
ax4.set_title(r'$\theta$ (K), mean increment',zorder=10)
#
# Fix axis labels
#
for ax in [ax2,ax3,ax4]:
p.setp(ax.get_yticklabels(), visible=False)
#
if exp.nr.is_bwind:
#
# u wind component
#
ax5 = fig.add_subplot(ii,4,5)
ax5.set_title(r'$u$ (m/s), nature run')
ax5.set_ylabel(r'$z$ (m)')
ax5.set_ylim([0,zmax])
#
ax6 = fig.add_subplot(ii,4,6,sharey=ax1)
ensplot(ax6,b,nz,nz*2,zt)
ax6.set_title(r'$u$ (m/s), prior')
#
ax7 = fig.add_subplot(ii,4,7,sharey=ax1)
ensplot(ax7,a,nz,nz*2,zt)
ax7.set_title(r'$u$ (m/s), posterior')
#
ax8 = fig.add_subplot(ii,4,8,sharey=ax1)
ax8.plot(a[nz:nz*2].mean(axis=1)-b[nz:nz*2].mean(axis=1),zt,color=ensmean_color)
ax8.set_title(r'$u$ (m/s), mean increment',zorder=10)
#
# v wind component
#
ax9 = fig.add_subplot(ii,4,9)
ax9.set_title(r'$v$ (m/s), nature run')
ax9.set_ylabel(r'$z$ (m)')
ax9.set_ylim([0,zmax])
#
axa = fig.add_subplot(ii,4,10,sharey=ax1)
ensplot(axa,b,nz*2,nz*3,zt)
axa.set_title(r'$v$ (m/s), prior')
#
axb = fig.add_subplot(ii,4,11,sharey=ax1)
ensplot(axb,a,nz*2,nz*3,zt)
axb.set_title(r'$v$ (m/s), posterior')
#
axc = fig.add_subplot(ii,4,12,sharey=ax1)
axc.plot(a[nz*2:nz*3].mean(axis=1)-b[nz*2:nz*3].mean(axis=1),zt,color=ensmean_color)
axc.set_title(r'$v$ (m/s), mean increment',zorder=10)
#
# Fix axis labels
#
for ax in [ax6,ax7,ax8,axa,axb,axc]:
p.setp(ax.get_yticklabels(), visible=False)
#
fig.savefig(figtitle,format='png',dpi=300)
p.close(fig)
else:
ensplot(ax,b,0,nz,zt,colors = ['red','salmon'],label='prior')
ensplot(ax,a,0,nz,zt,colors = ['royalblue','lightskyblue'],label='posterior')
ax.scatter(o,ocoords,marker="o",s=20,color=naturerun_color,zorder=10)
#ax.set_title(r'$\theta$ (K), posterior')
return ax
def plot_CBL_PE(exp,figtitle,parameter_id=0,plot_spread=False,ax=None):
if exp.nr.do_parameter_estimation:
spread_color = 'lightskyblue'
mean_color = 'royalblue'
mean_color_2 = 'orange'
if parameter_id==0:
true_value = exp.nr.pfac
pt = exp.nr.parameter_transform
parameter_number = exp.nr.parameter_number
b = exp.da.backgrounds[:,-parameter_number:,:][:,parameter_id,:]
par_b = pt[parameter_id](exp.da.backgrounds[:,-parameter_number:,:][:,parameter_id,:],kind='inv')
t = exp.da.time/3600.
if ax is None:
printfigure = True
else:
printfigure = False
if printfigure:
if plot_spread:
ii=2
else:
ii=1
fig = p.figure(151)
fig.set_size_inches(4,3*ii)
ax1 = fig.add_subplot(ii,1,1)
else:
ax1 = ax
ax1.step(t,par_b.mean(axis=1),color=mean_color)
for i in range(t.size-1):
bmin = np.percentile(par_b, 10, axis=1)[i+1]
bmax = np.percentile(par_b, 90, axis=1)[i+1]
ax1.fill_between(t[i:i+2],\
y1=[bmin,bmin],\
y2=[bmax,bmax],\
color=spread_color,
edgecolor=None,
alpha=0.3)
ax1.axhline(y=true_value,linestyle='--',color='black')
ax1.set_xlim([0,t.max()])
if printfigure and plot_spread:
ax2 = fig.add_subplot(ii,1,2)
ax2.plot(t,par_b.std(axis=1),color=mean_color,label='physical')
ax2.plot(t,b.std(axis=1),color=mean_color_2,label='transformed', linestyle='--')
ax2.legend()
ax2.set_title('Parameter spread')
if printfigure:
fig.tight_layout()
fig.savefig(figtitle,format='png',dpi=300)
p.close(fig)
else:
return ax1
def plot_CBL_identifiability(cbl,sigma_o,figtitle,ax=None):
# Read relevant dimensions
times = cbl.history['0000']['time']
zt = cbl.zt
nz = cbl.nz
ntimes = times.size
nens = cbl.nens
# Reconstruct_history
theta = np.zeros((nens,nz,ntimes))+np.nan
for k in range(nens):
theta[k,:,:]= cbl.history['%04u'%k]['theta']
# Read parameter values
if hasattr(cbl,'initial_perturbed_parameters'):
pfac = cbl.parameter_transform[0](cbl.initial_perturbed_parameters[0],kind='inv')
else:
pfac = np.ones(nens)*cbl.pfac
# Compute Kalman Gain
# (a little noise is added to the denominator to avoid singularities)
kalman_gain = np.zeros((nz,ntimes))+np.nan
for i in range(nz):
for j in range(ntimes):
covmat=np.cov(pfac,theta[:,i,j])
#p_x_correlation[i,j] = covmat[0,1] / (1e-9+np.sqrt(covmat[0,0]*covmat[1,1]))
kalman_gain[i,j] = covmat[0,1] / (1e-9+covmat[1,1]+sigma_o**2)
# Make plot
zmax = 2000
ncont = 13
if ax is None:
fig = p.figure(151)
fig.set_size_inches(9,3)
#
ax1 = fig.add_subplot(1,3,1)
c1 = ax1.pcolormesh(times,zt,theta.mean(axis=0),vmin=cbl.theta_0,vmax=cbl.theta_0+cbl.gamma*zmax)
ax1.contour(times,zt,theta.mean(axis=0),
np.linspace(cbl.theta_0,cbl.theta_0+cbl.gamma*zmax,ncont),
colors='white',
linestyles='--',
linewidths=0.75)
ax1.set_ylim([0,zmax])
ax1.set_xlabel(r'time')
ax1.set_ylabel(r'$z$ (m)')
ax1.set_title(r'ensemble mean $\theta$ (K)')
p.colorbar(c1)
#
ax2 = fig.add_subplot(1,3,2)
c2 = ax2.pcolormesh(times,zt,theta.std(axis=0))
ax2.set_ylim([0,zmax])
ax2.set_xlabel(r'time')
ax2.set_ylabel(r'$z$ (m)')
ax2.set_title(r'ensemble sdev $\theta$ (K)')
p.colorbar(c2)
#
ax3 = fig.add_subplot(1,3,3,sharey=ax1)
c3 = ax3.pcolormesh(times,zt,p_x_correlation,vmin=-1,vmax=1,cmap='RdBu_r')
ax3.set_xlabel(r'time')
ax3.set_title(r'$p$-$\theta$ correlation')
p.colorbar(c3)
#
p.setp(ax2.get_yticklabels(), visible=False)
p.setp(ax3.get_yticklabels(), visible=False)
fig.tight_layout()
fig.savefig(figtitle,format='png',dpi=300)
p.close(fig)
else:
c = ax.pcolormesh(times/3600.,zt,kalman_gain,norm=colors.CenteredNorm(),cmap='RdBu_r')
#c = ax.pcolormesh(times/3600.,zt,kalman_gain,vmin=-5, vmax=5,cmap='RdBu_r')
ax.contour(times/3600.,zt,theta.mean(axis=0),
np.linspace(cbl.theta_0,cbl.theta_0+cbl.gamma*zmax,ncont),
colors='black',
linestyles='--',
linewidths=0.75)
ax.set_ylim([0,zmax])
return ax,c
def plot_p(p_factors,theta_profiles,zt,figtitle,ax=None):
zoverh = np.linspace(0,1,101)
if ax is None:
fig = p.figure(151)
fig.set_size_inches(6,3)
#
ax1 = fig.add_subplot(1,2,1)
for pfac in p_factors:
Koverkws = zoverh*(1-zoverh)**pfac
ax1.plot(Koverkws,zoverh,label='$p=%4.1f$'%pfac)
ax1.set_xlabel('$K/(\kappa w_s z_i)$')
ax1.set_ylabel('$z/z_i$')
#
ax2 = fig.add_subplot(1,2,2)
for i in range(len(p_factors)):
ax2.plot(theta_profiles[i],zt,label='$p=%4.1f$'%p_factors[i])
ax2.set_xlabel(r'$\theta$ (K)')
ax2.set_ylabel(r'$z$ (m)')
ax2.set_ylim([0,1500])
ax2.set_xlim([291,297])
ax2.legend(loc=4)
#
fig.tight_layout()
fig.savefig(figtitle,format='png',dpi=300)
p.close(fig)
else:
for i in range(len(p_factors)):
ax.plot(theta_profiles[i],zt,label='$p=%4.1f$'%p_factors[i])
return ax