From 95158ec4a203e5cd91f6b1faf6baf9877478ed45 Mon Sep 17 00:00:00 2001
From: Stefano Serafin <serafin@jet02.img.univie.ac.at>
Date: Thu, 1 Aug 2024 16:47:41 +0200
Subject: [PATCH] extensive changes to identifiability figure; added code for
 DA diagnostics plotting

---
 graphics.py | 92 +++++++++++++++++++++++++++++++++++++++--------------
 1 file changed, 68 insertions(+), 24 deletions(-)

diff --git a/graphics.py b/graphics.py
index c31698e..568af05 100644
--- a/graphics.py
+++ b/graphics.py
@@ -243,7 +243,6 @@ def plot_CBL_identifiability(cbl,sigma_o,figtitle,ax=None):
     zmax = 2000
     ncont = 13
 
-    work_in_model_space = False
     work_in_observation_space = True
 
     if work_in_observation_space:
@@ -260,18 +259,14 @@ def plot_CBL_identifiability(cbl,sigma_o,figtitle,ax=None):
             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')
+            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)
+                correlation = covariance/(sigma2_yb*sigma2_p)**0.5
+                innov = cbl.da.innov
+                deltap = np.squeeze(cbl.da.increments) # ntim*nobs(*npar)
 
-            # Compute Pearson correlation coefficient and Kalman Gain
-            # (a little noise is added to the denominator to avoid singularities)
-            p_x_correlation = np.zeros((nobs,ntimes))+np.nan
-            kalman_gain = np.zeros((nobs,ntimes))+np.nan
-            sigma_yb = np.zeros((nobs,ntimes))+np.nan
-            for i in range(nobs):
-                for j in range(ntimes):
-                    covmat=np.cov(pfac[j,:],theta[j,i,:])
-                    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)
-                    sigma_yb[i,j] = covmat[1,1]
         else:
             pass
 
@@ -283,7 +278,7 @@ def plot_CBL_identifiability(cbl,sigma_o,figtitle,ax=None):
             pass
         elif type(ax) is p.Axes:
             # Make plots
-            c = ax.pcolormesh(times/3600.,zobs,p_x_correlation,norm=colors.CenteredNorm(),cmap='RdBu_r')
+            c = ax.pcolormesh(times/3600.,zobs,correlation,norm=colors.CenteredNorm(),cmap='RdBu_r')
             ax.contour(times/3600.,zobs,theta.mean(axis=2).T,
                         contours,
                         colors='black',
@@ -293,7 +288,7 @@ def plot_CBL_identifiability(cbl,sigma_o,figtitle,ax=None):
             return ax,c
         elif type(ax) is list and len(ax)==2:
             # Make plots
-            c0 = ax[0].pcolormesh(times/3600.,zobs,p_x_correlation,norm=colors.CenteredNorm(),cmap='RdBu_r')
+            c0 = ax[0].pcolormesh(times/3600.,zobs,correlation,norm=colors.CenteredNorm(),cmap='RdBu_r')
             ax[0].contour(times/3600.,zobs,theta.mean(axis=2).T,
                         contours,
                         colors='black',
@@ -301,7 +296,7 @@ def plot_CBL_identifiability(cbl,sigma_o,figtitle,ax=None):
                         linewidths=0.75)
             ax[0].set_ylim([0,zmax])
 
-            c1 = ax[1].pcolormesh(times/3600.,zobs,sigma_yb,cmap=new_cmap)
+            c1 = ax[1].pcolormesh(times/3600.,zobs,np.sqrt(sigma2_p[None,:]/sigma2_yb),cmap=new_cmap)
             ax[1].contour(times/3600.,zobs,theta.mean(axis=2).T,
                         contours,
                         colors='black',
@@ -311,7 +306,11 @@ def plot_CBL_identifiability(cbl,sigma_o,figtitle,ax=None):
             return ax,c0,c1
         elif type(ax) is list and len(ax)==3:
             # Make plots
-            c0 = ax[0].pcolormesh(times/3600.,zobs,p_x_correlation,norm=colors.CenteredNorm(),cmap='RdBu_r')
+            A = correlation.T
+            B = ( innov*(sigma2_p*sigma2_yb)**0.5/(sigma2_yb+sigma_o**2) ).T
+            C = deltap.T
+
+            c0 = ax[0].pcolormesh(times/3600.,zobs,A,vmin=-1,vmax=1,cmap='RdBu_r')
             ax[0].contour(times/3600.,zobs,theta.mean(axis=2).T,
                         contours,
                         colors='black',
@@ -319,7 +318,7 @@ def plot_CBL_identifiability(cbl,sigma_o,figtitle,ax=None):
                         linewidths=0.75)
             ax[0].set_ylim([0,zmax])
 
-            c1 = ax[1].pcolormesh(times/3600.,zobs,sigma_yb,cmap=new_cmap)
+            c1 = ax[1].pcolormesh(times/3600.,zobs,B,norm=colors.CenteredNorm(),cmap='RdBu_r')#,cmap=new_cmap)
             ax[1].contour(times/3600.,zobs,theta.mean(axis=2).T,
                         contours,
                         colors='black',
@@ -327,16 +326,17 @@ def plot_CBL_identifiability(cbl,sigma_o,figtitle,ax=None):
                         linewidths=0.75)
             ax[1].set_ylim([0,zmax])
 
-            c2 = ax[2].pcolormesh(times/3600.,zobs,kalman_gain,norm=colors.CenteredNorm(),cmap='RdBu_r')
+            c2 = ax[2].pcolormesh(times/3600.,zobs,C,norm=colors.CenteredNorm(),cmap='RdBu_r')
             ax[2].contour(times/3600.,zobs,theta.mean(axis=2).T,
                         contours,
                         colors='black',
                         linestyles='--',
                         linewidths=0.75)
             ax[2].set_ylim([0,zmax])
+
             return ax,c0,c1,c2
 
-    if work_in_model_space:
+    else:
         if isinstance(cbl,experiment):
             # Read relevant dimensions
             zt = cbl.da.zt
@@ -367,12 +367,12 @@ def plot_CBL_identifiability(cbl,sigma_o,figtitle,ax=None):
 
             # Compute Pearson correlation coefficient and Kalman Gain
             # (a little noise is added to the denominator to avoid singularities)
-            p_x_correlation = np.zeros((nz,ntimes))+np.nan
+            correlation = np.zeros((nz,ntimes))+np.nan
             kalman_gain = np.zeros((nz,ntimes))+np.nan
             for i in range(nz):
                 for j in range(ntimes):
                     covmat=np.cov(pfac[:,j],theta[:,i,j])
-                    p_x_correlation[i,j] = covmat[0,1] / (1e-9+np.sqrt(covmat[0,0]*covmat[1,1]))
+                    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)
 
             # For plotting
@@ -399,12 +399,12 @@ def plot_CBL_identifiability(cbl,sigma_o,figtitle,ax=None):
 
             # Compute proxies of Pearson correlation coefficient and Kalman Gain
             # (a little noise is added to the denominator to avoid singularities)
-            p_x_correlation = np.zeros((nz,ntimes))+np.nan
+            correlation = np.zeros((nz,ntimes))+np.nan
             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]))
+                    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)
             
             # For plotting
@@ -438,7 +438,7 @@ def plot_CBL_identifiability(cbl,sigma_o,figtitle,ax=None):
             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')
+            c3 = ax3.pcolormesh(times,zt,correlation,vmin=-1,vmax=1,cmap='RdBu_r')
             ax3.set_xlabel(r'time')
             ax3.set_title(r'$p$-$\theta$ correlation')
             p.colorbar(c3)
@@ -536,3 +536,47 @@ def plot_spread(cbl,ax=None):
                     linewidths=0.75)
         ax.set_ylim([0,zmax])
         return ax,c
+
+def plot_diagnostics(experiments_pe,experiments_nope,labels,filename):
+    
+    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>1000] = 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)')
+    #
+    #ax2.legend(frameon=False)
+    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)
+    #
+    fig.savefig(filename,format='png',dpi=300)
+    p.close(fig)
-- 
GitLab