From 28bf0f1701a1740a4e1a9ccc99c94d17c71f31f6 Mon Sep 17 00:00:00 2001
From: Fabian Joswig <fabian.joswig@ed.ac.uk>
Date: Thu, 18 Nov 2021 11:17:20 +0000
Subject: [PATCH] feat: linalg.jack_matmul now also works with irregular monte
 carlo chains

---
 pyerrors/linalg.py | 20 +++++++++++++-------
 pyerrors/obs.py    |  4 ++--
 2 files changed, 15 insertions(+), 9 deletions(-)

diff --git a/pyerrors/linalg.py b/pyerrors/linalg.py
index e9f5dbd5..075b9605 100644
--- a/pyerrors/linalg.py
+++ b/pyerrors/linalg.py
@@ -187,40 +187,46 @@ def jack_matmul(*operands):
     """
 
     if any(isinstance(o[0, 0], CObs) for o in operands):
+        name = operands[0][0, 0].real.names[0]
+        idl = operands[0][0, 0].real.idl[name]
+
         def _exp_to_jack(matrix):
             base_matrix = np.empty_like(matrix)
             for (n, m), entry in np.ndenumerate(matrix):
                 base_matrix[n, m] = entry.real.export_jackknife() + 1j * entry.imag.export_jackknife()
             return base_matrix
 
-        def _imp_from_jack(matrix, name):
+        def _imp_from_jack(matrix):
             base_matrix = np.empty_like(matrix)
             for (n, m), entry in np.ndenumerate(matrix):
-                base_matrix[n, m] = CObs(import_jackknife(entry.real, name),
-                                         import_jackknife(entry.imag, name))
+                base_matrix[n, m] = CObs(import_jackknife(entry.real, name, [idl]),
+                                         import_jackknife(entry.imag, name, [idl]))
             return base_matrix
 
         r = _exp_to_jack(operands[0])
         for op in operands[1:]:
             r = r @ _exp_to_jack(op)
-        return _imp_from_jack(r, op.ravel()[0].real.names[0])
+        return _imp_from_jack(r)
     else:
+        name = operands[0][0, 0].names[0]
+        idl = operands[0][0, 0].idl[name]
+
         def _exp_to_jack(matrix):
             base_matrix = np.empty_like(matrix)
             for (n, m), entry in np.ndenumerate(matrix):
                 base_matrix[n, m] = entry.export_jackknife()
             return base_matrix
 
-        def _imp_from_jack(matrix, name):
+        def _imp_from_jack(matrix):
             base_matrix = np.empty_like(matrix)
             for (n, m), entry in np.ndenumerate(matrix):
-                base_matrix[n, m] = import_jackknife(entry, name)
+                base_matrix[n, m] = import_jackknife(entry, name, [idl])
             return base_matrix
 
         r = _exp_to_jack(operands[0])
         for op in operands[1:]:
             r = r @ _exp_to_jack(op)
-        return _imp_from_jack(r, op.ravel()[0].names[0])
+        return _imp_from_jack(r)
 
 
 def inv(x):
diff --git a/pyerrors/obs.py b/pyerrors/obs.py
index 1e594ea4..56a67e57 100644
--- a/pyerrors/obs.py
+++ b/pyerrors/obs.py
@@ -1559,7 +1559,7 @@ def load_object(path):
         return pickle.load(file)
 
 
-def import_jackknife(jacks, name):
+def import_jackknife(jacks, name, idl=None):
     """Imports jackknife samples and returns an Obs
 
     Parameters
@@ -1573,7 +1573,7 @@ def import_jackknife(jacks, name):
     length = len(jacks) - 1
     prj = (np.ones((length, length)) - (length - 1) * np.identity(length))
     samples = jacks[1:] @ prj
-    new_obs = Obs([samples], [name])
+    new_obs = Obs([samples], [name], idl=idl)
     new_obs._value = jacks[0]
     return new_obs