From f1150f09c826ae78b828866401bc6fda1b668edd Mon Sep 17 00:00:00 2001 From: Fabian Joswig Date: Mon, 17 Jul 2023 11:48:57 +0100 Subject: [PATCH] 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 * feat: exception type for correlator trace of 1dim correlator changed. * tests: trace N=1 exception tested. --------- Co-authored-by: Matteo Di Carlo --- pyerrors/correlators.py | 63 +++++++++++++++++++- tests/correlators_test.py | 122 ++++++++++++++++++++++++++++++++++++++ tests/fits_test.py | 6 +- 3 files changed, 185 insertions(+), 6 deletions(-) diff --git a/pyerrors/correlators.py b/pyerrors/correlators.py index 85e28f4d..d78a3c28 100644 --- a/pyerrors/correlators.py +++ b/pyerrors/correlators.py @@ -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): diff --git a/tests/correlators_test.py b/tests/correlators_test.py index c5b530cb..0502462c 100644 --- a/tests/correlators_test.py +++ b/tests/correlators_test.py @@ -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) diff --git a/tests/fits_test.py b/tests/fits_test.py index 80b6de5a..dbc227dc 100644 --- a/tests/fits_test.py +++ b/tests/fits_test.py @@ -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():