Adds the warning (#388)

* Adds the warning

* passes black

* Apply suggestions from code review

Co-authored-by: rachelchadwick <33873960+rachelchadwick@users.noreply.github.com>

* Update thewalrus/internal_modes/fock_density_matrices.py

---------

Co-authored-by: Nicolas Quesada <nquesada@pop-os.localdomain>
Co-authored-by: rachelchadwick <33873960+rachelchadwick@users.noreply.github.com>

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

@@ -16,5 +16,5 @@
 """
 
 from .pnr_statistics import pnr_prob, haf_blocked
-from .fock_density_matrices import density_matrix_single_mode
+from .fock_density_matrices import density_matrix_single_mode, check_probabilities
 from .prepare_cov import *

thewalrus/internal_modes/fock_density_matrices.py

@@ -15,6 +15,7 @@
 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
@@ -33,7 +34,7 @@
 # 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
+    cov, pattern, LO_overlap=None, cutoff=13, hbar=2
 ):  # pragma: no cover
     """
     numba function (use the wrapper function: density_matrix_multimode)
@@ -44,7 +45,6 @@ def _density_matrix_single_mode(
     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)
@@ -139,16 +139,44 @@ def _density_matrix_single_mode(
 
     rho = rho[:cutoff, :cutoff]
 
-    if normalize:
-        return rho / np.trace(rho).real
     return rho
 
 
+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"
+    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:
@@ -158,7 +186,8 @@ def density_matrix_single_mode(
         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" or "direct"
+        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
     """
@@ -185,7 +214,15 @@ def density_matrix_single_mode(
         cov = cov[:, double_perm][double_perm]
 
     if method == "recursive":
-        return _density_matrix_single_mode(cov, N_nums, normalize, LO_overlap, cutoff, hbar)
+        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)
         num_modes = len(cov) // 2
@@ -224,8 +261,10 @@ def density_matrix_single_mode(
         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))
+                dm[i, i] = (
+                    pref
+                    * haf_blocked(A, blocks=blocks, repeats=patt_long)
+                    / np.prod(factorial(patt_long))
                 )
         if normalize:
             dm = dm / np.trace(dm)

thewalrus/tests/test_internal_modes.py

@@ -50,7 +50,12 @@
     squeezing,
 )
 
-from thewalrus.internal_modes import pnr_prob, density_matrix_single_mode, haf_blocked
+from thewalrus.internal_modes import (
+    pnr_prob,
+    density_matrix_single_mode,
+    haf_blocked,
+    check_probabilities,
+)
 from thewalrus.internal_modes.prepare_cov import (
     O_matrix,
     orthonormal_basis,
@@ -1448,6 +1453,82 @@ def test_haf_blocked(n1, n2):
     assert np.allclose(haf_sum, haf_val)
 
 
+@pytest.mark.parametrize("atol", [1e-08, 1e-07, 1e-06])
+def test_check_probabilities(atol):
+    """Tests for check_probabilities"""
+    probs = np.random.rand(10) + 0j
+    assert check_probabilities(probs, atol)
+    probs[0] = -1
+    assert not check_probabilities(probs, atol)
+    probs[0] = -0.1 * atol * np.sum(probs[0:])
+    assert check_probabilities(probs, atol)
+    probs[0] = -10 * atol * np.sum(probs[0:])
+    assert not check_probabilities(probs, atol)
+    probs = np.random.rand(10) + 0j
+    probs[0] = -1j
+    assert not check_probabilities(probs, atol)
+    probs[0] = -1j * 0.1 * atol * np.sum(probs[0:])
+    assert check_probabilities(probs, atol)
+    probs[0] = -1j * 10 * atol * np.sum(probs[0:])
+    assert not check_probabilities(probs, atol)
+
+
+@pytest.mark.parametrize("cutoff", [43, 44])
+@pytest.mark.parametrize("method", ["recursive", "non-recursive", "diagonals"])
+def test_warning_non_recursive_gives_negative_probs(cutoff, method):
+    """"""
+    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
+    if method == "recursive":
+        with pytest.warns(
+            UserWarning,
+            match="Some of the diagonal elements of the density matrix are significantly",
+        ):
+            result = density_matrix_single_mode(
+                cov, {1: m1, 2: m2}, cutoff=cutoff, normalize=False, method=method
+            )
+        assert not check_probabilities(np.diag(result))
+    else:
+        result = density_matrix_single_mode(
+            cov, {1: m1, 2: m2}, cutoff=cutoff, normalize=False, method=method
+        )
+        assert check_probabilities(np.diag(result))
+
+
 @pytest.mark.parametrize("cutoff", [8, 9])
 @pytest.mark.parametrize("method", ["non-recursive"])
 def test_density_matrix(cutoff, method):