Off diagonal

thewalrus/_hafnian.py

@@ -273,7 +273,7 @@ def f_loop_odd(AX, AX_S, XD_S, D_S, n, oddloop, oddVX_S):  # pragma: no cover
 
 
 @numba.jit(nopython=True, cache=True)
-def f_from_powertrace(powertraces, n):
+def f_from_powertrace(powertraces, n):  # pragma: no cover
     """Evaluate the polynomial coefficients of the function in the eigenvalue-trace formula from the powertraces.
 
     Args:

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
-from .fock_density_matrices import density_matrix_single_mode
+from .pnr_statistics import pnr_prob, haf_blocked
+from .fock_density_matrices import density_matrix_single_mode, check_probabilities
 from .prepare_cov import *

thewalrus/internal_modes/fock_density_matrices.py

@@ -15,22 +15,27 @@
 Set of functions for calculating Fock basis density matrices for heralded states created by PNR measurements on Gaussian states with multiple internal modes
 """
 
+import warnings
 import numpy as np
 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,
     spatial_reps_to_schmidt_reps,
     fact,
     project_onto_local_oscillator,
 )
+from .pnr_statistics import haf_blocked
+from ..quantum import Qmat, Amat
 
 
 # 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, LO_overlap=None, cutoff=13, hbar=2
+):  # pragma: no cover
     """
     numba function (use the wrapper function: density_matrix_multimode)
 
@@ -40,7 +45,6 @@ def _density_matrix_single_mode(cov, pattern, normalize=False, LO_overlap=None,
     Args:
         cov (array): 2MK x 2MK covariance matrix
         pattern (array): M-1 length array of the heralding pattern
-        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)
@@ -75,13 +79,13 @@ def _density_matrix_single_mode(cov, pattern, normalize=False, LO_overlap=None,
     x = np.array(x)
     Ax = nb_ix(A, x, x)  # A[np.ix_(x, x)]
 
-    edge_reps = np.array([half_c, half_c, 1] + list(pattern))
+    edge_reps = np.array((half_c, half_c, 1) + pattern)
     n_edges = 3 + K * len(pattern)
 
     assert n_edges == Ax.shape[0] // 2 == 3 + K * (M - 1)
 
     N_max = 2 * edge_reps.sum()
-    N_fixed = 2 * np.sum(pattern)
+    N_fixed = 2 * sum(pattern)
 
     steps = np.prod(edge_reps + 1)
 
@@ -126,7 +130,9 @@ def _density_matrix_single_mode(cov, pattern, normalize=False, LO_overlap=None,
         haf_arr += haf_arr_new
 
     rho = (
-        (-1) ** pattern.sum() * haf_arr / (np.sqrt(np.linalg.det(Q).real) * np.prod(fact[pattern]))
+        (-1) ** sum(pattern)
+        * haf_arr
+        / (np.sqrt(np.linalg.det(Q).real) * np.prod(fact[np.array(list(pattern))]))
     )
 
     for n in range(cutoff):
@@ -135,14 +141,44 @@ def _density_matrix_single_mode(cov, pattern, normalize=False, LO_overlap=None,
 
     rho = rho[:cutoff, :cutoff]
 
-    if normalize:
-        return rho / np.trace(rho).real
     return rho
 
 
-def density_matrix_single_mode(cov, pattern, normalize=False, LO_overlap=None, cutoff=13, hbar=2):
+def check_probabilities(probs, atol=1e-08):
+    """
+    Convenience function for checking that the input is close enough to a probability distribution.
+
+    Args:
+        probs (array): probabilities to be tested.
+        atol (float): absolute tolerance relative to the normalization.
+
+    Returns:
+        (boolean): whether the test passed or not.
+    """
+    real_probs = probs.real
+    imag_probs = probs.imag
+    pos_probs = real_probs[real_probs > 0]
+    neg_probs = real_probs[real_probs < 0]
+    net_prob = sum(pos_probs)
+    if np.any(np.abs(imag_probs) > atol * net_prob):
+        return False
+    if np.any(np.abs(neg_probs) > atol * net_prob):
+        return False
+    return True
+
+
+def density_matrix_single_mode(
+    cov,
+    pattern,
+    normalize=False,
+    LO_overlap=None,
+    cutoff=13,
+    hbar=2,
+    method="recursive",
+    atol=1e-08,
+):
     """
-    calculates density matrix of first mode when heralded by pattern on a zero-displaced, M-mode Gaussian state
+    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:
@@ -152,16 +188,18 @@ def density_matrix_single_mode(cov, pattern, normalize=False, LO_overlap=None, c
         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)
+        method (str): which method to use, "recursive", "non-recursive" or "diagonals"
+        atol (float): value for raising warning when testing for valid probabilities
     Returns:
         array[complex]: (cutoff+1, cutoff+1) dimension density matrix
     """
 
-    cov = np.array(cov).astype(np.float64)
+    cov = np.array(cov.real).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()))
+    N_nums = tuple(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]:
@@ -170,12 +208,72 @@ def density_matrix_single_mode(cov, pattern, normalize=False, LO_overlap=None, c
             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)
-    return _density_matrix_single_mode(cov, N_nums, normalize, LO_overlap, cutoff, hbar)
+    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":
+        vals = _density_matrix_single_mode(cov, N_nums, LO_overlap, cutoff, hbar)
+        if check_probabilities(np.diag(vals), atol=atol) is False:
+            warnings.warn(
+                "Some of the diagonal elements of the density matrix are significantly negative or have significant imaginary parts. Try using the `non-recursive` method instead.",
+                UserWarning,
+            )
+        if normalize:
+            vals = vals / np.trace(vals).real
+        return vals
+    if method in ["non-recursive", "diagonals"]:
+        cov = project_onto_local_oscillator(cov, M, LO_overlap=LO_overlap, hbar=hbar)
+        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)
+        if method == "non-recursive":
+            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([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
+        else:
+            for i in range(cutoff):
+                patt_long = (i,) + N_nums
+                dm[i, i] = (
+                    pref
+                    * haf_blocked(A, blocks=blocks, repeats=patt_long)
+                    / np.prod(factorial(patt_long))
+                )
+        if check_probabilities(np.diag(dm)) is False:
+            warnings.warn(
+                "Some of the diagonal elements of the density matrix are significantly negative or have significant imaginary parts.",
+                UserWarning,
+            )
+        if normalize:
+            dm = dm / np.trace(dm)
+        return dm
+
+    raise ValueError("Unknown method for density_matrix_single_mode")

thewalrus/internal_modes/pnr_statistics.py

@@ -22,20 +22,18 @@
 
 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,
 )
 
 
 @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 +127,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.
 
@@ -165,12 +163,12 @@ 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)
 
 
 @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:
@@ -195,64 +193,5 @@ def _haf_blocked_numba(A, blocks, repeats_p):
             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
-
-
-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

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: