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