passes black

thewalrus/internal_modes/__init__.py

@@ -15,6 +15,6 @@
 Functions to do internal modes/distinguishable GBS
 """
 
-from .pnr_statistics import pnr_prob, probabilities_single_mode, haf_blocked
+from .pnr_statistics import pnr_prob, haf_blocked
 from .fock_density_matrices import density_matrix_single_mode
 from .prepare_cov import *

thewalrus/internal_modes/fock_density_matrices.py

@@ -188,36 +188,46 @@ 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)
-    cov = project_onto_local_oscillator(cov, M, LO_overlap=LO_overlap, hbar=hbar)
-    num_modes = len(cov) // 2
-    A = Amat(cov)
-    Q = Qmat(cov)
-    pref = 1 / np.sqrt(np.linalg.det(Q).real)
-    blocks = np.arange(K * M).reshape([M, K])
-    dm = np.zeros([cutoff, cutoff], dtype=np.complex128)
-    num_modes = M * K
-    block_size = K
-    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))
-                )
-                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])
+    elif method == "non-recursive" or method == "diagonals":
+        cov = project_onto_local_oscillator(cov, M, LO_overlap=LO_overlap, hbar=hbar)
+        num_modes = len(cov) // 2
+        A = Amat(cov)
+        Q = Qmat(cov)
+        pref = 1 / np.sqrt(np.linalg.det(Q).real)
+        blocks = np.arange(K * M).reshape([M, K])
+        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:
-                dm[i, j] = 0
-                dm[j, i] = 0
-    if normalize:
-        dm = dm / np.trace(dm)
-    return dm
+                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))
+                        )
+                    )
+                    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
+    else:
+        raise ValueError("Unknown method for density_matrix_single_mode")

thewalrus/internal_modes/pnr_statistics.py

@@ -198,136 +198,3 @@ def _haf_blocked_numba(A, blocks, repeats_p):
 
         netsum += coeff_pref * f_from_matrix(AX_S, 2 * n)[n]
     return netsum
-
-
-def probabilities_single_mode(cov, pattern, normalize=False, LO_overlap=None, cutoff=13, hbar=2):
-    """
-    Calculates the diagonal of the density matrix, hence the name probabilities, of first mode when heralded by pattern on a zero-displaced, M-mode Gaussian state
-    where each mode contains K internal modes.
-
-    Args:
-        cov (array): 2MK x 2MK covariance matrix
-        pattern (dict): heralding pattern total photon number in the spatial modes (int), indexed by spatial mode
-        normalize (bool): whether to normalise the output density matrix
-        LO_overlap (array): overlap between internal modes and local oscillator
-        cutoff (int): photon number cutoff. Should be odd. Even numbers will be rounded up to an odd number
-        hbar (float): the value of hbar (default 2)
-    Returns:
-        array[complex]: (cutoff+1, cutoff+1) dimension density matrix
-    """
-
-    # The lines until A = Amat(...) are copies from density_matrix_single_mode
-    cov = np.array(cov).astype(np.float64)
-    M = len(pattern) + 1
-    K = cov.shape[0] // (2 * M)
-    if not set(list(pattern.keys())).issubset(set(list(np.arange(M)))):
-        raise ValueError("Keys of pattern must correspond to all but one spatial mode")
-    N_nums = list(pattern.values())
-    HM = list(set(list(np.arange(M))).difference(list(pattern.keys())))[0]
-    if LO_overlap is not None:
-        if not K == LO_overlap.shape[0]:
-            raise ValueError("Number of overlaps with LO must match number of internal modes")
-        if not (np.linalg.norm(LO_overlap) < 1 or np.allclose(np.linalg.norm(LO_overlap), 1)):
-            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)
-
-    cov = project_onto_local_oscillator(cov, M, LO_overlap=LO_overlap, hbar=hbar)
-
-    A = Amat(cov)
-    Q = Qmat(cov)
-    fact = 1 / np.sqrt(np.linalg.det(Q).real)
-    blocks = np.arange(K * M).reshape([M, K])
-    probs = np.zeros([cutoff])
-    for i in range(cutoff):
-        patt_long = [i] + N_nums
-        probs[i] = fact * np.real(
-            haf_blocked(A, blocks=blocks, repeats=patt_long) / np.prod(fac(patt_long))
-        )
-
-    if normalize:
-        probs = probs / np.sum(probs)
-    return probs
-
-
-def dm_single_mode(cov, pattern, normalize=False, LO_overlap=None, cutoff=13, hbar=2):
-    """
-    Calculates the diagonal of the density matrix, hence the name probabilities, of first mode when heralded by pattern on a zero-displaced, M-mode Gaussian state
-    where each mode contains K internal modes.
-
-    Args:
-        cov (array): 2MK x 2MK covariance matrix
-        pattern (dict): heralding pattern total photon number in the spatial modes (int), indexed by spatial mode
-        normalize (bool): whether to normalise the output density matrix
-        LO_overlap (array): overlap between internal modes and local oscillator
-        cutoff (int): photon number cutoff. Should be odd. Even numbers will be rounded up to an odd number
-        hbar (float): the value of hbar (default 2)
-    Returns:
-        array[complex]: (cutoff+1, cutoff+1) dimension density matrix
-    """
-
-    # The lines until A = Amat(...) are copies from density_matrix_single_mode
-    cov = np.array(cov).astype(np.float64)
-    M = len(pattern) + 1
-    K = cov.shape[0] // (2 * M)
-    if not set(list(pattern.keys())).issubset(set(list(np.arange(M)))):
-        raise ValueError("Keys of pattern must correspond to all but one spatial mode")
-    N_nums = list(pattern.values())
-    HM = list(set(list(np.arange(M))).difference(list(pattern.keys())))[0]
-    if LO_overlap is not None:
-        if not K == LO_overlap.shape[0]:
-            raise ValueError("Number of overlaps with LO must match number of internal modes")
-        if not (np.linalg.norm(LO_overlap) < 1 or np.allclose(np.linalg.norm(LO_overlap), 1)):
-            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)
-
-    cov = project_onto_local_oscillator(cov, M, LO_overlap=LO_overlap, hbar=hbar)
-    num_modes = len(cov) // 2
-    A = Amat(cov)
-    Q = Qmat(cov)
-    fact = 1 / np.sqrt(np.linalg.det(Q).real)
-    blocks = np.arange(K * M).reshape([M, K])
-    dm = np.zeros([cutoff, cutoff], dtype=np.complex128)
-    num_modes = M * K
-    block_size = K
-    for i in range(cutoff):
-        for j in range(i + 1):
-            if (i - j) % 2 == 0:
-                patt_long = [j] + N_nums + [(i - j) // 2]
-                new_blocks = np.concatenate((blocks, np.array([1 + 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[i, j] = (
-                    fact
-                    * haf_blocked(Aperm, blocks=new_blocks, repeats=patt_long)
-                    / (np.prod(fac(patt_long[1:-1])) * np.sqrt(fac(i) * fac(j)))
-                )
-                dm[j, i] = np.conj(dm[i, j])
-            else:
-                dm[i, j] = 0
-                dm[j, i] = 0
-    if normalize:
-        dm = dm / np.trace(dm)
-    return dm

thewalrus/tests/test_internal_modes.py

@@ -50,7 +50,7 @@
     squeezing,
 )
 
-from thewalrus.internal_modes import pnr_prob, density_matrix_single_mode, probabilities_single_mode
+from thewalrus.internal_modes import pnr_prob, density_matrix_single_mode
 from thewalrus.internal_modes.prepare_cov import (
     O_matrix,
     orthonormal_basis,
@@ -1159,11 +1159,15 @@ def test_mixed_heralded_photon(nh, method):
         cov, {1: nh}, normalize=True, LO_overlap=LO_overlapb, cutoff=nh + 1, method=method
     )
 
-    p_a = probabilities_single_mode(
-        cov, {1: nh}, normalize=True, LO_overlap=LO_overlapa, cutoff=nh + 1
+    p_a = np.diag(
+        density_matrix_single_mode(
+            cov, {1: nh}, normalize=True, LO_overlap=LO_overlapa, cutoff=nh + 1, method="diagonals"
+        )
     )
-    p_b = probabilities_single_mode(
-        cov, {1: nh}, normalize=True, LO_overlap=LO_overlapb, cutoff=nh + 1
+    p_b = np.diag(
+        density_matrix_single_mode(
+            cov, {1: nh}, normalize=True, LO_overlap=LO_overlapb, cutoff=nh + 1, method="diagonals"
+        )
     )
 
     assert np.allclose(dm_modea, rho_a)
@@ -1233,8 +1237,12 @@ def test_pure_gkp(method):
     assert np.allclose(
         rho2, rho3, atol=4.8e-4
     )  # For the method "non-recursive" the absolute max difference is 1e-8
-    # probs = probabilities_single_mode(cov, {1: m1, 2: m2}, cutoff=cutoff, normalize=True)
-    # assert np.allclose(np.diag(rho1), probs)
+    probs = np.diag(
+        density_matrix_single_mode(
+            cov, {1: m1, 2: m2}, cutoff=cutoff, normalize=True, method="diagonals"
+        )
+    )
+    assert np.allclose(np.diag(rho1), probs)
 
     #### Note that the tolerances are higher than they should be.
 
@@ -1294,7 +1302,11 @@ def test_lossy_gkp(method):
     rho_loss2 = density_matrix_single_mode(cov_lossy, {1: m1, 2: m2}, cutoff=cutoff, method=method)
     rho_loss2 /= np.trace(rho_loss2)
     assert np.allclose(rho_loss1, rho_loss2, atol=2.7e-4)
-    probs = probabilities_single_mode(cov_lossy, {1: m1, 2: m2}, cutoff=cutoff, normalize=True)
+    probs = np.diag(
+        density_matrix_single_mode(
+            cov_lossy, {1: m1, 2: m2}, cutoff=cutoff, normalize=True, method="diagonals"
+        )
+    )
     assert np.allclose(np.diag(rho_loss1), probs)
 
 
@@ -1400,11 +1412,12 @@ def test_density_matrix_error(method):
 
 
 @pytest.mark.parametrize("cutoff", [8, 9])
-@pytest.mark.parametrize("method", ["recursive", "non-recursive"])
+@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}
@@ -1440,7 +1453,9 @@ def test_density_matrix(cutoff, method):
     rho2 = density_matrix_single_mode(cov, N, cutoff=cutoff, method=method)
     rho2_norm = rho2 / np.trace(rho2).real
 
-    probs = probabilities_single_mode(cov, N, cutoff=cutoff, normalize=True)
+    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)