From ee6bc0c62a1339823b1a12d41db7eec70757d926 Mon Sep 17 00:00:00 2001
From: Nicolas Quesada <nquesada@pop-os.localdomain>
Date: Tue, 23 Apr 2024 15:59:46 -0400
Subject: [PATCH 01/15] Now the non-recursive function do not produce warning

---
 thewalrus/internal_modes/pnr_statistics.py |  7 ++-----
 thewalrus/tests/test_internal_modes.py     | 20 ++++++++++++++++++--
 2 files changed, 20 insertions(+), 7 deletions(-)

diff --git a/thewalrus/internal_modes/pnr_statistics.py b/thewalrus/internal_modes/pnr_statistics.py
index e7732c66..518ae9a7 100644
--- a/thewalrus/internal_modes/pnr_statistics.py
+++ b/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
diff --git a/thewalrus/tests/test_internal_modes.py b/thewalrus/tests/test_internal_modes.py
index c8d57f17..e3e57a10 100644
--- a/thewalrus/tests/test_internal_modes.py
+++ b/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,19 @@ 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)
\ No newline at end of file

From a03c145c6fb4889afde5babc4c754f23a5fb6e9f Mon Sep 17 00:00:00 2001
From: Nicolas Quesada <nquesada@pop-os.localdomain>
Date: Tue, 23 Apr 2024 16:00:40 -0400
Subject: [PATCH 02/15] Now the non-recursive function do not produce warning

---
 thewalrus/tests/test_internal_modes.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/thewalrus/tests/test_internal_modes.py b/thewalrus/tests/test_internal_modes.py
index e3e57a10..d9e96187 100644
--- a/thewalrus/tests/test_internal_modes.py
+++ b/thewalrus/tests/test_internal_modes.py
@@ -1547,4 +1547,4 @@ def test_haf_blocked():
     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)
\ No newline at end of file
+    assert np.allclose(haf_sum, haf_val)

From f4aa6b4d6d763ce4a2c766c51c6ea3b45a2e97f2 Mon Sep 17 00:00:00 2001
From: Nicolas Quesada <nquesada@pop-os.localdomain>
Date: Tue, 23 Apr 2024 16:09:58 -0400
Subject: [PATCH 03/15] Passes black

---
 thewalrus/tests/test_internal_modes.py | 11 ++++++-----
 1 file changed, 6 insertions(+), 5 deletions(-)

diff --git a/thewalrus/tests/test_internal_modes.py b/thewalrus/tests/test_internal_modes.py
index d9e96187..7653c57f 100644
--- a/thewalrus/tests/test_internal_modes.py
+++ b/thewalrus/tests/test_internal_modes.py
@@ -1533,18 +1533,19 @@ def test_unknown_method_in_density_matrix_single_mode():
     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)
+    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)]
+    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))
+    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)

From 2598599f3d210a41546695e42d679433a60dc627 Mon Sep 17 00:00:00 2001
From: Nicolas Quesada <nquesada@pop-os.localdomain>
Date: Tue, 23 Apr 2024 21:50:07 -0400
Subject: [PATCH 04/15] Simplifies the diagonal method

---
 .../internal_modes/fock_density_matrices.py   | 47 +++++++++++--------
 1 file changed, 27 insertions(+), 20 deletions(-)

diff --git a/thewalrus/internal_modes/fock_density_matrices.py b/thewalrus/internal_modes/fock_density_matrices.py
index 3ab14dc1..945d187e 100644
--- a/thewalrus/internal_modes/fock_density_matrices.py
+++ b/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

From 665125887e91b0667476fdc9d297635f998a1380 Mon Sep 17 00:00:00 2001
From: Nicolas Quesada <nquesada@pop-os.localdomain>
Date: Tue, 23 Apr 2024 21:56:52 -0400
Subject: [PATCH 05/15] More tests, always

---
 thewalrus/tests/test_internal_modes.py | 10 +++++++---
 1 file changed, 7 insertions(+), 3 deletions(-)

diff --git a/thewalrus/tests/test_internal_modes.py b/thewalrus/tests/test_internal_modes.py
index 7653c57f..f9513589 100644
--- a/thewalrus/tests/test_internal_modes.py
+++ b/thewalrus/tests/test_internal_modes.py
@@ -1534,18 +1534,22 @@ def test_unknown_method_in_density_matrix_single_mode():
         density_matrix_single_mode(cov, N, cutoff=cutoff, normalize=True, method="Coo coo ca choo")
 
 
-def test_haf_blocked():
+@pytest.mark.parametrize("n1", [0, 1, 2, 3])
+@pytest.mark.parametrize("n2", [0, 1, 2, 3, 4])
+def test_haf_blocked(n1, n2):
     """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)]
+    n1 = 3
+    n2 = 4
+    reps_list = [[i, n1 - i, n2] for i in range(n1 + 1)]
     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)
+    repeats = (n1, n2)
     haf_val = haf_blocked(A, blocks=blocks, repeats=repeats) / np.product(factorial(repeats))
     assert np.allclose(haf_sum, haf_val)

From 085d9bdf15b6f36f165f98f8f6ada1167bcdb430 Mon Sep 17 00:00:00 2001
From: Nicolas Quesada <nquesada@pop-os.localdomain>
Date: Wed, 24 Apr 2024 07:22:44 -0400
Subject: [PATCH 06/15] More tests, always

---
 thewalrus/tests/test_internal_modes.py | 2 --
 1 file changed, 2 deletions(-)

diff --git a/thewalrus/tests/test_internal_modes.py b/thewalrus/tests/test_internal_modes.py
index f9513589..674071d9 100644
--- a/thewalrus/tests/test_internal_modes.py
+++ b/thewalrus/tests/test_internal_modes.py
@@ -1541,8 +1541,6 @@ def test_haf_blocked(n1, n2):
     n = 6
     B = np.random.rand(n, n) + 1j * np.random.rand(n, n)
     A = B + B.T
-    n1 = 3
-    n2 = 4
     reps_list = [[i, n1 - i, n2] for i in range(n1 + 1)]
     haf_sum = 0j
     for reps in reps_list:

From 6fd264e7065230584b6cf1030728497bcc83ef24 Mon Sep 17 00:00:00 2001
From: Nicolas Quesada <nquesada@pop-os.localdomain>
Date: Wed, 24 Apr 2024 09:14:05 -0400
Subject: [PATCH 07/15] minor pylint improvements

---
 thewalrus/internal_modes/fock_density_matrices.py | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/thewalrus/internal_modes/fock_density_matrices.py b/thewalrus/internal_modes/fock_density_matrices.py
index 945d187e..c995026f 100644
--- a/thewalrus/internal_modes/fock_density_matrices.py
+++ b/thewalrus/internal_modes/fock_density_matrices.py
@@ -190,7 +190,7 @@ def density_matrix_single_mode(
     # N_nums = list(pattern.values())
     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)
@@ -238,5 +238,5 @@ def density_matrix_single_mode(
         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")

From 94cf74b46dbe43623a35fb6441583d6ac8bdb6ba Mon Sep 17 00:00:00 2001
From: Nicolas Quesada <nquesada@pop-os.localdomain>
Date: Wed, 24 Apr 2024 09:51:09 -0400
Subject: [PATCH 08/15] Adds extra test

---
 .../internal_modes/fock_density_matrices.py     | 17 ++++++-----------
 thewalrus/tests/test_internal_modes.py          | 11 ++++++-----
 2 files changed, 12 insertions(+), 16 deletions(-)

diff --git a/thewalrus/internal_modes/fock_density_matrices.py b/thewalrus/internal_modes/fock_density_matrices.py
index c995026f..f2345f02 100644
--- a/thewalrus/internal_modes/fock_density_matrices.py
+++ b/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"]:
diff --git a/thewalrus/tests/test_internal_modes.py b/thewalrus/tests/test_internal_modes.py
index 674071d9..78aeefd0 100644
--- a/thewalrus/tests/test_internal_modes.py
+++ b/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"
         )
     )
 

From 575d0d0ef997a1dbcd4a15a0dad8be47e0d85bae Mon Sep 17 00:00:00 2001
From: Nicolas Quesada <nquesada@pop-os.localdomain>
Date: Wed, 24 Apr 2024 09:53:04 -0400
Subject: [PATCH 09/15] black

---
 thewalrus/tests/test_internal_modes.py | 14 ++++++++++++--
 1 file changed, 12 insertions(+), 2 deletions(-)

diff --git a/thewalrus/tests/test_internal_modes.py b/thewalrus/tests/test_internal_modes.py
index 78aeefd0..8afeb4ce 100644
--- a/thewalrus/tests/test_internal_modes.py
+++ b/thewalrus/tests/test_internal_modes.py
@@ -1162,12 +1162,22 @@ def test_mixed_heralded_photon(nh, method, herald):
 
     p_a = np.diag(
         density_matrix_single_mode(
-            cov, {herald: 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, {herald: 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",
         )
     )
 

From b9ce7ae35cbe7bce19c978f02f474a1dc6b53235 Mon Sep 17 00:00:00 2001
From: Nicolas Quesada <nquesada@pop-os.localdomain>
Date: Wed, 24 Apr 2024 10:02:46 -0400
Subject: [PATCH 10/15] minor simplification

---
 thewalrus/internal_modes/fock_density_matrices.py | 11 ++++++-----
 1 file changed, 6 insertions(+), 5 deletions(-)

diff --git a/thewalrus/internal_modes/fock_density_matrices.py b/thewalrus/internal_modes/fock_density_matrices.py
index f2345f02..a2a16e9d 100644
--- a/thewalrus/internal_modes/fock_density_matrices.py
+++ b/thewalrus/internal_modes/fock_density_matrices.py
@@ -177,11 +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
-    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 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)

From 62cd625b18ec602e5ee15f71cf8dd9dcad2c9f3b Mon Sep 17 00:00:00 2001
From: Nicolas Quesada <nquesada@pop-os.localdomain>
Date: Wed, 24 Apr 2024 10:20:44 -0400
Subject: [PATCH 11/15] some tests pass

---
 .../internal_modes/fock_density_matrices.py   | 64 +++++++++----------
 1 file changed, 30 insertions(+), 34 deletions(-)

diff --git a/thewalrus/internal_modes/fock_density_matrices.py b/thewalrus/internal_modes/fock_density_matrices.py
index a2a16e9d..ebb1313d 100644
--- a/thewalrus/internal_modes/fock_density_matrices.py
+++ b/thewalrus/internal_modes/fock_density_matrices.py
@@ -196,41 +196,37 @@ 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 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))
+        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
-                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))
-                    )
+                            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
+        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

From d62716d8fc7bed9b391b46c906509b2c19a2e208 Mon Sep 17 00:00:00 2001
From: Nicolas Quesada <nquesada@pop-os.localdomain>
Date: Wed, 24 Apr 2024 10:21:08 -0400
Subject: [PATCH 12/15] some tests pass

---
 .../internal_modes/fock_density_matrices.py   | 45 +++++++++----------
 1 file changed, 22 insertions(+), 23 deletions(-)

diff --git a/thewalrus/internal_modes/fock_density_matrices.py b/thewalrus/internal_modes/fock_density_matrices.py
index ebb1313d..a50ff4e4 100644
--- a/thewalrus/internal_modes/fock_density_matrices.py
+++ b/thewalrus/internal_modes/fock_density_matrices.py
@@ -199,33 +199,32 @@ def density_matrix_single_mode(
         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))
+                    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))
                             )
-                            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
+                        )
+                        dm[i, j] = np.conj(dm[j, i])
+                    else:
+                        dm[i, j] = 0
+                        dm[j, i] = 0
         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))
+                haf_blocked(A, blocks=blocks, repeats=patt_long) / np.prod(factorial(patt_long))
             )
         if normalize:
             dm = dm / np.trace(dm)

From 828aeee4610fb3344cb492c18e81a2b778614922 Mon Sep 17 00:00:00 2001
From: Nicolas Quesada <nquesada@pop-os.localdomain>
Date: Wed, 24 Apr 2024 10:47:08 -0400
Subject: [PATCH 13/15] Moves the tests that take forever to the end of the
 test file

---
 thewalrus/tests/test_internal_modes.py | 234 +++++++++++++------------
 1 file changed, 118 insertions(+), 116 deletions(-)

diff --git a/thewalrus/tests/test_internal_modes.py b/thewalrus/tests/test_internal_modes.py
index 8afeb4ce..f320c3b2 100644
--- a/thewalrus/tests/test_internal_modes.py
+++ b/thewalrus/tests/test_internal_modes.py
@@ -1320,73 +1320,6 @@ def test_lossy_gkp(method):
     )
     assert np.allclose(np.diag(rho_loss1), probs)
 
-
-@pytest.mark.parametrize("method", ["recursive", "non-recursive"])
-def test_vac_schmidt_modes_gkp(method):
-    """
-    add vacuum schmidt modes and check it doesn't change the state
-    """
-    m1, m2 = 5, 7
-    params = np.array(
-        [
-            -1.38155106,
-            -1.21699567,
-            0.7798817,
-            1.04182349,
-            0.87702211,
-            0.90243916,
-            1.48353639,
-            1.6962906,
-            -0.24251599,
-            0.1958,
-        ]
-    )
-    sq_r = params[:3]
-    bs_theta1, bs_theta2, bs_theta3 = params[3:6]
-    bs_phi1, bs_phi2, bs_phi3 = params[6:9]
-    sq_virt = params[9]
-
-    S1 = squeezing(np.abs(sq_r), phi=np.angle(sq_r))
-    BS1, BS2, BS3 = (
-        beam_splitter(bs_theta1, bs_phi1),
-        beam_splitter(bs_theta2, bs_phi2),
-        beam_splitter(bs_theta3, bs_phi3),
-    )
-    Usymp1, Usymp2, Usymp3 = (
-        expand(BS1, [0, 1], 3),
-        expand(BS2, [1, 2], 3),
-        expand(BS3, [0, 1], 3),
-    )
-    Usymp = Usymp3 @ Usymp2 @ Usymp1
-    r2 = np.array([0, 0, sq_virt])
-    S2 = squeezing(np.abs(r2), phi=np.angle(r2))
-    Z = S2 @ Usymp @ S1
-    cov = Z @ Z.T
-    mu = np.zeros([len(cov)])
-
-    cutoff = 26
-    psi = state_vector(mu, cov, post_select={1: m1, 2: m2}, cutoff=cutoff)
-
-    rho1 = np.outer(psi, psi.conj())
-    rho1 /= np.trace(rho1)
-
-    M = 3
-    K = 5
-    big_cov = np.eye(2 * M * K, dtype=np.complex128)
-    big_cov[::K, ::K] = cov
-
-    rho_big = density_matrix_single_mode(big_cov, {1: m1, 2: m2}, cutoff=cutoff, method=method)
-    rho_big /= np.trace(rho_big)
-
-    assert np.allclose(rho1, rho_big, atol=4e-4)
-    probs = np.diag(
-        density_matrix_single_mode(
-            big_cov, {1: m1, 2: m2}, cutoff=cutoff, normalize=True, method="diagonals"
-        )
-    )
-    assert np.allclose(np.diag(rho1), probs)
-
-
 @pytest.mark.parametrize("method", ["recursive", "non-recursive"])
 def test_density_matrix_error(method):
     """Testing value errors in density_matrix_single_mode"""
@@ -1426,55 +1359,6 @@ def test_density_matrix_error(method):
         density_matrix_single_mode(cov, N, LO_overlap=LO_overlap2)
 
 
-@pytest.mark.parametrize("cutoff", [8, 9])
-@pytest.mark.parametrize("method", ["non-recursive"])
-def test_density_matrix(cutoff, method):
-    """
-    test generation of heralded density matrix against combinatorial calculation
-    """
-    # Note that the recursive method fails this
-    U = unitary_group.rvs(2)
-
-    N = {0: 3}
-
-    efficiency = 1 * np.ones(2)
-
-    noise = None
-
-    n0 = 2.9267754749886055
-    n1 = 2.592138225047742
-    zs0 = np.array([np.arcsinh(np.sqrt(n0))])
-    zs1 = np.array([np.arcsinh(np.sqrt(n1))])
-    rjs = [zs0, zs1]
-
-    O = np.identity(2, dtype=np.complex128)
-    S = 0.8 * np.exp(0 * 1j)
-    O[0, 1] = S.conj()
-    O[1, 0] = S
-
-    dm = heralded_density_matrix(
-        rjs, O, U, N, efficiency=efficiency, noise=noise, Ncutoff=cutoff, thresh=5e-3
-    )
-
-    rho = np.zeros((cutoff, cutoff), dtype=np.complex128)
-    for i in range(cutoff):
-        for j in range(cutoff):
-            rho[i, j] = sum(dm[i, j, m, m] for m in range(cutoff))
-
-    rho_norm = rho / np.trace(rho)
-
-    cov = prepare_cov(rjs, U, O=O, thresh=5e-3)
-
-    rho2 = density_matrix_single_mode(cov, N, cutoff=cutoff, method=method)
-    rho2_norm = rho2 / np.trace(rho2).real
-
-    probs = np.diag(
-        density_matrix_single_mode(cov, N, cutoff=cutoff, normalize=True, method="diagonals")
-    )
-
-    assert np.allclose(rho_norm, rho2_norm, atol=1e-6, rtol=1e-6)
-    assert np.allclose(np.diag(rho_norm), probs)
-
 
 def test_density_matrix_LO():
     """
@@ -1562,3 +1446,121 @@ def test_haf_blocked(n1, n2):
     repeats = (n1, n2)
     haf_val = haf_blocked(A, blocks=blocks, repeats=repeats) / np.product(factorial(repeats))
     assert np.allclose(haf_sum, haf_val)
+
+
+@pytest.mark.parametrize("cutoff", [8, 9])
+@pytest.mark.parametrize("method", ["non-recursive"])
+def test_density_matrix(cutoff, method):
+    """
+    test generation of heralded density matrix against combinatorial calculation
+    """
+    # Note that the recursive method fails this
+    # Note this test is super slow
+    U = unitary_group.rvs(2)
+
+    N = {0: 3}
+
+    efficiency = 1 * np.ones(2)
+
+    noise = None
+
+    n0 = 2.9267754749886055
+    n1 = 2.592138225047742
+    zs0 = np.array([np.arcsinh(np.sqrt(n0))])
+    zs1 = np.array([np.arcsinh(np.sqrt(n1))])
+    rjs = [zs0, zs1]
+
+    O = np.identity(2, dtype=np.complex128)
+    S = 0.8 * np.exp(0 * 1j)
+    O[0, 1] = S.conj()
+    O[1, 0] = S
+
+    dm = heralded_density_matrix(
+        rjs, O, U, N, efficiency=efficiency, noise=noise, Ncutoff=cutoff, thresh=5e-3
+    )
+
+    rho = np.zeros((cutoff, cutoff), dtype=np.complex128)
+    for i in range(cutoff):
+        for j in range(cutoff):
+            rho[i, j] = sum(dm[i, j, m, m] for m in range(cutoff))
+
+    rho_norm = rho / np.trace(rho)
+
+    cov = prepare_cov(rjs, U, O=O, thresh=5e-3)
+
+    rho2 = density_matrix_single_mode(cov, N, cutoff=cutoff, method=method)
+    rho2_norm = rho2 / np.trace(rho2).real
+
+    probs = np.diag(
+        density_matrix_single_mode(cov, N, cutoff=cutoff, normalize=True, method="diagonals")
+    )
+
+    assert np.allclose(rho_norm, rho2_norm, atol=1e-6, rtol=1e-6)
+    assert np.allclose(np.diag(rho_norm), probs)
+
+
+
+@pytest.mark.parametrize("method", ["recursive", "non-recursive"])
+def test_vac_schmidt_modes_gkp(method):
+    """
+    add vacuum schmidt modes and check it doesn't change the state
+    """
+    m1, m2 = 5, 7
+    params = np.array(
+        [
+            -1.38155106,
+            -1.21699567,
+            0.7798817,
+            1.04182349,
+            0.87702211,
+            0.90243916,
+            1.48353639,
+            1.6962906,
+            -0.24251599,
+            0.1958,
+        ]
+    )
+    sq_r = params[:3]
+    bs_theta1, bs_theta2, bs_theta3 = params[3:6]
+    bs_phi1, bs_phi2, bs_phi3 = params[6:9]
+    sq_virt = params[9]
+
+    S1 = squeezing(np.abs(sq_r), phi=np.angle(sq_r))
+    BS1, BS2, BS3 = (
+        beam_splitter(bs_theta1, bs_phi1),
+        beam_splitter(bs_theta2, bs_phi2),
+        beam_splitter(bs_theta3, bs_phi3),
+    )
+    Usymp1, Usymp2, Usymp3 = (
+        expand(BS1, [0, 1], 3),
+        expand(BS2, [1, 2], 3),
+        expand(BS3, [0, 1], 3),
+    )
+    Usymp = Usymp3 @ Usymp2 @ Usymp1
+    r2 = np.array([0, 0, sq_virt])
+    S2 = squeezing(np.abs(r2), phi=np.angle(r2))
+    Z = S2 @ Usymp @ S1
+    cov = Z @ Z.T
+    mu = np.zeros([len(cov)])
+
+    cutoff = 26
+    psi = state_vector(mu, cov, post_select={1: m1, 2: m2}, cutoff=cutoff)
+
+    rho1 = np.outer(psi, psi.conj())
+    rho1 /= np.trace(rho1)
+
+    M = 3
+    K = 5
+    big_cov = np.eye(2 * M * K, dtype=np.complex128)
+    big_cov[::K, ::K] = cov
+
+    rho_big = density_matrix_single_mode(big_cov, {1: m1, 2: m2}, cutoff=cutoff, method=method)
+    rho_big /= np.trace(rho_big)
+
+    assert np.allclose(rho1, rho_big, atol=4e-4)
+    probs = np.diag(
+        density_matrix_single_mode(
+            big_cov, {1: m1, 2: m2}, cutoff=cutoff, normalize=True, method="diagonals"
+        )
+    )
+    assert np.allclose(np.diag(rho1), probs)

From da4913d5bc28dff0288a0e39538bad64d4b8d59b Mon Sep 17 00:00:00 2001
From: Nicolas Quesada <nquesada@pop-os.localdomain>
Date: Wed, 24 Apr 2024 10:54:31 -0400
Subject: [PATCH 14/15] Passes black

---
 thewalrus/tests/test_internal_modes.py | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/thewalrus/tests/test_internal_modes.py b/thewalrus/tests/test_internal_modes.py
index f320c3b2..496c25c8 100644
--- a/thewalrus/tests/test_internal_modes.py
+++ b/thewalrus/tests/test_internal_modes.py
@@ -1320,6 +1320,7 @@ def test_lossy_gkp(method):
     )
     assert np.allclose(np.diag(rho_loss1), probs)
 
+
 @pytest.mark.parametrize("method", ["recursive", "non-recursive"])
 def test_density_matrix_error(method):
     """Testing value errors in density_matrix_single_mode"""
@@ -1359,7 +1360,6 @@ def test_density_matrix_error(method):
         density_matrix_single_mode(cov, N, LO_overlap=LO_overlap2)
 
 
-
 def test_density_matrix_LO():
     """
     test generation of heralded density matrix in LO basis against combinatorial calculation
@@ -1499,7 +1499,6 @@ def test_density_matrix(cutoff, method):
     assert np.allclose(np.diag(rho_norm), probs)
 
 
-
 @pytest.mark.parametrize("method", ["recursive", "non-recursive"])
 def test_vac_schmidt_modes_gkp(method):
     """

From e9a1b76b20dac34b5a28444ea788849d77a3a7d8 Mon Sep 17 00:00:00 2001
From: Nicolas Quesada <nquesada@pop-os.localdomain>
Date: Wed, 24 Apr 2024 11:06:07 -0400
Subject: [PATCH 15/15] Adds else

---
 thewalrus/internal_modes/fock_density_matrices.py | 11 ++++++-----
 1 file changed, 6 insertions(+), 5 deletions(-)

diff --git a/thewalrus/internal_modes/fock_density_matrices.py b/thewalrus/internal_modes/fock_density_matrices.py
index a50ff4e4..3b097d99 100644
--- a/thewalrus/internal_modes/fock_density_matrices.py
+++ b/thewalrus/internal_modes/fock_density_matrices.py
@@ -221,11 +221,12 @@ def density_matrix_single_mode(
                     else:
                         dm[i, j] = 0
                         dm[j, i] = 0
-        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))
-            )
+        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