Adds extra test

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,16 +177,12 @@ 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())
+    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)
     if method in ["non-recursive", "diagonals"]:

thewalrus/tests/test_internal_modes.py

@@ -1132,7 +1132,8 @@ def test_LO_overlaps(r, S, phi):
 
 @pytest.mark.parametrize("nh", [1, 2, 3, 4])
 @pytest.mark.parametrize("method", ["recursive", "non-recursive"])
-def test_mixed_heralded_photon(nh, method):
+@pytest.mark.parametrize("herald", [0, 1])
+def test_mixed_heralded_photon(nh, method, herald):
     """test code for generating heralded fock states from squeezed states with 2 internal modes"""
     na = 1
     nb = 0.5
@@ -1153,20 +1154,20 @@ def test_mixed_heralded_photon(nh, method):
     LO_overlapa = LO_overlaps(chis, chis[0])
     LO_overlapb = LO_overlaps(chis, chis[1])
     rho_a = density_matrix_single_mode(
-        cov, {1: nh}, normalize=True, LO_overlap=LO_overlapa, cutoff=nh + 1, method=method
+        cov, {herald: nh}, normalize=True, LO_overlap=LO_overlapa, cutoff=nh + 1, method=method
     )
     rho_b = density_matrix_single_mode(
-        cov, {1: nh}, normalize=True, LO_overlap=LO_overlapb, cutoff=nh + 1, method=method
+        cov, {herald: nh}, normalize=True, LO_overlap=LO_overlapb, cutoff=nh + 1, method=method
     )
 
     p_a = np.diag(
         density_matrix_single_mode(
-            cov, {1: nh}, normalize=True, LO_overlap=LO_overlapa, cutoff=nh + 1, method="diagonals"
+            cov, {herald: nh}, normalize=True, LO_overlap=LO_overlapa, cutoff=nh + 1, method="diagonals"
         )
     )
     p_b = np.diag(
         density_matrix_single_mode(
-            cov, {1: nh}, normalize=True, LO_overlap=LO_overlapb, cutoff=nh + 1, method="diagonals"
+            cov, {herald: nh}, normalize=True, LO_overlap=LO_overlapb, cutoff=nh + 1, method="diagonals"
         )
     )