Rachels code review

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

@@ -17,7 +17,7 @@
 
 import numpy as np
 import numba
-from scipy.special import factorial as fac
+from scipy.special import factorial
 
 from ..symplectic import passive_transformation
 from .._hafnian import nb_binom, nb_ix, find_kept_edges, f_from_matrix
@@ -33,7 +33,9 @@
 
 # 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):
+def _density_matrix_single_mode(
+    cov, pattern, normalize=False, LO_overlap=None, cutoff=13, hbar=2
+):  # pragma: no cover
     """
     numba function (use the wrapper function: density_matrix_multimode)
 
@@ -188,36 +190,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)
-    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 = 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] = (
-                    fact
-                    * haf_blocked(Aperm, blocks=new_blocks, repeats=patt_long)
-                    / (np.prod(fac(patt_long[1:-1])) * np.sqrt(fac(i) * fac(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

@@ -35,7 +35,7 @@
 
 
 @numba.jit(nopython=True, parallel=True, cache=True)
-def hafkd(As, edge_reps, K=1):
+def hafkd(As, edge_reps, K=1):  # pragma: no cover
     r"""
     generalised version of hafnian to include multiple internal modes
 
@@ -129,7 +129,7 @@ def pnr_prob(covs, i, hbar=2):
 
 
 @numba.jit(nopython=True, cache=True)
-def finite_difference_operator_coeffs(der_order, m, u=None, v=None):
+def finite_difference_operator_coeffs(der_order, m, u=None, v=None):  # pragma: no cover
     """Returns the mth coefficient of the finite difference operator of given derivative order.
     For details see:  E. T. Bax, Finite-difference algorithms for counting problems. PhD thesis, Caltech, 1998.
 
@@ -170,7 +170,7 @@ def haf_blocked(A, blocks, repeats):
 
 
 @numba.jit(nopython=True)
-def _haf_blocked_numba(A, blocks, repeats_p):
+def _haf_blocked_numba(A, blocks, repeats_p):  # pragma: no cover
     """Calculates the hafnian of the matrix with a given block and repetition pattern.
 
     Args:
@@ -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/internal_modes/utils.py

@@ -20,7 +20,7 @@
 
 
 @numba.jit(nopython=True, cache=True)
-def spatial_modes_to_schmidt_modes(spatial_modes, K):
+def spatial_modes_to_schmidt_modes(spatial_modes, K):  # pragma: no cover
     """
     returns index of schmidt modes corresponding to the give spatial modes.
     e.g. if there are K=3 schmidt modes and spatial_modes=[0,2]
@@ -42,7 +42,7 @@ def spatial_modes_to_schmidt_modes(spatial_modes, K):
 
 
 @numba.jit(nopython=True, cache=True)
-def spatial_reps_to_schmidt_reps(spatial_reps, K):
+def spatial_reps_to_schmidt_reps(spatial_reps, K):  # pragma: no cover
     """
     returns reps of schmidt modes corresponding to the give spatial reps.
     e.g. if there are K=3 schmidt modes and spatial_reps=[1,2]
@@ -107,7 +107,7 @@ def nb_block(X):  # pragma: no cover
 
 
 @numba.jit(nopython=True, parallel=True, cache=True)
-def project_onto_local_oscillator(cov, M, LO_overlap=None, hbar=2):
+def project_onto_local_oscillator(cov, M, LO_overlap=None, hbar=2):  # pragma: no cover
     """Projects a given covariance matrix into the relevant internal mode in the first external mode.
 
     Args: