diff --git a/LinAdvFunc/da_functions.py b/LinAdvFunc/da_functions.py
new file mode 100644
index 0000000000000000000000000000000000000000..7bc059a5aa7e8b13608c733d9f2f3cfa70c53995
--- /dev/null
+++ b/LinAdvFunc/da_functions.py
@@ -0,0 +1,3128 @@
+#!/usr/bin/env python
+
+#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 LinAdvFunc.model_functions import *
+from LinAdvFunc.misc_functions import *
+from numba import jit
+
+
+
+def set_da_constants_22(ncyc=100,nt=1,dt=550,u_std=2.0,dhdt_std=0.,
+ True_std_obs=0.15,used_std_obs=0.15,pert_std_obs=0.15,
+ obs_loc=np.arange(49,299,100).astype(int),
+ True_std_obs_sat=0.15,used_std_obs_sat=0.15,pert_std_obs_sat=0.15,
+ obs_loc_sat=np.arange(7,299,15),
+ nens=32,nexp=1,
+ init_noise=0.,init_spread = False, init_spread_h=0.5,init_spread_x = 500.,
+ ens_seed=0, obs_seed=0, fixed_seed = False,
+ loc=True,loc_length = 3200,loc_type='gaspari_cohn',
+ method='sqEnKF'):
+ """
+ Sets constants needed for data assimilation and stores them in a dictionary.
+ Is now expanded to also use "satellite" observations, so that all observation statements are doubled, once for the default h point measurements, and once for the sat measurements.
+
+ There is some confusting misnaming going on, e.g. "u_std_ens" = u_std, but nothing that needs to be changed immediately
+
+ nt is currently really dangerous because of how the u perubation of the linear advection modell are rerolled, please leave at 1 for now.
+
+
+
+ TODO:
+
+ """
+ DA_const = {}
+
+
+ DA_const["ncyc"] = ncyc # Number of assimilation cycles
+ DA_const["nens"] = nens # number of ensemble members
+ DA_const["nexp"] = nexp # number of parallel experiments to run
+
+
+ DA_const["nt"] = nt # Number of model timesteps between observations
+ DA_const["dt"] = dt # time duration of model timesteps
+
+ #Assimilation Method
+ DA_const['method'] = method # Options so far include EnKF and LETKF, both with and without localization.
+
+ #Ensemble Errors and ensemble
+ DA_const["u_std_ens"] = u_std # Standard deviation of model u
+ DA_const["dhdt_std_ens"] = dhdt_std # Standard deviation of model dydt
+ DA_const["ens_seed"] = ens_seed # Is added to the seed when generating ensemble paramter deviations
+ DA_const["fixed_seed"] = fixed_seed # If True then the ensemble errors will always be constant
+
+ DA_const["True_std_obs_sat"] = True_std_obs_sat # Standard deviation of true observation error of y
+ DA_const["used_std_obs_sat"] = used_std_obs_sat # Standard deviation of assumed observation error used to calculate R
+ DA_const["pert_std_obs_sat"] = pert_std_obs_sat # Standard deviation of pertubations added to the true observation for each ensemble member individually when updating
+
+ #Ensemble initialization
+ DA_const["init_noise"] = init_noise # spread of noise added to initial conditions to avoid singular matrices
+ DA_const["init_spread"] = init_spread # If True, adds an x and h displacement to initial ensemble spread
+ DA_const["init_spread_h"] = init_spread_h # initial spread of ensemble in h
+ DA_const["init_spread_x"] = init_spread_x # initial spread of ensemble in x
+
+ #Point measurement observations
+ DA_const["True_std_obs"] = True_std_obs # Standard deviation of true observation error of y
+ DA_const["used_std_obs"] = used_std_obs # Standard deviation of assumed observation error used to calculate R
+ DA_const["pert_std_obs"] = pert_std_obs # Standard deviation of pertubations added to the true observation for each ensemble member individually when updating
+ DA_const["obs_loc"] = obs_loc # index array of which state elements are observed
+
+ # Satellite observations
+ DA_const["True_std_obs_sat"] = True_std_obs_sat # Standard deviation of true observation error of y
+ DA_const["used_std_obs_sat"] = used_std_obs_sat # Standard deviation of assumed observation error used to calculate R
+ DA_const["pert_std_obs_sat"] = pert_std_obs_sat # Standard deviation of pertubations added to the true observation for each ensemble member individually when updating
+ DA_const["obs_loc_sat"] = obs_loc_sat # index array of which state elements are observed
+
+ DA_const["obs_seed"] = obs_seed # Is added to the seed when generating observations
+
+ #Localization
+ DA_const["loc"] = loc # localization yes or no, no for now
+ DA_const["loc_length"] = loc_length # localization yes or no, no for now
+ DA_const["loc_type"] = loc_type # localization yes or no, no for now
+
+ #Observations variance matrix
+ 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
+ r =np.ones(n_obs)
+ r[:n_obs_h] = DA_const["used_std_obs"]**2.
+ r[n_obs_h:] = DA_const["used_std_obs_sat"]**2.
+ DA_const['R'] = np.diag(r)
+ DA_const['n_obs_h'] = n_obs_h
+ DA_const['n_obs_sat'] = n_obs_sat
+ return DA_const
+
+def set_da_constants(ncyc=10,nt=1,dt=500,u_std=0.5,dhdt_std=1e-4,True_std_obs=0.1,used_std_obs=0.1,pert_std_obs=0.1,obs_loc_h=np.arange(5,101,15),nens=20,nexp=1,init_noise=0.,fixed_seed=True,
+ loc=None,init_spread = False, init_spread_h=0.5,init_spread_x = 500.,
+ loc_length = 1000,loc_type='gaspari_cohn',method='EnKF'):
+ """
+ Sets constants needed for data assimilation and stores them in a dictionary.
+ There is some confusting misnaming going on, e.g. "u_std_ens" = u_std, but nothing that needs to be changed immediately
+
+ nt is not really relevant as long as we use the linear advection model, but I left it in from Yvonne in case we start applying it to different models.
+
+ TODO:
+
+ """
+ DA_const = {}
+
+
+ DA_const["ncyc"] = ncyc # Number of assimilation cycles
+ DA_const["nens"] = nens # number of ensemble members
+ DA_const["nexp"] = nexp # number of parallel experiments to run
+
+
+ DA_const["nt"] = nt # Number of model timesteps between observations
+ DA_const["dt"] = dt # time duration of model timesteps
+
+ #Assimilation Method
+ DA_const['method'] = method # Options so far include EnKF and LETKF, both with and without localization.
+
+ #Ensemble Errors and ensemble
+ DA_const["u_std_ens"] = u_std # Standard deviation of model u
+ DA_const["dhdt_std_ens"] = dhdt_std # Standard deviation of model dydt
+ DA_const["fixed_seed"] = fixed_seed # If True, then the ensemble number is used as a seed, so that the u and dhdt values are always randomized the same way
+
+ #Ensemble initialization
+ DA_const["init_noise"] = init_noise # spread of noise added to initial conditions to avoid singular matrices
+ DA_const["init_spread"] = init_spread # If True, adds an x and h displacement to initial ensemble spread
+ DA_const["init_spread_h"] = init_spread_h # initial spread of ensemble in h
+ DA_const["init_spread_x"] = init_spread_x # initial spread of ensemble in x
+
+ #Observations
+ DA_const["True_std_obs"] = True_std_obs # Standard deviation of true observation error of y
+ DA_const["used_std_obs"] = used_std_obs # Standard deviation of assumed observation error used to calculate R
+ DA_const["pert_std_obs"] = pert_std_obs # Standard deviation of pertubations added to the true observation for each ensemble member individually when updating
+ DA_const["obs_loc"] = obs_loc_h # index array of which state elements are observed
+
+
+ #Localization
+ DA_const["loc"] = loc # localization yes or no, no for now
+ DA_const["loc_length"] = loc_length # localization yes or no, no for now
+ DA_const["loc_type"] = loc_type # localization yes or no, no for now
+
+ return DA_const
+
+
+def sine_initial_condition(x,n):
+ """
+ Generates a sin curve on the x grid with overlapping waves defined by n.
+ """
+ mu = 0
+ y = np.zeros_like(x)
+ for ni in n:
+ y = y+ np.sin(x/np.max(x)*2*np.pi*ni)
+
+ return y
+
+
+
+def create_states_dict(j,states,m_const,da_const):
+ """
+ Generates the initial analysis and truth.
+ Also creates the "states" dictionary where the analysis ensemble, background ensemble, truth and observations are all going to be stored.
+ Very memory hungry, as everything from all assimilation time steps and all experiments is stored.
+ Works find for simple model though.
+
+ A fixed seed is used by default so that the model errors of all ensemble members is constant. But this can also be randomized.
+ Alternative version would be to generate the model errors and store them, but this has not happened yet. Would be necessary to test some parameter estimation tests.
+
+ Modified version of the Yvonne setup to work with the linear advection model
+
+ Todo:
+ - Describe states dictionary here.
+ - make model errors stored variables so enable parameter estimation
+ - Maybe make a more sensible name, such as init_da_ensembles
+ - I am not convinced this is the best way to generate the initial analysis.
+ - Might be best to change to dictionary time axis. Forecast and anaylis at the same time have difference of 1 in the time integer.
+ """
+ nx = m_const["nx"]
+
+ #initial conditions
+ if m_const['init_func']=='gaus': h_initial = gaussian_initial_condition(m_const["x_grid"],m_const["h_init_std"])
+ if m_const['init_func']=='sine': h_initial = sine_initial_condition( m_const["x_grid"],m_const["sine_init"])
+
+
+ #Generate truth
+ if da_const["fixed_seed"]==True: np.random.seed(j)
+
+ u_truth = np.random.normal(m_const["u_ref"],m_const["u_std_truth"])
+ dhdt_truth = np.random.normal(m_const["dhdt_ref"],m_const["dhdt_std_truth"])
+ truth = linear_advection_model(h_initial,u_truth,dhdt_truth,m_const["dx"],da_const["dt"],da_const["nt"])
+
+ an = np.zeros((nx,da_const["nens"]))
+
+ #First rows full of nans :)
+ bg = np.zeros((nx,da_const["nens"]))
+ bg[:] = np.nan
+ obs = np.zeros(nx)
+ obs[:] = np.nan
+
+ for i in range(da_const["nens"]):
+ if da_const["fixed_seed"]==True: np.random.seed(i+j*da_const["nens"])
+ if da_const["init_noise"]>0:
+ h_ens = np.random.normal(h_initial,da_const["init_noise"])
+ else:
+ h_ens = h_initial
+ if da_const["init_spread"]>0:
+ #initial spread generated by moving waves forward and backward and up and down using
+ #da_const["init_spread_h"] and da_const["init_spread_x"]
+ x_displace = np.random.normal(0.,da_const["init_spread_x"])
+ h_displace = np.random.normal(0.,da_const["init_spread_h"])
+ u_tmp = x_displace/da_const["dt"]
+ h_ens = semi_lagrangian_advection(h_ens,m_const["dx"],u_tmp,da_const["dt"])
+ h_ens = h_ens+h_displace
+
+ u_ens = np.random.normal(m_const["u_ref"],da_const["u_std_ens"])
+ dhdt_ens = np.random.normal(m_const["dhdt_ref"],da_const["dhdt_std_ens"])
+ an[:,i] = linear_advection_model(h_ens,u_ens,dhdt_ens,m_const["dx"],da_const["dt"],da_const["nt"])
+
+ states[j]={}
+ states[j]['bg']=[bg]
+ states[j]['an']=[an]
+ states[j]['truth']=[truth]
+ states[j]['obs']=[obs]
+ return an, truth, states
+
+
+ return DA_const
+
+
+def generate_obs(truth,states,m_const,da_const,j,t):
+ """
+ Generates the truth and observations for the next assimilation step.
+ If a fixed_seed is used the model "errors" of the truth remain the same, but noise of the obs is still reseeded differently each time.
+
+ To avoid diffusion over time in the truth I take advantage of the linearity of the model to compute it directly from the initial conditions.
+ This only works if the truth is not perturbed each timestep though.
+ """
+ #Generate new truth constants and integrate in time
+ if da_const["fixed_seed"]==True: np.random.seed(j)
+ u_truth = np.random.normal(m_const["u_ref"],m_const["u_std_truth"])
+ dhdt_truth = np.random.normal(m_const["dhdt_ref"],m_const["dhdt_std_truth"])
+ if da_const["fixed_seed"]==True:
+ #t starts at zero, that is why we need a plus 1
+ truth = linear_advection_model(states['truth'][0],u_truth,dhdt_truth,m_const["dx"],da_const["dt"]*(t+1),da_const["nt"])
+
+ else:
+ truth = linear_advection_model(truth,u_truth,dhdt_truth,m_const["dx"],da_const["dt"],da_const["nt"])
+
+ #make obs by adding some noise, a fixed seed should change over time so that the differences are not always the same for each measurement location
+ if da_const["fixed_seed"]==True: np.random.seed((j+1)*(t+1))
+ obs = truth + np.random.normal(0,da_const["True_std_obs"],m_const["nx"])
+ states["truth"] = states["truth"]+[truth]
+ states["obs"] = states["obs"]+[obs]
+ return truth, obs, states
+
+
+
+def generate_obs_22(truth,truth_init,m_const,da_const,j,t,sat_operator):
+ """
+ Generates the truth and observations for the next assimilation step to be used when running the model forward with assimilation steps.
+ In contrast to the old generate_obs this now includes the possibility to generate sat_obs.
+
+ Another new difference is that the new da_const has seeds for the obs and the ensembles, so that we can set them independantly without messing around with flags
+
+ To avoid diffusion over time in the truth I take advantage of the linearity of the model to compute it directly from the initial conditions.
+ This only works if the truth is not perturbed each timestep though.
+
+ Should replace the standard function at some point
+ """
+ #Generate new truth constants and integrate in time
+ if m_const["u_std_truth"]+m_const["dhdt_std_truth"]==0:
+ #t starts at zero, that is why we need a plus 1
+ truth = linear_advection_model(truth_init,m_const["u_ref"],m_const["u_std_truth"],m_const["dx"],da_const["dt"]*(t+1),da_const["nt"])
+
+ else:
+ np.random.seed((j+1)*(t+1)+1000+m_const['truth_seed'])
+ u_truth = np.random.normal(m_const["u_ref"],m_const["u_std_truth"])
+ dhdt_truth = np.random.normal(m_const["dhdt_ref"],m_const["dhdt_std_truth"])
+ truth = linear_advection_model(truth,u_truth,dhdt_truth,m_const["dx"],da_const["dt"],da_const["nt"])
+
+ #make obs by adding some noise, a fixed seed should change over time so that the differences are not always the same for each measurement location
+ np.random.seed((j+1)*(t+1)+10000+da_const['obs_seed'])
+ obs = truth + np.random.normal(0,da_const["True_std_obs"],m_const["nx"])
+ if len(da_const['obs_loc_sat'])>0:
+ truth_sat = sat_operator(truth)
+ obs_sat = truth_sat + np.random.normal(0,da_const["True_std_obs_sat"],m_const["nx"])
+ obs_sat[obs_sat>1.]= 1.
+ obs_sat[obs_sat<0.]= 0.
+ else:
+ obs_sat = np.zeros(m_const['nx'])
+
+ return truth, obs, obs_sat
+
+
+
+
+
+def predict(analysis,states,da_const,m_const,j):
+ """
+ Runs the analysis ensemble forward in time using the model to generate the next forecast/background ensemble for the next assimilation step.
+
+ Outlook:
+ - Add different models than the linear_advection_model?
+ """
+ bg = np.zeros((m_const["nx"],da_const["nens"]))
+ for i in range(da_const["nens"]):
+ if da_const["fixed_seed"]==True: np.random.seed(i+j*da_const["nens"])
+ u_ens = np.random.normal(m_const["u_ref"],da_const["u_std_ens"])
+ dhdt_ens = np.random.normal(m_const["dhdt_ref"],da_const["dhdt_std_ens"])
+ bg[:,i] = linear_advection_model(analysis[:,i],u_ens,dhdt_ens,m_const["dx"],da_const["dt"],da_const["nt"])
+
+ states["bg"] = states["bg"]+[bg]
+ return bg, states
+
+def predict_LA_22(analysis,states,da_const,m_const,j,t):
+ """
+ Runs the analysis ensemble forward in time using the model to generate the next forecast/background ensemble for the next assimilation step.
+ """
+ bg = np.zeros((m_const["nx"],da_const["nens"]))
+ for i in range(da_const["nens"]):
+ if da_const["fixed_seed"]==True:
+ np.random.seed(i+j*da_const["nens"]+da_const['ens_seed'])
+ else:
+ np.random.seed(i+j*da_const["nens"]+(da_const['ens_seed']+1)*t)
+ u_ens = np.random.normal(m_const["u_ref"],da_const["u_std_ens"])
+ dhdt_ens = np.random.normal(m_const["dhdt_ref"],da_const["dhdt_std_ens"])
+ bg[:,i] = linear_advection_model(analysis[:,i],u_ens,dhdt_ens,m_const["dx"],da_const["dt"],da_const["nt"])
+
+ states["bg"] = states["bg"]+[bg]
+ return bg, states
+
+
+
+
+def update(background,obs,R,H,C,states,da_const,m_const):
+ """
+ Computes the analysis by individually changing each ensemble member of the forecast/background through the shared ensemble Kalman Gain and the observations.
+ Now also used the localization matrix,C
+ """
+ # Compute the background error covariance matrix
+ P = np.cov(background)*C
+
+ # define relative weights of observations and background in terms of the Kalman Gain matrix of size
+ K = KalmanGain(P, R, H)
+
+ # Compute the analysis for each ensemble members
+ an = np.zeros((m_const["nx"],da_const["nens"]))
+ for i in range(da_const["nens"]):
+ an[:,i] = get_analysis(background[:,i],obs,K,H,da_const)
+ states["an"] = states["an"]+[an]
+ return an, states
+
+#def update_noloc(background,obs,R,H,states,da_const,m_const):
+# """
+# Computes the analysis by individually changing each ensemble member of the forecast/background through the shared ensemble Kalman Gain and the observations.
+#
+# Todo:
+# -Figure out exactly which version of the EnKF is used here, so that I can refer to it properly
+# """
+# # Compute the background error covariance matrix
+# P = np.cov(background)
+#
+# # define relative weights of observations and background in terms of the Kalman Gain matrix of size
+# K = KalmanGain(P, R, H)
+#
+# # Compute the analysis for each ensemble members
+# an = np.zeros((m_const["nx"],da_const["nens"]))
+# for i in range(da_const["nens"]):
+# an[:,i] = get_analysis(background[:,i],obs,K,H,da_const)
+# states["an"] = states["an"]+[an]
+# return an, states
+
+
+
+def KalmanGain(P,R,H):
+ """
+ Computes the Kalman gain matrix: K = PH^T(HPH^T+R)^(-1)
+ """
+ P_H = np.dot(P,H.transpose())
+ S = np.dot(H,P_H)+R
+ K = np.dot(P_H,np.linalg.inv(S))
+
+ return K
+
+def get_analysis(bg,obs,K,H,da_const):
+ """
+ Computes analysis: an = bg + K(H*bg-obs_pert), where obs_pert are perturbed observations
+ """
+ obs_pert = np.dot(H,obs+np.random.normal(0,da_const["pert_std_obs"],len(obs)))
+ bg_obs = np.dot(H,bg)
+ #obs_orig =np.dot(H,obs)
+ #print('obs_orig-bg_obs',np.sum(obs_orig-bg_obs),'obs_pert-bg_obs',np.sum(obs_pert-bg_obs),'update/np.sum(x)',np.sum(np.dot(K,obs_pert-bg_obs) )/np.sum(bg),'max(K)',np.max(K))
+ an = bg + np.dot(K,obs_pert-bg_obs)
+ return an
+
+def get_analysis_v2(bg, obs, K, H, obs_error_vec):
+ """
+ Computes analysis: an = bg + K(H*bg-obs_pert), where obs_pert are perturbed observations. v2 comes with a perturb_obs vector to enable differing observation erros.
+ Mostly made to deal with the shallow water model, but I guess it could be used to make some fun tests.
+ """
+ #print(obs)
+ #print(obs_error_vec)
+ obs_pert = np.dot(H,obs+np.random.normal(0,obs_error_vec,len(obs)))
+ bg_obs = np.dot(H,bg)
+ #obs_orig =np.dot(H,obs)
+ #print('obs_orig-bg_obs',np.sum(obs_orig-bg_obs),'obs_pert-bg_obs',np.sum(obs_pert-bg_obs),'update/np.sum(x)',np.sum(np.dot(K,obs_pert-bg_obs) )/np.sum(bg),'max(K)',np.max(K))
+ an = bg + np.dot(K,obs_pert-bg_obs)
+ return an
+
+
+def run_linear_advection_KF(m_const,da_const):
+ """
+ The heart and soul of the whole linear advection EnKF filter.
+ Steps:
+ - Computes constant Matrices H and R
+ - Initializes states dictionary with first truth and analysis ensemble
+ - And so on
+
+ Todo: Improve the documentation
+ """
+
+ """
+ constant matrices that follow from previously defined constants
+ """
+ H = np.identity(m_const["nx"])[da_const["obs_loc"],:] # Observation operator
+ R = da_const["used_std_obs"]**2*np.identity(len(da_const["obs_loc"])) # Observation error corvariance matrix
+
+
+
+
+ """
+ create dictionary to store background, analysis, truth and observations for verification
+ """
+ states = {}
+
+ #Construct localization matrix C if loc!=None
+ if da_const["loc"]:
+ C = loc_matrix(da_const,m_const)
+ else:
+ C=np.ones([m_const["nx"],m_const["nx"]])
+
+ """
+ Loop over experiments
+ """
+ for j in range(da_const["nexp"]):
+
+ """
+ Initialize a random truth and a random analysis ensemble stored in "states"
+ """
+ analysis, truth, states = create_states_dict(j,states,m_const,da_const)
+
+
+ """
+ START DATA ASSIMILATION CYCLE
+ """
+ for t in range(0,da_const["ncyc"]):
+ """
+ Run the truth forward in time until the next assimilation step and create observations.
+ Note that this step is usually provided by nature and measurements obtained from weather stations,
+ wind profilers, radiosondes, aircraft reports, radars, satellites, etc.
+ """
+ truth, obs, states[j] = generate_obs(truth,states[j],m_const,da_const,j,t)
+
+ """
+ Predict
+ """
+ # Predict the state at the next assimilation step by running the analysis forward in time
+ background, states[j] = predict(analysis,states[j],da_const,m_const,j)
+
+ """
+ Update
+ """
+ ## Combine background and observations to get the analysis
+ np.random.seed(t)
+ if da_const['method'] == 'EnKF':
+ analysis, states[j] = update(background,obs,R,H,C,states[j],da_const,m_const)
+ if da_const['method'] == 'LETKF':
+ analysis,bla = LETKF_analysis(background,obs,m_const,da_const)
+ states[j]["an"] = states[j]["an"]+[analysis]
+ """
+ END DATA ASSIMILATION CYCLE
+ """
+ """
+ end loop over experiments
+ """
+ return states
+
+
+def run_linear_advection_KF_22(m_const,da_const,sat_operator):
+ """
+ The heart and soul of the whole linear advection EnKF filter, now updated in the 22 version to work with satelitte data and the seeds.
+
+ Todo: Improve the documentation
+ """
+
+ """
+ constant matrices that follow from previously defined constants
+ """
+ # should not be needed anymore H = np.identity(m_const["nx"])[da_const["obs_loc"],:] # Observation operator
+ R = da_const["R"] # Observation error corvariance matrix
+
+
+
+
+ """
+ create dictionary to store background, analysis, truth and observations for verification
+ """
+ states = {}
+
+ #Construct localization matrix C if loc!=None
+ if da_const["loc"]:
+ C = loc_matrix(da_const,m_const)
+ else:
+ C=np.ones([m_const["nx"],m_const["nx"]])
+
+ """
+ Loop over experiments
+ """
+ for j in range(da_const["nexp"]):
+
+ """
+ Initialize a random truth and a random analysis ensemble stored in "states"
+ """
+ analysis, truth, states = create_states_dict_22(j,states,m_const,da_const,sat_operator)
+
+
+ """
+ START DATA ASSIMILATION CYCLE
+ """
+ for t in range(0,da_const["ncyc"]):
+ """
+ Run the truth forward in time until the next assimilation step and create observations.
+ Note that this step is usually provided by nature and measurements obtained from weather stations,
+ wind profilers, radiosondes, aircraft reports, radars, satellites, etc.
+ """
+ truth, obs, obs_sat = generate_obs_22(truth,states[j]['truth'][0],m_const,da_const,j,t,sat_operator)
+ states[j]["truth"] = states[j]["truth"] +[truth]
+ states[j]["obs"] = states[j]["obs"] +[obs]
+ states[j]["obs_sat"] = states[j]["obs_sat"]+[obs_sat]
+
+ """
+ Predict
+ """
+ # Predict the state at the next assimilation step by running the analysis forward in time
+ background, states[j] = predict_LA_22(analysis,states[j],da_const,m_const,j,t)
+
+ """
+ Update
+ """
+ ## Combine background and observations to get the analysis
+ np.random.seed(t)
+ if da_const['method'] == 'EnKF':
+ np.random.seed((j+3)**2+(t+4)**3)
+ analysis = ENKF_analysis_22(background,obs,obs_sat,da_const,m_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,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':
+ analysis = sqEnKF_analysis_22(background,obs,obs_sat,da_const,m_const,sat_operator)
+ states[j]["an"] = states[j]["an"]+[analysis]
+ """
+ END DATA ASSIMILATION CYCLE
+ """
+ """
+ end loop over experiments
+ """
+ return states
+
+
+def sqEnKF_analysis_22(background,obs,obs_sat,da_const,m_const,sat_operator):
+ """
+ Computes the analysis using the square root kalman Gain as introduced by Whitaker 2002. It used the modified Kalman gain to compute the change in deviations,
+ and the normal Kalman gain to compute the change in the mean.
+ """
+ #Step 1, preparing all the basics
+ an = np.zeros((m_const["nx"],da_const["nens"]))
+ bg_obs_deviations,bg_obs_ol = state_to_observation_space(background,m_const,da_const,sat_operator)
+ bg_obs = (bg_obs_deviations.T+bg_obs_ol).T
+ bg_ol = np.mean(background,axis=1)
+ if da_const['n_obs_sat']* da_const['n_obs_h']>0:
+ obs_stack = np.hstack([obs[da_const['obs_loc']],obs_sat[da_const['obs_loc_sat']]])
+ else:
+ if da_const['n_obs_sat']>0:
+ obs_stack = obs_sat[da_const['obs_loc_sat']]
+ if da_const['n_obs_h']>0:
+ obs_stack = obs[da_const['obs_loc']]
+
+ dbg= (background.T-bg_ol).T
+ # calculating the analysis ensemble mean using the normal Kalman gain
+ K = Kalman_gain_observation_deviations(background,m_const,da_const,sat_operator)
+ an_ol = bg_ol + np.dot(K,obs_stack-bg_obs_ol)
+
+ # Calculating change in ensemble deviations using squar root Kalman gain
+ sqK = square_root_Kalman_gain_observation_deviations(background,m_const,da_const,sat_operator)
+
+ dan = dbg -np.dot(sqK,bg_obs_deviations)
+
+ an = (dan.T + an_ol).T
+ #print(obs_stack)
+ #print(mean)
+ #print(bg_obs[0,:])
+ #for i in range(da_const["nens"]):
+ # an[:,i] = background[:,i] + np.dot(K,obs_stack-bg_obs[:,i])
+ return an
+
+
+def ENKF_analysis_22(background,obs,obs_sat,da_const,m_const,sat_operator):
+ """
+ Computes the analysis by individually changing each ensemble member of the forecast/background through the shared ensemble Kalman Gain and the observations.
+ """
+ # define relative weights of observations and background in terms of the Kalman Gain matrix of size
+ K = Kalman_gain_observation_deviations(background,m_const,da_const,sat_operator)
+
+ # Compute the analysis for each ensemble members
+ an = np.zeros((m_const["nx"],da_const["nens"]))
+ bg_obs_deviations,mean = state_to_observation_space(background,m_const,da_const,sat_operator)
+ bg_obs = (bg_obs_deviations.T+mean).T
+ for i in range(da_const["nens"]):
+ obs_pert_h = obs[da_const['obs_loc']] +np.random.normal(0,da_const["pert_std_obs"] ,len(da_const['obs_loc']))
+ obs_pert_sat = obs_sat[da_const['obs_loc_sat']]+np.random.normal(0,da_const["pert_std_obs_sat"],len(da_const['obs_loc_sat']))
+ obs_pert = np.hstack([obs_pert_h,obs_pert_sat])
+ an[:,i] = background[:,i] + np.dot(K,obs_pert-bg_obs[:,i])
+ return an
+
+
+
+def get_spread_and_rmse(states,da_const,m_const):
+ """
+ computes RMSE over space and time respectively, averaged over all experiments for the background ensmeble mean and analysis ensemble mean.
+ Structure of the output is as follows: rmse['dim']['quan'] where dim = space or time, quan = bg or an, bf will also be calculated when available
+
+ Returns spread and rmse over time and space averaged over all experiments. So increasing the number of experiments acts smooths out the time errors.
+ """
+
+
+ n = m_const["nx"]
+ ncyc = da_const["ncyc"]+1
+ nexp = da_const["nexp"]
+ nens = da_const["nens"]
+ time = np.arange(0,ncyc)
+
+ quantities = ['an','bg']
+ if 'bf' in states[0].keys(): quantities.append('bf')
+ tr_matrix = np.zeros([nexp,ncyc,n])
+ for j in range(nexp):
+ tr_matrix[j,:,:] = np.array(states[j]["truth"][:])
+ rmse = {}
+ spread = {}
+
+
+ for dim in ['time','space']: # Choose dim to average results
+ rmse[dim]={}
+ spread[dim]={}
+ for quan in quantities: # Loop over fields
+
+
+ rmse[dim][quan]={}
+ spread[dim][quan] = {}
+ rmse[dim][quan]['mean']=0.
+ spread[dim][quan]['mean']=0.
+
+ for j in range(nexp): # Loop over experiments
+ ens_matrix = np.array(states[j][quan][:])
+ if dim == "space": # Average over time to get one value per gridpoint
+ rmse[dim][quan][j] = np.zeros(n)
+ spread[dim][quan][j] = np.zeros(n)
+ for x in range(n): # Loop DA cycles
+ rmse[dim][quan][j][x] = np.nanmean(L2norm(ens_matrix[:,x,:].T-tr_matrix[j,:,x]))
+ spread[dim][quan][j][x] = np.nanmean(np.std(ens_matrix[:,x,:], axis=1, ddof=1))
+
+ if dim == "time": # Average over space so that one value remains per timestep
+ rmse[dim][quan][j] = np.zeros(ncyc)
+ spread[dim][quan][j] = np.zeros(ncyc)
+ for t in time: # Loop DA cycles
+ rmse[dim][quan][j][t] = np.mean(L2norm(ens_matrix[t,:,:].T-tr_matrix[j,t,:]))
+ spread[dim][quan][j][t] = np.mean(np.std(ens_matrix[t,:,:], axis=1, ddof=1))
+ rmse[dim][quan]['mean'] = rmse[dim][quan]['mean'] + rmse[dim][quan][j]/float(nexp)
+ spread[dim][quan]['mean'] = spread[dim][quan]['mean'] + spread[dim][quan][j] / float(nexp)
+
+ rmse[dim][quan] = rmse[dim][quan]['mean']
+ spread[dim][quan] = spread[dim][quan]['mean']
+ return rmse, spread
+
+
+def L2norm(error):
+ """
+ Computes the l2 norm over the first dimension of the input array
+ """
+ return np.sqrt(np.average(error**2,axis=0))
+
+
+
+def predict_blind(background,states,da_const,m_const,j):
+ """
+ Runs the background ensemble forward in time using the model to predict the truth at the next assimilation step.
+ This is than saved as 'bf', for blind forecast.
+ """
+ bf = np.zeros((m_const["nx"],da_const["nens"]))
+ for i in range(da_const["nens"]):
+ if da_const["fixed_seed"]==True: np.random.seed(i+j*da_const["nens"])
+ u_ens = np.random.normal(m_const["u_ref"],da_const["u_std_ens"])
+ dhdt_ens = np.random.normal(m_const["dhdt_ref"],da_const["dhdt_std_ens"])
+ bf[:,i] = linear_advection_model(background[:,i],u_ens,dhdt_ens,m_const["dx"],da_const["dt"],da_const["nt"])
+
+ #bg[:,i] = linear_advection_model(analysis[:,i],da_const["dt"],da_const["std_model"],m_const["C"],m_const["a"])
+ states["bf"] = states["bf"]+[bf]
+ return bf, states
+
+
+
+def add_blind_forecast(states,m_const,da_const):
+ """
+ Takes a states dictionary created by run_linear_advection_EnKF and adds a 'bf' field which stands for blind forecast, also known as free forecast.
+ This blind forecast is generated directly from the background ensemble, without assimilating first.
+ Only makes sense if the randomized errors are carefully generated from the same seed, otherwise the model errors used to make the background ensemble will change when creating the blind forecast.
+
+ This feels dangerously close to philipp making a custom solution for everything himself again.
+ """
+ """
+ Loop over experiments
+ """
+ for j in range(da_const["nexp"]):
+
+ #initializing new dictionary field, this is where things can easily go wrong, so pay attention
+ states[j]['bf'] = []
+ #add two first initial nan states, which we take from the initial forecast field
+ states[j]['bf'].append(states[j]['bg'][0])
+ states[j]['bf'].append(states[j]['bg'][0])
+
+ """
+ START blind forecasting
+ """
+ for t in range(2,da_const["ncyc"]+1):
+ """
+ Predict
+ """
+ #print('blind forecast time',t)
+ # Predict the state at the next assimilation step by running the analysis forward in time
+ blind_forecast, states[j] = predict_blind(states[j]['bg'][t-1],states[j],da_const,m_const,j)
+
+ """
+ END blind forecasting
+ """
+ """
+ end loop over experiments
+ """
+ return states
+
+def sum_mid_tri(x):
+ """
+ Default response function, is a simple sum over the middle of the state vector
+ """
+ nx = len(x)
+ idx_str = int(nx/3.)
+ idx_end = int(2*nx/3.)
+ return np.sum(x[idx_str:idx_end])
+
+def sum_mid_tri_abs(x):
+ """
+ Is a sum over the absolute values over the middle of the state
+ """
+ nx = len(x)
+ idx_str = int(nx/3.)
+ idx_end = int(2*nx/3.)
+ x_mid = x[idx_str:idx_end]
+ return np.sum(np.abs(x_mid))
+
+
+def mid_ma(x):
+ """
+ mean absolute values over the middle
+ """
+ nx = len(x)
+ idx_str = int(nx/3.)
+ idx_end = int(2*nx/3.)
+ x_mid = x[idx_str:idx_end]
+ return np.sum(np.abs(x_mid))/np.float(len(x_mid))
+
+def mid_rms(x):
+ """
+ root mean square ober mid
+ """
+ nx = len(x)
+ idx_str = int(nx/3.)
+ idx_end = int(2*nx/3.)
+ x_mid = x[idx_str:idx_end]
+ ms = np.sum(np.power(x_mid,2.))/np.float(len(x_mid))
+ return np.sqrt(ms)
+
+def randomized_obs_loc(n_obs,start=0,end=100,seed=0):
+ """randomly assignes n_obs within the window given"""
+ np.random.seed(seed)
+ randomized_obs_loc = np.random.choice(np.arange(start,end), n_obs, replace=False)
+ randomized_obs_loc.sort()
+ return randomized_obs_loc
+
+def add_response(states,func_J=sum_mid_tri):
+ """
+ Goes through the analysis, background, truth, and blind forecast if it exists and applies the given response function to the ensemble members.
+ Isn't pretty, could be made more elegantly to loop over the given variables instead of hard coding them.
+ Also poorly handles that there are differing number of bg, an, and bf fields.
+ Also so far doesn't allow computing the truth as a response function, as only the individual ensemble members are passed to func_J with no other information
+
+ ToDo:
+ - Allow using more diverse func_Js
+ - tidy up, is a mess
+ """
+ nexp = len(states)
+ ncyc = len(states[0]["bg"])
+ nx = states[0]["bg"][0][:,0].shape[0]
+ nens = states[0]["bg"][0][0,:].shape[0]
+
+ #Poor way of checking if a blind forecast also exists
+ bf_flag = 0
+ if 'bf' in states[0]: bf_flag = 1
+
+ for e in range(nexp):
+ states[e]["response"]={}
+
+ #Initialize dictionary fields, not the cleanest solution, but fuck it
+ c = 0
+
+ an_response = np.zeros(nens)
+ bg_response = np.zeros(nens)
+ if bf_flag: bf_response = np.zeros(nens)
+ for n in range(nens):
+ an_response[n] = func_J(states[e]["an"][0][:,n])
+ bg_response[n] = func_J(states[e]["bg"][0][:,n])
+ if bf_flag: bf_response[n] = func_J(states[e]["bf"][0][:,n])
+
+
+ states[e]["response"]["an"] =[an_response]
+ states[e]["response"]["bg"] =[bg_response]
+ states[e]["response"]["truth"]=[func_J(states[e]["truth"][0])]
+
+ if bf_flag : states[e]["response"]["bf"] =[bf_response]
+
+
+ for c in range(1,ncyc):
+ an_response = np.zeros(nens)
+ bg_response = np.zeros(nens)
+ bf_response = np.zeros(nens)
+ for n in range(nens):
+ an_response[n] = func_J(states[e]["an"][c][:,n])
+ bg_response[n] = func_J(states[e]["bg"][c][:,n])
+ if bf_flag: bf_response[n] = func_J(states[e]["bf"][c][:,n])
+
+ states[e]["response"]["bg"] =states[e]["response"]["bg"]+[bg_response]
+ if bf_flag: states[e]["response"]["bf"] =states[e]["response"]["bf"]+[bf_response]
+
+ states[e]["response"]["an"] =states[e]["response"]["an"]+[an_response]
+ states[e]["response"]["truth"]=states[e]["response"]["truth"]+[func_J(states[e]["truth"][c])]
+
+
+
+ return states
+
+def var_reduction_estimate(states,m_const,da_const,j=0,obs = [],ncyc=0,n_start = 10):
+ """
+ just loops over a full experiment
+ Calculates an estimated var reduction of the response fucntion for all observations, as well as for each observation individually
+ if a sepcific obs list is applied, that will be used instead of all observations
+
+ following naming of Hakim 2020
+
+ """
+
+ #For now I am not worried about being efficient
+ #j=0
+ #t=3
+ nobs = len(da_const["obs_loc"])
+ if obs == []:
+ obs = np.arange(nobs)
+
+
+ nens = da_const["nens"]
+ if ncyc ==0: ncyc = len(states[j]['response']['bf'])
+
+
+ var_reduction_total = []
+ var_reduction_individual = []
+ real_reduction = []
+
+ for t in range(n_start,ncyc-1):
+
+ R = da_const["used_std_obs"]**2*np.identity(len(obs)) # Observation error corvariance matrix
+ H = np.identity(m_const["nx"])[da_const["obs_loc"][obs],:]
+
+ x = states[j]['bg'][t][:,:]
+ dx = x.T-np.mean(x,axis=1)
+ dx = dx.T
+ A = np.dot(dx,dx.T)/(dx.shape[1]-1)
+
+ J= states[j]["response"]['bf'][t+1][:]
+ dJ = J-np.mean(J)
+
+ HAHt = np.dot(H,np.dot(A,H.T))
+ HAHtRinv= np.linalg.inv(HAHt+R)
+ dJHdxt = np.dot(dJ,np.dot(H,dx).T)/(nens-1)
+
+ vr_total = -np.dot(dJHdxt,np.dot(HAHtRinv,dJHdxt))
+ var_reduction_total.append(vr_total)
+
+ #And now the loop for the indivudual reductions
+ vr_individual = np.zeros(nobs)
+ for o in range(nobs):
+ R = da_const["used_std_obs"]**2*np.identity(1) # Observation error corvariance matrix
+ H = np.identity(m_const["nx"])[da_const["obs_loc"][o],:]
+ HAHt = np.dot(H,np.dot(A,H.T))
+ HAHtRinv= np.linalg.inv(HAHt+R)
+ dJHdxt = np.dot(dJ,np.dot(H,dx).T)/(nens-1)
+ vr_individual[o] = - np.dot(dJHdxt,np.dot(HAHtRinv,dJHdxt))
+ var_reduction_individual.append(vr_individual)
+
+ J_f=states[j]["response"]['bg'][t+1][:]
+ dJ_f = J_f-np.mean(J_f)
+
+ real_reduction.append(np.var(dJ_f) - np.var(dJ))
+ return var_reduction_total,var_reduction_individual,real_reduction
+
+def var_reduction_estimate_iterative(states,m_const,da_const,j=0,dJ_update_flag=0):
+ #from scipy.linalg import sqrtm
+
+ """
+ Iteratively goes through all observations, estimates the impact on the response function, updates dx and dJ, and continues on.
+ Currently starts with the location with the largest estimated impact, and then goes to the next highest, and so on. But the ordering should be unimportant.
+ It return the total variance reduction estimate for all the individual observations.
+
+ Following naming of Hakim 2020.
+
+ It was a bit unclear in the paper, but we decided that we would update dx and dJ simultaneously by bringing dJ into the state vector. This was also confirmed by Hakim via email.
+ This results in reduction estimates which are pretty much identical down to machine precision with those reached by including all observations at once.
+
+ Is coded very inefficiently, does not make benefit of only looking at the problem in observation space "reduced model space".
+ Currently is not tested for inputing multiple observations at once.
+
+ Output:
+ Individual variance reduction per each observation
+
+ Options, dJ_update_flag
+ 0: Hakim version, updates dJ using "appended state" approach
+ 1: Modified Hakim. Updates dJ, but then scales dJ to perfectly match the estimated variance reduction perfectly. Should be a very small tweak
+ 2: Scales dJ to match the reduced variance estimate. Should be extremely cheap and robust, but totally neglects covariance information.
+ 2: Doesn't update dJ at all, only used to check with what Tanya did
+
+
+ """
+
+ #For now I am not worried about being efficient
+ nobs = len(da_const["obs_loc"])
+ #if obs == []:
+ obs = np.arange(nobs)
+
+ j=0
+
+ nens = da_const["nens"]
+ ncyc = len(states[j]['response']['bf'])
+
+
+ var_reduction_total = []
+ var_reduction_individual = []
+ observation_ranking = []
+ real_reduction = []
+
+ #R is always the same
+ R = da_const["used_std_obs"]**2*np.identity(1) # Observation error corvariance matrix
+
+ for t in range(1,ncyc-1):
+ #for t in [10]:
+
+
+ x = states[j]['bg'][t][:,:]
+ dx = x.T-np.mean(x,axis=1)
+ dx = dx.T
+ A = np.dot(dx,dx.T)/(dx.shape[1]-1)
+
+ J= states[j]["response"]['bf'][t+1][:]
+ dJ = J-np.mean(J)
+
+
+ obs_ordered = []
+ obs_remain = list(obs)
+ vr_individual = np.zeros(nobs)
+ for o in range(nobs): #loop over the number of observations
+ #for o in range(1): #loop over the number of observations
+
+ vr_ind_tmp = np.zeros(nobs)
+ for oo in obs_remain: #loop over all observation locastions which have not been used yet used
+
+ H = np.identity(m_const["nx"])[da_const["obs_loc"][oo],:]
+ HAHt = np.dot(H,np.dot(A,H.T))
+ HAHtRinv= np.linalg.inv(HAHt+R)
+ dJHdxt = np.dot(dJ,np.dot(H,dx).T)/(nens-1)
+ vr_ind_tmp[oo] = - np.dot(dJHdxt,np.dot(HAHtRinv,dJHdxt))
+ #print(vr_ind_tmp)
+
+ ind_min = np.where(vr_ind_tmp==np.min(vr_ind_tmp))[0][0]
+ #print(ind_min)
+ vr_individual[ind_min] = vr_ind_tmp[ind_min]
+ obs_remain.remove(ind_min)
+
+ H = np.identity(m_const["nx"])[da_const["obs_loc"][ind_min],:]
+
+ #New we add in Tanyas reduction
+ E = np.matmul(A,H.T)
+ E = np.matmul(H,E)
+ E = E + R
+ alpha = 1./(1.+np.sqrt(R/E))
+
+
+
+ if dJ_update_flag==0 or dJ_update_flag==1:
+
+ #Moded Tanyas code because I had en error.
+ #Turned out it was a matrix multiplication issue I solved by switching the final matmul with an np.outer. It has the right dimension, but I fear my version will only work if its single point measurements.
+ #Now we include dJ into dx to include it in the update
+ dxJ = np.vstack([dx,dJ])
+ AxJ = np.dot(dxJ,dxJ.T)/(dxJ.shape[1]-1)
+ HxJ = np.hstack([H,np.array(0)])
+ HAHt = np.dot(HxJ,np.dot(AxJ,HxJ.T))
+ HAHtRinv= np.linalg.inv(HAHt+R)
+
+
+ K = np.dot(AxJ,HxJ.T)*HAHtRinv
+
+
+ # Update state vector
+ HdxJ = np.matmul(HxJ,dxJ)
+ dxJ = dxJ - alpha*np.outer(K,HdxJ)
+
+ dx = dxJ[:-1,:]
+ old_var_dJ = np.var(dJ,ddof=1)
+ dJ = dxJ[-1,:]
+ new_var_dJ = np.var(dJ,ddof=1)
+
+ if dJ_update_flag==1:
+ var_scaling = (old_var_dJ+ vr_ind_tmp[ind_min])/new_var_dJ
+ dJ=np.sqrt(var_scaling)*dJ
+ #print(old_var_dJ+vr_ind_tmp[ind_min],new_var_dJ,np.var(dJ))
+
+ if dJ_update_flag==2 or dJ_update_flag==3:
+
+ if dJ_update_flag==2:
+ var_scaling=(np.var(dJ,ddof=1)+vr_ind_tmp[ind_min])/np.var(dJ,ddof=1)
+ #New dJ
+ dJ=np.sqrt(var_scaling)*dJ
+
+ #Update dx
+ #Moded Tanyas code because I had en error.
+ #Turned out it was a matrix multiplication issue I solved by switching the final matmul with an np.outer. It has the right dimension, but I fear my version will only work if its single point measurements.
+ HAHt = np.dot(H,np.dot(A,H.T))
+ HAHtRinv= np.linalg.inv(HAHt+R)
+ K = np.dot(A,H.T)*HAHtRinv
+
+ # Update state vector
+ Hdx = np.matmul(H,dx)
+ dx = dx - alpha*np.outer(K,Hdx)
+
+ #Recalculating A makes things worse if you don't rescale dJ as well
+ A = np.dot(dx,dx.T)/(dx.shape[1]-1)
+
+ var_reduction_individual.append(vr_individual)
+ return var_reduction_individual
+
+
+
+def loc_matrix(da_const,m_const):
+ """
+ Creates the localization matrix, using either a gaussian or gaspari cohn function
+ """
+ C = np.zeros([m_const['nx'],m_const['nx']])
+ # I cheat a bit by mirroring the functions to accomodate for the repeating boundary conditions, but this should only lead to a maximum dx/2 error.
+ if da_const['loc_type']=='gaussian':
+ C[:,0] = gaussian_initial_condition(m_const["x_grid"],da_const["loc_length"])
+ if da_const['loc_type']=='gaspari_cohn':
+ C[:,0] = gaspari_cohn(m_const["x_grid"],da_const["loc_length"])
+
+ for i in range(1,m_const['nx']):
+ C[:,i] = np.hstack([C[-1,i-1],C[:-1,i-1]])
+ 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.
+ 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]
+ 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]
+
+ #Now we mirror things to zero for the periodic boundary domain
+ half_idx = int(ra.shape[0]/2)
+ 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):
+ """Gaspari-Cohn function, with no mirroring."""
+
+ ra = np.abs(x)/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]
+
+ return gp
+
+# def single_step_analysis_forecast(state,da_const,m_const,j,t,seed_add=1):
+# """
+# This function takes a given background state, to which is computes an analysis for provided observations. A blind forecast is computed from the background, and a forecast is computed from the analysis.
+
+# The main purpose of this approach of having a little isolated timestep is to enable applying different test to the same backgound state.
+
+# """
+
+# """
+# constant matrices that follow from previously defined constants
+# """
+# H = np.identity(m_const["nx"])[da_const["obs_loc"],:] # Observation operator
+# R = da_const["used_std_obs"]**2*np.identity(len(da_const["obs_loc"])) # Observation error corvariance matrix
+
+
+# #Construct localization matrix C if loc!=None
+# if da_const["loc"]:
+# C = loc_matrix(da_const,m_const)
+# else:
+# C=np.ones([m_const["nx"],m_const["nx"]])
+
+# """
+# Generate obs
+# """
+# #make obs by adding some noise, a fixed seed should change over time so that the differences are not always the same for each measurement location
+# #if da_const["fixed_seed"]==True:
+# #Changed to the following
+# np.random.seed((j+1)*(t+1)*(seed_add+1))
+
+
+# #changing obs to only exist where observations are to enable having multiple observations at the same space
+# #obs = np.zeros(len(da_const['obs_loc']))
+# #for o in range(len(da_const['obs_loc'])):
+# # obs[o] = state[j]['truth'][t][0] + np.random.normal(0,da_const["True_std_obs"])
+
+# obs = state[j]['truth'][t] + np.random.normal(0,da_const["True_std_obs"],m_const["nx"])
+
+# #Generate new truth constants and integrate in time
+# if da_const["fixed_seed"]==True:
+# np.random.seed(j)
+# else:
+# np.random.seed((j+1)*(t+1)*(seed_add+1))
+
+
+# u_truth = np.random.normal(m_const["u_ref"],m_const["u_std_truth"])
+# dhdt_truth = np.random.normal(m_const["dhdt_ref"],m_const["dhdt_std_truth"])
+# truth_forecast = linear_advection_model(state[j]['truth'][t],u_truth,dhdt_truth,m_const["dx"],da_const["dt"],da_const["nt"])
+
+
+# """
+# EnKF
+# """
+# if da_const['method'] == 'EnKF':
+# # Compute the background error covariance matrix
+# P = np.cov(state[j]['bg'][t])*C
+
+# # define relative weights of observations and background in terms of the Kalman Gain matrix of size
+# K = KalmanGain(P, R, H)
+
+# # Compute the analysis for each ensemble members
+# an = np.zeros((m_const["nx"],da_const["nens"]))
+# for i in range(da_const["nens"]):
+# an[:,i] = get_analysis(state[j]['bg'][t][:,i],obs,K,H,da_const)
+# """
+# LETKF
+# """
+# if da_const['method'] == 'LETKF':
+# an,bla = LETKF_analysis(state[j]['bg'][t],state[j]['obs'][t],m_const,da_const)
+
+
+# """
+# Predict blind forecast and forecast
+# """
+# bf = np.zeros((m_const["nx"],da_const["nens"]))
+# fc = np.zeros((m_const["nx"],da_const["nens"]))
+# for i in range(da_const["nens"]):
+# if da_const["fixed_seed"]==True: np.random.seed(i+j*da_const["nens"])
+# u_ens = np.random.normal(m_const["u_ref"],da_const["u_std_ens"])
+# dhdt_ens = np.random.normal(m_const["dhdt_ref"],da_const["dhdt_std_ens"])
+# bf[:,i] = linear_advection_model(state[j]['bg'][t][:,i],u_ens,dhdt_ens,m_const["dx"],da_const["dt"],da_const["nt"])
+# 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'] = state[j]['bg'][t]
+# quad_state['an'] = an
+# quad_state['bf'] = bf
+# quad_state['fc'] = fc
+# quad_state['tr_fc'] = truth_forecast
+# quad_state['tr_bg'] = state[j]['truth'][t]
+# quad_state['obs'] = obs
+# return quad_state
+
+def LETKF_analysis(bg,obs,m_const,da_const):
+ """
+ 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.
+ The application when not localized is pretty straight forward.
+ The localized version is coded maximully inefficiently. If needed two available options to speed it up are parallelization and batch processing.
+ Currently also only uses gaspari-cohn.
+ """
+ from scipy.linalg import sqrtm #Needed to calculate the matrix square root. Can lead to complex values due to numrical noise, which I deal with by only using real values.
+
+
+ L = da_const['nens']
+ H = np.identity(m_const["nx"])[da_const["obs_loc"],:] # Observation operator
+ R = da_const["used_std_obs"]**2*np.identity(len(da_const["obs_loc"]))
+ 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
+
+ y_obs = obs[da_const["obs_loc"]]
+ y_b = np.dot(H,x_b)
+ delta_y = y_obs-np.mean(y_b,axis=1)
+ Y_b = y_b.T-np.mean(y_b,axis=1)
+ Y_b = Y_b.T
+
+ 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 ensemlbe 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:
+ """
+ Localization baby!
+ Works by recalculating the whole process for every single grid point. For each point the inverse observation error is multiplied with the gaspari-cohn function.
+ Only implemented for diagonal R.
+ Accordingly, observations outside 2*the loc radius have no impact.
+ As mentioned, this is not computationally efficient, at all. While it does skip grid points that are completely unaffected by observations,
+ localization still increases the cost linearly with grid size.
+
+ """
+ x_grid_neg = -m_const['x_grid'][::-1]-m_const['dx']
+ x_grid_ext = m_const['x_grid']+(m_const['dx']*m_const['nx'])
+ N = m_const['nx']
+ x_a_loc = np.zeros((N,L))
+ x_ol_a_loc = np.zeros((N))
+
+ for g in range(N):
+
+ # first step is to get the disance of the current grid point to the observations.
+ # I know this isn't the quickest way, but when I coded it I was not feeling really smart and it works.
+ dist_reg = np.abs(m_const['x_grid'][da_const['obs_loc']]-m_const['x_grid'][g])
+ dist_neg = m_const['x_grid'][g]-x_grid_neg[da_const['obs_loc']]
+ dist_ext = x_grid_ext[da_const['obs_loc']]-m_const['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
+ gc_loc = gaspari_cohn_non_mirrored(dist,da_const['loc_length'])
+ 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,:]
+ 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+np.dot(X_b,w_ol_a)
+ 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
+
+ x_ol_a_loc[g] = x_ol_a[g]
+ x_a_loc[g,:] = x_a[g,:]
+ x_a = x_a_loc
+ x_ol_a = x_ol_a_loc
+
+
+
+ return x_a,x_ol_a
+
+
+
+
+def L2_regularized_inversion(A, b, alpha_init=1,alpha=None,mismatch_threshold=0.05):
+ """Instead of solving for Ax=b, which isn't possible if A is not invertible, the regularization instead minimizes ||Ax-b||^2 + || alpha x ||^2.
+ The solution is unique and well defined: x = (AA.T + alpha*alpha I )^-1 A.T b.
+
+ While this works fine with np.linalg.inv, I use .solve instead because it is roughly a factor 3 quicker.
+
+ I tried to find a easy way to link the starting alpha to the model space or ensemble size, but after not finding anything easy I just start with a prescribed value
+ which is 0.1 by default. It is checked that the mismatch between sum(Ax-b)/sum(b) does not exceed the mismatch_threshold, and if it does alpha is reduced by a factor of 2 until it does.
+ The mismatch_threshold is given in percent, and is used to check if ||(Ax-b)||/||b|| falls below the theshold.
+
+ alternative would be to try to use the things presented by Shu-Chih Yang's talk at the ISDA-online
+ kappa_req = 10000.
+ n = numbers of non-diagonal componentts of correlation matrix C
+ r = average of non-diagonal componentts of correlation matrix C
+ lamda_max = 1+(n_cols-1)*r
+ alpha = lamda_max/kappa_req
+ print(kappa_req,r,lamda_max,alpha)
+
+
+ """
+ n_cols = A.shape[1]
+ I = np.identity(n_cols)
+ #n_ens = b.size
+ #this is just a guess, but alpha should decrease with size, is not applied if a value other than 1 is applied
+ if alpha == None:
+ #alpha= 1./np.sqrt(n_cols)
+ alpha= alpha_init #1#n_cols/n_ens
+ #x = np.linalg.inv(A.T.dot(A) + alpha**2 *I).dot(A.T).dot(b)
+ x = np.linalg.solve((A.T.dot(A) + alpha**2 *I),(A.T).dot(b))#,rcond=-1)
+ while np.sum(np.abs(A.dot(x)-b))/np.sum(np.abs(b))> mismatch_threshold:
+ alpha = alpha/2.
+ x = np.linalg.solve((A.T.dot(A) + alpha**2 *I),(A.T).dot(b))#,rcond=-1)
+ #x = np.linalg.inv(A.T.dot(A) + alpha**2 *I).dot(A.T).dot(b)
+ #print('reducing regularization:',alpha,np.sum(np.abs(A.dot(x)-b))/np.sum(np.abs(b)))
+ else:
+ x = np.linalg.solve((A.T.dot(A) + alpha**2 *I),(A.T).dot(b))#,rcond=-1)
+ #x = np.linalg.inv(A.T.dot(A) + alpha**2 *I).dot(A.T).dot(b)
+
+
+ return x
+
+
+
+def single_step_analysis_forecast_v2(background,truth,da_const,m_const,model_seed=0,obs_seed=0):
+ """
+ The idea is that this should be merged into single_step_analysis_forecast, for a uniform processing chain for the paper. So all SWM changes will be included in flags.
+
+ What is a bit annoying is that the original was built around the states dictionary, I am not really a fan of the states setup, will switch to background and truth matrixes.
+
+ for now the perturbed observations in the EnKF for the SWM are the same as the actual errors.
+ """
+
+ """
+ constant matrices that follow from previously defined constants
+ """
+ if m_const['model'] == 'LA':
+ H = np.identity(m_const["nx"])[da_const["obs_loc"],:] # Observation operator
+ R = da_const["used_std_obs"]**2*np.identity(len(da_const["obs_loc"])) # Observation error corvariance matrix
+ if m_const['model'] == 'SWM':
+ H = np.identity(m_const["nx"]*3)[da_const["obs_loc"],:] # Observation operator
+ obs_error_vec = np.ones(len(da_const["obs_loc"]))*da_const['used_h_std_obs']
+ obs_error_vec[obs_error_vec<m_const['nx']] = da_const['used_u_std_obs']
+ obs_error_vec[obs_error_vec>2*m_const['nx']] = da_const['used_r_std_obs']
+ R = np.diag(obs_error_vec**2)
+
+
+ #Construct localization matrix C if loc!=None
+ if da_const["loc"]:
+ C = loc_matrix(da_const,m_const)
+ else:
+ C=np.ones([m_const["nx"],m_const["nx"]])
+ if m_const['model'] == 'SWM':
+ C = np.hstack([C,C,C])
+ C = np.vstack([C,C,C])
+
+ """
+ Generate obs
+ """
+ #make obs by adding some noise, a fixed seed should change over time so that the differences are not always the same for each measurement location
+ np.random.seed(obs_seed)
+
+ if m_const['model']=='LA': obs = truth + np.random.normal(0,da_const["True_std_obs"],m_const["nx"])
+ if m_const['model']=='SWM':
+ u_obs_noise =np.random.normal(0,da_const["True_u_std_obs"],m_const["nx"])
+ h_obs_noise =np.random.normal(0,da_const["True_h_std_obs"],m_const["nx"])
+ r_obs_noise =np.random.normal(0,da_const["True_r_std_obs"],m_const["nx"])
+ obs = truth + np.hstack([u_obs_noise,h_obs_noise,r_obs_noise])
+
+ #Generate new truth constants and integrate in time
+ if m_const['model']=='LA':
+ if da_const["fixed_seed"]==True and model_seed==0: np.random.seed(model_seed)
+ if da_const["fixed_seed"]==False and model_seed!=0 : np.random.seed(model_seed)
+
+ u_truth = np.random.normal(m_const["u_ref"],m_const["u_std_truth"])
+ dhdt_truth = np.random.normal(m_const["dhdt_ref"],m_const["dhdt_std_truth"])
+ truth_forecast = linear_advection_model(truth,u_truth,dhdt_truth,m_const["dx"],da_const["dt"],da_const["nt"])
+
+ if m_const['model']=='SWM':
+ #I assume we are going to have to add an empty second dimension here
+ truth_double = np.vstack([truth,truth]).T
+ truth_double_forecast = shallow_water(truth_double,m_const)
+ truth_forecast = truth_double_forecast[:,0]
+
+
+ """
+ EnKF
+ """
+ if da_const['method'] == 'EnKF':
+ # Compute the background error covariance matrix
+ P = np.cov(background)*C
+
+ # define relative weights of observations and background in terms of the Kalman Gain matrix of size
+ K = KalmanGain(P, R, H)
+
+ # Compute the analysis for each ensemble members
+ an = np.zeros_like(background)
+ if m_const['model'] == 'LA': obs_pert_vec = np.ones(len(obs))*da_const['pert_std_obs']
+ if m_const['model'] == 'SWM':
+ obs_pert_vec = np.hstack([np.ones(m_const['nx'])*da_const['pert_u_std_obs'],
+ np.ones(m_const['nx'])*da_const['pert_h_std_obs'],
+ np.ones(m_const['nx'])*da_const['pert_r_std_obs']])
+ for i in range(da_const["nens"]):
+ an[:,i] = get_analysis_v2(background[:,i],obs,K,H,obs_pert_vec)
+ """
+ LETKF
+ """
+ if da_const['method'] == 'LETKF':
+ an,bla = LETKF_analysis(background,obs,m_const,da_const)
+
+
+ """
+ Predict blind forecast and forecast
+ """
+ bf = np.zeros_like(background)
+ fc = np.zeros_like(background)
+
+ if m_const['model'] == 'LA':
+ for i in range(da_const["nens"]):
+ if da_const["fixed_seed"]==True: np.random.seed(i+model_seed*da_const["nens"])
+ u_ens = np.random.normal(m_const["u_ref"],da_const["u_std_ens"])
+ dhdt_ens = np.random.normal(m_const["dhdt_ref"],da_const["dhdt_std_ens"])
+ bf[:,i] = linear_advection_model(background[:,i],u_ens,dhdt_ens,m_const["dx"],da_const["dt"],da_const["nt"])
+ fc[:,i] = linear_advection_model(an[:,i],u_ens,dhdt_ens,m_const["dx"],da_const["dt"],da_const["nt"])
+ if m_const['model'] == 'SWM':
+ bf = shallow_water(background,m_const)
+ fc = shallow_water(an,m_const)
+
+
+ """
+ create dictionary to store this single step
+ """
+ quad_state = {}
+ quad_state['bg'] = background
+ quad_state['an'] = an
+ quad_state['bf'] = bf
+ quad_state['fc'] = fc
+ quad_state['tr_fc'] = truth_forecast
+ quad_state['tr_bg'] = truth
+ quad_state['obs'] = obs
+ return quad_state
+
+
+
+def vr_reloaded(background,truth,m_const,da_const,func_J=sum_mid_tri,
+ reduc = 1,reg_flag=1,
+ quad_state = None,dJdx_inv=None,alpha=None,mismatch_threshold=0.1,
+ iterative_flag=0,explicit_sens_flag = 1,exp_num=0,obs_seed=0,model_seed=0):
+
+ """
+ New version takes only the background and truth, then caculates the quad, finaly calculates the response function and the variance reduction.
+
+ The quad can be supplied to save time, which makes sense when you are chaning the response function.
+ Also is the way to go if you don't need the real reduction value, only the estimate.
+
+ Should also return the dJ values of all 4 ensembles (analysis, background, forecast, blind forecast)
+
+ exp_num x(experiment number j) needs to passed on to the forecast of LA somehow, because it determines the random model error of the ensembles
+ This will become a major part of the paper, and will see lots of refining. Important options are:
+
+ not implemented:
+ - state model reduction
+
+ response functions planned:
+ - mid tri, sum over middle of domain.
+ - sum over first and last third
+ - Triangle sum
+ - right/left rain amount.
+
+ Speed up options:
+ - If the same forecast can be recycled, the quad can be calculated once and than passed on.
+ - If the sensitivity can be recycled, that can be calculated once and then reused.
+
+ implemented:
+ -iteratively and all at once
+ -explicit vs implicit
+ -reduced model spacing by subsampling the grid. Reduc is the spacing, eg reduc 3 means every third grid point is used.
+ should find the nearest reduced model grid point for each observation.
+ -reg_flag added to use L2 regularization
+ """
+ if iterative_flag ==0: from scipy.linalg import sqrtm
+ ###########################################################################
+ #First, need we calculate the quad (analysis, forecast, blind forecast)
+ #A precomputed quad_state can be supplied to quickly use different response functions
+ ###########################################################################
+ if type(quad_state)== type(None):
+ quad_state = single_step_analysis_forecast_v2(background,truth,da_const,m_const,obs_seed=obs_seed,model_seed=model_seed)
+
+ ###########################################################################
+ # Next we need the response functions.
+ # We only really need the forecast and blind forecast, but for fun I'll calculate all for now
+ ###########################################################################
+
+ #For now I am not worried about being efficient
+ nobs = len(da_const["obs_loc"])
+ obs = np.arange(nobs)
+ nens = da_const["nens"]
+ nstate = len(truth)
+
+ bf_response = np.zeros(nens)
+ fc_response = np.zeros(nens)
+ an_response = np.zeros(nens)
+ bg_response = np.zeros(nens)
+ for n in range(da_const["nens"]):
+ bf_response[n] = func_J(quad_state["bf"][:,n])
+ fc_response[n] = func_J(quad_state["fc"][:,n])
+ an_response[n] = func_J(quad_state["an"][:,n])
+ bg_response[n] = func_J(quad_state["bg"][:,n])
+
+ J_dict = {}
+ J_dict['bf'] = bf_response
+ J_dict['bg'] = bg_response
+ J_dict['an'] = an_response
+ J_dict['fc'] = fc_response
+ J_dict['tr_bg'] = func_J(quad_state['tr_bg'])
+ J_dict['tr_fc'] = func_J(quad_state['tr_fc'])
+
+ ###########################################################################
+ # Creating the R,H, and C matrices we will need for the VR estimate
+ ###########################################################################
+
+ if m_const['model'] == 'LA':
+ H = np.identity(m_const["nx"])[da_const["obs_loc"],:] # Observation operator
+ R = da_const["used_std_obs"]**2*np.identity(len(da_const["obs_loc"])) # Observation error corvariance matrix
+ if m_const['model'] == 'SWM':
+ H = np.identity(m_const["nx"]*3)[da_const["obs_loc"],:] # Observation operator
+ obs_error_vec = np.ones(len(da_const["obs_loc"]))*da_const['used_h_std_obs']
+ obs_error_vec[obs_error_vec<m_const['nx']] = da_const['used_u_std_obs']
+ obs_error_vec[obs_error_vec>2*m_const['nx']] = da_const['used_r_std_obs']
+ R = np.diag(obs_error_vec**2)
+
+ if da_const["loc"]:
+ C = loc_matrix(da_const,m_const)
+ else:
+ C=np.ones([m_const["nx"],m_const["nx"]])
+ if m_const['model'] == 'SWM':
+ C = np.hstack([C,C,C])
+ C = np.vstack([C,C,C])
+
+
+
+
+
+
+ if reduc>1: #Does currently not work for
+ #Defining reduced model domain/state, and then defining obs location on new reduced grid
+ #For the SWM model this only really makes sense if nx is a multiple of reduc
+ reduced_model_domain = np.arange(0,nstate,reduc)
+ reduced_model_size = len(reduced_model_domain)
+ reduced_obs_loc = np.zeros(nobs).astype(int)
+ for o in range(nobs):
+ reduced_obs_loc[o] = (np.abs(reduced_model_domain-da_const['obs_loc'][o])).argmin()
+ #print('reduced_model_domain:',reduced_model_domain)
+ #print('reduced_model_size :',reduced_model_size )
+ #print('reduced_obs_loc :',reduced_obs_loc )
+
+ #Getting the localization matrix in reduced model space
+ H = np.identity(m_const["nx"])[reduced_model_domain,:]
+ C = np.dot(H,np.dot(C,H.T))
+ x = quad_state['bg'][reduced_model_domain,:]
+
+ else:
+ x = quad_state['bg'][:,:]
+
+ #x = states[j]['bg'][t][:,:]
+ dx = x.T-np.mean(x,axis=1)
+ dx = dx.T
+ dx_orig = dx+0
+
+ A = np.dot(dx,dx.T)/(dx.shape[1]-1)
+
+ J = bf_response
+ dJ = J-np.mean(J)
+ dJ_orig = J-np.mean(J)
+
+
+ ###############################################################################################
+ # Sensitivity
+ ###############################################################################################
+ #Covarianz between reponse function dJ and state ensemble dx
+ cov_dJdx_vec = np.dot(dJ,dx.T)/(dx.shape[1]-1)
+
+ if explicit_sens_flag==1:
+ #If a sensitivity is provided it is used instead of calculating it
+ if type(dJdx_inv) == np.ndarray:
+ #testing supplied sensitivity
+ rel_error_sens = np.sum(np.abs(A.dot(dJdx_inv)-cov_dJdx_vec))/np.sum(np.abs(cov_dJdx_vec))
+ if rel_error_sens>0.05:
+ print('using supplied sensitivity has a relative error of:',rel_error_sens)
+ #Computing the sensitivity, highly recommend using the regularized version before doing so.
+ else:
+ if reg_flag == 1:
+ dJdx_inv = L2_regularized_inversion(A,cov_dJdx_vec,alpha=alpha,mismatch_threshold=mismatch_threshold)
+ else:
+ A_inv = np.linalg.pinv(A)
+ dJdx_inv = np.dot(A_inv,cov_dJdx_vec)
+
+ estimated_J = bf_response + 0.
+
+
+ if iterative_flag ==0:
+
+ vr_individual = 0.
+
+ H = np.identity(m_const["nx"])[da_const["obs_loc"],:] # Observation operator
+ R = da_const["used_std_obs"]**2*np.identity(len(da_const["obs_loc"])) # Observation error corvariance matrix
+
+
+ R_obs = R # Observation error corvariance matrix
+ H_obs = H # Not the cleanest coding I know
+
+ if explicit_sens_flag ==1:
+ #Tanjas approach of calculating the square root K following Kalman gain, formula (10) Whitaker and Hamil 2002
+ E = np.matmul(C*A,H_obs.T)
+ E = np.matmul(H_obs,E)
+ E = E + R_obs
+ Esqrt = sqrtm(E)
+ #alpha = 1./(1.+np.sqrt(R_obs/E))
+ # Kalman gain, formula (10) Whitaker and Hamil 2002
+ K1 = np.matmul(C*A,H_obs.T)
+ K2 = np.linalg.inv(Esqrt)
+ K2 = K2.T
+ K3 = np.linalg.inv(Esqrt + sqrtm(R_obs))
+ K = np.matmul(K1,K2)
+ K = np.matmul(K,K3)
+
+ Hdx = np.matmul(H_obs,dx)
+ dx_prime = dx - np.dot(K,Hdx)
+
+ #estimated J calculated by updating the original dJ
+ estimated_J = estimated_J -np.dot(dJdx_inv,np.dot(K,Hdx))
+
+ #Using the cheaper variance calculation instead of going to A-B
+ new_J = np.dot(dJdx_inv.T,dx_prime)
+ vr_individual = np.var(new_J,ddof=1)-np.var(np.dot(dJdx_inv.T,dx),ddof=1)
+ vr_total = vr_individual
+ #print('all at once:',vr_individual)
+ dx = dx_prime
+
+ if explicit_sens_flag ==0:
+ HAHt = np.dot(H_obs,np.dot(C*A,H_obs.T))
+ HAHtRinv= np.linalg.inv(HAHt+R_obs)
+
+ dJHdxt = np.dot(dJ,np.dot(H_obs,dx).T)/(nens-1)
+ vr_individual = -np.dot(dJHdxt,np.dot(HAHtRinv,dJHdxt))
+ vr_total = vr_individual
+
+
+ if iterative_flag:
+ vr_individual = np.zeros(nobs)
+ for o in range(nobs): #loop over each observation individually
+
+ #New A after dx was updated
+ A = np.cov(dx,ddof=1)
+
+ #Selecting the single R value for this observation
+ R_obs = R[o,o]*np.identity(1) # Observation error corvariance matrix
+ H_obs = H[o,:] #Not sure about this :(
+ #if
+ ##H_rms = np.identity(nobs)[o,:]
+ #H_rms = np.identity(reduced_model_size)[reduced_obs_loc[o],:]
+ ##H = np.identity(m_const["nx"])[da_const["obs_loc"][o],:]
+#
+ ##HAHt = np.dot(H,np.dot(A,H.T))
+ ##HAHtRinv= np.linalg.inv(HAHt+R)
+ ##dJHdxt = np.dot(dJ,np.dot(H,dx).T)/(nens-1)
+
+ #Now we get the change in dx, which uses the localized A matrix and a square root
+ E = np.matmul(C*A,H_obs.T)
+ E = np.matmul(H_obs,E)
+ E = E + R_obs
+ alpha = 1./(1.+np.sqrt(R_obs/E))
+
+
+ HAHt = np.dot(H_obs,np.dot(C*A,H_obs.T))
+ HAHtRinv= np.linalg.inv(HAHt+R_obs)
+ #print((C*A).shape,H_obs.T.shape,np.dot(C*A,H_obs.T).shape,HAHtRinv.shape)
+ ##
+
+ if explicit_sens_flag ==1:
+
+ K = np.dot(C*A,H_obs.T)*HAHtRinv
+ # Update state vector
+ Hdx = np.matmul(H_obs,dx)
+ dx_prime = dx - alpha*np.outer(K,Hdx)
+
+ #Update expected J for each ensemble member
+ estimated_J = estimated_J -np.dot(dJdx_inv,alpha*np.outer(K,Hdx))
+
+ #A_new = np.dot(dx,dx.T)/(dx.shape[1]-1)
+ #Now we get the variance estimate using the difference between new and old B
+ #dSigma = np.matmul(A_new-A,dJdx_inv.T)
+ #print('wtf: np.sum(np.abs(A_new-A_old)) ',o,np.sum(np.abs(A_new-A)))
+ #dSigma = np.matmul(dJdx_inv,dSigma)
+ #vr_individual[o] = dSigma#np.matmul(np.atleast_2d(np.diag(A_new-A)),(dJdx_inv.T**2))
+
+ vr_individual[o] = np.var(np.dot(dJdx_inv.T,dx_prime),ddof=1)-np.var(np.dot(dJdx_inv.T,dx),ddof=1)
+ dx = dx_prime
+ if explicit_sens_flag==0:
+ dJHdxt = np.dot(dJ,np.dot(H_obs,dx).T)/(nens-1)
+ vr_individual[o] = -np.dot(dJHdxt,np.dot(HAHtRinv,dJHdxt))
+
+ #Still needs localization!
+
+ #Now we include dJ into dx to include it in the update
+ dxJ = np.vstack([dx,dJ])
+ AxJ = np.dot(dxJ,dxJ.T)/(dxJ.shape[1]-1)
+ HxJ = np.hstack([H_obs,np.array(0)])
+ CxJ = np.ones([nstate+1,nstate+1])
+ CxJ[:nstate,:nstate] = C
+ HAHt = np.dot(HxJ,np.dot(CxJ*AxJ,HxJ.T))
+ HAHtRinv= np.linalg.inv(HAHt+R_obs)
+
+ #print(AxJ.shape,CxJ.shape)
+ #print(AxJ.shape,HxJ.T.shape,np.dot(AxJ,HxJ.T).shape,HAHtRinv.shape)
+ K = np.dot(CxJ*AxJ,HxJ.T)*HAHtRinv
+
+
+ # Update state vector
+ HdxJ = np.matmul(HxJ,dxJ)
+ dxJ = dxJ - alpha*np.outer(K,HdxJ)
+
+ dx = dxJ[:-1,:]
+ #old_var_dJ = np.var(dJ)
+ dJ = dxJ[-1,:]
+ estimated_J = dJ+np.mean(bf_response)
+ #new_var_dJ = np.var(dJ)
+
+ #Final var reduciton estimate
+ if explicit_sens_flag==0: vr_total=np.sum(vr_individual)
+
+ if explicit_sens_flag==1:
+ vr_total=np.var(np.dot(dJdx_inv.T,dx),ddof=1)-np.var(np.dot(dJdx_inv.T,dx_orig),ddof=1)
+
+ #Checking different formulation
+
+
+ #
+
+ J_dict['es'] = estimated_J
+
+ J_fc= fc_response
+ dJ_fc = J_fc-np.mean(J_fc)
+ real_reduction=np.var(dJ_fc,ddof=1) - np.var(dJ_orig,ddof=1)
+
+ return vr_total,vr_individual,real_reduction,J_dict,dJdx_inv,quad_state,dx
+
+
+def vr_individual_loc(background,truth,m_const,da_const,response_func=1,abs_flag=0,
+ quad_state = None,exp_num=0,
+ advect_flag = 0,obs_seed=0,random_samples=0,seed_samples=0,error_u=100,error_u_ens=100):
+
+ """
+ Based on vr_reloaded, but adapted to individually localize the covariances of the response functions at each point.
+
+ important parameters:
+ response_func: controls the response, 1 results in the default sum over the mean, 2 in the mean over the mean
+ abs_flag: determines if the absolute value of the state vector is used for the response function.
+
+ If advect_flag==1, the localization matrix is advected using the mo_const advection speed which should be close to the ensemble mean. The same advection function used as in the linear advection model.
+ If advect_flag==2, the localization matrix is advected using the individual ensemble members, so that in the end the localization vector is broadened accoring to ensemble disperstion.
+
+ random_samples=0,seed_samples=0,error_u=100,error_u_ens=100: these all are various ways of adding an error to the advection term.
+
+
+ For now only using the all at once method, and it probably won't work for only a single observation.
+
+ The equations should be in the paper draft, but I made quite the mistake by thinking that the A at the end of A'-A = -KHA is also localized.
+
+ Possible speed ups not used so far:
+ * I could speed this up a bit by only calculating the localized correlations where they will be needed (so where H says they should be), but that is a bit of work to do well so for now I will leave it out.
+ * Given that we use Gaspari-Cohn localization most of the time, I might be able to save some time by only calculating covariances where C_i is not zero. But probably won't make a difference for toy model
+ * We currently assign j_i a value at every grid point, even though we skip calculations where response_s = 0. It would be more memory efficient to allocate J_i in a restricted area, but again I think this won't matter in a toymodel
+ * Possibly the biggest speed up would be to avoid the loop when calculating the response functions, which I haven't bothered with yet.
+
+
+
+ """
+ ##########################################################################
+ #We begin by computing the vecors which determine the response function
+ ###########################################################################
+ response_s = np.zeros(m_const['nx'])
+ response_c = np.ones(m_const['nx'])
+ nx = m_const['nx']
+ if response_func ==1:
+ idx_str = int(nx/3.)
+ idx_end = int(2*nx/3.)
+ response_s[idx_str:idx_end] = 1
+ if response_func ==2:
+ idx_str = int(nx/3.)
+ idx_end = int(2*nx/3.)
+ response_s[idx_str:idx_end] = 1./np.float(len(response_s[idx_str:idx_end]))
+
+
+
+ ###########################################################################
+ #First, need we calculate the quad (analysis, forecast, blind forecast)
+ #A precomputed quad_state can be supplied to quickly use different response functions
+ ###########################################################################
+ if type(quad_state)== type(None):
+ quad_state = single_step_analysis_forecast_v2(background,truth,da_const,m_const,obs_seed=obs_seed)
+
+ ###########################################################################
+ # Next we need the response functions.
+ # We only really need the forecast and blind forecast, but for fun I'll calculate all for now
+ ###########################################################################
+
+ #For now I am not worried about being efficient
+ nobs = len(da_const["obs_loc"])
+ obs = np.arange(nobs)
+ nens = da_const["nens"]
+ nstate = len(truth)
+
+ bf_response = np.zeros(nens)
+ fc_response = np.zeros(nens)
+ an_response = np.zeros(nens)
+ bg_response = np.zeros(nens)
+ for n in range(da_const["nens"]):
+ if abs_flag==0:
+ bf_response[n] = np.sum(response_s * np.power(quad_state["bf"][:,n],response_c))
+ fc_response[n] = np.sum(response_s * np.power(quad_state["fc"][:,n],response_c))
+ an_response[n] = np.sum(response_s * np.power(quad_state["an"][:,n],response_c))
+ bg_response[n] = np.sum(response_s * np.power(quad_state["bg"][:,n],response_c))
+ if abs_flag==1:
+ bf_response[n] = np.sum(response_s * np.power(np.abs(quad_state["bf"][:,n]),response_c))
+ fc_response[n] = np.sum(response_s * np.power(np.abs(quad_state["fc"][:,n]),response_c))
+ an_response[n] = np.sum(response_s * np.power(np.abs(quad_state["an"][:,n]),response_c))
+ bg_response[n] = np.sum(response_s * np.power(np.abs(quad_state["bg"][:,n]),response_c))
+
+ J_dict = {}
+ J_dict['bf'] = bf_response
+ J_dict['bg'] = bg_response
+ J_dict['an'] = an_response
+ J_dict['fc'] = fc_response
+ J_dict['tr_bg'] = np.sum(response_s * np.power(quad_state["tr_bg"][:],response_c))
+ J_dict['tr_fc'] = np.sum(response_s * np.power(quad_state["tr_fc"][:],response_c))
+ if abs_flag==1:
+ J_dict['tr_bg'] = np.sum(response_s * np.power(np.abs(quad_state["tr_bg"][:]),response_c))
+ J_dict['tr_fc'] = np.sum(response_s * np.power(np.abs(quad_state["tr_fc"][:]),response_c))
+
+ ###########################################################################
+ # Creating the R,H, and C matrices we will need for the VR estimate
+ ###########################################################################
+
+ if m_const['model'] == 'LA':
+ H = np.identity(m_const["nx"])[da_const["obs_loc"],:] # Observation operator
+ R = da_const["used_std_obs"]**2*np.identity(len(da_const["obs_loc"])) # Observation error corvariance matrix
+ if m_const['model'] == 'SWM':
+ H = np.identity(m_const["nx"]*3)[da_const["obs_loc"],:] # Observation operator
+ obs_error_vec = np.ones(len(da_const["obs_loc"]))*da_const['used_h_std_obs']
+ obs_error_vec[obs_error_vec<m_const['nx']] = da_const['used_u_std_obs']
+ obs_error_vec[obs_error_vec>2*m_const['nx']] = da_const['used_r_std_obs']
+ R = np.diag(obs_error_vec**2)
+
+ if da_const["loc"]:
+ C = loc_matrix(da_const,m_const)
+ else:
+ C=np.ones([m_const["nx"],m_const["nx"]])
+ if m_const['model'] == 'SWM':
+ C = np.hstack([C,C,C])
+ C = np.vstack([C,C,C])
+
+
+ x = quad_state['bg'][:,:]
+
+ dx = x.T-np.mean(x,axis=1)
+ dx = dx.T
+
+ A = np.dot(dx,dx.T)/(dx.shape[1]-1)
+
+ #Now we need the vector of individual js for all ensembles
+ J = np.zeros_like(dx)
+ for n in range(da_const["nens"]):
+ if abs_flag==0:
+ J[:,n] = response_s * np.power(quad_state["bf"][:,n],response_c)
+ if abs_flag==1:
+ J[:,n] = response_s * np.power(np.abs(quad_state["bf"][:,n]),response_c)
+ #J[:,n] = response_s * np.power(quad_state["bg"][:,n],response_c)
+ dJ = J.T-np.mean(J,axis=1)
+ dJ = dJ.T
+
+ dJ_sum = bf_response - np.mean(bf_response)
+
+
+ H = np.identity(m_const["nx"])[da_const["obs_loc"],:] # Observation operator
+ R = da_const["used_std_obs"]**2*np.identity(len(da_const["obs_loc"])) # Observation error corvariance matrix
+
+
+ R_obs = R # Observation error corvariance matrix
+ H_obs = H # Not the cleanest coding I know
+
+
+ HAHt = np.dot(H_obs,np.dot(C*A,H_obs.T))
+ HAHtRinv= np.linalg.inv(HAHt+R_obs)
+
+ dJHdxt = np.dot(dJ_sum,np.dot(H_obs,dx).T)/(nens-1)
+ CdJdX = np.zeros(m_const['nx'])
+
+ #Since we no longer need C for the HCAHT which was calculated above, we can just advect the C matrix for the correlations and not make a new one.
+ if advect_flag > 0:
+ if advect_flag ==1:
+ u_advect =m_const['u_ref']*error_u/100.
+ C[:,0] = semi_lagrangian_advection(C[:,0],m_const['dx'],u_advect,da_const['dt'])
+ #for c in range(C.shape[0]):
+ # C[c,:] = semi_lagrangian_advection(C[c,:],m_const['dx'],-m_const['u_ref'],da_const['dt'])
+
+ if advect_flag ==2:
+ C_adv = C[:,0]*0.
+ if random_samples == 0:
+ for i in range(da_const["nens"]):
+ if da_const["fixed_seed"]==True: np.random.seed(i)
+ u_ens = np.random.normal(m_const["u_ref"],da_const["u_std_ens"])
+ u_ens = m_const["u_ref"]*error_u/100.+(u_ens-m_const["u_ref"])*error_u_ens/100.
+ #print('u_ens',u_ens)
+ C_adv = C_adv + semi_lagrangian_advection(C[:,0],m_const['dx'],u_ens,da_const['dt'])
+ #C_adv[:,0] = C_adv[:,0] + semi_lagrangian_advection(C[:,0],m_const['dx'],10.,da_const['dt'])
+ C[:,0] = C_adv/np.float(da_const["nens"])
+ else:
+ np.random.seed(random_samples+seed_samples)
+ randomized_ens = np.random.choice(np.arange(0,da_const["nens"]), random_samples, replace=False)
+ for r in randomized_ens:
+ #print(r)
+ if da_const["fixed_seed"]==True: np.random.seed(r)
+ u_ens = np.random.normal(m_const["u_ref"],da_const["u_std_ens"])
+ u_ens = m_const["u_ref"]*error_u/100.+(u_ens-m_const["u_ref"])*error_u_ens/100.
+ #print('u_ens',u_ens)
+ C_adv = C_adv + semi_lagrangian_advection(C[:,0],m_const['dx'],u_ens,da_const['dt'])
+ #C_adv[:,0] = C_adv[:,0] + semi_lagrangian_advection(C[:,0],m_const['dx'],10.,da_const['dt'])
+ C[:,0] = C_adv/np.float(random_samples)
+
+ for i in range(1,m_const['nx']):
+ C[:,i] = np.hstack([C[-1,i-1],C[:-1,i-1]])
+
+
+ #for n in range(m_const["nx"]): #Not sure if the if loop is much quicker than just calculating the dot product everywhere ;)
+ # if response_s[n] != 0: CdJdX += C[n,:]*np.dot(dJ[n,:],dx.T)/(da_const['nens']-1)
+ # Somewhat quicker way that avoids the loop
+ CdJdX=np.sum(C*np.dot(dJ,dx.T)/(da_const['nens']-1.),axis=0)
+
+ CdJdXH = np.dot(CdJdX,H.T)
+ vr_total = - np.dot(CdJdXH,np.dot(HAHtRinv,dJHdxt).T)
+ J_fc= fc_response
+ dJ_fc = J_fc-np.mean(J_fc)
+ real_reduction=np.var(dJ_fc,ddof=1) - np.var(bf_response,ddof=1)
+
+
+ return vr_total,real_reduction,quad_state,J_dict#,C
+
+
+#==========================================================
+# 22 version with satelitte data
+#==========================================================
+
+
+
+def LETKF_analysis_22(bg,obs,obs_sat,m_const,da_const,sat_operator):
+ """
+ Revised version of LETKF_analysis that allows for a satelitte observation
+
+ 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.
+ The application when not localized is pretty straight forward.
+ The localized version is coded maximully inefficiently. If needed two available options to speed it up are parallelization and batch processing.
+ Currently also only uses gaspari-cohn.
+ """
+ from scipy.linalg import sqrtm #Needed to calculate the matrix square root. Can lead to complex values due to numrical noise, which I deal with by only using real values.
+
+
+
+ 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_sat = sat_operator(quad['tr_bg'],window=window)+ np.random.normal(0,da_const["used_std_obs_sat"],size=bg.shape[0])
+ 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 =np.ones(n_obs)
+# r[:n_obs_h] = da_const["used_std_obs"]**2.
+# r[n_obs_h:] = da_const["used_std_obs_sat"]**2.
+ R = da_const['R'] #np.diag(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:
+ x_grid_neg = -m_const['x_grid'][::-1]-m_const['dx']
+ x_grid_ext = m_const['x_grid']+(m_const['dx']*m_const['nx'])
+ N = m_const['nx']
+ x_a_loc = np.zeros((N,L))
+ x_ol_a_loc = np.zeros((N))
+ for g in range(N):
+ #for g in range(1):
+ dist_reg = np.abs(m_const['x_grid'][obs_loc]-m_const['x_grid'][g])
+ dist_neg = m_const['x_grid'][g]-x_grid_neg[obs_loc]
+ dist_ext = x_grid_ext[obs_loc]-m_const['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
+ gc_loc = gaspari_cohn_non_mirrored(dist,da_const['loc_length'])
+ 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,:]
+ 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+np.dot(X_b,w_ol_a)
+ 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
+
+ x_ol_a_loc[g] = x_ol_a[g]
+ x_a_loc[g,:] = x_a[g,:]
+ x_a = x_a_loc
+ x_ol_a = x_ol_a_loc
+ 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):
+ """
+ Computes Y_b, so the ensembles transfered to observation space, when a satetellite operator is used.
+
+ Should be used for LETKF, as well as for the variance reduction estimate using the modified Kalman gain
+
+ Important, X must not be the ensemble deviations, but the full ensemble state
+
+ Returns the ensemble deviations as well as the mean
+
+ """
+ if np.max(np.mean(X,axis=1))<1e-10:
+ print('I think you accidentally used the ensemble deviations when calculating from state to observation space')
+ 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
+
+ if n_obs_h>0:
+ H = np.identity(m_const["nx"])[da_const["obs_loc"],:] # Observation operator
+ y_b = np.dot(H,X)
+
+ if n_obs_sat>0:
+ #here we will use sat instead of Y and y
+ sat_b = sat_operator(X)[da_const["obs_loc_sat"],:]
+
+ if n_obs_h>0 and n_obs_sat>0 :
+ y_b = np.hstack([y_b.T,sat_b.T])
+ y_b = y_b.T
+
+ if n_obs_sat>0 and n_obs_h==0 :
+ y_b = sat_b
+
+ y_ol_b =np.mean(y_b,axis=1)
+ Y_b = y_b.T-y_ol_b
+ Y_b = Y_b.T
+
+ return Y_b, y_ol_b
+
+
+
+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
+
+ 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"])
+ 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']
+
+ dY_b, y_ol_b = state_to_observation_space(bg,m_const,da_const,sat_operator)
+ 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']
+ 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 = 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)
+ return K
+
+
+def square_root_Kalman_gain_observation_deviations(bg,m_const,da_const,sat_operator):
+ """
+ Calculates modified square root 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
+ """
+ from scipy.linalg import sqrtm #Needed to calculate the matrix square root. Can lead to complex values due to numrical noise, which I deal with by only using real values.
+ 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']
+
+ dY_b, y_ol_b = state_to_observation_space(bg,m_const,da_const,sat_operator)
+ 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']
+
+
+ #Tryig to do this directly from the formula because Tanyas implementation works but slightly confuses me.
+ if da_const['loc']:
+ 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([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)
+ # Separating last two terms is modified Kalman gain into terms I and II
+ I = (np.linalg.inv(sqrtm(YYR))).T
+ II = np.linalg.inv(sqrtm(YYR)+sqrtm(R))
+ K = np.dot(XY,np.dot(I,II))
+
+ return K
+
+
+def localization_matrices_observation_space(m_const,da_const):
+ """
+ Creates the two localization matrices needed to calculate the localized Kalman gain with oberservation space ensemble pertubations.
+
+ The distance calculation is still far from elegant, but I can't be bothered to make it faster as long as the LETKF is still the numerical bottleneck.
+ """
+ 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_obs = np.zeros([n_obs,n_obs])
+ L_obs_state = np.zeros([m_const['nx'],n_obs])
+
+ x_grid_neg = -m_const['x_grid'][::-1]-m_const['dx']
+ x_grid_ext = m_const['x_grid']+(m_const['dx']*m_const['nx'])
+
+ for o in range(n_obs):
+ g = obs_loc[o]
+ # First the one using the distance between observations and the state
+ dist_reg = np.abs(m_const['x_grid']-m_const['x_grid'][g])
+ dist_neg = m_const['x_grid'][g]-x_grid_neg
+ dist_ext = x_grid_ext-m_const['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)
+ gc_loc = gaspari_cohn_non_mirrored(dist,da_const['loc_length'])
+ L_obs_state[:,o] = gc_loc
+
+
+ L_obs[o,:] = L_obs_state[obs_loc,o]
+
+ return L_obs, L_obs_state
+
+def generate_obs_22_single(truth,m_const,da_const,sat_operator,obs_seed):
+ """
+ Generates the observations for the single forecast analysis, including the sat obs, for a provided seed
+
+ """
+ np.random.seed(obs_seed)
+ obs = truth + np.random.normal(0,da_const["True_std_obs"],m_const["nx"])
+ if len(da_const['obs_loc_sat'])>0:
+ truth_sat = sat_operator(truth)
+ obs_sat = truth_sat + np.random.normal(0,da_const["True_std_obs_sat"],m_const["nx"])
+
+ obs_sat[obs_sat>1.]= 1.
+ obs_sat[obs_sat<0.]= 0.
+ else:
+ obs_sat = np.zeros(m_const['nx'])
+
+ return obs, obs_sat
+
+def single_step_analysis_forecast_22(background,truth,da_const,m_const,sat_operator,model_seed=0,obs_seed=0):
+ """
+ Revisted version of single step analysis forecast to deal with sat data as well.
+
+ For now I will remove the SWM stuff, which can still be found in the older version
+ """
+ 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,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)
+
+
+
+
+
+ """
+ Predict blind forecast and forecast
+ """
+ bf = np.zeros_like(background)
+ fc = np.zeros_like(background)
+
+ for i in range(da_const["nens"]):
+ np.random.seed(i+model_seed)
+ u_ens = np.random.normal(m_const["u_ref"],da_const["u_std_ens"])
+ dhdt_ens = np.random.normal(m_const["dhdt_ref"],da_const["dhdt_std_ens"])
+ bf[:,i] = linear_advection_model(background[:,i],u_ens,dhdt_ens,m_const["dx"],da_const["dt"],da_const["nt"])
+ fc[:,i] = linear_advection_model(an[:,i],u_ens,dhdt_ens,m_const["dx"],da_const["dt"],da_const["nt"])
+
+
+ quad_state['bg'] = background
+ quad_state['an'] = an
+ quad_state['bf'] = bf
+ quad_state['fc'] = fc
+ quad_state['tr_bg'] = truth
+ quad_state['obs'] = obs
+ return quad_state
+
+
+def create_states_dict_22(j,states,m_const,da_const,sat_operator):
+ """
+ Generates the initial analysis and truth.
+ Also creates the "states" dictionary where the analysis ensemble, background ensemble, truth and observations are all going to be stored.
+ Very memory hungry, as everything from all assimilation time steps and all experiments is stored.
+ Works find for simple model though.
+
+ A fixed seed is used by default so that the model errors of all ensemble members is constant. But this can also be randomized.
+ Alternative version would be to generate the model errors and store them, but this has not happened yet. Would be necessary to test some parameter estimation tests.
+
+ Modified version of the Yvonne setup to work with the linear advection model
+
+ Todo:
+ - Describe states dictionary here.
+ - make model errors stored variables so enable parameter estimation
+ - Maybe make a more sensible name, such as init_da_ensembles
+ - I am not convinced this is the best way to generate the initial analysis.
+ - Might be best to change to dictionary time axis. Forecast and anaylis at the same time have difference of 1 in the time integer.
+ """
+ nx = m_const["nx"]
+
+ #initial conditions
+ if m_const['init_func']=='gaus': h_initial = gaussian_initial_condition(m_const["x_grid"],m_const["h_init_std"])
+ if m_const['init_func']=='sine': h_initial = sine_initial_condition( m_const["x_grid"],m_const["sine_init"])
+
+
+ #Generate truth and first observations
+ truth, obs, obs_sat = generate_obs_22(h_initial,h_initial,m_const,da_const,j,0,sat_operator)
+
+
+ an = np.zeros((nx,da_const["nens"]))
+
+ #First rows full of nans :)
+ bg = np.zeros((nx,da_const["nens"]))
+ bg[:] = np.nan
+ obs = np.zeros(nx)
+ obs[:] = np.nan
+
+ for i in range(da_const["nens"]):
+ np.random.seed(i+j*da_const["nens"]+da_const['ens_seed'])
+ if da_const["init_noise"]>0:
+ h_ens = np.random.normal(h_initial,da_const["init_noise"])
+ else:
+ h_ens = h_initial
+ if da_const["init_spread"]>0:
+ #initial spread generated by moving waves forward and backward and up and down using
+ #da_const["init_spread_h"] and da_const["init_spread_x"]
+ x_displace = np.random.normal(0.,da_const["init_spread_x"])
+ h_displace = np.random.normal(0.,da_const["init_spread_h"])
+ u_tmp = x_displace/da_const["dt"]
+ h_ens = semi_lagrangian_advection(h_ens,m_const["dx"],u_tmp,da_const["dt"])
+ h_ens = h_ens+h_displace
+
+ u_ens = np.random.normal(m_const["u_ref"],da_const["u_std_ens"])
+ dhdt_ens = np.random.normal(m_const["dhdt_ref"],da_const["dhdt_std_ens"])
+ an[:,i] = linear_advection_model(h_ens,u_ens,dhdt_ens,m_const["dx"],da_const["dt"],da_const["nt"])
+
+ states[j]={}
+ states[j]['bg']=[bg]
+ states[j]['an']=[an]
+ states[j]['truth']=[truth]
+ states[j]['obs']=[obs]
+ states[j]['obs_sat']=[obs_sat]
+ return an, truth, states
+
+
+ return DA_const
+
+
+def vr_reloaded_22(background,truth,m_const,da_const,sat_operator,
+ func_J=sum_mid_tri,
+ reduc = 1,reg_flag=1,
+ quad_state = None,dJdx_inv=None,alpha=0.01,mismatch_threshold=0.1,
+ iterative_flag=0,explicit_sens_flag = 1,exp_num=0,obs_seed=0,model_seed=0):
+
+ """
+ Version of vr_reloaded that can now also use sat obs.
+ Takes only the background and truth, then caculates the quad, finaly calculates the response function and the variance reduction.
+
+ The quad can be supplied to save time, which makes sense when you comparing different methods for the same experiment calculating.
+
+ Should also return the dJ values of all 4 ensembles (analysis, background, forecast, blind forecast)
+
+
+ -reg_flag added to use L2 regularization
+
+ not implemented:
+ - state model reduction
+
+
+ Speed up options:
+ - If the same forecast can be recycled, the quad can be calculated once and than passed on.
+ - If the sensitivity can be recycled, that can be calculated once and then reused.
+ - Don't use LETKF
+
+ First getting the explicit all at once verion working, will then add the various approaches
+
+ to be implemented:
+ -iteratively and all at once
+ -explicit vs implicit
+ -reduced model spacing by subsampling the grid. Reduc is the spacing, eg reduc 3 means every third grid point is used.
+ should find the nearest reduced model grid point for each observation.
+
+
+
+ """
+ if iterative_flag ==0: from scipy.linalg import sqrtm
+ ###########################################################################
+ #First, need we calculate the quad (analysis, forecast, blind forecast)
+ #A precomputed quad_state can be supplied to quickly use different response functions
+ ###########################################################################
+ if type(quad_state)== type(None):
+ quad_state = single_step_analysis_forecast_22(background,truth,da_const,m_const,sat_operator,obs_seed=obs_seed,model_seed=model_seed)
+
+ ###########################################################################
+ # Next we need the response functions.
+ # We only really need the forecast and blind forecast, but for fun I'll calculate all for now
+ ###########################################################################
+
+ #For now I am not worried about being efficient
+ nobs = len(da_const["obs_loc"])
+ obs = np.arange(nobs)
+ nens = da_const["nens"]
+ nstate = len(truth)
+
+ bf_response = np.zeros(nens)
+ fc_response = np.zeros(nens)
+ an_response = np.zeros(nens)
+ bg_response = np.zeros(nens)
+ for n in range(da_const["nens"]):
+ bf_response[n] = func_J(quad_state["bf"][:,n])
+ fc_response[n] = func_J(quad_state["fc"][:,n])
+ an_response[n] = func_J(quad_state["an"][:,n])
+ bg_response[n] = func_J(quad_state["bg"][:,n])
+
+ J_dict = {}
+ J_dict['bf'] = bf_response
+ J_dict['bg'] = bg_response
+ J_dict['an'] = an_response
+ J_dict['fc'] = fc_response
+ J_dict['tr_bg'] = func_J(quad_state['tr_bg'])
+
+
+
+
+
+
+
+ #if reduc>1: #Does currently not work for
+ # #Defining reduced model domain/state, and then defining obs location on new reduced grid
+ # #For the SWM model this only really makes sense if nx is a multiple of reduc
+ # reduced_model_domain = np.arange(0,nstate,reduc)
+ # reduced_model_size = len(reduced_model_domain)
+ # reduced_obs_loc = np.zeros(nobs).astype(int)
+ # for o in range(nobs):
+ # reduced_obs_loc[o] = (np.abs(reduced_model_domain-da_const['obs_loc'][o])).argmin()
+ # #print('reduced_model_domain:',reduced_model_domain)
+ # #print('reduced_model_size :',reduced_model_size )
+ # #print('reduced_obs_loc :',reduced_obs_loc )
+ #
+ # #Getting the localization matrix in reduced model space
+ # H = np.identity(m_const["nx"])[reduced_model_domain,:]
+ # C = np.dot(H,np.dot(C,H.T))
+ # x = quad_state['bg'][reduced_model_domain,:]
+ #
+ #else:
+ x = quad_state['bg'][:,:]
+
+ dx = x.T-np.mean(x,axis=1)
+ dx = dx.T
+ dx_orig = dx+0
+
+ A = np.dot(dx,dx.T)/(dx.shape[1]-1)
+
+ J = bf_response
+ dJ = J-np.mean(J)
+ dJ_orig = J-np.mean(J)
+
+
+ ###############################################################################################
+ # Sensitivity
+ ###############################################################################################
+ #Covarianz between reponse function dJ and state ensemble dx
+ cov_dJdx_vec = np.dot(dJ,dx.T)/(dx.shape[1]-1)
+
+ if explicit_sens_flag==1:
+ #If a sensitivity is provided it is used instead of calculating it
+ if type(dJdx_inv) == np.ndarray:
+ #testing supplied sensitivity
+ rel_error_sens = np.sum(np.abs(A.dot(dJdx_inv)-cov_dJdx_vec))/np.sum(np.abs(cov_dJdx_vec))
+ if rel_error_sens>0.05:
+ print('using supplied sensitivity has a relative error of:',rel_error_sens)
+ #Computing the sensitivity, highly recommend using the regularized version before doing so.
+ else:
+ if reg_flag == 1:
+ dJdx_inv = L2_regularized_inversion(A,cov_dJdx_vec,alpha=alpha,mismatch_threshold=mismatch_threshold)
+ else:
+ A_inv = np.linalg.pinv(A)
+ dJdx_inv = np.dot(A_inv,cov_dJdx_vec)
+
+ estimated_J = bf_response + 0.
+
+
+ #if iterative_flag ==0:
+
+ vr_individual = 0.
+
+
+ if explicit_sens_flag ==1:
+ #Tanjas approach of calculating the square root K following Kalman gain, formula (10) Whitaker and Hamil 2002
+ sqK = square_root_Kalman_gain_observation_deviations(x,m_const,da_const,sat_operator)
+ bg_obs_deviations,bg_obs_ol = state_to_observation_space(x,m_const,da_const,sat_operator)
+ dx_prime = dx - np.dot(sqK,bg_obs_deviations)
+
+ #estimated J calculated by updating the original dJ
+ estimated_J = estimated_J -np.dot(dJdx_inv,np.dot(sqK,bg_obs_deviations))
+
+ #Using the cheaper variance calculation instead of going to A-B
+ new_J = np.dot(dJdx_inv.T,dx_prime)
+ vr_total = np.var(new_J,ddof=1)-np.var(np.dot(dJdx_inv.T,dx),ddof=1)
+ #print('all at once:',vr_individual)
+ dx = dx_prime
+
+ #if explicit_sens_flag ==0:
+ # HAHt = np.dot(H_obs,np.dot(C*A,H_obs.T))
+ # HAHtRinv= np.linalg.inv(HAHt+R_obs)
+
+ # dJHdxt = np.dot(dJ,np.dot(H_obs,dx).T)/(nens-1)
+ # vr_individual = -np.dot(dJHdxt,np.dot(HAHtRinv,dJHdxt))
+ # vr_total = vr_individual
+ #
+ #
+ #if iterative_flag:
+ # vr_individual = np.zeros(nobs)
+ # for o in range(nobs): #loop over each observation individually
+
+ # #New A after dx was updated
+ # A = np.cov(dx,ddof=1)
+
+ # #Selecting the single R value for this observation
+ # R_obs = R[o,o]*np.identity(1) # Observation error corvariance matrix
+ # H_obs = H[o,:] #Not sure about this :(
+ # #if
+ # ##H_rms = np.identity(nobs)[o,:]
+ # #H_rms = np.identity(reduced_model_size)[reduced_obs_loc[o],:]
+ # ##H = np.identity(m_const["nx"])[da_const["obs_loc"][o],:]
+#
+ # ##HAHt = np.dot(H,np.dot(A,H.T))
+ # ##HAHtRinv= np.linalg.inv(HAHt+R)
+ # ##dJHdxt = np.dot(dJ,np.dot(H,dx).T)/(nens-1)
+ #
+ # #Now we get the change in dx, which uses the localized A matrix and a square root
+ # E = np.matmul(C*A,H_obs.T)
+ # E = np.matmul(H_obs,E)
+ # E = E + R_obs
+ # alpha = 1./(1.+np.sqrt(R_obs/E))
+
+
+ # HAHt = np.dot(H_obs,np.dot(C*A,H_obs.T))
+ # HAHtRinv= np.linalg.inv(HAHt+R_obs)
+ # #print((C*A).shape,H_obs.T.shape,np.dot(C*A,H_obs.T).shape,HAHtRinv.shape)
+ ###
+
+ # if explicit_sens_flag ==1:
+ #
+ # K = np.dot(C*A,H_obs.T)*HAHtRinv
+ # # Update state vector
+ # Hdx = np.matmul(H_obs,dx)
+ # dx_prime = dx - alpha*np.outer(K,Hdx)
+ #
+ # #Update expected J for each ensemble member
+ # estimated_J = estimated_J -np.dot(dJdx_inv,alpha*np.outer(K,Hdx))
+
+ # #A_new = np.dot(dx,dx.T)/(dx.shape[1]-1)
+ # #Now we get the variance estimate using the difference between new and old B
+ # #dSigma = np.matmul(A_new-A,dJdx_inv.T)
+ # #print('wtf: np.sum(np.abs(A_new-A_old)) ',o,np.sum(np.abs(A_new-A)))
+ # #dSigma = np.matmul(dJdx_inv,dSigma)
+ # #vr_individual[o] = dSigma#np.matmul(np.atleast_2d(np.diag(A_new-A)),(dJdx_inv.T**2))
+ #
+ # vr_individual[o] = np.var(np.dot(dJdx_inv.T,dx_prime),ddof=1)-np.var(np.dot(dJdx_inv.T,dx),ddof=1)
+ # dx = dx_prime
+ # if explicit_sens_flag==0:
+ # dJHdxt = np.dot(dJ,np.dot(H_obs,dx).T)/(nens-1)
+ # vr_individual[o] = -np.dot(dJHdxt,np.dot(HAHtRinv,dJHdxt))
+ #
+ # #Still needs localization!
+ #
+ # #Now we include dJ into dx to include it in the update
+ # dxJ = np.vstack([dx,dJ])
+ # AxJ = np.dot(dxJ,dxJ.T)/(dxJ.shape[1]-1)
+ # HxJ = np.hstack([H_obs,np.array(0)])
+ # CxJ = np.ones([nstate+1,nstate+1])
+ # CxJ[:nstate,:nstate] = C
+ # HAHt = np.dot(HxJ,np.dot(CxJ*AxJ,HxJ.T))
+ # HAHtRinv= np.linalg.inv(HAHt+R_obs)
+
+ # #print(AxJ.shape,CxJ.shape)
+ # #print(AxJ.shape,HxJ.T.shape,np.dot(AxJ,HxJ.T).shape,HAHtRinv.shape)
+ # K = np.dot(CxJ*AxJ,HxJ.T)*HAHtRinv
+
+
+ # # Update state vector
+ # HdxJ = np.matmul(HxJ,dxJ)
+ # dxJ = dxJ - alpha*np.outer(K,HdxJ)
+
+ # dx = dxJ[:-1,:]
+ # #old_var_dJ = np.var(dJ)
+ # dJ = dxJ[-1,:]
+ # estimated_J = dJ+np.mean(bf_response)
+ # #new_var_dJ = np.var(dJ)
+ #
+ # #Final var reduciton estimate
+ # if explicit_sens_flag==0: vr_total=np.sum(vr_individual)
+ #
+ # if explicit_sens_flag==1:
+ # vr_total=np.var(np.dot(dJdx_inv.T,dx),ddof=1)-np.var(np.dot(dJdx_inv.T,dx_orig),ddof=1)
+ #
+ # #Checking different formulation
+
+
+ #
+
+ J_dict['es'] = estimated_J
+
+ J_fc= fc_response
+ dJ_fc = J_fc-np.mean(J_fc)
+ real_reduction=np.var(dJ_fc,ddof=1) - np.var(dJ_orig,ddof=1)
+
+ return vr_total,vr_individual,real_reduction,J_dict,dJdx_inv,quad_state,dx
+
+def vr_individual_loc_22(background,truth,m_const,da_const,sat_operator,response_func=1,abs_flag=0,
+ quad_state = None,exp_num=0,
+ advect_flag = 0,obs_seed=0,model_seed=0,
+ random_samples=0,seed_samples=0,error_u=100,error_u_ens=100):
+
+ """
+ Based on vr_individual_loc, but adapted to take sat measurements as well.
+ The main purpose is to do variance reduction estimates by individually localizing the covariances of the response functions at each point.
+
+ important parameters:
+ response_func: controls the response, 1 results in the default sum over the mean, 2 in the mean over the mean
+ abs_flag: determines if the absolute value of the state vector is used for the response function.
+
+ If advect_flag==1, the localization matrix is advected using the mo_const advection speed which should be close to the ensemble mean. The same advection function used as in the linear advection model.
+ If advect_flag==2, the localization matrix is advected using the individual ensemble members, so that in the end the localization vector is broadened accoring to ensemble disperstion.
+
+ random_samples=0,seed_samples=0,error_u=100,error_u_ens=100: these all are various ways of adding an error to the advection term.
+
+
+ For now only using the all at once method, and it probably won't work for only a single observation.
+
+ The equations should be in the paper draft, but I made quite the mistake by thinking that the A at the end of A'-A = -KHA is also localized.
+
+ Possible speed ups not used so far:
+ * I could speed this up a bit by only calculating the localized correlations where they will be needed (so where H says they should be), but that is a bit of work to do well so for now I will leave it out.
+ * Given that we use Gaspari-Cohn localization most of the time, I might be able to save some time by only calculating covariances where C_i is not zero. But probably won't make a difference for toy model
+ * We currently assign j_i a value at every grid point, even though we skip calculations where response_s = 0. It would be more memory efficient to allocate J_i in a restricted area, but again I think this won't matter in a toymodel
+ * Possibly the biggest speed up would be to avoid the loop when calculating the response functions, which I haven't bothered with yet.
+
+
+
+ """
+ ##########################################################################
+ #We begin by computing the vecors which determine the response function
+ ###########################################################################
+ response_s = np.zeros(m_const['nx'])
+ response_c = np.ones(m_const['nx'])
+ nx = m_const['nx']
+ if response_func ==1:
+ idx_str = int(nx/3.)
+ idx_end = int(2*nx/3.)
+ response_s[idx_str:idx_end] = 1
+ if response_func ==2:
+ idx_str = int(nx/3.)
+ idx_end = int(2*nx/3.)
+ response_s[idx_str:idx_end] = 1./np.float(len(response_s[idx_str:idx_end]))
+
+
+
+
+ ###########################################################################
+ #First, need we calculate the quad (analysis, forecast, blind forecast)
+ #A precomputed quad_state can be supplied to quickly use different response functions
+ ###########################################################################
+ if type(quad_state)== type(None):
+ quad_state = single_step_analysis_forecast_22(background,truth,da_const,m_const,sat_operator,obs_seed=obs_seed,model_seed=model_seed)
+
+ ###########################################################################
+ # Next we need the response functions.
+ # We only really need the forecast and blind forecast, but for fun I'll calculate all for now
+ ###########################################################################
+
+ #For now I am not worried about being efficient
+ nobs = da_const["n_obs_h"] + da_const["n_obs_sat"]
+ obs = np.arange(nobs)
+ nens = da_const["nens"]
+ nstate = len(truth)
+
+ bf_response = np.zeros(nens)
+ fc_response = np.zeros(nens)
+ an_response = np.zeros(nens)
+ bg_response = np.zeros(nens)
+ for n in range(da_const["nens"]):
+ if abs_flag==0:
+ bf_response[n] = np.sum(response_s * np.power(quad_state["bf"][:,n],response_c))
+ fc_response[n] = np.sum(response_s * np.power(quad_state["fc"][:,n],response_c))
+ an_response[n] = np.sum(response_s * np.power(quad_state["an"][:,n],response_c))
+ bg_response[n] = np.sum(response_s * np.power(quad_state["bg"][:,n],response_c))
+ if abs_flag==1:
+ bf_response[n] = np.sum(response_s * np.power(np.abs(quad_state["bf"][:,n]),response_c))
+ fc_response[n] = np.sum(response_s * np.power(np.abs(quad_state["fc"][:,n]),response_c))
+ an_response[n] = np.sum(response_s * np.power(np.abs(quad_state["an"][:,n]),response_c))
+ bg_response[n] = np.sum(response_s * np.power(np.abs(quad_state["bg"][:,n]),response_c))
+
+ J_dict = {}
+ J_dict['bf'] = bf_response
+ J_dict['bg'] = bg_response
+ J_dict['an'] = an_response
+ J_dict['fc'] = fc_response
+ J_dict['tr_bg'] = np.sum(response_s * np.power(quad_state["tr_bg"][:],response_c))
+ if abs_flag==1:
+ J_dict['tr_bg'] = np.sum(response_s * np.power(np.abs(quad_state["tr_bg"][:]),response_c))
+
+ ###########################################################################
+ # Creating the localization matrix which we will need for the VR estimate
+ ###########################################################################
+
+
+ x = quad_state['bg'][:,:]
+
+ dx = x.T-np.mean(x,axis=1)
+ dx = dx.T
+ dx_obs,x_obs_ol = state_to_observation_space(x,m_const,da_const,sat_operator)
+
+ A = np.dot(dx,dx.T)/(dx.shape[1]-1)
+
+ #Now we need the vector of individual js for all ensembles
+ J = np.zeros_like(dx)
+ for n in range(da_const["nens"]):
+ if abs_flag==0:
+ J[:,n] = response_s * np.power(quad_state["bf"][:,n],response_c)
+ if abs_flag==1:
+ J[:,n] = response_s * np.power(np.abs(quad_state["bf"][:,n]),response_c)
+ #J[:,n] = response_s * np.power(quad_state["bg"][:,n],response_c)
+ dJ = J.T-np.mean(J,axis=1)
+ dJ = dJ.T
+
+ dJ_sum = bf_response - np.mean(bf_response)
+
+ L_obs, L_obs_state = localization_matrices_observation_space(m_const,da_const)
+
+ # calculating HAHt with the background deviations in obs space
+ HAHt = L_obs*np.dot(dx_obs,dx_obs.T)/(nens-1)
+ HAHtRinv= np.linalg.inv(HAHt+da_const['R'])
+ # Again replacing Hdx with the background deviations in obs space
+ dJHdxt = np.dot(dJ_sum,dx_obs.T)/(nens-1)
+
+ #missing is all the advection shit, for now we use only a prescribed error
+ if advect_flag ==1:
+ u_advect =m_const['u_ref']*error_u/100.
+ for o in range(nobs):
+ L_obs_state[:,o] = semi_lagrangian_advection(L_obs_state[:,o],m_const['dx'],u_advect,da_const['dt'])
+
+
+
+ dJdYT=np.dot(dJ,dx_obs.T)/(nens-1)
+ LdJdYT=np.sum(L_obs_state*dJdYT,axis=0)
+
+ vr_total = - np.dot(LdJdYT,np.dot(HAHtRinv,dJHdxt).T)
+ J_fc= fc_response
+ dJ_fc = J_fc-np.mean(J_fc)
+ real_reduction=np.var(dJ_fc,ddof=1) - np.var(bf_response,ddof=1)
+
+ return vr_total,real_reduction,quad_state,J_dict#,C
+ #return dJ,dx_obs,dx
+
+ # # calculating HAHt with the background deviations in obs space
+ # HAHt = L_obs*np.dot(dx_obs,dx_obs.T)/(nens-1)
+ # HAHtRinv= np.linalg.inv(HAHt+da_const['R'])
+ #
+ # CdJdX = np.zeros(m_const['nx'])
+ #
+ # #Since we no longer need C for the HCAHT which was calculated above, we can just advect the C matrix for the correlations and not make a new one.
+ # if advect_flag > 0:
+ # if advect_flag ==1:
+ # u_advect =m_const['u_ref']*error_u/100.
+ # C[:,0] = semi_lagrangian_advection(C[:,0],m_const['dx'],u_advect,da_const['dt'])
+ # #for c in range(C.shape[0]):
+ # # C[c,:] = semi_lagrangian_advection(C[c,:],m_const['dx'],-m_const['u_ref'],da_const['dt'])
+ #
+ # if advect_flag ==2:
+ # C_adv = C[:,0]*0.
+ # if random_samples == 0:
+ # for i in range(da_const["nens"]):
+ # if da_const["fixed_seed"]==True: np.random.seed(i)
+ # u_ens = np.random.normal(m_const["u_ref"],da_const["u_std_ens"])
+ # u_ens = m_const["u_ref"]*error_u/100.+(u_ens-m_const["u_ref"])*error_u_ens/100.
+ # #print('u_ens',u_ens)
+ # C_adv = C_adv + semi_lagrangian_advection(C[:,0],m_const['dx'],u_ens,da_const['dt'])
+ # #C_adv[:,0] = C_adv[:,0] + semi_lagrangian_advection(C[:,0],m_const['dx'],10.,da_const['dt'])
+ # C[:,0] = C_adv/np.float(da_const["nens"])
+ # else:
+ # np.random.seed(random_samples+seed_samples)
+ # randomized_ens = np.random.choice(np.arange(0,da_const["nens"]), random_samples, replace=False)
+ # for r in randomized_ens:
+ # #print(r)
+ # if da_const["fixed_seed"]==True: np.random.seed(r)
+ # u_ens = np.random.normal(m_const["u_ref"],da_const["u_std_ens"])
+ # u_ens = m_const["u_ref"]*error_u/100.+(u_ens-m_const["u_ref"])*error_u_ens/100.
+ # #print('u_ens',u_ens)
+ # C_adv = C_adv + semi_lagrangian_advection(C[:,0],m_const['dx'],u_ens,da_const['dt'])
+ # #C_adv[:,0] = C_adv[:,0] + semi_lagrangian_advection(C[:,0],m_const['dx'],10.,da_const['dt'])
+ # C[:,0] = C_adv/np.float(random_samples)
+ #
+ # for i in range(1,m_const['nx']):
+ # C[:,i] = np.hstack([C[-1,i-1],C[:-1,i-1]])
+ #
+ #
+ # #for n in range(m_const["nx"]): #Not sure if the if loop is much quicker than just calculating the dot product everywhere ;)
+ # # if response_s[n] != 0: CdJdX += C[n,:]*np.dot(dJ[n,:],dx.T)/(da_const['nens']-1)
+ # # Somewhat quicker way that avoids the loop
+ # CdJdX=np.sum(C*np.dot(dJ,dx.T)/(da_const['nens']-1.),axis=0)
+ #
+ # #This is where shit gets real!
+ # CdJdXH = np.dot(CdJdX,H.T)
+ # vr_total = - np.dot(CdJdXH,np.dot(HAHtRinv,dJHdxt).T)
+ # J_fc= fc_response
+ # dJ_fc = J_fc-np.mean(J_fc)
+ # real_reduction=np.var(dJ_fc,ddof=1) - np.var(bf_response,ddof=1)
+
+ #
+ #
+
+########################################################################################################################
+# 2023
+########################################################################################################################
+def vr_reloaded_22_locsens(background,truth,m_const,da_const,sat_operator,
+ func_J=sum_mid_tri,
+ sens_loc_flag=0,sens_loc_length = 2000,sens_loc_adv_error=100,
+ reduc = 1,reg_flag=1,
+ quad_state = None,dJdx_inv=None,alpha=0.01,mismatch_threshold=0.1,
+ iterative_flag=0,explicit_sens_flag = 1,exp_num=0,obs_seed=0,model_seed=0):
+
+ """
+ Version of vr_reloaded_22 that includes the possibility to apply localization to the the sensitivity.
+ Takes only the background and truth, then caculates the quad, finaly calculates the response function and the variance reduction.
+
+ The quad can be supplied to save time, which makes sense when you comparing different methods for the same experiment calculating.
+
+ Should also return the dJ values of all 4 ensembles (analysis, background, forecast, blind forecast)
+
+
+ -reg_flag added to use L2 regularization
+
+ not implemented:
+ - state model reduction
+
+
+ Speed up options:
+ - If the same forecast can be recycled, the quad can be calculated once and than passed on.
+ - If the sensitivity can be recycled, that can be calculated once and then reused.
+ - Don't use LETKF
+
+ First getting the explicit all at once verion working, will then add the various approaches
+
+ to be implemented:
+ -iteratively and all at once
+ -explicit vs implicit
+ -reduced model spacing by subsampling the grid. Reduc is the spacing, eg reduc 3 means every third grid point is used.
+ should find the nearest reduced model grid point for each observation.
+
+
+
+ """
+ if iterative_flag ==0: from scipy.linalg import sqrtm
+ ###########################################################################
+ #First, need we calculate the quad (analysis, forecast, blind forecast)
+ #A precomputed quad_state can be supplied to quickly use different response functions
+ ###########################################################################
+ if type(quad_state)== type(None):
+ quad_state = single_step_analysis_forecast_22(background,truth,da_const,m_const,sat_operator,obs_seed=obs_seed,model_seed=model_seed)
+
+ ###########################################################################
+ # Next we need the response functions.
+ # We only really need the forecast and blind forecast, but for fun I'll calculate all for now
+ ###########################################################################
+
+ #For now I am not worried about being efficient
+ nobs = len(da_const["obs_loc"])
+ obs = np.arange(nobs)
+ nens = da_const["nens"]
+ nstate = len(truth)
+
+ bf_response = np.zeros(nens)
+ fc_response = np.zeros(nens)
+ an_response = np.zeros(nens)
+ bg_response = np.zeros(nens)
+ for n in range(da_const["nens"]):
+ bf_response[n] = func_J(quad_state["bf"][:,n])
+ fc_response[n] = func_J(quad_state["fc"][:,n])
+ an_response[n] = func_J(quad_state["an"][:,n])
+ bg_response[n] = func_J(quad_state["bg"][:,n])
+
+ J_dict = {}
+ J_dict['bf'] = bf_response
+ J_dict['bg'] = bg_response
+ J_dict['an'] = an_response
+ J_dict['fc'] = fc_response
+ J_dict['tr_bg'] = func_J(quad_state['tr_bg'])
+
+
+
+ x = quad_state['bg'][:,:]
+
+ dx = x.T-np.mean(x,axis=1)
+ dx = dx.T
+ dx_orig = dx+0
+
+ A = np.dot(dx,dx.T)/(dx.shape[1]-1)
+
+ J = bf_response
+ dJ = J-np.mean(J)
+ dJ_orig = J-np.mean(J)
+
+ if sens_loc_flag==1:
+ X_J =quad_state['bf'][:,:]
+ dX_J = X_J.T - np.mean(X_J,axis=1)
+ dX_J = dX_J.T
+
+ dji = dX_J*1
+ dji[0:100,:] = 0.
+ dji[200:300,:] = 0.
+
+
+ ###############################################################################################
+ # Sensitivity
+ ###############################################################################################
+ #Covarianz between reponse function dJ and state ensemble dx
+ cov_dJdx_vec = np.dot(dJ,dx.T)/(dx.shape[1]-1)
+
+ if explicit_sens_flag==1:
+ #If a sensitivity is provided it is used instead of calculating it
+ if type(dJdx_inv) == np.ndarray:
+ #testing supplied sensitivity
+ rel_error_sens = np.sum(np.abs(A.dot(dJdx_inv)-cov_dJdx_vec))/np.sum(np.abs(cov_dJdx_vec))
+ if rel_error_sens>0.05:
+ print('using supplied sensitivity has a relative error of:',rel_error_sens)
+ #Computing the sensitivity, highly recommend using the regularized version before doing so.
+ else:
+ if reg_flag == 1:
+ dJdx_inv = L2_regularized_inversion(A,cov_dJdx_vec,alpha=alpha,mismatch_threshold=mismatch_threshold)
+ else:
+ A_inv = np.linalg.pinv(A)
+ dJdx_inv = np.dot(A_inv,cov_dJdx_vec)
+ if sens_loc_flag==1:
+ da_const_wide = da_const.copy()
+ da_const_wide['loc_length'] = sens_loc_length
+ C_sens = loc_matrix(da_const_wide,m_const)
+ C_adv = C_sens*1.
+ for nn in range(m_const['nx']):
+ C_adv[:,nn] =semi_lagrangian_advection(C_sens[:,nn],m_const['dx'],+m_const['u_ref']*sens_loc_adv_error/100. ,da_const['dt'])
+
+ sum_loc_cov_adv_djidX=np.sum(C_adv*np.dot(dji,dx.T),axis=0)/(da_const['nens']-1)
+ dJdx_inv = L2_regularized_inversion(C_sens*A,sum_loc_cov_adv_djidX,alpha=alpha)
+
+ estimated_J = bf_response + 0.
+
+
+ #if iterative_flag ==0:
+
+ vr_individual = 0.
+
+
+ if explicit_sens_flag ==1:
+ #Tanjas approach of calculating the square root K following Kalman gain, formula (10) Whitaker and Hamil 2002
+ sqK = square_root_Kalman_gain_observation_deviations(x,m_const,da_const,sat_operator)
+ bg_obs_deviations,bg_obs_ol = state_to_observation_space(x,m_const,da_const,sat_operator)
+ dx_prime = dx - np.dot(sqK,bg_obs_deviations)
+
+ #estimated J calculated by updating the original dJ
+ estimated_J = estimated_J -np.dot(dJdx_inv,np.dot(sqK,bg_obs_deviations))
+
+ #Using the cheaper variance calculation instead of going to A-B
+ new_J = np.dot(dJdx_inv.T,dx_prime)
+ vr_total = np.var(new_J,ddof=1)-np.var(np.dot(dJdx_inv.T,dx),ddof=1)
+ #print('all at once:',vr_individual)
+ dx = dx_prime
+
+
+ J_dict['es'] = estimated_J
+
+ J_fc= fc_response
+ dJ_fc = J_fc-np.mean(J_fc)
+ real_reduction=np.var(dJ_fc,ddof=1) - np.var(dJ_orig,ddof=1)
+
+ return vr_total,vr_individual,real_reduction,J_dict,dJdx_inv,quad_state,dx
\ No newline at end of file
diff --git a/LinAdvFunc/efsoi_functions.py b/LinAdvFunc/efsoi_functions.py
new file mode 100644
index 0000000000000000000000000000000000000000..3fedbf4e588211a79f51e948ce999dc680832038
--- /dev/null
+++ b/LinAdvFunc/efsoi_functions.py
@@ -0,0 +1,575 @@
+#!/usr/bin/env python
+
+#Contains all functions generated for the EFSOI project.
+
+from LinAdvFunc.model_functions import *
+from LinAdvFunc.misc_functions import *
+from LinAdvFunc.da_functions import *
+import numpy as np
+import seaborn as sns
+import matplotlib.pyplot as plt
+
+
+def set_obs_v_dict(
+ obs_loc = np.arange(10,299,10),
+ obs_loc_sat = np.array([]),
+ True_std_obs_h = 0.01,
+ True_std_obs_sat= 0.01,
+ loc =True,
+ loc_length=2000):
+
+ """
+ This function sets up the dictionary for the observations.
+ obs_loc is the location of the observations
+ obs_loc_sat is the location of the satellite observations
+ True_std_obs_h is the standard deviation of the observations
+ True_std_obs_sat is the standard deviation of the satellite observations
+ loc is a boolean that determines whether or not to use localization
+ loc_length is the length of the localization
+ loc length stuff is for localizing the explicit sensitivity
+
+ """
+
+
+ obs_v_dict = {}
+ # obs_loc_sat = np.arange(30,299,40)
+
+ obs_v_dict["obs_loc"] = obs_loc
+ obs_v_dict["obs_loc_sat"] = obs_loc_sat
+ obs_v_dict["True_std_obs"] = True_std_obs_h
+ obs_v_dict["True_std_obs_sat"] = True_std_obs_sat
+ obs_v_dict['n_obs_h'] = len(obs_loc)
+ obs_v_dict['n_obs_sat'] = len(obs_loc_sat)
+ obs_v_dict['loc_length'] = loc_length
+ obs_v_dict['loc'] = loc
+ obs_v_dict['loc_type'] = 'gaspari_cohn'
+ return obs_v_dict
+
+
+def localization_matrices_obs_obs(m_const, da_const, obs_v_dict, sigprop_error=1.0):
+ """
+ Creates the localization matrix needed to localize the YaYf covariances used in the EFSOI.
+
+ _a is for analysis, _v for verification (forecast)
+
+ The distance calculation is still far from elegant.
+
+ Annoyingly, this version doesn't work with the advection routine, because that only works for vectors defined on the regular grid.
+ So in a break with tradition we achieve things by moving the x_grid first
+
+ 2024-02-07: fixed something in the x_grid_neg definition, messed up an indes
+ """
+ n_obs_a = da_const['n_obs_h'] + da_const['n_obs_sat']
+ obs_loc_a = np.hstack([da_const["obs_loc"], da_const["obs_loc_sat"]]).astype(int)
+ n_obs_v = obs_v_dict['n_obs_h'] + obs_v_dict['n_obs_sat']
+ obs_loc_v = np.hstack([obs_v_dict["obs_loc"], obs_v_dict["obs_loc_sat"]]).astype(int)
+ L = np.zeros([n_obs_a, n_obs_v])
+
+ # Now we need the positions of the verification observations
+ x_grid_v_obs = m_const['x_grid'][obs_loc_v]
+ x_grid_neg = -m_const['x_grid'][::-1][obs_loc_v] - m_const['dx']
+ # x_grid_neg = -x_grid_v_obs[::-1] - m_const['dx'] # This is wrong
+ x_grid_ext = x_grid_v_obs + (m_const['dx'] * m_const['nx'])
+
+ # Add on advection here with the option to add a multiplicative factor to the advection as a signal propagation error
+ x_grid_v_obs = x_grid_v_obs - m_const['u_ref'] *sigprop_error* da_const['dt']
+ x_grid_neg = x_grid_neg - m_const['u_ref'] *sigprop_error* da_const['dt']
+ x_grid_ext = x_grid_ext - m_const['u_ref'] *sigprop_error* da_const['dt']
+
+ for o in range(n_obs_a):
+ g = obs_loc_a[o]
+ # First the one using the distance between observations and the state
+ dist_reg = np.abs(x_grid_v_obs - m_const['x_grid'][g])
+ dist_neg = m_const['x_grid'][g] - x_grid_neg
+ dist_ext = x_grid_ext - m_const['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)
+ gc_loc = gaspari_cohn_non_mirrored(dist, da_const['loc_length'])
+ L[o, :] = gc_loc
+
+
+ return L
+
+
+def quad_plotter_paper_efsoi(quad_state,m_const,da_const,obs_v_dict):
+ """
+ Plots the initial background and blind forecast, as well as the analysis and forecast.
+ Both the assimilated and verification observations are plotted with a standard deviation error bar.
+
+
+ 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=(10,6),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.,np.mean(quad_state['bg'][:,:],axis=1),'w',alpha =1,zorder=1,lw=8,label='ens mean')
+ # ax[0,1].plot(m_const['x_grid']/1000.,np.mean(quad_state['an'][:,:],axis=1),'w',alpha =1,zorder=1,lw=8,label='ens mean')
+ # ax[1,0].plot(m_const['x_grid']/1000.,np.mean(quad_state['bf'][:,:],axis=1),'w',alpha =1,zorder=1,lw=8,label='ens mean')
+ # ax[1,1].plot(m_const['x_grid']/1000.,np.mean(quad_state['fc'][:,:],axis=1),'w',alpha =1,zorder=1,lw=8,label='ens mean')
+ # ax[0,0].plot(m_const['x_grid']/1000.,np.mean(quad_state['bg'][:,:],axis=1),'lightgrey',alpha =1,zorder=2,lw=6,label='ens mean')
+ # ax[0,1].plot(m_const['x_grid']/1000.,np.mean(quad_state['an'][:,:],axis=1),'lightgrey',alpha =1,zorder=2,lw=6,label='ens mean')
+ # ax[1,0].plot(m_const['x_grid']/1000.,np.mean(quad_state['bf'][:,:],axis=1),'lightgrey',alpha =1,zorder=2,lw=6,label='ens mean')
+ # ax[1,1].plot(m_const['x_grid']/1000.,np.mean(quad_state['fc'][:,:],axis=1),'lightgrey',alpha =1,zorder=2,lw=6,label='ens mean')
+ # ax[0,0].plot(m_const['x_grid']/1000.,np.mean(quad_state['bg'][:,:],axis=1),'r' ,alpha =1,zorder=2,lw=3,ls='--',label='ens mean')
+ # ax[0,1].plot(m_const['x_grid']/1000.,np.mean(quad_state['an'][:,:],axis=1),'magenta',alpha =1,zorder=2,lw=3,ls='--',label='ens mean')
+ # ax[1,0].plot(m_const['x_grid']/1000.,np.mean(quad_state['bf'][:,:],axis=1),'b' ,alpha =1,zorder=2,lw=3,ls='--',label='ens mean')
+ # ax[1,1].plot(m_const['x_grid']/1000.,np.mean(quad_state['fc'][:,:],axis=1),'c' ,alpha =1,zorder=2,lw=3,ls='--',label='ens mean')
+ ax[0,0].plot(m_const['x_grid']/1000.,np.mean(quad_state['bg'][:,:],axis=1),'k',alpha =1,zorder=2,lw=2,ls='-',label='ens mean')
+ ax[0,1].plot(m_const['x_grid']/1000.,np.mean(quad_state['an'][:,:],axis=1),'k',alpha =1,zorder=2,lw=2,ls='-',label='ens mean')
+ ax[1,0].plot(m_const['x_grid']/1000.,np.mean(quad_state['bf'][:,:],axis=1),'k',alpha =1,zorder=2,lw=2,ls='-',label='ens mean')
+ ax[1,1].plot(m_const['x_grid']/1000.,np.mean(quad_state['fc'][:,:],axis=1),'k',alpha =1,zorder=2,lw=2,ls='-',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='verification 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='verification 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].set_title('background')
+ ax[1,0].set_title('free-forecast')
+ ax[0,1].set_title('analysis')
+ ax[1,1].set_title('forecast')
+ ax[1,0].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$')
+
+
+ if obs_v_dict['n_obs_h']:
+ ax[1,0].scatter(m_const["x_grid"][obs_v_dict["obs_loc"]] /1000.,quad_state['obs_h_v'] ,zorder=3,s=10,label='verification obs',color='k')
+ ax[1,1].scatter(m_const["x_grid"][obs_v_dict["obs_loc"]] /1000.,quad_state['obs_h_v'] ,zorder=3,s=10,label='verification obs',color='k')
+ ax[1,1].errorbar(m_const["x_grid"][obs_v_dict["obs_loc"]]/1000.,quad_state['obs_h_v'] ,yerr=obs_v_dict['True_std_obs'] ,zorder=3,color='k',capsize=3,capthick=2,ls='none')
+ ax[1,0].errorbar(m_const["x_grid"][obs_v_dict["obs_loc"]]/1000.,quad_state['obs_h_v'] ,yerr=obs_v_dict['True_std_obs'] ,zorder=3,color='k',capsize=3,capthick=2,ls='none')
+
+
+ return fig,ax
+
+def single_step_analysis_forecast_23(background,truth,da_const,m_const,sat_operator,model_seed=0,obs_seed=0):
+ """
+ Revisited version of single step analysis forecast to deal with sat data as well.
+
+ If the LETKF is used, the ensemble weighting matrix W_a is also returned.
+ """
+ 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,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)
+
+
+
+
+
+ """
+ Predict blind forecast and forecast
+ """
+ bf = np.zeros_like(background)
+ fc = np.zeros_like(background)
+
+ for i in range(da_const["nens"]):
+ np.random.seed(i+model_seed)
+ u_ens = np.random.normal(m_const["u_ref"],da_const["u_std_ens"])
+ dhdt_ens = np.random.normal(m_const["dhdt_ref"],da_const["dhdt_std_ens"])
+ bf[:,i] = linear_advection_model(background[:,i],u_ens,dhdt_ens,m_const["dx"],da_const["dt"],da_const["nt"])
+ fc[:,i] = linear_advection_model(an[:,i],u_ens,dhdt_ens,m_const["dx"],da_const["dt"],da_const["nt"])
+
+
+ quad_state['bg'] = background
+ quad_state['an'] = an
+ quad_state['bf'] = bf
+ quad_state['fc'] = fc
+ quad_state['tr_bg'] = truth
+ quad_state['obs'] = obs
+ quad_state['obs_sat'] = obs_sat
+ return quad_state
+
+
+def quad_plotter_sat_a_efsoi(quad_state,m_const,da_const,obs_v_dict,sat_operator):
+ """
+ Not sure how to do this, but will start off trying to plot the assimilated non-linear observations in reflectance space, and the linear verification observations.
+ Both the assimilated and verification observations are plotted with a standard deviation error bar.
+
+
+ Returns:
+ figure and axes
+ """
+
+
+ sns.set()
+ sns.set_style('whitegrid')
+
+
+ alpha = np.sqrt(1/da_const['nens'])+0.1
+
+ # for plotting reflectance
+ sat_an = sat_operator(quad_state['an'])
+ sat_bg = sat_operator(quad_state['bg'])
+ n_obs_h =len(da_const["obs_loc"])
+ n_obs_sat =len(da_const["obs_loc_sat"])
+
+ if n_obs_sat==0: print('No satellite observations, wrong function buddy')
+
+ dY_b, y_ol_b = state_to_observation_space(quad_state['bg'],m_const,da_const,sat_operator)
+ y_b = dY_b.T + y_ol_b
+ y_b = y_b.T
+
+
+ # for plotting the pixels
+ window = m_const['x_grid'][da_const["obs_loc_sat"][1]]-m_const['x_grid'][da_const["obs_loc_sat"][0]]
+ xpixmin = m_const['x_grid'][da_const["obs_loc_sat"]]-window/2
+ xpixmax = m_const['x_grid'][da_const["obs_loc_sat"]]+window/2
+
+
+
+ fig, ax = plt.subplots(2,2,figsize=(10,6),sharex='all',sharey='row')
+
+
+
+ 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,1].hlines(
+ sat_an [da_const["obs_loc_sat"],i],
+ xpixmin/1000,xpixmax/1000,
+ color='magenta',alpha=0.2,lw=3)
+ ax[0,0].hlines(
+ sat_bg [da_const["obs_loc_sat"],i],
+ xpixmin/1000,xpixmax/1000,
+ color='r',alpha=0.2,lw=3)
+ ax[0,0].hlines(sat_bg [da_const["obs_loc_sat"],:].mean(axis=1),xpixmin/1000,xpixmax/1000,color='white',alpha=1,lw=4,zorder=1)
+ ax[0,0].hlines(sat_bg [da_const["obs_loc_sat"],:].mean(axis=1),xpixmin/1000,xpixmax/1000,color='black',alpha=1,lw=2,zorder=2)
+ ax[0,1].hlines(sat_an [da_const["obs_loc_sat"],:].mean(axis=1),xpixmin/1000,xpixmax/1000,color='white',alpha=1,lw=4,zorder=1)
+ ax[0,1].hlines(sat_an [da_const["obs_loc_sat"],:].mean(axis=1),xpixmin/1000,xpixmax/1000,color='black',alpha=1,lw=2,zorder=2)
+
+
+ ax[1,0].plot(m_const['x_grid']/1000.,np.mean(quad_state['bf'][:,:],axis=1),'k' ,alpha =1,zorder=2,lw=2,ls='-',label='ens mean')
+ ax[1,1].plot(m_const['x_grid']/1000.,np.mean(quad_state['fc'][:,:],axis=1),'k' ,alpha =1,zorder=2,lw=2,ls='-',label='ens mean')
+
+ if da_const['n_obs_sat']:
+ # ax[0,0].scatter( m_const["x_grid"][da_const["obs_loc_sat"]]/1000.,quad_state['obs_sat'][da_const["obs_loc_sat"]] ,zorder=3,s=30,label='verification obs',color='k',marker='_')
+ # ax[0,1].scatter( m_const["x_grid"][da_const["obs_loc_sat"]]/1000.,quad_state['obs_sat'][da_const["obs_loc_sat"]] ,zorder=3,s=30,label='verification obs',color='k',marker='_')
+ ax[0,0].scatter( m_const["x_grid"][da_const["obs_loc_sat"]]/1000.,quad_state['obs_sat'][da_const["obs_loc_sat"]] ,zorder=3,s=10,label='verification obs',color='k')
+ ax[0,1].scatter( m_const["x_grid"][da_const["obs_loc_sat"]]/1000.,quad_state['obs_sat'][da_const["obs_loc_sat"]] ,zorder=3,s=10,label='verification obs',color='k')
+ ax[0,1].errorbar(m_const["x_grid"][da_const["obs_loc_sat"]]/1000.,quad_state['obs_sat'][da_const["obs_loc_sat"]] ,yerr=da_const['used_std_obs_sat'] ,zorder=3,color='k',capsize=3,capthick=2,ls='none')
+ ax[0,0].errorbar(m_const["x_grid"][da_const["obs_loc_sat"]]/1000.,quad_state['obs_sat'][da_const["obs_loc_sat"]] ,yerr=da_const['used_std_obs_sat'] ,zorder=3,color='k',capsize=3,capthick=2,ls='none')
+
+
+ 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[1,0].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('reflectance')
+ ax[1,0].set_ylabel(r'$\phi$')
+
+ ax[0,0].set_ylim(0,1)
+ ax[0,0].set_yticks([0,0.3,0.7,1.0])
+ # ax[0,1].yaxis.tick_right()
+ # ax[0,1].tick_params(labelright=True)
+ # ax[0,0].set_yticklabels(['','no cloud','cloudy',''],rotation=70)
+
+ if obs_v_dict['n_obs_h']:
+ ax[1,0].scatter(m_const["x_grid"][obs_v_dict["obs_loc"]] /1000.,quad_state['obs_h_v'] ,zorder=3,s=10,label='verification obs',color='k')
+ ax[1,1].scatter(m_const["x_grid"][obs_v_dict["obs_loc"]] /1000.,quad_state['obs_h_v'] ,zorder=3,s=10,label='verification obs',color='k')
+ ax[1,1].errorbar(m_const["x_grid"][obs_v_dict["obs_loc"]]/1000.,quad_state['obs_h_v'] ,yerr=obs_v_dict['True_std_obs'] ,zorder=3,color='k',capsize=3,capthick=2,ls='none')
+ ax[1,0].errorbar(m_const["x_grid"][obs_v_dict["obs_loc"]]/1000.,quad_state['obs_h_v'] ,yerr=obs_v_dict['True_std_obs'] ,zorder=3,color='k',capsize=3,capthick=2,ls='none')
+
+
+ return fig,ax
+
+
+def quad_plotter_sat_v_efsoi(quad_state,m_const,da_const,obs_v_dict,sat_operator):
+ """
+ For experiment with non-linear verification observations
+ Both the assimilated and verification observations are plotted with a standard deviation error bar.
+
+
+ Returns:
+ figure and axes
+ """
+
+
+ sns.set()
+ sns.set_style('whitegrid')
+
+
+ alpha = np.sqrt(1/da_const['nens'])+0.1
+
+ # for plotting reflectance
+ sat_fc = sat_operator(quad_state['fc'])
+ sat_ff = sat_operator(quad_state['bf'])
+ n_obs_h = len(obs_v_dict["obs_loc"])
+ n_obs_sat =len(obs_v_dict["obs_loc_sat"])
+
+ if n_obs_sat==0: print('No satellite observations, wrong function buddy')
+
+ dY_ff, y_ol_ff = state_to_observation_space(quad_state['bf'],m_const,da_const,sat_operator)
+ y_ff = dY_ff.T + y_ol_ff
+ y_ff = y_ff.T
+
+
+ # for plotting the pixels
+ window = m_const['x_grid'][obs_v_dict["obs_loc_sat"][1]]-m_const['x_grid'][obs_v_dict["obs_loc_sat"][0]]
+ xpixmin = m_const['x_grid'][obs_v_dict["obs_loc_sat"]]-window/2
+ xpixmax = m_const['x_grid'][obs_v_dict["obs_loc_sat"]]+window/2
+
+
+
+ fig, ax = plt.subplots(2,2,figsize=(10,6),sharex='all',sharey='row')
+
+
+
+ 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[1,1].hlines(
+ sat_fc [obs_v_dict["obs_loc_sat"],i],
+ xpixmin/1000,xpixmax/1000,
+ color='cyan',alpha=0.2,lw=3)
+ ax[1,0].hlines(
+ sat_ff [obs_v_dict["obs_loc_sat"],i],
+ xpixmin/1000,xpixmax/1000,
+ color='blue',alpha=0.2,lw=3)
+
+ ax[1,0].hlines(sat_ff [obs_v_dict["obs_loc_sat"],:].mean(axis=1),xpixmin/1000,xpixmax/1000,color='white',alpha=1,lw=4,zorder=1)
+ ax[1,0].hlines(sat_ff [obs_v_dict["obs_loc_sat"],:].mean(axis=1),xpixmin/1000,xpixmax/1000,color='black',alpha=1,lw=2,zorder=2)
+ ax[1,1].hlines(sat_fc [obs_v_dict["obs_loc_sat"],:].mean(axis=1),xpixmin/1000,xpixmax/1000,color='white',alpha=1,lw=4,zorder=1)
+ ax[1,1].hlines(sat_fc [obs_v_dict["obs_loc_sat"],:].mean(axis=1),xpixmin/1000,xpixmax/1000,color='black',alpha=1,lw=2,zorder=2)
+ # ax[1,0].hlines(sat_ff [obs_v_dict["obs_loc_sat"],:].mean(axis=1),xpixmin/1000,xpixmax/1000,color='blue',alpha=1,lw=2,zorder=3)
+
+ ax[0,0].plot(m_const['x_grid']/1000.,np.mean(quad_state['bg'][:,:],axis=1),'k',alpha =1,zorder=2,lw=2,ls='-',label='ens mean')
+ ax[0,1].plot(m_const['x_grid']/1000.,np.mean(quad_state['an'][:,:],axis=1),'k',alpha =1,zorder=2,lw=2,ls='-',label='ens mean')
+
+ if obs_v_dict['n_obs_sat']:
+ ax[1,0].scatter( m_const["x_grid"][obs_v_dict["obs_loc_sat"]]/1000.,quad_state['obs_sat_v'] ,zorder=3,s=10,label='verification obs',color='k')
+ ax[1,1].scatter( m_const["x_grid"][obs_v_dict["obs_loc_sat"]]/1000.,quad_state['obs_sat_v'] ,zorder=3,s=10,label='verification obs',color='k')
+ ax[1,1].errorbar(m_const["x_grid"][obs_v_dict["obs_loc_sat"]]/1000.,quad_state['obs_sat_v'] ,yerr=obs_v_dict['True_std_obs_sat'] ,zorder=3,color='k',capsize=3,capthick=2,ls='none')
+ ax[1,0].errorbar(m_const["x_grid"][obs_v_dict["obs_loc_sat"]]/1000.,quad_state['obs_sat_v'] ,yerr=obs_v_dict['True_std_obs_sat'] ,zorder=3,color='k',capsize=3,capthick=2,ls='none')
+
+ 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='verification 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='verification 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].set_title('background')
+ ax[1,0].set_title('free-forecast')
+ ax[0,1].set_title('analysis')
+ ax[1,1].set_title('forecast')
+ ax[1,0].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[1,0].set_ylabel('reflectance')
+ ax[0,0].set_ylabel(r'$\phi$')
+
+ # ax[1,0].set_ylim(0,1)
+ ax[1,0].set_yticks([0.3,0.7])
+
+
+
+ return fig,ax
+
+
+def EFSOI_LETKF(background,truth,m_const,da_const,obs_v_dict,sat_operator,
+ quad_state = None,exp_num=0,obs_seed=0,model_seed=0,obs_v_seed=1000,alpha=0.1,sigprop_error=1.0):
+
+ """
+ EFSOI, specifically for the LETKF. It has 3 different versions of computing Ya
+
+ 1. Ya = H Xa (HXa, not great for non-linear H)
+ 2. Ya = Wa_l Yb_l (WlYbl, notation is a bit strange, but the _l indicates the l gridpoint of where the observation is. )
+ 2. Ya_l = Wa_l Yb (WYb, this Ya needs to calculated separately for each observation point.)
+
+ Shows how saving the ensemble weights leads to the correct answer for dt=0, because the Kalman gain is correctly captured. Is like calculating the Kalman
+ from the background instead of from the analysis.
+
+ If the observation operator is non-linear, option 2 is better than option 1, option 3 is always better than 2, but not feasible for large systems.
+
+ For now I've thrown out all the localization stuff for the explicit sensitivity. Not sure if or how it should be included in the paper
+
+ """
+
+ if obs_v_dict['n_obs_h']*obs_v_dict['n_obs_sat']>0:
+ print('This function was not tested for both h and sat obs, only for one of them. It could work, but it was not tested')
+ return
+ ###########################################################################
+ #First, need we calculate the basic quad (analysis, forecast, blind forecast)
+ #A precomputed quad_state can be supplied to quickly use different response functions
+ ###########################################################################
+ if quad_state==None:
+ quad_state = single_step_analysis_forecast_23(background,truth,da_const,m_const,sat_operator,obs_seed=obs_seed,model_seed=model_seed)
+
+ # if random_truth_u_flag:
+ # np.random.seed(model_seed+da_const['nens'])
+ # u_tr = np.random.normal(m_const["u_ref"],da_const["u_std_ens"])
+ # dhdt_tr = np.random.normal(m_const["dhdt_ref"],da_const["dhdt_std_ens"])
+ # tr_ff = linear_advection_model(truth,u_tr,dhdt_tr,m_const["dx"],da_const["dt"],da_const["nt"])
+ # else:
+ tr_ff = linear_advection_model(truth,m_const['u_ref'],m_const['dhdt_ref'],m_const["dx"],da_const["dt"],da_const["nt"])
+ quad_state['tr_ff'] = tr_ff
+
+ ###########################################################################
+ # generating verification obs, and model equivalents
+ ###########################################################################
+ obs_h_v, obs_sat_v = generate_obs_22_single(quad_state['tr_ff'],m_const,obs_v_dict,sat_operator,obs_v_seed)
+ if obs_v_dict['n_obs_h']*obs_v_dict['n_obs_sat']>0:
+ obs_v = np.hstack([obs_h_v[obs_v_dict['obs_loc']],obs_sat_v[obs_v_dict['obs_loc_sat']]])
+ quad_state['obs_h_v'] = obs_h_v[obs_v_dict['obs_loc']]
+ quad_state['obs_sat_v'] = obs_sat_v[obs_v_dict['obs_loc_sat']]
+ if obs_v_dict['n_obs_h']>1 and obs_v_dict['n_obs_sat']==0:
+ quad_state['obs_h_v'] = obs_h_v[obs_v_dict['obs_loc']]
+ quad_state['obs_sat_v'] = []
+ obs_v = obs_h_v[obs_v_dict['obs_loc']]
+ if obs_v_dict['n_obs_sat']>1 and obs_v_dict['n_obs_h']==0:
+ quad_state['obs_sat_v'] = obs_sat_v[obs_v_dict['obs_loc_sat']]
+ quad_state['obs_h_v'] = []
+ obs_v = obs_sat_v[obs_v_dict['obs_loc_sat']]
+
+
+ #model equivalents of verification observation
+ bf_obs_v_deviations,bf_obs_v_ol = state_to_observation_space(quad_state['bf'],m_const,obs_v_dict,sat_operator)
+ bf_obs_v = (bf_obs_v_deviations.T+bf_obs_v_ol).T
+ fc_obs_v_deviations,fc_obs_v_ol = state_to_observation_space(quad_state['fc'],m_const,obs_v_dict,sat_operator)
+ fc_obs_v = (fc_obs_v_deviations.T+fc_obs_v_ol).T
+ quad_state['y_bf'] = bf_obs_v
+ quad_state['y_bf_ol'] = bf_obs_v_ol
+ quad_state['y_fc'] = fc_obs_v
+ quad_state['y_fc_ol'] = fc_obs_v_ol
+ quad_state['y_obs_v'] = obs_v
+
+
+ #And observation equivalent of the observation that is being denied
+ an_obs_deviations,an_obs_ol = state_to_observation_space(quad_state['an'],m_const,da_const,sat_operator)
+ an_obs = (an_obs_deviations.T+an_obs_ol).T
+ bg_obs_deviations,bg_obs_ol = state_to_observation_space(quad_state['bg'],m_const,da_const,sat_operator)
+
+ #calculating error of mean, and for now we use as error simply the difference to the obs
+ quad_state['e_bf'] = bf_obs_v_ol-obs_v
+ quad_state['e_fc'] = fc_obs_v_ol-obs_v
+ quad_state['e_sq_bf'] = (bf_obs_v_ol-obs_v)**2
+ quad_state['e_sq_fc'] = (fc_obs_v_ol-obs_v)**2
+
+ quad_state['Jall'] = quad_state['e_sq_fc']-quad_state['e_sq_bf']
+
+
+
+ # Ok, this was all just foreplay. Now the real fun starts, estimating Jd' via efsoi, and then looping over the observations
+
+ n_obs_v =obs_v_dict['n_obs_h']+obs_v_dict['n_obs_sat']
+ n_obs_a =da_const['n_obs_h']+da_const['n_obs_sat']
+
+ # Initializing the matrices that will be used to calculate the explicit sensitivity of the observation
+ efsoi_Jd_WYb = np.zeros([n_obs_a,n_obs_v])
+ efsoi_Jd_WlYbl = np.zeros([n_obs_a,n_obs_v])
+ efsoi_Jd_HXa = np.zeros([n_obs_a,n_obs_v])
+
+
+ if da_const['loc']: L= localization_matrices_obs_obs(m_const,da_const,obs_v_dict,sigprop_error=sigprop_error)
+ else: L = np.ones([n_obs_a,n_obs_v]) #Already includes advection
+
+
+ # Just saving the locactions into a single array
+ obs_loc_a = np.hstack([da_const["obs_loc"],da_const["obs_loc_sat"]]).astype(int)
+ obs_loc_v = np.hstack([obs_v_dict["obs_loc"],obs_v_dict["obs_loc_sat"]]).astype(int)
+
+
+ Yf = fc_obs_v_deviations
+
+ #Now the 3 different ways of computing computing Ya
+ W_a = quad_state['W_a']
+
+ #1. Directly from the analysis state stuff for Ya = H Xa
+ HXa = an_obs_deviations
+ #2. From W_a_ll. This could be done more efficiently by including it in the LETKF analysis function, but for now we do it here so we don't have to change the functions to also pass Ya
+ WlYbl = bg_obs_deviations*0.
+ for a in range(n_obs_a):
+ g = obs_loc_a[a]
+ WlYbl[a,:] = np.dot(bg_obs_deviations[a,:],W_a[g,:,:])
+
+
+
+ for o in range(n_obs_v):
+
+
+ #LETKF version, which useses the ya vector, and a localized R^-1
+ # first order calculate localization matrix for R
+ g = obs_loc_v[o]
+
+
+ # We need the Wa where the observation was assimilated
+ # Finding the upstream index of the verification observation
+ g_upstream = (g-m_const['u_ref']*da_const['dt']/m_const['dx']).astype(int)
+ WYb = np.dot(bg_obs_deviations,W_a[g_upstream,:])
+ # print(g,g_upstream)
+
+ # We also need the advected localization matrix
+ gc_loc=L[:,o]
+
+
+ WYbRY_loc = np.dot(np.dot(np.diag(gc_loc)*np.linalg.inv(da_const['R']),WYb),Yf[o,:])
+ WlYblRY_loc = np.dot(np.dot(np.diag(gc_loc)*np.linalg.inv(da_const['R']),WlYbl),Yf[o,:])
+ HXaRY_loc = np.dot(np.dot(np.diag(gc_loc)*np.linalg.inv(da_const['R']),HXa),Yf[o,:])
+
+ init_error =quad_state['e_bf'][o]+quad_state['e_fc'][o]
+ for a in range(n_obs_a):
+ if a < da_const['n_obs_h']:
+ dbg = quad_state['obs'][da_const['obs_loc'][a]]-bg_obs_ol[a]
+ else:
+ dbg = quad_state['obs_sat'][da_const['obs_loc_sat'][a-da_const['n_obs_h']]]-bg_obs_ol[a]
+
+ efsoi_Jd_WYb[a,o] = init_error*WYbRY_loc[a]*dbg/(da_const['nens']-1)
+ efsoi_Jd_WlYbl[a,o] = init_error*WlYblRY_loc[a]*dbg/(da_const['nens']-1)
+ efsoi_Jd_HXa[a,o] = init_error*HXaRY_loc[a]*dbg/(da_const['nens']-1)
+
+
+ quad_state['efsoi_WYb'] = efsoi_Jd_WYb
+ quad_state['efsoi_WlYbl'] = efsoi_Jd_WlYbl
+ quad_state['efsoi_HXa'] = efsoi_Jd_HXa
+
+
+ return quad_state
\ No newline at end of file
diff --git a/LinAdvFunc/model_functions.py b/LinAdvFunc/model_functions.py
new file mode 100644
index 0000000000000000000000000000000000000000..c6b6ccd9220458a77e6abaa379ecdc38e6a4124b
--- /dev/null
+++ b/LinAdvFunc/model_functions.py
@@ -0,0 +1,118 @@
+#!/usr/bin/env python
+#collection of all functions which determine the model setup and time evolution
+
+import numpy as np
+
+def set_model_constants_22(nx=300,dx=100,u=10,u_std=0,dhdt=0,dhdt_std=0,h_init_std = 3000, init_func='gaus',sine_init=[1],truth_seed=0):
+ """
+ Sets all the constants used to describe the model. The 2022 version (_22) is only new in that a truth seed is added so that
+ truth perturbations can be added systematically, but that hasn't been tested yet anyway.
+
+ """
+ const = {}
+
+ #grid
+ const["nx"] = nx
+ const["dx"] = dx
+ const["x_grid"] = np.arange(nx)*dx
+
+ #properties of truth wave
+ const["h_init_std"] = h_init_std # width of the initial gaussian distribution, or number of sine waves
+ const["sine_init"] = sine_init # width of the initial gaussian distribution, or number of sine waves
+ const["init_func"] = init_func # currently a gaussian ('gaus') or a sine wave ('sine') are available
+ const["u_ref"] = u # reference u around which ensemble is generated, truth can
+ const["u_std_truth"] = u_std # standard deviation of u for truth
+ const["dhdt_ref"] = dhdt # reference dhdt around which ensemble is generated
+ const["dhdt_std_truth"] = dhdt_std # standard deviation of dhdt for truth
+ const["truth_seed"] = truth_seed # added to seed used to generate truth deviations
+
+ #model
+ const["model"] = 'LA'
+
+ return const
+
+def set_model_constants(nx=101,dx=100,u=2.33333333,u_std=0,dhdt=0,dhdt_std=0,h_init_std = 500, init_func='gaus',sine_init=[1]):
+ """
+ Sets all the constants used to describe the model.
+ Currently has only gaussian initial conditions.
+
+ ToDo: Add different initial conditions as an option.
+ """
+ const = {}
+
+ #grid
+ const["nx"] = nx
+ const["dx"] = dx
+ const["x_grid"] = np.arange(nx)*dx
+
+ #properties of truth wave
+ const["h_init_std"] = h_init_std # width of the initial gaussian distribution, or number of sine waves
+ const["sine_init"] = sine_init # width of the initial gaussian distribution, or number of sine waves
+ const["init_func"] = init_func # currently a gaussian ('gaus') or a sine wave ('sine') are available
+ const["u_ref"] = u # reference u around which ensemble is generated, truth can
+ const["u_std_truth"] = u_std # standard deviation of u for truth
+ const["dhdt_ref"] = dhdt # reference dhdt around which ensemble is generated
+ const["dhdt_std_truth"] = dhdt_std # standard deviation of dhdt for truth
+
+ #model
+ const["model"] = 'LA'
+
+ return const
+
+def gaussian_initial_condition(x,sig):
+ """
+ Generates a gaussian which is commonly used for the initial conditions.
+ It always uses a mean of zero, and then mirrors values along the x axis to account for periodic boundary conditions.
+ This is not very flexible, and also not super clean because it just assumes that the number of x grid points is even
+ But since it works making it nice is very far down the priority list.
+ """
+ mu = 0
+ y = np.exp(-np.power(x - mu, 2.) / (2 * np.power(sig, 2.)))
+ half_idx = int(x.shape[0]/2)
+ y[-half_idx:] = y[half_idx-1::-1]
+
+ return y
+
+def linear_advection_model(y,u,dydt,dx,dt,nt):
+ """
+ This is a simple wrapper for the advection and growth routines which make up the time evolution of the model.
+ """
+ for n in range(nt):
+ y = semi_lagrangian_advection(y,dx,u,dt)
+ y = linear_growth(y,dydt,dt)
+
+ return y
+
+
+def semi_lagrangian_advection(y,dx,u,dt):
+ """
+ Advects a 1D field for a given time difference and speed.
+ A linear polation is applied between the nearest grid points. Getting the periodic boundary right is a bit of annoying book keeping.
+ Works by calculating the fractional displacement and the integer displacement.
+ Only works with a constant dx.
+ """
+ dx_frac = u*dt/dx-np.floor(u*dt/dx)
+ dx_int = int(np.floor(u*dt/dx))
+
+ ##initial quality check, shouldn't be needed anymore
+ #if np.abs(u*dt - (dx_frac+dx_int)*dx)>0.00001:
+ # print(u,u*dt,(dx_frac+dx_int)*dx)
+
+ if abs(dx_int)>len(y):
+ #this is what to do if things move through the whole domain in one step
+ if dx_int >0 :
+ dx_int = dx_int - int(dx_int/float(len(y)))*len(y)
+ if dx_int<0:
+ dx_int = dx_int + int(-dx_int/float(len(y)))*len(y)
+
+
+ y = np.hstack([(1-dx_frac)*y[0]+(dx_frac)*y[-1],(1-dx_frac)*y[1:]+(dx_frac)*y[:-1]])
+ y = np.hstack([y[-dx_int:],y[:-dx_int]])
+ return y
+
+
+def linear_growth(y,dydt,dt):
+ """
+ Applies linear growth rate. Very highly advanced math.
+ """
+ return y + dydt*dt
diff --git a/LinAdvFunc/plot_functions.py b/LinAdvFunc/plot_functions.py
new file mode 100644
index 0000000000000000000000000000000000000000..0f1ccff19405de3884d5c2a4b74a6876aad04178
--- /dev/null
+++ b/LinAdvFunc/plot_functions.py
@@ -0,0 +1,1125 @@
+#!/usr/bin/env python
+
+# Collection of various plot scripts.
+# All plot scripts return the figure and axes object. Final tweaks and saving are intended to happen in notebooks
+import seaborn as sns
+import matplotlib.pyplot as plt
+import numpy as np
+
+from LinAdvFunc.da_functions import *
+
+def ensemble_plotter(states,m_const,da_const,j=0,t_start=0,t_end=0):
+ """
+ Plots the full forecast/background ensemble with observations, as well as the resulting analysis.
+ Plots all timesteps from t_start till t_end.
+ Initial starting conditions have no background state
+
+
+ ToDo:
+ intital state needs to be added as an exception with no backbround state.
+
+ Input:
+ j : experiment number
+ t_start and t_end: time frame to plot
+
+ Returns:
+ figure and axes
+ """
+ sns.set()
+ sns.set_style('whitegrid')
+ nc = da_const['ncyc']
+ if t_end ==0: t_end = nc
+
+ na = t_end-t_start
+
+ alpha = np.sqrt(1/da_const['nens'])
+
+ fig, ax = plt.subplots(na,2,figsize=(7.5,na*1.5),sharex='all',sharey='all')
+
+ #left bg, right analysis
+
+
+ #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,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,0].scatter(m_const["x_grid"][da_const["obs_loc"]],states[j]['obs'][c+t_start][da_const["obs_loc"]],zorder=3,label='observations')
+ 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')
+ ax[c,0].set_ylabel('assim step '+str(c+t_start) + '\n h [m]')
+
+
+ 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]])
+ 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]')
+ plt.legend()
+ return fig,ax
+
+
+def quad_plotter(quad_state,m_const,da_const):
+ """
+ Plots the initial background and blind forecast, as well as the analysis and forecast with observations
+
+ Input:
+ j : experiment number
+ t_start and t_end: time frame to plot
+
+ Returns:
+ figure and axes
+ """
+ sns.set()
+ sns.set_style('whitegrid')
+
+
+ alpha = np.sqrt(1/da_const['nens'])
+
+ fig, ax = plt.subplots(2,2,figsize=(10,5),sharex='all',sharey='all')
+
+
+
+ for i in range(da_const["nens"]):
+ ax[0,0].plot(m_const['x_grid'],quad_state['bg'][:,i],'r',alpha =alpha,zorder=1)
+ ax[0,1].plot(m_const['x_grid'],quad_state['bf'][:,i],'b',alpha =alpha,zorder=1)
+ ax[1,0].plot(m_const['x_grid'],quad_state['an'][:,i],'magenta',alpha =alpha,zorder=1)
+ ax[1,1].plot(m_const['x_grid'],quad_state['fc'][:,i],'c',alpha =alpha,zorder=1)
+
+
+ ax[0,0].plot(m_const["x_grid"],quad_state['tr_bg'],'k',zorder=10,label='truth')
+ ax[1,0].plot(m_const["x_grid"],quad_state['tr_bg'],'k',zorder=10,label='truth')
+ #ax[0,1].plot(m_const["x_grid"],quad_state['tr_fc'],'k',zorder=10,label='truth')
+ #ax[1,1].plot(m_const["x_grid"],quad_state['tr_fc'],'k',zorder=10,label='truth')
+
+ ax[0,0].plot(m_const['x_grid'],np.mean(quad_state['bg'][:,:],axis=1),'k--',alpha =1,zorder=2,label='ens mean')
+ ax[1,0].plot(m_const['x_grid'],np.mean(quad_state['an'][:,:],axis=1),'k--',alpha =1,zorder=2,label='ens mean')
+ ax[0,1].plot(m_const['x_grid'],np.mean(quad_state['bf'][:,:],axis=1),'k--',alpha =1,zorder=2,label='ens mean')
+ ax[1,1].plot(m_const['x_grid'],np.mean(quad_state['fc'][:,:],axis=1),'k--',alpha =1,zorder=2,label='ens mean')
+
+ ax[0,0].scatter(m_const["x_grid"][da_const["obs_loc"]],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"]],quad_state['obs'][da_const["obs_loc"]],zorder=3,label='observations',color='k')
+
+
+ ax[0,0].set_title('background with obs')
+ ax[0,1].set_title('free forecast')
+ ax[1,0].set_title('analysis with obs')
+ ax[1,1].set_title('forecast')
+ ax[0,1].set_xlim([m_const["x_grid"][0],m_const["x_grid"][-1]])
+ plt.subplots_adjust(wspace=0.03,hspace=0.30)
+ ax[1,1].set_xlabel('x [m]')
+ ax[1,0].set_xlabel('x [m]')
+ ax[0,0].set_ylabel('h [m]')
+ ax[1,0].set_ylabel('h [m]')
+ ax[1,0].legend()
+ return fig,ax
+
+def quad_plotter_v2(quad_state,m_const,da_const):
+ """
+ Plots the initial background and blind forecast, as well as the analysis and forecast with observations
+
+ Input:
+ j : experiment number
+ t_start and t_end: time frame to plot
+
+ 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[0,1].plot(m_const['x_grid']/1000.,quad_state['bf'][:,i],'b',alpha =alpha,zorder=1)
+ ax[1,0].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[1,0].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[1,0].plot(m_const['x_grid']/1000.,np.mean(quad_state['an'][:,:],axis=1),'k--',alpha =1,zorder=2,label='ens mean')
+ ax[0,1].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')
+
+ 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 with observations')
+ ax[0,1].set_title('free-forecast')
+ ax[1,0].set_title('analysis with observations')
+ 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('h [m]')
+ ax[1,0].set_ylabel('h [m]')
+ ax[1,0].legend()
+
+ # now to add in shading for response domain
+ ylimits = ax[0,0].get_ylim()
+ ax[0,1].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_paper(quad_state,m_const,da_const):
+ """
+ Plots the initial background and blind forecast, as well as the analysis and forecast with observations.
+
+ Swapped free forecast and analysis figure locations, also changed h to phi
+
+ 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,label='observations',color='k')
+ ax[0,1].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[1,0].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[0,1].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 B_plotter(states,j=0,ncyc=0,matrix='bg'):
+ """
+ Plots the covariance plotter of either the forecast/background ensemble ('bg'), analysis 'an', or the free/blind forecast 'bf'
+
+ """
+
+
+ fig, ax = plt.subplots(1,1,figsize=(5,5))
+
+ ax.set_ylabel('[h1,h2,h3,...]')
+ ax.set_xlabel('[h1,h2,h3,...]')
+
+ #B = np.cov(ensemble_diff.transpose())
+ B = np.cov(states[j][matrix][ncyc],ddof=1)
+ limit = np.max(np.abs(B))/4.*3.
+ pm = ax.pcolormesh(B,vmin=-limit,vmax=limit,cmap='RdBu_r')
+ plt.colorbar(pm,shrink = 0.8)
+ ax.set_aspect('equal')
+
+ return fig,ax
+
+def plot_scores_spread(rmse, spread):
+ """
+ Plots the evolution of the RMSE and spread over time and space.
+ Could be expanded to also plot the blind forecast if available
+
+ Todo:
+ plot next to each other with a shared y axis
+ """
+ sns.set_style('whitegrid')
+ figs = {}
+ for dim in ["space", "time"]:
+ figs[dim], ax = plt.subplots(1, sharex='col', figsize=(10, 4))
+
+ #for i, var in enumerate(['RMSE over ' + dim + ' h', 'RMSE over ' + dim + ' u']):
+ ax.set_title('RMSE over ' + dim, fontsize=24)
+ ax.tick_params(labelsize=18)
+ if 'bf' in rmse['time'].keys():
+ ax.plot(rmse[dim]['bf'], 'b', label="bf rmse="+str(round(np.nanmean(rmse[dim]['bf']),3)))
+ ax.plot(spread[dim]['bf'], 'b', label="bf spread="+str(round(np.nanmean(spread[dim]['bf']),3)), ls='--')
+ ax.plot(rmse[dim]['bg'] , 'r' , label="bg rmse=" +str(round(np.nanmean(rmse[dim]['bg']),3)),lw=3)
+ ax.plot(spread[dim]['bg'], 'r' , label="bg spread="+str(round(np.nanmean(spread[dim]['bg']),3)), ls='--',lw=3)
+ ax.plot(rmse[dim]['an'] , 'magenta', label="an rmse=" +str(round(np.mean(rmse[dim]['an']),3)))
+ ax.plot(spread[dim]['an'], 'magenta', label="an spread="+str(round(np.mean(spread[dim]['an']),3)), ls='--')
+
+
+ ax.set_xlim(left=0)
+ ax.legend()
+ if dim == "space":
+ ax.set_xlabel("space in grid points", fontsize=18)
+ if dim == "time":
+ ax.set_xlabel("time in assimilation cycles", fontsize=18)
+ return figs
+
+
+def plot_response_function_pdf_split(states,j=0,t_start=0,t_end=0,left_var='bg',right_var='an'):
+ """
+ Makes a split Seaborn violin plot of the respone function values. By default plots the analysis and forecast, but can also plot the free/blind forecast
+ Is not super efficient because it first transforms the response values into a panda dataframe to get the seaborn plot to work.
+ """
+ import seaborn as sns
+ import pandas as pd
+ sns.set()
+ fig,ax = plt.subplots(1,1,figsize=(6,4))
+
+ ncyc,nens = np.array(states[j]['response']['bg']).shape
+ if t_end==0: t_end=ncyc
+
+ #Got to figure out how to make sure the times line up, since the time integers of an, bg, and bf are all 1 step delayed
+ dict_label ={'an':'analysis','bg':'forecast-t1','bf':'forecast-t2'}
+ dict_color ={'an':'magenta','bg':'r','bf':'b'}
+ #t_diff = t_adjust[right_var]-t_adjust[left_var]
+
+ nt = t_end-t_start
+
+
+ plot_data = {
+ 'response' : np.hstack([np.array(states[j]['response'][left_var ][t_start:t_end]).ravel(),
+ np.array(states[j]['response'][right_var][t_start:t_end]).ravel()]),
+ 'cyc' : np.hstack([np.repeat(np.arange(nt),nens),np.repeat(np.arange(nt),nens)])+t_start,
+ 'type' : [dict_label[left_var]]*nt*nens+[dict_label[right_var]]*nt*nens}
+
+
+
+ sns.violinplot(data=plot_data,y='response',x='cyc',hue='type', inner=None, orient="v",split=True,palette={dict_label[left_var]:dict_color[left_var],dict_label[right_var]:dict_color[right_var]},cut=0,bw=0.25)
+
+
+ plt.plot(states[j]['response']['truth'][t_start:t_end],'k--',marker='o',label='truth')
+
+ plt.legend(loc='lower center')
+ ax.set_xlabel('assimilation step')
+ ax.set_ylabel('PDF of response function J')
+
+ #Moves the left and right pdf a bit appart so you can see where they end. Thank you stackoverflow
+ inner=None
+ width = 0.75
+ delta = 0.05
+ final_width = width - delta
+ offset_violinplot_halves(ax, delta, final_width, inner, 'vertical')
+ return fig,ax
+
+
+
+
+def offset_violinplot_halves(ax, delta, width, inner, direction):
+ """
+ From Stackoverflow!
+
+ This function offsets the halves of a violinplot to compare tails
+ or to plot something else in between them. This is specifically designed
+ for violinplots by Seaborn that use the option `split=True`.
+
+ For lines, this works on the assumption that Seaborn plots everything with
+ integers as the center.
+
+ Args:
+ <ax> The axis that contains the violinplots.
+ <delta> The amount of space to put between the two halves of the violinplot
+ <width> The total width of the violinplot, as passed to sns.violinplot()
+ <inner> The type of inner in the seaborn
+ <direction> Orientation of violinplot. 'hotizontal' or 'vertical'.
+
+ Returns:
+ - NA, modifies the <ax> directly
+ """
+ import matplotlib.collections
+ # offset stuff
+ if inner == 'sticks':
+ lines = ax.get_lines()
+ for line in lines:
+ if direction == 'horizontal':
+ data = line.get_ydata()
+ print(data)
+ if int(data[0] + 1)/int(data[1] + 1) < 1:
+ # type is top, move neg, direction backwards for horizontal
+ data -= delta
+ else:
+ # type is bottom, move pos, direction backward for hori
+ data += delta
+ line.set_ydata(data)
+ elif direction == 'vertical':
+ data = line.get_xdata()
+ print(data)
+ if int(data[0] + 1)/int(data[1] + 1) < 1:
+ # type is left, move neg
+ data -= delta
+ else:
+ # type is left, move pos
+ data += delta
+ line.set_xdata(data)
+
+
+ for ii, item in enumerate(ax.collections):
+ # axis contains PolyCollections and PathCollections
+ if isinstance(item, matplotlib.collections.PolyCollection):
+ # get path
+ path, = item.get_paths()
+ vertices = path.vertices
+ half_type = _wedge_dir(vertices, direction)
+ # shift x-coordinates of path
+ if half_type in ['top','bottom']:
+ if inner in ["sticks", None]:
+ if half_type == 'top': # -> up
+ vertices[:,1] -= delta
+ elif half_type == 'bottom': # -> down
+ vertices[:,1] += delta
+ elif half_type in ['left', 'right']:
+ if inner in ["sticks", None]:
+ if half_type == 'left': # -> left
+ vertices[:,0] -= delta
+ elif half_type == 'right': # -> down
+ vertices[:,0] += delta
+
+def _wedge_dir(vertices, direction):
+ """
+ Args:
+ <vertices> The vertices from matplotlib.collections.PolyCollection
+ <direction> Direction must be 'horizontal' or 'vertical' according to how
+ your plot is laid out.
+ Returns:
+ - a string in ['top', 'bottom', 'left', 'right'] that determines where the
+ half of the violinplot is relative to the center.
+ """
+ if direction == 'horizontal':
+ result = (direction, len(set(vertices[1:5,1])) == 1)
+ elif direction == 'vertical':
+ result = (direction, len(set(vertices[-3:-1,0])) == 1)
+ outcome_key = {('horizontal', True): 'bottom',
+ ('horizontal', False): 'top',
+ ('vertical', True): 'left',
+ ('vertical', False): 'right'}
+ # if the first couple x/y values after the start are the same, it
+ # is the input direction. If not, it is the opposite
+ return outcome_key[result]
+
+
+#Little subplot labeling script found online from https://gist.github.com/tacaswell/9643166
+#Put here for no good reason
+#Warning, if you define your colorbar using an additional axes it will give that a label as well, which is pretty funny but also very annoying.
+
+import string
+from itertools import cycle
+from six.moves import zip
+def label_axes_abcd(fig, labels=None, loc=None, **kwargs):
+ """
+ Walks through axes and labels each.
+ kwargs are collected and passed to `annotate`
+ Parameters
+ ----------
+ fig : Figure
+ Figure object to work on
+ labels : iterable or None
+ iterable of strings to use to label the axes.
+ If None, lower case letters are used.
+
+ loc : len=2 tuple of floats
+ Where to put the label in axes-fraction units
+ """
+ if labels is None:
+ labels = string.ascii_lowercase
+
+ # re-use labels rather than stop labeling
+ labels = cycle(labels)
+ if loc is None:
+ loc = (.9, .9)
+ for ax, lab in zip(fig.axes, labels):
+ ax.annotate(lab, xy=loc,
+ xycoords='axes fraction',
+ **kwargs)
+###########################################################################################
+# 2022 baby, now with sat data
+###########################################################################################
+
+
+def ensemble_plotter_22(states,m_const,da_const,j=0,t_start=0,t_end=0,h_c=None):
+ """
+ Plots the full forecast/background ensemble with observations, as well as the resulting analysis.
+ Plots all timesteps from t_start till t_end.
+ Initial starting conditions have no background state
+
+
+
+ Input:
+ j : experiment number
+ t_start and t_end: time frame to plot
+
+ Returns:
+ figure and axes
+ """
+ sns.set()
+ sns.set_style('whitegrid')
+ nc = da_const['ncyc']
+ if t_end ==0: t_end = nc
+
+ na = t_end-t_start
+
+ alpha = np.sqrt(1/da_const['nens'])
+
+ fig, ax = plt.subplots(na,2,figsize=(7.5,na*1.5),sharex='all',sharey='all')
+
+ #left bg, right analysis
+
+
+ #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,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')
+
+ 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 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')
+
+ #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[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]])
+ 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]')
+ plt.legend()
+ return fig,ax
+
+
+
+
+def quad_plotter_22(quad_state,m_const,da_const):
+ """
+ Plots the initial background and blind forecast, as well as the analysis and forecast with observations.
+
+ Only difference to previous versions is that it can deal with no obs being present
+
+
+ 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[0,1].plot(m_const['x_grid']/1000.,quad_state['bf'][:,i],'b',alpha =alpha,zorder=1)
+ ax[1,0].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[1,0].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[1,0].plot(m_const['x_grid']/1000.,np.mean(quad_state['an'][:,:],axis=1),'k--',alpha =1,zorder=2,label='ens mean')
+ ax[0,1].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,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[0,1].set_title('free-forecast')
+ ax[1,0].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('h [m]')
+ ax[1,0].set_ylabel('h [m]')
+ ax[1,0].legend()
+
+ # now to add in shading for response domain
+ ylimits = ax[0,0].get_ylim()
+ ax[0,1].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 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')
+ """
+ Plots the background and analysis ensemble, as well as the satetlitte obs equivalents. Includes observations as well.
+
+ Is still quite rough around the edges, and will hopefully be improced upon
+ """
+
+ sat_an = sat_operator(an)
+ sat_bg = sat_operator(bg)
+ sat_tr_bg = sat_operator(truth)
+ n_obs_h =len(da_const["obs_loc"])
+ n_obs_sat =len(da_const["obs_loc_sat"])
+
+ dY_b, y_ol_b = state_to_observation_space(bg,m_const,da_const,sat_operator)
+ y_b = dY_b.T + y_ol_b
+ y_b = y_b.T
+
+
+ # for plotting the pixels
+ window = m_const['x_grid'][da_const["obs_loc_sat"][1]]-m_const['x_grid'][da_const["obs_loc_sat"][0]]
+ xpixmin = m_const['x_grid'][da_const["obs_loc_sat"]]-window/2
+ xpixmax = m_const['x_grid'][da_const["obs_loc_sat"]]+window/2
+
+ for i in range(bg.shape[1]):
+ #ax[0].plot(m_const['x_grid'],x_a[:,i],'r',alpha=0.2)
+ ax[1,0].plot(m_const['x_grid']/1000,an [:,i],'magenta',alpha=0.2)
+ ax[0,0].plot(m_const['x_grid']/1000,bg [:,i],'r',alpha=0.2)
+
+ ax[1,1].hlines(
+ sat_an [da_const["obs_loc_sat"],i],
+ xpixmin/1000,xpixmax/1000,
+ color='magenta',alpha=0.2,lw=3)
+ ax[0,1].hlines(
+ sat_bg [da_const["obs_loc_sat"],i],
+ xpixmin/1000,xpixmax/1000,
+ color='r',alpha=0.2,lw=3)
+ ax[1,0].plot(m_const['x_grid']/1000+1000,an [:,-1],'magenta',alpha=1,label='background ensemble')
+ ax[0,0].plot(m_const['x_grid']/1000+1000,bg [:,-1],'r',alpha=1,label='analysis ensemble')
+
+ #ax[0].plot(m_const['x_grid'],x_ol_a,'r')
+ ax[1,0].plot(m_const['x_grid']/1000,truth,'k',linewidth=2,label='truth')
+ ax[0,0].plot(m_const['x_grid']/1000,truth,'k',linewidth=2)
+
+
+ if n_obs_h>0:
+ ax[1,0].scatter(m_const['x_grid'][da_const['obs_loc']]/1000,obs[da_const['obs_loc']],
+ c='k',label='h point observation')
+ ax[0,0].scatter(m_const['x_grid'][da_const['obs_loc']]/1000,obs[da_const['obs_loc']],c='k')#,label='h point obs')
+ if n_obs_sat>0:
+ ax[1,1].scatter(m_const['x_grid'][da_const['obs_loc_sat']]/1000,obs_sat[da_const['obs_loc_sat']],
+ marker='x',s=50,c='k',label='reflectance observation')
+ ax[0,1].scatter(m_const['x_grid'][da_const['obs_loc_sat']]/1000,obs_sat[da_const['obs_loc_sat']],marker='x',s=50,c='k')
+
+
+ ax[0,0].plot(m_const['x_grid']/1000,np.mean(bg,axis=1),'k--',lw=2,label='ensemble mean')
+ ax[1,0].plot(m_const['x_grid']/1000,np.mean(an,axis=1),'k--',lw=2)#,label='ens mean')
+
+ if h_c is not None:
+ ax[0,0].hlines(h_c,m_const['x_grid'][0]/1000,m_const['x_grid'][-1]/1000,color='k',ls=':')
+ ax[1,0].hlines(h_c,m_const['x_grid'][0]/1000,m_const['x_grid'][-1]/1000,color='k',ls=':',label=r'cloud threshold $h_c$')
+
+ label_axes_abcd(fig)
+
+ ax1 = ax[0,1].twinx()
+ ax1.set_yticks([0.3,0.7])
+ ax1.set_yticklabels(['clear \n sky','cloud'])
+ ax2 = ax[1,1].twinx()
+ ax2.set_yticks([0.3,0.7])
+ ax2.set_yticklabels(['clear \n sky','cloud'])
+ ax1.set_ylim(-0.05,1.05)
+ ax2.set_ylim(-0.05,1.05)
+ ax[1,1].set_ylim(-0.05,1.05)
+ ax[0,0].set_yticks([0,0.5,1])
+ ax[0,1].set_yticks([0,0.3,0.7,1.0])
+ ax[0,0].set_yticklabels(['0',r'$h_c$','1'])
+
+ ax[1,0].set_xlabel('x [km]')
+ ax[1,1].set_xlabel('x [km]')
+ ax[1,1].set_xlim(m_const['x_grid'][0]/1000,30)#m_const['x_grid'][-1]/1000)
+ ax[0,0].set_ylabel('background \n h [m]')
+ ax[1,0].set_ylabel('analysis \n h[m]')
+ ax[0,1].set_ylabel('reflectance')
+ ax[1,1].set_ylabel('reflectance')
+ plt.subplots_adjust(wspace=0.2,hspace=0.1)
+
+ ax[0,1].step(m_const['x_grid'][da_const['obs_loc_sat']]/1000,np.mean(sat_bg[da_const['obs_loc_sat']],axis=1),'k--',lw=2,where='mid')#,marker='s',markersize=4,label='ens mean')
+ ax[1,1].step(m_const['x_grid'][da_const['obs_loc_sat']]/1000,np.mean(sat_an[da_const['obs_loc_sat']],axis=1),'k--',lw=2,where='mid')#,marker='s',markersize=4,label='ens mean')
+
+ lines_labels = [ax.get_legend_handles_labels() for ax in fig.axes]
+ lines, labels = [sum(lol, []) for lol in zip(*lines_labels)]
+ fig.legend(lines, labels, loc='upper center',ncol=4)
+ return fig, ax
+
+
+
+def plot_ensemble_sat_analysis_paper(bg,an,obs,obs_sat,truth,sat_operator,m_const,da_const,h_c=None):
+ fig ,ax = plt.subplots(2,2,figsize=(7.5,5),sharey='row',sharex='all')
+ """
+ Plots the background and analysis ensemble of the state, as well as the satetlitte obs equivalents. Includes observations as well.
+ """
+
+ sat_an = sat_operator(an)
+ sat_bg = sat_operator(bg)
+ sat_tr_bg = sat_operator(truth)
+ n_obs_h =len(da_const["obs_loc"])
+ n_obs_sat =len(da_const["obs_loc_sat"])
+
+ dY_b, y_ol_b = state_to_observation_space(bg,m_const,da_const,sat_operator)
+ y_b = dY_b.T + y_ol_b
+ y_b = y_b.T
+
+
+ # for plotting the pixels
+ window = m_const['x_grid'][da_const["obs_loc_sat"][1]]-m_const['x_grid'][da_const["obs_loc_sat"][0]]
+ xpixmin = m_const['x_grid'][da_const["obs_loc_sat"]]-window/2
+ xpixmax = m_const['x_grid'][da_const["obs_loc_sat"]]+window/2
+
+ for i in range(bg.shape[1]):
+ #ax[0].plot(m_const['x_grid'],x_a[:,i],'r',alpha=0.2)
+ ax[0,1].plot(m_const['x_grid']/1000,an [:,i],'magenta',alpha=0.2)
+ ax[0,0].plot(m_const['x_grid']/1000,bg [:,i],'r',alpha=0.2)
+
+ # plotting the pixels
+
+ #ax[1,1].plot(m_const['x_grid'],sat_an [:,i],'magenta',alpha=0.2)
+ #ax[0,1].plot(m_const['x_grid'],sat_bg [:,i],'r',alpha=0.2)
+ ax[1,1].hlines(
+ sat_an [da_const["obs_loc_sat"],i],
+ xpixmin/1000,xpixmax/1000,
+ color='magenta',alpha=0.2,lw=3)
+ ax[1,0].hlines(
+ sat_bg [da_const["obs_loc_sat"],i],
+ xpixmin/1000,xpixmax/1000,
+ color='r',alpha=0.2,lw=3)
+ ax[0,1].plot(m_const['x_grid']/1000+1000,an [:,-1],'magenta',alpha=1,label='analysis ensemble')
+ ax[0,0].plot(m_const['x_grid']/1000+1000,bg [:,-1],'r',alpha=1,label='background ensemble')
+
+ #ax[0].plot(m_const['x_grid'],x_ol_a,'r')
+ ax[0,1].plot(m_const['x_grid']/1000,truth,'k',linewidth=2,label='truth')
+ ax[0,0].plot(m_const['x_grid']/1000,truth,'k',linewidth=2)
+
+
+ if n_obs_h>0:
+ ax[0,1].scatter(m_const['x_grid'][da_const['obs_loc']]/1000,obs[da_const['obs_loc']],
+ c='k',label=r'$\phi$ point observation')
+ ax[0,0].scatter(m_const['x_grid'][da_const['obs_loc']]/1000,obs[da_const['obs_loc']],c='k')#,label='h point obs')
+ if n_obs_sat>0:
+ ax[1,1].scatter(m_const['x_grid'][da_const['obs_loc_sat']]/1000,obs_sat[da_const['obs_loc_sat']],
+ marker='x',s=50,c='k',label='reflectance observation')
+ ax[1,0].scatter(m_const['x_grid'][da_const['obs_loc_sat']]/1000,obs_sat[da_const['obs_loc_sat']],marker='x',s=50,c='k')
+
+ ax[0,0].plot(m_const['x_grid']/1000,np.mean(bg,axis=1),'k--',lw=2,label='ensemble mean')
+ ax[0,1].plot(m_const['x_grid']/1000,np.mean(an,axis=1),'k--',lw=2)#,label='ens mean')
+
+ if h_c is not None:
+ ax[0,0].hlines(h_c,m_const['x_grid'][0]/1000,m_const['x_grid'][-1]/1000,color='k',ls=':')
+ ax[0,1].hlines(h_c,m_const['x_grid'][0]/1000,m_const['x_grid'][-1]/1000,color='k',ls=':',label=r'cloud threshold $\phi_c$')
+
+ label_axes_abcd(fig)
+
+ ax1 = ax[0,1].twinx()
+ ax1.set_yticks([0.,0.5,1])
+ ax1.set_yticklabels(['0',r'$\phi_c$','1'])
+ ax1.set_ylim(-0.3,1.3)
+ ax[0,0].set_ylim(-0.3,1.3)
+ ax2 = ax[1,1].twinx()
+ ax2.set_yticks([0.3,0.7])
+ ax2.set_yticklabels(['clear \n sky','cloud'])
+ ax2.set_ylim(-0.05,1.05)
+ ax[1,1].set_ylim(-0.05,1.05)
+ ax[0,0].set_yticks([0,0.5,1])
+ ax[1,0].set_yticks([0,0.3,0.7,1.0])
+ ax[0,0].set_yticklabels(['0',r'$\phi_c$','1'])
+
+
+ ax[1,0].set_xlabel('x [km]')
+ ax[1,1].set_xlabel('x [km]')
+ ax[0,1].set_xlim(m_const['x_grid'][0]/1000,30)#m_const['x_grid'][-1]/1000)
+ ax[0,0].set_title('background')
+ ax[0,1].set_title('analysis')
+ ax[1,0].set_ylabel('reflectance')
+ ax[0,0].set_ylabel(r'$\phi$')
+ plt.subplots_adjust(wspace=0.05,hspace=0.05)
+
+ ax[1,0].step(m_const['x_grid'][da_const['obs_loc_sat']]/1000,np.mean(sat_bg[da_const['obs_loc_sat']],axis=1),'k--',lw=2,where='mid')#,marker='s',markersize=4,label='ens mean')
+ ax[1,1].step(m_const['x_grid'][da_const['obs_loc_sat']]/1000,np.mean(sat_an[da_const['obs_loc_sat']],axis=1),'k--',lw=2,where='mid')#,marker='s',markersize=4,label='ens mean')
+
+ lines_labels = [ax.get_legend_handles_labels() for ax in fig.axes]
+ lines, labels = [sum(lol, []) for lol in zip(*lines_labels)]
+ fig.legend(lines, labels, loc='upper center',bbox_to_anchor=(0.5,1.1),ncol=3)
+ return fig, ax
+
+
+def plot_ensemble_sat_analysis_abstract(bg,an,obs,obs_sat,truth,sat_operator,m_const,da_const,h_c=None):
+ fig ,ax = plt.subplots(2,2,figsize=(5.5,4.5),sharey='row',sharex='all')
+ """
+ Modified for paper visual abstract
+ Plots the background and analysis ensemble, as well as the satetlitte obs equivalents. Includes observations as well.
+ """
+
+ sat_an = sat_operator(an)
+ sat_bg = sat_operator(bg)
+ sat_tr_bg = sat_operator(truth)
+ n_obs_h =len(da_const["obs_loc"])
+ n_obs_sat =len(da_const["obs_loc_sat"])
+
+ dY_b, y_ol_b = state_to_observation_space(bg,m_const,da_const,sat_operator)
+ y_b = dY_b.T + y_ol_b
+ y_b = y_b.T
+
+
+ # for plotting the pixels
+ window = m_const['x_grid'][da_const["obs_loc_sat"][1]]-m_const['x_grid'][da_const["obs_loc_sat"][0]]
+ xpixmin = m_const['x_grid'][da_const["obs_loc_sat"]]-window/2
+ xpixmax = m_const['x_grid'][da_const["obs_loc_sat"]]+window/2
+
+ for i in range(bg.shape[1]):
+ #ax[0].plot(m_const['x_grid'],x_a[:,i],'r',alpha=0.2)
+ ax[0,1].plot(m_const['x_grid']/1000,an [:,i],'magenta',alpha=0.2)
+ ax[0,0].plot(m_const['x_grid']/1000,bg [:,i],'r',alpha=0.2)
+
+ ax[1,1].hlines(
+ sat_an [da_const["obs_loc_sat"],i],
+ xpixmin/1000,xpixmax/1000,
+ color='magenta',alpha=0.2,lw=3)
+ ax[1,0].hlines(
+ sat_bg [da_const["obs_loc_sat"],i],
+ xpixmin/1000,xpixmax/1000,
+ color='r',alpha=0.2,lw=3)
+ ax[0,1].plot(m_const['x_grid']/1000+1000,an [:,-1],'magenta',alpha=1,label='analysis ensemble')
+ ax[0,0].plot(m_const['x_grid']/1000+1000,bg [:,-1],'r',alpha=1,label='background ensemble')
+
+ #ax[0].plot(m_const['x_grid'],x_ol_a,'r')
+ ax[0,1].plot(m_const['x_grid']/1000,truth,'k',linewidth=2,label='truth')
+ ax[0,0].plot(m_const['x_grid']/1000,truth,'k',linewidth=2)
+
+ if n_obs_h>0:
+ ax[0,1].scatter(m_const['x_grid'][da_const['obs_loc']]/1000,obs[da_const['obs_loc']],
+ c='k',label=r'$\phi$ point observation')
+ ax[0,0].scatter(m_const['x_grid'][da_const['obs_loc']]/1000,obs[da_const['obs_loc']],c='k')#,label='h point obs')
+ if n_obs_sat>0:
+ ax[1,1].scatter(m_const['x_grid'][da_const['obs_loc_sat']]/1000,obs_sat[da_const['obs_loc_sat']],
+ marker='x',s=50,c='k',label='reflectance observation')
+ ax[1,0].scatter(m_const['x_grid'][da_const['obs_loc_sat']]/1000,obs_sat[da_const['obs_loc_sat']],marker='x',s=50,c='k')
+
+
+ ax[0,0].plot(m_const['x_grid']/1000,np.mean(bg,axis=1),'k--',lw=2,label='ensemble mean')
+ ax[0,1].plot(m_const['x_grid']/1000,np.mean(an,axis=1),'k--',lw=2)#,label='ens mean')
+
+ if h_c is not None:
+ ax[0,0].hlines(h_c,m_const['x_grid'][0]/1000,m_const['x_grid'][-1]/1000,color='k',ls=':')
+ ax[0,1].hlines(h_c,m_const['x_grid'][0]/1000,m_const['x_grid'][-1]/1000,color='k',ls=':',label=r'cloud threshold $\phi_c$')
+
+
+ ax1 = ax[0,1].twinx()
+ ax1.set_yticks([0.,0.5,1])
+ ax1.set_yticklabels(['0',r'$\phi_c$','1'])
+ ax1.set_ylim(-0.3,1.3)
+ ax[0,0].set_ylim(-0.3,1.3)
+ ax2 = ax[1,1].twinx()
+ ax2.set_yticks([0.3,0.7])
+ ax2.set_yticklabels(['clear \n sky','cloud'])
+ ax2.set_ylim(-0.05,1.05)
+ ax[1,1].set_ylim(-0.05,1.05)
+ ax[0,0].set_yticks([0,0.5,1])
+ ax[1,0].set_yticks([0.3,0.7])
+ ax[1,0].set_yticklabels([])
+ ax[0,0].set_yticklabels([])
+ ax[1,0].set_xticklabels([])
+ ax[1,1].set_xticklabels([])
+
+
+ ax[1,0].set_xlabel('x',labelpad=-8)
+ ax[1,1].set_xlabel('x',labelpad=-8)
+ ax[0,1].set_xlim(m_const['x_grid'][0]/1000,30)#m_const['x_grid'][-1]/1000)
+ ax[0,0].set_title('background')
+ ax[0,1].set_title('analysis')
+ ax[0,0].set_ylabel(r'$\phi$',labelpad=-10)
+ ax[1,0].set_ylabel('reflectance',labelpad=-8)
+ plt.subplots_adjust(wspace=0.05,hspace=0.05)
+
+ ax[1,0].step(m_const['x_grid'][da_const['obs_loc_sat']]/1000,np.mean(sat_bg[da_const['obs_loc_sat']],axis=1),'k--',lw=2,where='mid')#,marker='s',markersize=4,label='ens mean')
+ ax[1,1].step(m_const['x_grid'][da_const['obs_loc_sat']]/1000,np.mean(sat_an[da_const['obs_loc_sat']],axis=1),'k--',lw=2,where='mid')#,marker='s',markersize=4,label='ens mean')
+
+ lines_labels = [ax.get_legend_handles_labels() for ax in fig.axes]
+ lines, labels = [sum(lol, []) for lol in zip(*lines_labels)]
+ return fig, ax
+
+def plot_J_quad_paper(J_dict,quad,sens,dx,bw=0.3,dJ=True):
+ """
+ Plots the forecast metric distributions of the free forecast, forecast, and their linear approximations for the given sensitivity
+ """
+
+ fig = plt.figure(figsize=(4,3))
+ nens = len(J_dict['bf'])
+ dX_bg=(quad['bg'].T-np.mean(quad['bg'],axis=1)).T
+ dX_an=(quad['an'].T-np.mean(quad['an'],axis=1)).T
+ 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 ))
+ #'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 = {
+ 'response' : np.hstack([J_dict['bf']-np.mean(J_dict['bf']),dJ_ff,J_dict['fc']-np.mean(J_dict['fc']),dJ_fc]),
+ 'x_pos' : np.hstack([np.zeros(nens)+0,np.zeros(nens)+1,np.zeros(nens)+2,np.zeros(nens)+3]),
+ 'type' : ['blindforecast']*nens+['estimated']*nens+['forecast']*nens}
+ else:
+ plot_data = {
+ 'response' : np.hstack([J_dict['bf'],dJ_ff+np.mean(J_dict['bf']),J_dict['fc'],dJ_fc+np.mean(J_dict['fc'])]),
+ 'x_pos' : np.hstack([np.zeros(nens)+0,np.zeros(nens)+1,np.zeros(nens)+2,np.zeros(nens)+3]),
+ 'type' : ['blindforecast']*nens+['estimated']*nens+['forecast']*nens}
+
+ my_pal = ["blue", "peru","cyan","orange" ]
+
+ PROPS = {
+ '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.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')
+
+ #if dJ == False: ax.hlines(J_dict['tr_fc'],-0.5,2.5,'k',ls='--',label='truth'); plt.legend()
+ #if dJ: ax.hlines(0,-0.5,3.5,'k',ls='--')
+ ax.set_xlim(-0.5,3.5)
+ if dJ == False: ax.set_ylabel(r'$j$')
+ if dJ: ax.set_ylabel(r'$\delta j$')
+ ax.set_xticklabels(['free-\nforecast','estimated \n free-forecast','\n forecast','estimated \n forecast'])
+ return fig, ax
+
+def vr_scatter_v6(vr_tot1,vr_rea1,vr_tot2,vr_rea2,vr_tot3,vr_rea3,vr_tot4,vr_tot5,vr_tot6,alpha=0.3,alpha2=0.5,color1='blueviolet',color2=plt.cm.viridis(0.8),
+ label1='',label2='',label3='',llabel1='explicit sens',llabel2='implicit sens'):
+ """
+ Just a 2x3 scatterplot with shared axis and a linear regressions """
+
+ fig, ax = plt.subplots(2,3,figsize=(8,5.5),sharex='all',sharey='all')
+
+ color = color1
+ vr_rea = vr_rea1
+ vr_tot = vr_tot1
+ m, b = np.polyfit(vr_rea, vr_tot, 1)
+ ax[0,0].plot([-1000,1000],[-1000,1000],'k--',zorder=0)
+ ax[0,0].scatter(vr_rea,vr_tot,c=color,alpha=alpha,s=5,zorder=1,label=llabel1)
+ ax[0,0].plot(vr_rea, m*np.array(vr_rea) + b,'k',zorder=2)
+ ax[0,0].scatter(np.mean(vr_rea),np.mean(vr_tot),c='w',s=100,zorder=2)
+ ax[0,0].scatter(np.mean(vr_rea),np.mean(vr_tot),c='k',s=50,zorder=3)
+ ax[0,0].scatter(np.mean(vr_rea),np.mean(vr_tot),c='w',s=10,zorder=4)
+ ax[0,0].set_aspect('equal', 'box')
+ ax[0,0].legend(loc='lower center')
+
+
+ vr_rea = vr_rea2
+ vr_tot = vr_tot2
+ m, b = np.polyfit(vr_rea, vr_tot, 1)
+ ax[0,1].plot([-1000,1000],[-1000,1000],'k--',zorder=0)
+ ax[0,1].scatter(vr_rea,vr_tot,c=color,alpha=alpha,s=5,zorder=1)
+ ax[0,1].plot(vr_rea, m*np.array(vr_rea) + b,'k',zorder=2)
+ ax[0,1].scatter(np.mean(vr_rea),np.mean(vr_tot),c='w',s=100,zorder=2)
+ ax[0,1].scatter(np.mean(vr_rea),np.mean(vr_tot),c='k',s=50,zorder=3)
+ ax[0,1].scatter(np.mean(vr_rea),np.mean(vr_tot),c='w',s=10,zorder=4)
+ ax[0,1].set_aspect('equal', 'box')
+
+ vr_rea = vr_rea3
+ vr_tot = vr_tot3
+ m, b = np.polyfit(vr_rea, vr_tot, 1)
+ ax[0,2].plot([-1000,1000],[-1000,1000],'k--',zorder=0)
+ ax[0,2].scatter(vr_rea,vr_tot,c=color,alpha=alpha,s=5,zorder=1)
+ ax[0,2].plot(vr_rea, m*np.array(vr_rea) + b,'k',zorder=2)
+ ax[0,2].scatter(np.mean(vr_rea),np.mean(vr_tot),c='w',s=100,zorder=2)
+ ax[0,2].scatter(np.mean(vr_rea),np.mean(vr_tot),c='k',s=50,zorder=3)
+ ax[0,2].scatter(np.mean(vr_rea),np.mean(vr_tot),c='w',s=10,zorder=4)
+ ax[0,2].set_aspect('equal', 'box')
+
+ color = color2
+ vr_rea = vr_rea1
+ vr_tot = vr_tot4
+ m, b = np.polyfit(vr_rea, vr_tot, 1)
+ ax[1,0].plot([-1000,1000],[-1000,1000],'k--',zorder=0)
+ ax[1,0].scatter(vr_rea,vr_tot,c=color,alpha=alpha2,s=5,zorder=1,label=llabel2)
+ ax[1,0].plot(vr_rea, m*np.array(vr_rea) + b,'k',zorder=2)
+ ax[1,0].scatter(np.mean(vr_rea),np.mean(vr_tot),c='w',s=100,zorder=2)
+ ax[1,0].scatter(np.mean(vr_rea),np.mean(vr_tot),c='k',s=50,zorder=3)
+ ax[1,0].scatter(np.mean(vr_rea),np.mean(vr_tot),c='w',s=10,zorder=4)
+ ax[1,0].set_aspect('equal', 'box')
+ ax[1,0].legend(loc='lower center')
+
+
+ vr_rea = vr_rea2
+ vr_tot = vr_tot5
+ m, b = np.polyfit(vr_rea, vr_tot, 1)
+ ax[1,1].plot([-1000,1000],[-1000,1000],'k--',zorder=0)
+ ax[1,1].scatter(vr_rea,vr_tot,c=color,alpha=alpha2,s=5,zorder=1)
+ ax[1,1].plot(vr_rea, m*np.array(vr_rea) + b,'k',zorder=2)
+ ax[1,1].scatter(np.mean(vr_rea),np.mean(vr_tot),c='w',s=100,zorder=2)
+ ax[1,1].scatter(np.mean(vr_rea),np.mean(vr_tot),c='k',s=50,zorder=3)
+ ax[1,1].scatter(np.mean(vr_rea),np.mean(vr_tot),c='w',s=10,zorder=4)
+ ax[1,1].set_aspect('equal', 'box')
+
+ vr_rea = vr_rea3
+ vr_tot = vr_tot6
+ m, b = np.polyfit(vr_rea, vr_tot, 1)
+ ax[1,2].plot([-1000,1000],[-1000,1000],'k--',zorder=0)
+ ax[1,2].scatter(vr_rea,vr_tot,c=color,alpha=alpha2,s=5,zorder=1)
+ ax[1,2].plot(vr_rea, m*np.array(vr_rea) + b,'k',zorder=2)
+ ax[1,2].scatter(np.mean(vr_rea),np.mean(vr_tot),c='w',s=100,zorder=2)
+ ax[1,2].scatter(np.mean(vr_rea),np.mean(vr_tot),c='k',s=50,zorder=3)
+ ax[1,2].scatter(np.mean(vr_rea),np.mean(vr_tot),c='w',s=10,zorder=4)
+ ax[1,2].set_aspect('equal', 'box')
+
+
+
+ plt.subplots_adjust(wspace=0.05,hspace=0.05)
+
+ ax[1,0].set_xlabel('variance reduction')
+ ax[1,1].set_xlabel('variance reduction')
+ ax[1,2].set_xlabel('variance reduction')
+ ax[1,0].set_ylabel('estimated var reduction')
+ ax[0,0].set_ylabel('estimated var reduction')
+
+ ax[0,0].set_title(label1)
+ ax[0,1].set_title(label2)
+ ax[0,2].set_title(label3)
+
+ x_max = np.max(np.hstack([vr_rea1,vr_rea2,vr_tot1,vr_tot2]))
+ x_min = np.min(np.hstack([vr_rea1,vr_rea2,vr_tot1,vr_tot2]))
+
+
+ ax[0,0].set_xlim(x_min,x_max)
+ ax[0,0].set_ylim(x_min,x_max)
+ plt.locator_params(axis='y', nbins=4)
+ plt.locator_params(axis='x', nbins=4)
+ return fig, ax
+
+
+def vr_scatter_v4(vr_tot1,vr_rea1,vr_tot2,vr_rea2,vr_tot3,vr_tot4,alpha=0.3,alpha2=0.5,color1='blueviolet',color2=plt.cm.viridis(0.8),
+ label1='',label2='',label3='',llabel1='explicit sens',llabel2='implicit sens'):
+ """
+ Just a 2x2 scatterplot with shared axis and a linear regressions """
+
+ fig, ax = plt.subplots(2,2,figsize=(5.5,5.5),sharex='all',sharey='all')
+
+ color = color1
+ vr_rea = vr_rea1
+ vr_tot = vr_tot1
+ m, b = np.polyfit(vr_rea, vr_tot, 1)
+ ax[0,0].plot([-1000,1000],[-1000,1000],'k--',zorder=0)
+ ax[0,0].scatter(vr_rea,vr_tot,c=color,alpha=alpha,s=5,zorder=1,label=llabel1)
+ ax[0,0].plot(vr_rea, m*np.array(vr_rea) + b,'k',zorder=2)
+ ax[0,0].scatter(np.mean(vr_rea),np.mean(vr_tot),c='w',s=100,zorder=2)
+ ax[0,0].scatter(np.mean(vr_rea),np.mean(vr_tot),c='k',s=50,zorder=3)
+ ax[0,0].scatter(np.mean(vr_rea),np.mean(vr_tot),c='w',s=10,zorder=4)
+ ax[0,0].set_aspect('equal', 'box')
+ ax[0,0].legend(loc='lower center')
+
+
+ vr_rea = vr_rea2
+ vr_tot = vr_tot2
+ m, b = np.polyfit(vr_rea, vr_tot, 1)
+ ax[0,1].plot([-1000,1000],[-1000,1000],'k--',zorder=0)
+ ax[0,1].scatter(vr_rea,vr_tot,c=color,alpha=alpha,s=5,zorder=1)
+ ax[0,1].plot(vr_rea, m*np.array(vr_rea) + b,'k',zorder=2)
+ ax[0,1].scatter(np.mean(vr_rea),np.mean(vr_tot),c='w',s=100,zorder=2)
+ ax[0,1].scatter(np.mean(vr_rea),np.mean(vr_tot),c='k',s=50,zorder=3)
+ ax[0,1].scatter(np.mean(vr_rea),np.mean(vr_tot),c='w',s=10,zorder=4)
+ ax[0,1].set_aspect('equal', 'box')
+
+ color = color2
+ vr_rea = vr_rea1
+ vr_tot = vr_tot3
+ m, b = np.polyfit(vr_rea, vr_tot, 1)
+ ax[1,0].plot([-1000,1000],[-1000,1000],'k--',zorder=0)
+ ax[1,0].scatter(vr_rea,vr_tot,c=color,alpha=alpha2,s=5,zorder=1,label=llabel2)
+ ax[1,0].plot(vr_rea, m*np.array(vr_rea) + b,'k',zorder=2)
+ ax[1,0].scatter(np.mean(vr_rea),np.mean(vr_tot),c='w',s=100,zorder=2)
+ ax[1,0].scatter(np.mean(vr_rea),np.mean(vr_tot),c='k',s=50,zorder=3)
+ ax[1,0].scatter(np.mean(vr_rea),np.mean(vr_tot),c='w',s=10,zorder=4)
+ ax[1,0].set_aspect('equal', 'box')
+ ax[1,0].legend(loc='lower center')
+
+
+ vr_rea = vr_rea2
+ vr_tot = vr_tot4
+ m, b = np.polyfit(vr_rea, vr_tot, 1)
+ ax[1,1].plot([-1000,1000],[-1000,1000],'k--',zorder=0)
+ ax[1,1].scatter(vr_rea,vr_tot,c=color,alpha=alpha2,s=5,zorder=1)
+ ax[1,1].plot(vr_rea, m*np.array(vr_rea) + b,'k',zorder=2)
+ ax[1,1].scatter(np.mean(vr_rea),np.mean(vr_tot),c='w',s=100,zorder=2)
+ ax[1,1].scatter(np.mean(vr_rea),np.mean(vr_tot),c='k',s=50,zorder=3)
+ ax[1,1].scatter(np.mean(vr_rea),np.mean(vr_tot),c='w',s=10,zorder=4)
+ ax[1,1].set_aspect('equal', 'box')
+
+
+ plt.subplots_adjust(wspace=0.05,hspace=0.05)
+
+ ax[1,0].set_xlabel('variance reduction')
+ ax[1,1].set_xlabel('variance reduction')
+ ax[1,0].set_ylabel('estimated var reduction')
+ ax[0,0].set_ylabel('estimated var reduction')
+
+ ax[0,0].set_title(label1)
+ ax[0,1].set_title(label2)
+
+ x_max = np.max(np.hstack([vr_rea1,vr_rea2,vr_tot1,vr_tot2]))
+ x_min = np.min(np.hstack([vr_rea1,vr_rea2,vr_tot1,vr_tot2]))
+
+
+ ax[0,0].set_xlim(x_min,x_max)
+ ax[0,0].set_ylim(x_min,x_max)
+ plt.locator_params(axis='y', nbins=4)
+ plt.locator_params(axis='x', nbins=4)
+ return fig, ax
\ No newline at end of file
diff --git a/LinAdvFunc/setup.py b/LinAdvFunc/setup.py
new file mode 100644
index 0000000000000000000000000000000000000000..913bca279ce54afdcafe81e00d4c81a229cebd4d
--- /dev/null
+++ b/LinAdvFunc/setup.py
@@ -0,0 +1,7 @@
+from setuptools import setup, find_packages
+
+setup(
+ name="LinAdvFunc",
+ version="0.1",
+ packages=find_packages(),
+)