def __init__(self, matrix):
self.F_grid = matrix.F_grid
ftmp = gpt.vspincolor(self.F_grid)
def _Mpc(dst, src):
matrix.G5M.mat(ftmp, src)
matrix.G5M.mat(dst, ftmp)
def _L(dst, src):
gpt.copy(dst, src)
def _R(dst, src):
dst @= matrix.G5M * gpt.gamma[5] * src
def _S(dst, src):
dst[:] = 0
self.Mpc = gpt.matrix_operator(mat=_Mpc,
otype=matrix.otype,
grid=self.F_grid)
self.L = gpt.matrix_operator(mat=_L,
otype=matrix.otype,
grid=self.F_grid)
self.R = gpt.matrix_operator(mat=_R,
otype=matrix.otype,
grid=self.F_grid)
self.S = gpt.matrix_operator(mat=_S,
otype=matrix.otype,
grid=self.F_grid)
def __init__(self, matrix, pc):
self.matrix = matrix
self.F_grid_eo = matrix.F_grid_eo
self.F_grid = matrix.F_grid
self.U_grid = matrix.U_grid
self.otype = matrix.otype[0]
self.F_tmp = gpt.lattice(self.F_grid, self.otype)
self.F_tmp_2 = gpt.lattice(self.F_grid, self.otype)
def _L(dst, src):
pc.L.mat(self.F_tmp, src)
self.matrix.ExportPhysicalFermionSolution(dst, self.F_tmp)
def _R_adj(dst, src):
pc.R.adj_mat(self.F_tmp, src)
self.matrix.Dminus.adj_mat(self.F_tmp_2, self.F_tmp)
self.matrix.ExportPhysicalFermionSource(dst, self.F_tmp_2)
self.L = gpt.matrix_operator(
mat=_L,
otype=self.otype,
accept_guess=(False, False),
grid=(self.U_grid, self.F_grid_eo),
)
self.R = gpt.matrix_operator(
mat=None,
adj_mat=_R_adj,
otype=self.otype,
accept_guess=(False, False),
grid=(self.F_grid_eo, self.U_grid),
)
def __init__(self, U, params):
covariant_shift.__init__(self, U, params)
Nc = U[0].otype.Nc
otype = g.ot_vector_spin_color(4, Nc)
grid = U[0].grid
if "mass" in params:
assert "kappa" not in params
self.kappa = 1.0 / (params["mass"] + 4.0) / 2.0
else:
self.kappa = params["kappa"]
self.Meooe = g.matrix_operator(lambda dst, src: self._Meooe(dst, src),
otype=otype,
grid=grid)
self.Mooee = g.matrix_operator(lambda dst, src: self._Mooee(dst, src),
otype=otype,
grid=grid)
matrix_operator.__init__(self,
lambda dst, src: self._M(dst, src),
otype=otype,
grid=grid)
self.G5M = g.matrix_operator(lambda dst, src: self._G5M(dst, src),
otype=otype,
grid=grid)
def __init__(self, matrix):
self.F_grid = matrix.F_grid
ftmp = gpt.vspincolor(self.F_grid)
def _Mpc(dst, src):
matrix.G5M.mat(ftmp, src)
matrix.G5M.mat(dst, ftmp)
def _ident(dst, src):
gpt.copy(dst, src)
def _R(dst, src):
dst @= matrix.G5M * gpt.gamma[5] * src
def _S(dst, src):
dst[:] = 0
self.Mpc = gpt.matrix_operator(mat=_Mpc,
vector_space=matrix.vector_space)
self.L = gpt.matrix_operator(mat=_ident,
inv_mat=_ident,
vector_space=matrix.vector_space)
self.R = gpt.matrix_operator(mat=_R, vector_space=matrix.vector_space)
self.S = gpt.matrix_operator(mat=_S, vector_space=matrix.vector_space)
def __call__(self, mat):
inverter_mat = [
self.inverter(
g.matrix_operator(
mat=lambda dst, src, s_val=s: g.eval(dst, mat * src + s_val * src),
accept_list=True,
)
)
for s in self.shifts
]
@self.timed_function
def inv(dst, src, t):
for j, i in enumerate(inverter_mat):
i(dst[j * len(src) : (j + 1) * len(src)], src)
vector_space = None
if isinstance(mat, g.matrix_operator):
vector_space = mat.vector_space
return g.matrix_operator(
mat=inv,
vector_space=vector_space,
accept_guess=(True, False),
accept_list=lambda src: len(src) * len(self.shifts),
)
def __call__(self, dwf_outer):
dwf_inner = self.dwf_inner
dwf_inner_pv = self.dwf_inner_pv
dwf_outer_pv = dwf_outer.modified(mass=1.0)
inv_dwf_outer_pv = self.solver_pv(dwf_outer_pv)
inv_dwf_inner = self.solver(dwf_inner)
def sep(x):
return g.separate(x, dimension=0)
def mrg(x):
return g.merge(x, dimension=0)
def _P(dst, src, offset):
src_s = sep(src)
Ls = len(src_s)
Pplus = 0.5 * (g.gamma["I"] + g.gamma[5])
Pminus = 0.5 * (g.gamma["I"] - g.gamma[5])
dst @= mrg([
g(Pminus * src_s[s] + Pplus * src_s[(s + Ls + offset) % Ls])
for s in range(Ls)
])
def _P_mat(dst, src):
_P(dst, src, 1)
def _P_inv_mat(dst, src):
_P(dst, src, -1)
P = g.matrix_operator(mat=_P_mat,
adj_mat=_P_inv_mat,
adj_inv_mat=_P_mat,
inv_mat=_P_inv_mat)
def inv(dst_outer, src_outer):
Ls_inner = dwf_inner.F_grid.fdimensions[0]
zero4d = g.lattice(dwf_outer.U_grid, src_outer.otype)
zero4d[:] = 0
c_s = sep(g.adj(P) * inv_dwf_outer_pv * src_outer)
y0prime = sep(
g.adj(P) * inv_dwf_inner * dwf_inner_pv * P *
mrg([c_s[0]] + [zero4d] * (Ls_inner - 1)))[0]
dst_outer @= P * mrg([y0prime] + sep(
g.adj(P) * inv_dwf_outer_pv * dwf_outer * P *
mrg([g(-y0prime)] + c_s[1:]))[1:])
return g.matrix_operator(
mat=inv,
inv_mat=dwf_outer,
otype=dwf_outer.otype,
accept_guess=(True, False),
grid=dwf_outer.F_grid,
cb=None,
)
def __init__(self, coarse_grid, basis, mask=None, basis_n_block=8):
assert type(coarse_grid) == gpt.grid
assert len(basis) > 0
if mask is None:
mask = gpt.complex(basis[0].grid)
mask.checkerboard(basis[0].checkerboard())
mask[:] = 1
else:
assert basis[0].grid is mask.grid
assert len(mask.v_obj) == 1
c_otype = gpt.ot_vector_complex_additive_group(len(basis))
basis_size = c_otype.v_n1[0]
self.coarse_grid = coarse_grid
self.basis = basis
self.obj = cgpt.create_block_map(
coarse_grid.obj,
basis,
basis_size,
basis_n_block,
mask.v_obj[0],
)
def _project(coarse, fine):
assert fine[0].checkerboard().__name__ == basis[0].checkerboard(
).__name__
cgpt.block_project(self.obj, coarse, fine)
def _promote(fine, coarse):
assert fine[0].checkerboard().__name__ == basis[0].checkerboard(
).__name__
cgpt.block_promote(self.obj, coarse, fine)
self.project = gpt.matrix_operator(
mat=_project,
vector_space=(
gpt.vector_space.explicit_grid_otype(coarse_grid, c_otype),
gpt.vector_space.explicit_lattice(basis[0]),
),
accept_list=True,
)
self.promote = gpt.matrix_operator(
mat=_promote,
vector_space=(
gpt.vector_space.explicit_lattice(basis[0]),
gpt.vector_space.explicit_grid_otype(coarse_grid, c_otype),
),
accept_list=True,
)
def __call__(self, mat):
def inv(dst, src):
for i in range(len(dst)):
eps = g.norm2(mat * dst[i] - src[i]) ** 0.5
nrm = g.norm2(src[i]) ** 0.5
if nrm != 0.0:
g.message(
f"{self.tag}| mat * dst[{i}] - src[{i}] | / | src | = {eps/nrm}, | src[{i}] | = {nrm}"
)
else:
g.message(
f"{self.tag}| mat * dst[{i}] - src[{i}] | = {eps}, | src[{i}] | = {nrm}"
)
vector_space = None
if isinstance(mat, g.matrix_operator):
vector_space = mat.vector_space
return g.matrix_operator(
mat=inv,
inv_mat=mat,
vector_space=vector_space,
accept_guess=(True, False),
accept_list=True,
)
def __call__(self, mat):
vector_space = None
if type(mat) == g.matrix_operator:
vector_space = mat.vector_space
inv_mat = self.inverter(mat)
@self.timed_function
def inv(psi, src, t):
t("inverter")
inv_mat(psi, src)
t("update")
space = self.solution_space
if len(space) == self.N:
space.pop()
space.insert(0, psi)
self.log(
f"solution space now has {len(self.solution_space)} / {self.N} elements"
)
return g.matrix_operator(
mat=inv,
inv_mat=mat,
vector_space=vector_space,
accept_guess=(True, False),
)
def __call__(self, mat):
if type(mat) == float or type(mat) == complex or type(mat) == int:
return self.eval(mat)
else:
otype, grid, cb = None, None, None
if type(mat) == g.matrix_operator:
otype, grid, cb = mat.otype, mat.grid, mat.cb
mat = mat.mat # unwrap for performance benefit
def evalOp(dst, src):
dst = make_list(dst)
xscale = 2.0 / (self.hi - self.lo)
mscale = -(self.hi + self.lo) / (self.hi - self.lo)
T0, T1, T2, y = (
g.copy(src),
g.lattice(src),
g.lattice(src),
g.lattice(src),
)
Tnm, Tn, Tnp = T0, T1, T2
mat(y, T0)
T1 @= y * xscale + src * mscale
for i in range(self.n):
dst[i] @= (0.5 * self.coeffs[i][0]) * T0 + self.coeffs[i][1] * T1
for n in range(2, self.morder):
mat(y, Tn)
y @= xscale * y + mscale * Tn
Tnp @= 2.0 * y - Tnm
for i in range(self.n):
if len(self.coeffs[i]) > n:
dst[i] += self.coeffs[i][n] * Tnp
Tnm, Tn, Tnp = Tn, Tnp, Tnm
return g.matrix_operator(evalOp, grid=grid, otype=otype, cb=cb)
def __call__(self, op):
# for now ad-hoc treatment of 5d fermions
if op.F_grid.nd == len(self.bs) + 1:
F_bs = [op.F_grid.fdimensions[0]] + self.bs
else:
F_bs = self.bs
# create domains
F_domains = [
gpt.domain.even_odd_blocks(op.F_grid, F_bs, gpt.even),
gpt.domain.even_odd_blocks(op.F_grid, F_bs, gpt.odd),
]
U_domains = [
gpt.domain.even_odd_blocks(op.U_grid, self.bs, gpt.even),
gpt.domain.even_odd_blocks(op.U_grid, self.bs, gpt.odd),
]
solver = [
self.blk_solver(domain_fermion_operator(op, domain))
for domain in U_domains
]
def inv(dst, src):
dst[:] = 0
eta = gpt.copy(src)
ws = [gpt.copy(src) for _ in range(2)]
for eo in range(2):
ws[0][:] = 0
src_blk = F_domains[eo].lattice(op.otype)
dst_blk = F_domains[eo].lattice(op.otype)
F_domains[eo].project(src_blk, eta)
dst_blk[:] = 0 # for now
solver[eo](dst_blk, src_blk)
F_domains[eo].promote(ws[0], dst_blk)
if eo == 0:
op(ws[1], ws[0])
eta -= ws[1]
dst += ws[0]
gpt.message(
f"SAP cycle; |rho|^2 = {gpt.norm2(eta):g}; |dst|^2 = {gpt.norm2(dst):g}"
)
return gpt.matrix_operator(
mat=inv,
inv_mat=op,
adj_inv_mat=op.adj(),
adj_mat=None,
vector_space=op.vector_space,
accept_guess=(True, False),
)
def __init__(self, op, parity):
super().__init__(op, parity)
def _R(op, i):
self.import_parity(i)
self.op.Mooee.inv_mat(self.tmp, self.in_np)
self.op.Meooe.mat(op, self.tmp)
op @= self.in_p - op
self.op.Mooee.inv_mat(self.tmp, op)
self.N.adj_mat(op, self.tmp)
def _RDag(o, ip):
# R^dag = ( 1 - EO OO^-1 )^dag EE^-1^dag N
self.N.mat(self.out_np, ip)
self.op.Mooee.adj_inv_mat(self.out_p, self.out_np)
self.op.Meooe.adj_mat(self.tmp, self.out_p)
self.op.Mooee.adj_inv_mat(self.out_np, self.tmp)
self.out_np @= -self.out_np
self.export_parity(o)
self.R = gpt.matrix_operator(
mat=_R,
adj_mat=_RDag,
otype=op.otype,
grid=(self.F_grid_eo, self.F_grid),
cb=(self.parity, None),
)
self.Mpc = self.NDagN
def __call__(self, mat):
matrix = self.preconditioner(mat)
inv_mat = self.inverter(matrix.Mpc)
@self.timed_function
def inv(dst, src, t):
t("prepare")
pc_src = g(matrix.R * src)
pc_dst = g(matrix.L.inv() * dst)
t("inv mat")
inv_mat(pc_dst, pc_src)
t("combine")
# TODO: further improve this, maybe understand why eval is not optimal
tmp = g.lattice(dst[0])
# g.eval(dst, matrix.L * pc_dst + matrix.S * src)
for i in range(len(dst)):
matrix.L.mat(dst[i], pc_dst[i])
matrix.S.mat(tmp, src[i])
g.axpy(dst[i], 1.0, dst[i], tmp)
return g.matrix_operator(
mat=inv,
inv_mat=mat,
adj_inv_mat=mat.adj(),
adj_mat=None, # implement adj_mat when needed
otype=mat.otype,
accept_guess=(True, False),
grid=matrix.F_grid,
cb=None,
accept_list=True,
)
def fine_operator(self, coarse_operator):
verbose = gpt.default.is_verbose("block_operator")
coarse_otype = gpt.ot_vector_complex_additive_group(len(self.basis))
def mat(dst, src):
csrc = [gpt.lattice(self.coarse_grid, coarse_otype) for x in src]
cdst = [gpt.lattice(self.coarse_grid, coarse_otype) for x in src]
t0 = gpt.time()
self.project(csrc, src)
t1 = gpt.time()
coarse_operator(cdst, csrc)
t2 = gpt.time()
self.promote(dst, cdst)
t3 = gpt.time()
if verbose:
gpt.message(
"fine_operator acting on %d vector(s) in %g s (project %g s, coarse_operator %g s, promote %g s)"
% (len(src), t3 - t0, t1 - t0, t2 - t1, t3 - t2))
return gpt.matrix_operator(
mat=mat,
vector_space=gpt.vector_space.explicit_lattice(self.basis[0]),
accept_list=True,
)
def __call__(self, mat):
if g.util.is_num(mat):
return self.eval(mat)
else:
vector_space = None
if type(mat) == g.matrix_operator:
vector_space = mat.vector_space
mat = mat.mat
# remove wrapper for performance benefits
if self.inverter is None:
raise NotImplementedError()
pf = self.partial_fractions(mat)
def operator(dst, src):
chi = pf(src)
dst @= src
if self.pf0 == 0.0:
dst[:] = 0
for i, c in enumerate(chi):
dst += self.r[i] * c
if self.norm != 1.0:
dst *= self.norm
return g.matrix_operator(mat=operator, vector_space=vector_space)
def __init__(self, U, boundary_phases):
self.nd = len(U)
self.U = [gpt.copy(u) for u in U]
self.L = U[0].grid.fdimensions
if boundary_phases is not None:
for mu in range(self.nd):
last_slice = tuple([
self.L[mu] - 1 if mu == nu else slice(None, None, None)
for nu in range(self.nd)
])
self.U[mu][
last_slice] = self.U[mu][last_slice] * boundary_phases[mu]
# now take boundary_phase from params and apply here
self.Udag = [gpt.eval(gpt.adj(u)) for u in self.U]
def _forward(mu):
def wrap(dst, src):
dst @= self.U[mu] * gpt.cshift(src, mu, +1)
return wrap
def _backward(mu):
def wrap(dst, src):
dst @= gpt.cshift(self.Udag[mu] * src, mu, -1)
return wrap
self.forward = [
gpt.matrix_operator(mat=_forward(mu), inv_mat=_backward(mu))
for mu in range(self.nd)
]
self.backward = [o.inv() for o in self.forward]
def coarse_operator(self, fine_operator):
verbose = gpt.default.is_verbose("block_operator")
def mat(dst_coarse, src_coarse):
src_fine = [gpt.lattice(self.basis[0]) for x in src_coarse]
dst_fine = [gpt.lattice(self.basis[0]) for x in src_coarse]
t0 = gpt.time()
self.promote(src_fine, src_coarse)
t1 = gpt.time()
fine_operator(dst_fine, src_fine)
t2 = gpt.time()
self.project(dst_coarse, dst_fine)
t3 = gpt.time()
if verbose:
gpt.message(
"coarse_operator acting on %d vector(s) in %g s (promote %g s, fine_operator %g s, project %g s)"
% (len(src_coarse), t3 - t0, t1 - t0, t2 - t1, t3 - t2))
otype = gpt.ot_vector_complex_additive_group(len(self.basis))
return gpt.matrix_operator(
mat=mat,
vector_space=gpt.vector_space.explicit_grid_otype(
self.coarse_grid, otype),
accept_list=True,
)
def kappa(self):
b = self.params["b"]
c = self.params["c"]
M5 = self.params["M5"]
bs = [
1.0 / 2.0 * (1.0 / omega_s * (b + c) + (b - c))
for omega_s in self.params["omega"]
]
kappa = np.array(
[1.0 / (2.0 * (bsi * (4.0 - M5) + 1.0)) for bsi in bs],
np.complex128)
adj_kappa = np.conj(kappa)
inv_kappa = 1.0 / kappa
adj_inv_kappa = np.conj(inv_kappa)
def _mat(dst, src):
gpt.scale_per_coordinate(dst, src, kappa, 0)
def _inv_mat(dst, src):
gpt.scale_per_coordinate(dst, src, inv_kappa, 0)
def _adj_mat(dst, src):
gpt.scale_per_coordinate(dst, src, adj_kappa, 0)
def _adj_inv_mat(dst, src):
gpt.scale_per_coordinate(dst, src, adj_inv_kappa, 0)
return gpt.matrix_operator(mat=_mat,
inv_mat=_inv_mat,
adj_mat=_adj_mat,
adj_inv_mat=_adj_inv_mat)
def fine_operator(self, coarse_operator):
verbose = gpt.default.is_verbose("block_operator")
coarse_otype = gpt.ot_vector_singlet(len(self.basis))
otype = self.basis[0].otype
grid = self.basis[0].grid
cb = self.basis[0].checkerboard()
def mat(dst, src):
csrc = [gpt.lattice(self.coarse_grid, coarse_otype) for x in src]
cdst = [gpt.lattice(self.coarse_grid, coarse_otype) for x in src]
t0 = gpt.time()
self.project(csrc, src)
t1 = gpt.time()
coarse_operator(cdst, csrc)
t2 = gpt.time()
self.promote(dst, cdst)
t3 = gpt.time()
if verbose:
gpt.message(
"fine_operator acting on %d vector(s) in %g s (project %g s, coarse_operator %g s, promote %g s)"
% (len(src), t3 - t0, t1 - t0, t2 - t1, t3 - t2)
)
return gpt.matrix_operator(
mat=mat, otype=otype, grid=grid, cb=cb, accept_list=True
)
def __call__(self, matrix=None):
# ignore matrix
left = self.left
right = self.right
evals = self.evals
f_evals = [self.f(x) for x in evals]
otype = (left[0].otype, right[0].otype)
grid = (left[0].grid, right[0].grid)
cb = (left[0].checkerboard(), right[0].checkerboard())
def approx(dst, src):
assert src != dst
verbose = g.default.is_verbose("modes")
t0 = g.time()
dst[:] = 0
for i, x in enumerate(left):
dst += f_evals[i] * x * g.innerProduct(right[i], src)
if verbose:
t1 = g.time()
g.message("Approximation by %d modes took %g s" %
(len(left), t1 - t0))
return g.matrix_operator(mat=approx,
otype=otype,
zero=(False, False),
grid=grid,
cb=cb)
def __call__(self, mat):
otype, grid, cb = None, None, None
if type(mat) == g.matrix_operator:
otype, grid, cb = mat.otype, mat.grid, mat.cb
mat = mat.mat
# remove wrapper for performance benefits
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
return g.matrix_operator(mat=inv,
inv_mat=mat,
otype=otype,
zero=(True, False),
grid=grid,
cb=cb)
def __init__(self, coarse_grid, basis, mask=None, basis_n_block=8):
assert type(coarse_grid) == gpt.grid
assert len(basis) > 0
if mask is None:
mask = gpt.complex(basis[0].grid)
mask.checkerboard(basis[0].checkerboard())
mask[:] = 1
else:
assert basis[0].grid is mask.grid
assert len(mask.v_obj) == 1
c_otype = gpt.ot_vsinglet(len(basis))
basis_size = c_otype.v_n1[0]
self.coarse_grid = coarse_grid
self.basis = basis
self.obj = cgpt.create_block_map(
coarse_grid.obj,
basis,
basis_size,
basis_n_block,
mask.v_obj[0],
)
def _project(coarse, fine):
assert fine[0].checkerboard().__name__ == basis[0].checkerboard().__name__
cgpt.block_project(self.obj, coarse, fine)
def _promote(fine, coarse):
assert fine[0].checkerboard().__name__ == basis[0].checkerboard().__name__
cgpt.block_promote(self.obj, coarse, fine)
self.project = gpt.matrix_operator(
mat=_project,
otype=(c_otype, basis[0].otype),
grid=(coarse_grid, basis[0].grid),
cb=(None, basis[0].checkerboard()),
accept_list=True,
)
self.promote = gpt.matrix_operator(
mat=_promote,
otype=(basis[0].otype, c_otype),
grid=(basis[0].grid, coarse_grid),
cb=(basis[0].checkerboard(), None),
accept_list=True,
)
def __call__(self, mat):
vector_space = None
if type(mat) == g.matrix_operator:
vector_space = mat.vector_space
mat = mat.mat
# remove wrapper for performance benefits
@self.timed_function
def inv(psi, src, t):
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.log_convergence(k, r2, rsq)
if r2 <= rsq:
self.log(f"converged in {k+1} iterations")
return
self.log(
f"NOT converged in {k+1} iterations; squared residual {r2:e} / {rsq:e}"
)
return g.matrix_operator(
mat=inv,
inv_mat=mat,
vector_space=vector_space,
accept_guess=(True, False),
)
def Mdir(self, mu, fb):
op = gpt.matrix_operator(
mat=lambda dst, src: self._Mdir(dst, src, mu, fb),
vector_space=self.vector_space_F,
)
if self.daggered:
op = op.adj()
return op
def wrap(Mpc, L, R):
tmp = Mpc.vector_space[0].lattice()
def _Mpc(o_d, i_d):
V.inv_mat(o_d, i_d)
Mpc.mat(tmp, o_d)
V.mat(o_d, tmp)
def _Mpc_dag(o_d, i_d):
V.adj_mat(o_d, i_d)
Mpc.adj_mat(tmp, o_d)
V.adj_inv_mat(o_d, tmp)
def _R(o_d, i):
R.mat(tmp, i)
V.mat(o_d, tmp)
def _R_dag(o, i_d):
V.adj_mat(tmp, i_d)
R.adj_mat(o, tmp)
def _L(o, i_d):
V.inv_mat(tmp, i_d)
L.mat(o, tmp)
def _L_inv(o_d, i):
L.inv_mat(tmp, i)
V.mat(o_d, tmp)
wrapped_R = gpt.matrix_operator(
mat=_R,
adj_mat=_R_dag,
vector_space=R.vector_space,
)
wrapped_L = gpt.matrix_operator(
mat=_L,
inv_mat=_L_inv,
vector_space=L.vector_space,
)
wrapped_Mpc = gpt.matrix_operator(mat=_Mpc,
adj_mat=_Mpc_dag,
vector_space=Mpc.vector_space)
return wrapped_Mpc, wrapped_L, wrapped_R
def laplace(cov, dimensions):
def mat(dst, src):
assert dst != src
dst[:] = 0.0
for mu in dimensions:
dst += g.eval(-2.0 * src + cov.forward[mu] * src + cov.backward[mu] * src)
return g.matrix_operator(mat=mat)
def __init__(self, name, U, params, otype=None):
differentiable_fine_operator.__init__(self, name, U, params, otype)
def _G5(dst, src):
dst @= gpt.gamma[5] * src
gauge_independent_g5_hermitian.__init__(
self, gpt.matrix_operator(_G5, vector_space=self.vector_space)
)
def __call__(self, mat):
otype, grid, cb = None, None, None
if type(mat) == g.matrix_operator:
otype, grid, cb = mat.otype, mat.grid, mat.cb
mat = mat.mat
# remove wrapper for performance benefits
@self.timed_function
def inv(psi, src, t):
assert src != psi
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(self.maxiter):
c = cp
t("matrix")
mat(mmp, p)
t("inner_product")
dc = g.inner_product(p, mmp)
d = dc.real
a = c / d
t("axpy_norm2")
cp = g.axpy_norm2(r, -a, mmp, r)
t("linear combination")
b = cp / c
psi += a * p
p @= b * p + r
t("other")
self.log_convergence(k, cp, rsq)
if cp <= rsq:
self.log(f"converged in {k+1} iterations")
return
self.log(
f"NOT converged in {k+1} iterations; squared residual {cp:e} / {rsq:e}"
)
return g.matrix_operator(
mat=inv,
inv_mat=mat,
otype=otype,
accept_guess=(True, False),
grid=grid,
cb=cb,
)
def __init__(self, grid, local_coordinates):
self.local_coordinates = local_coordinates
self.grid = grid
# kernel to avoid circular references through captures below
kernel = sparse_kernel(grid, local_coordinates)
def _project(dst, src):
for d, s in zip(dst, src):
d[kernel.embedded_coordinates,
get_cache(kernel.embedded_cache, d)] = s[
kernel.local_coordinates,
get_cache(kernel.local_cache, s)]
def _promote(dst, src):
for d, s in zip(dst, src):
d[kernel.local_coordinates,
get_cache(kernel.local_cache, d)] = s[
kernel.embedded_coordinates,
get_cache(kernel.embedded_cache, s)]
self.project = gpt.matrix_operator(
mat=_project,
vector_space=(
gpt.vector_space.explicit_grid(kernel.embedding_grid),
gpt.vector_space.explicit_grid(kernel.grid),
),
accept_list=True,
accept_guess=True,
)
self.promote = gpt.matrix_operator(
mat=_promote,
vector_space=(
gpt.vector_space.explicit_grid(kernel.grid),
gpt.vector_space.explicit_grid(kernel.embedding_grid),
),
accept_list=True,
accept_guess=True,
)
self.kernel = kernel
def __call__(self, op):
sap = sap_instance(op, self.bs)
otype = sap.op.otype[0]
src_blk = gpt.lattice(sap.op_blk[0].F_grid, otype)
dst_blk = gpt.lattice(sap.op_blk[0].F_grid, otype)
solver = [self.blk_solver(op) for op in sap.op_blk]
def inv(dst, src):
dst[:] = 0
eta = gpt.copy(src)
ws = [gpt.copy(src) for _ in range(2)]
dt_solv = dt_distr = dt_hop = 0.0
for eo in range(2):
ws[0][:] = 0
dt_distr -= gpt.time()
src_blk[sap.pos] = eta[
sap.coor[eo]
] # reminder view interface eta[[pos]], ... eta[...,idx]
dt_distr += gpt.time()
dt_solv -= gpt.time()
dst_blk[:] = 0 # for now
solver[eo](dst_blk, src_blk)
dt_solv += gpt.time()
dt_distr -= gpt.time()
ws[0][sap.coor[eo]] = dst_blk[sap.pos]
dt_distr += gpt.time()
dt_hop -= gpt.time()
if eo == 0:
sap.op(ws[1], ws[0])
eta -= ws[1]
dst += ws[0]
dt_hop += gpt.time()
gpt.message(
f"SAP cycle; |rho|^2 = {gpt.norm2(eta):g}; |dst|^2 = {gpt.norm2(dst):g}"
)
gpt.message(
f"SAP Timings: distr {dt_distr:g} secs, blk_solver {dt_solv:g} secs, hop+update {dt_hop:g} secs"
)
return gpt.matrix_operator(
mat=inv,
inv_mat=sap.op,
adj_inv_mat=sap.op.adj(),
adj_mat=None, # implement adj_mat when needed
otype=otype,
accept_guess=(True, False),
grid=sap.op.F_grid,
cb=None,
)
