.gitignore

+# Ignore Jupyter Notebook checkpoints
+.ipynb_checkpoints/
+
+#Ignore pycache
+*.pyc
+__pycache__/
+
+#And all figures
+fig*/
+
+func_bk/
+
+#Also no pkl files
+*.pkl
+
+plot-data/
+theresa-data/
+obsolete/
+

LinAdvFunc/LinAdvFunc.egg-info/PKG-INFO

+Metadata-Version: 2.1
+Name: LinAdvFunc
+Version: 0.1

LinAdvFunc/LinAdvFunc.egg-info/SOURCES.txt

+setup.py
+LinAdvFunc.egg-info/PKG-INFO
+LinAdvFunc.egg-info/SOURCES.txt
+LinAdvFunc.egg-info/dependency_links.txt
+LinAdvFunc.egg-info/top_level.txt
\ No newline at end of file

LinAdvFunc/LinAdvFunc.egg-info/dependency_links.txt

+

LinAdvFunc/LinAdvFunc.egg-info/top_level.txt

+

LinAdvFunc/__init__.py

+from LinAdvFunc.da_functions import *
+from LinAdvFunc.misc_functions import *    
+from LinAdvFunc.model_functions import *
+from LinAdvFunc.plot_functions import *
+from LinAdvFunc.efsoi_functions import *
\ No newline at end of file

da_functions.py → LinAdvFunc/da_functions.py

@@ -3,8 +3,9 @@
 #Contains all functions related to setting up and running the ensemble filter. There is quite a bit of overlap with the model functions, but I felt it was a bit cleaner to try to separate them a bit. 
 
 import numpy as np
-from model_functions import *
-from misc_functions import *
+from LinAdvFunc.model_functions import *
+from LinAdvFunc.misc_functions import *
+from numba import jit
 
 
 
@@ -556,7 +557,8 @@ def run_linear_advection_KF_22(m_const,da_const,sat_operator):
                 states[j]["an"] = states[j]["an"]+[analysis]
             
             if da_const['method'] == 'LETKF':
-                analysis,bla = LETKF_analysis_22(background,obs,obs_sat,m_const,da_const,sat_operator)
+                # analysis,bla = LETKF_analysis_22(background,obs,obs_sat,m_const,da_const,sat_operator)
+                analysis,bla,W_a = LETKF_analysis_23(background,obs,obs_sat,m_const,da_const,sat_operator)
                 states[j]["an"] = states[j]["an"]+[analysis]
             
             if da_const['method'] == 'sqEnKF':
@@ -759,6 +761,14 @@ def sum_mid_tri(x):
     idx_end = int(2*nx/3.)
     return np.sum(x[idx_str:idx_end])
 
+def mid_hh(x):
+    """
+    For Herbert Hartl, the response function is simply the value at the middle of the state vector
+    """
+    nx = len(x)
+    idx_str = int(nx/2.)
+    return x[idx_str]
+
 def sum_mid_tri_abs(x):
     """
     Is a sum over the absolute values over the middle of the state 
@@ -1084,8 +1094,12 @@ def loc_matrix(da_const,m_const):
     return C
 
 def gaspari_cohn(x,loc_length):
-    """Gaspari-Cohn function intended to be applied to the x grid, which mirrors the matrix to account for periodic boundary domains."""
+    """
+    Gaspari-Cohn function intended to be applied to the x grid, which mirrors the matrix to account for periodic boundary domains.
+    2024-02-07: fixed the index in the mirroring that mean one side of the mirroring was one index off
+    """
     
+
     ra = x/loc_length
     gp = np.zeros_like(ra)
     i=np.where(ra<=1.)[0]
@@ -1095,7 +1109,8 @@ def gaspari_cohn(x,loc_length):
         
     #Now we mirror things to zero for the periodic boundary domain
     half_idx = int(ra.shape[0]/2)
-    gp[-half_idx:] = gp[half_idx-1::-1]
+    gp[-half_idx+1:] = gp[half_idx-1:0:-1]
+    # gp[-half_idx:] = gp[half_idx-1::-1]
     return gp
 
 def gaspari_cohn_non_mirrored(x,loc_length):
@@ -2081,6 +2096,134 @@ def LETKF_analysis_22(bg,obs,obs_sat,m_const,da_const,sat_operator):
     return x_a, x_ol_a
 
 
+def LETKF_analysis_23(bg,obs,obs_sat,m_const,da_const,sat_operator):
+    """
+    Same as LETKF_analysis_22 but speed things up by putting the middle loop into a separate numba function.
+
+    It seems roughly 10 times faster than the 22 version.
+
+    Further speed up could probably be achieved by reducing the size of the state arrays to not include areas outside the localized observations influence. 
+    
+    Follows the recipe and notation of Hunt 2007, e.g. what is normally called B is now P_a, and _ol refers to over line, so the mean.
+    P_tilde_a refers to B but in ensemble space.
+
+    Has an additional outputs needed for the EFSOI, namely the full weight matrix W_a
+    """
+
+
+
+    n_obs_h  =len(da_const["obs_loc"])
+    n_obs_sat =len(da_const["obs_loc_sat"])
+    n_obs = n_obs_h + n_obs_sat
+    obs_loc = np.hstack([da_const["obs_loc"],da_const["obs_loc_sat"]]).astype(int)
+    
+    
+    L = da_const['nens']
+    x_b = bg
+    x_ol_b = np.mean(x_b,axis=1)
+    X_b = x_b.T-x_ol_b
+    X_b = X_b.T
+
+    R = da_const['R'] 
+   
+    # this is the function that ports the background ensemble state to pertubations in observation space
+    Y_b, y_ol_b = state_to_observation_space(bg,m_const,da_const,sat_operator)
+    
+    # ineligant way to merge the obs vector depending on which observations occur
+    if n_obs_h>0:  y_obs = obs[da_const["obs_loc"]]
+    if n_obs_sat>0: sat_obs = obs_sat[da_const["obs_loc_sat"]]
+    if n_obs_h>0 and n_obs_sat>0 :  y_obs = np.hstack([y_obs,sat_obs])
+    if n_obs_sat>0 and n_obs_h==0 : y_obs   = sat_obs
+    
+    delta_y = y_obs-y_ol_b
+        
+    if da_const['loc']==False:
+        """ Now that all the variables are set, we start by computing the covariance matrix in ensemble state """
+        YRY = np.dot(Y_b.T,np.dot(np.linalg.inv(R),Y_b))
+        P_tilde_a = np.linalg.inv((L-1)*np.identity(L)+YRY)
+
+        """Next step, computing the enesemble mean analysis via the weighting vector w_ol_a"""
+        w_ol_a = np.dot(P_tilde_a,np.dot(Y_b.T,np.dot(np.linalg.inv(R),delta_y)))
+        x_ol_a = x_ol_b+np.dot(X_b,w_ol_a)
+
+        """We now get the ensemble by calculating the weighting matrix through a square root of the error covariance matrix, and adding the mean values to the ensemble deviations"""
+        W_a =  np.real(sqrtm((L-1)*P_tilde_a))
+        w_a = W_a+w_ol_a
+
+        x_a =  np.dot(X_b,w_a).T+ x_ol_a
+        x_a = x_a.T
+
+    else:
+        # print('making sure that the fix is being used')
+        # print('using the old version')
+        # print('trying the third version')
+        # print('sigh, what is not working now?')
+        # bla, blub,bleurgh=LETKF_numba_loop(m_const['x_grid'],da_const['loc_length'],obs_loc,x_ol_b,x_b,R,Y_b,L,delta_y,X_b,m_const['nx'],m_const['dx'])
+        # print('this is annoying, I hate numba')
+        x_a, x_ol_a,W_a=LETKF_numba_loop(m_const['x_grid'],da_const['loc_length'],obs_loc,x_ol_b,x_b,R,Y_b,L,delta_y,X_b,m_const['nx'],m_const['dx'])
+    return x_a, x_ol_a,W_a
+
+# @jit
+def LETKF_numba_loop(x_grid,loc_length,obs_loc,x_ol_b,x_b,R,Y_b,L,delta_y,X_b,N,dx):
+    x_grid_neg = -x_grid[::-1]-dx
+    x_grid_ext = x_grid+dx*N
+    x_a_loc = np.zeros((N,L))
+    W_a_loc = np.zeros((N,L,L))
+    x_ol_a_loc = np.zeros((N))
+    
+    for g in range(N):
+        dist_reg = np.abs(x_grid[obs_loc]-x_grid[g])
+        dist_neg = x_grid[g]-x_grid_neg[obs_loc]
+        dist_ext = x_grid_ext[obs_loc]-x_grid[g]
+        dist = np.minimum(dist_reg,dist_ext)#apparently minimum doesn't like 3 variables   ,dist_neg)
+        dist = np.minimum(dist,dist_neg)
+        #And now we calculate the gaspari cohn weighting and apply it to inverse R 
+        """Gaspari-Cohn function, with no mirroring."""
+    
+        ra = np.abs(dist)/loc_length
+    
+        gp = np.zeros_like(ra)
+        i=np.where(ra<=1.)[0]
+        gp[i]=-0.25*ra[i]**5+0.5*ra[i]**4+0.625*ra[i]**3-5./3.*ra[i]**2+1.
+        i=np.where((ra>1.)*(ra<=2.))[0]
+        gp[i]=1./12.*ra[i]**5-0.5*ra[i]**4+0.625*ra[i]**3+5./3.*ra[i]**2-5.*ra[i]+4.-2./3./ra[i]
+
+        gc_loc=gp 
+
+        if np.max(gc_loc) == 0.:
+            #If no observations are within the double gc radius no need to bother computing things
+            x_ol_a_loc[g] = x_ol_b[g]
+            x_a_loc[g,:]  = x_b[g,:]
+            W_a_loc[g,:,:] = np.identity(L)
+        else:
+            R_inv =  np.linalg.inv(R) * np.diag(gc_loc)
+
+            YRY = np.dot(Y_b.T,np.dot(R_inv,Y_b))
+            """ Now that all the variables are set, we start by computing the covariance matrix in ensemble state """
+            P_tilde_a = np.linalg.inv((L-1)*np.identity(L)+YRY)
+
+            """Next step, computing the enesemble mean analysis via the weighting vector w_ol_a"""
+            w_ol_a = np.dot(P_tilde_a,np.dot(Y_b.T,np.dot(R_inv,delta_y)))
+            x_ol_a = x_ol_b[g]+np.dot(X_b[g],w_ol_a)
+            # Trying to replace sqrtm with numpy eigenvectors so jit can work.
+            evalues, evectors = np.linalg.eigh((L-1.)*P_tilde_a)
+            W_a = evectors * np.sqrt(evalues) @ np.linalg.inv(evectors)
+            # w_a = W_a+w_ol_a
+            # x_a =  np.dot(X_b[g,:],w_a).T+ x_ol_a #this seems to be wrong
+            # x_a =  np.dot(X_b[g,:],w_a).T+ x_ol_b[g] # this should be right but isn't returning the right stuff
+            # This fix somehow makes things worse. I'll try something slightly different
+            # Well try Xa = Xb Wa, and then add x_ol_a to that
+            X_a = np.dot(X_b[g,:],W_a)
+            x_a = X_a+x_ol_a
+            x_ol_a_loc[g]   = x_ol_a
+            x_a_loc[g,:]    = x_a
+            W_a_loc[g,:,:]  = W_a
+            # print(W_a_loc.shape)
+            # print(x_a_loc.shape)
+            # print(x_ol_a_loc.shape)
+
+    return x_a_loc,x_ol_a_loc, W_a_loc
+
 
 def state_to_observation_space(X,m_const,da_const,sat_operator):
     """
@@ -2125,7 +2268,9 @@ def state_to_observation_space(X,m_const,da_const,sat_operator):
 def Kalman_gain_observation_deviations(bg,m_const,da_const,sat_operator):
     """
     Calculates Kalman gain, but using YYT instead of HBH, and XYT intead of BHT,
-    This is important for satellite data because we don't need a linearized version of the satellite operator 
+    This is important for satellite data because we don't need a linearized version of the satellite operator
+
+    Important fix! Localization was applied incorrectly. In stead of localizing BH the whole Kalman gain was localized :( 
     """
     n_obs_h  =len(da_const["obs_loc"])
     n_obs_sat =len(da_const["obs_loc_sat"])
@@ -2143,7 +2288,8 @@ def Kalman_gain_observation_deviations(bg,m_const,da_const,sat_operator):
     if da_const['loc']:
         L_obs, L_obs_state = localization_matrices_observation_space(m_const,da_const)
         YYlocR_inv = np.linalg.inv(L_obs*np.dot(dY_b,dY_b.T)/(L-1)+R)
-        K = L_obs_state*np.dot(X_b,np.dot(dY_b.T,YYlocR_inv))/(L-1)
+        # K = L_obs_state*np.dot(X_b,np.dot(dY_b.T,YYlocR_inv))/(L-1)
+        K = np.dot(L_obs_state*np.dot(X_b,dY_b.T),YYlocR_inv)/(L-1)
     else:
         YYR_inv = np.linalg.inv(np.dot(dY_b,dY_b.T)/(L-1)+R)
         K = np.dot(X_b,np.dot(dY_b.T,YYR_inv))/(L-1)
@@ -2176,7 +2322,7 @@ def square_root_Kalman_gain_observation_deviations(bg,m_const,da_const,sat_opera
         L_obs, L_obs_state = localization_matrices_observation_space(m_const,da_const)
     else:
         L_obs = np.ones([n_obs,n_obs])
-        L_obs_state = np.ones([n_x,n_obs])
+        L_obs_state = np.ones([m_const['nx'],n_obs])
         
     YYR = L_obs*np.dot(dY_b,dY_b.T)/(L-1)+R
     XY  = L_obs_state*np.dot(X_b,dY_b.T)/(L-1)
@@ -2246,13 +2392,20 @@ def single_step_analysis_forecast_22(background,truth,da_const,m_const,sat_opera
     """
     obs, obs_sat = generate_obs_22_single(truth,m_const,da_const,sat_operator,obs_seed) 
     
+    """
+    create dictionary to store this single step 
+    """
+    quad_state = {}
+    
     #Getting the analysis
     if da_const['method'] == 'EnKF':
         np.random.seed(obs_seed+100000)
         an = ENKF_analysis_22(background,obs,obs_sat,da_const,m_const,sat_operator)
     
     if da_const['method'] == 'LETKF':
-        an,bla = LETKF_analysis_22(background,obs,obs_sat,m_const,da_const,sat_operator)
+        # an,bla = LETKF_analysis_22(background,obs,obs_sat,m_const,da_const,sat_operator)
+        an,bla,W_a = LETKF_analysis_23(background,obs,obs_sat,m_const,da_const,sat_operator)
+        quad_state['W_a'] = W_a
     
     if da_const['method'] == 'sqEnKF':
         an = sqEnKF_analysis_22(background,obs,obs_sat,da_const,m_const,sat_operator)
@@ -2275,10 +2428,6 @@ def single_step_analysis_forecast_22(background,truth,da_const,m_const,sat_opera
         fc[:,i]    = linear_advection_model(an[:,i],u_ens,dhdt_ens,m_const["dx"],da_const["dt"],da_const["nt"])
 
     
-    """
-    create dictionary to store this single step 
-    """
-    quad_state = {}
     quad_state['bg'] = background
     quad_state['an'] = an
     quad_state['bf'] = bf

plot_functions.py → LinAdvFunc/plot_functions.py

@@ -6,7 +6,7 @@ import seaborn as sns
 import matplotlib.pyplot as plt
 import numpy as np
 
-from da_functions import *
+from LinAdvFunc.da_functions import *
 
 def ensemble_plotter(states,m_const,da_const,j=0,t_start=0,t_end=0):
     """
@@ -253,8 +253,8 @@ def B_plotter(states,j=0,ncyc=0,matrix='bg'):
     
     fig, ax = plt.subplots(1,1,figsize=(5,5))
 
-    ax.set_ylabel('[h1,h2,h3,...]')
-    ax.set_xlabel('[h1,h2,h3,...]')
+    ax.set_ylabel(r'$[\phi_1,\phi_2,\phi_3,...]$')
+    ax.set_xlabel(r'$[\phi_1,\phi_2,\phi_3,...]$')
 
     #B = np.cov(ensemble_diff.transpose())
     B = np.cov(states[j][matrix][ncyc],ddof=1)
@@ -510,35 +510,34 @@ def ensemble_plotter_22(states,m_const,da_const,j=0,t_start=0,t_end=0,h_c=None):
     #for c in range(nc):
     for c in range(na):
         for i in range(da_const["nens"]):
-            ax[c,0].plot(m_const['x_grid'],states[j]['bg'][c+t_start][:,i],'r',alpha =alpha,zorder=1)
-            ax[c,1].plot(m_const['x_grid'],states[j]['an'][c+t_start][:,i],'magenta',alpha =alpha,zorder=1)
-        ax[c,0].plot(m_const['x_grid'],np.mean(states[j]['bg'][c+t_start][:,:],axis=1),'k--',alpha =1,zorder=2,label='ens mean')
-        ax[c,0].plot(m_const["x_grid"],states[j]['truth'][c+t_start],'k',zorder=10,label='truth')
+            ax[c,0].plot(m_const['x_grid']/1000.,states[j]['bg'][c+t_start][:,i],'r',alpha =alpha,zorder=1)
+            ax[c,1].plot(m_const['x_grid']/1000.,states[j]['an'][c+t_start][:,i],'magenta',alpha =alpha,zorder=1)
+        ax[c,0].plot(m_const['x_grid']/1000.,np.mean(states[j]['bg'][c+t_start][:,:],axis=1),'k--',alpha =1,zorder=2,label='ens mean')
+        ax[c,0].plot(m_const["x_grid"]/1000.,states[j]['truth'][c+t_start],'k',zorder=10,label='truth')
         
-        ax[c,1].plot(m_const['x_grid'],np.mean(states[j]['an'][c+t_start][:,:],axis=1),'k--',alpha =1,zorder=2,label='ens mean')
-        ax[c,1].plot(m_const["x_grid"],states[j]['truth'][c+t_start],'k',zorder=10,label='truth')
+        ax[c,1].plot(m_const['x_grid']/1000.,np.mean(states[j]['an'][c+t_start][:,:],axis=1),'k--',alpha =1,zorder=2,label='ens mean')
+        ax[c,1].plot(m_const["x_grid"]/1000.,states[j]['truth'][c+t_start],'k',zorder=10,label='truth')
         
-        if da_const['n_obs_h']>0:
-            ax[c,0].scatter(m_const["x_grid"][da_const["obs_loc"]],states[j]['obs'][c+t_start][da_const["obs_loc"]],zorder=3,label='observations',c='k')
-            ax[c,1].scatter(m_const["x_grid"][da_const["obs_loc"]],states[j]['obs'][c+t_start][da_const["obs_loc"]],zorder=3,label='observations',c='k')
+        # if da_const['n_obs_h']>0:
+        if da_const['n_obs_h']:
+            ax[c,0].scatter( m_const["x_grid"][da_const["obs_loc"]]/1000.,states[j]['obs'][c+t_start][da_const["obs_loc"]]   ,zorder=3,s=10,label='obs',color='k')
+            ax[c,1].scatter( m_const["x_grid"][da_const["obs_loc"]]/1000.,states[j]['obs'][c+t_start][da_const["obs_loc"]]   ,zorder=3,s=10,label='obs',color='k')
+            ax[c,1].errorbar(m_const["x_grid"][da_const["obs_loc"]]/1000.,states[j]['obs'][c+t_start][da_const["obs_loc"]]   ,yerr=da_const['used_std_obs'] ,zorder=3,color='k',capsize=3,capthick=2,ls='none')
+            ax[c,0].errorbar(m_const["x_grid"][da_const["obs_loc"]]/1000.,states[j]['obs'][c+t_start][da_const["obs_loc"]]   ,yerr=da_const['used_std_obs'] ,zorder=3,color='k',capsize=3,capthick=2,ls='none')
         if h_c is not None:
-            ax[c,0].hlines(h_c,m_const['x_grid'][0],m_const['x_grid'][-1],color='k',ls=':',label='sat threshold')
-            ax[c,1].hlines(h_c,m_const['x_grid'][0],m_const['x_grid'][-1],color='k',ls=':',label='sat threshold')
+            ax[c,0].hlines(h_c,m_const['x_grid'][0]/1000.,m_const['x_grid'][-1]/1000.,color='k',ls=':',label='sat threshold')
+            ax[c,1].hlines(h_c,m_const['x_grid'][0]/1000.,m_const['x_grid'][-1]/1000.,color='k',ls=':',label='sat threshold')
 
-        #if da_const['n_obs_sat']>0:
-        #    #sat obs location will be added as well, I'll add them in with their value for now Maybe I'll get fancy and use a full circle where cloudy and empty circle where empty, because why not
-        #    ax[c,0].scatter(m_const["x_grid"][da_const["obs_loc_sat"]],states[j]['obs_sat'][c+t_start][da_const["obs_loc_sat"]],zorder=3,label='sat obs loc',marker='x',c='k')
-        #    ax[c,1].scatter(m_const["x_grid"][da_const["obs_loc_sat"]],states[j]['obs_sat'][c+t_start][da_const["obs_loc_sat"]],zorder=3,label='sat obs loc',marker='x',c='k')
-        ax[c,0].set_ylabel('assim step '+str(c+t_start) + '\n h [m]')
-        
+        # ax[c,0].set_ylabel('assim step: ' + str(c+t_start)+ '\n' +r'$\phi$') 
+        ax[c,0].set_ylabel('step: ' + str(c+t_start)+ '\n' +r'$\phi$') 
     
-    ax[0,0].set_title('background ensemble and observations')
-    ax[0,1].set_title('analysis ensemble after update')
-    ax[0,1].set_xlim([m_const["x_grid"][0],m_const["x_grid"][-1]])
+    ax[0,0].set_title('background ensemble \n and observations')
+    ax[0,1].set_title('analysis ensemble \n after assimilation')
+    ax[0,1].set_xlim([m_const["x_grid"][0]/1000.,m_const["x_grid"][-1]/1000.])
     plt.subplots_adjust(wspace=0.01,hspace=0.01)
 #     ax[0].set_xticklabels([])
-    ax[na-1,1].set_xlabel('x [m]')
-    ax[na-1,0].set_xlabel('x [m]')
+    ax[na-1,1].set_xlabel('x [km]')
+    ax[na-1,0].set_xlabel('x [km]')
     plt.legend()
     return fig,ax
 
@@ -608,6 +607,134 @@ def quad_plotter_22(quad_state,m_const,da_const):
     return fig,ax
 
 
+def quad_plotter_25(quad_state,m_const,da_const):
+    """
+    Plots the initial background and blind forecast, as well as the analysis and forecast with observations.
+    
+    Difference to earlier _22 is the arrangement of the subplots, and the addition of uncertainity bars for the observations. 
+ 
+
+    Returns: 
+    figure and axes
+    """
+    sns.set()
+    sns.set_style('whitegrid')
+
+
+    alpha = np.sqrt(1/da_const['nens'])+0.1
+    
+    fig, ax = plt.subplots(2,2,figsize=(7.5,4),sharex='all',sharey='all')
+  
+    
+    
+    for i in range(da_const["nens"]):
+        ax[0,0].plot(m_const['x_grid']/1000.,quad_state['bg'][:,i],'r',alpha =alpha,zorder=1)
+        ax[1,0].plot(m_const['x_grid']/1000.,quad_state['bf'][:,i],'b',alpha =alpha,zorder=1)
+        ax[0,1].plot(m_const['x_grid']/1000.,quad_state['an'][:,i],'magenta',alpha =alpha,zorder=1)
+        ax[1,1].plot(m_const['x_grid']/1000.,quad_state['fc'][:,i],'c',alpha =alpha,zorder=1)
+    
+
+    ax[0,0].plot(m_const["x_grid"]/1000.,quad_state['tr_bg'],'k',zorder=10,label='truth')
+    ax[0,1].plot(m_const["x_grid"]/1000.,quad_state['tr_bg'],'k',zorder=10,label='truth')
+    #ax[0,1].plot(m_const["x_grid"],quad_stat/1000.e['tr_fc'],'k',zorder=10,label='truth')
+    #ax[1,1].plot(m_const["x_grid"],quad_stat/1000.e['tr_fc'],'k',zorder=10,label='truth')
+    
+    ax[0,0].plot(m_const['x_grid']/1000.,np.mean(quad_state['bg'][:,:],axis=1),'k--',alpha =1,zorder=2,label='ens mean')
+    ax[0,1].plot(m_const['x_grid']/1000.,np.mean(quad_state['an'][:,:],axis=1),'k--',alpha =1,zorder=2,label='ens mean')
+    ax[1,0].plot(m_const['x_grid']/1000.,np.mean(quad_state['bf'][:,:],axis=1),'k--',alpha =1,zorder=2,label='ens mean')
+    ax[1,1].plot(m_const['x_grid']/1000.,np.mean(quad_state['fc'][:,:],axis=1),'k--',alpha =1,zorder=2,label='ens mean')
+
+    if da_const['n_obs_h']:
+        ax[0,0].scatter( m_const["x_grid"][da_const["obs_loc"]]/1000.,quad_state['obs'][da_const["obs_loc"]]   ,zorder=3,s=10,label='obs',color='k')
+        ax[0,1].scatter( m_const["x_grid"][da_const["obs_loc"]]/1000.,quad_state['obs'][da_const["obs_loc"]]   ,zorder=3,s=10,label='obs',color='k')
+        ax[0,1].errorbar(m_const["x_grid"][da_const["obs_loc"]]/1000.,quad_state['obs'][da_const["obs_loc"]]   ,yerr=da_const['used_std_obs'] ,zorder=3,color='k',capsize=3,capthick=2,ls='none')
+        ax[0,0].errorbar(m_const["x_grid"][da_const["obs_loc"]]/1000.,quad_state['obs'][da_const["obs_loc"]]   ,yerr=da_const['used_std_obs'] ,zorder=3,color='k',capsize=3,capthick=2,ls='none')
+        # ax[0,0].scatter(m_const["x_grid"][da_const["obs_loc"]]/1000.,quad_state['obs'][da_const["obs_loc"]],zorder=3,label='observations',color='k')
+        # ax[1,0].scatter(m_const["x_grid"][da_const["obs_loc"]]/1000.,quad_state['obs'][da_const["obs_loc"]],zorder=3,label='observations',color='k')
+    
+    
+    ax[0,0].set_title('background')
+    ax[1,0].set_title('free-forecast')
+    ax[0,1].set_title('analysis')
+    ax[1,1].set_title('forecast')
+    ax[0,1].set_xlim([m_const["x_grid"][0],m_const["x_grid"][-1]/1000])
+    plt.subplots_adjust(wspace=0.03,hspace=0.20)
+    ax[1,1].set_xlabel('x [km]')
+    ax[1,0].set_xlabel('x [km]')
+    ax[0,0].set_ylabel(r'$\phi$')
+    ax[1,0].set_ylabel(r'$\phi$')
+    ax[1,0].legend()
+    
+    # now to add in shading for response domain
+    ylimits = ax[0,0].get_ylim()
+    ax[1,0].fill_between([10,20], ylimits[0], y2=ylimits[1],color='grey',alpha=0.2,zorder=-1)
+    ax[1,1].fill_between([10,20], ylimits[0], y2=ylimits[1],color='grey',alpha=0.2,zorder=-1)
+    ax[1,1].set_ylim(ylimits)
+    
+    return fig,ax
+
+def quad_plotter_hh(quad_state,m_const,da_const):
+    """
+    Plots the initial background and blind forecast, as well as the analysis and forecast with observations.
+    
+    Intended for Herbert Hartl, but only difference to _25 is that the grey shading used to show the middle of the domain is not used. 
+ 
+
+    Returns: 
+    figure and axes
+    """
+    sns.set()
+    sns.set_style('whitegrid')
+
+
+    alpha = np.sqrt(1/da_const['nens'])+0.1
+    
+    fig, ax = plt.subplots(2,2,figsize=(7.5,4),sharex='all',sharey='all')
+  
+    
+    
+    for i in range(da_const["nens"]):
+        ax[0,0].plot(m_const['x_grid']/1000.,quad_state['bg'][:,i],'r',alpha =alpha,zorder=1)
+        ax[1,0].plot(m_const['x_grid']/1000.,quad_state['bf'][:,i],'b',alpha =alpha,zorder=1)
+        ax[0,1].plot(m_const['x_grid']/1000.,quad_state['an'][:,i],'magenta',alpha =alpha,zorder=1)
+        ax[1,1].plot(m_const['x_grid']/1000.,quad_state['fc'][:,i],'c',alpha =alpha,zorder=1)
+    
+
+    ax[0,0].plot(m_const["x_grid"]/1000.,quad_state['tr_bg'],'k',zorder=10,label='truth')
+    ax[0,1].plot(m_const["x_grid"]/1000.,quad_state['tr_bg'],'k',zorder=10,label='truth')
+    #ax[0,1].plot(m_const["x_grid"],quad_stat/1000.e['tr_fc'],'k',zorder=10,label='truth')
+    #ax[1,1].plot(m_const["x_grid"],quad_stat/1000.e['tr_fc'],'k',zorder=10,label='truth')
+    
+    ax[0,0].plot(m_const['x_grid']/1000.,np.mean(quad_state['bg'][:,:],axis=1),'k--',alpha =1,zorder=2,label='ens mean')
+    ax[0,1].plot(m_const['x_grid']/1000.,np.mean(quad_state['an'][:,:],axis=1),'k--',alpha =1,zorder=2,label='ens mean')
+    ax[1,0].plot(m_const['x_grid']/1000.,np.mean(quad_state['bf'][:,:],axis=1),'k--',alpha =1,zorder=2,label='ens mean')
+    ax[1,1].plot(m_const['x_grid']/1000.,np.mean(quad_state['fc'][:,:],axis=1),'k--',alpha =1,zorder=2,label='ens mean')
+
+    if da_const['n_obs_h']:
+        ax[0,0].scatter( m_const["x_grid"][da_const["obs_loc"]]/1000.,quad_state['obs'][da_const["obs_loc"]]   ,zorder=3,s=10,label='obs',color='k')
+        ax[0,1].scatter( m_const["x_grid"][da_const["obs_loc"]]/1000.,quad_state['obs'][da_const["obs_loc"]]   ,zorder=3,s=10,label='obs',color='k')
+        ax[0,1].errorbar(m_const["x_grid"][da_const["obs_loc"]]/1000.,quad_state['obs'][da_const["obs_loc"]]   ,yerr=da_const['used_std_obs'] ,zorder=3,color='k',capsize=3,capthick=2,ls='none')
+        ax[0,0].errorbar(m_const["x_grid"][da_const["obs_loc"]]/1000.,quad_state['obs'][da_const["obs_loc"]]   ,yerr=da_const['used_std_obs'] ,zorder=3,color='k',capsize=3,capthick=2,ls='none')
+        # ax[0,0].scatter(m_const["x_grid"][da_const["obs_loc"]]/1000.,quad_state['obs'][da_const["obs_loc"]],zorder=3,label='observations',color='k')
+        # ax[1,0].scatter(m_const["x_grid"][da_const["obs_loc"]]/1000.,quad_state['obs'][da_const["obs_loc"]],zorder=3,label='observations',color='k')
+    
+    
+    ax[0,0].set_title('background')
+    ax[1,0].set_title('free-forecast')
+    ax[0,1].set_title('analysis')
+    ax[1,1].set_title('forecast')
+    ax[0,1].set_xlim([m_const["x_grid"][0],m_const["x_grid"][-1]/1000])
+    plt.subplots_adjust(wspace=0.03,hspace=0.20)
+    ax[1,1].set_xlabel('x [km]')
+    ax[1,0].set_xlabel('x [km]')
+    ax[0,0].set_ylabel(r'$\phi$')
+    ax[1,0].set_ylabel(r'$\phi$')
+    ax[1,0].legend()
+    
+
+    
+    return fig,ax
+
 def plot_ensemble_sat_analysis(bg,an,obs,obs_sat,truth,sat_operator,m_const,da_const,h_c=None):
     fig ,ax = plt.subplots(2,2,figsize=(10,6),sharey='col',sharex='all')
     """
@@ -909,8 +1036,8 @@ def plot_J_quad_paper(J_dict,quad,sens,dx,bw=0.3,dJ=True):
     dX_an=dx
     dJ_ff=np.dot(sens,dX_bg)
     dJ_fc=np.dot(sens,dX_an)
-    print('vr_reductions:',np.var(dJ_fc,ddof=1)-np.var(dJ_ff,ddof=1 ),np.var(J_dict['fc'],ddof=1)-np.var(J_dict['bf'],ddof=1))
-    print('variance:',np.var(J_dict['bf'],ddof=1),np.var(dJ_ff,ddof=1),np.var(J_dict['fc'],ddof=1),np.var(dJ_fc,ddof=1 ))
+    print('estimated and real var reductions:',np.var(dJ_fc,ddof=1)-np.var(dJ_ff,ddof=1 ),np.var(J_dict['fc'],ddof=1)-np.var(J_dict['bf'],ddof=1))
+    print('variance of free forecast, estimated variance of ff, and the same for the forecast:',np.var(J_dict['bf'],ddof=1),np.var(dJ_ff,ddof=1),np.var(J_dict['fc'],ddof=1),np.var(dJ_fc,ddof=1 ))
             #'response' : np.hstack([J_dict['bf']-np.mean(J_dict['bf']),dJ_ff,J_dict['fc']-np.mean(J_dict['fc']),J_dict['es']-np.mean(J_dict['es'])]),
     if dJ:
         plot_data = {
@@ -929,7 +1056,8 @@ def plot_J_quad_paper(J_dict,quad,sens,dx,bw=0.3,dJ=True):
     'boxprops':{'facecolor':'none', 'edgecolor':'black'},
     }
     #ax = sns.violinplot(data=plot_data, inner='quartile', orient="v",cut=0,bw=bw,y='response',x='x_pos',palette=my_pal)#sns.color_palette('cool',n_colors=3))#,x='type')#,y='response',x='cyc',hue='type',,split=True,palette={dict_label[left_var]:dict_color[left_var],dict_label[right_var]:dict_color[right_var]}
-    ax = sns.stripplot(data=plot_data, y='response',x='x_pos',alpha=0.7,jitter=0.15,size=5,palette=my_pal)#color='0.0')#
+    ax = sns.stripplot(data=plot_data, y='response', x='x_pos', hue='x_pos', alpha=0.7, jitter=0.15, size=5, palette=my_pal, legend=False)#color='0.0')#
+    # ax = sns.stripplot(data=plot_data, y='response',x='x_pos',alpha=0.7,jitter=0.15,size=5,palette=my_pal)#color='0.0')#
     #ax = sns.boxplot(data=plot_data, y='response',x='x_pos',showfliers=False,**PROPS)#,patch_artist=False)#color='0.0')#,palette=my_pal
     #plot errorbars
     plt.errorbar(np.arange(4),np.zeros(4),[np.std(J_dict['bf'],ddof=1),np.std(dJ_ff,ddof=1),np.std(J_dict['fc'],ddof=1),np.std(dJ_fc,ddof=1 )],fmt='.',capsize=15,lw=3,color='k') 

LinAdvFunc/setup.py

+from setuptools import setup, find_packages
+
+setup(
+    name="LinAdvFunc",
+    version="0.1",
+    packages=find_packages(),
+)