def assimilate(self, HMM, xx, yy):
Nx = HMM.Dyn.M
R = HMM.Obs.noise.C.full
Q = 0 if HMM.Dyn.noise.C == 0 else HMM.Dyn.noise.C.full
mu = np.zeros((HMM.tseq.K+1, Nx))
P = np.zeros((HMM.tseq.K+1, Nx, Nx))
# Forecasted values
muf = np.zeros((HMM.tseq.K+1, Nx))
Pf = np.zeros((HMM.tseq.K+1, Nx, Nx))
Ff = np.zeros((HMM.tseq.K+1, Nx, Nx))
mu[0] = HMM.X0.mu
P[0] = HMM.X0.C.full
self.stats.assess(0, mu=mu[0], Cov=P[0])
# Forward pass
for k, ko, t, dt in progbar(HMM.tseq.ticker, 'ExtRTS->'):
mu[k] = HMM.Dyn(mu[k-1], t-dt, dt)
F = HMM.Dyn.linear(mu[k-1], t-dt, dt)
P[k] = self.infl**(dt)*(F@P[k-1]@F.T) + dt*Q
# Store forecast and Jacobian
muf[k] = mu[k]
Pf[k] = P[k]
Ff[k] = F
if ko is not None:
self.stats.assess(k, ko, 'f', mu=mu[k], Cov=P[k])
H = HMM.Obs.linear(mu[k], t)
KG = mrdiv(P[k] @ H.T, H@P[k]@H.T + R)
y = yy[ko]
mu[k] = mu[k] + KG@(y - HMM.Obs(mu[k], t))
KH = KG@H
P[k] = (np.eye(Nx) - KH) @ P[k]
self.stats.assess(k, ko, 'a', mu=mu[k], Cov=P[k])
# Backward pass
for k in progbar(range(HMM.tseq.K)[::-1], 'ExtRTS<-'):
J = mrdiv(P[k]@Ff[k+1].T, Pf[k+1])
J *= self.DeCorr
mu[k] = mu[k] + J @ (mu[k+1] - muf[k+1])
P[k] = P[k] + J @ (P[k+1] - Pf[k+1]) @ J.T
for k in progbar(range(HMM.tseq.K+1), desc='Assess'):
self.stats.assess(k, mu=mu[k], Cov=P[k])
def assimilate(self, HMM, xx, yy):
Dyn, Obs, chrono, stats = HMM.Dyn, HMM.Obs, HMM.t, self.stats
# Compute "climatological" Kalman gain
muC = np.mean(xx, 0)
AC = xx - muC
PC = (AC.T @ AC) / (xx.shape[0] - 1)
# Setup scalar "time-series" covariance dynamics.
# ONLY USED FOR DIAGNOSTICS, not to affect the Kalman gain.
L = series.estimate_corr_length(AC.ravel(order='F'))
SM = fit_sigmoid(1/2, L, 0)
# Init
mu = muC
stats.assess(0, mu=mu, Cov=PC)
for k, kObs, t, dt in progbar(chrono.ticker):
# Forecast
mu = Dyn(mu, t-dt, dt)
if kObs is not None:
stats.assess(k, kObs, 'f', mu=muC, Cov=PC)
# Analysis
H = Obs.linear(muC, t)
KG = mrdiv([email protected], H@[email protected] + Obs.noise.C.full)
mu = muC + KG@(yy[kObs] - Obs(muC, t))
P = (np.eye(Dyn.M) - KG@H) @ PC
SM = fit_sigmoid(P.trace()/PC.trace(), L, k)
stats.assess(k, kObs, mu=mu, Cov=2*PC*SM(k))
def assimilate(self, HMM, xx, yy):
R = HMM.Obs.noise.C.full
Q = 0 if HMM.Dyn.noise.C == 0 else HMM.Dyn.noise.C.full
mu = HMM.X0.mu
P = HMM.X0.C.full
self.stats.assess(0, mu=mu, Cov=P)
for k, ko, t, dt in progbar(HMM.tseq.ticker):
mu = HMM.Dyn(mu, t-dt, dt)
F = HMM.Dyn.linear(mu, t-dt, dt)
P = self.infl**(dt)*(F@[email protected]) + dt*Q
# Of academic interest? Higher-order linearization:
# mu_i += 0.5 * (Hessian[f_i] * P).sum()
if ko is not None:
self.stats.assess(k, ko, 'f', mu=mu, Cov=P)
H = HMM.Obs.linear(mu, t)
KG = mrdiv(P @ H.T, H@[email protected] + R)
y = yy[ko]
mu = mu + KG@(y - HMM.Obs(mu, t))
KH = KG@H
P = (np.eye(HMM.Dyn.M) - KH) @ P
self.stats.trHK[ko] = KH.trace()/HMM.Dyn.M
self.stats.assess(k, ko, mu=mu, Cov=P)
def assimilate(self, HMM, xx, yy):
Dyn, Obs, chrono, X0, stats = HMM.Dyn, HMM.Obs, HMM.t, HMM.X0, self.stats
R = Obs.noise.C.full
Q = 0 if Dyn.noise.C == 0 else Dyn.noise.C.full
mu = X0.mu
P = X0.C.full
stats.assess(0, mu=mu, Cov=P)
for k, kObs, t, dt in progbar(chrono.ticker):
mu = Dyn(mu, t - dt, dt)
F = Dyn.linear(mu, t - dt, dt)
P = self.infl**(dt) * (F @ P @ F.T) + dt * Q
# Of academic interest? Higher-order linearization:
# mu_i += 0.5 * (Hessian[f_i] * P).sum()
if kObs is not None:
stats.assess(k, kObs, 'f', mu=mu, Cov=P)
H = Obs.linear(mu, t)
KG = mrdiv(P @ H.T, H @ P @ H.T + R)
y = yy[kObs]
mu = mu + KG @ (y - Obs(mu, t))
KH = KG @ H
P = (np.eye(Dyn.M) - KH) @ P
stats.trHK[kObs] = KH.trace() / Dyn.M
stats.assess(k, kObs, mu=mu, Cov=P)
def genOG_modified(M, opts=(0, 1.0)):
"""Do `genOG` with modifications.
Caution: although 'degree' ∈ (0,1) for all versions,
they're not supposed going to be strictly equivalent.
Testing: scripts/sqrt_rotations.py
"""
# Parse opts
if not opts:
# Shot-circuit in case of False or 0
return np.eye(M)
elif isinstance(opts, bool) or opts == 1:
return genOG(M)
elif isinstance(opts, float):
ver = 1
degree = opts
else:
ver = opts[0]
degree = opts[1]
if ver == 1:
# Only rotate "once in a while"
dc = 1 / degree # = "while"
# Retrieve/store persistent variable
counter = getattr(genOG_modified, "counter", 0) + 1
genOG_modified.counter = counter
# Compute rot or skip
if np.mod(counter, dc) < 1:
Q = genOG(M)
else:
Q = np.eye(M)
elif ver == 2:
# Decompose and reduce angle of (complex) diagonal. Background:
# https://stackoverflow.com/q/38426349
# https://en.wikipedia.org/wiki/Orthogonal_matrix
Q = genOG(M)
s, U = sla.eig(Q)
s2 = np.exp(1j * np.angle(s) * degree) # reduce angles
Q = mrdiv(U * s2, U)
Q = Q.real
elif ver == 3:
# Reduce Given's rotations in QR algo
raise NotImplementedError
elif ver == 4:
# Introduce correlation between columns of randn(M,M)
raise NotImplementedError
elif ver == 5:
# https://stats.stackexchange.com/q/25552
raise NotImplementedError
else:
raise KeyError
return Q
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
def assimilate(self, HMM, xx, yy):
Dyn, Obs, chrono, X0, stats = \
HMM.Dyn, HMM.Obs, HMM.t, HMM.X0, self.stats
N, Nx, R = self.N, Dyn.M, Obs.noise.C.full
E = X0.sample(N)
w = 1 / N * np.ones(N)
stats.assess(0, E=E, w=w)
for k, kObs, t, dt in progbar(chrono.ticker):
E = Dyn(E, t - dt, dt)
if Dyn.noise.C != 0:
E += np.sqrt(dt) * (rnd.randn(N, Nx) @ Dyn.noise.C.Right)
if kObs is not None:
stats.assess(k, kObs, 'f', E=E, w=w)
y = yy[kObs]
Eo = Obs(E, t)
innovs = y - Eo
# EnKF-ish update
s = self.Qs * auto_bandw(N, Nx)
As = s * raw_C12(E, w)
Ys = s * raw_C12(Eo, w)
C = Ys.T @ Ys + R
KG = As.T @ mrdiv(Ys, C)
E += sample_quickly_with(As)[0]
D = Obs.noise.sample(N)
dE = KG @ (y - Obs(E, t) - D).T
E = E + dE.T
# Importance weighting
chi2 = innovs * mldiv(C, innovs.T).T
logL = -0.5 * np.sum(chi2, axis=1)
w = reweight(w, logL=logL)
# Resampling
if trigger_resampling(w, self.NER, [stats, E, k, kObs]):
C12 = self.reg * auto_bandw(N, Nx) * raw_C12(E, w)
idx, w = resample(w, self.resampl, wroot=self.wroot)
E, _ = regularize(C12, E, idx, self.nuj)
stats.assess(k, kObs, 'u', E=E, w=w)
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 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 assimilate(self, HMM, xx, yy):
N, xN, Nx = self.N, self.xN, HMM.Dyn.M
Rm12, Ri = HMM.Obs.noise.C.sym_sqrt_inv, HMM.Obs.noise.C.inv
E = HMM.X0.sample(N)
w = 1/N*np.ones(N)
DD = None
self.stats.assess(0, E=E, w=w)
for k, ko, t, dt in progbar(HMM.tseq.ticker):
E = HMM.Dyn(E, t-dt, dt)
if HMM.Dyn.noise.C != 0:
E += np.sqrt(dt)*(rnd.randn(N, Nx)@HMM.Dyn.noise.C.Right)
if ko is not None:
self.stats.assess(k, ko, 'f', E=E, w=w)
y = yy[ko]
Eo = HMM.Obs(E, t)
wD = w.copy()
# Importance weighting
innovs = (y - Eo) @ Rm12.T
w = reweight(w, innovs=innovs)
# Resampling
if trigger_resampling(w, self.NER, [self.stats, E, k, ko]):
# Weighted covariance factors
Aw = raw_C12(E, wD)
Yw = raw_C12(Eo, wD)
# EnKF-without-pertubations update
if N > Nx:
C = Yw.T @ Yw + HMM.Obs.noise.C.full
KG = mrdiv(Aw.T@Yw, C)
cntrs = E + (y-Eo)@KG.T
Pa = Aw.T@Aw - [email protected]@Aw
P_cholU = funm_psd(Pa, np.sqrt)
if DD is None or not self.re_use:
DD = rnd.randn(N*xN, Nx)
chi2 = np.sum(DD**2, axis=1) * Nx/N
log_q = -0.5 * chi2
else:
V, sig, UT = svd0(Yw @ Rm12.T)
dgn = pad0(sig**2, N) + 1
Pw = (V * dgn**(-1.0)) @ V.T
cntrs = E + (y-Eo)@[email protected]@Pw@Aw
P_cholU = (V*dgn**(-0.5)).T @ Aw
# Generate N·xN random numbers from NormDist(0,1),
# and compute log(q(x))
if DD is None or not self.re_use:
rnk = min(Nx, N-1)
DD = rnd.randn(N*xN, N)
chi2 = np.sum(DD**2, axis=1) * rnk/N
log_q = -0.5 * chi2
# NB: the DoF_linalg/DoF_stoch correction
# is only correct "on average".
# It is inexact "in proportion" to [email protected],
# where V,s,UT = tsvd(Aw).
# Anyways, we're computing the tsvd of Aw below,
# so might as well compute q(x) instead of q(xi).
# Duplicate
ED = cntrs.repeat(xN, 0)
wD = wD.repeat(xN) / xN
# Sample q
AD = DD@P_cholU
ED = ED + AD
# log(prior_kernel(x))
s = self.Qs*auto_bandw(N, Nx)
innovs_pf = AD @ tinv(s*Aw)
# NB: Correct: innovs_pf = (ED-E_orig) @ tinv(s*Aw)
# But it seems to make no difference on well-tuned performance !
log_pf = -0.5 * np.sum(innovs_pf**2, axis=1)
# log(likelihood(x))
innovs = (y - HMM.Obs(ED, t)) @ Rm12.T
log_L = -0.5 * np.sum(innovs**2, axis=1)
# Update weights
log_tot = log_L + log_pf - log_q
wD = reweight(wD, logL=log_tot)
# Resample and reduce
wroot = 1.0
while wroot < self.wroot_max:
idx, w = resample(wD, self.resampl, wroot=wroot, N=N)
dups = sum(mask_unique_of_sorted(idx))
if dups == 0:
E = ED[idx]
break
else:
wroot += 0.1
self.stats.assess(k, ko, 'u', E=E, w=w)
