Calculation of n-body marginals of a GBS distribution

thewalrus/quantum/__init__.py

@@ -56,6 +56,7 @@
     update_probabilities_with_noise
     find_classical_subsystem
     tvd_cutoff_bounds
+    n_body_marginals
 
 Adjacency matrices
 ^^^^^^^^^^^^^^^^^^
@@ -132,6 +133,7 @@
     update_probabilities_with_noise,
     find_classical_subsystem,
     tvd_cutoff_bounds,
+    n_body_marginals,
 )
 
 from .adjacency_matrices import (
@@ -210,6 +212,7 @@
     "mean_clicks",
     "variance_clicks",
     "pure_state_distribution",
+    "n_body_marginals",
 ]
 
 def deprecate(new_func):

thewalrus/quantum/fock_tensors.py

@@ -19,6 +19,8 @@
 
 from itertools import count, product, chain
 
+from collections import OrderedDict
+
 import numpy as np
 import dask
 
@@ -45,6 +47,7 @@
 )
 
 
+
 def pure_state_amplitude(mu, cov, i, include_prefactor=True, tol=1e-10, hbar=2, check_purity=True):
     r"""Returns the :math:`\langle i | \psi\rangle` element of the state ket
     of a Gaussian state defined by covariance matrix cov.
@@ -632,3 +635,62 @@ def tvd_cutoff_bounds(mu, cov, cutoff, hbar=2, check_is_valid_cov=True, rtol=1e-
         ps = np.real_if_close(np.diag(density_matrix(mu_red, cov_red, cutoff=cutoff, hbar=hbar)))
         bounds += 1 - np.cumsum(ps)
     return bounds
+
+def n_body_marginals(mean, cov, cutoff, n, hbar=2):
+    r"""Calculates the first n-body marginals of a Gaussian state.
+
+
+    For an M-mode Gaussian state there exists a photon number distribution with probability mass function
+    :math:`p[i_0,i_1,\ldots, i_{M-1}]`. The function ``n_body_marginals`` calculates the first n-body marginals
+    of the (all-mode) probability distribution :math:`p`. The :math:`n=1` marginals or single body marginals
+    are simply the probability that mode :math:`k` has :math:`i` photons, i.e. :math:`p_k[i]`.
+    For :math:`n=2` one obtains the two-body probabilities. For two modes :math:`k` and :math:`l` this is a
+    two dimensional probability distribution :math:`p_{k,l}[i,j]`.
+    Note that these marginals have interesting permutation properties, for example :math:`p_{k,l}[i,j] = p_{l,k}[j,i]`.
+
+    The function provided here takes advantage of these symmetries to minimize the amount of calculations.
+    The return of this function is a list of tensors where the first entry contains the one-body marginals of the :math:`M` modes
+    (giving a tensor of shape ``(M, cutoff)``), the second entry contains the two-body marginals (giving a tensor of shape ``(M,M,cutoff, cutoff)``)
+    and so on and so forth.
+
+    To be clever about not calculating things that can be obtained by permutations it checks whether the index vector representing the modes is sorted.
+    From the way ``itertools.product`` works we know that it will always produce a sorted index vector before generating any of its unordered permutations.
+    Thus whenever the index vector is ordered we perform the numerical calculation.
+
+    If it is an unsorted index vector it realizes, in the ``if`` statement, that it can be obtained by permuting the
+    marginal distribution of something that has already been calculated.
+
+    Args:
+        mean (array): length-:math:`2N` quadrature displacement vector
+        cov (array): length-:math:`2N` covariance matrix
+        cutoff (int): cutoff in Fock space
+        n (int): order of the correlations
+        hbar (float): the value of :math:`\hbar` in the commutation
+            relation :math:`[\x,\p]=i\hbar`.
+
+    Returns:
+        list[array]: List with arrays containing the :math:`1,..,n` body marginal
+            distributions of the modes
+    """
+    M = len(mean)
+    if (M, M) != cov.shape:
+        raise ValueError("The covariance matrix and vector of means have incompatible dimensions")
+    if M % 2 != 0:
+        raise ValueError("The vector of means is not of even dimensions")
+    M = M // 2
+    if M < n:
+        raise ValueError("The order of the correlations is higher than the number of modes")
+
+    marginal = [np.zeros(([M] * i) + ([cutoff] * i)) for i in range(1, n + 1)]
+
+    for ind in product(range(M), repeat=n):
+        modes = list(set(ind))
+        acc = len(modes) - 1
+        if list(ind) == sorted(ind):
+            sub_mean, sub_cov = reduced_state(mean, cov, modes)  # this happens in phase space
+            marginal[acc][tuple(modes)] = probabilities(sub_mean, sub_cov, cutoff, hbar=hbar)
+        else:
+            modes_usrt = list(OrderedDict.fromkeys(ind))
+            perm = np.argsort(modes_usrt)
+            marginal[acc][tuple(modes_usrt)] = marginal[acc][tuple(modes)].transpose(perm)
+    return marginal

thewalrus/tests/test_quantum.py

@@ -72,6 +72,7 @@
     tvd_cutoff_bounds,
     total_photon_number_distribution,
     characteristic_function,
+    n_body_marginals
 )
 
 @pytest.mark.parametrize("n", [0, 1, 2])
@@ -1631,3 +1632,89 @@ def test_charactetistic_function_no_loss(s, k):
     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)
+
+
+
+@pytest.mark.parametrize("r", [0.5, 0.7, 2])
+@pytest.mark.parametrize("phi", [0.5, 0.7, 2])
+def test_two_mode_squeezed_vacuum_marginals(r, phi):
+    """Check correct answer for two-mode squeezed vacuum"""
+    cov_tmsv = two_mode_squeezing(2 * r, phi)
+    means = np.zeros(len(cov_tmsv))
+    cutoff = 20
+    nbar = np.sinh(r) ** 2
+    ratio = nbar / (1 + nbar)
+    expected = (ratio ** np.arange(20)) * 1 / (1 + nbar)
+    result = n_body_marginals(means, cov_tmsv, cutoff, 2)
+    assert np.allclose(result[0][0], expected)
+    assert np.allclose(result[0][1], expected)
+    assert np.allclose(result[1][0, 1], np.diag(expected))
+    assert np.allclose(result[1][1, 0], np.diag(expected))
+    assert np.allclose(result[1][0, 0], 0)
+    assert np.allclose(result[1][1, 1], 0)
+
+
+def test_marginal_probs_are_normalized():
+    """Check that for an arbitrary states all the marginal are proper probabilities"""
+    cutoff = 10
+    r = 1.58
+    eta = 0.3
+    M = 5
+    N = 3
+    U = random_interferometer(M)
+    sq = np.array([r] * N + [0] * (M - N))
+    cov_in = eta * np.concatenate([np.exp(2 * sq), np.exp(-2 * sq)]) + (1 - eta)
+    O = interferometer(U)
+    cov = O @ np.diag(cov_in) @ O.T
+    means = np.zeros(len(cov))
+
+    result = n_body_marginals(means, cov, cutoff, 3)
+    for tensor in result:
+        assert np.alltrue(np.round(tensor, 14) >= 0.0)
+        assert np.alltrue(tensor <= 1.0)
+
+
+def test_correct_indexing_n_body_marginals():
+    """Check the indexing is correct by checking correct composition rules for a product state"""
+    r1 = np.arcsinh(1)
+    r2 = np.arcsinh(np.sqrt(2))
+    cov3 = squeezing([2 * r1, 0, 2 * r2])
+    marg = n_body_marginals(np.zeros(len(cov3)), cov3, 4, 3)
+    assert np.allclose(marg[1][2, 0][0, 2], marg[0][2][0] * marg[0][0][2])
+    assert np.allclose(marg[2][0, 1, 2][:, 0, :], marg[2][1, 0, 2][0, :, :])
+    assert np.allclose(marg[2][0, 1, 2][:, 0, :], marg[2][2, 0, 1][0, :, :].T)
+    assert np.allclose(marg[2][0, 1, 2][:, 1, :], 0)
+
+
+def test_n_body_marginals_mismatch_shape():
+    """Check the correct error is raised when shapes of means of cov do not match"""
+    r1 = np.arcsinh(1)
+    r2 = np.arcsinh(np.sqrt(2))
+    cov = squeezing([2 * r1, 0, 2 * r2])
+    mu = np.zeros([4])
+    with pytest.raises(
+        ValueError, match="The covariance matrix and vector of means have incompatible dimensions"
+    ):
+        n_body_marginals(mu, cov, 4, 3)
+
+
+def test_n_body_marginals_not_even_shape():
+    """Check the correct error is raised when shapes of means of cov are of odd size"""
+    r1 = np.arcsinh(1)
+    r2 = np.arcsinh(np.sqrt(2))
+    cov = squeezing([2 * r1, 0, 2 * r2])
+    mu = np.zeros([5])
+    with pytest.raises(ValueError, match="The vector of means is not of even dimensions"):
+        n_body_marginals(mu, cov[:5, :5], 4, 3)
+
+
+def test_n_body_marginals_too_high_correlation():
+    """Check the correct error is raised when the order of the correlation is larger than the number of modes"""
+    r1 = np.arcsinh(1)
+    r2 = np.arcsinh(np.sqrt(2))
+    cov = squeezing([2 * r1, 0, 2 * r2])
+    mu = np.zeros([6])
+    with pytest.raises(
+        ValueError, match="The order of the correlations is higher than the number of modes"
+    ):
+        n_body_marginals(mu, cov, 4, 4)