From 207678513ce1f1595e91e1cbfc95cf4f33c1efac Mon Sep 17 00:00:00 2001
From: Fabian Joswig <fabian.joswig@ed.ac.uk>
Date: Tue, 19 Oct 2021 23:53:44 +0100
Subject: [PATCH] Npr_matrix class added

---
 pyerrors/input/hadrons.py |  4 +-
 pyerrors/npr.py           | 96 +++++++++++++++++++++++++++++++++++++--
 2 files changed, 94 insertions(+), 6 deletions(-)

diff --git a/pyerrors/input/hadrons.py b/pyerrors/input/hadrons.py
index ffbf4437..b903b3bb 100644
--- a/pyerrors/input/hadrons.py
+++ b/pyerrors/input/hadrons.py
@@ -6,7 +6,7 @@ import h5py
 import numpy as np
 from ..pyerrors import Obs, CObs
 from ..correlators import Corr
-from ..npr import Propagator
+from ..npr import Npr_matrix
 
 
 def _get_files(path, filestem):
@@ -103,4 +103,4 @@ def read_ExternalLeg_hd5(path, filestem, ens_id, order='C'):
         matrix[si, sj, ci, cj] = CObs(real, imag)
         matrix[si, sj, ci, cj].gamma_method()
 
-    return Propagator(matrix.reshape((12,12), order=order), mom)
+    return Npr_matrix(matrix.reshape((12,12), order=order), mom_in=mom)
diff --git a/pyerrors/npr.py b/pyerrors/npr.py
index 2eb8ce2a..71b222ad 100644
--- a/pyerrors/npr.py
+++ b/pyerrors/npr.py
@@ -1,13 +1,101 @@
 import numpy as np
 
 
-class Propagator(np.ndarray):
+_gamma = ['gammas', 0, 0, 0, 0, 0]
+_gamma[1] = np.array(
+    [[0, 0, 0, 1j], [0, 0, 1j, 0], [0, -1j, 0, 0], [-1j, 0, 0, 0]],
+    dtype=complex)
+_gamma[2] = np.array(
+    [[0, 0, 0, -1], [0, 0, 1, 0], [0, 1, 0, 0], [-1, 0, 0, 0]],
+    dtype=complex)
+_gamma[3] = np.array(
+    [[0, 0, 1j, 0], [0, 0, 0, -1j], [-1j, 0, 0, 0], [0, 1j, 0, 0]],
+    dtype=complex)
+_gamma[4] = np.array(
+    [[0, 0, 1, 0], [0, 0, 0, 1], [1, 0, 0, 0], [0, 1, 0, 0]],
+    dtype=complex)
+_gamma[5] = np.array(
+    [[1, 0, 0, 0], [0, 1, 0, 0], [0, 0, -1, 0], [0, 0, 0, -1]],
+    dtype=complex)
+_imat = np.array(
+    [[1, 0, 0, 0], [0, 1, 0, 0], [0, 0, 1, 0], [0, 0, 0, 1]],
+    dtype=complex)
 
-    def __new__(cls, input_array, mom=None):
+
+def gamma_matrix(gamma_tag):
+    """Returns gamma matrix in Grid labeling."""
+    if gamma_tag == 'Identity':
+        g = _imat
+    elif gamma_tag == 'Gamma5':
+        g = _gamma[5]
+    elif gamma_tag == 'GammaX':
+        g = _gamma[1]
+    elif gamma_tag == 'GammaY':
+        g = _gamma[2]
+    elif gamma_tag == 'GammaZ':
+        g = _gamma[3]
+    elif gamma_tag == 'GammaT':
+        g = _gamma[4]
+    elif gamma_tag == 'GammaXGamma5':
+        g = _gamma[1] @ _gamma[5]
+    elif gamma_tag == 'GammaYGamma5':
+        g = _gamma[2] @ _gamma[5]
+    elif gamma_tag == 'GammaZGamma5':
+        g = _gamma[3] @ _gamma[5]
+    elif gamma_tag == 'GammaTGamma5':
+        g = _gamma[4] @ _gamma[5]
+    elif gamma_tag == 'SigmaXT':
+        g = 0.5 * (_gamma[1] @ _gamma[4] - _gamma[4] @ _gamma[1])
+    elif gamma_tag == 'SigmaXY':
+        g = 0.5 * (_gamma[1] @ _gamma[2] - _gamma[2] @ _gamma[1])
+    elif gamma_tag == 'SigmaXZ':
+        g = 0.5 * (_gamma[1] @ _gamma[3] - _gamma[3] @ _gamma[1])
+    elif gamma_tag == 'SigmaYT':
+        g = 0.5 * (_gamma[2] @ _gamma[4] - _gamma[4] @ _gamma[2])
+    elif gamma_tag == 'SigmaYZ':
+        g = 0.5 * (_gamma[2] @ _gamma[3] - _gamma[3] @ _gamma[2])
+    elif gamma_tag == 'SigmaZT':
+        g = 0.5 * (_gamma[3] @ _gamma[4] - _gamma[4] @ _gamma[3])
+    else:
+        raise Exception('Unkown gamma structure', gamma_tag)
+    return g
+
+class Npr_matrix(np.ndarray):
+
+    g5 = np.array([[1, 0, 0, 0], [0, 1, 0, 0], [0, 0, -1, 0], [0, 0, 0, -1]],
+                  dtype=complex)
+
+    def __new__(cls, input_array, mom_in=None, mom_out=None):
         obj = np.asarray(input_array).view(cls)
-        obj.mom = mom
+        obj.mom_in = mom_in
+        obj.mom_out = mom_out
         return obj
 
+    @property
+    def g5H(self):
+        new_matrix = Npr_matrix.g5 @ self.conj().T @ Npr_matrix.g5
+        new_matrix.mom_in = self.mom_out
+        new_matrix.mom_out = self.mom_in
+        return new_matrix
+
+    def __matmul__(self, other):
+        if hasattr(other, 'mom_in'):
+            if self.mom_in != other.mom_in and self.mom_in and other.mom_in:
+                    raise Exception('mom_in does not match.')
+            mom_in = self.mom_in if self.mom_in else other.mom_in
+        else:
+            mom_in = self.mom_in
+
+        if hasattr(other, 'mom_out'):
+            if self.mom_out != other.mom_out and self.mom_out and other.mom_out:
+                    raise Exception('mom_out does not match.')
+            mom_out = self.mom_out if self.mom_out else other.mom_out
+        else:
+            mom_out = self.mom_out
+
+        return self.__new__(Npr_matrix, super().__matmul__(other), mom_in, mom_out)
+
     def __array_finalize__(self, obj):
         if obj is None: return
-        self.mom = getattr(obj, 'mom', None)
+        self.mom_in = getattr(obj, 'mom_in', None)
+        self.mom_out = getattr(obj, 'mom_out', None)