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)