From 94e3ebe7f03d949e3162ed32084b957671c75911 Mon Sep 17 00:00:00 2001
From: Fabian Joswig <fabian.joswig@ed.ac.uk>
Date: Fri, 22 Oct 2021 18:29:28 +0100
Subject: [PATCH] read bilinear added to input/hadrons

---
 pyerrors/input/hadrons.py | 57 +++++++++++++++++++++++++++++++++++++++
 1 file changed, 57 insertions(+)

diff --git a/pyerrors/input/hadrons.py b/pyerrors/input/hadrons.py
index 509c0922..2e9b52d4 100644
--- a/pyerrors/input/hadrons.py
+++ b/pyerrors/input/hadrons.py
@@ -109,3 +109,60 @@ def read_ExternalLeg_hd5(path, filestem, ens_id, order='F'):
         matrix[si, sj, ci, cj].gamma_method()
 
     return Npr_matrix(matrix.swapaxes(1, 2).reshape((12, 12), order=order), mom_in=mom)
+
+
+def read_Bilinear_hd5(path, filestem, ens_id, order='F'):
+    """Read hadrons Bilinear hdf5 file and output an array of CObs
+
+    Parameters
+    -----------------
+    path -- path to the files to read
+    filestem -- namestem of the files to read
+    ens_id -- name of the ensemble, required for internal bookkeeping
+    order -- order in which the array is to be reshaped,
+             'F' for the first index changing fastest (9 4x4 matrices) default.
+             'C' for the last index changing fastest (16 3x3 matrices),
+    """
+
+    files = _get_files(path, filestem)
+
+    mom_in = None
+    mom_out = None
+
+    corr_data = {}
+    for hd5_file in files:
+        file = h5py.File(path + '/' + hd5_file, "r")
+        for i in range(16):
+            name = file['Bilinear/Bilinear_' + str(i) + '/info'].attrs['gamma'][0].decode('UTF-8')
+            if name not in corr_data:
+                corr_data[name] = []
+            raw_data = file['Bilinear/Bilinear_' + str(i) + '/corr'][0][0].view('complex')
+            corr_data[name].append(raw_data)
+            if mom_in is not None:
+                assert np.allclose(mom_in, np.array(str(file['Bilinear/Bilinear_' + str(i) + '/info'].attrs['pIn'])[3:-2].strip().split(' '), dtype=int))
+            else:
+                mom_in = np.array(str(file['Bilinear/Bilinear_' + str(i) + '/info'].attrs['pIn'])[3:-2].strip().split(' '), dtype=int)
+            if mom_out is not None:
+                assert np.allclose(mom_out, np.array(str(file['Bilinear/Bilinear_' + str(i) + '/info'].attrs['pIn'])[3:-2].strip().split(' '), dtype=int))
+            else:
+                mom_out = np.array(str(file['Bilinear/Bilinear_' + str(i) + '/info'].attrs['pIn'])[3:-2].strip().split(' '), dtype=int)
+
+        file.close()
+
+    result_dict = {}
+
+    for key, data in corr_data.items():
+        local_data = np.array(data)
+
+        rolled_array = np.rollaxis(local_data, 0, 5)
+
+        matrix = np.empty((rolled_array.shape[:-1]), dtype=object)
+        for si, sj, ci, cj in np.ndindex(rolled_array.shape[:-1]):
+            real = Obs([rolled_array[si, sj, ci, cj].real], [ens_id])
+            imag = Obs([rolled_array[si, sj, ci, cj].imag], [ens_id])
+            matrix[si, sj, ci, cj] = CObs(real, imag)
+            matrix[si, sj, ci, cj].gamma_method()
+
+        result_dict[key] = Npr_matrix(matrix.swapaxes(1, 2).reshape((12, 12), order=order), mom_in=mom_in, mom_out=mom_out)
+
+    return result_dict