def loop(niter,
beta,
betan,
phi,
Acond,
cs,
dbarn,
eplnn,
rnorm,
sn,
Tnorm,
rnorml,
xnorm,
Dnorm,
gamma,
pnorm,
gammal,
Axnorm,
relrnorm,
relArnorml,
Anorm,
flag,
*args):
#-----------------------------------------------------------------
## Obtain quantities for the next Lanczos vector vk+1, k = 1, 2,...
# The general iteration is similar to the case k = 1 with v0 = 0:
#
# p1 = Operator * v1 - beta1 * v0,
# alpha1 = v1'p1,
# q2 = p2 - alpha1 * v1,
# beta2^2 = q2'q2,
# v2 = (1/beta2) q2.
#
# Again, p = betak P vk, where P = C**(-1).
# .... more description needed.
#-----------------------------------------------------------------
xs = args[0 * n_params: 1 * n_params]
r1s = args[1 * n_params: 2 * n_params]
r2s = args[2 * n_params: 3 * n_params]
r3s = args[3 * n_params: 4 * n_params]
dls = args[4 * n_params: 5 * n_params]
ds = args[5 * n_params: 6 * n_params]
betal = beta
beta = betan
vs = [r3 / beta for r3 in r3s]
r3s, upds = compute_Av(*vs)
r3s = [r3 - shift * v for r3, v in zip(r3s, vs)]
r3s = [TT.switch(TT.ge(niter, constantX(1.)),
r3 - (beta / betal) * r1,
r3) for r3, r1 in zip(r3s, r1s)]
alpha = inner_product(r3s, vs)
r3s = [r3 - (alpha / beta) * r2 for r3, r2 in zip(r3s, r2s)]
r1s = [r2 for r2 in r2s]
r2s = [r3 for r3 in r3s]
if Ms is not None:
r3s = [r3 / M for r3, M in zip(r3s, Ms)]
betan = sqrt_inner_product(r2s, r3s)
else:
betan = sqrt_inner_product(r3s)
pnorml = pnorm
pnorm = TT.switch(TT.eq(niter, constantX(0.)),
TT.sqrt(TT.sqr(alpha) + TT.sqr(betan)),
TT.sqrt(TT.sqr(alpha) + TT.sqr(betan) +
TT.sqr(beta)))
#-----------------------------------------------------------------
## Apply previous rotation Qk-1 to get
# [dlta_k epln_{k+1}] = [cs sn][dbar_k 0 ]
# [gbar_k dbar_{k+1} ] [sn -cs][alpha_k beta_{k+1}].
#-----------------------------------------------------------------
dbar = dbarn
epln = eplnn
dlta = cs * dbar + sn * alpha
gbar = sn * dbar - cs * alpha
eplnn = sn * betan
dbarn = -cs * betan
## Compute the current plane rotation Qk
gammal2 = gammal
gammal = gamma
cs, sn, gamma = symGivens2(gbar, betan)
tau = cs * phi
phi = sn * phi
Axnorm = TT.sqrt(TT.sqr(Axnorm) + TT.sqr(tau))
# Update d
dl2s = [dl for dl in dls]
dls = [d for d in ds]
ds = [TT.switch(TT.neq(gamma, constantX(0.)),
(v - epln * dl2 - dlta * dl) / gamma,
v)
for v, dl2, dl in zip(vs, dl2s, dls)]
d_norm = TT.switch(TT.neq(gamma, constantX(0.)),
sqrt_inner_product(ds),
constantX(numpy.inf))
# Update x except if it will become too big
xnorml = xnorm
dl2s = [x for x in xs]
xs = [x + tau * d for x, d in zip(xs, ds)]
xnorm = sqrt_inner_product(xs)
xs = [TT.switch(TT.ge(xnorm, maxxnorm),
dl2, x)
for dl2, x in zip(dl2s, xs)]
flag = TT.switch(TT.ge(xnorm, maxxnorm),
constantX(6.), flag)
# Estimate various norms
rnorml = rnorm # ||r_{k-1}||
Anorml = Anorm
Acondl = Acond
relrnorml = relrnorm
flag_no_6 = TT.neq(flag, constantX(6.))
Dnorm = TT.switch(flag_no_6,
TT.sqrt(TT.sqr(Dnorm) + TT.sqr(d_norm)),
Dnorm)
xnorm = TT.switch(flag_no_6, sqrt_inner_product(xs), xnorm)
rnorm = TT.switch(flag_no_6, phi, rnorm)
relrnorm = TT.switch(flag_no_6,
rnorm / (Anorm * xnorm + bnorm),
relrnorm)
Tnorm = TT.switch(flag_no_6,
TT.switch(TT.eq(niter, constantX(0.)),
TT.sqrt(TT.sqr(alpha) + TT.sqr(betan)),
TT.sqrt(TT.sqr(Tnorm) +
TT.sqr(beta) +
TT.sqr(alpha) +
TT.sqr(betan))),
Tnorm)
Anorm = TT.maximum(Anorm, pnorm)
Acond = Anorm * Dnorm
rootl = TT.sqrt(TT.sqr(gbar) + TT.sqr(dbarn))
Anorml = rnorml * rootl
relArnorml = rootl / Anorm
#---------------------------------------------------------------
# See if any of the stopping criteria are satisfied.
# In rare cases, flag is already -1 from above (Abar = const*I).
#---------------------------------------------------------------
epsx = Anorm * xnorm * eps
epsr = Anorm * xnorm * rtol
#Test for singular Hk (hence singular A)
# or x is already an LS solution (so again A must be singular).
t1 = constantX(1) + relrnorm
t2 = constantX(1) + relArnorml
flag = TT.switch(
TT.bitwise_or(TT.eq(flag, constantX(0)),
TT.eq(flag, constantX(6))),
multiple_switch(TT.le(t1, constantX(1)),
constantX(3),
TT.le(t2, constantX(1)),
constantX(4),
TT.le(relrnorm, rtol),
constantX(1),
TT.le(Anorm, constantX(1e-20)),
constantX(12),
TT.le(relArnorml, rtol),
constantX(10),
TT.ge(epsx, beta1),
constantX(5),
TT.ge(xnorm, maxxnorm),
constantX(6),
TT.ge(niter, TT.cast(maxit,
theano.config.floatX)),
constantX(8),
flag),
flag)
flag = TT.switch(TT.lt(Axnorm, rtol * Anorm * xnorm),
constantX(11.),
flag)
return [niter + constantX(1.),
beta,
betan,
phi,
Acond,
cs,
dbarn,
eplnn,
rnorm,
sn,
Tnorm,
rnorml,
xnorm,
Dnorm,
gamma,
pnorm,
gammal,
Axnorm,
relrnorm,
relArnorml,
Anorm,
flag] + xs + r1s + r2s + r3s + dls + ds, upds, \
theano.scan_module.scan_utils.until(TT.neq(flag, 0))
def loop(
niter,
beta,
betan,
phi,
Acond,
cs,
dbarn,
eplnn,
rnorm,
sn,
Tnorm,
rnorml,
xnorm,
Dnorm,
gamma,
pnorm,
gammal,
Axnorm,
relrnorm,
relArnorml,
Anorm,
flag,
*args
):
# -----------------------------------------------------------------
## Obtain quantities for the next Lanczos vector vk+1, k = 1, 2,...
# The general iteration is similar to the case k = 1 with v0 = 0:
#
# p1 = Operator * v1 - beta1 * v0,
# alpha1 = v1'p1,
# q2 = p2 - alpha1 * v1,
# beta2^2 = q2'q2,
# v2 = (1/beta2) q2.
#
# Again, p = betak P vk, where P = C**(-1).
# .... more description needed.
# -----------------------------------------------------------------
xs = args[0 * n_params : 1 * n_params]
r1s = args[1 * n_params : 2 * n_params]
r2s = args[2 * n_params : 3 * n_params]
r3s = args[3 * n_params : 4 * n_params]
dls = args[4 * n_params : 5 * n_params]
ds = args[5 * n_params : 6 * n_params]
betal = beta
beta = betan
vs = [r3 / beta for r3 in r3s]
r3s, upds = compute_Av(*vs)
r3s = [r3 - shift * v for r3, v in zip(r3s, vs)]
r3s = [TT.switch(TT.ge(niter, constantX(1.0)), r3 - (beta / betal) * r1, r3) for r3, r1 in zip(r3s, r1s)]
alpha = inner_product(r3s, vs)
r3s = [r3 - (alpha / beta) * r2 for r3, r2 in zip(r3s, r2s)]
r1s = [r2 for r2 in r2s]
r2s = [r3 for r3 in r3s]
if Ms is not None:
r3s = [r3 / M for r3, M in zip(r3s, Ms)]
betan = sqrt_inner_product(r2s, r3s)
else:
betan = sqrt_inner_product(r3s)
pnorml = pnorm
pnorm = TT.switch(
TT.eq(niter, constantX(0.0)),
TT.sqrt(TT.sqr(alpha) + TT.sqr(betan)),
TT.sqrt(TT.sqr(alpha) + TT.sqr(betan) + TT.sqr(beta)),
)
# -----------------------------------------------------------------
## Apply previous rotation Qk-1 to get
# [dlta_k epln_{k+1}] = [cs sn][dbar_k 0 ]
# [gbar_k dbar_{k+1} ] [sn -cs][alpha_k beta_{k+1}].
# -----------------------------------------------------------------
dbar = dbarn
epln = eplnn
dlta = cs * dbar + sn * alpha
gbar = sn * dbar - cs * alpha
eplnn = sn * betan
dbarn = -cs * betan
## Compute the current plane rotation Qk
gammal2 = gammal
gammal = gamma
cs, sn, gamma = symGivens2(gbar, betan)
tau = cs * phi
phi = sn * phi
Axnorm = TT.sqrt(TT.sqr(Axnorm) + TT.sqr(tau))
# Update d
dl2s = [dl for dl in dls]
dls = [d for d in ds]
ds = [
TT.switch(TT.neq(gamma, constantX(0.0)), (v - epln * dl2 - dlta * dl) / gamma, v)
for v, dl2, dl in zip(vs, dl2s, dls)
]
d_norm = TT.switch(TT.neq(gamma, constantX(0.0)), sqrt_inner_product(ds), constantX(numpy.inf))
# Update x except if it will become too big
xnorml = xnorm
dl2s = [x for x in xs]
xs = [x + tau * d for x, d in zip(xs, ds)]
xnorm = sqrt_inner_product(xs)
xs = [TT.switch(TT.ge(xnorm, maxxnorm), dl2, x) for dl2, x in zip(dl2s, xs)]
flag = TT.switch(TT.ge(xnorm, maxxnorm), constantX(6.0), flag)
# Estimate various norms
rnorml = rnorm # ||r_{k-1}||
Anorml = Anorm
Acondl = Acond
relrnorml = relrnorm
flag_no_6 = TT.neq(flag, constantX(6.0))
Dnorm = TT.switch(flag_no_6, TT.sqrt(TT.sqr(Dnorm) + TT.sqr(d_norm)), Dnorm)
xnorm = TT.switch(flag_no_6, sqrt_inner_product(xs), xnorm)
rnorm = TT.switch(flag_no_6, phi, rnorm)
relrnorm = TT.switch(flag_no_6, rnorm / (Anorm * xnorm + bnorm), relrnorm)
Tnorm = TT.switch(
flag_no_6,
TT.switch(
TT.eq(niter, constantX(0.0)),
TT.sqrt(TT.sqr(alpha) + TT.sqr(betan)),
TT.sqrt(TT.sqr(Tnorm) + TT.sqr(beta) + TT.sqr(alpha) + TT.sqr(betan)),
),
Tnorm,
)
Anorm = TT.maximum(Anorm, pnorm)
Acond = Anorm * Dnorm
rootl = TT.sqrt(TT.sqr(gbar) + TT.sqr(dbarn))
Anorml = rnorml * rootl
relArnorml = rootl / Anorm
# ---------------------------------------------------------------
# See if any of the stopping criteria are satisfied.
# In rare cases, flag is already -1 from above (Abar = const*I).
# ---------------------------------------------------------------
epsx = Anorm * xnorm * eps
epsr = Anorm * xnorm * rtol
# Test for singular Hk (hence singular A)
# or x is already an LS solution (so again A must be singular).
t1 = constantX(1) + relrnorm
t2 = constantX(1) + relArnorml
flag = TT.switch(
TT.bitwise_or(TT.eq(flag, constantX(0)), TT.eq(flag, constantX(6))),
multiple_switch(
TT.le(t1, constantX(1)),
constantX(3),
TT.le(t2, constantX(1)),
constantX(4),
TT.le(relrnorm, rtol),
constantX(1),
TT.le(Anorm, constantX(1e-20)),
constantX(12),
TT.le(relArnorml, rtol),
constantX(10),
TT.ge(epsx, beta1),
constantX(5),
TT.ge(xnorm, maxxnorm),
constantX(6),
TT.ge(niter, TT.cast(maxit, theano.config.floatX)),
constantX(8),
flag,
),
flag,
)
flag = TT.switch(TT.lt(Axnorm, rtol * Anorm * xnorm), constantX(11.0), flag)
return (
[
niter + constantX(1.0),
beta,
betan,
phi,
Acond,
cs,
dbarn,
eplnn,
rnorm,
sn,
Tnorm,
rnorml,
xnorm,
Dnorm,
gamma,
pnorm,
gammal,
Axnorm,
relrnorm,
relArnorml,
Anorm,
flag,
]
+ xs
+ r1s
+ r2s
+ r3s
+ dls
+ ds,
upds,
theano.scan_module.scan_utils.until(TT.neq(flag, 0)),
)