added code for figure 2

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]

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,6 +28,8 @@ def plot_CBL_assimilation(exp,figtitle,which_cycle=1):
         ii=1
 
     zmax = 2000
+
+    if ax is None:
         fig = p.figure(151)
         fig.set_size_inches(9,3*ii)
         #
@@ -106,7 +108,15 @@ def plot_CBL_assimilation(exp,figtitle,which_cycle=1):
         fig.savefig(figtitle,format='png',dpi=300)
         p.close(fig)
     
-def plot_CBL_PE(exp,figtitle,parameter_id=0,plot_spread=False):
+    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:
 
@@ -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 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]
@@ -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')
-        #
+
+        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):