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
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)
mu = np.mean(E, 0)
A = E - mu
Eo = Obs(E, t)
xo = np.mean(Eo, 0)
YR = (Eo - xo) @ Rm12.T
yR = (yy[kObs] - xo) @ Rm12.T
state_batches, obs_taperer = Obs.localizer(
self.loc_rad, 'x2y', t, self.taper)
for ii in state_batches:
# Localize obs
jj, tapering = obs_taperer(ii)
if len(jj) == 0:
return
Y_jj = YR[:, jj] * np.sqrt(tapering)
dy_jj = yR[jj] * np.sqrt(tapering)
# NETF:
# This "paragraph" is the only difference to the LETKF.
innovs = (dy_jj - Y_jj) / self.Rs
if 'laplace' in str(type(Obs.noise)).lower():
w = laplace_lklhd(innovs)
else: # assume Gaussian
w = reweight(np.ones(self.N), innovs=innovs)
dmu = w @ A[:, ii]
AT = np.sqrt(self.N) * funm_psd(
np.diag(w) - np.outer(w, w), np.sqrt) @ A[:, ii]
E[:, ii] = mu[ii] + dmu + AT
E = post_process(E, self.infl, self.rot)
stats.assess(k, kObs, E=E)
def assimilate(self, HMM, xx, yy):
# Unpack
Dyn, Obs, chrono, X0, stats = \
HMM.Dyn, HMM.Obs, HMM.t, HMM.X0, self.stats
R, N, N1 = HMM.Obs.noise.C, self.N, self.N-1
# Init
E = X0.sample(N)
stats.assess(0, E=E)
# Loop
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)
# Analysis
if kObs is not None:
stats.assess(k, kObs, 'f', E=E)
Eo = Obs(E, t)
y = yy[kObs]
mu = np.mean(E, 0)
A = E - mu
xo = np.mean(Eo, 0)
Y = Eo-xo
dy = y - xo
V, s, UT = svd0(Y @ R.sym_sqrt_inv.T)
du = UT @ (dy @ R.sym_sqrt_inv.T)
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)
if self.dual:
# Make dual cost function (in terms of l1)
def pad_rk(arr): return pad0(arr, min(N, Obs.M))
def dgn_rk(l1): return pad_rk((l1*s)**2) + N1
def J(l1):
val = np.sum(du**2/dgn_rk(l1)) \
+ eN/l1**2 \
+ cL*np.log(l1**2)
return val
# Derivatives (not required with minimize_scalar):
def Jp(l1):
val = -2*l1 * np.sum(pad_rk(s**2) * du**2/dgn_rk(l1)**2) \
+ -2*eN/l1**3 + 2*cL/l1
return val
def Jpp(l1):
val = 8*l1**2 * np.sum(pad_rk(s**4) * du**2/dgn_rk(l1)**3) \
+ 6*eN/l1**4 + -2*cL/l1**2
return val
# Find inflation factor (optimize)
l1 = Newton_m(Jp, Jpp, 1.0)
# l1 = fmin_bfgs(J, x0=[1], gtol=1e-4, disp=0)
# l1 = minimize_scalar(J, bracket=(sqrt(prior_mode), 1e2),
# tol=1e-4).x
else:
# Primal form, in a fully linearized version.
def za(w): return zeta_a(eN, cL, w)
def J(w): return \
.5*np.sum(((dy-w@Y)@R.sym_sqrt_inv.T)**2) + \
.5*N1*cL*np.log(eN + w@w)
# Derivatives (not required with fmin_bfgs):
def Jp(w): return [email protected]@(dy-w@Y) + w*za(w)
# Jpp = lambda w: [email protected]@Y.T + \
# za(w)*(eye(N) - 2*np.outer(w,w)/(eN + w@w))
# Approx: no radial-angular cross-deriv:
# Jpp = lambda w: [email protected]@Y.T + za(w)*eye(N)
def nvrs(w):
# inverse of Jpp-approx
return (V * (pad0(s**2, N) + za(w)) ** -1.0) @ V.T
# Find w (optimize)
wa = Newton_m(Jp, nvrs, zeros(N), is_inverted=True)
# wa = Newton_m(Jp,Jpp ,zeros(N))
# wa = fmin_bfgs(J,zeros(N),Jp,disp=0)
l1 = sqrt(N1/za(wa))
# Uncomment to revert to ETKF
# l1 = 1.0
# Explicitly inflate prior
# => formulae look different from `bib.bocquet2015expanding`.
A *= l1
Y *= l1
# Compute sqrt update
Pw = (V * dgn_N(l1)**(-1.0)) @ V.T
w = [email protected]@Y.T@Pw
# For the anomalies:
if not self.Hess:
# Regular ETKF (i.e. sym sqrt) update (with inflation)
T = (V * dgn_N(l1)**(-0.5)) @ V.T * sqrt(N1)
# = ([email protected]@Y.T/N1 + eye(N))**(-0.5)
else:
# Also include angular-radial co-dependence.
# Note: denominator not squared coz
# unlike `bib.bocquet2015expanding` we have inflated Y.
Hw = [email protected]@Y.T/N1 + eye(N) - 2*np.outer(w, w)/(eN + w@w)
T = funm_psd(Hw, lambda x: x**-.5) # is there a sqrtm Woodbury?
E = mu + w@A + T@A
E = post_process(E, self.infl, self.rot)
stats.infl[kObs] = l1
stats.trHK[kObs] = (((l1*s)**2 + N1)**(-1.0)*s**2).sum()/HMM.Ny
stats.assess(k, kObs, E=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)
