mirror of
https://github.com/fjosw/pyerrors.git
synced 2025-03-15 14:50:25 +01:00
Matmul overloaded for correlator class. (#199)
* feat: matmul method added to correlator class. * feat: corr, corr matmul and correlator matrix trace added. * tests: tests for matmul and trace added. * tests: slightly reduced tolerance and good guess bad guess test. * feat: rmatmul added and __array_priority__ set. * tests: additional tests for rmatmul added. * tests: one more tests for rmatmul added. * docs: docstring added to Corr.trace. * tests: associative property test added for complex Corr matmul. * fix: Corr.roll method now also works for correlator matrices by explicitly specifying the axis. Co-authored-by: Matteo Di Carlo <matteo.dicarlo93@gmail.com> * feat: exception type for correlator trace of 1dim correlator changed. * tests: trace N=1 exception tested. --------- Co-authored-by: Matteo Di Carlo <matteo.dicarlo93@gmail.com>
This commit is contained in:
parent
7d1858f6c4
commit
f1150f09c8
3 changed files with 185 additions and 6 deletions
|
@ -220,7 +220,7 @@ class Corr:
|
|||
def anti_symmetric(self):
|
||||
"""Anti-symmetrize the correlator around x0=0."""
|
||||
if self.N != 1:
|
||||
raise Exception('anti_symmetric cannot be safely applied to multi-dimensional correlators.')
|
||||
raise TypeError('anti_symmetric cannot be safely applied to multi-dimensional correlators.')
|
||||
if self.T % 2 != 0:
|
||||
raise Exception("Can not symmetrize odd T")
|
||||
|
||||
|
@ -242,7 +242,7 @@ class Corr:
|
|||
def is_matrix_symmetric(self):
|
||||
"""Checks whether a correlator matrices is symmetric on every timeslice."""
|
||||
if self.N == 1:
|
||||
raise Exception("Only works for correlator matrices.")
|
||||
raise TypeError("Only works for correlator matrices.")
|
||||
for t in range(self.T):
|
||||
if self[t] is None:
|
||||
continue
|
||||
|
@ -254,6 +254,18 @@ class Corr:
|
|||
return False
|
||||
return True
|
||||
|
||||
def trace(self):
|
||||
"""Calculates the per-timeslice trace of a correlator matrix."""
|
||||
if self.N == 1:
|
||||
raise ValueError("Only works for correlator matrices.")
|
||||
newcontent = []
|
||||
for t in range(self.T):
|
||||
if _check_for_none(self, self.content[t]):
|
||||
newcontent.append(None)
|
||||
else:
|
||||
newcontent.append(np.trace(self.content[t]))
|
||||
return Corr(newcontent)
|
||||
|
||||
def matrix_symmetric(self):
|
||||
"""Symmetrizes the correlator matrices on every timeslice."""
|
||||
if self.N == 1:
|
||||
|
@ -405,7 +417,7 @@ class Corr:
|
|||
dt : int
|
||||
number of timeslices
|
||||
"""
|
||||
return Corr(list(np.roll(np.array(self.content, dtype=object), dt)))
|
||||
return Corr(list(np.roll(np.array(self.content, dtype=object), dt, axis=0)))
|
||||
|
||||
def reverse(self):
|
||||
"""Reverse the time ordering of the Corr"""
|
||||
|
@ -1020,6 +1032,8 @@ class Corr:
|
|||
# This is because Obs*Corr checks Obs.__mul__ first and does not catch an exception.
|
||||
# One could try and tell Obs to check if the y in __mul__ is a Corr and
|
||||
|
||||
__array_priority__ = 10000
|
||||
|
||||
def __add__(self, y):
|
||||
if isinstance(y, Corr):
|
||||
if ((self.N != y.N) or (self.T != y.T)):
|
||||
|
@ -1076,6 +1090,49 @@ class Corr:
|
|||
else:
|
||||
raise TypeError("Corr * wrong type")
|
||||
|
||||
def __matmul__(self, y):
|
||||
if isinstance(y, np.ndarray):
|
||||
if y.ndim != 2 or y.shape[0] != y.shape[1]:
|
||||
raise ValueError("Can only multiply correlators by square matrices.")
|
||||
if not self.N == y.shape[0]:
|
||||
raise ValueError("matmul: mismatch of matrix dimensions")
|
||||
newcontent = []
|
||||
for t in range(self.T):
|
||||
if _check_for_none(self, self.content[t]):
|
||||
newcontent.append(None)
|
||||
else:
|
||||
newcontent.append(self.content[t] @ y)
|
||||
return Corr(newcontent)
|
||||
elif isinstance(y, Corr):
|
||||
if not self.N == y.N:
|
||||
raise ValueError("matmul: mismatch of matrix dimensions")
|
||||
newcontent = []
|
||||
for t in range(self.T):
|
||||
if _check_for_none(self, self.content[t]) or _check_for_none(y, y.content[t]):
|
||||
newcontent.append(None)
|
||||
else:
|
||||
newcontent.append(self.content[t] @ y.content[t])
|
||||
return Corr(newcontent)
|
||||
|
||||
else:
|
||||
return NotImplemented
|
||||
|
||||
def __rmatmul__(self, y):
|
||||
if isinstance(y, np.ndarray):
|
||||
if y.ndim != 2 or y.shape[0] != y.shape[1]:
|
||||
raise ValueError("Can only multiply correlators by square matrices.")
|
||||
if not self.N == y.shape[0]:
|
||||
raise ValueError("matmul: mismatch of matrix dimensions")
|
||||
newcontent = []
|
||||
for t in range(self.T):
|
||||
if _check_for_none(self, self.content[t]):
|
||||
newcontent.append(None)
|
||||
else:
|
||||
newcontent.append(y @ self.content[t])
|
||||
return Corr(newcontent)
|
||||
else:
|
||||
return NotImplemented
|
||||
|
||||
def __truediv__(self, y):
|
||||
if isinstance(y, Corr):
|
||||
if not ((self.N == 1 or y.N == 1 or self.N == y.N) and self.T == y.T):
|
||||
|
|
|
@ -570,3 +570,125 @@ def test_corr_symmetric():
|
|||
assert scorr[1] == scorr[3]
|
||||
assert scorr[2] == corr[2]
|
||||
assert scorr[0] == corr[0]
|
||||
|
||||
|
||||
def test_two_matrix_corr_inits():
|
||||
T = 4
|
||||
rn = lambda : np.random.normal(0.5, 0.1)
|
||||
|
||||
# Generate T random CObs in a list
|
||||
list_of_timeslices =[]
|
||||
for i in range(T):
|
||||
re = pe.pseudo_Obs(rn(), rn(), "test")
|
||||
im = pe.pseudo_Obs(rn(), rn(), "test")
|
||||
list_of_timeslices.append(pe.CObs(re, im))
|
||||
|
||||
# First option: Correlator of matrix of correlators
|
||||
corr = pe.Corr(list_of_timeslices)
|
||||
mat_corr1 = pe.Corr(np.array([[corr, corr], [corr, corr]]))
|
||||
|
||||
# Second option: Correlator of list of arrays per timeslice
|
||||
list_of_arrays = [np.array([[elem, elem], [elem, elem]]) for elem in list_of_timeslices]
|
||||
mat_corr2 = pe.Corr(list_of_arrays)
|
||||
|
||||
for el in mat_corr1 - mat_corr2:
|
||||
assert np.all(el == 0)
|
||||
|
||||
|
||||
def test_matmul_overloading():
|
||||
N = 4
|
||||
rn = lambda : np.random.normal(0.5, 0.1)
|
||||
|
||||
# Generate N^2 random CObs and assemble them in an array
|
||||
ll =[]
|
||||
for i in range(N ** 2):
|
||||
re = pe.pseudo_Obs(rn(), rn(), "test")
|
||||
im = pe.pseudo_Obs(rn(), rn(), "test")
|
||||
ll.append(pe.CObs(re, im))
|
||||
mat = np.array(ll).reshape(N, N)
|
||||
|
||||
# Multiply with gamma matrix
|
||||
corr = pe.Corr([mat] * 4, padding=[0, 1])
|
||||
|
||||
# __matmul__
|
||||
mcorr = corr @ pe.dirac.gammaX
|
||||
comp = mat @ pe.dirac.gammaX
|
||||
for i in range(4):
|
||||
assert np.all(mcorr[i] == comp)
|
||||
|
||||
# __rmatmul__
|
||||
mcorr = pe.dirac.gammaX @ corr
|
||||
comp = pe.dirac.gammaX @ mat
|
||||
for i in range(4):
|
||||
assert np.all(mcorr[i] == comp)
|
||||
|
||||
test_mat = pe.dirac.gamma5 + pe.dirac.gammaX
|
||||
icorr = corr @ test_mat @ np.linalg.inv(test_mat)
|
||||
tt = corr - icorr
|
||||
for i in range(4):
|
||||
assert np.all(tt[i] == 0)
|
||||
|
||||
# associative property
|
||||
tt = (corr.real @ pe.dirac.gammaX + corr.imag @ (pe.dirac.gammaX * 1j)) - corr @ pe.dirac.gammaX
|
||||
for el in tt:
|
||||
if el is not None:
|
||||
assert np.all(el == 0)
|
||||
|
||||
corr2 = corr @ corr
|
||||
for i in range(4):
|
||||
np.all(corr2[i] == corr[i] @ corr[i])
|
||||
|
||||
|
||||
def test_matrix_trace():
|
||||
N = 4
|
||||
rn = lambda : np.random.normal(0.5, 0.1)
|
||||
|
||||
# Generate N^2 random CObs and assemble them in an array
|
||||
ll =[]
|
||||
for i in range(N ** 2):
|
||||
re = pe.pseudo_Obs(rn(), rn(), "test")
|
||||
im = pe.pseudo_Obs(rn(), rn(), "test")
|
||||
ll.append(pe.CObs(re, im))
|
||||
mat = np.array(ll).reshape(N, N)
|
||||
|
||||
corr = pe.Corr([mat] * 4)
|
||||
|
||||
# Explicitly check trace
|
||||
for el in corr.trace():
|
||||
el == np.sum(np.diag(mat))
|
||||
|
||||
# Trace is cyclic
|
||||
for one, two in zip((pe.dirac.gammaX @ corr).trace(), (corr @ pe.dirac.gammaX).trace()):
|
||||
assert np.all(one == two)
|
||||
|
||||
# Antisymmetric matrices are traceless.
|
||||
mat = (mat - mat.T) / 2
|
||||
corr = pe.Corr([mat] * 4)
|
||||
for el in corr.trace():
|
||||
assert el == 0
|
||||
|
||||
|
||||
with pytest.raises(ValueError):
|
||||
corr.item(0, 0).trace()
|
||||
|
||||
|
||||
def test_corr_roll():
|
||||
T = 4
|
||||
rn = lambda : np.random.normal(0.5, 0.1)
|
||||
|
||||
ll = []
|
||||
for i in range(T):
|
||||
re = pe.pseudo_Obs(rn(), rn(), "test")
|
||||
im = pe.pseudo_Obs(rn(), rn(), "test")
|
||||
ll.append(pe.CObs(re, im))
|
||||
|
||||
# Rolling by T should produce the same correlator
|
||||
corr = pe.Corr(ll)
|
||||
tt = corr - corr.roll(T)
|
||||
for el in tt:
|
||||
assert np.all(el == 0)
|
||||
|
||||
mcorr = pe.Corr(np.array([[corr, corr + 0.1], [corr - 0.1, 2 * corr]]))
|
||||
tt = mcorr.roll(T) - mcorr
|
||||
for el in tt:
|
||||
assert np.all(el == 0)
|
||||
|
|
|
@ -235,12 +235,12 @@ def test_fit_corr_independent():
|
|||
|
||||
|
||||
def test_linear_fit_guesses():
|
||||
for err in [10, 0.1, 0.001]:
|
||||
for err in [1.2, 0.1, 0.001]:
|
||||
xvals = []
|
||||
yvals = []
|
||||
for x in range(1, 8, 2):
|
||||
xvals.append(x)
|
||||
yvals.append(pe.pseudo_Obs(x + np.random.normal(0.0, err), err, 'test1') + pe.pseudo_Obs(0, err / 100, 'test2', samples=87))
|
||||
yvals.append(pe.pseudo_Obs(x + np.random.normal(0.0, err), err, 'test1') + pe.pseudo_Obs(0, err / 97, 'test2', samples=87))
|
||||
lin_func = lambda a, x: a[0] + a[1] * x
|
||||
with pytest.raises(Exception):
|
||||
pe.least_squares(xvals, yvals, lin_func)
|
||||
|
@ -251,7 +251,7 @@ def test_linear_fit_guesses():
|
|||
bad_guess = pe.least_squares(xvals, yvals, lin_func, initial_guess=[999, 999])
|
||||
good_guess = pe.least_squares(xvals, yvals, lin_func, initial_guess=[0, 1])
|
||||
assert np.isclose(bad_guess.chisquare, good_guess.chisquare, atol=1e-8)
|
||||
assert np.all([(go - ba).is_zero(atol=1e-6) for (go, ba) in zip(good_guess, bad_guess)])
|
||||
assert np.all([(go - ba).is_zero(atol=5e-5) for (go, ba) in zip(good_guess, bad_guess)])
|
||||
|
||||
|
||||
def test_total_least_squares():
|
||||
|
|
Loading…
Add table
Reference in a new issue