Adds functions for calculating the total photon distribution of k lossy squeezed states

thewalrus/quantum/__init__.py

@@ -110,6 +110,9 @@
 .. autosummary::
 
     pure_state_distribution
+    total_photon_number_distribution
+    characteristic_function
+
 
 Details
 ^^^^^^^
@@ -168,6 +171,8 @@
 
 from .photon_number_distributions import (
     pure_state_distribution,
+    total_photon_number_distribution,
+    characteristic_function,
 )
 
 __all__ = [

thewalrus/quantum/fock_tensors.py

@@ -608,7 +608,7 @@ def tvd_cutoff_bounds(mu, cov, cutoff, hbar=2, check_is_valid_cov=True, rtol=1e-
     and the user provided cutoff.
 
     For the derivation see Appendix B of `'Exact simulation of Gaussian boson sampling in polynomial space and exponential time',
-    Quesada and Arrazola et al. <10.1103/PhysRevResearch.2.023005>`_.
+    Quesada and Arrazola <https://journals.aps.org/prresearch/abstract/10.1103/PhysRevResearch.2.023005>`_.
 
     Args:
         mu (array): vector of means of the Gaussian state

thewalrus/quantum/photon_number_distributions.py

@@ -17,6 +17,7 @@
 
 import numpy as np
 from scipy.stats import nbinom
+from scipy.special import binom, hyp2f1
 
 from .conversions import Amat
 from .gaussian_checks import is_pure_cov
@@ -89,3 +90,91 @@ def _convolve_squeezed_state_distribution(s, cutoff=50, padding_factor=2):
     for s_val in s:
         ps = np.convolve(ps, _squeezed_state_distribution(s_val, cutoff_sc))[0:cutoff_sc]
     return ps
+
+
+def total_photon_number_distribution(n, k, s, eta, pref=1.0):
+    r"""Probability of observing a total of :math:`n` photons when :math:`k` identical
+    single-mode squeezed vacua with squeezing parameter :math:`s` undergo loss by transmission :math:`\eta`.
+
+    For the derivation see Appendix E of `'Quantum Computational Supremacy via High-Dimensional Gaussian Boson Sampling',
+    Deshpande et al. <https://arxiv.org/abs/2102.12474>`_.
+
+
+    Args:
+        n (int): number of photons
+        k (int): number of squeezed modes
+        s (float): squeezing parameter
+        eta (float): transmission parameter, between 0 and 1 inclusive
+        pref (float): use to return the probability times ``pref**n``
+    Returns:
+        (float): probability of observing a total of ``n`` photons or the probability times ``pref ** n``.
+    """
+    if n % 2 == 0:
+        peven = (
+            binom(-1 + k / 2.0 + n / 2.0, n / 2.0)
+            * hyp2f1(0.5 + n / 2.0, k / 2.0 + n / 2.0, 0.5, (1 - eta) ** 2 * np.tanh(s) ** 2)
+            * (1 / np.cosh(s)) ** k
+            * (pref * eta * np.tanh(s)) ** n
+        )
+        return peven
+
+    podd = (
+        (1 + n)
+        * (1 - eta)
+        * binom((-1 + k + n) / 2.0, (1 + n) / 2.0)
+        * hyp2f1((2 + n) / 2.0, (1 + k + n) / 2.0, 1.5, (1 - eta) ** 2 * np.tanh(s) ** 2)
+        * (1 / np.cosh(s)) ** k
+        * np.tanh(s)
+        * (pref * eta * np.tanh(s)) ** n
+    )
+    return podd
+
+
+def characteristic_function(
+    k, s, eta, mu, max_iter=10000, delta=1e-14, poly_corr=None
+):  # pylint: disable=too-many-arguments
+    r"""Calculates the expectation value of the characteristic function
+    :math:`\langle n^m \exp(mu n) \rangle` where :math:`n` is the total photon number of :math:`k` identical
+    single-mode squeezed vacua with squeezing parameter :math:`s` undergoing loss by
+    transmission :math:`\eta`.
+
+    Args:
+        k (int): number of squeezed modes
+        s (float): squeezing parameter
+        eta (float): transmission parameter, between 0 and 1 inclusive
+        mu (float): value at which to evaluate the characteristic function
+        max_iter (int): maximum number of terms allowed in the sum
+        delta (float): fractional change in the sum after which the sum is stopped
+        poly_corr (int): give the value of the exponent :math:`m` of the polynomial correction
+
+    Returns:
+        (float): the expected value of the moment generation function
+    """
+
+    if poly_corr is None or poly_corr == 0:
+        f = lambda x: 1
+    else:
+        f = lambda x: x ** poly_corr
+
+    if s == 0 or eta == 0:
+        return f(0)
+
+    pref = np.exp(mu)
+    tot_sum = f(0) * total_photon_number_distribution(0, k, s, eta, pref=pref)
+    converged = False
+
+    i = 1
+    prev_addend = tot_sum
+    while converged is False and i < max_iter:
+        old_tot_sum = tot_sum
+        addend = f(i) * total_photon_number_distribution(i, k, s, eta, pref=pref)
+        tot_sum += addend
+        i += 1
+        # Note that we check that the sum of the last *two* values does not change the net
+        # sum much, this is because for eta=0 the distrobution does not have support over
+        # the odd integers.
+        ratio = (addend + prev_addend) / old_tot_sum
+        if ratio < delta:
+            converged = True
+        prev_addend = addend
+    return tot_sum

thewalrus/tests/test_quantum.py

@@ -70,9 +70,10 @@
     variance_clicks,
     mean_clicks,
     tvd_cutoff_bounds,
+    total_photon_number_distribution,
+    characteristic_function,
 )
 
-
 @pytest.mark.parametrize("n", [0, 1, 2])
 def test_reduced_gaussian(n):
     """test that reduced gaussian returns the correct result"""
@@ -1521,3 +1522,48 @@ def test_big_displaced_squeezed_matelem():
     cutoff = 100
     probs_displaced_squeezed = probabilities(mu, cov, cutoff, hbar=1.0)
     assert np.allclose(np.sum(probs_displaced_squeezed), 1.0)
+
+@pytest.mark.parametrize("s", [0.5, 0.7, 0.9, 1.1])
+@pytest.mark.parametrize("k", [4, 6, 10, 12])
+def test_total_photon_number_distribution_values(s, k):
+    """Test that the total photon number distribution is correct when there are no losses"""
+    cutoff = 300
+    eta = 1.0
+    expected_probs = _squeezed_state_distribution(s, cutoff, N=k)
+    probs = np.array([total_photon_number_distribution(i, k, s, eta) for i in range(cutoff)])
+    assert np.allclose(expected_probs, probs)
+
+
+@pytest.mark.parametrize("s", [0.5, 0.7, 0.9, 1.1])
+@pytest.mark.parametrize("k", [4, 6, 10, 12])
+@pytest.mark.parametrize("eta", [0, 0.1, 0.5, 1])
+def test_total_photon_number_distribution_moments(s, k, eta):
+    """Test that the total photon number distribution has the correct mean and variance"""
+    expectation_n = characteristic_function(s=s, k=k, eta=eta, poly_corr=1, mu=0)
+    expectation_n2 = characteristic_function(s=s, k=k, eta=eta, poly_corr=2, mu=0)
+    var_n = expectation_n2 - expectation_n ** 2
+    expected_n = eta * k * np.sinh(s) ** 2
+    expected_var = expected_n * (1 + eta * (1 + 2 * np.sinh(s) ** 2))
+    assert np.allclose(expectation_n, expected_n)
+    assert np.allclose(expected_var, var_n)
+
+
+@pytest.mark.parametrize("s", [0.5, 0.7, 0.9, 1.1])
+@pytest.mark.parametrize("k", [4, 6, 10, 12])
+@pytest.mark.parametrize("eta", [0, 0.1, 0.5, 1])
+def test_characteristic_function_is_normalized(s, k, eta):
+    r"""Check that when evaluated at \mu=0 the characteristic function gives 1"""
+    val = characteristic_function(k=k, s=s, eta=eta, mu=0)
+    assert np.allclose(val, 1.0)
+
+
+@pytest.mark.parametrize("s", [0.5, 0.6, 0.7, 0.8])
+@pytest.mark.parametrize("k", [4, 6, 10, 12])
+def test_charactetistic_function_no_loss(s, k):
+    """Check the values of the characteristic function when there is no loss"""
+    mu = mu = 0.5 * np.log(2)
+    # Note that s must be less than np.arctanh(np.sqrt(0.5)) ~= 0.88
+    p = np.tanh(s) ** 2
+    val = characteristic_function(k=k, s=s, eta=1.0, mu=mu)
+    expected = ((1 - p) / (1 - 2 * p)) ** (k / 2)
+    assert np.allclose(val, expected)