import numpy as np
from ..conesFunctions import printCones
from .conesEnKF import conesEnKF
from ..conesErrors import conesKalmanGainInversionError
from scipy import stats
from sklearn.preprocessing import StandardScaler
rng = np.random.default_rng(seed=1)
[docs]
class conesMFEnKF(conesEnKF):
"""
Multi-Fidelity Ensemble Kalman Filter.
Extends conesEnKF (principal ensemble, base class) with two extra
ensembles already projected onto the principal (fine) mesh:
- control: same size as principal, paired 1:1 (en_i <-> control_i)
- ancillary: separate, cheaper ensemble, all projected onto en0
"""
def __init__(self, region):
"""
:param region: conesMFRegion where the analysis phase will be performed
:type region: conesMFRegion
"""
super().__init__(region)
self.set_x_anomaly()
self.set_R()
self.set_y_matrix()
self.stateControlProj = region.stateControlProj
self.paramsControlProj = region.paramControlProj
self.xControlProj = np.vstack((self.stateControlProj, self.paramsControlProj))
self.samplingControlProj = region.samplingControlProj
assert self.samplingControlProj.shape[1] == self.y.shape[1], \
"control and principal ensembles must be paired member-wise for eta*Uhat = eta^X"
self.yControl = self.y.copy()
self.stateAncillaryProj = region.stateAncillaryProj
self.paramsAncillaryProj = region.paramAncillaryProj
self.xAncillaryProj = np.vstack((self.stateAncillaryProj, self.paramsAncillaryProj))
self.samplingAncillaryProj = region.samplingAncillaryProj
self.yAncillary = self.setObs(nens=self.samplingAncillaryProj.shape[1], variance_factor=3.0)
self.upxControlProj = np.empty(1, dtype="float")
self.upxAncillaryProj = np.empty(1, dtype="float")
# Anomaly Matrices for the three ensembles
self.XControlProj = self.anomaly(self.xControlProj)
self.SControlProj = self.anomaly(self.samplingControlProj)
self.XAncillaryProj = self.anomaly(self.xAncillaryProj)
self.SAncillaryProj = self.anomaly(self.samplingAncillaryProj)
self.upxControlProj = np.empty(1, dtype="float")
self.upxAncillaryProj = np.empty(1, dtype="float")
self.upStateControlProj = np.empty(1, dtype="float")
self.upStateAncillaryProj = np.empty(1, dtype="float")
self.upParamsControlProj = np.empty(1, dtype="float")
self.upParamsAncillaryProj = np.empty(1, dtype="float")
[docs]
def KalmanGainTotal(self):
"""
Compute the Multi-Fidelity Kalman Gain from the combined covariances
of the principal, control and ancillary ensembles
:returns: None
"""
XTST = (
self.X @ self.S.T
+ 0.25 * (self.XControlProj @ self.SControlProj.T + self.XAncillaryProj @ self.SAncillaryProj.T)
- 0.5 * (self.X @ self.SControlProj.T + self.XControlProj @ self.S.T)
)
STST = (
self.S @ self.S.T
+ 0.25 * (self.SControlProj @ self.SControlProj.T + self.SAncillaryProj @ self.SAncillaryProj.T)
- 0.5 * (self.S @ self.SControlProj.T + self.SControlProj @ self.S.T)
)
if self.Ls.size > 0:
XTST = self.Ls * XTST
STST = self.Lo * STST
inv = np.linalg.inv(STST + self.R)
inv_check = np.sum((STST + self.R) @ inv - np.eye(inv.shape[0]))
if (np.abs(inv_check) > 1):
printCones("Check inversion", inv_check)
printCones("STST+R eigenvalues", np.linalg.eigvalsh(STST + self.R))
printCones("diag(R)", np.diag(self.R))
raise conesKalmanGainInversionError(inv_check, 1e-3)
self.K = XTST @ inv
return
[docs]
def updateTotal(self):
"""
Update the state and parameters of the principal, control-projected
and ancillary-projected ensembles using the Multi-Fidelity Kalman Gain
:returns: None
"""
self.upx = self.x + self.K @ (self.y - self.H)
self.upxControlProj = self.xControlProj + self.K @ (self.yControl - self.samplingControlProj)
self.dxAncillaryProj = self.K @ (self.yAncillary - self.samplingAncillaryProj)
self.upxAncillaryProj = self.xAncillaryProj + self.dxAncillaryProj
printCones("check eta: max|yControl -y|", np.max(np.abs(self.yControl -self.y)))
printCones("check eta: std ratio anc/princ",
float(np.mean(np.std(self.yAncillary, axis=1) / np.std(self.y, axis=1))))
printCones("check innovation: max|y - H|", np.max(np.abs(self.y - self.H)))
printCones("check spread: param bg std",
np.std(self.x[self.state.shape[0]:, :], axis=1))
self.recenter_total()
return
[docs]
def recenter_total(self):
"""
Re-center the principal and ancillary analyses on the mean of the analysis, Popov Remark 11 equation (4.17)
The control variate needs no explicit treatment: it is assigned Phi*_r X^a
by the solver-side projection, so re-centering the principal propagates to
it and mu_Uhat^a <- Phi*_r mu_Z^a holds by construction.
:returns: None
"""
muZb = (np.mean(self.x, axis=1)
- 0.5 * (np.mean(self.xControlProj, axis=1)
- np.mean(self.xAncillaryProj, axis=1)))
muHZb = (np.mean(self.H, axis=1)
- 0.5 * (np.mean(self.samplingControlProj, axis=1)
- np.mean(self.samplingAncillaryProj, axis=1)))
muZa = muZb + self.K @ (self._obs_mean - muHZb)
self.upx += (muZa - np.mean(self.upx, axis=1))[:, None]
self.dxAncillaryProj += (muZa - np.mean(self.upxAncillaryProj, axis=1))[:, None]
self.upxAncillaryProj = self.xAncillaryProj + self.dxAncillaryProj
return
[docs]
def updateStateControlProj(self):
"""
Set the updated state of the control-projected ensemble
:returns: None
"""
self.upStateControlProj = self.upxControlProj[:self.stateControlProj.shape[0]]
return
[docs]
def updateParamsControlProj(self):
"""
Set the updated parameters of the control-projected ensemble
:returns: None
"""
self.upParamsControlProj = self.upxControlProj[self.stateControlProj.shape[0]:]
return
[docs]
def updateStateAncillaryProj(self):
"""
Set the updated state of the ancillary-projected ensemble
:returns: None
"""
self.upStateAncillaryProj = self.upxAncillaryProj[:self.stateAncillaryProj.shape[0]]
self.dStateAncillaryProj = self.dxAncillaryProj[:self.stateAncillaryProj.shape[0]]
return
[docs]
def updateParamsAncillaryProj(self):
"""
Set the updated parameters of the ancillary-projected ensemble
:returns: None
"""
self.upParamsAncillaryProj = self.upxAncillaryProj[self.stateAncillaryProj.shape[0]:]
return