def calc_hess(w, *args): xc, pf, y, rmat, htype = args rinv = rmat @ rmat.transpose() x = xc + pf @ w if htype["perturbation"] == "gradw": dh = obs.dhdx(x, htype["operator"]) @ pf else: dh = obs.h_operator(x[:, None] + pf, htype["operator"]) - obs.h_operator( x, htype["operator"])[:, None] return np.eye(w.size) + dh.transpose() @ rinv @ dh
def calc_grad_j(w, *args): xc, pf, y, rmat, htype = args rinv = rmat @ rmat.transpose() x = xc + pf @ w hx = obs.h_operator(x, htype["operator"]) ob = y - hx if htype["perturbation"] == "gradw": dh = obs.dhdx(x, htype["operator"]) @ pf else: dh = obs.h_operator(x[:, None] + pf, htype["operator"]) - hx[:, None] return w - dh.transpose() @ rinv @ ob
def calc_grad_j(zeta, *args): xc, pf, y, tmat, gmat, heinv, rinv, htype = args x = xc + gmat @ zeta hx = obs.h_operator(x, htype["operator"]) ob = y - hx if htype["perturbation"] == "grad05": #if htype["perturbation"] == "grad": dh = obs.dhdx(x, htype["operator"]) @ pf else: dh = obs.h_operator(x[:, None] + pf, htype["operator"]) - hx[:, None] return heinv @ zeta - tmat @ dh.transpose() @ rinv @ ob
def calc_hess(zeta, *args): xc, pf, y, tmat, gmat, heinv, rinv, htype = args x = xc + gmat @ zeta if htype["perturbation"] == "grad05": #if htype["perturbation"] == "grad": dh = obs.dhdx(x, htype["operator"]) @ pf else: dh = obs.h_operator(x[:, None] + pf, htype["operator"]) - obs.h_operator(x, htype["operator"])[:, None] hess = np.eye(zeta.size) + dh.transpose() @ rinv @ dh hess = tmat @ hess @ tmat logger.debug(f"hess={hess}") return hess
def calc_step(alpha_t, zeta, dk, *args): xc, pf, y, tmat, gmat, heinv, rinv, htype = args zeta_t = zeta + alpha_t * dk x = xc + gmat @ zeta x_t = xc + gmat @ zeta_t hx = obs.h_operator(x, htype["operator"]) ob = y - hx if htype["perturbation"] == "grad05": #if htype["perturbation"] == "grad": dh = obs.dhdx(x, htype["operator"]) @ pf else: dh = obs.h_operator(x[:, None] + pf, htype["operator"]) - hx[:, None] cost_t = (zeta_t - zeta).transpose() @ tmat @ dh.transpose() @ rinv @ ob \ - (zeta_t - zeta).transpose() @ heinv @ zeta cost_b = (zeta_t - zeta).transpose() @ heinv @ (zeta_t - zeta) \ + (zeta_t - zeta).transpose() @ tmat @ dh.transpose() @ rinv @ dh @ tmat @ (zeta_t - zeta) return alpha_t * cost_t / cost_b
def calc_j(zeta, *args): xc, pf, y, tmat, gmat, heinv, rinv, htype = args x = xc + gmat @ zeta ob = y - obs.h_operator(x, htype["operator"]) j = 0.5 * (zeta.transpose() @ heinv @ zeta + ob.transpose() @ rinv @ ob) # logger.debug("zeta.shape={}".format(zeta.shape)) # logger.debug("j={} zeta={}".format(j, zeta)) return j
def calc_j(w, *args): xc, pf, y, rmat, htype = args rinv = rmat @ rmat.transpose() x = xc + pf @ w ob = y - obs.h_operator(x, htype["operator"]) j = 0.5 * (w.transpose() @ w + ob.transpose() @ rinv @ ob) # logger.debug("zeta.shape={}".format(zeta.shape)) # logger.debug("j={} zeta={}".format(j, zeta)) return j
def calc_grad_j(zeta, *args): #xc, pf, y, tmat, gmat, heinv, rinv, htype = args xc, pf, y, precondition, rmat, htype = args nmem = zeta.size tmat = np.load("tmat.npy") heinv = tmat.transpose() @ tmat gmat = pf @ tmat x = xc + gmat @ zeta hx = obs.h_operator(x, htype["operator"], htype["gamma"]) ob = y - hx if htype["perturbation"] == "grad": dh = obs.dhdx(x, htype["operator"], htype["gamma"]) @ pf else: dh = obs.h_operator(x[:, None] + pf, htype["operator"], htype["gamma"]) - hx[:, None] zmat = rmat @ dh g = heinv @ zeta - tmat.transpose() @ zmat.transpose() @ rmat @ ob tmat = precondition(zmat) np.save("tmat.npy", tmat) return g
def calc_j(zeta, *args): xc, pf, y, ytype, tmat, gmat, heinv, rinv, tlm, l_jhi = args x = xc + gmat @ zeta ob = np.zeros_like(y) for i in range(len(ytype)): ob[i] = y[i] - obs.h_operator(x, ytype[i]) j = 0.5 * (zeta.transpose() @ heinv @ zeta + ob.transpose() @ rinv @ ob) # logger.debug("zeta.shape={}".format(zeta.shape)) # logger.debug("j={} zeta={}".format(j, zeta)) return j
def calc_grad_j(w, *args): xf_, dxf, y, precondition, eps, rmat, htype = args #xf_, dxf, y, tmat, gmat, heinv, rinv, htype = args nmem = w.size x = xf_ + dxf @ w emat = x[:, None] + eps * dxf hemat = obs.h_operator(emat, htype["operator"], htype["gamma"]) dy = (hemat - np.mean(hemat, axis=1)[:, None]) / eps ob = y - np.mean(hemat, axis=1) rinv = rmat @ rmat.T return w - dy.transpose() @ rinv @ ob
def calc_hess(zeta, *args): xc, pf, y, ytype, tmat, gmat, heinv, rinv, tlm, l_jhi = args x = xc + gmat @ zeta dh = np.zeros((y.size,pf.shape[1])) if tlm: #if htype["perturbation"] == "grad": for i in range(len(ytype)): op = ytype[i] if not l_jhi: logger.debug("linearized about ensemble mean") dh[i,:] = obs.dhdx(x, op) @ pf else: logger.debug("linearized about each ensemble member") for j in range(pf.shape[1]): jhi = obs.dhdx(x+pf[:,j], op) dh[i,j] = jhi @ pf[:, j] else: for i in range(len(ytype)): dh[i,:] = obs.h_operator(x[:, None] + pf, ytype[i]) - obs.h_operator(x, ytype[i])[:, None] hess = np.eye(zeta.size) + dh.transpose() @ rinv @ dh hess = tmat @ hess @ tmat #logger.debug(f"hess={hess}") return hess
def calc_j(zeta, *args): #xc, pf, y, tmat, gmat, heinv, rinv, htype = args xc, pf, y, precondition, rmat, htype = args nmem = zeta.size tmat = np.load("tmat.npy") gmat = pf @ tmat heinv = tmat.transpose() @ tmat rinv = rmat.transpose() @ rmat x = xc + gmat @ zeta ob = y - obs.h_operator(x, htype["operator"], htype["gamma"]) j = 0.5 * (zeta.transpose() @ heinv @ zeta + ob.transpose() @ rinv @ ob) #j = 0.5 * ((nmem-1)*zeta.transpose() @ heinv @ zeta + nob.transpose() @ rinv @ ob) # logger.debug("zeta.shape={}".format(zeta.shape)) # logger.debug("j={} zeta={}".format(j, zeta)) return j
def calc_grad_j(zeta, *args): xc, pf, y, ytype, tmat, gmat, heinv, rinv, tlm, l_jhi = args x = xc + gmat @ zeta hx = np.zeros_like(y) for i in range(len(ytype)): hx[i] = obs.h_operator(x, ytype[i]) ob = y - hx dh = np.zeros((y.size,pf.shape[1])) if tlm: #if htype["perturbation"] == "grad": for i in range(len(ytype)): op = ytype[i] if not l_jhi: logger.debug("linearized about ensemble mean") dh[i,:] = obs.dhdx(x, op) @ pf else: logger.debug("linearized about each ensemble member") for j in range(pf.shape[1]): jhi = obs.dhdx(x+pf[:,j], op) dh[i,j] = jhi @ pf[:, j] else: for i in range(len(ytype)): dh[i,:] = obs.h_operator(x[:, None] + pf, ytype[i]) - obs.h_operator(x, ytype[i])[:, None] return heinv @ zeta - tmat @ dh.transpose() @ rinv @ ob
def calc_j(w, *args): xf_, dxf, y, precondition, eps, rmat, htype = args #xf_, dxf, y, tmat, gmat, heinv, rinv, htype = args nmem = w.size x = xf_ + dxf @ w emat = x[:, None] + eps * dxf hemat = obs.h_operator(emat, htype["operator"], htype["gamma"]) dy = (hemat - np.mean(hemat, axis=1)[:, None]) / eps ob = y - np.mean(hemat, axis=1) rinv = rmat @ rmat.T j = 0.5 * (w.transpose() @ w + ob.transpose() @ rinv @ ob) #j = 0.5 * ((nmem-1)*zeta.transpose() @ heinv @ zeta + nob.transpose() @ rinv @ ob) # logger.debug("zeta.shape={}".format(zeta.shape)) # logger.debug("j={} zeta={}".format(j, zeta)) return j
def calc_grad_j(w, *args): xf_, dxf, y, calc_hess, rinv, htype = args #xf_, dxf, y, tmat, gmat, heinv, rinv, htype = args nmem = w.size x = xf_ + dxf @ w tmat = np.load("tmat.npy") tinv = np.load("tinv.npy") emat = x[:, None] + np.sqrt(nmem - 1) * dxf @ tmat hemat = obs.h_operator(emat, htype["operator"], htype["gamma"]) dy = (hemat - np.mean(hemat, axis=1)[:, None]) @ tinv / np.sqrt(nmem - 1) # update tmat & tinv tmat, tinv = calc_hess(dy, rinv) np.save("tmat.npy", tmat) np.save("tinv.npy", tinv) ob = y - np.mean(hemat, axis=1) return w - dy.transpose() @ rinv @ ob
def calc_j(w, *args): xf_, dxf, y, calc_hess, rinv, htype = args #xf_, dxf, y, tmat, gmat, heinv, rinv, htype = args nmem = w.size x = xf_ + dxf @ w tmat = np.load("tmat.npy") tinv = np.load("tinv.npy") emat = x[:, None] + np.sqrt(nmem - 1) * dxf @ tmat hemat = obs.h_operator(emat, htype["operator"], htype["gamma"]) dy = (hemat - np.mean(hemat, axis=1)[:, None]) @ tinv / np.sqrt(nmem - 1) ob = y - np.mean(hemat, axis=1) j = 0.5 * (w.transpose() @ w + ob.transpose() @ rinv @ ob) #j = 0.5 * ((nmem-1)*zeta.transpose() @ heinv @ zeta + nob.transpose() @ rinv @ ob) # logger.debug("zeta.shape={}".format(zeta.shape)) # logger.debug("j={} zeta={}".format(j, zeta)) return j
def get_true_and_obs(na, nx, sigma, op, gamma=1): truth = pd.read_csv("data.csv") xt = truth.values.reshape(na, nx) #obs = pd.read_csv("observation_data.csv") #y = obs.values.reshape(na,nx) f = "obs_{}.npy".format(op) if not os.path.isfile(f): seed = 514 y = add_noise(h_operator(xt, op, gamma), sigma) #, seed=seed) np.save(f, y) else: y = np.load(f) #y = np.load("obs_{}.npy".format(op)) #y = np.zeros((na, nx, 2)) #obsloc = np.arange(nx) #for k in range(na): # y[k,:,0] = obsloc[:] # y[k,:,1] = obs.add_noise(obs.h_operator(obsloc, xt[k])) return xt, y
def analysis(xf, xc, y, rmat, rinv, htype, gtol=1e-6, method="LBFGS", cgtype=None, maxiter=None, disp=False, save_hist=False, save_dh=False, infl=False, loc=False, infl_parm=1.0, model="z08", icycle=100): global zetak zetak = [] op = htype["operator"] pt = htype["perturbation"] ga = htype["gamma"] nmem = xf.shape[1] pf = xf - xc[:, None] # pf = (xf - xc[:, None]) / np.sqrt(nmem) #condh = np.zeros(2) if infl: """ if op == "linear" or op == "test": if pt == "mlef": alpha = 1.2 else: alpha = 1.2 elif op == "quadratic" or op == "quadratic-nodiff": if pt == "mlef": alpha = 1.3 else: alpha = 1.35 elif op == "cubic" or op == "cubic-nodiff": if pt == "mlef": alpha = 1.6 else: alpha = 1.65 """ alpha = infl_parm logger.info("==inflation==, alpha={}".format(alpha)) #print("==inflation==, alpha={}".format(alpha)) pf *= alpha # logger.debug("norm(pf)={}".format(la.norm(pf))) if loc: l_sig = 4.0 logger.info("==localization==, l_sig={}".format(l_sig)) #print("==localization==, l_sig={}".format(l_sig)) pf = pfloc(pf, l_sig, save_dh, model, op, pt, icycle) if pt == "grad": logger.debug("dhdx={}".format(obs.dhdx(xc, op, ga))) dh = obs.dhdx(xc, op, ga) @ pf else: dh = obs.h_operator(xf, op, ga) - obs.h_operator(xc, op, ga)[:, None] if save_dh: #np.save("{}_dh_{}_{}_cycle{}.npy".format(model, op, pt, icycle), dh) np.save("{}_dy_{}_{}_cycle{}.npy".format(model, op, pt, icycle), dh) ob = y - obs.h_operator(xc, op, ga) np.save("{}_d_{}_{}_cycle{}.npy".format(model, op, pt, icycle), ob) logger.info("save_dh={}".format(save_dh)) # print("save_dh={}".format(save_dh)) zmat = rmat @ dh # logger.debug("cond(zmat)={}".format(la.cond(zmat))) tmat, heinv, condh = precondition(zmat) #if icycle == 0: # print("tmat={}".format(tmat)) # print("heinv={}".format(heinv)) # logger.debug("pf.shape={}".format(pf.shape)) # logger.debug("tmat.shape={}".format(tmat.shape)) # logger.debug("heinv.shape={}".format(heinv.shape)) gmat = pf @ tmat #gmat = np.sqrt(nmem-1) * pf @ tmat # logger.debug("gmat.shape={}".format(gmat.shape)) x0 = np.zeros(xf.shape[1]) args_j = (xc, pf, y, tmat, gmat, heinv, rinv, htype) iprint = np.zeros(2, dtype=np.int32) iprint[0] = 1 minimize = Minimize(x0.size, calc_j, jac=calc_grad_j, args=args_j, iprint=iprint, method=method, cgtype=cgtype) logger.info("save_hist={}".format(save_hist)) # print("save_hist={} cycle{}".format(save_hist, icycle)) cg = spo.check_grad(calc_j, calc_grad_j, x0, *args_j) logger.info("check_grad={}".format(cg)) if save_hist: x = minimize(x0, callback=callback) logger.debug(zetak) # print(len(zetak)) jh = np.zeros(len(zetak)) gh = np.zeros(len(zetak)) for i in range(len(zetak)): jh[i] = calc_j(np.array(zetak[i]), *args_j) g = calc_grad_j(np.array(zetak[i]), *args_j) gh[i] = np.sqrt(g.transpose() @ g) np.savetxt("{}_jh_{}_{}_cycle{}.txt".format(model, op, pt, icycle), jh) np.savetxt("{}_gh_{}_{}_cycle{}.txt".format(model, op, pt, icycle), gh) if model == "z08": if len(zetak) > 0: xmax = max(np.abs(np.min(zetak)), np.max(zetak)) else: xmax = max(np.abs(np.min(x)), np.max(x)) logger.debug("resx max={}".format(xmax)) # print("resx max={}".format(xmax)) #if xmax < 1000: #cost_j(1000, xf.shape[1], model, x, icycle, *args_j) #cost_j2d(1000, xf.shape[1], model, x, icycle, *args_j) # costJ.cost_j2d(1000, xf.shape[1], model, x, icycle, # htype, calc_j, calc_grad_j, *args_j) #else: #xmax = int(xmax*0.01+1)*100 xmax = np.ceil(xmax * 0.01) * 100 logger.debug("resx max={}".format(xmax)) # print("resx max={}".format(xmax)) #cost_j(xmax, xf.shape[1], model, x, icycle, *args_j) #cost_j2d(xmax, xf.shape[1], model, x, icycle, *args_j) costJ.cost_j2d(xmax, xf.shape[1], model, x, icycle, htype, calc_j, calc_grad_j, *args_j) elif model == "l96": cost_j(200, xf.shape[1], model, x, icycle, *args_j) else: x = minimize(x0) xa = xc + gmat @ x if save_dh: np.save("{}_dx_{}_{}_cycle{}.npy".format(model, op, pt, icycle), gmat @ x) if pt == "grad": dh = obs.dhdx(xa, op, ga) @ pf else: dh = obs.h_operator(xa[:, None] + pf, op, ga) - obs.h_operator( xa, op, ga)[:, None] zmat = rmat @ dh # logger.debug("cond(zmat)={}".format(la.cond(zmat))) tmat, heinv, dum = precondition(zmat) pa = pf @ tmat # pa *= np.sqrt(nmem) if save_dh: np.save("{}_pa_{}_{}_cycle{}.npy".format(model, op, pt, icycle), pa) ua = np.zeros((xa.size, nmem + 1)) ua[:, 0] = xa ua[:, 1:] = xa[:, None] + pa np.save("{}_ua_{}_{}_cycle{}.npy".format(model, op, pt, icycle), ua) #print("xa={}".format(xa)) d = y - obs.h_operator(xa, op, ga) chi2 = chi2_test(zmat, heinv, rmat, d) ds = dof(zmat) logger.info("DOF = {}".format(ds)) return xa, pa, chi2, ds, condh
def analysis(xf, binv, y, ytype, rinv, htype, gtol=1e-6, method="LBFGS", cgtype=None, maxiter=None, restart=False, maxrest=20, disp=False, save_hist=False, save_dh=False, model="model", icycle=0): global zetak, alphak zetak = [] alphak = [] op = htype["operator"] pt = htype["perturbation"] ga = htype["gamma"] JH = np.zeros((y.size, xf.size)) ob = np.zeros_like(y) for i in range(len(ytype)): op = ytype[i] JH[i, :] = obs.dhdx(xf, op, ga) ob[i] = y[i] - obs.h_operator(xf, op, ga) nobs = ob.size x0 = np.zeros_like(xf) args_j = (binv, JH, rinv, ob) iprint = np.zeros(2, dtype=np.int32) iprint[0] = 1 minimize = Minimize(x0.size, calc_j, jac=calc_grad_j, hess=calc_hess, args=args_j, iprint=iprint, method=method, cgtype=cgtype, maxiter=maxiter, restart=restart) cg = spo.check_grad(calc_j, calc_grad_j, x0, *args_j) logger.info("check_grad={}".format(cg)) if save_hist: x, flg = minimize(x0, callback=callback) jh = np.zeros(len(zetak)) gh = np.zeros(len(zetak)) for i in range(len(zetak)): jh[i] = calc_j(np.array(zetak[i]), *args_j) g = calc_grad_j(np.array(zetak[i]), *args_j) gh[i] = np.sqrt(g.transpose() @ g) np.savetxt("{}_jh_{}_{}_cycle{}.txt".format(model, op, pt, icycle), jh) np.savetxt("{}_gh_{}_{}_cycle{}.txt".format(model, op, pt, icycle), gh) else: x, flg = minimize(x0) xa = xf + x fun = calc_j(x, *args_j) chi2 = fun / nobs return xa, chi2
def analysis(xf, xf_, y, ytype, sig, dx, htype, \ infl=False, loc=False, tlm=True, infl_parm=1.1, \ infl_r=False, l_jhi=False,\ save_dh=False, model="z08", icycle=0): #op = htype["operator"] da = htype["perturbation"] ga = htype["gamma"] #obs = Obs(op, sig) #JH = obs.dhdx(xf_, op, ga) ##JH = obs.dh_operator(y[:,0], xf_) JH = np.zeros((y.size, xf_.size)) for i in range(len(ytype)): op = ytype[i] JH[i, :] = obs.dhdx(xf_, op, ga) logger.debug(f"JH {JH}") R = np.eye(y.size) * sig * sig rinv = np.eye(y.size) / sig / sig #R, rmat, rinv = obs.set_r(y.shape[0]) nmem = xf.shape[1] dxf = xf - xf_[:, None] logger.debug(dxf.shape) dy = np.zeros((y.size, nmem)) dyj = np.zeros((nmem, y.size, nmem)) xa = np.zeros_like(xf) dxa = np.zeros_like(dxf) if tlm: for i in range(len(ytype)): op = ytype[i] logger.debug(op) if not l_jhi: logger.debug("linearized about ensemble mean") dy[i, :] = obs.dhdx(xf_, op) @ dxf else: logger.debug("linearized about each ensemble member") for j in range(nmem): jhi = obs.dhdx(xf[:, j], op, ga) dyj[j, i, :] = jhi @ dxf dy[i, :] = obs.dhdx(xf_, op) @ dxf else: for i in range(len(ytype)): op = ytype[i] logger.debug(op) dy[i, :] = obs.h_operator(xf, op) - np.mean(obs.h_operator(xf, op), axis=1)[:, None] #dy[i, :] = obs.h_operator(xf, op) - obs.h_operator(xf_, op)[:, None] ##dy = obs.h_operator(xf, op, ga) - np.mean(obs.h_operator(xf, op, ga), axis=1)[:, None] ##dy = obs.h_operator(y[:,0], xf) - np.mean(obs.h_operator(y[:,0], xf), axis=1)[:, None] #dy = obs.h_operator(xf, op, ga) - obs.h_operator(xf_, op, ga)[:, None] logger.debug(f"HPfH {dy @ dy.T / (nmem-1)}") d = np.zeros_like(y) for i in range(len(ytype)): op = ytype[i] d[i] = y[i] - np.mean(obs.h_operator(xf, op, ga), axis=1) #d = y[:,1] - np.mean(obs.h_operator(y[:,0], xf), axis=1) #d = y - obs.h_operator(xf_, op, ga) # inflation parameter alpha = infl_parm if infl: if da != "letkf" and da != "etkf": logger.info("==inflation==, alpha={}".format(alpha)) dxf *= alpha pf = dxf @ dxf.T / (nmem - 1) if save_dh: #logger.info("save pf") np.save("{}_pf_{}_{}_cycle{}.npy".format(model, op, da, icycle), pf) #if loc: # B-localization # if da == "etkf": # print("==B-localization==") # dist, l_mat = loc_mat(sigma=2.0, nx=xf_.size, ny=xf_.size) # pf = pf * l_mat if save_dh: #logger.info("save dxf") np.save("{}_dxf_{}_{}_cycle{}.npy".format(model, op, da, icycle), dxf) np.save("{}_dhdx_{}_{}_cycle{}.npy".format(model, op, da, icycle), JH) np.save("{}_dh_{}_{}_cycle{}.npy".format(model, op, da, icycle), dy) np.save("{}_d_{}_{}_cycle{}.npy".format(model, op, da, icycle), d) logger.info("save_dh={} cycle{}".format(save_dh, icycle)) # logger.debug("norm(pf)={}".format(la.norm(pf))) if da == "etkf": if infl: logger.info("==inflation==, alpha={}".format(alpha)) A = np.eye(nmem) / alpha else: A = np.eye(nmem) A = (nmem - 1) * A + dy.T @ rinv @ dy #TT = la.inv( np.eye(nmem) + dy.T @ rinv @ dy / (nmem-1) ) lam, v = la.eigh(A) #print("eigenvalue(sqrt)={}".format(np.sqrt(lam))) Dinv = np.diag(1.0 / lam) TT = v @ Dinv @ v.T T = v @ np.sqrt((nmem - 1) * Dinv) @ v.T K = dxf @ TT @ dy.T @ rinv if save_dh: #logger.info("save K") np.save("{}_K_{}_{}_cycle{}.npy".format(model, op, da, icycle), K) if loc: # K-localization logger.info("==K-localization==") dist, l_mat = loc_mat(sigma=4.0, nx=xf_.size, ny=y.size) #print(l_mat) K = K * l_mat #print(K) if save_dh: np.save( "{}_Kloc_{}_{}_cycle{}.npy".format(model, op, da, icycle), K) if save_dh: #logger.info("save increment") np.save("{}_dx_{}_{}_cycle{}.npy".format(model, op, da, icycle), K @ d) xa_ = xf_ + K @ d dxa = dxf @ T xa = dxa + xa_[:, None] elif da == "po": Y = np.zeros((y.size, nmem)) #Y = np.zeros((y.shape[0],nmem)) #np.random.seed(514) #err = np.random.normal(0, scale=sig, size=Y.size) #err_ = np.mean(err.reshape(Y.shape), axis=1) err = np.zeros_like(Y) for j in range(nmem): err[:, j] = np.random.normal(0.0, scale=sig, size=err.shape[0]) err_ = np.mean(err, axis=1) stdv = np.sqrt(np.mean((err - err_[:, None])**2, axis=1)) logger.debug("err mean {}".format(err_)) logger.debug("err stdv {}".format(stdv)) #Y = y[:,None] + err.reshape(Y.shape) Y = y[:, None] + err #Y = y[:,1].reshape(-1,1) + err.reshape(Y.shape) #d_ = y + err_ - obs.h_operator(xf_, op, ga) d_ = d + err_ #if tlm: # K = pf @ JH.T @ la.inv(JH @ pf @ JH.T + R) #else: K = dxf @ dy.T @ la.inv(dy @ dy.T + (nmem - 1) * R) if loc: logger.info("==localization==") dist, l_mat = loc_mat(sigma=4.0, nx=xf_.size, ny=y.size) K = K * l_mat xa_ = xf_ + K @ d_ #if tlm: # xa = xf + K @ (Y - JH @ xf) # pa = (np.eye(xf_.size) - K @ JH) @ pf #else: HX = np.zeros((y.size, nmem)) for i in range(len(ytype)): op = ytype[i] HX[i, :] = obs.h_operator(xf, op, ga) #HX = obs.h_operator(y[:,0], xf) if not tlm or not l_jhi: xa = xf + K @ (Y - HX) else: for j in range(nmem): Kj = dxf @ dyj[j].T @ la.inv(dyj[j] @ dyj[j].T + (nmem - 1) * R) xa[:, j] = xf[:, j] + Kj @ (Y[:, j] - HX[:, j]) dxa = xa - xa_[:, None] #pa = pf - K @ dy @ dxf.T / (nmem-1) elif da == "srf": I = np.eye(xf_.size) #p0 = np.zeros_like(pf) #p0 = pf[:,:] dx0 = np.zeros_like(dxf) dx0 = dxf[:, :] dy0 = np.zeros_like(dy) dy0 = dy[:, :] x0_ = np.zeros_like(xf_) x0_ = xf_[:] d0 = np.zeros_like(d) d0 = d[:] if loc: logger.info("==localization==") dist, l_mat = loc_mat(sigma=4.0, nx=xf_.size, ny=y.size) #for i in range(y.size): i = 0 while i < y.size: op = ytype[i] #for i in range(y.shape[0]): #hrow = JH[i].reshape(1,-1) dyi = dy0[i, :].reshape(1, -1) d1 = dyi @ dyi.T / (nmem - 1) + sig * sig #k1 = p0 @ hrow.T /d1 k1 = dx0 @ dyi.T / d1 / (nmem - 1) if loc: k1 = k1 * l_mat[:, i].reshape(k1.shape) xa_ = x0_.reshape(k1.shape) + k1 * d0[i] #d1 = hrow @ p0 @ hrow.T + sig*sig if not tlm or not l_jhi: k1_ = k1 / (1.0 + sig / np.sqrt(d1)) #dxa = (I - k1_@hrow) @ dx0 dxa = dx0 - k1_ @ dyi else: for j in range(nmem): dyji = dyj[j, i, :].reshape(1, -1) d1 = dyji @ dyji.T / (nmem - 1) + sig * sig kj = dx0 @ dyji.T / (nmem - 1) / (d1 + sig * np.sqrt(d1)) dxa[:, j] = dx0[:, j] - kj @ dyji[:, j] x0_ = xa_[:] dx0 = dxa[:, :] x0 = x0_ + dx0 i += 1 if i < y.size: if tlm: op = ytype[i] logger.debug(op) if not l_jhi: logger.debug("linearized about ensemble mean") dy0[i, :] = obs.dhdx(x0_, op) @ dx0 else: logger.debug("linearized about each ensemble member") for j in range(nmem): jhi = obs.dhdx(x0[:, j], op, ga) dyj[j, i, :] = jhi @ dx0 dy0[i, :] = obs.dhdx(x0_, op) @ dx0 else: op = ytype[i] logger.debug(op) dy0[i, :] = obs.h_operator(x0, op) - np.mean( obs.h_operator(x0, op), axis=1)[:, None] d0 = y - np.mean(obs.h_operator(x0, op), axis=1) #pa = [email protected]/(nmem-1) xa = dxa + xa_ xa_ = np.squeeze(xa_) elif da == "letkf": #r0 = dx * 10.0 # local observation max distance sigma = 7.5 #r0 = 2.0 * np.sqrt(10.0/3.0) * sigma r0 = 100.0 # all if loc: r0 = 5.0 nx = xf_.size dist, l_mat = loc_mat(sigma, nx, ny=y.size) #dist, l_mat = loc_mat(sigma, nx, ny=y.shape[0]) logger.debug(dist[0]) xa = np.zeros_like(xf) xa_ = np.zeros_like(xf_) dxa = np.zeros_like(dxf) E = np.eye(nmem) #hx = obs.h_operator(xf_, op, ga) hx = np.mean(obs.h_operator(xf, op, ga), axis=1) #hx = np.mean(obs.h_operator(y[:,0], xf), axis=1) if infl: logger.info("==inflation==") E /= alpha for i in range(nx): far = np.arange(y.size) #far = np.arange(y.shape[0]) far = far[dist[i] > r0] logger.info("number of assimilated obs.={}".format(y.size - len(far))) #print("number of assimilated obs.={}".format(y.shape[0] - len(far))) #yi = np.delete(y,far) #if tlm: # Hi = np.delete(JH,far,axis=0) # di = yi - Hi @ xf_ # dyi = Hi @ dxf #else: # hxi = np.delete(hx,far) # di = yi - hxi dyi = np.delete(dy, far, axis=0) di = np.delete(d, far) Ri = np.delete(R, far, axis=0) Ri = np.delete(Ri, far, axis=1) if loc: logger.info("==localization==") diagR = np.diag(Ri) l = np.delete(l_mat[i], far) Ri = np.diag(diagR / l) R_inv = la.inv(Ri) A = (nmem - 1) * E + dyi.T @ R_inv @ dyi lam, v = la.eigh(A) D_inv = np.diag(1.0 / lam) pa_ = v @ D_inv @ v.T xa_[i] = xf_[i] + dxf[i] @ pa_ @ dyi.T @ R_inv @ di sqrtPa = v @ np.sqrt(D_inv) @ v.T * np.sqrt(nmem - 1) dxa[i] = dxf[i] @ sqrtPa xa[i] = np.full(nmem, xa_[i]) + dxa[i] if infl_r: sigb = np.sum(np.diag(dy @ dy.T / (nmem - 1))) sigo = sig * sig * y.size logger.debug(f"{sigb}, {sigo}") if sigb / sigo < 1e-3: #|HPfHt| << |R| beta = 0.5 else: beta = 1.0 / (1.0 + np.sqrt(sigo / (sigo + sigb))) logger.info("relaxation factor = {}".format(beta)) dxa = beta * dxa + (1.0 - beta) * dxf xa = xa_[:, None] + dxa pa = dxa @ dxa.T / (nmem - 1) if save_dh: #logger.info("save pa") np.save("{}_pa_{}_{}_cycle{}.npy".format(model, op, da, icycle), pa) np.save("{}_ua_{}_{}_cycle{}.npy".format(model, op, da, icycle), xa) #innv = y - obs.h_operator(xa_, op, ga) innv = y - np.mean(obs.h_operator(xa, op, ga), axis=1) #innv = y[:,1] - np.mean(obs.h_operator(y[:,0], xa), axis=1) p = innv.size G = dy @ dy.T + R chi2 = innv.T @ la.inv(G) @ innv / p ds = dof(dy, sig, nmem) logger.info("ds={}".format(ds)) #if len(innv.shape) > 1: # Rsim = innv.reshape(innv.size,1) @ d[None,:] #else: # Rsim = innv[:,None] @ d[None,:] #chi2 = np.mean(np.diag(Rsim)) / sig / sig return xa, xa_, pa #, chi2, ds, condh
def gen_obs(u, sigma, op): seed = 514 y = add_noise(h_operator(u, op), sigma, seed=seed) return y
def analysis(xf, xc, y, ytype, rmat, rinv, htype, gtol=1e-6, method="CG", cgtype=None, maxiter=None, restart=True, maxrest=20, disp=False, save_hist=False, save_dh=False, infl=False, loc = False, infl_parm=1.0, tlm=False, l_jhi=False, model="z08", icycle=0): global zetak, alphak zetak = [] alphak = [] #op = htype["operator"] pt = htype["perturbation"] pf = xf - xc[:, None] nmem = pf.shape[1] # logger.debug("norm(pf)={}".format(la.norm(pf))) if infl: logger.info("==inflation==, alpha={}".format(infl_parm)) pf *= infl_parm #if pt == "grad": dh = np.zeros((y.size,pf.shape[1])) if tlm: # logger.debug("dhdx.shape={}".format(obs.dhdx(xc, op).shape)) for i in range(len(ytype)): op = ytype[i] logger.debug(op) if not l_jhi: logger.debug("linearized about control") dh[i,:] = obs.dhdx(xc, op) @ pf else: logger.debug("linearized about each ensemble member") for j in range(nmem): jhi = obs.dhdx(xf[:,j], op) dh[i,j] = jhi @ pf[:, j] #dh = obs.dhdx(xc, op) @ pf else: for i in range(len(ytype)): op = ytype[i] logger.debug(op) dh[i,:] = obs.h_operator(xf, op) - obs.h_operator(xc, op)[:, None] #dh = obs.h_operator(xf, op) - obs.h_operator(xc, op)[:, None] zmat = rmat @ dh # logger.debug("cond(zmat)={}".format(la.cond(zmat))) tmat, heinv = precondition(zmat) # logger.debug("pf.shape={}".format(pf.shape)) # logger.debug("tmat.shape={}".format(tmat.shape)) # logger.debug("heinv.shape={}".format(heinv.shape)) gmat = pf @ tmat # logger.debug("gmat.shape={}".format(gmat.shape)) x0 = np.zeros(xf.shape[1]) x = x0 args_j = (xc, pf, y, ytype, tmat, gmat, heinv, rinv, tlm, l_jhi) iprint = np.zeros(2, dtype=np.int32) irest = 0 # restart counter flg = -1 # optimization result flag #if icycle != 0: # logger.info("calculate gradient") minimize = Minimize(x0.size, calc_j, jac=calc_grad_j, hess=calc_hess, args=args_j, iprint=iprint, method=method, cgtype=cgtype, maxiter=maxiter, restart=restart) #else: # logger.info("finite difference approximation") # minimize = Minimize(x0.size, calc_j, jac=None, hess=None, # args=args_j, iprint=iprint, method=Method, cgtype=cgtype, # maxiter=maxiter) logger.info("save_hist={}".format(save_hist)) # print("save_hist={}".format(save_hist)) cg = spo.check_grad(calc_j, calc_grad_j, x0, *args_j) logger.info("check_grad={}".format(cg)) if restart: if save_hist: jh = [] gh = [] while irest < maxrest: zetak = [] xold = x if save_hist: x, flg = minimize(x0, callback=callback) else: x, flg = minimize(x0) irest += 1 if save_hist: for i in range(len(zetak)): jh.append(calc_j(np.array(zetak[i]), *args_j)) g = calc_grad_j(np.array(zetak[i]), *args_j) gh.append(np.sqrt(g.transpose() @ g)) if flg == 0: logger.info(f"Converged at {irest}th restart") break xup = x - xold #logger.debug(f"update : {xup}") if np.sqrt(np.dot(xup,xup)) < 1e-10: logger.info(f"Stagnation at {irest}th : solution not updated") break xa = xc + gmat @ x xc = xa if tlm: #if pt == "grad": for i in range(len(ytype)): op = ytype[i] logger.debug(op) if not l_jhi: logger.debug("linearized about control") dh[i,:] = obs.dhdx(xc, op) @ pf else: logger.debug("linearized about each ensemble member") for j in range(nmem): jhi = obs.dhdx(xf[:,j], op) dh[i,j] = jhi @ pf[:, j] #dh = obs.dhdx(xa, op) @ pf else: for i in range(len(ytype)): op = ytype[i] logger.debug(op) dh[i,:] = obs.h_operator(xc[:, None] + pf, op) - obs.h_operator(xc, op)[:, None] zmat = rmat @ dh tmat, heinv = precondition(zmat) #pa = pf @ tmat #pf = pa gmat = pf @ tmat args_j = (xc, pf, y, ytype, tmat, gmat, heinv, rinv, tlm, l_jhi) x0 = np.zeros(xf.shape[1]) #reload(minimize) minimize = Minimize(x0.size, calc_j, jac=calc_grad_j, hess=calc_hess, args=args_j, iprint=iprint, method=method, cgtype=cgtype, maxiter=maxiter, restart=restart) else: #x, flg = minimize_cg(calc_j, x0, args=args_j, jac=calc_grad_j, # calc_step=calc_step, callback=callback) #x, flg = minimize_bfgs(calc_j, x0, args=args_j, jac=calc_grad_j, # calc_step=calc_step, callback=callback) if save_hist: x, flg = minimize(x0, callback=callback) jh = np.zeros(len(zetak)) gh = np.zeros(len(zetak)) for i in range(len(zetak)): jh[i] = calc_j(np.array(zetak[i]), *args_j) g = calc_grad_j(np.array(zetak[i]), *args_j) gh[i] = np.sqrt(g.transpose() @ g) logger.debug(f"zetak={len(zetak)} alpha={len(alphak)}") np.savetxt("{}_jh_{}_{}_cycle{}.txt".format(model, op, pt, icycle), jh) np.savetxt("{}_gh_{}_{}_cycle{}.txt".format(model, op, pt, icycle), gh) np.savetxt("{}_alpha_{}_{}_cycle{}.txt".format(model, op, pt, icycle), alphak) else: x, flg = minimize(x0) xa = xc + gmat @ x if tlm: #if pt == "grad": for i in range(len(ytype)): op = ytype[i] logger.debug(op) if not l_jhi: logger.debug("linearized about control") dh[i,:] = obs.dhdx(xa, op) @ pf else: logger.debug("linearized about each ensemble member") for j in range(nmem): jhi = obs.dhdx(xa+pf[:,j], op) dh[i,j] = jhi @ pf[:, j] #dh = obs.dhdx(xa, op) @ pf else: for i in range(len(ytype)): op = ytype[i] logger.debug(op) dh[i,:] = obs.h_operator(xa[:, None] + pf, op) - obs.h_operator(xa, op)[:, None] zmat = rmat @ dh # logger.debug("cond(zmat)={}".format(la.cond(zmat))) tmat, heinv = precondition(zmat)#, plot=True) d = np.zeros_like(y) for i in range(len(ytype)): op = ytype[i] d[i] = y[i] - obs.h_operator(xa, op) chi2 = chi2_test(zmat, heinv, rmat, d) pa = pf @ tmat return xa, pa, chi2
def plot_wind(xf, xf_, xa, xa_, y, sig, ytype, htype, var): theta = np.linspace(0.0, 2.0 * np.pi, 360) ind_speed = ytype.index("speed") ind_u = ytype.index("u") rmin = y[ind_speed] - sig rmax = y[ind_speed] + sig xmin = rmin * np.cos(theta) ymin = rmin * np.sin(theta) xmax = rmax * np.cos(theta) ymax = rmax * np.sin(theta) x = np.arange(-5, 11) fig, ax = plt.subplots(1, 2) ax[0].scatter(xf[0, :], xf[1, :], s=5) ax[0].scatter(xf_[0], xf_[1], s=30, marker='^') #ax[0].quiver(xf_[0], xf_[1], -djx[0], -djx[1], angles="xy", color="red", scale_units="xy", scale=5.0) ax[0].set_xlabel("u") ax[0].set_ylabel("v") ax[0].set_aspect("equal") ax[0].set_xticks(x[::5]) ax[0].set_yticks(x[::5]) ax[0].set_xticks(x, minor=True) ax[0].set_yticks(x, minor=True) ax[0].grid(which="major") ax[0].grid(which="minor", linestyle="dashed") ax[0].set_title("initial ensemble") ax[1].plot(xmin, ymin, color="black", linestyle="dashed") ax[1].plot(xmax, ymax, color="black", linestyle="dashed") if y.size > 1: umin = y[ind_u] - sig umax = y[ind_u] + sig ax[1].axvline(x=umin, color="black", linestyle="dashed") ax[1].axvline(x=umax, color="black", linestyle="dashed") ax[1].scatter(xa[0, :], xa[1, :], s=10, marker='*') #, zorder=2.5) ax[1].scatter(xa_[0], xa_[1], s=30, marker='^') #, zorder=2.5) ax[1].set_xlabel("u") ax[1].set_ylabel("v") ax[1].set_aspect("equal") ax[1].set_xticks(x[::5]) ax[1].set_yticks(x[::5]) ax[1].set_xticks(x, minor=True) ax[1].set_yticks(x, minor=True) ax[1].grid(which="major") ax[1].grid(which="minor", linestyle="dashed") ax[1].set_title(htype["perturbation"] + " analysis") #ax[1].set_title(r"etkf analysis $\mathbf{H}_i \delta \mathbf{X}^f$") fig.tight_layout() fig.savefig("{}_speed_{}.png".format(htype["perturbation"], var)) plt.close() sa = np.zeros((y.size, xa.shape[1])) for i in range(y.size): op = ytype[i] sa[i] = obs.h_operator(xa, op) sa -= y[:, None] #[0] vlim = 3.0 #max(np.max(sa), -np.min(sa)) if y.size <= 1: plt.hist(sa[0], bins=100, density=True, range=(-vlim, vlim)) plt.title(r"$y-H(x^a_i)$") plt.savefig("{}_hist_{}.png".format(htype["perturbation"], var)) plt.close() else: fig, ax = plt.subplots(nrows=2) for i in range(y.size): ax[i].hist(sa[i], bins=100, density=True, range=(-vlim, vlim)) ax[i].set_title(ytype[i]) fig.suptitle(r"$y-H(x^a_i)$") fig.tight_layout() fig.savefig("{}_hist_{}.png".format(htype["perturbation"], var)) plt.close() for i in range(y.size): print(ytype[i]) print("O-A mean={}".format(np.mean(sa[i]))) stdv = np.sqrt(np.mean(sa[i]**2) - np.mean(sa[i])**2) print("O-A stdv={}".format(stdv)) ut = y[ind_u] st = y[ind_speed] vtp = np.sqrt(st**2 - ut**2) vtm = -vtp print(vtp, vtm) tan = xa[0] / xa[1] fig, ax = plt.subplots() ax.hist(tan, bins=100, density=True, range=(-vlim, vlim)) ax.axvline(x=1.0 / vtp, color="red", linestyle="dashed") ax.axvline(x=1.0 / vtm, color="red", linestyle="dashed") ax.set_title(r"$\tan \theta$") fig.savefig("{}_tan_{}.png".format(htype["perturbation"], var)) plt.close()
def minimize_newton(fun, w0, args=(), jac=None, hess=None, callback=None, gtol=1e-5, maxiter=None, disp=False, delta=1e-10, mu=0.5, c1=1e-3, c2=0.9): _status_message = { 'success': 'Optimization terminated successfully.', 'maxfev': 'Maximum number of function evaluations has ' 'been exceeded.', 'maxiter': 'Maximum number of iterations has been ' 'exceeded.', 'pr_loss': 'Desired error not necessarily achieved due ' 'to precision loss.', 'nan': 'NaN result encountered.', 'out_of_bounds': 'The result is outside of the provided ' 'bounds.' } w0 = np.asarray(w0).flatten() if maxiter is None: maxiter = 100 old_fval = fun(w0, *args) gfk = jac(w0, *args) nfev = 1 ngev = 1 k = 0 wk = w0 warnflag = 0 gnorm = np.sqrt(np.dot(gfk, gfk)) phik = min(mu, np.sqrt(gnorm)) xc, pf, y, rmat, htype = args if htype["perturbation"] == "gradw": dh = obs.dhdx(xc, htype["operator"]) @ pf else: dh = obs.h_operator(xc[:, None] + pf, htype["operator"]) - obs.h_operator( xc, htype["operator"])[:, None] zmat = rmat @ dh tmat, heinv, Mk = precondition(zmat) #Mk = np.eye(w0.size) logger.debug(f"preconditioner={Mk}") while (gnorm > gtol) and (k < maxiter): Hk = hess(wk, *args) lam, v = la.eigh(Hk) #lam[:len(lam)-1] = 0.0 #Mk = v @ np.diag(lam) @ v.transpose() logger.debug(f"Hessian eigen values={lam.max(), lam.min()}") #pk = la.solve(Hk, -gfk) #eps = phik * gnorm eps = 1e-8 pk = pcg(gfk, Hk, Mk, delta=delta, eps=eps) rk = Hk @ pk + gfk logger.debug(f"residual:{np.sqrt(np.dot(rk, rk))}") alphak = back_tracking(fun, wk, pk, gfk, old_fval, nfev, args, c1, c2) #alphak = 1.0 wkp1 = wk + alphak * pk gfkp1 = jac(wkp1, *args) gfk = gfkp1 wk = wkp1 ngev += 1 gnorm = np.sqrt(np.dot(gfk, gfk)) old_fval = fun(wkp1, *args) nfev += 1 if callback is not None: callback(wk, alphak) logger.debug( f"current:{k} gnorm={gnorm} step-length={alphak} phik={eps}") k += 1 phik = min(mu / k, np.sqrt(gnorm)) fval = fun(wk, *args) nfev += 1 if warnflag == 2: msg = _status_message['pr_loss'] elif k >= maxiter: warnflag = 1 msg = _status_message['maxiter'] elif np.isnan(gnorm) or np.isnan(fval) or np.isnan(wk).any(): warnflag = 3 msg = _status_message['nan'] else: msg = _status_message['success'] if disp: logger.info("%s%s" % ("Warning: " if warnflag != 0 else "", msg)) logger.info(" Current function value: %f" % fval) logger.info(" Iterations: %d" % k) logger.info(" Function evaluations: %d" % nfev) logger.info(" Gradient evaluations: %d" % ngev) else: logger.info("success={} message={}".format(warnflag == 0, msg)) logger.info("J={:7.3e} dJ={:7.3e} nit={}".format( \ fval, gnorm, k)) return wk
y = np.array([3.0]) #observation [m/s] sig = 0.3 #observation error [m/s] rmat = np.array([1.0 / sig]).reshape(-1, 1) rinv = np.array([1.0 / sig / sig]).reshape(-1, 1) np.random.seed(514) wind = random.multivariate_normal(wmean, np.diag(wstdv), size=nmem) xf = wind.transpose() xfc = wmean xf_ = np.mean(xf, axis=1) pf = xf - xfc[:, None] dxf = (xf - xf_[:, None]) / np.sqrt(nmem - 1) x0 = np.zeros(nmem) # mlef dh = obs.h_operator(xf, htype["operator"], htype["gamma"]) - obs.h_operator( xfc, htype["operator"], htype["gamma"])[:, None] zmat = rmat @ dh tmat, heinv, condh = precondition(zmat) gmat = pf @ tmat args = (xfc, pf, y, tmat, gmat, heinv, rinv, htype) g = calc_grad_j(x0, *args) alpha = np.linspace(-0.01, 0.01, 101, endpoint=True).tolist() f = f_ls(x0, -g, alpha, calc_j, *args) print(f"max={f.max()} min={f.min()}") fig = plt.figure(figsize=[8, 8]) ax = fig.add_subplot() levs = (np.arange(10) + 1) * 5000.0 cntr = ax.contourf(alpha, alpha, f, levs) ax.quiver(0.0, 0.0, 1.0, 0.0, color="red") ax.set_xticks(alpha[::10])
def analysis(xf, xf_, y, sig, dx, htype, infl=False, loc=False, tlm=True, infl_parm=1.1, \ save_dh=False, model="z08", icycle=0): op = htype["operator"] da = htype["perturbation"] ga = htype["gamma"] #obs = Obs(op, sig) JH = obs.dhdx(xf_, op, ga) #JH = obs.dh_operator(y[:,0], xf_) R = np.eye(y.size) * sig * sig rinv = np.eye(y.size) / sig / sig #R, rmat, rinv = obs.set_r(y.shape[0]) nmem = xf.shape[1] dxf = xf - xf_[:, None] logger.debug(dxf.shape) if tlm: dy = JH @ dxf #dy = np.zeros((y.size,nmem)) #for i in range(nmem): # jhi = obs.dhdx(xf[:,i], op, ga) # dy[:,i] = jhi @ dxf[:,i] else: dy = obs.h_operator(xf, op, ga) - np.mean(obs.h_operator(xf, op, ga), axis=1)[:, None] #dy = obs.h_operator(y[:,0], xf) - np.mean(obs.h_operator(y[:,0], xf), axis=1)[:, None] #dy = obs.h_operator(xf, op, ga) - obs.h_operator(xf_, op, ga)[:, None] d = y - np.mean(obs.h_operator(xf, op, ga), axis=1) #d = y[:,1] - np.mean(obs.h_operator(y[:,0], xf), axis=1) #d = y - obs.h_operator(xf_, op, ga) hes = np.eye(nmem) + dy.T @ rinv @ dy condh = la.cond(hes) # inflation parameter alpha = infl_parm if infl: if da != "letkf" and da != "etkf": logger.info("==inflation==, alpha={}".format(alpha)) dxf *= alpha pf = dxf @ dxf.T / (nmem - 1) if save_dh: #logger.info("save pf") np.save("{}_pf_{}_{}_cycle{}.npy".format(model, op, da, icycle), pf) #if loc: # B-localization # if da == "etkf": # print("==B-localization==") # dist, l_mat = loc_mat(sigma=2.0, nx=xf_.size, ny=xf_.size) # pf = pf * l_mat if save_dh: #logger.info("save dxf") np.save("{}_dxf_{}_{}_cycle{}.npy".format(model, op, da, icycle), dxf) np.save("{}_dhdx_{}_{}_cycle{}.npy".format(model, op, da, icycle), JH) np.save("{}_dh_{}_{}_cycle{}.npy".format(model, op, da, icycle), dy) np.save("{}_d_{}_{}_cycle{}.npy".format(model, op, da, icycle), d) logger.info("save_dh={} cycle{}".format(save_dh, icycle)) # logger.debug("norm(pf)={}".format(la.norm(pf))) if da == "etkf": #if tlm: # K1 = pf @ JH.T # K2 = JH @ pf @ JH.T + R # #K = pf @ JH.T @ la.inv(JH @ pf @ JH.T + R) #else: """ K1 = dxf @ dy.T / (nmem-1) K2 = dy @ dy.T / (nmem-1) + R #K = dxf @ dy.T @ la.inv(dy @ dy.T + (nmem-1)*R) K2inv = la.inv(K2) K = K1 @ K2inv if save_dh: print("save K") np.save("{}_K_{}_{}_cycle{}.npy".format(model, op, da, icycle), K) np.save("{}_K1_{}_{}_cycle{}.npy".format(model, op, da, icycle), K1) np.save("{}_K2_{}_{}_cycle{}.npy".format(model, op, da, icycle), K2) np.save("{}_K2i_{}_{}_cycle{}.npy".format(model, op, da, icycle), K2inv) if loc: # K-localization print("==K-localization==") dist, l_mat = loc_mat(sigma=2.0, nx=xf_.size, ny=y.size) #print(l_mat) K = K * l_mat #print(K) if save_dh: np.save("{}_Kloc_{}_{}_cycle{}.npy".format(model, op, da, icycle), K) #xa_ = xf_ + K @ d if save_dh: print("save increment") np.save("{}_dx_{}_{}_cycle{}.npy".format(model, op, da, icycle), K@d) #TT = la.inv( np.eye(nmem) + dy.T @ la.inv(R) @ dy / (nmem-1) ) #lam, v = la.eigh(TT) #print("eigenvalue(sqrt)={}".format(np.sqrt(lam))) #D = np.diag(np.sqrt(lam)) #T = v @ D """ if infl: logger.info("==inflation==, alpha={}".format(alpha)) A = np.eye(nmem) / alpha else: A = np.eye(nmem) A = (nmem - 1) * A + dy.T @ rinv @ dy #TT = la.inv( np.eye(nmem) + dy.T @ rinv @ dy / (nmem-1) ) lam, v = la.eigh(A) #print("eigenvalue(sqrt)={}".format(np.sqrt(lam))) Dinv = np.diag(1.0 / lam) TT = v @ Dinv @ v.T T = v @ np.sqrt((nmem - 1) * Dinv) @ v.T K = dxf @ TT @ dy.T @ rinv if save_dh: #logger.info("save K") np.save("{}_K_{}_{}_cycle{}.npy".format(model, op, da, icycle), K) if loc: # K-localization logger.info("==K-localization==") dist, l_mat = loc_mat(sigma=4.0, nx=xf_.size, ny=y.size) #print(l_mat) K = K * l_mat #print(K) if save_dh: np.save( "{}_Kloc_{}_{}_cycle{}.npy".format(model, op, da, icycle), K) if save_dh: #logger.info("save increment") np.save("{}_dx_{}_{}_cycle{}.npy".format(model, op, da, icycle), K @ d) xa_ = xf_ + K @ d dxa = dxf @ T xa = dxa + xa_[:, None] # pa = [email protected]/(nmem-1) elif da == "po": Y = np.zeros((y.size, nmem)) #Y = np.zeros((y.shape[0],nmem)) #np.random.seed(514) #err = np.random.normal(0, scale=sig, size=Y.size) #err_ = np.mean(err.reshape(Y.shape), axis=1) err = np.zeros_like(Y) for j in range(nmem): err[:, j] = np.random.normal(0.0, scale=sig, size=err.shape[0]) err_ = np.mean(err, axis=1) stdv = np.sqrt(np.mean((err - err_[:, None])**2, axis=1)) logger.debug("err mean {}".format(err_)) logger.debug("err stdv {}".format(stdv)) #Y = y[:,None] + err.reshape(Y.shape) Y = y[:, None] + err #Y = y[:,1].reshape(-1,1) + err.reshape(Y.shape) #d_ = y + err_ - obs.h_operator(xf_, op, ga) d_ = d #+ err_ if tlm: K = pf @ JH.T @ la.inv(JH @ pf @ JH.T + R) else: K = dxf @ dy.T @ la.inv(dy @ dy.T + (nmem - 1) * R) if loc: logger.info("==localization==") dist, l_mat = loc_mat(sigma=4.0, nx=xf_.size, ny=y.size) K = K * l_mat xa_ = xf_ + K @ d_ if tlm: xa = xf + K @ (Y - JH @ xf) # pa = (np.eye(xf_.size) - K @ JH) @ pf else: HX = obs.h_operator(xf, op, ga) #HX = obs.h_operator(y[:,0], xf) xa = xf + K @ (Y - HX) # pa = pf - K @ dy @ dxf.T / (nmem-1) dxa = xa - xa_[:, None] elif da == "srf": I = np.eye(xf_.size) #p0 = np.zeros_like(pf) #p0 = pf[:,:] dx0 = np.zeros_like(dxf) dx0 = dxf[:, :] dy0 = np.zeros_like(dy) dy0 = dy[:, :] x0_ = np.zeros_like(xf_) x0_ = xf_[:] d0 = np.zeros_like(d) d0 = d[:] if loc: logger.info("==localization==") dist, l_mat = loc_mat(sigma=4.0, nx=xf_.size, ny=y.size) for i in range(y.size): #for i in range(y.shape[0]): #hrow = JH[i].reshape(1,-1) dyi = dy0[i, :].reshape(1, -1) #d1 = hrow @ p0 @ hrow.T + sig*sig d1 = dyi @ dyi.T / (nmem - 1) + sig * sig #k1 = p0 @ hrow.T /d1 k1 = dx0 @ dyi.T / d1 / (nmem - 1) k1_ = k1 / (1.0 + sig / np.sqrt(d1)) if loc: k1_ = k1_ * l_mat[:, i].reshape(k1_.shape) #xa_ = x0_.reshape(k1_.shape) + k1_ * (y[i] - hrow@x0_) #xa_ = x0_.reshape(k1_.shape) + k1_ * d0[i] xa_ = x0_.reshape(k1.shape) + k1 * d0[i] #dxa = (I - k1_@hrow) @ dx0 dxa = dx0 - k1_ @ dyi x0_ = xa_[:] dx0 = dxa[:, :] x0 = x0_ + dx0 #if tlm: # logger.debug(x0_.shape) # logger.debug(obs.dhdx(np.squeeze(x0_), op, ga).shape) # logger.debug(dx0.shape) # dy0 = obs.dhdx(np.squeeze(x0_), op, ga) @ dx0 # #dy0 = obs.dh_operator(y[:,0], x0_) @ dx0 # #dy0 = np.zeros((y.size,nmem)) # #for i in range(nmem): # # jhi = obs.dhdx(x0[:,i], op, ga) # # dy0[:,i] = jhi @ dx0[:,i] #else: # dy0 = obs.h_operator(x0, op, ga) - np.mean(obs.h_operator(x0, op, ga), axis=1)[:, None] # #dy0 = obs.h_operator(y[:,0], x0) - np.mean(obs.h_operator(y[:,0], x0), axis=1)[:, None] # #dy0 = obs.h_operator(x0, op, ga) - obs.h_operator(x0_, op, ga)#[:, None] #d0 = y - np.mean(obs.h_operator(x0, op, ga), axis=1) #d0 = y[:,1] - np.mean(obs.h_operator(y[:,0], x0), axis=1) #d0 = y - np.squeeze(obs.h_operator(x0_, op, ga)) #p0 = pa[:,:] #print(x0_.shape) #print(dx0.shape) #print(dy0.shape) #print(d0.shape) #print(p0.shape) #pa = [email protected]/(nmem-1) xa = dxa + xa_ xa_ = np.squeeze(xa_) elif da == "letkf": #r0 = dx * 10.0 # local observation max distance sigma = 7.5 #r0 = 2.0 * np.sqrt(10.0/3.0) * sigma r0 = 100.0 # all if loc: r0 = 5.0 nx = xf_.size dist, l_mat = loc_mat(sigma, nx, ny=y.size) #dist, l_mat = loc_mat(sigma, nx, ny=y.shape[0]) logger.debug(dist[0]) xa = np.zeros_like(xf) xa_ = np.zeros_like(xf_) dxa = np.zeros_like(dxf) E = np.eye(nmem) #hx = obs.h_operator(xf_, op, ga) hx = np.mean(obs.h_operator(xf, op, ga), axis=1) #hx = np.mean(obs.h_operator(y[:,0], xf), axis=1) if infl: logger.info("==inflation==") E /= alpha for i in range(nx): far = np.arange(y.size) #far = np.arange(y.shape[0]) far = far[dist[i] > r0] logger.info("number of assimilated obs.={}".format(y.size - len(far))) #print("number of assimilated obs.={}".format(y.shape[0] - len(far))) #yi = np.delete(y,far) #if tlm: # Hi = np.delete(JH,far,axis=0) # di = yi - Hi @ xf_ # dyi = Hi @ dxf #else: # hxi = np.delete(hx,far) # di = yi - hxi dyi = np.delete(dy, far, axis=0) di = np.delete(d, far) Ri = np.delete(R, far, axis=0) Ri = np.delete(Ri, far, axis=1) if loc: logger.info("==localization==") diagR = np.diag(Ri) l = np.delete(l_mat[i], far) Ri = np.diag(diagR / l) R_inv = la.inv(Ri) A = (nmem - 1) * E + dyi.T @ R_inv @ dyi lam, v = la.eigh(A) D_inv = np.diag(1.0 / lam) pa_ = v @ D_inv @ v.T xa_[i] = xf_[i] + dxf[i] @ pa_ @ dyi.T @ R_inv @ di sqrtPa = v @ np.sqrt(D_inv) @ v.T * np.sqrt(nmem - 1) dxa[i] = dxf[i] @ sqrtPa xa[i] = np.full(nmem, xa_[i]) + dxa[i] #pa = [email protected]/(nmem-1) pa = dxa @ dxa.T / (nmem - 1) if save_dh: #logger.info("save pa") np.save("{}_pa_{}_{}_cycle{}.npy".format(model, op, da, icycle), pa) np.save("{}_ua_{}_{}_cycle{}.npy".format(model, op, da, icycle), xa) #innv = y - obs.h_operator(xa_, op, ga) innv = y - np.mean(obs.h_operator(xa, op, ga), axis=1) #innv = y[:,1] - np.mean(obs.h_operator(y[:,0], xa), axis=1) p = innv.size G = dy @ dy.T + R chi2 = innv.T @ la.inv(G) @ innv / p ds = dof(dy, sig, nmem) logger.info("ds={}".format(ds)) #if len(innv.shape) > 1: # Rsim = innv.reshape(innv.size,1) @ d[None,:] #else: # Rsim = innv[:,None] @ d[None,:] #chi2 = np.mean(np.diag(Rsim)) / sig / sig return xa, xa_, pa, chi2, ds, condh
def analysis(xf, xc, y, rmat, rinv, htype, gtol=1e-6, method="CG", cgtype=None, maxiter=None, restart=True, maxrest=20, disp=False, save_hist=False, save_dh=False, infl=False, loc = False, infl_parm=1.0, model="z08", icycle=0): global zetak, alphak zetak = [] alphak = [] op = htype["operator"] pt = htype["perturbation"] pf = xf - xc[:, None] # logger.debug("norm(pf)={}".format(la.norm(pf))) if infl: logger.info("==inflation==, alpha={}".format(infl_parm)) pf *= infl_parm #if pt == "grad": if pt == "grad05": # logger.debug("dhdx.shape={}".format(obs.dhdx(xc, op).shape)) dh = obs.dhdx(xc, op) @ pf else: dh = obs.h_operator(xf, op) - obs.h_operator(xc, op)[:, None] if save_dh: np.save("{}_dxf_{}_{}_cycle{}.npy".format(model, op, pt, icycle), pf) np.save("{}_dh_{}_{}_cycle{}.npy".format(model, op, pt, icycle), dh) ob = y - obs.h_operator(xc, op) np.save("{}_d_{}_{}_cycle{}.npy".format(model, op, pt, icycle), ob) logger.info("save_dh={}".format(save_dh)) # print("save_dh={}".format(save_dh)) zmat = rmat @ dh # logger.debug("cond(zmat)={}".format(la.cond(zmat))) tmat, heinv = precondition(zmat) # logger.debug("pf.shape={}".format(pf.shape)) # logger.debug("tmat.shape={}".format(tmat.shape)) # logger.debug("heinv.shape={}".format(heinv.shape)) gmat = pf @ tmat # logger.debug("gmat.shape={}".format(gmat.shape)) if save_dh: np.save("{}_tmat_{}_{}_cycle{}.npy".format(model, op, pt, icycle), tmat) np.save("{}_pf_{}_{}_cycle{}.npy".format(model, op, pt, icycle), pf) np.save("{}_gmat_{}_{}_cycle{}.npy".format(model, op, pt, icycle), gmat) x0 = np.zeros(xf.shape[1]) x = x0 htype_c = htype.copy() Method = method #if icycle == 0: # htype_c["perturbation"] = "grad05" # Method="dogleg" logger.debug(f"original {htype}") logger.debug(f"copy {htype_c}") args_j = (xc, pf, y, tmat, gmat, heinv, rinv, htype_c) iprint = np.zeros(2, dtype=np.int32) irest = 0 # restart counter flg = -1 # optimization result flag #if icycle != 0: # logger.info("calculate gradient") minimize = Minimize(x0.size, calc_j, jac=calc_grad_j, hess=calc_hess, args=args_j, iprint=iprint, method=Method, cgtype=cgtype, maxiter=maxiter, restart=restart) #else: # logger.info("finite difference approximation") # minimize = Minimize(x0.size, calc_j, jac=None, hess=None, # args=args_j, iprint=iprint, method=Method, cgtype=cgtype, # maxiter=maxiter) logger.info("save_hist={}".format(save_hist)) # print("save_hist={}".format(save_hist)) cg = spo.check_grad(calc_j, calc_grad_j, x0, *args_j) logger.info("check_grad={}".format(cg)) if restart: if save_hist: jh = [] gh = [] while irest < maxrest: zetak = [] xold = x if save_hist: x, flg = minimize(x0, callback=callback) else: x, flg = minimize(x0) irest += 1 if save_hist: for i in range(len(zetak)): jh.append(calc_j(np.array(zetak[i]), *args_j)) g = calc_grad_j(np.array(zetak[i]), *args_j) gh.append(np.sqrt(g.transpose() @ g)) if flg == 0: logger.info(f"Converged at {irest}th restart") break xup = x - xold #logger.debug(f"update : {xup}") if np.sqrt(np.dot(xup,xup)) < 1e-10: logger.info(f"Stagnation at {irest}th : solution not updated") break xa = xc + gmat @ x xc = xa if pt == "grad05": dh = obs.dhdx(xc, op) @ pf else: dh = obs.h_operator(xc[:, None] + pf, op) - obs.h_operator(xc, op)[:, None] zmat = rmat @ dh tmat, heinv = precondition(zmat) #pa = pf @ tmat #pf = pa gmat = pf @ tmat args_j = (xc, pf, y, tmat, gmat, heinv, rinv, htype_c) x0 = np.zeros(xf.shape[1]) #reload(minimize) minimize = Minimize(x0.size, calc_j, jac=calc_grad_j, hess=calc_hess, args=args_j, iprint=iprint, method=Method, cgtype=cgtype, maxiter=maxiter, restart=restart) if save_hist: logger.debug(f"zetak={len(zetak)} alpha={len(alphak)}") np.savetxt("{}_jh_{}_{}_cycle{}.txt".format(model, op, pt, icycle), jh) np.savetxt("{}_gh_{}_{}_cycle{}.txt".format(model, op, pt, icycle), gh) np.savetxt("{}_alpha_{}_{}_cycle{}.txt".format(model, op, pt, icycle), alphak) if model=="z08": if len(zetak) > 0: xmax = max(np.abs(np.min(zetak)),np.max(zetak)) else: xmax = max(np.abs(np.min(x)),np.max(x)) logger.debug("resx max={}".format(xmax)) # print("resx max={}".format(xmax)) #if xmax < 1000: #cost_j(1000, xf.shape[1], model, x, icycle, *args_j) #cost_j2d(1000, xf.shape[1], model, x, icycle, *args_j) # costJ.cost_j2d(1000, xf.shape[1], model, x, icycle, # htype, calc_j, calc_grad_j, *args_j) #else: #xmax = int(xmax*0.01+1)*100 xmax = np.ceil(xmax*0.01)*100 logger.debug("resx max={}".format(xmax)) # print("resx max={}".format(xmax)) #cost_j(xmax, xf.shape[1], model, x, icycle, *args_j) #cost_j2d(xmax, xf.shape[1], model, x, icycle, *args_j) costJ.cost_j2d(xmax, xf.shape[1], model, np.array(zetak), icycle, htype, calc_j, calc_grad_j, *args_j) elif model=="l96": cost_j(200, xf.shape[1], model, x, icycle, *args_j) else: #x, flg = minimize_cg(calc_j, x0, args=args_j, jac=calc_grad_j, # calc_step=calc_step, callback=callback) #x, flg = minimize_bfgs(calc_j, x0, args=args_j, jac=calc_grad_j, # calc_step=calc_step, callback=callback) if save_hist: x, flg = minimize(x0, callback=callback) jh = np.zeros(len(zetak)) gh = np.zeros(len(zetak)) for i in range(len(zetak)): jh[i] = calc_j(np.array(zetak[i]), *args_j) g = calc_grad_j(np.array(zetak[i]), *args_j) gh[i] = np.sqrt(g.transpose() @ g) logger.debug(f"zetak={len(zetak)} alpha={len(alphak)}") np.savetxt("{}_jh_{}_{}_cycle{}.txt".format(model, op, pt, icycle), jh) np.savetxt("{}_gh_{}_{}_cycle{}.txt".format(model, op, pt, icycle), gh) np.savetxt("{}_alpha_{}_{}_cycle{}.txt".format(model, op, pt, icycle), alphak) if model=="z08": if len(zetak) > 0: xmax = max(np.abs(np.min(zetak)),np.max(zetak)) else: xmax = max(np.abs(np.min(x)),np.max(x)) logger.debug("resx max={}".format(xmax)) # print("resx max={}".format(xmax)) #if xmax < 1000: #cost_j(1000, xf.shape[1], model, x, icycle, *args_j) #cost_j2d(1000, xf.shape[1], model, x, icycle, *args_j) # costJ.cost_j2d(1000, xf.shape[1], model, x, icycle, # htype, calc_j, calc_grad_j, *args_j) #else: #xmax = int(xmax*0.01+1)*100 xmax = np.ceil(xmax*0.01)*100 logger.debug("resx max={}".format(xmax)) # print("resx max={}".format(xmax)) #cost_j(xmax, xf.shape[1], model, x, icycle, *args_j) #cost_j2d(xmax, xf.shape[1], model, x, icycle, *args_j) costJ.cost_j2d(xmax, xf.shape[1], model, np.array(zetak), icycle, htype, calc_j, calc_grad_j, *args_j) elif model=="l96": cost_j(200, xf.shape[1], model, x, icycle, *args_j) else: x, flg = minimize(x0) xa = xc + gmat @ x if pt == "grad05": #if pt == "grad": dh = obs.dhdx(xa, op) @ pf else: dh = obs.h_operator(xa[:, None] + pf, op) - obs.h_operator(xa, op)[:, None] zmat = rmat @ dh # logger.debug("cond(zmat)={}".format(la.cond(zmat))) tmat, heinv = precondition(zmat) d = y - obs.h_operator(xa, op) chi2 = chi2_test(zmat, heinv, rmat, d) pa = pf @ tmat if save_dh: logger.info("save ua") ua = np.zeros((xa.size,pf.shape[1]+1)) ua[:,0] = xa ua[:,1:] = xa[:, None] + pa np.save("{}_ua_{}_{}_cycle{}.npy".format(model, op, pt, icycle), ua) return xa, pa, chi2
def analysis(xf, xc, y, rmat, rinv, htype, gtol=1e-6, method="CG", cgtype=None, maxiter=None, restart=True, maxrest=20, disp=False, save_hist=False, save_dh=False, infl=False, loc=False, infl_parm=1.0, model="z08", icycle=0): global wk, alphak wk = [] alphak = [] op = htype["operator"] pt = htype["perturbation"] nmem = xf.shape[1] pf = xf - xc[:, None] #pf = (xf - xc[:, None]) / np.sqrt(nmem) # logger.debug("norm(pf)={}".format(la.norm(pf))) if infl: logger.info("==inflation==, alpha={}".format(infl_parm)) pf *= infl_parm if pt == "gradw": # logger.debug("dhdx.shape={}".format(obs.dhdx(xc, op).shape)) dh = obs.dhdx(xc, op) @ pf else: dh = obs.h_operator(xf, op) - obs.h_operator(xc, op)[:, None] if save_dh: np.save("{}_dh_{}_{}.npy".format(model, op, pt), dh) #logger.info("save_dh={}".format(save_dh)) zmat = rmat @ dh # logger.debug("cond(zmat)={}".format(la.cond(zmat))) tmat, heinv, prec = precondition(zmat) # print("save_dh={}".format(save_dh)) x0 = np.zeros(xf.shape[1]) args_j = (xc, pf, y, rmat, htype) iprint = np.zeros(2, dtype=np.int32) minimize = Minimize(x0.size, calc_j, jac=calc_grad_j, hess=calc_hess, prec=prec, args=args_j, iprint=iprint, method=method, cgtype=cgtype, maxiter=maxiter, restart=restart) logger.info("save_hist={}".format(save_hist)) cg = spo.check_grad(calc_j, calc_grad_j, x0, *args_j) logger.info("check_grad={}".format(cg)) # print("save_hist={}".format(save_hist)) if save_hist: #g = calc_grad_j(x0, *args_j) #print("g={}".format(g)) x, flg = minimize(x0, callback=callback) #x = minimize_newton(calc_j, x0, args=args_j, jac=calc_grad_j, hess=calc_hess, # callback=callback, maxiter=maxiter, disp=disp) logger.debug(f"wk={len(wk)} alpha={len(alphak)}") jh = np.zeros(len(wk)) gh = np.zeros(len(wk)) for i in range(len(wk)): jh[i] = calc_j(np.array(wk[i]), *args_j) g = calc_grad_j(np.array(wk[i]), *args_j) gh[i] = np.sqrt(g.transpose() @ g) np.savetxt("{}_jh_{}_{}_cycle{}.txt".format(model, op, pt, icycle), jh) np.savetxt("{}_gh_{}_{}_cycle{}.txt".format(model, op, pt, icycle), gh) np.savetxt("{}_alpha_{}_{}_cycle{}.txt".format(model, op, pt, icycle), alphak) if model == "z08": xmax = max(np.abs(np.min(x)), np.max(x)) logger.debug("resx max={}".format(xmax)) # print("resx max={}".format(xmax)) xmax = np.ceil(xmax * 0.01) * 100 logger.debug("resx max={}".format(xmax)) # print("resx max={}".format(xmax)) #cost_j(xmax, xf.shape[1], model, x, icycle, *args_j) #cost_j2d(xmax, xf.shape[1], model, x, icycle, *args_j) costJ.cost_j2d(xmax, xf.shape[1], model, x, icycle, htype, calc_j, calc_grad_j, *args_j) elif model == "l96": cost_j(200, xf.shape[1], model, x, icycle, *args_j) else: x, flg = minimize(x0) #x = minimize_newton(calc_j, x0, args=args_j, jac=calc_grad_j, hess=calc_hess, # maxiter=maxiter, disp=disp) xa = xc + pf @ x if pt == "gradw": dh = obs.dhdx(xa, op) @ pf else: dh = obs.h_operator(xa[:, None] + pf, op) - obs.h_operator(xa, op)[:, None] zmat = rmat @ dh # logger.debug("cond(zmat)={}".format(la.cond(zmat))) tmat, heinv, prec = precondition(zmat) d = y - obs.h_operator(xa, op) chi2 = chi2_test(zmat, heinv, rmat, d) pa = pf @ tmat #* np.sqrt(nmem) return xa, pa, chi2
def analysis(xf, xf_, y, rmat, rinv, htype, gtol=1e-6, method="LBFGS", maxiter=None, disp=False, save_hist=False, save_dh=False, infl=False, loc=False, tlm=False, infl_parm=1.0, model="z08", icycle=100): global zetak zetak = [] op = htype["operator"] pt = htype["perturbation"] ga = htype["gamma"] nmem = xf.shape[1] dxf = (xf - xf_[:, None]) / np.sqrt(nmem - 1) #condh = np.zeros(2) eps = 1e-4 # Bundle parameter if infl: """ if op == "linear" or op == "test": if pt == "mlef": alpha = 1.2 else: alpha = 1.2 elif op == "quadratic" or op == "quadratic-nodiff": if pt == "mlef": alpha = 1.3 else: alpha = 1.35 elif op == "cubic" or op == "cubic-nodiff": if pt == "mlef": alpha = 1.6 else: alpha = 1.65 """ alpha = infl_parm logger.info("==inflation==, alpha={}".format(alpha)) # print("==inflation==, alpha={}".format(alpha)) dxf *= alpha # logger.debug("norm(pf)={}".format(la.norm(pf))) if loc: l_sig = 4.0 logger.info("==localization==, l_sig={}".format(l_sig)) # print("==localization==, l_sig={}".format(l_sig)) dxf = pfloc(dxf, l_sig, save_dh, model, op, pt, icycle) emat = xf_[:, None] + eps * dxf hemat = obs.h_operator(emat, op, ga) dy = (hemat - np.mean(hemat, axis=1)[:, None]) / eps #dh = obs.dhdx(xf_, op, ga) @ dxf * np.sqrt(nmem-1) if save_dh: #np.save("{}_dh_{}_{}_cycle{}.npy".format(model, op, pt, icycle), dh) np.save("{}_dy_{}_{}_cycle{}.npy".format(model, op, pt, icycle), dy) ob = y - np.mean(obs.h_operator(xf, op, ga), axis=1) np.save("{}_d_{}_{}_cycle{}.npy".format(model, op, pt, icycle), ob) logger.info("save_dh={}".format(save_dh)) # print("save_dh={}".format(save_dh)) zmat = rmat @ dy # logger.debug("cond(zmat)={}".format(la.cond(zmat))) tmat, tinv, heinv, condh = precondition(zmat) #if icycle == 0: # print("tmat={}".format(tmat)) # print("tinv={}".format(tinv)) # logger.debug("pf.shape={}".format(pf.shape)) # logger.debug("tmat.shape={}".format(tmat.shape)) # logger.debug("heinv.shape={}".format(heinv.shape)) #gmat = dxf @ tmat #gmat = np.sqrt(nmem-1) * pf @ tmat # logger.debug("gmat.shape={}".format(gmat.shape)) w0 = np.zeros(xf.shape[1]) args_j = (xf_, dxf, y, precondition, eps, rmat, htype) iprint = np.zeros(2, dtype=np.int32) minimize = Minimize(w0.size, 7, calc_j, calc_grad_j, args_j, iprint, method) logger.info("save_hist={}".format(save_hist)) # print("save_hist={} cycle{}".format(save_hist, icycle)) cg = spo.check_grad(calc_j, calc_grad_j, w0, *args_j) logger.info("check_grad={}".format(cg)) # print("check_grad={}".format(cg)) if save_hist: w = minimize(w0, callback=callback) logger.debug(zetak) # print(len(zetak)) jh = np.zeros(len(zetak)) gh = np.zeros(len(zetak)) for i in range(len(zetak)): jh[i] = calc_j(np.array(zetak[i]), *args_j) g = calc_grad_j(np.array(zetak[i]), *args_j) gh[i] = np.sqrt(g.transpose() @ g) np.savetxt("{}_jh_{}_{}_cycle{}.txt".format(model, op, pt, icycle), jh) np.savetxt("{}_gh_{}_{}_cycle{}.txt".format(model, op, pt, icycle), gh) if model == "z08": xmax = max(np.abs(np.min(w)), np.max(w)) logger.debug("resx max={}".format(xmax)) # print("resx max={}".format(xmax)) xmax = np.ceil(xmax * 0.001) * 1000 logger.debug("resx max={}".format(xmax)) # print("resx max={}".format(xmax)) #cost_j(xmax, xf.shape[1], model, w, icycle, *args_j) costJ.cost_j2d(xmax, xf.shape[1], model, w, icycle, htype, calc_j, calc_grad_j, *args_j) elif model == "l96": cost_j(200, xf.shape[1], model, w, icycle, *args_j) else: w = minimize(w0) xa_ = xf_ + dxf @ w emat = xa_[:, None] + eps * dxf hemat = obs.h_operator(emat, op, ga) dy = (hemat - np.mean(hemat, axis=1)[:, None]) / eps zmat = rmat @ dy # logger.debug("cond(zmat)={}".format(la.cond(zmat))) tmat, tinv, heinv, dum = precondition(zmat) dxa = dxf @ tmat xa = xa_[:, None] + dxa * np.sqrt(nmem - 1) pa = dxa @ dxa.T if save_dh: np.save("{}_pa_{}_{}_cycle{}.npy".format(model, op, pt, icycle), pa) ua = np.zeros((xa_.size, nmem + 1)) ua[:, 0] = xa_ ua[:, 1:] = xa np.save("{}_ua_{}_{}_cycle{}.npy".format(model, op, pt, icycle), ua) #print("xa={}".format(xa)) d = y - np.mean(obs.h_operator(xa, op, ga), axis=1) chi2 = chi2_test(zmat, heinv, rmat, d) ds = dof(zmat) logger.info("DOF = {}".format(ds)) return xa, xa_, pa, chi2, ds, condh
def gen_obs(u, sigma, op): y = obs.add_noise(obs.h_operator(u, sigma), op) return y