From a4d42bb09a22c3e8ae790aebe5376835b229591c Mon Sep 17 00:00:00 2001
From: Stefano Serafin <serafin@jet02.img.univie.ac.at>
Date: Tue, 9 Jul 2024 15:03:14 +0200
Subject: [PATCH] added code for figure 2
---
PE_CBL.py | 62 +++++++++++++++--
graphics.py | 194 +++++++++++++++++++++++++++++-----------------------
2 files changed, 165 insertions(+), 91 deletions(-)
diff --git a/PE_CBL.py b/PE_CBL.py
index d9731a6..04c2e7c 100644
--- a/PE_CBL.py
+++ b/PE_CBL.py
@@ -38,7 +38,7 @@ if __name__ == '__main__':
# perturbation_type can be "smooth", "random", "uniform"
# the adjective refers to vertical variability
'perturb_ensemble_state' : True,
- 'perturbations_type' : "smooth",
+ 'perturbations_type' : "uniform",
'perturbations_theta_amplitude' : .1,
'perturbations_uv_amplitude' : 0.1,
'perturbations_smooth_number' : 5,
@@ -73,7 +73,7 @@ if __name__ == '__main__':
cbl_det.run(output_full_history=True)
# Panel b) ensemble run
- cbl_settings['perturbations_theta_amplitude'] = 0.1
+ cbl_settings['perturbations_theta_amplitude'] = 1.0
cbl_settings['perturbations_type'] = "uniform"
cbl_settings['perturb_ensemble_state'] = True
cbl_settings['perturb_ensemble_parameters'] = True
@@ -119,12 +119,12 @@ if __name__ == '__main__':
linestyles='--',
linewidths=0.75)
- ax2,c2 = plot_CBL_identifiability(cbl_ens,None,ax2)
+ ax2,c2 = plot_CBL_identifiability(cbl_ens,None,ax=ax2)
ax2.set_xlabel(r'time (h)')
ax2.set_title(r'b) $\rho_{p,\theta}$')
p.colorbar(c2,orientation='horizontal')
- ax3 = plot_p(p_factors,theta_profiles,cbl_pf.zt,None,ax3)
+ ax3 = plot_p(p_factors,theta_profiles,cbl_pf.zt,None,ax=ax3)
ax3.set_xlabel(r'$\theta$ (K)')
ax3.set_xlim([291,297])
ax3.set_ylim([0,zmax])
@@ -136,6 +136,60 @@ if __name__ == '__main__':
fig.savefig('fig01.png',format='png',dpi=300)
p.close(fig)
+ elif CASE == 'fig02': # parameter estimation
+
+ cbl_settings['perturbations_theta_amplitude'] = 1.0
+ cbl_settings['perturbations_type'] = "uniform"
+ cbl_settings['perturb_ensemble_state'] = True
+ cbl_settings['perturb_ensemble_parameters'] = True
+ cbl_settings['perturbations_symmetric'] = True
+ cbl_settings['parameter_number'] = 1
+ cbl_settings['parameter_transform'] = [ parameter_transform_log ]
+ cbl_settings['parameter_ensemble_min'] = np.array([0.5])
+ cbl_settings['parameter_ensemble_max'] = np.array([4.5])
+ cbl_settings['parameter_true'] = np.array([1.5])
+ cbl_settings['do_parameter_estimation'] = True
+ cbl_settings['parameter_inflation_rtps_alpha'] = np.array([1])
+
+ nobs = 16
+ da_settings = {'cbl_settings' : cbl_settings,
+ 'tspinup' : 3600,
+ 'trun' : 7200,
+ 'assimilation_interval' : 300,
+ 'obs_coordinates' : np.linspace(0,2000,nobs),
+ 'obs_kinds' : np.array(['theta']*nobs),
+ 'obs_error_sdev_generate' : np.ones(nobs)*.0,
+ 'obs_error_sdev_assimilate' : np.ones(nobs)*.1,
+ 'localization_cutoff' : 400,
+ 'nens' : 100,
+ 'FILTER' : 'EAKF',
+ 'inflation_rtps_alpha' : 0.5,
+ }
+
+ # Safety check on time steps
+ if np.mod(da_settings['assimilation_interval'],integration_dt)==0:
+ pass
+ else:
+ raise ValueError("Assimilation interval must be an integer multiplier of model dt")
+
+ # Run experiment
+ exp = experiment(da_settings)
+
+ # Make plots
+ which_cycle = 0
+ tt = (which_cycle+1)*da_settings['assimilation_interval']
+ fig, [ax1, ax2] = p.subplots(1,2,constrained_layout=True)
+ fig.set_size_inches(6,3)
+ ax1 = plot_CBL_assimilation(exp,None,which_cycle=which_cycle,ax=ax1)
+ ax1.set_ylim([0,2000])
+ ax1.set_xlim([288,300])
+ ax1.set_ylabel(r'height (m)')
+ ax1.set_title(r'a) Analysis at $t=$%u s'%tt)
+ ax1.legend(frameon=False)
+ ax2 = plot_CBL_PE(exp,None,ax=ax2)
+ fig.savefig('fig02.png',format='png',dpi=300)
+ p.close(fig)
+
elif CASE == 0: # plot K profile examples
p_factors = [0.2,0.5,1,2,4,6]
diff --git a/graphics.py b/graphics.py
index 1bd2deb..b23ac17 100644
--- a/graphics.py
+++ b/graphics.py
@@ -3,16 +3,16 @@
import numpy as np
from matplotlib import pyplot as p
-def plot_CBL_assimilation(exp,figtitle,which_cycle=1):
+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):
+ 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=ensmean_color,zorder=5)
- for i in range(nens): ax.plot(x[x1:x2,i],z,color=ensmembers_color,alpha=0.2,zorder=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]
@@ -28,85 +28,95 @@ def plot_CBL_assimilation(exp,figtitle,which_cycle=1):
ii=1
zmax = 2000
- 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)
+
+ if ax is None:
+ fig = p.figure(151)
+ fig.set_size_inches(9,3*ii)
#
- # v wind component
+ # Potential temperature
#
- 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])
+ 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])
#
- 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')
+ 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')
#
- 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')
+ 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')
#
- 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)
+ 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 [ax6,ax7,ax8,axa,axb,axc]:
+ for ax in [ax2,ax3,ax4]:
p.setp(ax.get_yticklabels(), visible=False)
- #
- fig.savefig(figtitle,format='png',dpi=300)
- p.close(fig)
+ #
+ 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):
+def plot_CBL_PE(exp,figtitle,parameter_id=0,plot_spread=False,ax=None):
if exp.nr.do_parameter_estimation:
@@ -123,15 +133,22 @@ def plot_CBL_PE(exp,figtitle,parameter_id=0,plot_spread=False):
par_b = pt[parameter_id](exp.da.backgrounds[:,-parameter_number:,:][:,parameter_id,:],kind='inv')
t = exp.da.time/3600.
- if plot_spread:
- ii=2
+ if ax is None:
+ printfigure = True
else:
- ii=1
+ 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
- fig = p.figure(151)
- fig.set_size_inches(4,3*ii)
- #
- ax1 = fig.add_subplot(ii,1,1)
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]
@@ -143,20 +160,23 @@ def plot_CBL_PE(exp,figtitle,parameter_id=0,plot_spread=False):
edgecolor=None,
alpha=0.3)
ax1.axhline(y=true_value,linestyle='--',color='black')
- ax1.set_title('Parameter evolution')
+ ax1.set_title(r'b) Evolution of $p$')
ax1.set_xlabel('Time (h)')
ax1.set_xlim([0,t.max()])
- #
- if plot_spread:
+
+ 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')
- #
- fig.tight_layout()
- fig.savefig(figtitle,format='png',dpi=300)
- p.close(fig)
+
+ if printfigure:
+ fig.tight_layout()
+ fig.savefig(figtitle,format='png',dpi=300)
+ p.close(fig)
+ else:
+ return ax1
def plot_CBL_identifiability(cbl,figtitle,ax=None):
--
GitLab