Hafnian modifications

thewalrus/_hafnian.py

@@ -20,6 +20,7 @@
 from itertools import chain, combinations
 import numba
 import numpy as np
+from thewalrus import charpoly 
 
 
 @numba.jit(nopython=True, cache=True)
@@ -180,25 +181,24 @@ def find_kept_edges(j, reps):  # pragma: no cover
 
 
 @numba.jit(nopython=True, cache=True)
-def f(E, n):  # pragma: no cover
+def f(A, n):  # pragma: no cover
     """Evaluate the polynomial coefficients of the function in the eigenvalue-trace formula.
 
     Args:
-        E (array): eigenvalues of ``AX``
+        A (array): a two-dimensional matrix
         n (int): number of polynomial coefficients to compute
 
     Returns:
         array: polynomial coefficients
     """
-    E_k = E.copy()
     # Compute combinations in O(n^2log n) time
     # code translated from thewalrus matlab script
     count = 0
     comb = np.zeros((2, n // 2 + 1), dtype=np.complex128)
     comb[0, 0] = 1
+    powtrace = charpoly.powertrace(A, n // 2 + 1)
     for i in range(1, n // 2 + 1):
-        factor = E_k.sum() / (2 * i)
-        E_k *= E
+        factor = powtrace[i] / (2 * i)
         powfactor = 1
         count = 1 - count
         comb[count, :] = comb[1 - count, :]
@@ -210,11 +210,11 @@ def f(E, n):  # pragma: no cover
 
 
 @numba.jit(nopython=True, cache=True)
-def f_loop(E, AX_S, XD_S, D_S, n):  # pragma: no cover
+def f_loop(AX, AX_S, XD_S, D_S, n):  # pragma: no cover
     """Evaluate the polynomial coefficients of the function in the eigenvalue-trace formula.
 
     Args:
-        E (array): eigenvalues of ``AX``
+        AX (array): two-dimensional matrix
         AX_S (array): ``AX_S`` with weights given by repetitions and excluded rows removed
         XD_S (array): diagonal multiplied by ``X``
         D_S (array): diagonal
@@ -223,15 +223,14 @@ def f_loop(E, AX_S, XD_S, D_S, n):  # pragma: no cover
     Returns:
         array: polynomial coefficients
     """
-    E_k = E.copy()
     # Compute combinations in O(n^2log n) time
     # code translated from thewalrus matlab script
     count = 0
     comb = np.zeros((2, n // 2 + 1), dtype=np.complex128)
     comb[0, 0] = 1
+    powtrace = charpoly.powertrace(AX, n // 2 + 1)
     for i in range(1, n // 2 + 1):
-        factor = E_k.sum() / (2 * i) + (XD_S @ D_S) / 2
-        E_k *= E
+        factor = powtrace[i] / (2 * i) + (XD_S @ D_S) / 2
         XD_S = XD_S @ AX_S
         powfactor = 1
         count = 1 - count
@@ -245,12 +244,12 @@ def f_loop(E, AX_S, XD_S, D_S, n):  # pragma: no cover
 
 # pylint: disable = too-many-arguments
 @numba.jit(nopython=True, cache=True)
-def f_loop_odd(E, AX_S, XD_S, D_S, n, oddloop, oddVX_S):  # pragma: no cover
+def f_loop_odd(AX, AX_S, XD_S, D_S, n, oddloop, oddVX_S):  # pragma: no cover
     """Evaluate the polynomial coefficients of the function in the eigenvalue-trace formula
     when there is a self-edge in the fixed perfect matching.
 
     Args:
-        E (array): eigenvalues of ``AX``
+        AX (array): two-dimensional matrix
         AX_S (array): ``AX_S`` with weights given by repetitions and excluded rows removed
         XD_S (array): diagonal multiplied by ``X``
         D_S (array): diagonal
@@ -261,17 +260,16 @@ def f_loop_odd(E, AX_S, XD_S, D_S, n, oddloop, oddVX_S):  # pragma: no cover
     Returns:
         array: polynomial coefficients
     """
-    E_k = E.copy()
 
     count = 0
     comb = np.zeros((2, n + 1), dtype=np.complex128)
     comb[0, 0] = 1
+    powtrace = charpoly.powertrace(AX, n + 1)
     for i in range(1, n + 1):
         if i == 1:
             factor = oddloop
         elif i % 2 == 0:
-            factor = E_k.sum() / i + (XD_S @ D_S) / 2
-            E_k *= E
+            factor = powtrace[i // 2] / i + (XD_S @ D_S) / 2
         else:
             factor = oddVX_S @ D_S
             D_S = AX_S @ D_S
@@ -456,13 +454,11 @@ def _calc_hafnian(A, edge_reps, glynn=True):  # pragma: no cover
 
         AX_S = get_AX_S(kept_edges, A)
 
-        E = eigvals(AX_S)  # O(n^3) step
-
         prefac = (-1.0) ** (N // 2 - edge_sum) * binom_prod
 
         if glynn and kept_edges[0] == 0:
             prefac *= 0.5
-        Hnew = prefac * f(E, N)[N // 2]
+        Hnew = prefac * f(AX_S, N)[N // 2]
 
         H += Hnew
 
@@ -561,17 +557,16 @@ def _calc_loop_hafnian(A, D, edge_reps, oddloop=None, oddV=None, glynn=True):  #
             kept_edges = 2 * kept_edges - edge_reps
 
         AX_S, XD_S, D_S, oddVX_S = get_submatrices(kept_edges, A, D, oddV)
-
-        E = eigvals(AX_S)  # O(n^3) step
+        AX = AX_S.copy()
 
         prefac = (-1.0) ** (N // 2 - edge_sum) * binom_prod
 
         if oddloop is not None:
-            Hnew = prefac * f_loop_odd(E, AX_S, XD_S, D_S, N, oddloop, oddVX_S)[N]
+            Hnew = prefac * f_loop_odd(AX, AX_S, XD_S, D_S, N, oddloop, oddVX_S)[N]
         else:
             if glynn and kept_edges[0] == 0:
                 prefac *= 0.5
-            Hnew = prefac * f_loop(E, AX_S, XD_S, D_S, N)[N // 2]
+            Hnew = prefac * f_loop(AX, AX_S, XD_S, D_S, N)[N // 2]
 
         H += Hnew
 
@@ -894,7 +889,6 @@ def hafnian_repeated(A, rpt, mu=None, loop=False, rtol=1e-05, atol=1e-08, glynn=
 
         >>> hafnian_repeated(A, rpt) == hafnian(A)
 
-
     Args:
         A (array): a square, symmetric :math:`N\times N` array
         rpt (Sequence): a length-:math:`N` positive integer sequence, corresponding