From bf2dba8bf0de61873e863df0a18c34e45876d1c1 Mon Sep 17 00:00:00 2001
From: Philipp Griewank <philipp.griewank@uni-koeln.de>
Date: Fri, 18 Feb 2022 15:42:37 +0100
Subject: [PATCH] Satelitte data is here!
This is a very big change, that now allows having both point and satellite observations. A other very big change is how things are randomized. Now there are specific seed values for all aspects that can contain randomization.Things are still very preliminary, and instead of updating the old functions I created new ones labelled _22. Importantly, the new functions are not backwards compatible! Variance reduction code is also still only preliminary. I also dropped support of the modified shallow water model for now.
---
da_functions.py | 1543 +++++++++++++++++++++++++++++++++++++++++---
model_functions.py | 28 +
plot_functions.py | 320 +++++++++
3 files changed, 1807 insertions(+), 84 deletions(-)
diff --git a/da_functions.py b/da_functions.py
index 9194800..1e14c26 100644
--- a/da_functions.py
+++ b/da_functions.py
@@ -4,6 +4,93 @@
import numpy as np
from model_functions import *
+from misc_functions import *
+
+
+
+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.25,used_std_obs_sat=0.25,pert_std_obs_sat=0.25,
+ 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.,
@@ -140,6 +227,7 @@ def create_states_dict(j,states,m_const,da_const):
return DA_const
+
def generate_obs(truth,states,m_const,da_const,j,t):
"""
Generates the truth and observations for the next assimilation step.
@@ -168,6 +256,42 @@ def generate_obs(truth,states,m_const,da_const,j,t):
+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"])
+ else:
+ obs_sat = np.zeros(m_const['nx'])
+
+ return truth, obs, obs_sat
+
+
+
def predict(analysis,states,da_const,m_const,j):
@@ -187,6 +311,23 @@ def predict(analysis,states,da_const,m_const,j):
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
+
@@ -329,6 +470,7 @@ def run_linear_advection_KF(m_const,da_const):
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':
@@ -343,6 +485,147 @@ def run_linear_advection_KF(m_const,da_const):
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)
+ 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.
@@ -474,6 +757,44 @@ def sum_mid_tri(x):
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):
"""
@@ -710,9 +1031,9 @@ def var_reduction_estimate_iterative(states,m_const,da_const,j=0,dJ_update_flag=
dxJ = dxJ - alpha*np.outer(K,HdxJ)
dx = dxJ[:-1,:]
- old_var_dJ = np.var(dJ)
+ old_var_dJ = np.var(dJ,ddof=1)
dJ = dxJ[-1,:]
- new_var_dJ = np.var(dJ)
+ 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
@@ -722,7 +1043,7 @@ def var_reduction_estimate_iterative(states,m_const,da_const,j=0,dJ_update_flag=
if dJ_update_flag==2 or dJ_update_flag==3:
if dJ_update_flag==2:
- var_scaling=(np.var(dJ)+vr_ind_tmp[ind_min])/np.var(dJ)
+ 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
@@ -788,93 +1109,101 @@ def gaspari_cohn_non_mirrored(x,loc_length):
return gp
-def single_step_analysis_forecast(state,da_const,m_const,j,t):
- """
- 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.
+# 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.
+# 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
+# """
+# 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: np.random.seed((j+1)*(t+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)
- 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"])
-
+# #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"]])
- """
- EnKF
- """
- if da_const['method'] == 'EnKF':
- # Compute the background error covariance matrix
- P = np.cov(state[j]['bg'][t])*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
+# #if da_const["fixed_seed"]==True:
+# #Changed to the following
+# np.random.seed((j+1)*(t+1)*(seed_add+1))
- # 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)
+# #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"])
-
- """
- 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"])
-
+# obs = state[j]['truth'][t] + np.random.normal(0,da_const["True_std_obs"],m_const["nx"])
- """
- 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
+# #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):
"""
@@ -1066,7 +1395,9 @@ def single_step_analysis_forecast_v2(background,truth,da_const,m_const,model_see
#Generate new truth constants and integrate in time
if m_const['model']=='LA':
- if da_const["fixed_seed"]==True: np.random.seed(model_seed)
+ 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"])
@@ -1140,7 +1471,7 @@ def single_step_analysis_forecast_v2(background,truth,da_const,m_const,model_see
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=1,explicit_sens_flag = 1,exp_num=0):
+ 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.
@@ -1179,7 +1510,7 @@ def vr_reloaded(background,truth,m_const,da_const,func_J=sum_mid_tri,
#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)
+ 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.
@@ -1436,8 +1767,1052 @@ def vr_reloaded(background,truth,m_const,da_const,func_J=sum_mid_tri,
J_fc= fc_response
dJ_fc = J_fc-np.mean(J_fc)
- real_reduction=np.var(dJ_fc) - np.var(dJ_orig)
+ 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 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
+ """
+ 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)
+ 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([n_x,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"])
+ 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)
+
+ #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)
+
+ 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"])
+
+
+ """
+ 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_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=None,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)
+
+ #
+ #
diff --git a/model_functions.py b/model_functions.py
index 26505af..c6b6ccd 100644
--- a/model_functions.py
+++ b/model_functions.py
@@ -3,6 +3,34 @@
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.
diff --git a/plot_functions.py b/plot_functions.py
index ae06f44..a14d718 100644
--- a/plot_functions.py
+++ b/plot_functions.py
@@ -6,6 +6,8 @@ import seaborn as sns
import matplotlib.pyplot as plt
import numpy as np
+from 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.
@@ -119,6 +121,67 @@ def quad_plotter(quad_state,m_const,da_const):
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 B_plotter(states,j=0,ncyc=0,matrix='bg'):
"""
@@ -314,3 +377,260 @@ def _wedge_dir(vertices, direction):
# 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=(8,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
+ """
+ #ax[0].plot(m_const['x_grid'],x_ol_a,'r')
+ #ax[0].plot(m_const['x_grid'],x_ol_a_loc[:],'b')
+ #ax[0].plot(m_const['x_grid'],states[0]['truth'][t],'k')
+
+ 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'],an [:,i],'magenta',alpha=0.2)
+ ax[0,0].plot(m_const['x_grid'],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,xpixmax,
+ color='magenta',alpha=0.2,lw=3)
+ ax[0,1].hlines(
+ sat_bg [da_const["obs_loc_sat"],i],
+ xpixmin,xpixmax,
+ color='r',alpha=0.2,lw=3)
+
+ #ax[0].plot(m_const['x_grid'],x_ol_a,'r')
+ ax[1,0].plot(m_const['x_grid'],truth,'k',linewidth=2)
+ ax[0,0].plot(m_const['x_grid'],truth,'k',linewidth=2)
+ #ax[0,1].hlines(
+ # sat_tr_bg [da_const["obs_loc_sat"]],
+ # xpixmin,xpixmax,
+ # color='k')
+ #ax[1,1].plot(m_const['x_grid'],sat_tr_bg,'k',linewidth=2)
+ #ax[0,1].plot(m_const['x_grid'],sat_tr_bg,'k',linewidth=2)
+
+
+ if n_obs_h>0:
+ ax[1,0].scatter(m_const['x_grid'][da_const['obs_loc']],obs[da_const['obs_loc']],c='k')
+ ax[0,0].scatter(m_const['x_grid'][da_const['obs_loc']],obs[da_const['obs_loc']],c='k')
+ if n_obs_sat>0:
+ ax[1,1].scatter(m_const['x_grid'][da_const['obs_loc_sat']],obs_sat[da_const['obs_loc_sat']],marker='x',s=50,c='k')
+ ax[0,1].scatter(m_const['x_grid'][da_const['obs_loc_sat']],obs_sat[da_const['obs_loc_sat']],marker='x',s=50,c='k')
+
+ ax[0,1].plot(m_const['x_grid'],np.mean(sat_bg,axis=1),'k--',lw=2)
+ #ax[1,1].plot(m_const['x_grid'],np.mean(sat_an,axis=1),'magenta',lw=2)
+ #ax[0,1].plot(m_const['x_grid'],np.mean(sat_bg,axis=1),'r',lw=2)
+ ax[1,1].plot(m_const['x_grid'],np.mean(sat_an,axis=1),'k--',lw=2)
+
+ ax[0,0].plot(m_const['x_grid'],np.mean(bg,axis=1),'k--',lw=2)
+# ax[0,0].plot(m_const['x_grid'],np.mean(an,axis=1),'magenta',lw=2)
+ ax[1,0].plot(m_const['x_grid'],np.mean(an,axis=1),'k--',lw=2)
+ #ax[1,0].plot(m_const['x_grid'],np.mean(bg,axis=1),'r',lw=2)
+
+ if h_c is not None:
+ ax[0,0].hlines(h_c,m_const['x_grid'][0],m_const['x_grid'][-1],color='k',ls=':')
+ ax[1,0].hlines(h_c,m_const['x_grid'][0],m_const['x_grid'][-1],color='k',ls=':')
+
+
+ ax[0,0].set_title('height h')
+ ax[0,1].set_title('reflectance')
+ ax[1,0].set_xlabel('x [m]')
+ ax[1,1].set_xlabel('x [m]')
+ ax[1,1].set_xlim(m_const['x_grid'][0],m_const['x_grid'][-1])
+ #ax[1,0].set_ylabel('Analysis \n h [m]')
+ #ax[0,1].set_ylabel('Background \n cloud fraction')
+ ax[0,0].set_ylabel('Background')
+ ax[1,0].set_ylabel('Analysis')
+ #plt.subplots_adjust(wspace=0.1,hspace=0.1)
+ return fig, ax
\ No newline at end of file
--
GitLab