def EnKF_analysis(E, Eo, hnoise, y, upd_a, stats, kObs):
"""Perform the EnKF analysis update.
This implementation includes several flavours and forms,
specified by `upd_a`.
Main references: `bib.sakov2008deterministic`,
`bib.sakov2008implications`, `bib.hoteit2015mitigating`
"""
R = hnoise.C # Obs noise cov
N, Nx = E.shape # Dimensionality
N1 = N-1 # Ens size - 1
mu = np.mean(E, 0) # Ens mean
A = E - mu # Ens anomalies
xo = np.mean(Eo, 0) # Obs ens mean
Y = Eo-xo # Obs ens anomalies
dy = y - xo # Mean "innovation"
if 'PertObs' in upd_a:
# Uses classic, perturbed observations (Burgers'98)
C = Y.T @ Y + R.full*N1
D = mean0(hnoise.sample(N))
YC = mrdiv(Y, C)
KG = A.T @ YC
HK = Y.T @ YC
dE = (KG @ (y - D - Eo).T).T
E = E + dE
elif 'Sqrt' in upd_a:
# Uses a symmetric square root (ETKF)
# to deterministically transform the ensemble.
# The various versions below differ only numerically.
# EVD is default, but for large N use SVD version.
if upd_a == 'Sqrt' and N > Nx:
upd_a = 'Sqrt svd'
if 'explicit' in upd_a:
# Not recommended due to numerical costs and instability.
# Implementation using inv (in ens space)
Pw = sla.inv(Y @ R.inv @ Y.T + N1*eye(N))
T = sla.sqrtm(Pw) * sqrt(N1)
HK = R.inv @ Y.T @ Pw @ Y
# KG = R.inv @ Y.T @ Pw @ A
elif 'svd' in upd_a:
# Implementation using svd of Y R^{-1/2}.
V, s, _ = svd0(Y @ R.sym_sqrt_inv.T)
d = pad0(s**2, N) + N1
Pw = (V * d**(-1.0)) @ V.T
T = (V * d**(-0.5)) @ V.T * sqrt(N1)
# docs/snippets/trHK.jpg
trHK = np.sum((s**2+N1)**(-1.0) * s**2)
elif 'sS' in upd_a:
# Same as 'svd', but with slightly different notation
# (sometimes used by Sakov) using the normalization sqrt(N1).
S = Y @ R.sym_sqrt_inv.T / sqrt(N1)
V, s, _ = svd0(S)
d = pad0(s**2, N) + 1
Pw = (V * d**(-1.0))@V.T / N1 # = G/(N1)
T = (V * d**(-0.5))@V.T
# docs/snippets/trHK.jpg
trHK = np.sum((s**2 + 1)**(-1.0)*s**2)
else: # 'eig' in upd_a:
# Implementation using eig. val. decomp.
d, V = sla.eigh(Y @ R.inv @ Y.T + N1*eye(N))
T = V@diag(d**(-0.5))@V.T * sqrt(N1)
Pw = V@diag(d**(-1.0))@V.T
HK = R.inv @ Y.T @ (V @ diag(d**(-1)) @ V.T) @ Y
w = dy @ R.inv @ Y.T @ Pw
E = mu + w@A + T@A
elif 'Serial' in upd_a:
# Observations assimilated one-at-a-time:
inds = serial_inds(upd_a, y, R, A)
# Requires de-correlation:
dy = dy @ R.sym_sqrt_inv.T
Y = Y @ R.sym_sqrt_inv.T
# Enhancement in the nonlinear case:
# re-compute Y each scalar obs assim.
# But: little benefit, model costly (?),
# updates cannot be accumulated on S and T.
if any(x in upd_a for x in ['Stoch', 'ESOPS', 'Var1']):
# More details: Misc/Serial_ESOPS.py.
for i, j in enumerate(inds):
# Perturbation creation
if 'ESOPS' in upd_a:
# "2nd-O exact perturbation sampling"
if i == 0:
# Init -- increase nullspace by 1
V, s, UT = svd0(A)
s[N-2:] = 0
A = svdi(V, s, UT)
v = V[:, N-2]
else:
# Orthogonalize v wrt. the new A
#
# v = Zj - Yj (from paper) requires Y==HX.
# Instead: mult` should be c*ones(Nx) so we can
# project v into ker(A) such that v@A is null.
mult = (v@A) / (Yj@A) # noqa
v = v - mult[0]*Yj # noqa
v /= sqrt(v@v)
Zj = v*sqrt(N1) # Standardized perturbation along v
Zj *= np.sign(rnd.rand()-0.5) # Random sign
else:
# The usual stochastic perturbations.
Zj = mean0(rnd.randn(N)) # Un-coloured noise
if 'Var1' in upd_a:
Zj *= sqrt(N/(Zj@Zj))
# Select j-th obs
Yj = Y[:, j] # [j] obs anomalies
dyj = dy[j] # [j] innov mean
DYj = Zj - Yj # [j] innov anomalies
DYj = DYj[:, None] # Make 2d vertical
# Kalman gain computation
C = Yj@Yj + N1 # Total obs cov
KGx = Yj @ A / C # KG to update state
KGy = Yj @ Y / C # KG to update obs
# Updates
A += DYj * KGx
mu += dyj * KGx
Y += DYj * KGy
dy -= dyj * KGy
E = mu + A
else:
# "Potter scheme", "EnSRF"
# - EAKF's two-stage "update-regress" form yields
# the same *ensemble* as this.
# - The form below may be derived as "serial ETKF",
# but does not yield the same
# ensemble as 'Sqrt' (which processes obs as a batch)
# -- only the same mean/cov.
T = eye(N)
for j in inds:
Yj = Y[:, j]
C = Yj@Yj + N1
Tj = np.outer(Yj, Yj / (C + sqrt(N1*C)))
T -= Tj @ T
Y -= Tj @ Y
w = [email protected]@T/N1
E = mu + w@A + T@A
elif 'DEnKF' == upd_a:
# Uses "Deterministic EnKF" (sakov'08)
C = Y.T @ Y + R.full*N1
YC = mrdiv(Y, C)
KG = A.T @ YC
HK = Y.T @ YC
E = E + KG@dy - 0.5*([email protected]).T
else:
raise KeyError("No analysis update method found: '" + upd_a + "'.")
# Diagnostic: relative influence of observations
if 'trHK' in locals():
stats.trHK[kObs] = trHK / hnoise.M
elif 'HK' in locals():
stats.trHK[kObs] = HK.trace()/hnoise.M
return E
def add_noise(E, dt, noise, method):
"""Treatment of additive noise for ensembles.
Refs: `bib.raanes2014ext`
"""
if noise.C == 0:
return E
N, Nx = E.shape
A, mu = center(E)
Q12 = noise.C.Left
Q = noise.C.full
def sqrt_core():
T = np.nan # cause error if used
Qa12 = np.nan # cause error if used
A2 = A.copy() # Instead of using (the implicitly nonlocal) A,
# which changes A outside as well. NB: This is a bug in Datum!
if N <= Nx:
Ainv = tinv(A2.T)
Qa12 = Ainv@Q12
T = funm_psd(eye(N) + dt*(N-1)*([email protected]), sqrt)
A2 = T@A2
else: # "Left-multiplying" form
P = A2.T @ A2 / (N-1)
L = funm_psd(eye(Nx) + dt*mrdiv(Q, P), sqrt)
A2 = A2 @ L.T
E = mu + A2
return E, T, Qa12
if method == 'Stoch':
# In-place addition works (also) for empty [] noise sample.
E += sqrt(dt)*noise.sample(N)
elif method == 'none':
pass
elif method == 'Mult-1':
varE = np.var(E, axis=0, ddof=1).sum()
ratio = (varE + dt*diag(Q).sum())/varE
E = mu + sqrt(ratio)*A
E = svdi(*tsvd(E, 0.999)) # Explained in Datum
elif method == 'Mult-M':
varE = np.var(E, axis=0)
ratios = sqrt((varE + dt*diag(Q))/varE)
E = mu + A*ratios
E = svdi(*tsvd(E, 0.999)) # Explained in Datum
elif method == 'Sqrt-Core':
E = sqrt_core()[0]
elif method == 'Sqrt-Mult-1':
varE0 = np.var(E, axis=0, ddof=1).sum()
varE2 = (varE0 + dt*diag(Q).sum())
E, _, Qa12 = sqrt_core()
if N <= Nx:
A, mu = center(E)
varE1 = np.var(E, axis=0, ddof=1).sum()
ratio = varE2/varE1
E = mu + sqrt(ratio)*A
E = svdi(*tsvd(E, 0.999)) # Explained in Datum
elif method == 'Sqrt-Add-Z':
E, _, Qa12 = sqrt_core()
if N <= Nx:
Z = Q12 - A.T@Qa12
E += sqrt(dt)*([email protected](Z.shape[1], N)).T
elif method == 'Sqrt-Dep':
E, T, Qa12 = sqrt_core()
if N <= Nx:
# Q_hat12: reuse svd for both inversion and projection.
Q_hat12 = A.T @ Qa12
U, s, VT = tsvd(Q_hat12, 0.99)
Q_hat12_inv = (VT.T * s**(-1.0)) @ U.T
Q_hat12_proj = VT.T@VT
rQ = Q12.shape[1]
# Calc D_til
Z = Q12 - Q_hat12
D_hat = A.T@(T-eye(N))
Xi_hat = Q_hat12_inv @ D_hat
Xi_til = (eye(rQ) - Q_hat12_proj)@rnd.randn(rQ, N)
D_til = Z@(Xi_hat + sqrt(dt)*Xi_til)
E += D_til.T
else:
raise KeyError('No such method')
return E