Skip to content

Commit

Permalink
full gradients
Browse files Browse the repository at this point in the history
  • Loading branch information
@lehner
lehner committed Jun 7, 2024
1 parent c4ab7a6 commit 1a19be1
Show file tree
Hide file tree
Showing 3 changed files with 194 additions and 52 deletions.
124 changes: 110 additions & 14 deletions lib/gpt/qcd/gauge/stencil/algebra_laplace.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -20,6 +20,15 @@
import numpy as np


class fixed_gauge_algebra_laplace:
def __init__(self, parent, U):
self.parent = parent
self.U = U

def __call__(self, dst, src):
self.parent(self.U + dst, self.U + src)


class algebra_laplace:
def __init__(self, U):
self.U = U
Expand All @@ -37,39 +46,126 @@ def __init__(self, U):
_Sm = list(range(1 + self.nd, 1 + 2 * self.nd))

# indices
_dst = list(range(self.nd))
_src = [x + self.nd for x in _dst]
_U = [x + self.nd for x in _src]
_dU = list(range(self.nd))
_dV = [x + self.nd for x in _dU]
_sU = [x + self.nd for x in _dV]
_sV = [x + self.nd for x in _sU]

# code
code = []
for nu in range(self.nd):
code.append((_dU[nu], -1, 1.0, [(_sU[nu], _P, 0)]))
code.append((_dV[nu], -1, -2 * self.nd / 16, [(_sV[nu], _P, 0)]))

for mu in range(self.nd):
code.append(
(
_dst[nu],
-1 if mu == 0 else _dst[nu],
_dV[nu],
_dV[nu],
1 / 16,
[(_U[mu], _P, 0), (_src[nu], _Sp[mu], 0), (_U[mu], _P, 1)],
[(_sU[mu], _P, 0), (_sV[nu], _Sp[mu], 0), (_sU[mu], _P, 1)],
)
)

code.append(
(
_dst[nu],
_dst[nu],
_dV[nu],
_dV[nu],
1 / 16,
[(_U[mu], _Sm[mu], 1), (_src[nu], _Sm[mu], 0), (_U[mu], _Sm[mu], 0)],
[(_sU[mu], _Sm[mu], 1), (_sV[nu], _Sm[mu], 0), (_sU[mu], _Sm[mu], 0)],
)
)

code.append((_dst[nu], _dst[nu], -2 / 16, [(_src[nu], _P, 0)]))

# stencil
self.st = g.stencil.matrix(U[0], [origin] + evec + nevec, code)

# indices for gradient (ret, left, U, right)
_ret = list(range(self.nd))
_left = [x + self.nd for x in _ret]
_U = [x + self.nd for x in _left]
_right = [x + self.nd for x in _U]

# code
code = []
for mu in range(self.nd):
for nu in range(self.nd):
code.append(
(
_ret[mu],
_ret[mu] if nu > 0 else -1,
1j / 32,
[
(_U[mu], _P, 0),
(_right[nu], _Sp[mu], 0),
(_U[mu], _P, 1),
(_left[nu], _P, 1),
],
)
)
code.append(
(
_ret[mu],
_ret[mu],
-1j / 32,
[
(_right[nu], _P, 0),
(_U[mu], _P, 0),
(_left[nu], _Sp[mu], 1),
(_U[mu], _P, 1),
],
)
)
code.append(
(
_ret[mu],
_ret[mu],
-1j / 32,
[
(_left[nu], _P, 1),
(_U[mu], _P, 0),
(_right[nu], _Sp[mu], 0),
(_U[mu], _P, 1),
],
)
)
code.append(
(
_ret[mu],
_ret[mu],
1j / 32,
[
(_U[mu], _P, 0),
(_left[nu], _Sp[mu], 1),
(_U[mu], _P, 1),
(_right[nu], _P, 0),
],
)
)

# stencil
self.st_pg = g.stencil.matrix(U[0], [origin] + evec + nevec, code)

def __call__(self, dst, src):
assert len(src) == len(dst)
assert len(src) == self.nd
# src = U + V
assert len(dst) == 2 * self.nd
assert len(src) == 2 * self.nd

for i in range(len(dst)):
dst[i].otype = src[i].otype

self.st(*dst, *src)

def projected_gradient(self, left, U, right):
assert len(left) == self.nd
assert len(right) == self.nd
assert len(U) == self.nd

ret = [g.group.cartesian(u) for u in U]
self.st_pg(*ret, *left, *U, *right)
for mu in range(self.nd):
# could create a stencil version of project.traceless_hermitian that works on all ret[mu] at the same time
ret[mu] @= g.qcd.gauge.project.traceless_hermitian(ret[mu])
return ret

self.st(*dst, *src, *self.U)
def fixed_gauge(self, U):
return fixed_gauge_algebra_laplace(self, U)
64 changes: 45 additions & 19 deletions lib/gpt/qcd/scalar/action/mass_term.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -128,51 +128,77 @@ def gradient(self, pi, dpi):


class general_mass_term(differentiable_functional):
def __init__(self, M, sqrt_M):
def __init__(self, M, sqrt_M, M_projected_gradient):
self.M = M
self.sqrt_M = sqrt_M
self.inv_sqrt_M = sqrt_M.inv()
self.M_projected_gradient = M_projected_gradient
# Need:
# - sqrt_M^dag = sqrt_M
# - sqrt_M^2 = M

def __call__(self, pi):
pi_prime = self.M(pi)
n = len(pi)
# - M_projected_gradient = D[vec^dag M(e^{iTa eps} U) vec, eps] iTa

def __call__(self, fields):
# fields = U + pi
n = len(fields)
assert n % 2 == 0
n //= 2
pi = fields[n:]
pi_prime = self.M(fields)[n:]
A = 0.0
for mu in range(n):
A += g.inner_product(pi[mu], pi_prime[mu])
return A.real

def draw(self, pi, rng):
def draw(self, fields, rng):
# fields = U + pi
n = len(fields)
assert n % 2 == 0
n //= 2
pi = fields[n:]

# P(pi) = e^{-pi^dag sqrt_M^dag sqrt_M pi * 2 / 2}
rng.normal_element(pi) # pi = sqrt(2) sqrt_M pi_desired -> pi_desired = inv_sqrt_M pi

n = len(pi)

value = g.group.inner_product(pi, pi) * 0.5

pi_prime = self.inv_sqrt_M(pi)
for mu in range(n):
pi_prime[mu].otype = pi[mu].otype
pi[mu] @= g.project(pi_prime[mu], "defect")

return value

# U -> V(0) U V(1)^-1
# e^{i A} U -> V(0) e^{i A} U V(1)^{-1} -> A -> V(0) A V(0)^-1

@differentiable_functional.multi_field_gradient
def gradient(self, pi, dpi):
dS = []

pi_prime = self.M(pi)
def gradient(self, fields, dfields):
# fields = U + pi
n = len(fields)
assert n % 2 == 0
n //= 2
U = fields[0:n]
pi = fields[n:]

for _pi in dpi:
mu = pi.index(_pi)
dS = []

ret = pi_prime[mu]
pi_prime = self.M(fields)[n:]
dU = [f for f in dfields if f in U]
if len(dU) > 0:
grad_U = self.M_projected_gradient(U, pi)
else:
grad_U = None

for _field in dfields:
mu = fields.index(_field)

if mu >= n:
mu -= n
# pi[mu]
ret = pi_prime[mu]
ret = g(g.project(ret, "defect"))
else:
# U[mu]
ret = grad_U[mu]

ret = g(g.project(ret, "defect"))
dS.append(ret)

return dS
58 changes: 39 additions & 19 deletions tests/qcd/scalar.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -149,49 +149,69 @@

# first test laplacian
def __slap(dst, src):
assert len(src) == 4
for nu in range(len(dst)):
dst[nu][:] = 0
sU = src[0:4]
sV = src[4:]
dU = dst[0:4]
dV = dst[4:]
assert len(dst) == len(src)

for nu in range(len(dst) - 4):
dV[nu][:] = 0
for mu in range(4):
src_p = g.cshift(src[nu], mu, 1)
src_m = g.cshift(src[nu], mu, -1)
U_m = g.cshift(U[mu], mu, -1)
dst[nu] += (
1 / 16 * (U[mu] * src_p * g.adj(U[mu]) + g.adj(U_m) * src_m * U_m - 2 * src[nu])
sV_p = g.cshift(sV[nu], mu, 1)
sV_m = g.cshift(sV[nu], mu, -1)
sU_m = g.cshift(sU[mu], mu, -1)
dV[nu] += (
1 / 16 * (sU[mu] * sV_p * g.adj(sU[mu]) + g.adj(sU_m) * sV_m * sU_m - 2 * sV[nu])
)

g.copy(dU, sU)


lap = g.qcd.gauge.algebra_laplace(U)

tmp = g.copy(U_mom)
tmp2 = g.copy(U_mom)
tmp = g.copy(U + U_mom)
tmp2 = g.copy(U + U_mom)

lap(tmp, U_mom)
__slap(tmp2, U_mom)
rng.element(tmp + tmp2)

for mu in range(4):
lap(tmp, U + U_mom)
__slap(tmp2, U + U_mom)

for mu in range(8):
eps = g.norm2(tmp[mu] - tmp2[mu])
g.message(f"Test laplacian: {eps}")
assert eps < 1e-10

cg = g.algorithms.inverter.block_cg({"eps": 1e-12, "maxiter": 100})
slap = g.matrix_operator(mat=lap, inv_mat=cg(lap), accept_list=True, accept_guess=(False, True))
slap = g.matrix_operator(
mat=lap, inv_mat=cg(lap.fixed_gauge(U)), accept_list=True, accept_guess=(False, True)
)
slap2 = slap * slap

A2 = g.qcd.scalar.action.general_mass_term(M=slap2, sqrt_M=slap)

A2.assert_gradient_error(rng, U_mom, U_mom, 1e-3, 1e-8)
def slap2_pgrad(U, vec):
vec_prime = g(slap * (U + vec))[-len(vec) :]
grad1 = lap.projected_gradient(vec, U, vec_prime)
return [g(2 * x) for x in grad1]


A2 = g.qcd.scalar.action.general_mass_term(M=slap2, sqrt_M=slap, M_projected_gradient=slap2_pgrad)

A2.assert_gradient_error(rng, U + U_mom, U + U_mom, 1e-3, 1e-8)

# test distribution
r = A2.draw(U_mom, rng)
assert abs(A2(U_mom) / r - 1) < 1e-10
r = A2.draw(U + U_mom, rng)
eps = abs(A2(U + U_mom) / r - 1)
g.message(f"Draw test: {eps}")
assert eps < 1e-10
eps = g.group.defect(U_mom[0])
g.message("Group defect:", eps)
assert eps < 1e-10

ngen = len(U_mom[0].otype.generators(np.complex128))
x = 0.0
U_mom_prime = g(slap2 * U_mom)
U_mom_prime = g(slap2 * (U + U_mom))[4:]
for mu in range(4):
x += (
g.sum(g(g.trace(g.adj(U_mom[mu]) * U_mom_prime[mu]) * 2.0 / ngen)).real
Expand Down

0 comments on commit 1a19be1

Please sign in to comment.