diff --git a/thewalrus/_hafnian.py b/thewalrus/_hafnian.py index 39b4f4b2b..ad62d6da2 100644 --- a/thewalrus/_hafnian.py +++ b/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