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 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, mtools.center(E)[0])
for i, 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, t, dt 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 self.B in (None, 'clim'):
# Use climatological cov, ...
B = np.cov(xx.T) # ... estimated from truth
elif self.B == 'eye':
B = np.eye(HMM.Nx)
else:
B = self.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(mtools.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 = mtools.mrdiv([email protected], H@[email protected] + 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 auto_cov(xx, nlags=4, zero_mean=False, corr=False):
"""Auto covariance function, computed along axis 0.
- `nlags`: max lag (offset) for which to compute acf.
- `corr` : normalize acf by `acf[0]` so as to return auto-CORRELATION.
With `corr=True`, this is identical to
`statsmodels.tsa.stattools.acf(xx,True,nlags)`
"""
assert nlags < len(xx)
N = len(xx)
A = xx if zero_mean else center(xx)[0]
acovf = np.zeros((nlags + 1, ) + xx.shape[1:])
for i in range(nlags + 1):
Left = A[np.arange(N - i)]
Right = A[np.arange(i, N)]
acovf[i] = (Left * Right).sum(0) / (N - i)
if corr:
acovf /= acovf[0]
return acovf
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):
"""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 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 add_noise(E, dt, noise, method):
"""Treatment of additive noise for ensembles `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) * (Qa12 @ Qa12.T), 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) * (Z @ randn(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) @ randn(rQ, N)
D_til = Z @ (Xi_hat + sqrt(dt) * Xi_til)
E += D_til.T
else:
raise KeyError('No such method')
return E
# we will replicate it below by a covariance matrix approach.
# But, for strict equivalence, one would have to use
# uniform (i.e. not Gaussian) random numbers.
wnumQ = 25
sample_filename = dpr.rc.dirs.samples / ('LA_Q_wnum%d.npz' % wnumQ)
try:
# Load pre-generated
L = np.load(sample_filename)['Left']
except FileNotFoundError:
# First-time use
print('Did not find sample file', sample_filename,
'for experiment initialization. Generating...')
NQ = 20000 # Must have NQ > (2*wnumQ+1)
A = sinusoidal_sample(Nx, wnumQ, NQ)
A = 1 / 10 * center(A)[0] / np.sqrt(NQ)
Q = A.T @ A
U, s, _ = tsvd(Q)
L = U * np.sqrt(s)
np.savez(sample_filename, Left=L)
X0 = dpr.GaussRV(C=dpr.CovMat(np.sqrt(5) * L, 'Left'))
###################
# Forward model #
###################
damp = 0.98
Fm = Fmat(Nx, -1, 1, tseq.dt)
def step(x, t, dt):
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):
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(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(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)
