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
