Skip to content

Commit

Permalink
Adds the warning
Browse files Browse the repository at this point in the history
  • Loading branch information
Nicolas Quesada committed Apr 24, 2024
1 parent 5666166 commit 15e1292
Show file tree
Hide file tree
Showing 4 changed files with 121 additions and 13 deletions.
2 changes: 1 addition & 1 deletion thewalrus/_hafnian.py
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand Down
2 changes: 1 addition & 1 deletion thewalrus/internal_modes/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -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 *
56 changes: 46 additions & 10 deletions thewalrus/internal_modes/fock_density_matrices.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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)
Expand All @@ -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)
Expand Down Expand Up @@ -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):
"""
Convinience function for checking that the input is close enough to a probability distribution.
Args:
probs (array): probabilities to be tested.
atol (array): 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:
Expand All @@ -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", "direct" or "diagonals"
atol (float): value for raising warning when testing for valid probabilities
Returns:
array[complex]: (cutoff+1, cutoff+1) dimension density matrix
"""
Expand All @@ -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)) is False:
warnings.warn(
"Some of the diagonal elements of the density matrix are significantly negative or have significant imaginary parts",
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
Expand Down Expand Up @@ -224,8 +261,7 @@ 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)
Expand Down
74 changes: 73 additions & 1 deletion thewalrus/tests/test_internal_modes.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -1448,6 +1453,73 @@ 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):
Expand Down

0 comments on commit 15e1292

Please sign in to comment.