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, )