From 6bc8102f87f9e18d3cd65d31b1c34427a5b1d2b1 Mon Sep 17 00:00:00 2001
From: Fabian Joswig <fabian.joswig@ed.ac.uk>
Date: Wed, 1 Dec 2021 09:22:16 +0000
Subject: [PATCH] refactor: jackknife helper functions in linalg module
 refactored

---
 pyerrors/linalg.py | 62 ++++++++++++++++++++++++----------------------
 1 file changed, 33 insertions(+), 29 deletions(-)

diff --git a/pyerrors/linalg.py b/pyerrors/linalg.py
index dbbde944..e6ee9cef 100644
--- a/pyerrors/linalg.py
+++ b/pyerrors/linalg.py
@@ -174,6 +174,35 @@ def matmul(*operands):
         return derived_array(multi_dot, operands)
 
 
+def _exp_to_jack(matrix):
+    base_matrix = np.empty_like(matrix)
+    for index, entry in np.ndenumerate(matrix):
+        base_matrix[index] = entry.export_jackknife()
+    return base_matrix
+
+
+def _imp_from_jack(matrix, name, idl):
+    base_matrix = np.empty_like(matrix)
+    for index, entry in np.ndenumerate(matrix):
+        base_matrix[index] = import_jackknife(entry, name, [idl])
+    return base_matrix
+
+
+def _exp_to_jack_c(matrix):
+    base_matrix = np.empty_like(matrix)
+    for index, entry in np.ndenumerate(matrix):
+        base_matrix[index] = entry.real.export_jackknife() + 1j * entry.imag.export_jackknife()
+    return base_matrix
+
+
+def _imp_from_jack_c(matrix, name, idl):
+    base_matrix = np.empty_like(matrix)
+    for index, entry in np.ndenumerate(matrix):
+        base_matrix[index] = CObs(import_jackknife(entry.real, name, [idl]),
+                                  import_jackknife(entry.imag, name, [idl]))
+    return base_matrix
+
+
 def jack_matmul(*operands):
     """Matrix multiply both operands making use of the jackknife approximation.
 
@@ -190,49 +219,24 @@ def jack_matmul(*operands):
         name = operands[0].flat[0].real.names[0]
         idl = operands[0].flat[0].real.idl[name]
 
-        def _exp_to_jack(matrix):
-            base_matrix = np.empty_like(matrix)
-            for index, entry in np.ndenumerate(matrix):
-                base_matrix[index] = entry.real.export_jackknife() + 1j * entry.imag.export_jackknife()
-            return base_matrix
-
-        def _imp_from_jack(matrix):
-            base_matrix = np.empty_like(matrix)
-            for index, entry in np.ndenumerate(matrix):
-                base_matrix[index] = CObs(import_jackknife(entry.real, name, [idl]),
-                                          import_jackknife(entry.imag, name, [idl]))
-            return base_matrix
-
-        r = _exp_to_jack(operands[0])
+        r = _exp_to_jack_c(operands[0])
         for op in operands[1:]:
             if isinstance(op.flat[0], CObs):
-                r = r @ _exp_to_jack(op)
+                r = r @ _exp_to_jack_c(op)
             else:
                 r = r @ op
-        return _imp_from_jack(r)
+        return _imp_from_jack_c(r, name, idl)
     else:
         name = operands[0].flat[0].names[0]
         idl = operands[0].flat[0].idl[name]
 
-        def _exp_to_jack(matrix):
-            base_matrix = np.empty_like(matrix)
-            for index, entry in np.ndenumerate(matrix):
-                base_matrix[index] = entry.export_jackknife()
-            return base_matrix
-
-        def _imp_from_jack(matrix):
-            base_matrix = np.empty_like(matrix)
-            for index, entry in np.ndenumerate(matrix):
-                base_matrix[index] = import_jackknife(entry, name, [idl])
-            return base_matrix
-
         r = _exp_to_jack(operands[0])
         for op in operands[1:]:
             if isinstance(op.flat[0], Obs):
                 r = r @ _exp_to_jack(op)
             else:
                 r = r @ op
-        return _imp_from_jack(r)
+        return _imp_from_jack(r, name, idl)
 
 
 def inv(x):