# Source code for thewalrus.quantum.photon_number_distributions

# Copyright 2019-2020 Xanadu Quantum Technologies Inc.

# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at

# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""
Functions for calculating the photon number distributions of various states.
"""

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

[docs]
def pure_state_distribution(cov, cutoff=50, hbar=2, padding_factor=2):
r"""Calculates the total photon number distribution of a pure state
with zero mean.

Args:
cov (array): :math:2N\times 2N covariance matrix in xp-ordering
cutoff (int): Fock cutoff
tol (float): tolerance for determining if displacement is negligible
hbar (float): the value of :math:\hbar in the commutation
padding_factor (int): expanded size of the photon distribution to avoid accumulation of errors

Returns:
(array): Total photon number distribution
"""
if is_pure_cov(cov):
A = Amat(cov, hbar=hbar)
(n, _) = A.shape
N = n // 2
B = A[0:N, 0:N]
rs = np.arctanh(np.linalg.svd(B, compute_uv=False))
return _convolve_squeezed_state_distribution(
)[0:cutoff]
raise ValueError("The Gaussian state is not pure")

def _squeezed_state_distribution(s, cutoff=50, N=1):
"""Generate the photon number distribution of :math:N identical single mode squeezed states.

Args:
s (float): squeezing parameter
cutoff (int): Fock cutoff
N (float): number of squeezed states

Returns:
(array): Photon number distribution
"""
r = 0.5 * N
q = 1.0 - np.tanh(s) ** 2
N = cutoff // 2
ps_tot = np.zeros(cutoff)
if cutoff % 2 == 0:
ps = nbinom.pmf(np.arange(N), p=q, n=r)
ps_tot[0::2] = ps
else:
ps = nbinom.pmf(np.arange(N + 1), p=q, n=r)
ps_tot[0:-1][0::2] = ps[0:-1]
ps_tot[-1] = ps[-1]

return ps_tot

def _convolve_squeezed_state_distribution(s, cutoff=50, padding_factor=2):
"""Generates the total photon number distribution of single mode squeezed states with different squeezing values.

Args:
s (array): array of squeezing parameters
cutoff (int): Fock cutoff
padding_factor (int): expanded size of the photon distribution to avoid accumulation of errors
Returns:
(array[int]): total photon number distribution
"""
cutoff_sc = scale * cutoff
ps = np.zeros(cutoff_sc)
ps[0] = 1.0
for s_val in s:
ps = np.convolve(ps, _squeezed_state_distribution(s_val, cutoff_sc))[0:cutoff_sc]
return ps

[docs]
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

[docs]
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
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)