additional plots: B matrix, Desrozier statistics, overview diagrams,...

additional plots: B matrix, Desrozier statistics, overview diagrams, state/parameter linearity; + minor tweaks

graphics.py

@@ -130,7 +130,7 @@ def plot_CBL_assimilation(exp,figtitle,which_cycle=0,ax=None):
         #ax.set_title(r'$\theta$ (K), posterior')
         return ax
 
-def plot_CBL_PE(exp,figtitle,parameter_id=0,plot_spread=False,ax=None):
+def plot_CBL_PE(exp,figtitle,parameter_id=0,show_spread=False,show_true_parameter=False,show_transformed_parameter=False,plot_parameter_histogram = False,ax=None):
 
     if exp.nr.do_parameter_estimation:
 
@@ -138,17 +138,18 @@ def plot_CBL_PE(exp,figtitle,parameter_id=0,plot_spread=False,ax=None):
         mean_color = 'royalblue'
         mean_color_2 = 'orange'
 
+        pt = exp.nr.parameter_transform
         if parameter_id==0:
             true_value = exp.nr.pfac
+            true_value_transformed = pt[0](exp.nr.pfac,kind='dir')
 
-        pt = exp.nr.parameter_transform
         parameter_number = exp.nr.parameter_number
         par_tran = exp.da.backgrounds[:,-parameter_number:,:][:,parameter_id,:]
         par_phys = pt[parameter_id](par_tran,kind='inv')
         initpar_phys = pt[parameter_id](exp.da.initial_perturbed_parameters,kind='inv')
+        initpar_tran = exp.da.initial_perturbed_parameters
         t = exp.da.time/3600.
 
-        plot_parameter_histogram = True
         if plot_parameter_histogram:
             p10beg_tran = np.percentile(par_tran[0,:],  10)
             p90beg_tran = np.percentile(par_tran[0,:],  90)
@@ -190,7 +191,7 @@ def plot_CBL_PE(exp,figtitle,parameter_id=0,plot_spread=False,ax=None):
             printfigure = False
 
         if printfigure:
-            if plot_spread:
+            if show_spread:
                 ii=2
             else:
                 ii=1
@@ -199,34 +200,69 @@ def plot_CBL_PE(exp,figtitle,parameter_id=0,plot_spread=False,ax=None):
             ax1 = fig.add_subplot(ii,1,1)
         else:
             ax1 = ax
+        
+        if show_transformed_parameter:
+            par =  par_tran
+            initpar = initpar_tran
+        else:
+            par = par_phys
+            initpar = initpar_phys
 
         # Initial parameter distribution
-        ax1.step([0,t[0]],np.median(initpar_phys,axis=1)*np.array([1,1]),color=mean_color)
-        ax1.step([0,t[0]],np.mean(initpar_phys,axis=1)*np.array([1,1]),color=mean_color_2)
-        bmin = np.percentile(initpar_phys, 10)
-        bmax = np.percentile(initpar_phys, 90)
+        ax1.step([0,t[0]],np.median(initpar,axis=1)*np.array([1,1]),color=mean_color,label=r'ensemble median')
+        ax1.step([0,t[0]],np.mean(initpar,axis=1)*np.array([1,1]),color=mean_color_2,label=r'ensemble mean, $\mu$')
+        bmin = np.percentile(initpar, 10)
+        bmax = np.percentile(initpar, 90)
         ax1.fill_between([0,t[0]],\
                     y1=[bmin,bmin],\
                     y2=[bmax,bmax],\
                     color=spread_color,
                     edgecolor=None,
-                    alpha=0.3)
+                    alpha=0.3,
+                    label=r'10-90 percentile range')
         # Later times
-        ax1.step(t,np.median(par_phys,axis=1),color=mean_color)
-        ax1.step(t,np.mean(par_phys,axis=1),color=mean_color_2)
+        ax1.step(t,np.median(par,axis=1),color=mean_color)
+        ax1.step(t,np.mean(par,axis=1),color=mean_color_2)
         for i in range(t.size-1):
-            bmin = np.percentile(par_phys, 10, axis=1)[i+1]
-            bmax = np.percentile(par_phys, 90, axis=1)[i+1]
+            bmin = np.percentile(par, 10, axis=1)[i+1]
+            bmax = np.percentile(par, 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')
+        if show_true_parameter:
+            if show_transformed_parameter:
+                ax1.axhline(y=true_value_transformed,linestyle='--',color='black')
+            else:
+                ax1.axhline(y=true_value,linestyle='--',color='black')
         ax1.set_xlim([0,t.max()])
 
-        if printfigure and plot_spread:
+        if (not printfigure) and show_spread:
+            bmin = (np.mean(initpar)-np.std(initpar))
+            bmax = (np.mean(initpar)+np.std(initpar))
+            ax1.fill_between([0,t[0]],\
+                        y1=[bmin,bmin],\
+                        y2=[bmax,bmax],\
+                        color='lightgrey',
+                        edgecolor=None,
+                        alpha=0.3,
+                        zorder=10,
+                        label=r'$\mu\pm\sigma$')
+            for i in range(t.size-1):
+                bmin = (np.mean(par,axis=1)-np.std(par,axis=1))[i+1]
+                bmax = (np.mean(par,axis=1)+np.std(par,axis=1))[i+1]
+                ax1.fill_between(t[i:i+2],\
+                            y1=[bmin,bmin],\
+                            y2=[bmax,bmax],\
+                            color='lightgrey',
+                            edgecolor=None,
+                            alpha=0.3,
+                            zorder=10)
+        ax1.legend(loc=1,frameon=False,fontsize=6)
+
+        if printfigure and show_spread:
             ax2 = fig.add_subplot(ii,1,2)
             ax2.plot(t,par_phys.std(axis=1),color=mean_color,label='physical')
             ax2.plot(t,par_tran.std(axis=1),color=mean_color_2,label='transformed', linestyle='--')
@@ -261,39 +297,46 @@ def plot_CBL_identifiability(cbl,sigma_o,ax=None):
     times = times[:,None]+np.zeros((ntimes,nobs))
 
     # Read model equivalents and state vector elements
-    theta = cbl.da.model_equivalents  # ntim*nobs*nens
+    #theta = cbl.da.model_equivalents  # ntim*nobs*nens
     pfac = cbl.da.backgrounds[:,-npar,:] # ntim*nx*nens
     pfac = cbl.nr.parameter_transform[0](pfac,kind='inv')
 
     # Read observation coordinates
     if cbl.da.obs_coordinates.ndim == 1:
-        zobs = cbl.da.obs_coordinates
-        zobs = zobs[None,:]+np.zeros(times.shape)
+        zz = cbl.da.obs_coordinates
+        zz = zz[None,:]+np.zeros(times.shape)
     elif cbl.da.obs_coordinates.ndim == 2:
-        zobs = cbl.da.obs_coordinates
+        zz = cbl.da.obs_coordinates
 
     # At each time, check if the coordinate array is sorted
-    # if not, resort both coordinates and model equivalents
+    # if not, resort it.
+    theta      = np.zeros((ntimes,nobs,nens))+np.nan
+    zobs       = np.zeros((ntimes,nobs))+np.nan
+    innov      = np.zeros((ntimes,nobs))+np.nan
+    cov_yy     = np.zeros((ntimes,nobs))+np.nan
+    cov_pp     = np.zeros((ntimes,nobs,npar))+np.nan
+    cov_py     = np.zeros((ntimes,nobs,npar))+np.nan
+    increments = np.zeros((ntimes,nobs,npar))+np.nan
+
     for k in range(ntimes):
-        if is_sorted(zobs[k,:]):
-            pass
+        if is_sorted(zz[k,:]):
+            indices = np.arange(nobs)
         else:
-            indices = np.argsort(zobs[k,:])
-            zobs[k,:] = zobs[k,indices]
-            theta[k,:] = theta[k,indices]
-            cbl.da.innov[k,:] = cbl.da.innov[k,indices] # ntim*nobs
-            cbl.da.cov_yy[k,:] = cbl.da.cov_yy[k,indices] # ntim*nobs
-            cbl.da.cov_pp[k,:] = cbl.da.cov_pp[k,indices] # ntim*nobs(*npar)
-            cbl.da.cov_py[k,:] = cbl.da.cov_py[k,indices] # ntim*nobs(*npar)
-            cbl.da.increments[k,:] = cbl.da.increments[k,indices] # ntim*nobs(*npar)
+            indices = np.argsort(zz[k,:])
+        zobs[k,:] = zz[k,indices]
+        theta[k,:,:] = cbl.da.model_equivalents[k,indices,:]
+        innov[k,:] = cbl.da.innov[k,indices] # ntim*nobs
+        cov_yy[k,:] = cbl.da.cov_yy[k,indices] # ntim*nobs
+        cov_pp[k,:,:] = cbl.da.cov_pp[k,indices,:] # ntim*nobs(*npar)
+        cov_py[k,:,:] = cbl.da.cov_py[k,indices,:] # ntim*nobs(*npar)
+        increments[k,:,:] = cbl.da.increments[k,indices,:] # ntim*nobs(*npar)
 
     if hasattr(cbl.da,'cov_py'):
-        sigma2_yb = cbl.da.cov_yy # ntim*nobs
-        sigma2_p = np.squeeze(cbl.da.cov_pp) # ntim*nobs(*npar)
-        covariance = np.squeeze(cbl.da.cov_py) # ntim*nobs(*npar)
+        sigma2_yb = cov_yy # ntim*nobs
+        sigma2_p = np.squeeze(cov_pp) # ntim*nobs(*npar)
+        covariance = np.squeeze(cov_py) # ntim*nobs(*npar)
         correlation = covariance/(sigma2_yb*sigma2_p)**0.5
-        innov = cbl.da.innov
-        deltap = np.squeeze(cbl.da.increments) # ntim*nobs(*npar)
+        deltap = np.squeeze(increments) # ntim*nobs(*npar)
 
     # Shorthand for plots
     A = correlation.T
@@ -417,7 +460,10 @@ def plot_p(p_factors,theta_profiles,zt,figtitle,ax=None):
             ax.plot(theta_profiles[i],zt,label='$p=%4.1f$'%p_factors[i])
         return ax
 
-def plot_spread(cbl,plot='spread',ax=None):
+def plot_free_ensemble(cbl,show='spread',ax=None,cmap='viridis'):
+
+    # Read parameters
+    param = cbl.initial_perturbed_parameters
 
     # Read relevant dimensions
     times = cbl.history['0000']['time']
@@ -432,10 +478,19 @@ def plot_spread(cbl,plot='spread',ax=None):
         theta[k,:,:]= cbl.history['%04u'%k]['theta']
 
     # Compute spread
-    if plot=='spread':
+    if show=='spread':
         bkgd = theta.std(axis=0)
-    elif plot=='mean':
+        contour_color = 'white'
+    elif show=='mean':
         bkgd = theta.mean(axis=0)
+        contour_color = 'white'
+    if show=='cov':
+        cov = np.zeros((nz,ntimes))+np.nan
+        for ii in range(nz):
+            for jj in range(ntimes):
+                cov[ii,jj] = np.cov(param,theta[:,ii,jj])[0,1]
+        bkgd = cov
+        contour_color = 'black'
 
     # Plot
     zmax = 2000
@@ -443,54 +498,211 @@ def plot_spread(cbl,plot='spread',ax=None):
 
     if ax is not None:
         # Make plots
-        c = ax.pcolormesh(times/3600.,zt,bkgd)
+        c = ax.pcolormesh(times/3600.,zt,bkgd,cmap=cmap)
         ax.contour(times/3600.,zt,theta.mean(axis=0),
                     np.linspace(cbl.theta_0,cbl.theta_0+cbl.gamma*zmax,ncont),
-                    colors='white',
+                    colors=contour_color,
                     linestyles='--',
                     linewidths=0.75)
         ax.set_ylim([0,zmax])
         return ax,c
 
-def plot_diagnostics(experiments_pe,experiments_nope,labels,filename):
+def plot_diagnostics(exp,show='errors',axes=None,label='',linecolor='blue'):
+
+    # See if the experiment is an OSSE
+    if (not hasattr(exp.da,'truths')):
+        truth_exists = False
+    else:
+        truth_exists = True
+
+    # Get the data
+    sigma_o = exp.dg.varo**0.5
+    armsd = exp.dg.aRMSD_t
+    brmsd = exp.dg.bRMSD_t
+    bsprd = exp.dg.bSprd_t
+    if truth_exists:
+        armse = exp.dg.aRMSE_t
+        brmse = exp.dg.bRMSE_t
+
+    # Read observation coordinates
+    # In the 2D case, it is assumed that the set of 
+    # observation coordinates does not change
+    # from time to time. Anyway, coordinates need
+    # to be sorted. Diagnostic quantities are 
+    # already computed in a way that takes the correct
+    # spatial sorting into account.
+
+    if exp.da.obs_coordinates.ndim == 1:
+        z = np.sort(exp.da.obs_coordinates)
+    elif exp.da.obs_coordinates.ndim == 2:
+        z = np.sort(exp.da.obs_coordinates[0,:])
+
+    if axes is not None and len(axes)==4:
+
+        # Plot errors (wrt to true state) only if true state exists
+        if show=='errors' and truth_exists:
+
+            axes[0].plot(armse,z,label=label,color=linecolor)#,color=linecolors[i])
+            axes[1].plot(brmse,z,color=linecolor)
+            axes[2].plot(brmse-armse,z,color=linecolor)
+            axes[3].plot(np.sqrt(bsprd**2+sigma_o**2)/brmse,z,color=linecolor)
+
+        elif show=='deviations':
+
+            # Compute time-average mean bias of background and analyses
+            mean_bias_a_i1 = exp.dg.omina.mean(axis=0)
+            mean_bias_b_i1 = exp.dg.ominb.mean(axis=0)
+            # Plots
+            axes[0].plot(armsd,z,label=label,color=linecolor)
+            axes[1].plot(brmsd,z,color=linecolor)
+            axes[2].plot(brmsd-armsd,z,color=linecolor)
+            axes[3].plot(np.sqrt(bsprd**2+sigma_o**2)/brmsd,z,color=linecolor)
+
+    return axes
+
+def plot_linearity(cbl,ax=None):
     
-    linecolors = p.rcParams['axes.prop_cycle'].by_key()['color']
-
-    fig, [[ax1, ax2],[ax3, ax4]] = p.subplots(2,2,constrained_layout=True)
-    fig.set_size_inches(6,4)
-    z = experiments_pe[0].obs_coordinates
-    z_pbl = z*1.
-    z_pbl[z>1500] = np.nan
-    for i in range(len(experiments_pe)):
-        i1 = experiments_pe[i].dg
-        i2 = experiments_nope[i].dg
-        ax1.plot(i1.aRMSE_t,z,label=labels[i],color=linecolors[i])
-        ax1.plot(i2.aRMSE_t,z,color=linecolors[i],dashes=[3,1],alpha=0.3)
-        #
-        ax2.plot(i1.bRMSE_t,z,label=labels[i],color=linecolors[i])
-        ax2.plot(i2.bRMSE_t,z,color=linecolors[i],dashes=[3,1],alpha=0.3)
-        #
-        ax3.plot(i1.bRMSE_t-i1.aRMSE_t,z,label=labels[i],color=linecolors[i])
-        ax3.plot(i2.bRMSE_t-i2.aRMSE_t,z,color=linecolors[i],dashes=[3,1],alpha=0.3)
-        #
-        ax4.plot(i1.bSprd_t/i1.bRMSE_t,z_pbl,label=labels[i],color=linecolors[i])
-        ax4.plot(i2.bSprd_t/i2.bRMSE_t,z_pbl,color=linecolors[i],dashes=[3,1],alpha=0.3)
-    ax1.set_title('a) Analysis error')
-    ax1.set_xlabel(r'RMSE$^a_\theta$')
-    ax2.set_title('b) First-guess error')
-    ax2.set_xlabel(r'RMSE$^b_\theta$')
-    ax3.set_title('c) Error reduction')
-    ax3.set_xlabel(r'RMSE$^b_\theta-$RMSE$^a_\theta$')
-    ax4.set_title('d) Spread-error consistency')
-    ax4.set_xlabel(r'$\sigma^b_\theta$/RMSE$^b_\theta$')
-    ax1.set_ylabel('height (m)')
-    ax3.set_ylabel('height (m)')
+    # Read parameters
+    param = np.squeeze(cbl.parameter_transform[0](cbl.initial_perturbed_parameters,kind='inv'))
+
+    # Get the indices that sort parameters in ascending order
+    indices = np.argsort(param)
+
+    # 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']
+
+    # Make plot
+    if ax is not None:
+        tt = -1 # pick the last time in the history
+        zz = [0,10,20,30] # pick a few heights
+        hh = [int(cbl.dz/2)+i*cbl.dz for i in zz]
+        labels = [r'$z=%s$ m'%i for i in hh]
+        for z,label in zip(zz,labels):
+            #ax.scatter(param[indices],theta[indices,z,tt],s=5,label=label)
+            ax.plot(param[indices],theta[indices,z,tt],label=label)
+        ax.legend(loc=3,frameon=False,fontsize=6,facecolor='lightgrey')
+
+    return ax
+
+def plot_B_matrix(exp,timeid=-1):
+
+    # Read the background state for the requested analysis time
+    xx = exp.da.backgrounds[timeid]
+    nvar,nens = xx.shape
+
+    # Analysis time
+    tt = exp.da.time[timeid]
+
+    # Compute deviations from ensemble mean
+    X = xx.T-np.tile(xx.mean(axis=1),(nens,1))
+
+    # Standardize deviations from ensemble mean
+    X = X/np.tile(xx.std(axis=1),(nens,1))
+
+    # Compute covariance matrix
+    Pb = np.matmul(X.T,X)/(nens-1)
+    v = np.abs(Pb[:-1,:-1]).max()
+
+    # Plot
+    fig = p.figure(0)
+    fig.set_size_inches(4,3)
+    sp1 = fig.add_subplot(1,1,1)
+    c = sp1.pcolormesh(range(nvar),range(nvar),Pb,vmin=-v,vmax=v,cmap='RdBu_r',shading='nearest')
+    sp1.set_aspect('equal','box')
+    sp1.set_xlim(-0.5,nvar-0.5)
+    sp1.set_ylim(-0.5,nvar-0.5)
+    sp1.set_title('$1/(N-1)\cdot X^TX$, $t=%4.2f$ h'%(tt/3600))
+    p.colorbar(c)
+    fig.tight_layout()
+    fig.savefig('Covariance_matrix_%05u_s.png'%tt,format='png',dpi=150)
+    p.close(fig)
+
+def plot_desroziers(dgs,filename='desroziers_diagnostics.png',labels=None,logscale=False,ax=None):
+    from matplotlib.ticker import MaxNLocator
+
+    newfig = False    
+    if ax is None:
+        newfig = True
+        fig = p.figure(0)
+        fig.set_size_inches(4,3)
+        ax = fig.add_subplot(1,1,1)
+
+    for i,dg in zip( range(len(dgs)) , dgs ):
+        if labels is not None:
+            # Two formulations of the same thing; they give the same numbers!
+            #ax1.scatter( dg.dobdob_ta-dg.mean_ominb**2 , dg.varo+dg.varb_ta , s=5, label = labels[i])
+            ax.scatter( dg.sigma_ominb**2 , dg.varo+dg.varb_ta , s=5, label = labels[i])
+        else:
+            # Two formulations of the same thing; they give the same numbers!
+            #ax1.scatter( dg.dobdob_ta-dg.mean_ominb**2 , dg.varo+dg.varb_ta , s=5)
+            ax.scatter( dg.sigma_ominb**2 , dg.varo+dg.varb_ta , s=5) # 
+        #ax.set_title("Consistency of innovations",fontsize=8)
+        ax.set_xlabel(r'${\langle}\sigma_d^2{\rangle}_t$') # '${\langle}d_{ob}^2{\rangle}_t$'
+        ax.set_ylabel(r'$\sigma_o^2+{\langle}\sigma_b^2{\rangle}_t$')
+    
+    if labels is not None:
+        ax.legend(frameon=False, fontsize=6, ncol=2)
+
+    if logscale:
+        ax.axline((dgs[0].varo[0],dgs[0].varo[0]), slope=1, dashes=[3,2], color='k')
+        ax.set_xscale("log")
+        ax.set_yscale("log")
+    else:
+        ax.axline((0,0), slope=1, dashes=[3,2], color='k')
+        ax.axis("square")
+        ax.xaxis.set_major_locator(MaxNLocator(5))
+        ax.yaxis.set_major_locator(MaxNLocator(5))
+
+    if newfig:
+        fig.tight_layout()
+        fig.savefig(filename,format='png',dpi=150)
+        p.close(fig)
+
+    return ax
+
+def plot_overview(exp,label,filename,show_spread=False,show_transformed_parameter=False,show_true_parameter=True):
+
+    maxtime = exp.nr.maxtime
+    fig, [[ax0, ax1], [ax2, ax3]] = p.subplots(2,2,constrained_layout=True)
+    fig.set_size_inches(6,6)
     #
-    ax4.axvline(x=1,color='k',linewidth=0.5,dashes=[3,1])
-    ax2.sharey(ax1)
-    ax4.sharey(ax3)
-    p.setp(ax2.get_yticklabels(), visible=False)
-    p.setp(ax4.get_yticklabels(), visible=False)
+    [ax0,ax1,ax2],c0,c1,c2 = plot_CBL_identifiability(exp,exp.obs_error_sdev_assimilate[0],ax=[ax0,ax1,ax2])
+    ax0.set_title(r'a) Exp. %s, $\rho(p\prime\prime,y_b}$)'%label)
+    ax0.set_xlabel('Time (h)')
+    ax0.set_ylabel('Height (m)')
+    ax1.set_title(r'b) Exp. %s, $\delta\overline{y}\cdot(\sigma_{p\prime\prime}/\sigma_{y^b})$'%label)
+    ax1.set_xlabel('Time (h)')
+    ax1.set_ylabel('Height (m)')
+    ax2.set_title(r'c) Exp. %s, $\delta\overline{p}\prime\prime$'%label)
+    ax2.set_xlabel('Time (h)')
+    ax2.set_ylabel('Height (m)')
+    ax3 = plot_CBL_PE(exp,None,ax=ax3,show_spread=show_spread,show_transformed_parameter=show_transformed_parameter,show_true_parameter=show_true_parameter)
+    if show_transformed_parameter:
+        ax3.set_title(r'd) Exp. %s, evolution of $p\prime\prime$'%label)
+    else:
+        ax3.set_title(r'd) Exp. %s, evolution of $p$'%label)
+    ax3.set_xlabel('Time (h)')
+    ax3.set_yticks(np.arange(int(maxtime/3600)))
+    ax0.set_xlim([0,maxtime/3600])
+    ax1.set_xlim([0,maxtime/3600])
+    ax2.set_xlim([0,maxtime/3600])
+    ax3.set_xlim([0,maxtime/3600])
+    ax0.set_xticks(np.arange((maxtime+1)/3600))
+    ax1.set_xticks(np.arange((maxtime+1)/3600))
+    ax2.set_xticks(np.arange((maxtime+1)/3600))
+    ax3.set_xticks(np.arange((maxtime+1)/3600))
+    p.colorbar(c0,orientation='horizontal')
+    p.colorbar(c1,orientation='horizontal')
+    p.colorbar(c2,orientation='horizontal')
     #
     fig.savefig(filename,format='png',dpi=300)
     p.close(fig)