Hafnian modifications

.gitignore

@@ -26,3 +26,5 @@ docs/code/api/*
 .idea
 .mypy_cache/
 .DS_Store
+
+venv

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,25 @@ 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 (2d array): 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
+        # Maybe try powtrace[i-1] instead
+        factor = powtrace[i] / (2 * i)
         powfactor = 1
         count = 1 - count
         comb[count, :] = comb[1 - count, :]
@@ -210,11 +211,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 (2d array): 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 +224,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 +245,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 (2d array): 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 +261,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
@@ -366,6 +365,7 @@ def get_submatrix_batch_odd0(kept_edges, oddV0):  # pragma: no cover
     Args:
         kept_edges (array): number of repetitions of each edge
         oddV0 (array): Row of matrix at index of self-edge. ``None`` is no self-edge.
+
     Returns:
         array: scaled ``oddV0 @ X``
     """
@@ -456,13 +456,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
 
@@ -474,6 +472,7 @@ def _calc_hafnian(A, edge_reps, glynn=True):  # pragma: no cover
 
 def _haf(A, reps=None, glynn=True):
     r"""Calculate hafnian with (optional) repeated rows and columns.
+
     Code contributed by `Jake F.F. Bulmer <https://github.com/jakeffbulmer/gbs>`_ based on
     `arXiv:2108.01622 <https://arxiv.org/abs/2108.01622>`_.
 
@@ -518,6 +517,7 @@ def _haf(A, reps=None, glynn=True):
 def _calc_loop_hafnian(A, D, edge_reps, oddloop=None, oddV=None, glynn=True):  # pragma: no cover
     """Compute loop hafnian, using inputs as prepared by frontend loop_hafnian function
     compiled with Numba.
+
     Code contributed by `Jake F.F. Bulmer <https://github.com/jakeffbulmer/gbs>`_ based on
     `arXiv:2108.01622 <https://arxiv.org/abs/2108.01622>`_.
 
@@ -561,17 +561,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
 
@@ -587,6 +586,7 @@ def _calc_loop_hafnian(A, D, edge_reps, oddloop=None, oddV=None, glynn=True):  #
 # pylint: disable=redefined-outer-name
 def loop_hafnian(A, D=None, reps=None, glynn=True):
     """Calculate loop hafnian with (optional) repeated rows and columns.
+
     Code contributed by `Jake F.F. Bulmer <https://github.com/jakeffbulmer/gbs>`_ based on
     `arXiv:2108.01622 <https://arxiv.org/abs/2108.01622>`_.
 
@@ -640,7 +640,9 @@ def loop_hafnian(A, D=None, reps=None, glynn=True):
 
 def input_validation(A, rtol=1e-05, atol=1e-08):
     """Checks that the matrix A satisfies the requirements for Hafnian calculation.
+
     These include:
+
     * That the ``A`` is a NumPy array
     * That ``A`` is square
     * That ``A`` does not contain any NaNs
@@ -703,6 +705,7 @@ def powerset(iterable):
 
 def reduction(A, rpt):
     r"""Calculates the reduction of an array by a vector of indices.
+
     This is equivalent to repeating the ith row/column of :math:`A`, :math:`rpt_i` times.
 
     Args:
@@ -732,6 +735,7 @@ def hafnian(
     method="glynn",
 ):  # pylint: disable=too-many-arguments
     """Returns the hafnian of a matrix.
+
     Code contributed by `Jake F.F. Bulmer <https://github.com/jakeffbulmer/gbs>`_ based on
     `arXiv:2108.01622 <https://arxiv.org/abs/2108.01622>`_.
 
@@ -827,6 +831,7 @@ def hafnian(
 
 def hafnian_sparse(A, D=None, loop=False):
     r"""Returns the hafnian of a sparse symmetric matrix.
+
     This pure python implementation is very slow on full matrices, but faster the sparser a matrix is.
     As a rule of thumb, the crossover in runtime with respect to :func:`~.hafnian` happens around 50% sparsity.
 
@@ -894,7 +899,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
@@ -946,6 +950,7 @@ def hafnian_repeated(A, rpt, mu=None, loop=False, rtol=1e-05, atol=1e-08, glynn=
 
 def hafnian_banded(A, loop=False, rtol=1e-05, atol=1e-08):
     """Returns the loop hafnian of a banded matrix.
+
     For the derivation see Section V of `'Efficient sampling from shallow Gaussian quantum-optical
     circuits with local interactions', Qi et al. <https://arxiv.org/abs/2009.11824>`_.
 
@@ -989,7 +994,9 @@ def hafnian_banded(A, loop=False, rtol=1e-05, atol=1e-08):
 def recursive_hafnian(A):  # pragma: no cover
     r"""Computes the hafnian of the matrix with the recursive algorithm. It is an implementation of
     algorithm 2 in *Counting perfect matchings as fast as Ryser* :cite:`bjorklund2012counting`.
+
     This code is a modified version of the code found here:
+
     `Recursive hafnian
     <https://codegolf.stackexchange.com/questions/157049/calculate-the-hafnian-as-quickly-as-possible>`_.
 
@@ -1082,6 +1089,7 @@ def hafnian_approx(A, num_samples=1000):  # pragma: no cover
 
     The approximation follows the stochastic Barvinok's approximation allowing the
     hafnian can be approximated as the sum of determinants of matrices.
+
     The accuracy of the approximation increases with increasing number of iterations.
 
     Args: