Now the non-recursive function do not produce warning (#386)

* Now the non-recursive function do not produce warning

* Now the non-recursive function do not produce warning

* Passes black

* Simplifies the diagonal method

* More tests, always

* More tests, always

* minor pylint improvements

* Adds extra test

* black

* minor simplification

* some tests pass

* some tests pass

* Moves the tests that take forever to the end of the test file

* Passes black

* Adds else

---------

Co-authored-by: Nicolas Quesada <nquesada@pop-os.localdomain>

thewalrus/internal_modes/fock_density_matrices.py

@@ -19,7 +19,6 @@
 import numba
 from scipy.special import factorial
 
-from ..symplectic import passive_transformation
 from .._hafnian import nb_binom, nb_ix, find_kept_edges, f_from_matrix
 from .utils import (
     nb_Qmat,
@@ -178,19 +177,16 @@ def density_matrix_single_mode(
             raise ValueError("Norm of overlaps must not be greater than 1")
 
     # swapping the spatial modes around such that we are heralding in spatial mode 0
-    Uswap = np.zeros((M, M))
-    swapV = np.concatenate((np.array([HM]), np.arange(HM), np.arange(HM + 1, M)))
-    for j, k in enumerate(swapV):
-        Uswap[j][k] = 1
-    U_K = np.zeros((M * K, M * K))
-    for i in range(K):
-        U_K[i::K, i::K] = Uswap
-    _, cov = passive_transformation(np.zeros(cov.shape[0]), cov, U_K, hbar=hbar)
-
-    # N_nums = list(pattern.values())
+    if HM != 0:
+        swapV = list(range(M))
+        (swapV[0], swapV[HM]) = (swapV[HM], swapV[0])
+        perm = (np.arange(M * K).reshape(M, K))[swapV].flatten()
+        double_perm = np.concatenate([perm, perm + M * K])
+        cov = cov[:, double_perm][double_perm]
+
     if method == "recursive":
         return _density_matrix_single_mode(cov, N_nums, normalize, LO_overlap, cutoff, hbar)
-    elif method == "non-recursive" or method == "diagonals":
+    if method in ["non-recursive", "diagonals"]:
         cov = project_onto_local_oscillator(cov, M, LO_overlap=LO_overlap, hbar=hbar)
         num_modes = len(cov) // 2
         A = Amat(cov)
@@ -200,36 +196,39 @@ def density_matrix_single_mode(
         dm = np.zeros([cutoff, cutoff], dtype=np.complex128)
         num_modes = M * K
         block_size = K
-        for i in range(cutoff):
-            if method == "diagonals":
-                lower_limit = i
-            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":
+            for i in range(cutoff):
+                for j in range(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))
                         )
-                    )
-                    dm[i, j] = np.conj(dm[j, i])
-                else:
-                    dm[i, j] = 0
-                    dm[j, i] = 0
+                        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
+        else:
+            for i in range(cutoff):
+                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))
+                )
         if normalize:
             dm = dm / np.trace(dm)
         return dm
-    else:
-        raise ValueError("Unknown method for density_matrix_single_mode")
+
+    raise ValueError("Unknown method for density_matrix_single_mode")

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