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>

Author: fjosw

pyerrors/correlators.py

@@ -220,7 +220,7 @@ def symmetric(self):
     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 @@ def anti_symmetric(self):
     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 @@ def is_matrix_symmetric(self):
                         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 @@ def roll(self, dt):
         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 @@ def __str__(self):
     # 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 @@ def __mul__(self, y):
         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):

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)

tests/fits_test.py

@@ -235,12 +235,12 @@ def func(a, x):
 
 
 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():