def forecast(self):
"""
Forecast step of the filter.
Use the model object to propagate each ensemble member to the end of the given checkpoints to
produce and ensemble of forecasts.
Filter configurations dictionary is updated with the new results.
If the prior is assumed to be a Gaussian, we are all set, otherwise we have to estimate it's
parameters based on the provided forecast ensemble (e.g. in the case of 'prior_distribution' = GMM ).
Args:
None
Returns:
None
"""
# generate the forecast states
analysis_ensemble = self.filter_configs['analysis_ensemble']
timespan = self.filter_configs['timespan']
forecast_ensemble = utility.propagate_ensemble(
ensemble=analysis_ensemble,
model=self.filter_configs['model'],
checkpoints=timespan,
in_place=False)
self.filter_configs['forecast_ensemble'] = forecast_ensemble
forecast_state = utility.ensemble_mean(forecast_ensemble)
self.filter_configs['forecast_state'] = forecast_state.copy()
# Add more functionality after building the forecast ensemble if desired!
# Obtain information about the prior distribution from the forecast ensemble
# Covariance localization and hybridization are carried out if requested
self.generate_prior_info()
def get_model_info(timespan,
ensemble_size=25,
model=None,
load_ensemble=False,
ignore_existing_repo=False,
repo_file=None):
"""
"""
timespan = np.array(timespan).flatten()
# print("Requesting timespan: ", timespan)
# Check the experiment repository
if repo_file is None:
file_name = 'Coupled_Lorenz_Repository.pickle'
file_name = os.path.abspath(
os.path.join(os.path.dirname(__file__), file_name))
else:
file_name = os.path.abspath(repo_file)
# Check repo if needed, and found:
if ignore_existing_repo:
ref_IC_np = ref_trajectory_np = init_ensemble_np = observations_np = None
#
elif os.path.isfile(file_name):
print("Loading model repo from : %s" % file_name)
cont = pickle.load(open(file_name, 'rb'))
tspan = cont['tspan']
# Get the initial condition anyways
ref_IC_np = cont['ref_trajectory'][:, 0].copy()
# Verify timespan, and get reference trajectory
if tspan.size < timespan.size:
ref_trajectory_np = None
t_indexes = None
save_info = True
#
elif tspan.size > timespan.size:
# Look into the reference trajectory to find matches, if existing
t_indexes = []
for i, t in enumerate(timespan):
loc = np.where(np.isclose(tspan, t, rtol=1e-12))[0]
if loc.size > 0:
loc = loc[0]
t_indexes.append(loc)
else:
t_indexes = None
break
#
if t_indexes is not None:
t_indexes.sort()
t_indexes = np.asarray(t_indexes)
ref_trajectory_np = cont['ref_trajectory'][:, t_indexes]
else:
ref_trajectory_np = None
#
else:
if np.isclose(tspan, timespan).all():
ref_trajectory_np = cont['ref_trajectory']
t_indexes = np.arange(timespan.size)
#
# Observations:
if t_indexes is None or ref_trajectory_np is None: # TODO: this is overkill!
observations_np = None
else:
observations_np = cont['observations'][:, t_indexes]
# Read the ensemble
if load_ensemble:
init_ensemble_np = cont['init_ensemble']
if np.size(init_ensemble_np, 1) <= ensemble_size:
init_ensemble_np = init_ensemble_np[:, :ensemble_size]
else:
init_ensemble_np = None
else:
init_ensemble_np = None
else:
print(
"The coupled-lorenz model repo is not found. This should be the first time to run this function/script."
)
ref_IC_np = ref_trajectory_np = init_ensemble_np = observations_np = None
# Check the integrity of the loaded data:
if ref_trajectory_np is None or (load_ensemble and init_ensemble_np is None
) or observations_np is None:
#
# Generate stuff, and save them
if model is None:
print(
"YOU MUST pass a MODEL object because the repo is not found in %s, or some information is missing!"
% file_name)
raise ValueError
else:
print("Got the model >> Creating experiemnt info...")
# Reference IC
if ref_IC_np is None:
ref_IC = model._reference_initial_condition.copy()
ref_IC_np = ref_IC.get_numpy_array()
# Reference Trajectory (Truth)
if ref_trajectory_np is None:
if ref_IC_np is not None:
ref_IC = model.state_vector(ref_IC_np.copy())
else:
ref_IC = model._reference_initial_condition.copy()
# generate observations from coupled lorenz:
ref_trajectory = model.integrate_state(ref_IC, timespan)
ref_trajectory_np = utility.ensemble_to_np_array(ref_trajectory,
state_as_col=True)
# Initial (forecast) Ensmble:
if load_ensemble and init_ensemble_np is None:
prior_noise_model = model.background_error_model
prior_noise_sample = [
prior_noise_model.generate_noise_vec()
for i in xrange(ensemble_size)
]
noise_mean = utility.ensemble_mean(prior_noise_sample)
ic = prior_noise_model.generate_noise_vec().add(ref_IC)
for i in xrange(ensemble_size):
prior_noise_sample[i].axpy(-1.0, noise_mean)
prior_noise_sample[i].add(ic, in_place=True)
init_ensemble = prior_noise_sample
init_ensemble_np = utility.ensemble_to_np_array(init_ensemble,
state_as_col=True)
if observations_np is None:
# Create observations' and assimilation checkpoints:
if ref_trajectory is None:
if ref_trajectory_np is None:
ref_trajectory = model.integrate_state(ref_IC, timespan)
else:
ref_trajectory = [
model.state_vectory(ref_trajectory_np[i, :].copy())
for i in xrange(np.size(ref_trajectory, 1))
]
observations = [
model.evaluate_theoretical_observation(x)
for x in ref_trajectory
]
# Perturb observations:
obs_noise_model = model.observation_error_model
for obs_ind in xrange(len(timespan)):
observations[obs_ind].add(obs_noise_model.generate_noise_vec(),
in_place=True)
observations_np = utility.ensemble_to_np_array(observations,
state_as_col=True)
print("Experiment repo created; saving for later use...")
# save results for later use
out_dict = dict(tspan=timespan,
ref_trajectory=ref_trajectory_np,
init_ensemble=init_ensemble_np,
observations=observations_np)
pickle.dump(out_dict, open(file_name, 'wb'))
print("...done...")
#
else:
print("All information properly loaded from file")
pass
return timespan, ref_IC_np, ref_trajectory_np, init_ensemble_np, observations_np
def filtering_cycle(self, update_reference=False):
"""
Carry out a single assimilation cycle. This is just the analysis in this case
All required variables are obtained from 'filter_configs' dictionary.
Args:
update_reference (default True): bool,
A flag to decide whether to update the reference state in the filter or not.
This should be removed!
"""
#
model = self.model
try:
observation_time = self.filter_configs['observation_time']
except (NameError, KeyError, AttributeError):
observation_time = None
try:
forecast_time = self.filter_configs['forecast_time']
except (NameError, KeyError, AttributeError):
forecast_time = None
try:
analysis_time = self.filter_configs['analysis_time']
except (NameError, KeyError, AttributeError):
analysis_time = None
try:
reference_time = self.filter_configs['reference_time']
except (NameError, KeyError, AttributeError):
reference_time = None
#
try:
# Retrieve the reference state and evaluate initial root-mean-squared error
reference_state = self.filter_configs['reference_state'].copy()
except (NameError, KeyError, AttributeError):
reference_state = None
#
if self._verbose:
print("Couldn't retrieve the reference state/time! ")
print(
"True observations must be present in this case, and RMSE can't be evaluated!"
)
if reference_state is not None and reference_time is not None:
# If the reference time is below window final limit, march it forward:
if analysis_time is not None:
new_time = analysis_time
elif forecast_time is not None:
new_time = forecast_time
elif reference_time is not None:
new_time = reference_time
if (new_time - reference_time) > self._time_eps:
# propagate forward the reference state, and update it
local_trajectory = model.integrate_state(
reference_state, [reference_time, new_time])
reference_time = new_time
if isinstance(local_trajectory, list):
reference_state = local_trajectory[-1]
else:
reference_state = local_trajectory
elif (reference_time - new_time) > self._time_eps:
print(
"Can't have the reference time after the forecast/analysis time!"
)
raise ValueError
if update_reference:
self.filter_configs.update(
{'reference_state': reference_state})
self.filter_configs.update({'reference_time': reference_time})
# Check time instances:
times_array = np.array(
[observation_time, analysis_time, forecast_time, reference_time])
val_inds = np.where(times_array)[0]
if len(val_inds) > 0:
if not np.allclose(times_array[val_inds],
times_array[val_inds[0]],
atol=self._time_eps,
rtol=self._time_eps):
print(
"Given time instances, e.g. forecast time, observation time, etc. MUST match! "
)
print(
"[observation_time, analysis_time, forecast_time, reference_time]",
[
observation_time, analysis_time, forecast_time,
reference_time
])
raise AssertionError
else:
# Good to go!
pass
else:
# No times are given; just proceed
pass
#
try:
forecast_state = self.filter_configs['forecast_state'].copy()
except (NameError, KeyError, AttributeError):
print(
"Couldn't retrieve the forecast state! This can't work; Aborting!"
)
raise AssertionError
#
try:
observation = self.filter_configs['observation']
except (NameError, KeyError, AttributeError):
observation = None
finally:
if observation is None:
print("The observation is not found!")
raise AssertionError
#
# Forecast state/ensemble should be given:
if forecast_state is None:
try:
forecast_state = utility.ensemble_mean(
self.filter_configs['forecast_ensemble'])
except:
print(
"Couldn't find either forecast state or forecast ensemble while analysis should be done first!"
)
raise AssertionError
#
# calculate the initial/forecast RMSE: this is the RMSE before assimilation at the first entry of the time-span
state_size = model.state_size()
#
if forecast_state is not None and reference_state is not None:
f_rmse = initial_rmse = utility.calculate_rmse(
forecast_state, reference_state, state_size)
else:
f_rmse = initial_rmse = 0
#
# Start the filtering process: preprocessing -> filtering(forecast->+<-anslsysis) -> postprocessing
if self.filter_configs['apply_preprocessing']:
self.cycle_preprocessing()
#
# Analysis step (This calls the filter's analysis step.)
# > --------------------------------------------||
sep = "\n" + "*" * 80 + "\n"
if self._verbose:
print("%s ANALYSIS STEP %s" % (sep, sep))
self.analysis()
# > --------------------------------------------||
#
analysis_state = self.filter_configs['analysis_state'].copy()
# Analysis RMSE:
if reference_state is not None:
a_rmse = utility.calculate_rmse(reference_state, analysis_state,
state_size)
else:
a_rmse = 0
#
if True:
print("Initial (f_rmse, a_rmse) :", (f_rmse, a_rmse))
print("analysis_state", analysis_state)
print("forecast_state", forecast_state)
print("reference_state", reference_state)
#
# Apply post-processing if required
if self.filter_configs['apply_postprocessing']:
self.cycle_postprocessing()
#
# Update filter statistics (including RMSE)
if 'filter_statistics' not in self.output_configs:
self.output_configs.update(
dict(filter_statistics=dict(initial_rmse=None,
forecast_rmse=None,
analysis_rmse=None)))
else:
if 'analysis_rmse' not in self.output_configs['filter_statistics']:
self.output_configs['filter_statistics'].update(
dict(analysis_rmse=None))
if 'forecast_rmse' not in self.output_configs['filter_statistics']:
self.output_configs['filter_statistics'].update(
dict(forecast_rmse=None))
if 'initial_rmse' not in self.output_configs['filter_statistics']:
self.output_configs['filter_statistics'].update(
dict(initial_rmse=None))
# now update the RMSE's
self.output_configs['filter_statistics'].update(
{'initial_rmse': initial_rmse})
self.output_configs['filter_statistics'].update(
{'forecast_rmse': f_rmse})
#
self.output_configs['filter_statistics'].update(
{'analysis_rmse': a_rmse})
# output and save results if requested
if self.output_configs['scr_output']:
self.print_cycle_results()
if self.output_configs['file_output']:
self.save_cycle_results()
def analysis(self):
"""
Analysis step:
"""
#
model = self.filter_configs['model']
state_size = model.state_size()
#
# Check if the forecast ensemble should be inflated;
f = self.filter_configs['inflation_factor']
forecast_ensemble = utility.inflate_ensemble(
self.filter_configs['forecast_ensemble'], f, in_place=False)
forecast_state = utility.ensemble_mean(forecast_ensemble)
ensemble_size = len(forecast_ensemble)
# get the measurements vector
observation = self.filter_configs['observation']
observation_size = model.observation_vector_size()
#
normalized_weights = None
self.filter_configs.update({'normalized_weights': normalized_weights})
# Evaluate the likelihood for all particles
particles_likelihood = np.asarray(
self.evaluate_likelihood(model_states=forecast_ensemble,
observation=observation))
likelihood_sum = particles_likelihood.sum()
if likelihood_sum == 0: # This should not happen as long as all particles are passed to the likelihood
# print particles_likelihood
print(
'All particles received negligible weights rounded down to zero!. Filter diverged; Ensemble is not '
'updated!')
return None
else:
normalized_weights = particles_likelihood / likelihood_sum
# Monitors the effective particle number (effective sample size)
ess = 1.0 / (np.asarray(normalized_weights)**2).sum()
self.filter_configs['effective_sample_size'] = ess
self.filter_configs.update({
'particles_likelihood': particles_likelihood,
'normalized_weights': normalized_weights
})
# Check the resampling flag and resampling scheme
# Resampling step:
if self._resample:
resampling_scheme = self.filter_configs['resampling_scheme']
normalize_weights = self.filter_configs['normalize_weights']
resampling_indexes = self.resample_states(
weights=normalized_weights,
ensemble_size=ensemble_size,
strategy=resampling_scheme,
normalize_weights=normalize_weights)
# Generate Analysis ensemble by weighted resampling
print("Resampling...")
if self.filter_configs['analysis_ensemble'] is None:
analysis_ensemble = []
for ens_ind in xrange(ensemble_size):
res_ind = resampling_indexes[ens_ind]
# print("Ens ind: %d ; resampling from index %d" %(ens_ind, res_ind))
analysis_ensemble.append(forecast_ensemble[res_ind].copy())
else:
analysis_ensemble = self.filter_configs['analysis_ensemble']
for ens_ind in xrange(ensemble_size):
res_ind = resampling_indexes[ens_ind]
# print("Ens ind: %d ; resampling from index %d" %(ens_ind, res_ind))
analysis_ensemble[ens_ind] = forecast_ensemble[
res_ind].copy()
#
else:
resampling_indexes = np.arange(ensemble_size)
analysis_ensemble = [member.copy() for member in forecast_ensemble]
self.filter_configs['resampling_indexes'] = resampling_indexes
self.filter_configs['analysis_ensemble'] = analysis_ensemble
self.filter_configs['analysis_state'] = utility.ensemble_mean(
analysis_ensemble)
def filtering_cycle(self, update_reference=True):
"""
Carry out a single assimilation cycle. Forecast followed by analysis or the other way around.
All required variables are obtained from 'filter_configs' dictionary.
This base method is designed to work in both ensemble, and standard framework.
You can override it for sure in your filter.
Args:
update_reference (default True): bool,
A flag to decide whether to update the reference state in the filter or not.
Returns:
None
"""
model = self.filter_configs['model']
# print('reference_state', reference_state)
timespan = self.filter_configs['timespan']
observation_time = self.filter_configs['observation_time']
forecast_time = self.filter_configs['forecast_time']
analysis_time = self.filter_configs['analysis_time']
forecast_first = self.filter_configs['forecast_first']
apply_preprocessing = self.filter_configs['apply_preprocessing']
apply_postprocessing = self.filter_configs['apply_postprocessing']
# calculate the initial RMSE: this is the RMSE before assimilation at the first entry of the time-span
try:
reference_time = self.filter_configs['reference_time']
except (KeyError, ValueError, AttributeError, NameError):
raise ValueError(
"Couldn't find the reference time in the configurations dictionary"
)
else:
if reference_time is not None:
if reference_time != timespan[0]:
raise ValueError(
"Reference time does not match the initial time of the reference time-span!"
)
#
#
if forecast_first:
# Analysis ensemble should be given first in this case
try:
initial_state = self.filter_configs['analysis_state']
except (KeyError, ValueError, AttributeError, NameError):
try:
initial_state = utility.ensemble_mean(
self.filter_configs['analysis_ensemble'])
except:
raise ValueError(
"Couldn't find either analysis state or analysis ensemble while forecast should be done first!"
)
finally:
if initial_state is None:
try:
initial_state = utility.ensemble_mean(
self.filter_configs['analysis_ensemble'])
except:
raise ValueError(
"Couldn't find either analysis state or analysis ensemble while forecast should be done first!"
)
else:
# Forecast ensemble should be given first in this case
try:
initial_state = self.filter_configs['forecast_state']
except:
try:
initial_state = utility.ensemble_mean(
self.filter_configs['forecast_ensemble'])
except:
raise ValueError(
"Couldn't find either forecast state or forecast ensemble while analysis should be done first!"
)
if initial_state is None:
try:
initial_state = utility.ensemble_mean(
self.filter_configs['forecast_ensemble'])
except:
raise ValueError(
"Couldn't find either forecast state or forecast ensemble while analysis should be done first!"
)
# Retrieve the reference state and evaluate initial root-mean-squared error
reference_state = self.filter_configs[
'reference_state'] # always given at the initial time of the timespan
if reference_state is not None:
reference_time = self.filter_configs['reference_time']
try:
initial_rmse = utility.calculate_rmse(initial_state,
reference_state)
except:
reference_state = None
reference_time = None
print("Failed to calculate Initial RMSE!")
initial_rmse = 0.0
#
# Start the filtering process: preprocessing -> filtering(forecast->+<-anslsysis) -> postprocessing
if apply_preprocessing:
self.cycle_preprocessing()
if not forecast_first and analysis_time != min(
forecast_first, analysis_time, observation_time):
# this is a double check!
raise ValueError(
"While ANALYSIS should be done first, confusion occurred with times given!\n"
"\tCycle timespan:%s"
"\tObservation time: %f\n"
"\tForecast time: %f\n"
"\tAnalysis time: %f" % (repr(timespan), observation_time,
forecast_time, analysis_time))
elif not forecast_first:
#
# print("\n\n\n\n\n\n ANALYSIS FIRTS \n\n\n\n\n")
# analysis should be carried out first in this case
state_size = self.filter_configs['model'].state_size()
if reference_state is not None:
try:
analysis_rmse = utility.calculate_rmse(
reference_state, analysis_state, state_size)
except:
print("Failed to calculate Analysis RMSE!")
analysis_rmse = 0.0
analysis_time = self.filter_configs['analysis_time']
# Analysis step
self.analysis()
#
# update the reference state
if reference_time is not None:
if reference_time != timespan[
0] or reference_time != analysis_time:
raise ValueError(
"Reference time does not match the initial time of the reference time-span!"
)
local_checkpoints = [analysis_time, forecast_time]
if reference_state is not None:
tmp_trajectory = model.integrate_state(
initial_state=reference_state,
checkpoints=local_checkpoints)
if isinstance(tmp_trajectory, list):
reference_state = tmp_trajectory[-1].copy()
else:
reference_state = tmp_trajectory.copy()
reference_time = local_checkpoints[-1]
# forecast now:
self.forecast()
try:
forecast_rmse = utility.calculate_rmse(reference_state,
forecast_state,
state_size)
except:
print("Failed to calculate Forecast RMSE!")
forecast_rmse = 0.0
# update the reference state Moved to the process
if update_reference and reference_state is not None:
self.filter_configs['reference_state'] = reference_state.copy()
self.filter_configs['reference_time'] = reference_time
else:
# forecast should be done first
# print("\n\n\n\n\n\n FORECAST FIRTS \n\n\n\n\n")
state_size = self.filter_configs['model'].state_size()
if reference_time is not None:
if reference_time != timespan[0]:
raise ValueError(
"Reference time does not match the initial time of the reference time-span!"
)
local_checkpoints = [reference_time, timespan[-1]]
# Forecast step:
self.forecast()
#
try:
forecast_state = self.filter_configs['forecast_state']
except (NameError, AttributeError):
raise NameError("forecast_state must be updated by the filter "
"and added to 'self.filter_configs'!")
if reference_state is not None:
tmp_trajectory = model.integrate_state(
initial_state=reference_state, checkpoints=timespan)
if isinstance(tmp_trajectory, list):
up_reference_state = tmp_trajectory[-1].copy()
else:
up_reference_state = tmp_trajectory.copy()
reference_time = local_checkpoints[-1]
# Update the reference state
if update_reference:
self.filter_configs[
'reference_state'] = up_reference_state.copy()
self.filter_configs['reference_time'] = reference_time
# observation = self.filter_configs['observation']
#
try:
forecast_rmse = utility.calculate_rmse(up_reference_state,
forecast_state,
state_size)
except:
print("Failed to calculate Forecast RMSE!")
forecast_rmse = 0.0
# Analysis step:
self.analysis()
#
try:
analysis_state = self.filter_configs['analysis_state']
except (NameError, AttributeError):
raise NameError("analysis_state must be updated by the filter "
"and added to 'self.filter_configs'!")
try:
analysis_rmse = utility.calculate_rmse(up_reference_state,
analysis_state,
state_size)
except:
print("Failed to calculate Analysis RMSE!")
analysis_rmse = 0.0
#
# Apply post-processing if required
if apply_postprocessing:
self.cycle_postprocessing()
# Update filter statistics (including RMSE)
if 'filter_statistics' not in self.output_configs:
self.output_configs.update(
dict(filter_statistics=dict(initial_rmse=None,
forecast_rmse=None,
analysis_rmse=None)))
else:
if 'analysis_rmse' not in self.output_configs['filter_statistics']:
self.output_configs['filter_statistics'].update(
dict(analysis_rmse=None))
if 'forecast_rmse' not in self.output_configs['filter_statistics']:
self.output_configs['filter_statistics'].update(
dict(forecast_rmse=None))
if 'initial_rmse' not in self.output_configs['filter_statistics']:
self.output_configs['filter_statistics'].update(
dict(initial_rmse=None))
# now update the RMSE's
self.output_configs['filter_statistics']['initial_rmse'] = initial_rmse
self.output_configs['filter_statistics'][
'forecast_rmse'] = forecast_rmse
self.output_configs['filter_statistics'][
'analysis_rmse'] = analysis_rmse
# output and save results if requested
if self.output_configs['scr_output']:
self.print_cycle_results()
if self.output_configs['file_output']:
self.save_cycle_results()
def smoothing_cycle(self, update_reference=True):
"""
Carry out a single assimilation cycle. The analysis has to be carried out first given the forecast
at the assimilation time, and the observations (Real observations), or synthetic created using
the reference state of the model.
All required variables are obtained from 'smoother_configs' dictionary.
This base method is designed to work in both ensemble, and variational framework.
You can override it for sure in your smoother.
Args:
update_reference (default True): bool,
A flag to decide whether to update the reference state in the smoother or not.
"""
#
# > ===========================( Assimilation Cycle Starts )================================ <
#
model = self.model
#
try:
# Retrieve the reference state and evaluate initial root-mean-squared error
reference_state = self.smoother_configs['reference_state'].copy()
reference_time = self.smoother_configs['reference_time']
except (NameError, KeyError, AttributeError):
reference_state = reference_time = None
finally:
if reference_state is None or reference_time is None:
reference_state = reference_time = None
if self._verbose:
print("Couldn't retrieve the reference state/time! ")
print(
"True observations must be present in this case, and RMSE can't be evaluated!"
)
#
try:
forecast_state = self.smoother_configs['forecast_state'].copy()
except (NameError, KeyError, AttributeError):
print(
"Couldn't retrieve the forecast state! This can't work; Aborting!"
)
raise AssertionError
#
try:
observations_list = self.smoother_configs['observations_list']
except (NameError, KeyError, AttributeError):
observations_list = None
finally:
if observations_list is None:
print(
"Neither the reference state is passed, nor true observations list is found!"
)
raise AssertionError
#
try:
obs_checkpoints = self.smoother_configs['obs_checkpoints']
except:
print(
"Either analysis_time or obs_checkpoints or both is/are missing! Aborting!"
)
try:
analysis_timespan = self.smoother_configs['analysis_timespan']
except (NameError, KeyError, AttributeError):
analysis_timespan = None
finally:
if analysis_timespan is None:
obs_checkpoints = self.smoother_configs['obs_checkpoints']
bounds = self.smoother_configs['window_bounds']
t0, t1 = bounds[0], bounds[-1]
tf, ta = self.smoother_configs[
'forecast_time'], self.smoother_configs['analysis_time']
analysis_timespan = np.asarray(
list(set.union(set(obs_checkpoints), {t0, t1, tf, ta})))
analysis_timespan.sort()
self.smoother_configs.update(
{'analysis_timespan': analysis_timespan})
#
# Forecast state/ensemble should be given:
if forecast_state is None:
try:
forecast_state = utility.ensemble_mean(
self.smoother_configs['forecast_ensemble'])
except:
print(
"Couldn't find either forecast state or forecast ensemble while analysis should be done first!"
)
raise AssertionError
# Check forecast, and analysis times:
forecast_time = self.smoother_configs['forecast_time']
analysis_time = self.smoother_configs['analysis_time']
#
if forecast_time is not None and analysis_time is not None:
pass
#
elif forecast_time is None and analysis_time is None:
print(
"Both forecast and analysis time are not set. At least specify the forecast time!"
)
raise ValueError
#
elif analysis_time is None:
if analysis_timespan is not None:
analysis_time = analysis_timespan[0]
else:
# Analysis is carried out in the same place where forecast state is provided!
analysis_time = forecast_time
else:
print(
"You have to specify where the provided forecast state is calculated!"
)
raise ValueError
#
if analysis_time > forecast_time:
print(
"Analysis time > forecast time. Forward propagation of the forecast state is taking place now."
)
#
local_checkpoints = [forecast_time, analysis_time]
tmp_trajectory = model.integrate_state(
initial_state=forecast_state, checkpoints=local_checkpoints)
if isinstance(tmp_trajectory, list):
forecast_state = tmp_trajectory[-1].copy()
else:
forecast_state = tmp_trajectory.copy()
forecast_time = analysis_time
self.smoother_configs['forecast_time'] = local_checkpoints[-1]
self.smoother_configs['forecast_state'] = forecast_state
#
elif analysis_time < forecast_time:
print(
"Analysis time < Forecast time! Backward propagation of the state is Needed!"
)
raise AssertionError
#
else:
# All is Good; analysis_time = forecast_time
pass
#
if forecast_time is not None and reference_time is not None:
#
if reference_time < forecast_time:
local_checkpoints = [reference_time, forecast_time]
tmp_trajectory = model.integrate_state(
initial_state=reference_state,
checkpoints=local_checkpoints)
if isinstance(tmp_trajectory, list):
reference_state = tmp_trajectory[-1].copy()
else:
reference_state = tmp_trajectory.copy()
#
self.smoother_configs['reference_state'] = reference_state
self.smoother_configs['reference_time'] = local_checkpoints[-1]
#
elif reference_time == forecast_time:
pass
else:
print("Time settings are not conformable!")
print("Reference time:", reference_time)
print("Forecast time:", forecast_time)
print("Analysis time:", analysis_time)
raise AssertionError
#
#
# calculate the initial RMSE: this is the RMSE before assimilation at the first entry of the time-span
if reference_time is not None:
# print("analysis_timespan", analysis_timespan)
# print("reference_time", reference_time)
if (analysis_timespan[0] - reference_time) > self._time_eps:
print(
"Reference time < initial time of the assimilation time-span. Forward propagation of the reference is taking place now."
)
print("reference_time", reference_time)
print("analysis_timespan[0]", analysis_timespan[0])
#
local_checkpoints = [reference_time, analysis_timespan[0]]
tmp_trajectory = model.integrate_state(
initial_state=reference_state,
checkpoints=local_checkpoints)
if isinstance(tmp_trajectory, list):
reference_state = tmp_trajectory[-1].copy()
else:
reference_state = tmp_trajectory.copy()
reference_time = local_checkpoints[-1]
reference_time = analysis_timespan[0]
self.smoother_configs['reference_time'] = reference_time
self.smoother_configs['reference_state'] = reference_state
#
elif (reference_time - analysis_timespan[0]) > self._time_eps:
print(
"Reference time > the initial time of the assimilation time-span!"
)
raise ValueError
else:
# All good...
pass
else:
self.smoother_configs['reference_time'] = None
self.smoother_configs['reference_state'] = None
#
if abs(analysis_time - analysis_timespan[0]) > self._time_eps:
print(
"*" * 100 +
"\nWARNING: The analysis time is not at the beginning of the assimilation timespan!\n"
+ "*" * 100)
print("analysis_time", analysis_time)
print("analysis_timespan[0]", analysis_timespan[0])
print("\n" + "*" * 100 + "\n")
#
if (analysis_time - obs_checkpoints[0]) > self._time_eps:
print(
"First observation has to be at, or after, the analysis time")
raise AssertionError
if len(obs_checkpoints) != len(observations_list):
print(
"Observaions timespan must of be same length as the observations list!"
)
raise AssertionError
#
if forecast_state is not None and reference_state is not None:
if abs(analysis_time - reference_time) > self._time_eps:
print(
"While evaluating initial RMSE; analysis_time != reference_time"
)
print("reference_time: %f" % reference_time)
print("Analysis time: %f" % analysis_time)
raise ValueError
else:
initial_rmse = utility.calculate_rmse(forecast_state,
reference_state)
else:
initial_rmse = 0
#
# Start the smoothing process: preprocessing -> smoothing(forecast->+<-anslsysis) -> postprocessing
if self.smoother_configs['apply_preprocessing']:
self.cycle_preprocessing()
#
# Analysis step (This calls the smoother's analysis step.)
# > --------------------------------------------||
sep = "\n" + "*" * 80 + "\n"
if self._verbose:
print("%s ANALYSIS STEP %s" % (sep, sep))
self.analysis()
# > --------------------------------------------||
#
# Evaluate RMSEs: This is potential redundancy; just read analysis trajectory from the smoother)
reference_state = self.smoother_configs['reference_state'].copy()
reference_time = self.smoother_configs['reference_time']
#
forecast_state = self.smoother_configs['forecast_state'].copy()
forecast_time = self.smoother_configs['forecast_time']
#
analysis_state = self.smoother_configs['analysis_state'].copy()
analysis_time = self.smoother_configs['analysis_time']
#
analysis_trajectory = []
analysis_trajectory.append(analysis_state)
#
state_size = model.state_size()
# Forecast and analysis RMSE (at the beginning of the assimilation window)
if abs(analysis_time - reference_time) > self._time_eps:
print(
"While evaluating initial RMSE; analysis_time != reference_time"
)
print("reference_time: %f" % reference_time)
print("Analysis time: %f" % analysis_time)
raise ValueError
else:
f_rmse = utility.calculate_rmse(reference_state, forecast_state,
state_size)
forecast_rmse_list = [f_rmse]
forecast_times_list = [forecast_time]
a_rmse = utility.calculate_rmse(reference_state, analysis_state,
state_size)
analysis_rmse_list = [a_rmse]
analysis_times_list = [analysis_time]
#
if self._verbose:
print("Initial (f_rmse, a_rmse) :", (f_rmse, a_rmse))
print("analysis_state", analysis_state)
print("forecast_state", forecast_state)
print("reference_state", reference_state)
#
#
for t0, t1 in zip(analysis_timespan[:-1], analysis_timespan[1:]):
local_checkpoints = np.array([t0, t1])
#
# Propagate forecast
local_trajectory = model.integrate_state(forecast_state,
local_checkpoints)
if isinstance(local_trajectory, list):
forecast_state = local_trajectory[-1].copy()
else:
forecast_state = local_trajectory.copy()
# Propagate analysis
local_trajectory = model.integrate_state(analysis_state,
local_checkpoints)
if isinstance(local_trajectory, list):
analysis_state = local_trajectory[-1].copy()
else:
analysis_state = local_trajectory.copy()
analysis_trajectory.append(analysis_state)
#
#
if reference_state is not None:
if self._verbose:
print("\n Evaluating RMSE, inside smoothers_base...:")
print("Subinterval: ", [t0, t1])
print("reference_time: ", reference_time)
print("reference_state: ", reference_state)
# Propagate reference
local_trajectory = model.integrate_state(
reference_state, [reference_time, t1])
reference_time = t1
if isinstance(local_trajectory, list):
reference_state = local_trajectory[-1].copy()
else:
reference_state = local_trajectory.copy()
#
#
f_rmse = utility.calculate_rmse(reference_state,
forecast_state, state_size)
a_rmse = utility.calculate_rmse(reference_state,
analysis_state, state_size)
forecast_rmse_list.append(f_rmse)
analysis_rmse_list.append(a_rmse)
forecast_times_list.append(t1)
analysis_times_list.append(t1)
#
if self._verbose:
print("forecast_state: ", forecast_state)
print("analysis_state: ", analysis_state)
else:
# No reference state is provided. RMSE will be an empty list
if self._verbose:
print(
"WARNING: No reference state is provided. RMSE will be an empty list!"
)
#
if reference_state is not None and update_reference:
# If the reference time is below window final limit, march it forward:
if (self.smoother_configs['window_bounds'][-1] -
reference_time) > self._time_eps:
tf = self.smoother_configs['window_bounds'][-1]
local_trajectory = model.integrate_state(
reference_state, [reference_time, tf])
reference_time = tf
if isinstance(local_trajectory, list):
reference_state = local_trajectory[-1]
else:
reference_state = local_trajectory
#
self.smoother_configs.update({'reference_state': reference_state})
self.smoother_configs.update({'reference_time': reference_time})
#
#
# Apply post-processing if required
if self.smoother_configs['apply_postprocessing']:
self.cycle_postprocessing()
#
# forecast now (This calls the smoother's forecast step):
# > --------------------------------------------||
if self._verbose:
print("%s ANALYSIS STEP %s" % (sep, sep))
self.forecast()
# > --------------------------------------------||
#
# Update smoother statistics (including RMSE)
if 'smoother_statistics' not in self.output_configs:
self.output_configs.update(
dict(smoother_statistics=dict(initial_rmse=None,
forecast_rmse=None,
analysis_rmse=None)))
else:
if 'analysis_rmse' not in self.output_configs[
'smoother_statistics']:
self.output_configs['smoother_statistics'].update(
dict(analysis_rmse=None))
if 'forecast_rmse' not in self.output_configs[
'smoother_statistics']:
self.output_configs['smoother_statistics'].update(
dict(forecast_rmse=None))
if 'initial_rmse' not in self.output_configs[
'smoother_statistics']:
self.output_configs['smoother_statistics'].update(
dict(initial_rmse=None))
# now update the RMSE's
self.output_configs['smoother_statistics'].update(
{'initial_rmse': initial_rmse})
self.output_configs['smoother_statistics'].update(
{'forecast_rmse': forecast_rmse_list})
self.output_configs['smoother_statistics'].update(
{'forecast_times': forecast_times_list})
#
self.output_configs['smoother_statistics'].update(
{'analysis_rmse': analysis_rmse_list})
self.output_configs['smoother_statistics'].update(
{'analysis_times': analysis_times_list})
# output and save results if requested
if self.output_configs['scr_output']:
self.print_cycle_results()
if self.output_configs['file_output']:
self.save_cycle_results()
def EnKF_update(self, forecast_ensemble, observation, inflation_factor=1.0, in_place=False):
"""
Calculate the kalman updated matrix, i.e. the update matrix of ensemble anomalies
and calculate the (moving) analysis ensemble
Args:
state:
Returns:
analysis_ensemble: a list of model.state_vector objects
X5: the state anomalies update matrix
"""
model = self.smoother_configs['model']
ensemble_size = len(forecast_ensemble)
forecast_mean = utility.ensemble_mean(forecast_ensemble)
forecast_anomalies = []
if inflation_factor is None:
inflation_factor = 1.0
#
for state in forecast_ensemble:
if inflation_factor != 1:
forecast_anomalies.append(state.axpy(-1, forecast_mean, in_place=False).scale(inflation_factor))
else:
forecast_anomalies.append(state.axpy(-1, forecast_mean, in_place=False))
A_prime = utility.ensemble_to_np_array(forecast_anomalies)
observation_perturbations = [model.evaluate_theoretical_observation(state) for state in forecast_anomalies]
S = utility.ensemble_to_np_array(observation_perturbations)
C = S.dot(S.T)
try:
C += (ensemble_size - 1) * model.observation_error_model.R[...]
except:
C += (ensemble_size - 1) * model.observation_error_model.R.get_numpy_array()
C_eigs, Z = linalg.eigh(C)
C_eigs = 1.0 / C_eigs
obs_innov = observation.axpy(-1, model.evaluate_theoretical_observation(forecast_mean))
y = Z.T.dot(obs_innov.get_numpy_array())
y *= C_eigs
y = Z.dot(y)
y = S.T.dot(y)
# Update ensemble mean
analysis_mean = model.state_vector(A_prime.dot(y))
analysis_mean = analysis_mean.add(forecast_mean)
xa = analysis_mean.get_numpy_array()
inv_sqrt_Ceig = sparse.spdiags(np.sqrt(C_eigs), 0, C_eigs.size, C_eigs.size)
X2 = inv_sqrt_Ceig.dot(Z.T.dot(S))
_, Sig2, V2 = linalg.svd(X2, full_matrices=False)
# Calculate the square root of the matrix I-Sig2'*Sig2*theta
Sqrtmat = sparse.spdiags(np.sqrt(1.0 - np.power(Sig2,2)), 0, Sig2.size, Sig2.size)
Aa_prime = A_prime.dot(V2.T.dot(Sqrtmat.dot(V2)))
if self._verbose:
# For debugging
print(A_prime, type(A_prime))
# Update analysis ensemble:
analysis_ensemble = [model.state_vector(xa+A_prime[:, j]) for j in xrange(np.size(A_prime, 1)) ]
# Build the update kernel X5
xx = V2.T.dot(Sqrtmat.dot(V2))
for j in xrange(np.size(xx, 1)):
xx[:, j] += y[:]
IN = np.ones((ensemble_size, ensemble_size), dtype=np.float) / ensemble_size
X5 = IN + inflation_factor * np.dot( (sparse.spdiags(np.ones(ensemble_size), 0, ensemble_size, ensemble_size) - IN), xx)
return analysis_ensemble, X5
def analysis(self):
"""
Analysis step of the (Vanilla Ensemble Kalman Smoother (EnKS).
In this case, the given forecast ensemble is propagated to the observation time instances to create model
observations H(xk) that will be used in the assimilation process...
"""
model = self.model
# Timsespan info:
forecast_time = self.smoother_configs['forecast_time']
analysis_time = self.smoother_configs['analysis_time']
if self.smoother_configs['obs_checkpoints'] is None:
print("Couldn't find observation checkpoints in self.smoother_configs['obs_checkpoints']; None found!")
raise ValueError
else:
obs_checkpoints = np.asarray(self.smoother_configs['obs_checkpoints'])
if (analysis_time - obs_checkpoints[0] ) >= self._model_step_size:
print("Observations MUST follow the assimilation times in this implementation!")
raise ValueError
if (analysis_time - obs_checkpoints[0] ) >= self._model_step_size:
print("Observations MUST follow the assimilation times in this implementation!")
raise ValueError
# Get observations list:
observations_list = self.smoother_configs['observations_list']
#
assim_flags = [True] * len(obs_checkpoints)
if (forecast_time - obs_checkpoints[0]) >= self._model_step_size:
print("forecast time can't be after the first observation time instance!")
print("forecast_time", forecast_time)
print("obs_checkpoints[0]", obs_checkpoints[0])
raise AssertionError
#
elif abs(forecast_time - obs_checkpoints[0]) <= self._time_eps:
# an observation exists at the analysis time
pass
#
else:
obs_checkpoints = np.insert(obs_checkpoints, 0, forecast_time)
assim_flags.insert(0, False)
# state and observation vector dimensions
forecast_ensemble = self.smoother_configs['forecast_ensemble']
if forecast_ensemble is None:
print("Can't carry out the analysis step without a forecast ensemble; Found None!")
raise ValueError
if abs(forecast_time-analysis_time) > self._time_eps:
print("At this point, the analysis and forecast times have to be the SAME!")
print("Forecast Time=%f, Analysis time=%f" % (forecast_time, analysis_time))
raise ValueError
#
elif (analysis_time-obs_checkpoints[0]) > self._time_eps:
print("At this point, the analysis time should be equal to or less that the time of the first observation")
print("Analysis time=%f Time of the first observation=%f" %(analysis_time, obs_checkpoints[0]))
raise ValueError
else:
if len(observations_list) == 0:
print("An empty list of observations is detected! No observations to assimilate;")
print("Nothing to assimilate...")
print("Setting the analysis state to the forecast state...")
self.smoother_configs['analysis_state'] = self.smoother_configs['forecast_state'].copy()
return
#
elif len(observations_list) == 1:
if abs(obs_checkpoints[0]-forecast_time)<=self._time_eps and \
abs(analysis_time-forecast_time)<=self._time_eps:
print("A single observation is detected, and the assimilation time coincides with observation time;")
print("You are advised to use EnKF!") # Todo, redirect automatically to 3D-Var
raise ValueError
elif (obs_checkpoints[0]-forecast_time)>self._time_eps or (obs_checkpoints[0]-analysis_time)>self._time_eps:
pass
else:
print("This settings instances are not acceptable:")
print("Observation time instance: %f" % obs_checkpoints[0])
print("Forecast time: %f" % forecast_time)
print("Observation time instance: %f" % analysis_time)
print("Terminating ")
raise ValueError
else:
# Good to go!
pass
#
# > --------------------------------------------||
# START the Smoothing process:
# > --------------------------------------------||
# Forward propagation:
# 1- Apply EnKF at each observation timepoint,
# 2- Save Kalman gain matrix
#
cwd = os.getcwd()
saved_kernels = []
# Copy the initial forecast ensemble,
moving_ensemble = [state.copy() for state in forecast_ensemble]
ensemble_size = len(moving_ensemble)
num_obs_points = len(observations_list)
if self._verbose:
print("Started the Analysis step of EnKS; with [num_obs_points] observation/assimilation time points;")
inflation_factor = self.smoother_configs['inflation_factor']
#
# Forward to observation time instances and update observation term
obs_ind = 0
for iter_ind, t0, t1 in zip(xrange(num_obs_points), obs_checkpoints[: -1], obs_checkpoints[1: ]):
local_ckeckpoints = np.array([t0, t1])
# For initial subwindow only, check the flag on both and final assimilation flags
if iter_ind == 0:
assim_flag = assim_flags[iter_ind]
if assim_flag:
observation = observations_list[obs_ind]
obs_ind += 1
print("TODO: Applying EnKF step at the initial time of the whole assimilation window")
moving_ensemble, X5 = self.EnKF_update(moving_ensemble, observation, inflation_factor, in_place=True)
saved_kernels = self._save_EnKF_update_kernel(X5, saved_kernels)
# Copy initial ensemble; could be saved to file instead
analysis_ensemble_np = utility.ensemble_to_np_array(moving_ensemble)
# Now, let's do forecast/analysis step for each subwindow:
for ens_ind in xrange(ensemble_size):
state = moving_ensemble[ens_ind]
tmp_trajectory = model.integrate_state(initial_state=state, checkpoints=local_ckeckpoints)
if isinstance(tmp_trajectory, list):
moving_ensemble[ens_ind] = tmp_trajectory[-1].copy()
else:
moving_ensemble[ens_ind] = tmp_trajectory.copy()
#
assim_flag = assim_flags[iter_ind+1]
if assim_flag:
observation = observations_list[obs_ind]
obs_ind += 1
moving_ensemble, X5 = self.EnKF_update(moving_ensemble, observation, inflation_factor, in_place=True)
saved_kernels = self._save_EnKF_update_kernel(X5, saved_kernels)
#
# Backward propagation:
for iter_ind in xrange(num_obs_points-2, -1, -1):
assim_flag = assim_flags[iter_ind]
try:
next_assim_ind = iter_ind + 1 + assim_flags[iter_ind+1: ].index(True)
except ValueError:
# This is the latest assimilation cycle; proceed
continue
if assim_flag:
kernel_file = saved_kernels[iter_ind]
X5 = np.load(kernel_file)
next_kernel = saved_kernels[next_assim_ind]
X6 = np.load(next_kernel)
X5 = X5.dot(X6)
del X6
# save/overwrite the updated kernel; no need to keep both versions!
np.save(kernel_file, X5)
# Update analysis ensemble at the initial time of the window:
X5 = np.load(saved_kernels[0])
analysis_ensemble_np = analysis_ensemble_np.dot(X5)
del X5
analysis_ensemble = moving_ensemble
for ens_ind in xrange(ensemble_size):
analysis_ensemble[ens_ind][:] = analysis_ensemble_np[:, ens_ind].copy() # Need to copy?!
self.smoother_configs.update({'analysis_ensemble': analysis_ensemble})
self.smoother_configs.update({'analysis_state': utility.ensemble_mean(analysis_ensemble)})
num_Ys=num_Ys,
F=F,
observation_size=observation_size,
ensemble_size=ensemble_size,
observation_opertor_type=observation_opertor_type,
observation_noise_level=observation_noise_level,
background_noise_level=background_noise_level,
experiment_tspan=experiment_tspan
)
coupled_ref_trajectory, coupled_init_ensemble, coupled_observations = coupled_model_info
ref_trajectory, init_ensemble, observations = model_info
IC = ref_trajectory[0]
# Forecast/Free-Run Trajectory:
forecast_trajectory = model.integrate_state(utility.ensemble_mean(init_ensemble), checkpoints)
forecast_trajectory_fine = model.integrate_state(utility.ensemble_mean(init_ensemble), np.arange(checkpoints[0], checkpoints[-1], 0.005))
#
print("Model Info Generated; Started Plotting...")
enhance_plotter()
if not os.path.isdir(plots_dir):
os.makedirs(plots_dir)
#
# Coupled Lorenz:
coupled_state_size = len(coupled_ref_trajectory[1])
trajectory = np.asarray(coupled_ref_trajectory)
fig = plt.figure(facecolor='white')
ax = fig.add_subplot(111)
extent = [checkpoints[0], checkpoints[-1], 1, coupled_state_size]