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
def assimilate(self, HMM, xx, yy):
Dyn, Obs, chrono, X0, stats = HMM.Dyn, HMM.Obs, HMM.t, HMM.X0, self.stats
R, KObs = HMM.Obs.noise.C, HMM.t.KObs
Rm12 = R.sym_sqrt_inv
Nx = Dyn.M
# Set background covariance. Note that it is static (compare to iEnKS).
if self.B in (None, 'clim'):
# Use climatological cov, ...
B = np.cov(xx.T) # ... estimated from truth
elif self.B == 'eye':
B = np.eye(Nx)
else:
B = self.B
B *= self.xB
B12 = CovMat(B).sym_sqrt
# Init
x = X0.mu
stats.assess(0, mu=x, Cov=B)
# 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.
w = np.zeros(Nx) # Control vector for the mean state.
x0 = x.copy() # Increment reference.
for iteration in np.arange(self.nIter):
# Reconstruct smoothed state.
x = x0 + B12 @ w
X = B12 # Aggregate composite TLMs onto B12
# Forecast.
for kCycle in DAW:
for k, t, dt in chrono.cycle(kCycle): # noqa
X = Dyn.linear(x, t - dt, dt) @ X
x = Dyn(x, t - dt, dt)
# Assess forecast stats
if iteration == 0:
stats.assess(k, kObs, 'f', mu=x, Cov=X @ X.T)
# Observe.
Y = Obs.linear(x, t) @ X
xo = Obs(x, t)
# Analysis prep.
y = yy[kObs] # Get current obs.
dy = Rm12 @ (y - xo) # Transform obs space.
Y = Rm12 @ Y # Transform obs space.
V, s, UT = svd0(Y.T) # Decomp for lin-alg update comps.
# Post. cov (approx) of w,
# estimated at current iteration, raised to power.
Cow1 = (V * (pad0(s**2, Nx) + 1)**-1.0) @ V.T
# Compute analysis update.
grad = Y.T @ dy - w # Cost function gradient
dw = Cow1 @ grad # Gauss-Newton step
w += dw # Step
if dw @ dw < self.wtol * Nx:
break
# Assess (analysis) stats.
final_increment = X @ dw
stats.assess(k,
kObs,
'a',
mu=x + final_increment,
Cov=X @ Cow1 @ X.T)
stats.iters[kObs] = iteration + 1
# Final (smoothed) estimate at [kLag].
x = x0 + B12 @ w
X = B12
# Slide/shift DAW by propagating smoothed ('s') state from [kLag].
if -1 <= kLag < KObs:
if kLag >= 0:
stats.assess(chrono.kkObs[kLag],
kLag,
's',
mu=x,
Cov=X @ Cow1 @ X.T)
for k, t, dt in chrono.cycle(kLag + 1):
stats.assess(k - 1, None, 'u', mu=x, Cov=Y @ Y.T)
X = Dyn.linear(x, t - dt, dt) @ X
x = Dyn(x, t - dt, dt)
stats.assess(k, KObs, 'us', mu=x, Cov=X @ Cow1 @ X.T)
def Cowp(expo):
return (V * (pad0(s**2, N) + za)**-expo) @ V.T
def nvrs(w):
# inverse of Jpp-approx
return (V * (pad0(s**2, N) + za(w)) ** -1.0) @ V.T
def pad_rk(arr): return pad0(arr, min(N, Obs.M)) def dgn_rk(l1): return pad_rk((l1*s)**2) + N1
def dgn_N(l1): return pad0((l1*s)**2, N) + N1
# Adjust hyper-prior
# xN_ = noise_level(self.xN,stats,chrono,N1,kObs,A,
# locals().get('A_old',None))
eN, cL = hyperprior_coeffs(s, N, self.xN, self.g)
def hyperprior_coeffs(s, N, xN=1, g=0):
r"""Set EnKF-N inflation hyperparams.
The EnKF-N prior may be specified by the constants:
- eN: Effect of unknown mean
- cL: Coeff in front of log term
These are trivial constants in the original EnKF-N,
but are further adjusted (corrected and tuned) for the following reasons.
- Reason 1: mode correction.
These parameters bridge the Jeffreys (`xN=1`) and Dirac (`xN=Inf`) hyperpriors
for the prior covariance, B, as discussed in `bib.bocquet2015expanding`.
Indeed, mode correction becomes necessary when $$ R \rightarrow \infty $$
because then there should be no ensemble update (and also no inflation!).
More specifically, the mode of `l1`'s should be adjusted towards 1
as a function of $$ I - K H $$ ("prior's weight").
PS: why do we leave the prior mode below 1 at all?
Because it sets up "tension" (negative feedback) in the inflation cycle:
the prior pulls downwards, while the likelihood tends to pull upwards.
- Reason 2: Boosting the inflation prior's certainty from N to xN*N.
The aim is to take advantage of the fact that the ensemble may not
have quite as much sampling error as a fully stochastic sample,
as illustrated in section 2.1 of `bib.raanes2019adaptive`.
- Its damping effect is similar to work done by J. Anderson.
The tuning is controlled by:
- `xN=1`: is fully agnostic, i.e. assumes the ensemble is generated
from a highly chaotic or stochastic model.
- `xN>1`: increases the certainty of the hyper-prior,
which is appropriate for more linear and deterministic systems.
- `xN<1`: yields a more (than 'fully') agnostic hyper-prior,
as if N were smaller than it truly is.
- `xN<=0` is not meaningful.
"""
N1 = N-1
eN = (N+1)/N
cL = (N+g)/N1
# Mode correction (almost) as in eqn 36 of `bib.bocquet2015expanding`
prior_mode = eN/cL # Mode of l1 (before correction)
diagonal = pad0(s**2, N) + N1 # diag of [email protected]@Y + N1*I
# (Hessian of J)
I_KH = np.mean(diagonal**(-1))*N1 # ≈ 1/(1 + HBH/R)
# I_KH = 1/(1 + (s**2).sum()/N1) # Scalar alternative: use tr(HBH/R).
mc = sqrt(prior_mode**I_KH) # Correction coeff
# Apply correction
eN /= mc
cL *= mc
# Boost by xN
eN *= xN
cL *= xN
return eN, cL
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)
def pad_rk(arr):
return pad0(arr, min(N, HMM.Obs.M))
def dgn_N(l1):
return pad0((l1 * s)**2, N) + N1
def pad_rk(arr):
return pad0(arr, min(N, Ny))
