fixed plot_CBL_identifiability to deal with randomly sorted observations

graphics.py

@@ -5,6 +5,8 @@ from matplotlib import pyplot as p
 import matplotlib.colors as colors
 from ENDA import experiment
 
+is_sorted = lambda a: np.all(a[:-1] <= a[1:])
+
 def truncate_colormap(cmap, minval=0.0, maxval=1.0, n=100):
     #https://stackoverflow.com/questions/18926031/how-to-extract-a-subset-of-a-colormap-as-a-new-colormap-in-matplotlib
     new_cmap = colors.LinearSegmentedColormap.from_list(
@@ -27,7 +29,10 @@ def plot_CBL_assimilation(exp,figtitle,which_cycle=0,ax=None):
     o = exp.observations[which_cycle]
     b = exp.da.backgrounds[which_cycle,:,:]
     a = exp.da.analyses[which_cycle,:,:]
+    if exp.obs_coordinates.ndim == 1:
         ocoords = exp.obs_coordinates
+    elif exp.obs_coordinates.ndim == 2:
+        ocoords = exp.obs_coordinates[which_cycle]
     nz = exp.nr.nz
     zt = exp.nr.zt 
 
@@ -251,10 +256,37 @@ def plot_CBL_identifiability(cbl,sigma_o,ax=None):
     npar = cbl.nr.parameter_number
     ntimes = cbl.da.ncycles
     times = cbl.da.time
-    zobs = cbl.da.obs_coordinates
+
+    # Expand times array to 2D (matching the number of observations)
+    times = times[:,None]+np.zeros((ntimes,nobs))
+
+    # Read model equivalents and state vector elements
     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)
+    elif cbl.da.obs_coordinates.ndim == 2:
+        zobs = cbl.da.obs_coordinates
+
+    # At each time, check if the coordinate array is sorted
+    # if not, resort both coordinates and model equivalents
+    for k in range(ntimes):
+        if is_sorted(zobs[k,:]):
+            pass
+        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)
+
     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)
@@ -273,6 +305,10 @@ def plot_CBL_identifiability(cbl,sigma_o,ax=None):
     # For plotting
     contours = np.linspace(cbl.nr.theta_0,cbl.nr.theta_0+cbl.nr.gamma*zmax,ncont)
 
+    # Before plotting, transpose the 2D arrays of times and coordinates
+    times = times.T
+    zobs = zobs.T
+
     # Make plots
     if ax is None:
         pass
@@ -381,7 +417,7 @@ 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,ax=None):
+def plot_spread(cbl,plot='spread',ax=None):
 
     # Read relevant dimensions
     times = cbl.history['0000']['time']
@@ -396,7 +432,10 @@ def plot_spread(cbl,ax=None):
         theta[k,:,:]= cbl.history['%04u'%k]['theta']
 
     # Compute spread
-    spread = theta.std(axis=0)
+    if plot=='spread':
+        bkgd = theta.std(axis=0)
+    elif plot=='mean':
+        bkgd = theta.mean(axis=0)
 
     # Plot
     zmax = 2000
@@ -404,7 +443,7 @@ def plot_spread(cbl,ax=None):
 
     if ax is not None:
         # Make plots
-        c = ax.pcolormesh(times/3600.,zt,spread)
+        c = ax.pcolormesh(times/3600.,zt,bkgd)
         ax.contour(times/3600.,zt,theta.mean(axis=0),
                     np.linspace(cbl.theta_0,cbl.theta_0+cbl.gamma*zmax,ncont),
                     colors='white',