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
xa = xc[:, None] + pa plot_wind(xf, xfc, xa, xc, y, sig, htype, "mlef", var) 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) cg[0, i] = spo.check_grad(calc_j, calc_grad_j, x0, *args) # grad htype["perturbation"] = "grad" xc, pa, chi2, ds, condh = analysis(xf, xfc, y, rmat, rinv, htype) xa = xc[:, None] + pa plot_wind(xf, xfc, xa, xc, y, sig, htype, "grad", var) dh = obs.dhdx(xfc, htype["operator"], htype["gamma"]) @ pf zmat = rmat @ dh tmat, heinv, condh = precondition(zmat) gmat = pf @ tmat args = (xfc, pf, y, tmat, gmat, heinv, rinv, htype) cg[1, i] = spo.check_grad(calc_j, calc_grad_j, x0, *args) # mlef05 htype["perturbation"] = "mlef" xc, pa, chi2 = analysis05(xf, xfc, y, rmat, rinv, htype) xa = xc[:, None] + pa plot_wind(xf, xfc, xa, xc, y, sig, htype, "mlef05", var) 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)
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, binv, y, 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 = obs.dhdx(xf, op, ga) ob = y - 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 analysis(xf, xc, y, rmat, rinv, htype, gtol=1e-6, method="LBFGS", 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 = precondition(zmat) np.save("tmat.npy", tmat) #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, precondition, rmat, htype) iprint = np.zeros(2, dtype=np.int32) iprint[0] = 1 minimize = Minimize(x0.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, 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": 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) else: xmax = int(xmax*0.01+1)*100 logger.debug("resx max={}".format(xmax)) # print("resx max={}".format(xmax)) cost_j(xmax, xf.shape[1], model, x, icycle, *args_j) elif model=="l96": cost_j(200, xf.shape[1], model, x, icycle, *args_j) else: x = minimize(x0) tmat = np.load("tmat.npy") gmat = pf @ tmat xa = xc + 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 = precondition(zmat) heinv = tmat.transpose() @ tmat 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)) # print(ds) #if infl: # if op == "linear": # if pt == "mlef": # alpha = 1.1 # 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 # print("==inflation==, alpha={}".format(alpha)) # pa *= alpha return xa, pa, chi2, ds
dxf = (xf - xf_[:, None])/np.sqrt(nmem-1) bmat = dxf @ dxf.transpose() print(bmat) u, s, vt = np.linalg.svd(dxf) print(s.size) binv = u @ np.diag(1.0/s**2) @ u.transpose() #binv = np.diag(1.0/wstdv0**2) print(binv) xa = np.zeros_like(xf) dy = np.zeros((y.size, nmem)) for i in range(nmem): xa[:, i], chi2 = analysis_var(xf[:, i], binv, y, ytype, rinv, htype, gtol=1e-6, method="CGF", cgtype=1) for j in range(y.size): op = ytype[j] dy[j, i] = obs.dhdx(xf[:, i], op) @ dxf[:, i] xa_ = np.mean(xa, axis=1) print(np.sqrt(xa_[0]**2 + xa_[1]**2)) print(xa_[0]) if infl_r: dxa = xa - xa_[:, None] sigb = np.sum(np.diag([email protected])) sigo = sig*sig*y.size print(sigb, sigo) if sigb/sigo < 1e-3: #|HPfHt| << |R| beta = 0.5 else: beta = 1.0 / (1.0 + np.sqrt(sigo/(sigo + sigb))) print("relaxation factor = {}".format(beta)) dxa = beta * dxa + (1.0 - beta) * dxf * np.sqrt(nmem-1) xa = xa_[:, None] + dxa