Add function for generalized outer product

main.py

@@ -6,8 +6,49 @@ import sys
 import h5py
 import configparser
 from distutils.util import strtobool
+from functools import partial
+
+
+def calculate_outer_product_along_axis(a, b, axis):
+    """Return the outer product of arrays ``a`` and ``b`` along the
+    given ``axis``.
+
+    Calculate the result according to the following formular:
+
+        .. math::
+            c_{(i_0)...(i_{j-1})kl(i_{j+1})...(i_{N-1})}=
+            a_{(i_0)...(i_{j-1})k(i_{j+1})...(i_{N-1})}
+            b_{(i_0)...(i_{j-1})l(i_{j+1})...(i_{N-1})},
+
+    where indices in braces are not implicitly summed over.
+
+    :var a: N-dimensional numpy array :math:`a`
+    :var b: N-dimensional numpy array :math:`b`
+    :var axis: axis :math:`j\\in [0,N)` along which the outer product is
+        calculated
+    :return: N+1-dimensional numpy array :math:`c`
+    """
+
+    final_shape = list(a.shape)
+    final_shape.insert(axis, a.shape[axis])
+    if axis == 0:
+        a = a[np.newaxis, ...]
+        b = b[np.newaxis, ...]
+        axis = 1
+    if axis == a.ndim - 1:
+        a = a[..., np.newaxis]
+        b = b[..., np.newaxis]
+    new_shape = (np.product(a.shape[:axis]), a.shape[axis],
+                 np.product(a.shape[axis + 1:]))
+    a = a.reshape(new_shape)
+    b = b.reshape(new_shape)
+    result = np.empty((new_shape[0], new_shape[1], new_shape[1], new_shape[2]))
+    for i in range(new_shape[0]):
+        for j in range(new_shape[2]):
+            result[i, :, :, j] = np.outer(a[i, :, j], np.conjugate(b[i, :, j]))
+    return result.reshape(final_shape)
+
 
-# %%
 config_file_name = "jackknife.ini"
 n_processes = 1
 if len(sys.argv) == 1:
@@ -60,8 +101,8 @@ else:
 
 # %%
 # Do the Jackknife estimation
-jackknife = jk.Jackknife(x_generator, n, f,
-                         adga.outer_product_of_self_energies,
+outer_product = partial(calculate_outer_product_along_axis, axis=2)
+jackknife = jk.Jackknife(x_generator, n, f, outer_product,
                          general["output_file_prefix"],
                          strtobool(general["store_output_samples"]))
 jackknife.do_estimation()