diff --git a/thewalrus/quantum/__init__.py b/thewalrus/quantum/__init__.py index 33ebfd48c..64a4f8e87 100644 --- a/thewalrus/quantum/__init__.py +++ b/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): diff --git a/thewalrus/quantum/fock_tensors.py b/thewalrus/quantum/fock_tensors.py index a335b7cd3..5f8ddf4b8 100644 --- a/thewalrus/quantum/fock_tensors.py +++ b/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 diff --git a/thewalrus/tests/test_quantum.py b/thewalrus/tests/test_quantum.py index ebeaef4b5..c61c9c151 100644 --- a/thewalrus/tests/test_quantum.py +++ b/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)