Adds more tests

thewalrus/internal_modes/__init__.py

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

thewalrus/internal_modes/fock_density_matrices.py

@@ -25,40 +25,9 @@
     nb_Qmat,
     spatial_reps_to_schmidt_reps,
     fact,
+    project_onto_local_oscillator,
 )
 
-
-@numba.jit(nopython=True, parallel=True, cache=True)
-def project_onto_local_oscillator(cov, M, LO_overlap=None, hbar=2):
-    """Projects a given covariance matrix into the relevant internal mode in the first external mode.
-
-    Args:
-        cov (array): 2MK x 2MK covariance matrix
-        LO_overlap (array): overlap between internal modes and local oscillator
-
-    Returns:
-        (array): projected covariance matrix
-    """
-
-    K = cov.shape[0] // (2 * M)
-
-    # filter out all unwanted Schmidt modes in heralded spatial mode
-
-    # create passive transformation of filter
-    T = np.zeros((M * K, M * K), dtype=np.complex128)
-    if LO_overlap is not None:
-        T[0][:K] = LO_overlap
-    else:
-        T[0, 0] = 1
-    T[K:, K:] = np.eye((M - 1) * K, dtype=np.complex128)
-
-    # apply channel of filter
-    P = nb_block(((T.real, -T.imag), (T.imag, T.real)))
-    L = (hbar / 2) * (np.eye(P.shape[0]) - P @ P.T)
-    cov = P @ cov @ P.T + L
-
-    return cov
-
 # pylint: disable=too-many-arguments, too-many-statements
 @numba.jit(nopython=True, parallel=True, cache=True)
 def _density_matrix_single_mode(cov, pattern, normalize=False, LO_overlap=None, cutoff=13, hbar=2):

thewalrus/internal_modes/pnr_statistics.py

@@ -22,11 +22,12 @@
 
 from scipy.special import factorial as fac
 
-from ..quantum import Qmat
+from ..quantum import Qmat, Amat
+from thewalrus.symplectic import passive_transformation
 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
+from .utils import spatial_reps_to_schmidt_reps, spatial_modes_to_schmidt_modes, project_onto_local_oscillator
 
 
 @numba.jit(nopython=True, parallel=True, cache=True)
@@ -179,3 +180,62 @@ 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 density matrix 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_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 = np.array(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])
+    patt = [pattern[i] for i in range(1,1+len(pattern))]
+    for i in range(cutoff):
+        patt_long = [i]+patt
+        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

thewalrus/internal_modes/utils.py

@@ -104,3 +104,37 @@ def nb_block(X):  # pragma: no cover
     xtmp1 = np.concatenate(X[0], axis=1)
     xtmp2 = np.concatenate(X[1], axis=1)
     return np.concatenate((xtmp1, xtmp2), axis=0)
+
+
+
+
+@numba.jit(nopython=True, parallel=True, cache=True)
+def project_onto_local_oscillator(cov, M, LO_overlap=None, hbar=2):
+    """Projects a given covariance matrix into the relevant internal mode in the first external mode.
+
+    Args:
+        cov (array): 2MK x 2MK covariance matrix
+        LO_overlap (array): overlap between internal modes and local oscillator
+
+    Returns:
+        (array): projected covariance matrix
+    """
+
+    K = cov.shape[0] // (2 * M)
+
+    # filter out all unwanted Schmidt modes in heralded spatial mode
+
+    # create passive transformation of filter
+    T = np.zeros((M * K, M * K), dtype=np.complex128)
+    if LO_overlap is not None:
+        T[0][:K] = LO_overlap
+    else:
+        T[0, 0] = 1
+    T[K:, K:] = np.eye((M - 1) * K, dtype=np.complex128)
+
+    # apply channel of filter
+    P = nb_block(((T.real, -T.imag), (T.imag, T.real)))
+    L = (hbar / 2) * (np.eye(P.shape[0]) - P @ P.T)
+    cov = P @ cov @ P.T + L
+
+    return cov

thewalrus/tests/test_internal_modes.py

@@ -53,7 +53,7 @@
 from thewalrus.internal_modes import (
     pnr_prob,
     density_matrix_single_mode,
-    project_onto_local_oscillator,
+    probabilities_single_mode
 )
 from thewalrus.internal_modes.prepare_cov import (
     O_matrix,
@@ -63,6 +63,8 @@
     LO_overlaps,
 )
 
+from thewalrus.internal_modes.utils import project_onto_local_oscillator
+
 ### auxilliary functions for testing ###
 # if we want to have less auxilliary functions, we can remove a few tests and get rid of it all
 
@@ -1159,9 +1161,19 @@ def test_mixed_heralded_photon(nh):
         cov, {1: nh}, normalize=True, LO_overlap=LO_overlapb, cutoff=nh+1
     )
 
+    p_a = probabilities_single_mode(
+        cov, {1: nh}, normalize=True, LO_overlap=LO_overlapa, cutoff=nh+1
+    )
+    p_b = probabilities_single_mode(
+        cov, {1: nh}, normalize=True, LO_overlap=LO_overlapb, cutoff=nh+1
+    )
+
+
+
     assert np.allclose(dm_modea, rho_a)
     assert np.allclose(dm_modeb, rho_b)
-
+    assert np.allclose(np.diag(dm_modea), p_a)
+    assert np.allclose(np.diag(dm_modeb), p_b)
 
 def test_pure_gkp():
     """test pure gkp state density matrix using 2 methods from the walrus against
@@ -1221,6 +1233,9 @@ def test_pure_gkp():
     assert np.allclose(rho1, rho2, atol=2.5e-4)
     assert np.allclose(rho1, rho3, atol=4.7e-4)
     assert np.allclose(rho2, rho3, atol=4.8e-4)
+    probs = probabilities_single_mode(cov, {1: m1, 2: m2}, cutoff=cutoff,normalize=True)
+    assert np.allclose(np.diag(rho1), probs)
+
     #### Note that the tolerances are higher than they should be.
 
 
@@ -1278,7 +1293,8 @@ def test_lossy_gkp():
     rho_loss2 = density_matrix_single_mode(cov_lossy, {1: m1, 2: m2}, cutoff=cutoff)
     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)
+    assert np.allclose(np.diag(rho_loss1), probs)
 
 def test_vac_schmidt_modes_gkp():
     """
@@ -1337,6 +1353,8 @@ def test_vac_schmidt_modes_gkp():
     rho_big /= np.trace(rho_big)
 
     assert np.allclose(rho1, rho_big, atol=4e-4)
+    probs = probabilities_single_mode(big_cov, {1: m1, 2: m2}, cutoff=cutoff,normalize=True)
+    assert np.allclose(np.diag(rho1), probs)
 
 
 def test_density_matrix_error():
@@ -1417,7 +1435,10 @@ def test_density_matrix(cutoff):
     rho2 = density_matrix_single_mode(cov, N, cutoff=cutoff)
     rho2_norm = rho2 / np.trace(rho2).real
 
+    # probs = probabilities_single_mode(cov, N, cutoff=cutoff, normalize=True)
+
     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():