Now the non-recursive function do not produce warning

thewalrus/internal_modes/fock_density_matrices.py

@@ -206,28 +206,35 @@ def density_matrix_single_mode(
             else:
                 lower_limit = 0
             for j in range(lower_limit, i + 1):
-                if (i - j) % 2 == 0:
-                    patt_long = list((j,) + tuple(N_nums) + ((i - j) // 2,))
-                    new_blocks = np.concatenate((blocks, np.array([K + blocks[-1]])), axis=0)
-                    perm = (
-                        list(range(num_modes))
-                        + list(range(block_size))
-                        + list(range(num_modes, 2 * num_modes))
-                        + list(range(block_size))
-                    )
-                    Aperm = A[:, perm][perm]
-                    dm[j, i] = (
-                        pref
-                        * haf_blocked(Aperm, blocks=new_blocks, repeats=patt_long)
-                        / (
-                            np.prod(factorial(patt_long[1:-1]))
-                            * np.sqrt(factorial(i) * factorial(j))
+                if method == "non-recursive":
+                    if (i - j) % 2 == 0:
+                        patt_long = list((j,) + tuple(N_nums) + ((i - j) // 2,))
+                        new_blocks = np.concatenate((blocks, np.array([K + blocks[-1]])), axis=0)
+                        perm = (
+                            list(range(num_modes))
+                            + list(range(block_size))
+                            + list(range(num_modes, 2 * num_modes))
+                            + list(range(block_size))
+                        )
+                        Aperm = A[:, perm][perm]
+                        dm[j, i] = (
+                            pref
+                            * haf_blocked(Aperm, blocks=new_blocks, repeats=patt_long)
+                            / (
+                                np.prod(factorial(patt_long[1:-1]))
+                                * np.sqrt(factorial(i) * factorial(j))
+                            )
                         )
+                        dm[i, j] = np.conj(dm[j, i])
+                    else:
+                        dm[i, j] = 0
+                        dm[j, i] = 0
+                if method == "diagonals":
+                    patt_long = (i,) + tuple(N_nums)
+                    dm[i, i] = pref * np.real(
+                        haf_blocked(A, blocks=blocks, repeats=patt_long)
+                        / np.prod(factorial(patt_long))
                     )
-                    dm[i, j] = np.conj(dm[j, i])
-                else:
-                    dm[i, j] = 0
-                    dm[j, i] = 0
         if normalize:
             dm = dm / np.trace(dm)
         return dm

thewalrus/internal_modes/pnr_statistics.py

@@ -22,15 +22,13 @@
 
 from scipy.special import factorial as fac
 
-from ..quantum import Qmat, Amat
-from ..symplectic import passive_transformation
+from ..quantum import Qmat
 from .._hafnian import find_kept_edges, nb_binom, f_from_powertrace, nb_ix, f_from_matrix, get_AX_S
 from ..charpoly import powertrace
 
 from .utils import (
     spatial_reps_to_schmidt_reps,
     spatial_modes_to_schmidt_modes,
-    project_onto_local_oscillator,
 )
 
 
@@ -165,7 +163,7 @@ def haf_blocked(A, blocks, repeats):
     """
     # Note that the two lines below cannot be done inside numba hence the need for this function
     repeats = tuple(val + 1 for val in repeats)
-    blocks = [np.array(val, dtype=np.int32) for val in blocks]
+    blocks = numba.typed.List([np.array(val, dtype=np.int32) for val in blocks])
     return _haf_blocked_numba(A, blocks, repeats)
 
 
@@ -195,6 +193,5 @@ def _haf_blocked_numba(A, blocks, repeats_p):  # pragma: no cover
             for mode in block:
                 coeff_vect[mode] = val
         AX_S = get_AX_S(coeff_vect, A)
-
         netsum += coeff_pref * f_from_matrix(AX_S, 2 * n)[n]
     return netsum

thewalrus/tests/test_internal_modes.py

@@ -28,7 +28,7 @@
 
 from repoze.lru import lru_cache
 
-from thewalrus import low_rank_hafnian, reduction
+from thewalrus import low_rank_hafnian, reduction, hafnian
 
 from thewalrus.decompositions import takagi
 from thewalrus.random import random_covariance
@@ -50,7 +50,7 @@
     squeezing,
 )
 
-from thewalrus.internal_modes import pnr_prob, density_matrix_single_mode
+from thewalrus.internal_modes import pnr_prob, density_matrix_single_mode, haf_blocked
 from thewalrus.internal_modes.prepare_cov import (
     O_matrix,
     orthonormal_basis,
@@ -1532,3 +1532,20 @@ def test_unknown_method_in_density_matrix_single_mode():
     cutoff = 10
     with pytest.raises(ValueError, match="Unknown method for density_matrix_single_mode"):
         density_matrix_single_mode(cov, N, cutoff=cutoff, normalize=True, method="Coo coo ca choo")
+
+
+def test_haf_blocked():
+    """Tests that haf blocked is the sum of many hafnians"""
+    n = 6
+    B = np.random.rand(n, n) + 1j * np.random.rand(n, n)
+    A = B + B.T
+    reps_list = [[i, 3 - i, 4] for i in range(4)]
+    haf_sum = 0j
+    for reps in reps_list:
+        repreps = reps + reps
+        haf = hafnian(reduction(A, repreps))
+        haf_sum += haf / np.product(factorial(reps))
+    blocks = ((0, 1), (2,))
+    repeats = (3, 4)
+    haf_val = haf_blocked(A, blocks=blocks, repeats=repeats) / np.product(factorial(repeats))
+    assert np.allclose(haf_sum, haf_val)