import numpy as np
from ..conesFunctions import printCones
from ..conesErrors import conesKalmanGainInversionError, conesEnsembleSizeError
from scipy import stats
from sklearn.preprocessing import StandardScaler
SEED_ROOT = 1
_region_streams = {}
[docs]
def region_rng(region_id):
"""Persistent, independent random stream for one DA region.
Every MPI rank runs its own interpreter, so a module-level generator
seeded with a constant hands every rank the same stream: the perturbed
observations of the stochastic EnKF would be identical across the
regions owned by different ranks instead of independent. Keying the
stream on the region's global id removes that, and makes a run
independent of how the regions are distributed over the ranks, so two
runs with different rank counts draw the same numbers.
The dictionary is module level on purpose: conesEnKF is rebuilt at every
analysis cycle, so a generator stored on the instance would restart from
the same state each cycle and redraw the same perturbations.
"""
key = int(region_id)
if key not in _region_streams:
_region_streams[key] = np.random.default_rng(
np.random.SeedSequence(entropy=SEED_ROOT, spawn_key=(key,)))
return _region_streams[key]
[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.rng = region_rng(region.id)
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.Ls = np.empty(0, dtype="float")
self.Lo = 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()
self.dm_history = []
[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, variance_factor=1.0):
"""
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 * variance_factor) * 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=self.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
"""
if x.shape[1] < 2:
raise conesEnsembleSizeError(x.shape[1])
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
"""
if self.Ls.size > 0:
XST = self.Ls * ( self.X @ self.S.T )
SST = self.Lo * ( self.S @ self.S.T )
else:
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) > 1e-03):
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
"""
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=self.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)
# Calculation of the localisation matrix in the state space
if n > 0:
ds = self.region.get_state2Obs_distance_matrix()
Ws = np.exp(-(ds**2)/(2 * (correlation_scale ** 2)))
self.Ls = Ws
Ls = np.copy(self.Ls)
# Repeat localisation coefficients as many times as state variables
for i in range(1, n):
self.Ls = np.vstack((self.Ls, Ls))
else:
self.Ls = 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.Ls = np.concatenate((self.Ls, Lp))
# Calculation of the localisation matrix in observation space
do = self.region.get_obs2Obs_distance_matrix()
Wo = np.exp(-(do**2)/(2 * (correlation_scale ** 2)))
self.Lo = Wo
# # Repeat localisation coefficient column wise as many times as observation variables
self.Ls = np.repeat(self.Ls, no, axis=1)
self.Lo = np.repeat(self.Lo, no, axis=1)
self.Lo = np.repeat(self.Lo, no, axis=0)
assert self.Lo.shape[0] == self.S.shape[0], (
"localisation matrix and observation vector disagree: "
f"{self.Lo.shape[0]} weights for {self.S.shape[0]} observation entries")
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
self.dm_history.append(dm)
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("Ls", self.Ls)
print("Lo", self.Lo)
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("State localisation matrix: ", self.Ls)
print("Observation localisation matrix: ", self.Lo)
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