def assimilate(self, HMM, xx, yy):
E = zeros((HMM.tseq.K + 1, self.N, HMM.Dyn.M))
Ef = E.copy()
E[0] = HMM.X0.sample(self.N)
# Forward pass
for k, ko, t, dt in progbar(HMM.tseq.ticker):
E[k] = HMM.Dyn(E[k - 1], t - dt, dt)
E[k] = add_noise(E[k], dt, HMM.Dyn.noise, self.fnoise_treatm)
Ef[k] = E[k]
if ko is not None:
self.stats.assess(k, ko, 'f', E=E[k])
Eo = HMM.Obs(E[k], t)
y = yy[ko]
E[k] = EnKF_analysis(E[k], Eo, HMM.Obs.noise, y, self.upd_a,
self.stats, ko)
E[k] = post_process(E[k], self.infl, self.rot)
self.stats.assess(k, ko, 'a', E=E[k])
# Backward pass
for k in progbar(range(HMM.tseq.K)[::-1]):
A = center(E[k])[0]
Af = center(Ef[k + 1])[0]
J = tinv(Af) @ A
J *= self.DeCorr
E[k] += (E[k + 1] - Ef[k + 1]) @ J
for k, ko, _, _ in progbar(HMM.tseq.ticker, desc='Assessing'):
self.stats.assess(k, ko, 'u', E=E[k])
if ko is not None:
self.stats.assess(k, ko, 's', E=E[k])
def post_process(E, infl, rot):
"""Inflate, Rotate.
To avoid recomputing/recombining anomalies,
this should have been inside :func:`EnKF_analysis`
But it is kept as a separate function
- for readability;
- to avoid inflating/rotationg smoothed states (for the :func:`EnKS`).
"""
do_infl = infl != 1.0 and infl != '-N'
if do_infl or rot:
A, mu = center(E)
N, Nx = E.shape
T = eye(N)
if do_infl:
T = infl * T
if rot:
T = genOG_1(N, rot) @ T
E = mu + T@A
return E
def local_analyses(E,
Eo,
R,
y,
state_batches,
obs_taperer,
mp=map,
xN=None,
g=0):
"""Perform local analysis update for the LETKF."""
def local_analysis(ii):
"""Perform analysis, for state index batch `ii`."""
# Locate local domain
oBatch, tapering = obs_taperer(ii)
Eii = E[:, ii]
# No update
if len(oBatch) == 0:
return Eii, 1
# Localize
Yl = Y[:, oBatch]
dyl = dy[oBatch]
tpr = sqrt(tapering)
# Adaptive inflation estimation.
# NB: Localisation is not 100% compatible with the EnKF-N, since
# - After localisation there is much less need for inflation.
# - Tapered values (Y, dy) are too neat
# (the EnKF-N expects a normal amount of sampling error).
# One fix is to tune xN (maybe set it to 2 or 3). Thanks to adaptivity,
# this should still be easier than tuning the inflation factor.
infl1 = 1 if xN is None else sqrt(N1 / effective_N(Yl, dyl, xN, g))
Eii, Yl = inflate_ens(Eii, infl1), Yl * infl1
# Since R^{-1/2} was already applied (necesry for effective_N), now use R=Id.
# TODO 4: the cost of re-init this R might not always be insignificant.
R = GaussRV(C=1, M=len(dyl))
# Update
Eii = EnKF_analysis(Eii, Yl * tpr, R, dyl * tpr, "Sqrt")
return Eii, infl1
# Prepare analysis
N1 = len(E) - 1
Y, xo = center(Eo)
# Transform obs space
Y = Y @ R.sym_sqrt_inv.T
dy = (y - xo) @ R.sym_sqrt_inv.T
# Run
result = mp(local_analysis, state_batches)
# Assign
E_batches, infl1 = zip(*result)
# TODO: this overwrites E, possibly unbeknownst to caller
for ii, Eii in zip(state_batches, E_batches):
E[:, ii] = Eii
return E, dict(ad_inf=sqrt(np.mean(np.array(infl1)**2)))
def assimilate(self, HMM, xx, yy):
Dyn, Obs, chrono, X0, stats = HMM.Dyn, HMM.Obs, HMM.t, HMM.X0, self.stats
N1 = self.N-1
R = Obs.noise
Rm12 = Obs.noise.C.sym_sqrt_inv
E = X0.sample(self.N)
stats.assess(0, E=E)
for k, kObs, t, dt in progbar(chrono.ticker):
E = Dyn(E, t-dt, dt)
E = add_noise(E, dt, Dyn.noise, self.fnoise_treatm)
if kObs is not None:
stats.assess(k, kObs, 'f', E=E)
y = yy[kObs]
inds = serial_inds(self.ordr, y, R, center(E)[0])
state_taperer = Obs.localizer(self.loc_rad, 'y2x', t, self.taper)
for j in inds:
# Prep:
# ------------------------------------------------------
Eo = Obs(E, t)
xo = np.mean(Eo, 0)
Y = Eo - xo
mu = np.mean(E, 0)
A = E-mu
# Update j-th component of observed ensemble:
# ------------------------------------------------------
Y_j = Rm12[j, :] @ Y.T
dy_j = Rm12[j, :] @ (y - xo)
# Prior var * N1:
sig2_j = Y_j@Y_j
if sig2_j < 1e-9:
continue
# Update (below, we drop the locality subscript: _j)
sig2_u = 1/(1/sig2_j + 1/N1) # KG * N1
alpha = (N1/(N1+sig2_j))**(0.5) # Update contraction factor
dy2 = sig2_u * dy_j/N1 # Mean update
Y2 = alpha*Y_j # Anomaly update
# Update state (regress update from obs space, using localization)
# ------------------------------------------------------
ii, tapering = state_taperer(j)
if len(ii) == 0:
continue
Regression = (A[:, ii]*tapering).T @ Y_j/np.sum(Y_j**2)
mu[ii] += Regression*dy2
A[:, ii] += np.outer(Y2 - Y_j, Regression)
# Without localization:
# Regression = A.T @ Y_j/np.sum(Y_j**2)
# mu += Regression*dy2
# A += np.outer(Y2 - Y_j, Regression)
E = mu + A
E = post_process(E, self.infl, self.rot)
stats.assess(k, kObs, E=E)
def assimilate(self, HMM, xx, yy):
Dyn, Obs, chrono, X0, stats, N = \
HMM.Dyn, HMM.Obs, HMM.t, HMM.X0, self.stats, self.N
N1 = N - 1
step = 1 / N
cdf_grid = np.linspace(step / 2, 1 - step / 2, N)
R = Obs.noise
Rm12 = Obs.noise.C.sym_sqrt_inv
E = X0.sample(N)
stats.assess(0, E=E)
for k, kObs, t, dt in progbar(chrono.ticker):
E = Dyn(E, t - dt, dt)
E = add_noise(E, dt, Dyn.noise, self.fnoise_treatm)
if kObs is not None:
stats.assess(k, kObs, 'f', E=E)
y = yy[kObs]
inds = serial_inds(self.ordr, y, R, center(E)[0])
for _, j in enumerate(inds):
Eo = Obs(E, t)
xo = np.mean(Eo, 0)
Y = Eo - xo
mu = np.mean(E, 0)
A = E - mu
# Update j-th component of observed ensemble
dYf = Rm12[j, :] @ (y - Eo).T # NB: does Rm12 make sense?
Yj = Rm12[j, :] @ Y.T
Regr = A.T @ Yj / np.sum(Yj**2)
Sorted = np.argsort(dYf)
Revert = np.argsort(Sorted)
dYf = dYf[Sorted]
w = reweight(np.ones(N), innovs=dYf[:, None]) # Lklhd
w = w.clip(1e-10) # Avoid zeros in interp1
cw = w.cumsum()
cw /= cw[-1]
cw *= N1 / N
cdfs = np.minimum(np.maximum(cw[0], cdf_grid), cw[-1])
dhE = -dYf + np.interp(cdfs, cw, dYf)
dhE = dhE[Revert]
# Update state by regression
E += np.outer(-dhE, Regr)
E = post_process(E, self.infl, self.rot)
stats.assess(k, kObs, E=E)
def assimilate(self, HMM, xx, yy):
Dyn, Obs, chrono, X0, stats = \
HMM.Dyn, HMM.Obs, HMM.t, HMM.X0, self.stats
E = zeros((chrono.K + 1, self.N, Dyn.M))
Ef = E.copy()
E[0] = X0.sample(self.N)
# Forward pass
for k, kObs, t, dt in progbar(chrono.ticker):
E[k] = Dyn(E[k - 1], t - dt, dt)
E[k] = add_noise(E[k], dt, Dyn.noise, self.fnoise_treatm)
Ef[k] = E[k]
if kObs is not None:
stats.assess(k, kObs, 'f', E=E[k])
Eo = Obs(E[k], t)
y = yy[kObs]
E[k] = EnKF_analysis(E[k], Eo, Obs.noise, y, self.upd_a, stats,
kObs)
E[k] = post_process(E[k], self.infl, self.rot)
stats.assess(k, kObs, 'a', E=E[k])
# Backward pass
for k in progbar(range(chrono.K)[::-1]):
A = center(E[k])[0]
Af = center(Ef[k + 1])[0]
J = tinv(Af) @ A
J *= self.cntr
E[k] += (E[k + 1] - Ef[k + 1]) @ J
for k, kObs, _, _ in progbar(chrono.ticker, desc='Assessing'):
stats.assess(k, kObs, 'u', E=E[k])
if kObs is not None:
stats.assess(k, kObs, 's', E=E[k])
def assimilate(self, HMM, xx, yy):
Dyn, Obs, chrono, X0, stats = HMM.Dyn, HMM.Obs, HMM.t, HMM.X0, self.stats
if isinstance(self.B, np.ndarray):
# compare ndarray 1st to avoid == error for ndarray
B = self.B.astype(float)
elif self.B in (None, 'clim'):
# Use climatological cov, estimated from truth
B = np.cov(xx.T)
elif self.B == 'eye':
B = np.eye(HMM.Nx)
else:
raise ValueError("Bad input B.")
B *= self.xB
# ONLY USED FOR DIAGNOSTICS, not to change the Kalman gain.
CC = 2 * np.cov(xx.T)
L = series.estimate_corr_length(center(xx)[0].ravel(order='F'))
P = X0.C.full
SM = fit_sigmoid(P.trace() / CC.trace(), L, 0)
# Init
mu = X0.mu
stats.assess(0, mu=mu, Cov=P)
for k, kObs, t, dt in progbar(chrono.ticker):
# Forecast
mu = Dyn(mu, t - dt, dt)
P = CC * SM(k)
if kObs is not None:
stats.assess(k, kObs, 'f', mu=mu, Cov=P)
# Analysis
H = Obs.linear(mu, t)
KG = mrdiv(B @ H.T, H @ B @ H.T + Obs.noise.C.full)
mu = mu + KG @ (yy[kObs] - Obs(mu, t))
# Re-calibrate fit_sigmoid with new W0 = Pa/B
P = (np.eye(Dyn.M) - KG @ H) @ B
SM = fit_sigmoid(P.trace() / CC.trace(), L, k)
stats.assess(k, kObs, mu=mu, Cov=P)
def assimilate(self, HMM, xx, yy):
Dyn, Obs, chrono, X0, stats, N = \
HMM.Dyn, HMM.Obs, HMM.t, HMM.X0, self.stats, self.N
R, KObs, N1 = HMM.Obs.noise.C, HMM.t.KObs, N - 1
Rm12 = R.sym_sqrt_inv
assert Dyn.noise.C == 0, (
"Q>0 not yet supported."
" See Sakov et al 2017: 'An iEnKF with mod. error'")
if self.bundle:
EPS = 1e-4 # Sakov/Boc use T=EPS*eye(N), with EPS=1e-4, but I ...
else:
EPS = 1.0 # ... prefer using T=EPS*T, yielding a conditional cloud shape
# Initial ensemble
E = X0.sample(N)
# Loop over DA windows (DAW).
for kObs in progbar(np.arange(-1, KObs + self.Lag + 1)):
kLag = kObs - self.Lag
DAW = range(max(0, kLag + 1), min(kObs, KObs) + 1)
# Assimilation (if ∃ "not-fully-assimlated" obs).
if 0 <= kObs <= KObs:
# Init iterations.
X0, x0 = center(E) # Decompose ensemble.
w = np.zeros(N) # Control vector for the mean state.
T = np.eye(N) # Anomalies transform matrix.
Tinv = np.eye(N)
# Explicit Tinv [instead of tinv(T)] allows for merging MDA code
# with iEnKS/EnRML code, and flop savings in 'Sqrt' case.
for iteration in np.arange(self.nIter):
# Reconstruct smoothed ensemble.
E = x0 + (w + EPS * T) @ X0
# Forecast.
for kCycle in DAW:
for k, t, dt in chrono.cycle(kCycle): # noqa
E = Dyn(E, t - dt, dt)
# Observe.
Eo = Obs(E, t)
# Undo the bundle scaling of ensemble.
if EPS != 1.0:
E = inflate_ens(E, 1 / EPS)
Eo = inflate_ens(Eo, 1 / EPS)
# Assess forecast stats; store {Xf, T_old} for analysis assessment.
if iteration == 0:
stats.assess(k, kObs, 'f', E=E)
Xf, xf = center(E)
T_old = T
# Prepare analysis.
y = yy[kObs] # Get current obs.
Y, xo = center(Eo) # Get obs {anomalies, mean}.
dy = (y - xo) @ Rm12.T # Transform obs space.
Y = Y @ Rm12.T # Transform obs space.
Y0 = Tinv @ Y # "De-condition" the obs anomalies.
V, s, UT = svd0(Y0) # Decompose Y0.
# Set "cov normlzt fctr" za ("effective ensemble size")
# => pre_infl^2 = (N-1)/za.
if self.xN is None:
za = N1
else:
za = zeta_a(*hyperprior_coeffs(s, N, self.xN), w)
if self.MDA:
# inflation (factor: nIter) of the ObsErrCov.
za *= self.nIter
# Post. cov (approx) of w,
# estimated at current iteration, raised to power.
def Cowp(expo):
return (V * (pad0(s**2, N) + za)**-expo) @ V.T
Cow1 = Cowp(1.0)
if self.MDA: # View update as annealing (progressive assimilation).
Cow1 = Cow1 @ T # apply previous update
dw = dy @ Y.T @ Cow1
if 'PertObs' in self.upd_a: # == "ES-MDA". By Emerick/Reynolds
D = mean0(np.random.randn(*Y.shape)) * np.sqrt(
self.nIter)
T -= (Y + D) @ Y.T @ Cow1
elif 'Sqrt' in self.upd_a: # == "ETKF-ish". By Raanes
T = Cowp(0.5) * np.sqrt(za) @ T
elif 'Order1' in self.upd_a: # == "DEnKF-ish". By Emerick
T -= 0.5 * Y @ Y.T @ Cow1
# Tinv = eye(N) [as initialized] coz MDA does not de-condition.
else: # View update as Gauss-Newton optimzt. of log-posterior.
grad = Y0 @ dy - w * za # Cost function gradient
dw = grad @ Cow1 # Gauss-Newton step
# ETKF-ish". By Bocquet/Sakov.
if 'Sqrt' in self.upd_a:
# Sqrt-transforms
T = Cowp(0.5) * np.sqrt(N1)
Tinv = Cowp(-.5) / np.sqrt(N1)
# Tinv saves time [vs tinv(T)] when Nx<N
# "EnRML". By Oliver/Chen/Raanes/Evensen/Stordal.
elif 'PertObs' in self.upd_a:
D = mean0(np.random.randn(*Y.shape)) \
if iteration == 0 else D
gradT = -(Y + D) @ Y0.T + N1 * (np.eye(N) - T)
T = T + gradT @ Cow1
# Tinv= tinv(T, threshold=N1) # unstable
Tinv = sla.inv(T + 1) # the +1 is for stability.
# "DEnKF-ish". By Raanes.
elif 'Order1' in self.upd_a:
# Included for completeness; does not make much sense.
gradT = -0.5 * Y @ Y0.T + N1 * (np.eye(N) - T)
T = T + gradT @ Cow1
Tinv = tinv(T, threshold=N1)
w += dw
if dw @ dw < self.wtol * N:
break
# Assess (analysis) stats.
# The final_increment is a linearization to
# (i) avoid re-running the model and
# (ii) reproduce EnKF in case nIter==1.
final_increment = (dw + T - T_old) @ Xf
# See docs/snippets/iEnKS_Ea.jpg.
stats.assess(k, kObs, 'a', E=E + final_increment)
stats.iters[kObs] = iteration + 1
if self.xN:
stats.infl[kObs] = np.sqrt(N1 / za)
# Final (smoothed) estimate of E at [kLag].
E = x0 + (w + T) @ X0
E = post_process(E, self.infl, self.rot)
# Slide/shift DAW by propagating smoothed ('s') ensemble from [kLag].
if -1 <= kLag < KObs:
if kLag >= 0:
stats.assess(chrono.kkObs[kLag], kLag, 's', E=E)
for k, t, dt in chrono.cycle(kLag + 1):
stats.assess(k - 1, None, 'u', E=E)
E = Dyn(E, t - dt, dt)
stats.assess(k, KObs, 'us', E=E)
def assimilate(self, HMM, xx, yy):
Dyn, Obs, chrono, X0, stats, N = \
HMM.Dyn, HMM.Obs, HMM.t, HMM.X0, self.stats, self.N
R, N1 = HMM.Obs.noise.C, N-1
_map = mp.map if self.mp else map
E = X0.sample(N)
stats.assess(0, E=E)
for k, kObs, t, dt in progbar(chrono.ticker):
# Forecast
E = Dyn(E, t-dt, dt)
E = add_noise(E, dt, Dyn.noise, self.fnoise_treatm)
if kObs is not None:
stats.assess(k, kObs, 'f', E=E)
# Decompose ensmeble
mu = np.mean(E, 0)
A = E - mu
# Obs space variables
y = yy[kObs]
Y, xo = center(Obs(E, t))
# Transform obs space
Y = Y @ R.sym_sqrt_inv.T
dy = (y - xo) @ R.sym_sqrt_inv.T
# Local analyses
# Get localization configuration
state_batches, obs_taperer = \
Obs.localizer(self.loc_rad, 'x2y', t, self.taper)
# Avoid pickling self
xN, g, infl = self.xN, self.g, self.infl
def local_analysis(ii):
"""Do the local analysis.
Notation:
- ii: inds for the state batch defining the locality
- jj: inds for the associated obs
"""
# Locate local obs
jj, tapering = obs_taperer(ii)
if len(jj) == 0:
return E[:, ii], N1 # no update
Y_jj = Y[:, jj]
dy_jj = dy[jj]
# Adaptive inflation
za = effective_N(Y_jj, dy_jj, xN, g) if infl == '-N' else N1
# Taper
Y_jj *= sqrt(tapering)
dy_jj *= sqrt(tapering)
# Compute ETKF update
if len(jj) < N:
# SVD version
V, sd, _ = svd0(Y_jj)
d = pad0(sd**2, N) + za
Pw = (V * d**(-1.0)) @ V.T
T = (V * d**(-0.5)) @ V.T * sqrt(za)
else:
# EVD version
d, V = sla.eigh(Y_jj@Y_jj.T + za*eye(N))
T = V@diag(d**(-0.5))@V.T * sqrt(za)
Pw = V@diag(d**(-1.0))@V.T
AT = T @ A[:, ii]
dmu = dy_jj @ Y_jj.T @ Pw @ A[:, ii]
Eii = mu[ii] + dmu + AT
return Eii, za
# Run local analyses
EE, za = zip(*_map(local_analysis, state_batches))
for ii, Eii in zip(state_batches, EE):
E[:, ii] = Eii
# Global post-processing
E = post_process(E, self.infl, self.rot)
stats.infl[kObs] = sqrt(N1/np.mean(za))
stats.assess(k, kObs, E=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
def iEnKS_update(upd_a, E, DAW, HMM, stats, EPS, y, time, Rm12, xN, MDA, threshold):
"""Perform the iEnKS update.
This implementation includes several flavours and forms,
specified by `upd_a` (See `iEnKS`)
"""
# distribute variable
k, kObs, t = time
nIter, wtol = threshold
N, Nx = E.shape
# Init iterations.
N1 = N-1
X0, x0 = center(E) # Decompose ensemble.
w = np.zeros(N) # Control vector for the mean state.
T = np.eye(N) # Anomalies transform matrix.
Tinv = np.eye(N)
# Explicit Tinv [instead of tinv(T)] allows for merging MDA code
# with iEnKS/EnRML code, and flop savings in 'Sqrt' case.
for iteration in np.arange(nIter):
# Reconstruct smoothed ensemble.
E = x0 + (w + EPS*T)@X0
# Forecast.
for kCycle in DAW:
for k, t, dt in HMM.t.cycle(kCycle): # noqa
E = HMM.Dyn(E, t-dt, dt)
# Observe.
Eo = HMM.Obs(E, t)
# Undo the bundle scaling of ensemble.
if EPS != 1.0:
E = inflate_ens(E, 1/EPS)
Eo = inflate_ens(Eo, 1/EPS)
# Assess forecast stats; store {Xf, T_old} for analysis assessment.
if iteration == 0:
stats.assess(k, kObs, 'f', E=E)
Xf, xf = center(E)
T_old = T
# Prepare analysis.
Y, xo = center(Eo) # Get obs {anomalies, mean}.
dy = (y - xo) @ Rm12.T # Transform obs space.
Y = Y @ Rm12.T # Transform obs space.
Y0 = Tinv @ Y # "De-condition" the obs anomalies.
V, s, UT = svd0(Y0) # Decompose Y0.
# Set "cov normlzt fctr" za ("effective ensemble size")
# => pre_infl^2 = (N-1)/za.
if xN is None:
za = N1
else:
za = zeta_a(*hyperprior_coeffs(s, N, xN), w)
if MDA:
# inflation (factor: nIter) of the ObsErrCov.
za *= nIter
# Post. cov (approx) of w,
# estimated at current iteration, raised to power.
def Cowp(expo): return (V * (pad0(s**2, N) + za)**-expo) @ V.T
Cow1 = Cowp(1.0)
if MDA: # View update as annealing (progressive assimilation).
Cow1 = Cow1 @ T # apply previous update
dw = dy @ Y.T @ Cow1
if 'PertObs' in upd_a: # == "ES-MDA". By Emerick/Reynolds
D = mean0(np.random.randn(*Y.shape)) * np.sqrt(nIter)
T -= (Y + D) @ Y.T @ Cow1
elif 'Sqrt' in upd_a: # == "ETKF-ish". By Raanes
T = Cowp(0.5) * np.sqrt(za) @ T
elif 'Order1' in upd_a: # == "DEnKF-ish". By Emerick
T -= 0.5 * Y @ Y.T @ Cow1
# Tinv = eye(N) [as initialized] coz MDA does not de-condition.
else: # View update as Gauss-Newton optimzt. of log-posterior.
grad = Y0@dy - w*za # Cost function gradient
dw = grad@Cow1 # Gauss-Newton step
# ETKF-ish". By Bocquet/Sakov.
if 'Sqrt' in upd_a:
# Sqrt-transforms
T = Cowp(0.5) * np.sqrt(N1)
Tinv = Cowp(-.5) / np.sqrt(N1)
# Tinv saves time [vs tinv(T)] when Nx<N
# "EnRML". By Oliver/Chen/Raanes/Evensen/Stordal.
elif 'PertObs' in upd_a:
D = mean0(np.random.randn(*Y.shape)) \
if iteration == 0 else D
gradT = -(Y+D)@Y0.T + N1*(np.eye(N) - T)
T = T + gradT@Cow1
# Tinv= tinv(T, threshold=N1) # unstable
Tinv = sla.inv(T+1) # the +1 is for stability.
# "DEnKF-ish". By Raanes.
elif 'Order1' in upd_a:
# Included for completeness; does not make much sense.
gradT = -0.5*[email protected] + N1*(np.eye(N) - T)
T = T + gradT@Cow1
Tinv = tinv(T, threshold=N1)
w += dw
if dw@dw < wtol*N:
break
# Assess (analysis) stats.
# The final_increment is a linearization to
# (i) avoid re-running the model and
# (ii) reproduce EnKF in case nIter==1.
final_increment = (dw+T-T_old)@Xf
# See docs/snippets/iEnKS_Ea.jpg.
stats.assess(k, kObs, 'a', E=E+final_increment)
stats.iters[kObs] = iteration+1
if xN:
stats.infl[kObs] = np.sqrt(N1/za)
# Final (smoothed) estimate of E at [kLag].
E = x0 + (w+T)@X0
return E