diff --git a/thewalrus/_hafnian.py b/thewalrus/_hafnian.py index 7330b88e..8a7d2019 100644 --- a/thewalrus/_hafnian.py +++ b/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: diff --git a/thewalrus/internal_modes/__init__.py b/thewalrus/internal_modes/__init__.py index 75601b36..9d3a26d5 100644 --- a/thewalrus/internal_modes/__init__.py +++ b/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 * diff --git a/thewalrus/internal_modes/fock_density_matrices.py b/thewalrus/internal_modes/fock_density_matrices.py index 592348aa..49c1ac53 100644 --- a/thewalrus/internal_modes/fock_density_matrices.py +++ b/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) @@ -141,16 +141,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: @@ -160,7 +188,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 """ @@ -187,7 +216,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) A = Amat(cov) diff --git a/thewalrus/tests/test_internal_modes.py b/thewalrus/tests/test_internal_modes.py index 496c25c8..cbb9928e 100644 --- a/thewalrus/tests/test_internal_modes.py +++ b/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):