Sparse loop hafnian

.github/CHANGELOG.md

@@ -2,9 +2,11 @@
 
 ### New features
 
+* Adds the function `hafnian_sparse` to compute sparse loop hafnians (pure python implementation). [#245](https://github.com/XanaduAI/thewalrus/pull/245)
+
 ### Improvements
 
-* Speeds up the calculation of photon number variances/covariances [#224](https://github.com/XanaduAI/thewalrus/pull/224)
+* Speeds up the calculation of photon number variances/covariances [#244](https://github.com/XanaduAI/thewalrus/pull/244)
 
 ### Bug fixes
 
@@ -14,7 +16,7 @@
 
 This release contains contributions from (in alphabetical order):
 
-Nicolas Quesada
+Filippo Miatto, Nicolas Quesada
 
 ---
 

thewalrus/__init__.py

@@ -100,6 +100,7 @@
     hafnian,
     hafnian_repeated,
     reduction,
+    hafnian_sparse,
 )
 from ._low_rank_haf import low_rank_hafnian
 from ._hermite_multidimensional import hafnian_batched, hermite_multidimensional
@@ -112,6 +113,7 @@
     "hafnian",
     "hafnian_repeated",
     "hafnian_batched",
+    "hafnian_sparse",
     "tor",
     "perm",
     "permanent_repeated",

thewalrus/_hafnian.py

@@ -14,6 +14,8 @@
 """
 Hafnian Python interface
 """
+from functools import lru_cache
+from collections import Counter
 import numpy as np
 
 from .libwalrus import haf_complex, haf_int, haf_real, haf_rpt_complex, haf_rpt_real
@@ -178,6 +180,46 @@ def hafnian(
     )
 
 
+def hafnian_sparse(A, D = None):
+    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.
+
+    Args:
+        A (array): the symmetric matrix of which we want to compute the hafnian
+        D (set): set of indices that identify a submatrix. If ``None`` (default) it computes
+            the hafnian of the whole matrix.
+
+    Returns:
+        (float) hafnian of ``A`` or of the submatrix of ``A`` defined by the set of indices ``D``.
+    """
+    if D is None:
+        D = frozenset(range(len(A)))
+    else:
+        D = frozenset(D)
+
+    if np.allclose(A, np.zeros_like(A)):
+        return 0.0
+
+    r, _ = np.nonzero(A)
+    m = max(Counter(r).values())  # max nonzero values per row/column
+
+    @lru_cache(maxsize=2 ** m)
+    def indices(d, k):
+        return d.intersection(set(np.nonzero(A[k])[0]))
+
+    @lru_cache(maxsize=2 ** m)
+    def lhaf(d: frozenset) -> float:
+        if not d:
+            return 1
+        d_ = set(d)
+        k = d_.pop()
+        return sum(A[i, k] * lhaf(frozenset(d_).difference({i})) for i in indices(d, k))
+
+    return lhaf(D)
+
+
 def hafnian_repeated(A, rpt, mu=None, loop=False, rtol=1e-05, atol=1e-08):
     r"""Returns the hafnian of matrix with repeated rows/columns.
 

thewalrus/tests/conftest.py

@@ -41,10 +41,12 @@ def random_matrix(dtype):
     """Returns a random symmetric matrix of type dtype
     and of size n x n"""
 
-    def _wrapper(n):
+    def _wrapper(n, fill_factor=1.0):
         """wrapper function"""
         A = np.complex128(np.random.random([n, n]))
         A += 1j * np.random.random([n, n])
+        mask = np.random.binomial(1, p=fill_factor, size=(n, n))
+        A *= mask
         A += A.T
 
         if not np.issubdtype(dtype, np.complexfloating):

thewalrus/tests/test_hafnian.py

@@ -19,7 +19,7 @@
 from scipy.special import factorial as fac
 
 import thewalrus as hf
-from thewalrus import hafnian, reduction
+from thewalrus import hafnian, reduction, hafnian_sparse
 from thewalrus.libwalrus import haf_complex, haf_real, haf_int
 
 
@@ -277,3 +277,12 @@ def test_diag(self, n):
         haf = hafnian(A, loop=True)
         expected = np.prod(v)
         assert np.allclose(haf, expected)
+
+
+    @pytest.mark.parametrize("n", [2, 3, 5, 10, 15, 20])
+    @pytest.mark.parametrize("fill", [0.5, 0.2, 0.1, 0.05])
+    def test_valid_output(self, random_matrix, n, fill):
+        """Tests that sparse loop hafnian matches full implementation"""
+        A = random_matrix(n, fill_factor=fill)
+        assert np.allclose(hafnian_sparse(A), hafnian_sparse(A, D=set(range(len(A)))))
+        assert np.allclose(hafnian(A, loop=True), hafnian_sparse(A))