def orthogonalize(w, basis, ips=None, nblock=4):
# verbosity
t = gpt.timer("orthogonalize", verbose_performance)
n = len(basis)
if n == 0:
return
grid = basis[0].grid
i = 0
if verbose_performance:
cgpt.timer_begin()
for i in range(0, n, nblock):
t("rank_inner_product")
lip = gpt.rank_inner_product(basis[i : i + nblock], w)
t("global_sum")
grid.globalsum(lip)
t("create expression")
lip = [complex(x) for x in lip]
if ips is not None:
for j in range(len(lip)):
ips[i + j] = lip[j]
expr = w - lip[0] * basis[i + 0]
for j in range(1, len(lip)):
expr -= lip[j] * basis[i + j]
t("linear combination")
w @= expr
t()
if verbose_performance:
t_cgpt = gpt.timer("cgpt_orthogonalize", True)
t_cgpt += cgpt.timer_end()
gpt.message(f"\nPerformance of orthogonalize:\n{t}\n{t_cgpt}")
def __init__(self, name=None):
self.name = self.__class__.__name__ if name is None else name
self.verbose = gpt.default.is_verbose(self.name)
self.verbose_debug = gpt.default.is_verbose(self.name + "_debug")
self.verbose_performance = gpt.default.is_verbose(self.name +
"_performance")
self.timer = gpt.timer(self.name, self.verbose_performance)
def __init__(self, n, precision):
self.n = n
self.fdimensions = [2**n]
self.grid = g.grid(self.fdimensions, precision)
self.verbose = g.default.is_verbose("qis_map")
self.zero_coordinate = (0, ) # |00000 ... 0> state
t = g.timer("map_init")
t("coordinates")
# TODO: need to split over multiple dimensions, single dimension can hold at most 32 bits
self.coordinates = g.coordinates(self.grid)
self.not_coordinates = [
np.bitwise_xor(self.coordinates, 2**i) for i in range(n)
]
for i in range(n):
self.not_coordinates[i].flags["WRITEABLE"] = False
t("masks")
self.one_mask = []
self.zero_mask = []
for i in range(n):
proj = np.bitwise_and(self.coordinates, 2**i)
mask = g.complex(self.grid)
g.coordinate_mask(mask, proj != 0)
self.one_mask.append(mask)
mask = g.complex(self.grid)
g.coordinate_mask(mask, proj == 0)
self.zero_mask.append(mask)
t()
if self.verbose:
g.message(t)
def __call__(self, link, staple, mask):
verbose = g.default.is_verbose(
"metropolis"
) # need verbosity categories [ performance, progress ]
project_method = self.params["project_method"]
step_size = self.params["step_size"]
number_accept = 0
possible_accept = 0
t = g.timer("metropolis")
t("action")
action = g.component.real(g.eval(-g.trace(link * g.adj(staple)) *
mask))
t("lattice")
V = g.lattice(link)
V_eye = g.identity(link)
t("random")
self.rng.element(V, scale=step_size, normal=True)
t("update")
V = g.where(mask, V, V_eye)
link_prime = g.eval(V * link)
action_prime = g.component.real(
g.eval(-g.trace(link_prime * g.adj(staple)) * mask))
dp = g.component.exp(g.eval(action - action_prime))
rn = g.lattice(dp)
t("random")
self.rng.uniform_real(rn)
t("random")
accept = dp > rn
accept *= mask
number_accept += g.norm2(accept)
possible_accept += g.norm2(mask)
link @= g.where(accept, link_prime, link)
t()
g.project(link, project_method)
# g.message(t)
if verbose:
g.message(
f"Metropolis acceptance rate: {number_accept / possible_accept}"
)
def __call__(self, tau):
eps = tau / self.N
verbose = gpt.default.is_verbose(self.__name__)
time = gpt.timer(f"Symplectic integrator {self.__name__} [eps = {eps:.4e}]")
time(self.__name__)
for s in self.scheme:
s(tau)
if verbose:
time()
gpt.message(time)
def inv(psi, src):
self.history = []
verbose = g.default.is_verbose("mr")
t = g.timer("mr")
t("setup")
r, mmr = g.copy(src), g.copy(src)
mat(mmr, psi)
r @= src - mmr
r2 = g.norm2(r)
ssq = g.norm2(src)
# if ssq == 0.0:
# assert r2 != 0.0 # need either source or psi to not be zero
# ssq = r2
rsq = self.eps ** 2.0 * ssq
for k in range(self.maxiter):
t("mat")
mat(mmr, r)
t("inner")
ip, mmr2 = g.inner_product_norm2(mmr, r)
if mmr2 == 0.0:
continue
t("linearcomb")
alpha = ip.real / mmr2 * self.relax
psi += alpha * r
t("axpy_norm")
r2 = g.axpy_norm2(r, -alpha, mmr, r)
t("other")
self.history.append(r2)
if verbose:
g.message("mr: res^2[ %d ] = %g, target = %g" % (k, r2, rsq))
if r2 <= rsq:
if verbose:
t()
g.message(
"mr: converged in %d iterations, took %g s"
% (k + 1, t.dt["total"])
)
g.message(t)
break
def __call__(self, tau):
eps = tau / self.N
verbose = gpt.default.is_verbose(self.__name__)
time = gpt.timer(self.__name__)
time(self.__name__)
for s in self.scheme:
s(eps)
if verbose:
time()
gpt.message(
f"{self.__name__} [eps = {eps:.4e}] in {time.dt['total']:g} secs {time.dt['total']/self.N:g} secs/cycle"
)
def __init__(self, name, U, params, otype=None):
self.interface = interface()
super().__init__(name, U, params, otype, True)
self.tmp = gpt.lattice(self.F_grid, otype)
self.tmp_eo = gpt.lattice(self.F_grid_eo, otype)
self.U_self_inv = gpt.matrix.inv(self.U[8])
self.t = gpt.timer("coarse_operator")
self.params["U"] = [v_obj for u in U for v_obj in u.v_obj]
self.params["U_self_inv"] = self.U_self_inv.v_obj
self.interface.setup(name, self.U_grid, self.params)
def __init__(self, setup, params):
# save input
self.setup = setup
self.params = params
# aliases
s = self.setup
par = self.params
# parameters
self.smooth_solver = g.util.to_list(par["smooth_solver"], s.nlevel - 1)
self.wrapper_solver = g.util.to_list(par["wrapper_solver"],
s.nlevel - 2)
self.coarsest_solver = par["coarsest_solver"]
# verbosity
self.verbose = g.default.is_verbose("multi_grid_inverter")
# print prefix
self.print_prefix = ["mg: level %d:" % i for i in range(s.nlevel)]
# assertions
assert g.util.entries_have_length([self.smooth_solver], s.nlevel - 1)
assert g.util.entries_have_length([self.wrapper_solver], s.nlevel - 2)
assert g.util.is_callable(
[self.smooth_solver, self.coarsest_solver, self.wrapper_solver])
assert type(self.coarsest_solver) != list
assert not g.util.all_have_attribute(self.wrapper_solver, "inverter")
# timing
self.t = [
g.timer("mg_solve_lvl_%d" % (lvl)) for lvl in range(s.nlevel)
]
# temporary vectors
self.r, self.e = [None] * s.nlevel, [None] * s.nlevel
for lvl in range(s.finest + 1, s.nlevel):
nf_lvl = s.nf_lvl[lvl]
self.r[lvl] = g.vcomplex(s.grid[lvl], s.nbasis[nf_lvl])
self.e[lvl] = g.vcomplex(s.grid[lvl], s.nbasis[nf_lvl])
self.r[s.finest] = g.vspincolor(s.grid[s.finest])
# setup a history for all solvers
self.history = [None] * s.nlevel
for lvl in range(s.finest, s.coarsest):
self.history[lvl] = {"smooth": [], "wrapper": []}
self.history[s.coarsest] = {"coarsest": []}
def inv(psi, src):
assert src != psi
self.history = []
verbose = g.default.is_verbose("cg")
t = g.timer("cg")
t("setup")
p, mmp, r = g.copy(src), g.copy(src), g.copy(src)
mat(mmp, psi) # in, out
d = g.inner_product(psi, mmp).real
b = g.norm2(mmp)
r @= src - mmp
p @= r
a = g.norm2(p)
cp = a
ssq = g.norm2(src)
if ssq == 0.0:
assert a != 0.0 # need either source or psi to not be zero
ssq = a
rsq = self.eps ** 2.0 * ssq
for k in range(1, self.maxiter + 1):
c = cp
t("mat")
mat(mmp, p)
t("inner")
dc = g.inner_product(p, mmp)
d = dc.real
a = c / d
t("axpy_norm")
cp = g.axpy_norm2(r, -a, mmp, r)
t("linearcomb")
b = cp / c
psi += a * p
p @= b * p + r
t("other")
self.history.append(cp)
if verbose:
g.message("cg: res^2[ %d ] = %g, target = %g" % (k, cp, rsq))
if cp <= rsq:
if verbose:
t()
g.message(
"cg: converged in %d iterations, took %g s"
% (k, t.dt["total"])
)
g.message(t)
break
def inv(psi, src):
# verbosity
verbose = g.default.is_verbose("dci")
t = g.timer("dci")
t("setup")
# leading order
n = len(src)
_s = [g.copy(x) for x in src]
_d = [g.copy(x) for x in psi]
self.history = []
for i in range(self.maxiter):
# correction step
t("outer_mat")
for j in range(n):
_s[j] -= outer_mat * _d[j]
t("inner_inv")
_d = g.eval(inner_inv_mat * _s)
t("accum")
for j in range(n):
psi[j] += _d[j]
# true resid
eps = max([
g.norm2(outer_mat * psi[j] - src[j])**0.5 for j in range(n)
])
self.history.append(eps)
if verbose:
g.message("Defect-correcting inverter: res^2[ %d ] = %g" %
(i, eps))
if eps < self.eps:
if verbose:
t()
g.message(
"Defect-correcting inverter: converged in %d iterations, took %g s"
% (i + 1, t.dt["total"]))
g.message(t)
break
def exp(i):
t = gpt.timer("exp")
t("eval")
i = gpt.eval(i) # accept expressions
t("prep")
if i.grid.precision != gpt.double:
x = gpt.convert(i, gpt.double)
else:
x = gpt.copy(i)
n = gpt.norm2(x)**0.5 / x.grid.gsites
order = 19
maxn = 0.05
ns = 0
if n > maxn:
ns = int(numpy.log2(n / maxn))
x /= 2**ns
o = gpt.lattice(x)
t("mem")
o[:] = 0
nfac = 1.0
xn = gpt.copy(x)
t("id")
o @= gpt.identity(o)
t("add")
o += xn
t("loop")
for j in range(2, order + 1):
nfac /= j
xn @= xn * x
o += xn * nfac
t("reduce")
for j in range(ns):
o @= o * o
t("conv")
if i.grid.precision != gpt.double:
r = gpt.lattice(i)
gpt.convert(r, o)
o = r
t()
# gpt.message(t)
return o
def element(self, out, p={}):
if type(out) == list:
return [self.element(x, p) for x in out]
t = gpt.timer("element")
scale = p["scale"]
normal = p["normal"]
grid = out.grid
t("complex")
ca = gpt.complex(grid)
ca.checkerboard(out.checkerboard())
t("cartesian_space")
cartesian_space = gpt.group.cartesian(out)
t("csset")
cartesian_space[:] = 0
t("gen")
gen = cartesian_space.otype.generators(grid.precision.complex_dtype)
t()
for ta in gen:
t("rng")
if normal:
self.normal(ca)
else:
self.uniform_real(ca, {"min": -0.5, "max": 0.5})
t("lc")
cartesian_space += scale * ca * ta
t("conv")
gpt.convert(out, cartesian_space)
t()
# gpt.message(t)
return out
def __init__(self):
self.grad = {}
self.time = gpt.timer()
def reset(self):
self.time = gpt.timer()
for key in self.grad:
self.grad[key] = []
def inv(psi, src):
self.history = []
verbose = g.default.is_verbose("bicgstab")
t = g.timer("bicgstab")
t("setup")
r, rhat, p, s = g.copy(src), g.copy(src), g.copy(src), g.copy(src)
mmpsi, mmp, mms = g.copy(src), g.copy(src), g.copy(src)
rho, rhoprev, alpha, omega = 1.0, 1.0, 1.0, 1.0
mat(mmpsi, psi)
r @= src - mmpsi
rhat @= r
p @= r
mmp @= r
r2 = g.norm2(r)
ssq = g.norm2(src)
if ssq == 0.0:
assert r2 != 0.0 # need either source or psi to not be zero
ssq = r2
rsq = self.eps**2.0 * ssq
for k in range(self.maxiter):
t("inner")
rhoprev = rho
rho = g.inner_product(rhat, r).real
t("linearcomb")
beta = (rho / rhoprev) * (alpha / omega)
p @= r + beta * p - beta * omega * mmp
t("mat")
mat(mmp, p)
t("inner")
alpha = rho / g.inner_product(rhat, mmp).real
t("linearcomb")
s @= r - alpha * mmp
t("mat")
mat(mms, s)
t("inner")
ip, mms2 = g.inner_product_norm2(mms, s)
if mms2 == 0.0:
continue
t("linearcomb")
omega = ip.real / mms2
psi += alpha * p + omega * s
t("axpy_norm")
r2 = g.axpy_norm2(r, -omega, mms, s)
t("other")
self.history.append(r2)
if verbose:
g.message("bicgstab: res^2[ %d ] = %g, target = %g" %
(k, r2, rsq))
if r2 <= rsq:
if verbose:
t()
g.message(
"bicgstab: converged in %d iterations, took %g s" %
(k + 1, t.dt["total"]))
g.message(t)
break
def inv(psi, src):
# verbosity
verbose = g.default.is_verbose("dci")
t = g.timer("dci")
t("setup")
# inner source
n = len(src)
_s = [g.lattice(x) for x in src]
# norm of source
norm2_of_source = g.norm2(src)
for j in range(n):
if norm2_of_source[j] == 0.0:
norm2_of_source[j] = g.norm2(outer_mat * psi[j])
if norm2_of_source[j] == 0.0:
norm2_of_source[j] = 1.0
self.history = []
for i in range(self.maxiter):
t("outer matrix")
for j in range(n):
_s[j] @= src[j] - outer_mat * psi[j] # remaining src
t("norm2")
norm2_of_defect = g.norm2(_s)
# true resid
t("norm2")
eps = max(
[(norm2_of_defect[j] / norm2_of_source[j]) ** 0.5 for j in range(n)]
)
self.history.append(eps)
if verbose:
g.message("Defect-correcting inverter: res^2[ %d ] = %g" % (i, eps))
if eps < self.eps:
if verbose:
t()
g.message(
"Defect-correcting inverter: converged in %d iterations, took %g s"
% (i + 1, t.total)
)
g.message(t)
break
# normalize _s to avoid floating-point underflow in inner_inv_mat
t("linear algebra")
for j in range(n):
_s[j] /= norm2_of_source[j] ** 0.5
# correction step
t("inner inverter")
_d = g.eval(inner_inv_mat * _s)
t("linear algebra")
for j in range(n):
psi[j] += _d[j] * norm2_of_source[j] ** 0.5
def timed_start(self):
return gpt.timer(self.name) if self.verbose_performance else self.timer
def inv(psi, src):
self.history = []
# verbosity
verbose = g.default.is_verbose("fgmres")
# timing
t = g.timer("fgmres")
t("setup")
# parameters
rlen = self.restartlen
# tensors
dtype = g.double.complex_dtype
H = np.zeros((rlen + 1, rlen), dtype)
c = np.zeros((rlen + 1), dtype)
s = np.zeros((rlen + 1), dtype)
y = np.zeros((rlen + 1), dtype)
gamma = np.zeros((rlen + 1), dtype)
# fields
mmpsi, r = (
g.copy(src),
g.copy(src),
)
V = [g.lattice(src) for i in range(rlen + 1)]
Z = (
[g.lattice(src) for i in range(rlen + 1)] if prec is not None else None
) # save vectors if unpreconditioned
# initial residual
r2 = self.restart(mat, psi, mmpsi, src, r, V, Z, gamma)
# source
ssq = g.norm2(src)
if ssq == 0.0:
assert r2 != 0.0 # need either source or psi to not be zero
ssq = r2
# target residual
rsq = self.eps ** 2.0 * ssq
for k in range(self.maxiter):
# iteration within current krylov space
i = k % rlen
# iteration criteria
reached_maxiter = k + 1 == self.maxiter
need_restart = i + 1 == rlen
t("prec")
if prec is not None:
prec(Z[i], V[i])
t("mat")
if prec is not None:
mat(V[i + 1], Z[i])
else:
mat(V[i + 1], V[i])
t("ortho")
g.default.push_verbose("orthogonalize", False)
g.orthogonalize(V[i + 1], V[0 : i + 1], H[:, i])
g.default.pop_verbose()
t("linalg")
H[i + 1, i] = g.norm2(V[i + 1]) ** 0.5
if H[i + 1, i] == 0.0:
g.message("fgmres: breakdown, H[%d, %d] = 0" % (i + 1, i))
break
V[i + 1] /= H[i + 1, i]
t("qr")
self.qr_update(s, c, H, gamma, i)
t("other")
r2 = np.absolute(gamma[i + 1]) ** 2
self.history.append(r2)
if verbose:
g.message(
"fgmres: res^2[ %d, %d ] = %g, target = %g" % (k, i, r2, rsq)
)
if r2 <= rsq or need_restart or reached_maxiter:
t("update_psi")
if prec is not None:
self.update_psi(psi, gamma, H, y, Z, i)
else:
self.update_psi(psi, gamma, H, y, V, i)
comp_res = r2 / ssq
if r2 <= rsq:
if verbose:
t()
g.message(
"fgmres: converged in %d iterations, took %g s"
% (k + 1, t.dt["total"])
)
g.message(t)
if self.checkres:
res = self.calc_res(mat, psi, mmpsi, src, r) / ssq
g.message(
"fgmres: computed res = %g, true res = %g, target = %g"
% (comp_res ** 0.5, res ** 0.5, self.eps)
)
else:
g.message(
"fgmres: computed res = %g, target = %g"
% (comp_res ** 0.5, self.eps)
)
break
if reached_maxiter:
if verbose:
t()
g.message(
"fgmres: did NOT converge in %d iterations, took %g s"
% (k + 1, t.dt["total"])
)
g.message(t)
if self.checkres:
res = self.calc_res(mat, psi, mmpsi, src, r) / ssq
g.message(
"fgmres: computed res = %g, true res = %g, target = %g"
% (comp_res ** 0.5, res ** 0.5, self.eps)
)
else:
g.message(
"fgmres: computed res = %g, target = %g"
% (comp_res ** 0.5, self.eps)
)
break
if need_restart:
t("restart")
r2 = self.restart(mat, psi, mmpsi, src, r, V, Z, gamma)
if verbose:
g.message("fgmres: performed restart")
def hmc(tau, mom):
rng.normal_element(mom)
accrej = metro(U)
h0 = hamiltonian()
mdint(tau)
h1 = hamiltonian()
return [accrej(h1, h0), h1 - h0]
# thermalization
ntherm = 100
for i in range(1, 11):
h = []
timer = g.timer("hmc")
for _ in range(ntherm // 10):
timer("trajectory")
h += [hmc(tau, mom)]
h = numpy.array(h)
timer()
g.message(f"{i*10} % of thermalization completed")
g.message(
f'Average time per trajectory = {timer.dt["trajectory"]/ntherm*10:g} secs'
)
g.message(
f"Plaquette = {g.qcd.gauge.plaquette(U)}, Acceptance = {numpy.mean(h[:,0]):.2f}, |dH| = {numpy.mean(numpy.abs(h[:,1])):.4e}"
)
# production
history = []
]:
# Time
dt = 0.0
cgpt.timer_begin()
for it in range(N + Nwarmup):
access(one)
access(two)
if it >= Nwarmup:
dt -= g.time()
ip = g.rank_inner_product(one, two, use_accelerator)
if it >= Nwarmup:
dt += g.time()
# Report
GBPerSec = nbytes / dt / 1e9
cgpt_t = g.timer("rip")
cgpt_t += cgpt.timer_end()
g.message(
f"""{N} rank_inner_product
Object type : {tp.__name__}
Block : {n} x {n}
Data resides in : {access.__name__[7:]}
Performed on : {compute_name}
Time to complete : {dt:.2f} s
Effective memory bandwidth : {GBPerSec:.2f} GB/s
{cgpt_t}
"""
)
def inv(psi, src):
self.history = []
# verbosity
verbose = g.default.is_verbose("fgcr")
# timing
t = g.timer("fgcr")
t("setup")
# parameters
rlen = self.restartlen
# tensors
dtype_r, dtype_c = g.double.real_dtype, g.double.complex_dtype
alpha = np.empty((rlen), dtype_c)
beta = np.empty((rlen, rlen), dtype_c)
gamma = np.empty((rlen), dtype_r)
chi = np.empty((rlen), dtype_c)
# fields
r, mmpsi = g.copy(src), g.copy(src)
p = [g.lattice(src) for i in range(rlen)]
z = [g.lattice(src) for i in range(rlen)]
# initial residual
r2 = self.restart(mat, psi, mmpsi, src, r, p)
# source
ssq = g.norm2(src)
if ssq == 0.0:
assert r2 != 0.0 # need either source or psi to not be zero
ssq = r2
# target residual
rsq = self.eps**2.0 * ssq
for k in range(self.maxiter):
# iteration within current krylov space
i = k % rlen
# iteration criteria
reached_maxiter = k + 1 == self.maxiter
need_restart = i + 1 == rlen
t("prec")
if prec is not None:
prec(p[i], r)
else:
p[i] @= r
t("mat")
mat(z[i], p[i])
t("ortho")
g.default.push_verbose("orthogonalize", False)
g.orthogonalize(z[i], z[0:i], beta[:, i])
g.default.pop_verbose()
t("linalg")
ip, z2 = g.inner_product_norm2(z[i], r)
gamma[i] = z2**0.5
if gamma[i] == 0.0:
g.message("fgcr: breakdown, gamma[%d] = 0" % (i))
break
z[i] /= gamma[i]
alpha[i] = ip / gamma[i]
r2 = g.axpy_norm2(r, -alpha[i], z[i], r)
t("other")
self.history.append(r2)
if verbose:
g.message("fgcr: res^2[ %d, %d ] = %g, target = %g" %
(k, i, r2, rsq))
if r2 <= rsq or need_restart or reached_maxiter:
t("update_psi")
self.update_psi(psi, alpha, beta, gamma, chi, p, i)
comp_res = r2 / ssq
if r2 <= rsq:
if verbose:
t()
g.message(
"fgcr: converged in %d iterations, took %g s" %
(k + 1, t.total))
g.message(t)
if self.checkres:
res = self.calc_res(mat, psi, mmpsi, src,
r) / ssq
g.message(
"fgcr: computed res = %g, true res = %g, target = %g"
% (comp_res**0.5, res**0.5, self.eps))
else:
g.message(
"fgcr: computed res = %g, target = %g" %
(comp_res**0.5, self.eps))
break
if reached_maxiter:
if verbose:
t()
g.message(
"fgcr: did NOT converge in %d iterations, took %g s"
% (k + 1, t.dt["total"]))
g.message(t)
if self.checkres:
res = self.calc_res(mat, psi, mmpsi, src,
r) / ssq
g.message(
"fgcr: computed res = %g, true res = %g, target = %g"
% (comp_res**0.5, res**0.5, self.eps))
else:
g.message(
"fgcr: computed res = %g, target = %g" %
(comp_res**0.5, self.eps))
break
if need_restart:
t("restart")
r2 = self.restart(mat, psi, mmpsi, src, r, p)
if verbose:
g.message("fgcr: performed restart")
def expr_eval(first, second=None, ac=False):
t = gpt.timer("eval", verbose_performance)
# this will always evaluate to a (list of) lattice object(s)
# or remain an expression if it cannot do so
t("prepare")
if second is not None:
dst = gpt.util.to_list(first)
e = expr(second)
return_list = False
else:
assert ac is False
if gpt.util.is_list_instance(first, gpt.lattice):
return first
e = expr(first)
lat = get_lattice(e)
if lat is None:
# cannot evaluate to a lattice object, leave expression unevaluated
return first
return_list = type(lat) == list
lat = gpt.util.to_list(lat)
grid = lat[0].grid
nlat = len(lat)
dst = None
t("apply matrix ops")
# apply matrix_operators
e = apply_type_right_to_left(e, gpt.matrix_operator)
t("fast return")
# fast return if already a lattice
if dst is None:
if e.is_single(gpt.lattice):
ue, uf, v = e.get_single()
if uf == factor_unary.NONE and ue == expr_unary.NONE:
return v
# verbose output
if verbose:
gpt.message("eval: " + str(e))
if verbose_performance:
cgpt.timer_begin()
if dst is not None:
t("cgpt.eval")
for i, dst_i in enumerate(dst):
dst_i.update(cgpt.eval(dst_i.v_obj, e.val, e.unary, ac, i))
ret = dst
else:
assert ac is False
t("get otype")
# now find return type
otype = get_otype_from_expression(e)
ret = []
for idx in range(nlat):
t("cgpt.eval")
res = cgpt.eval(None, e.val, e.unary, False, idx)
t_obj, s_ot = (
[x[0] for x in res],
[x[1] for x in res],
)
assert s_ot == otype.v_otype
t("lattice")
ret.append(gpt.lattice(grid, otype, t_obj))
t()
if verbose_performance:
t_cgpt = gpt.timer("cgpt_eval", True)
t_cgpt += cgpt.timer_end()
gpt.message(t)
gpt.message(t_cgpt)
if not return_list:
return gpt.util.from_list(ret)
return ret
def __init__(self, mat_f, params):
# save parameters
self.params = params
# fine grid from fine matrix
if issubclass(type(mat_f), g.matrix_operator):
self.grid = [mat_f.grid[1]]
else:
self.grid = [mat_f.grid]
# grid sizes - allow specifying in two ways
if "grid" in params:
self.grid.extend(params["grid"])
elif "blocksize" in params:
for i, bs in enumerate(params["blocksize"]):
assert type(bs) == list
self.grid.append(g.block.grid(self.grid[i], bs))
else:
assert 0
# dependent sizes
self.nlevel = len(self.grid)
self.ncoarselevel = self.nlevel - 1
self.finest = 0
self.coarsest = self.nlevel - 1
# other parameters
self.nblockortho = g.util.to_list(params["nblockortho"],
self.nlevel - 1)
self.check_blockortho = g.util.to_list(params["check_blockortho"],
self.nlevel - 1)
self.nbasis = g.util.to_list(params["nbasis"], self.nlevel - 1)
self.make_hermitian = g.util.to_list(params["make_hermitian"],
self.nlevel - 1)
self.save_links = g.util.to_list(params["save_links"], self.nlevel - 1)
self.npreortho = g.util.to_list(params["npreortho"], self.nlevel - 1)
self.npostortho = g.util.to_list(params["npostortho"], self.nlevel - 1)
self.vector_type = g.util.to_list(params["vector_type"],
self.nlevel - 1)
self.distribution = g.util.to_list(params["distribution"],
self.nlevel - 1)
self.solver = g.util.to_list(params["solver"], self.nlevel - 1)
# verbosity
self.verbose = g.default.is_verbose("multi_grid_setup")
# print prefix
self.print_prefix = [
"mg_setup: level %d:" % i for i in range(self.nlevel)
]
# easy access to current level and neighbors
self.lvl = [i for i in range(self.nlevel)]
self.nf_lvl = [i - 1 for i in range(self.nlevel)]
self.nc_lvl = [i + 1 for i in range(self.nlevel)]
self.nf_lvl[self.finest] = None
self.nc_lvl[self.coarsest] = None
# halved nbasis
self.nb = []
for lvl, b in enumerate(self.nbasis):
assert b % 2 == 0
self.nb.append(b // 2)
# assertions
assert self.nlevel >= 2
assert g.util.entries_have_length(
[
self.nblockortho,
self.nbasis,
self.make_hermitian,
self.save_links,
self.npreortho,
self.npostortho,
self.vector_type,
self.distribution,
self.solver,
self.nb,
],
self.nlevel - 1,
)
# timing
self.t = [
g.timer("mg_setup_lvl_%d" % (lvl)) for lvl in range(self.nlevel)
]
# matrices (coarse ones initialized later)
self.mat = [mat_f] + [None] * (self.nlevel - 1)
# setup random basis vectors on all levels but coarsest
self.basis = [None] * self.nlevel
for lvl, grid in enumerate(self.grid):
if lvl == self.coarsest:
continue
elif lvl == self.finest:
self.basis[lvl] = [
g.vspincolor(grid) for __ in range(self.nbasis[lvl])
]
else:
self.basis[lvl] = [
g.vcomplex(grid, self.nbasis[self.nf_lvl[lvl]])
for __ in range(self.nbasis[lvl])
]
self.distribution[lvl](self.basis[lvl][0:self.nb[lvl]])
# setup a block map on all levels but coarsest
self.blockmap = [None] * self.nlevel
for lvl in self.lvl:
if lvl == self.coarsest:
continue
else:
self.blockmap[lvl] = g.block.map(self.grid[self.nc_lvl[lvl]],
self.basis[lvl])
# setup coarse link fields on all levels but finest
self.A = [None] * self.nlevel
for lvl in range(self.finest + 1, self.nlevel):
self.A[lvl] = [
g.mcomplex(self.grid[lvl], self.nbasis[self.nf_lvl[lvl]])
for __ in range(9)
]
# setup a solver history
self.history = [[None]] * (self.nlevel - 1)
# rest of setup
self.__call__()
def create_links(A, fmat, basis, params):
# NOTE: we expect the blocks in the basis vectors
# to already be orthogonalized!
# parameters
make_hermitian = params["make_hermitian"]
save_links = params["save_links"]
assert not (make_hermitian and not save_links)
# verbosity
verbose = gpt.default.is_verbose("coarsen")
# setup timings
t = gpt.timer("coarsen")
t("setup")
# get grids
f_grid = basis[0].grid
c_grid = A[0].grid
# directions/displacements we coarsen for
dirs = [1, 2, 3, 4] if f_grid.nd == 5 else [0, 1, 2, 3]
disp = +1
dirdisps_full = list(zip(dirs * 2, [+1] * 4 + [-1] * 4))
dirdisps_forward = list(zip(dirs, [disp] * 4))
nhops = len(dirdisps_full)
selflink = nhops
# setup fields
Mvr = [gpt.lattice(basis[0]) for i in range(nhops)]
tmp = gpt.lattice(basis[0])
oproj = gpt.vcomplex(c_grid, len(basis))
selfproj = gpt.vcomplex(c_grid, len(basis))
# setup masks
onemask, blockevenmask, blockoddmask = (
gpt.complex(f_grid),
gpt.complex(f_grid),
gpt.complex(f_grid),
)
dirmasks = [gpt.complex(f_grid) for p in range(nhops)]
# auxilliary stuff needed for masks
t("masks")
onemask[:] = 1.0
coor = gpt.coordinates(blockevenmask)
block = numpy.array(f_grid.ldimensions) / numpy.array(c_grid.ldimensions)
block_cb = coor[:, :] // block[:]
# fill masks for sites within even/odd blocks
gpt.coordinate_mask(blockevenmask, numpy.sum(block_cb, axis=1) % 2 == 0)
blockoddmask @= onemask - blockevenmask
# fill masks for sites on borders of blocks
dirmasks_forward_np = coor[:, :] % block[:] == block[:] - 1
dirmasks_backward_np = coor[:, :] % block[:] == 0
for mu in dirs:
gpt.coordinate_mask(dirmasks[mu], dirmasks_forward_np[:, mu])
gpt.coordinate_mask(dirmasks[mu + 4], dirmasks_backward_np[:, mu])
# save applications of matrix and coarsening if possible
dirdisps = dirdisps_forward if save_links else dirdisps_full
# create block maps
t("blockmap")
dirbms = [
gpt.block.map(c_grid, basis, dirmasks[p])
for p, (mu, fb) in enumerate(dirdisps)
]
fullbm = gpt.block.map(c_grid, basis)
for i, vr in enumerate(basis):
# apply directional hopping terms
# this triggers len(dirdisps) comms -> TODO expose DhopdirAll from Grid
# BUT problem with vector<Lattice<...>> in rhs
t("apply_hop")
[fmat.Mdir(*dirdisp)(Mvr[p], vr) for p, dirdisp in enumerate(dirdisps)]
# coarsen directional terms + write to link
for p, (mu, fb) in enumerate(dirdisps):
t("coarsen_hop")
dirbms[p].project(oproj, Mvr[p])
t("copy_hop")
A[p][:, :, :, :, :, i] = oproj[:]
# fast diagonal term: apply full matrix to both block cbs separately and discard hops into other cb
t("apply_self")
tmp @= (blockevenmask * fmat * vr * blockevenmask +
blockoddmask * fmat * vr * blockoddmask)
# coarsen diagonal term
t("coarsen_self")
fullbm.project(selfproj, tmp)
# write to self link
t("copy_self")
A[selflink][:, :, :, :, :, i] = selfproj[:]
if verbose:
gpt.message("coarsen: done with vector %d" % i)
# communicate opposite links
if save_links:
t("comm")
communicate_links(A, dirdisps_forward, make_hermitian)
t()
if verbose:
gpt.message(t)
def __call__(self, mat):
# now create a local matrix
U = mat.arguments()
U_dim = len(U)
U_grid = U[0].grid
F_grid = mat.vector_space[1].grid
F_dim = F_grid.nd
U_ldomain = g.domain.local(U_grid, [self.margin] * U_dim)
F_ldomain = g.domain.local(
F_grid,
[0] * (F_dim - U_dim) + [self.margin] * U_dim,
cb=mat.vector_space[1].lattice().checkerboard(),
)
U_local = []
colon = slice(None, None, None)
for i in range(U_dim):
u_local = U_ldomain.lattice(U[i].otype)
U_ldomain.project(u_local, U[i])
u_local[tuple([colon] * i + [u_local.grid.ldimensions[i] - 1] +
[colon] * (U_dim - i - 1))] = 0
U_local.append(u_local)
# get local matrix
mat_local = mat.updated(U_local)
F_ldomain_bulk = F_ldomain.bulk()
F_ldomain_margin = F_ldomain.margin()
# now need copy plans
z = mat_local.vector_space[0].lattice()
z[:] = 0
mat_local_mat = mat_local.mat
timer = [g.timer()]
def proj_mat_local(dst, src):
t = timer[0]
t("local matrix")
mat_local_mat(dst, src)
t("zero projection")
F_ldomain_margin.project(dst, z)
t()
# pml = g.matrix_operator(mat=proj_mat_local, vector_space=mat_local.vector_space).inherit(
pml_inv = self.inner_inverter(proj_mat_local)
# vector space
vector_space = None
if type(mat) == g.matrix_operator:
vector_space = mat.vector_space
_src = mat_local.vector_space[1].lattice()
_dst = mat_local.vector_space[0].lattice()
_src[:] = 0
_dst[:] = 0
def proj_mat(dst, src):
F_ldomain_bulk.project(_src, src)
proj_mat_local(_dst, _src)
F_ldomain_bulk.promote(dst, _dst)
@self.timed_function
def inv(dst, src, t):
timer[0] = t
t("project")
F_ldomain_bulk.project(_src, src)
F_ldomain_bulk.project(_dst, dst)
F_ldomain_margin.project(_dst, z)
t("inner inverter")
pml_inv(_dst, _src)
t("promote")
F_ldomain_bulk.promote(dst, _dst)
t()
return g.matrix_operator(mat=inv,
inv_mat=proj_mat,
accept_guess=(True, False),
vector_space=vector_space)