import numpy as np
from ..conesFunctions import printCones
from ..conesErrors import conesKalmanGainInversionError
from scipy import stats
from sklearn.preprocessing import StandardScaler
rng = np.random.default_rng(seed=1)
[docs]
class conesEnKF():
"""
A class containing all the basic elements to perform the Ensemble Kalman filter analysis phase
"""
def __init__(self, region):
"""
:param region: conesRegion where the analysis phase will be performed
:type region: conesRegion
"""
self.region = region
self.state = self.region.state
self.params = self.region.params
self.sampling = self.region.sampling
self.nens = self.state.shape[1]
self.x = np.vstack((self.state, self.params))
self.X = np.empty(1, dtype="float")
self.S = self.anomaly(self.sampling)
self._obs_mean = np.empty(1, dtype="float")
self._obs_std = np.empty(1, dtype="float")
self.set_obs_stats()
self.R = np.empty(1, dtype="float")
self.H = self.setH()
self.y = np.empty(1, dtype="float")
self.K = np.empty(1, dtype="float")
self.L = np.empty(0, dtype="float")
self.upx = np.empty(1, dtype="float")
self.upState = np.empty(1, dtype="float")
self.upParams = np.empty(1, dtype="float")
self.nAss = np.empty(1, dtype="float")
self.inflationType = self.region.inflationType
self.stateInflation = self.region.stateInflation
self.parametersInflation = self.region.parametersInflation
self.is_converged = bool()
[docs]
def setState(self, state):
"""
Set the EnKF state
:param state: The state vector
:type state: numpy ndarray
:return: None
"""
self.state = state
return
[docs]
def setParams(self, params):
"""
Set the EnKF parameters
:param params: The parameters vector
:type params: numpy ndarray
:returns: None
"""
self.params = params
return
[docs]
def setSampling(self, sampling):
"""
Set the EnKF sampling H(x)
:param sampling: The sampling matrix
:type sampling: numpy ndarray
:returns: None
"""
self.sampling = sampling
return
[docs]
def set_obs_stats(self):
"""
Read, once, the mean and standard deviation of every non-std observation
variable in the region at the current time index, caching them for
reuse across MDA iterations by set_R() and setObs()
:returns: None
"""
means = []
stds = []
_avg = getattr(self.region, "obsType", "instantaneous") == "timeAveraged"
for ii, obs in enumerate(self.region.obsList):
obs.set_current_time_index(self.region.current_time)
if _avg:
Twin = self.region.obsAverageWindow * self.region.deltaTSim
t_now = obs.time[obs.current_time_idx]
mask = (obs.time > t_now - Twin) & (obs.time <= t_now)
dt_rec = float(np.median(np.diff(obs.time)))
n_in = int(np.count_nonzero(mask))
n_exp = max(1, int(round(Twin / dt_rec))) if dt_rec > 0 else 1
if abs(n_in - n_exp) > 1:
raise RuntimeError("observation averaging window holds %d samples, but a window of %g on a record sampled every %g s should hold %d. The observation time axis is probably not in seconds." % (n_in, Twin, dt_rec, n_exp))
if n_exp < self.region.obsAverageWindow:
printCones("note: obsAverageWindow is %d simulation steps but the observation record only provides %d sample(s) over that interval: the observation is averaged over what the record contains." % (self.region.obsAverageWindow, n_exp))
for var in obs.var_list:
if not var.endswith("_std"):
series = getattr(obs,var)
if _avg:
seg = series[mask]
means.append(seg.mean())
s = np.abs(getattr(obs, var + "_std")[mask]).mean()
stds.append(s + 1e-06 if s < 1e-10 else s)
else:
means.append(series[obs.current_time_idx])
s = np.abs(getattr(obs, var + "_std")[obs.current_time_idx])
stds.append(s + 1e-06 if s < 1e-10 else s)
self._obs_mean = np.array(means)
self._obs_std = np.array(stds)
return
[docs]
def set_R(self, mda_iterations=1):
"""
inflates the observation covariance matrix for the currend MDA
iteration, reusing the std computed once by set_obs_stats().
:param mda_iterations: inflation coefficient for this iteration
:type mda_iterations: int
:returns: None
"""
r = mda_iterations * self._obs_std**2
self.R = np.diag(r)
return
[docs]
def setH(self):
"""
Returns the EnKF sampling matrix
:returns: H
:rtype: numpy ndarray
"""
H = self.sampling
return H
[docs]
def setObs(self, mda_iterations=1, nens=None):
"""
Draws a fresh perturbed observation vector for the currend MDA iteration,
using the mean/std already read by set_obs_stats()
:params nens: ensemble size to draw the perturbed observations for. Defaults to self.nens
:param mda_iterations: inflation coefficient for this iteration
:type mda_iterations: int
:type nens: int
:returns: y
:rtype: numpy ndarray
"""
if nens is None:
nens = self.nens
noise_scaled = np.sqrt(mda_iterations) * self._obs_std
y_noised = []
for y_mean, noise in zip(self._obs_mean, noise_scaled):
lower_bound, upper_bound = y_mean - 3*noise, y_mean + 3*noise
y_noised.append(stats.truncnorm.rvs(
(lower_bound - y_mean) / noise, (upper_bound - y_mean) / noise,
loc=y_mean, scale=noise, size=nens, random_state=rng))
y_noised = np.vstack(y_noised)
return y_noised
[docs]
def set_y_matrix(self, mda_iterations=1):
"""
Set function for setObs, called once per MDA iteration.
:returns: None
"""
self.y = self.setObs(mda_iterations=mda_iterations)
return
[docs]
def anomaly(self, x):
"""
Calculate an anomaly matrix
:param x: Matrix containing the ensemble members data
:type x: numpy ndarray
:returns: The anomaly matrix
:rtype: numpy ndarray
"""
xMean = x.mean(axis=1)
X = np.zeros_like(x)
for i in range(0, x.shape[1]):
X[:, i] = (x[:, i] - xMean)/np.sqrt(x.shape[1] - 1)
return X
[docs]
def set_x_anomaly(self):
"""
Recomputes the state-parameters anomaly matrix X. Must be called at every
MDA iteration since x is refreshed by update_background().
:returns: None
"""
self.X = self.anomaly(self.x)
return
[docs]
def update_background(self):
"""
Progagates the updated x as background for the next MDA iteration.
:returns: None
"""
self.x = self.upx
return
[docs]
def KalmanGain(self):
"""
Checks for matrix inversion and set the Kalman Gain
:returns: None
"""
XST = self.X @ self.S.T
SST = self.S @ self.S.T
inv = np.linalg.inv(SST + self.R)
inv_check = np.sum((SST+self.R) @ inv - np.eye(inv.shape[0]))
if (np.abs(inv_check) > 1):
printCones("Check inversion", np.sum((SST+self.R) @ inv - np.eye(inv.shape[0])))
printCones("SST+R eigenvalues", np.linalg.eigvalsh(SST+self.R))
printCones("diag(R)", np.diag(self.R))
raise conesKalmanGainInversionError(inv_check, 1e-3)
self.K = XST @ inv
return
[docs]
def update(self):
"""
Updates the state and parameters matrix
:returns: None
"""
if self.L.size > 0:
self.upx = self.x + self.L * self.K @ (self.y - self.H)
else:
self.upx = self.x + self.K @ (self.y - self.H)
return
[docs]
def calculateInflation(self, matrix, coefficient):
"""
Applies stochastic inflation on a matrix
:param arr: the matrix to be inflated
:type arr: numpy ndarray
:param infl: inflation coefficient
:type infl: float
:returns: arr
:rtype: numpy ndarray
"""
printCones("matrix =", matrix)
if self.inflationType == "deterministic":
mean = np.mean(matrix, axis=1, keepdims=True)
printCones("mean =", mean)
matrix = mean + (1+coefficient)*(matrix - mean)
printCones("Matrix deterministically inflated =", matrix)
else: # stochastic
for i in range(self.nens):
lower, upper = -2*coefficient, 2*coefficient
mu, sigma = 0, coefficient
inflationCoeff = stats.truncnorm.rvs(
(lower - mu) / sigma, (upper - mu) / sigma,
loc=mu, scale=sigma, random_state=rng)
printCones("The inflation coefficient is :", inflationCoeff)
matrix[:, i] = (1+inflationCoeff)*matrix[:, i]
printCones("Matrix stochastically inflated =", matrix)
return matrix
[docs]
def inflateState(self):
"""
Apply inflation on state and parameters depending on the settings
"""
# State Inflation
if self.stateInflation > 0:
printCones("Inflating the state")
self.upState = self.calculateInflation(self.upState, self.stateInflation)
else:
printCones("No state inflation applied")
return
[docs]
def inflateParams(self):
"""
Apply inflation on state and parameters depending on the settings
"""
# Parameters Inflation
if self.parametersInflation > 0:
printCones("Inflating the parameters with", self.inflationType, "inflation")
self.upParams = self.calculateInflation(self.upParams, self.parametersInflation)
else:
printCones("No parameters inflation applied")
return
[docs]
def updateState(self):
"""
Sets the updated state matrix
:returns: None
"""
self.upState = self.upx[:self.state.shape[0]]
return
[docs]
def updateParams(self):
"""
Sets the updated parameters matrix
:returns: None
"""
self.upParams = self.upx[self.state.shape[0]:]
return
[docs]
def setL(self, loc_length):
"""
Sets the state covariance localisation matrix
:param: loc_length localisation caracteristic length
"""
# Arbitrary percentage applied to the update at the border of localisation to minize discontinuities
correction_at_border = 0.01
correlation_scale = loc_length/np.sqrt(-2*np.log(correction_at_border))
n = self.region.stateNvar
# Here we assume all observations have the same number of variables
no = int(len(self.region.obsList[0].var_list)/2)
if n > 0:
d = self.region.get_distance_matrix()
W = np.exp(-(d**2)/(2 * (correlation_scale ** 2)))
somme = W.sum(axis=0, keepdims=True)
# Normalize gaussian distribution
self.L = np.divide(W, somme, out=np.zeros_like(W), where=somme != 0)
L = np.copy(self.L)
# Repeat localisation coefficients as many times as state variables
for i in range(1, n):
self.L = np.vstack((self.L, L))
else:
self.L = np.empty((0, len(self.region.obsList)), dtype="float")
# Add unitary weight for parameters (some could be localised, but not implemented for the time being)
Lp = np.ones((len(self.params), len(self.region.obsList)))
self.L = np.concatenate((self.L, Lp))
# Repeat localisation coefficient column wise as many times as observation variables
self.L = np.repeat(self.L, no, axis=1)
return
[docs]
def convergence_criteria(self, crit=0):
xall = np.vstack((self.x.T, self.upx.T))
scaler = StandardScaler().fit(xall)
x_std = scaler.transform(self.x.T)
xup_std = scaler.transform(self.upx.T)
dm = np.linalg.norm(x_std - xup_std, ord='fro')/np.sqrt(x_std.size)
printCones("distance between forecast and update =", dm)
self.is_converged = dm < crit
return
[docs]
def full_report(self):
"""
Prints and extended report of the EnKF main features.
:returns: None
"""
print("\n================ ENKF REPORT ==================\n")
print("Region: ", self.region.id)
print("state ", self.state)
print("params: ", self.params)
print("sampling: ", self.sampling)
print("x", self.x)
print("X", self.X)
print("S", self.S)
print("R", self.R)
print("H", self.H, type(self.H))
print("y", self.y, type(self.y))
print("K", self.K)
print("L", self.L)
print("upx", self.upx)
print("Updated State", self.upState)
print("Updated Parameters", self.upParams)
print("\n===============================================\n")
[docs]
def small_report(self):
"""
Prints to stdout a smaller report of the EnKF
:returns: None
"""
print("\n================ ENKF REPORT ==================\n")
print("Region: ", self.region.id)
print("Parameters: ", self.params)
print("sampling: ", self.sampling)
print("Kalman Gain: ", self.K)
print("Localisation matrix: ", self.L)
print("y", self.y)
print("y-Hx", self.y - self.H)
print("Updated Parameters", self.upParams)
print("avg = ", np.mean(self.upParams))
print("std = ", np.std(self.upParams))
print("RMSE", np.linalg.norm(np.abs(self.y - self.sampling))/np.linalg.norm(self.y))
print("\n===============================================\n")
[docs]
def dump_error_convergence(self):
"""
Dumps in file "err.log" the norm of the difference between observation and model sampling
:returns: None
"""
with open("err.log", "a") as file:
dump = str(np.linalg.norm(np.abs(self.y - self.sampling))) + "\n"
file.write(dump)
return
[docs]
def dump_params_convergence(self):
"""
Dumps in file "par.log" the updated parameter ensemble mean and standard deviation
:returns: None
"""
with open("par.log", "a") as file:
dump = str(np.mean(self.upParams)) + "\t" + str(np.std(self.upParams)) + "\n"
file.write(dump)
return