the Bristolian

.github/CHANGELOG.md

@@ -2,6 +2,8 @@
 
 ### New features
 
+* New functions for threshold detection probabilities of Fock states, the Bristolian (brs) and the Unitary Bristolian (ubrs) [#316](https://github.com/XanaduAI/thewalrus/pull/316)
+
 ### Improvements
 
 ### Bug fixes

thewalrus/__init__.py

@@ -118,7 +118,7 @@
     interferometer,
     grad_hermite_multidimensional,
 )
-from ._permanent import perm, permanent_repeated
+from ._permanent import perm, permanent_repeated, brs, ubrs
 
 from ._torontonian import (
     tor,

thewalrus/_permanent.py

@@ -21,9 +21,15 @@
   (2010), "The permanent of a square matrix", European Journal of Combinatorics, 31 (7): 1887-1891.
   <doi:10.1016/j.ejc.2010.01.010>`_
 """
+from itertools import chain
+
 import numpy as np
-from numba import jit
-from ._hafnian import hafnian_repeated
+from numba import jit, prange
+
+from scipy.special import factorial
+from scipy.linalg import sqrtm
+
+from ._hafnian import hafnian_repeated, find_kept_edges
 
 
 def perm(A, method="bbfg"):
@@ -178,3 +184,130 @@ def permanent_repeated(A, rpt):
     B = np.vstack([np.hstack([O, A]), np.hstack([A.T, O])])
 
     return hafnian_repeated(B, rpt * 2, loop=False)
+
+
+@jit(nopython=True, parallel=True)
+def brs(A, E):  # pragma: no cover
+    r"""
+    Calculates the Bristolian, a matrix function introduced for calculating the threshold detector
+    statistics on measurements of Fock states interfering in linear optical interferometers.
+
+    See the paper 'Threshold detector statistics of Bosonic states' for more detail (to be published soon)
+
+    Args:
+        A (array): matrix of size [m, n]
+        E (array): matrix of size [r, n]
+
+    Returns:
+        int or float or complex: the Bristol of matrices A and E
+    """
+    m = A.shape[0]
+
+    steps = 2 ** m
+    ones = np.ones(m, dtype=np.int8)
+    total = 0
+    for j in prange(steps):
+        kept_rows = np.where(find_kept_edges(j, ones) != 0)[0]
+        Ay = A[kept_rows, :]
+        plusminus = (-1) ** ((m - len(kept_rows)) % 2)
+        total += plusminus * perm_bbfg(Ay.conj().T @ Ay + E)
+    return total
+
+
+@jit(nopython=True, parallel=True)
+def ubrs(A):  # pragma: no cover
+    r"""
+    Calculates the Unitary Bristolian, a matrix function introduced for calculating the threshold detector
+    statistics on measurements of Fock states interfering in lossless linear optical interferometers.
+
+    See the paper 'Threshold detector statistics of Bosonic states' for more detail (to be published soon)
+
+    Args:
+        A (array): matrix of size [m, n]
+
+    Returns:
+        int or float or complex: the Unitary Bristol of matrix A
+    """
+    m = A.shape[0]
+    steps = 2 ** m
+    ones = np.ones(m, dtype=np.int8)
+    total = 0
+    for j in prange(steps):
+        kept_rows = np.where(find_kept_edges(j, ones) != 0)[0]
+        Az = A[kept_rows, :]
+        plusminus = (-1) ** ((m - len(kept_rows)) % 2)
+        total += plusminus * perm_bbfg(Az.conj().T @ Az)
+    return total
+
+
+def fock_prob(n, m, U):
+    r"""
+    Calculates the probability of a an input Fock state, n, scattering to an output Fock state, m, through
+    an interferometer described by matrix U.
+    The matrix U does not need to be unitary, but the total photon number at the input and the output must be equal.
+
+    Args:
+        n (sequence[int]): length-M list giving the input Fock state occupancy of each mode
+        m (sequence[int]): length-M list giving the output Fock state occupancy of each mode
+        U (array): M x M matrix describing the a linear optical transformation
+
+    Returns:
+        float: probability of Fock state, n, scattering to m, through an interferometer, U
+    """
+    assert sum(n) == sum(m)
+
+    in_modes = np.array(list(chain(*[[i] * j for i, j in enumerate(n) if j > 0])))
+    out_modes = np.array(list(chain(*[[i] * j for i, j in enumerate(m) if j > 0])))
+
+    Umn = U[np.ix_(out_modes, in_modes)]
+
+    n = np.array(n)
+    m = np.array(m)
+
+    return abs(perm(Umn)) ** 2 / (
+        np.prod(factorial(n), dtype=np.float64) * np.prod(factorial(m), dtype=np.float64)
+    )
+
+
+def fock_threshold_prob(n, d, T):
+    r"""
+    Calculates the probability of a an M_in mode input Fock state, n, scattering through an interferometer described by
+    T, being detected by M_out threshold detectors, with outcome given by M_out-length list, d.
+    T is an M_out x M_in matrix. It does not need to be unitary but M_out <= M_in.
+
+    Args:
+        n (sequence[int]): length-M_in list giving the input Fock state occupancy of each mode
+        d (sequence[int]): length-M_out list giving the outputs of threshold detectors
+        T (array): M_out x M_in matrix describing the a linear optical transformation, M_out <= M_in
+
+    Returns:
+        float: probability of Fock state, n, scattering through an interferometer, T, to give threshold detector outcome, d
+    """
+    n = np.array(n)
+    d = np.array(d)
+
+    assert len(n) == T.shape[1]
+    assert len(d) == T.shape[0]
+    assert T.shape[0] <= T.shape[1]
+
+    fac_prod = np.prod(factorial(n), dtype=np.float64)
+
+    in_modes = np.array(list(chain(*[[i] * j for i, j in enumerate(n) if j > 0])))
+    C = np.where(d > 0)[0]
+
+    A = T[np.ix_(C, in_modes)]
+
+    # if matrix is unitary, use the Unitary Bristolian
+    R2 = np.eye(T.shape[1]) - T.conj().T @ T
+    if np.allclose(R2, np.zeros((T.shape[1], T.shape[1]))):
+        U_dn = T[np.ix_(C, in_modes)]
+        return ubrs(U_dn).real / fac_prod
+
+    # Use the Bristolian for nonunitary transformations
+    if max(n) > 1:
+        R = sqrtm(R2)[:, in_modes]
+        E = R.conj().T @ R
+    else:
+        E = R2[np.ix_(in_modes, in_modes)]
+
+    return brs(A, E).real / fac_prod

thewalrus/tests/test_permanent.py

@@ -13,12 +13,19 @@
 # limitations under the License.
 """Tests for the Python permanent wrapper function"""
 # pylint: disable=no-self-use
+
+from itertools import chain, product
+
 import pytest
 
 import numpy as np
+
 from scipy.special import factorial as fac
+from scipy.linalg import sqrtm
+from scipy.stats import unitary_group
 
-from thewalrus import perm, permanent_repeated
+from thewalrus import perm, permanent_repeated, brs, ubrs
+from thewalrus._permanent import fock_prob, fock_threshold_prob
 
 perm_real = perm
 perm_complex = perm
@@ -161,3 +168,153 @@ def test_ones(self, n):
         A = np.array([[1]])
         p = permanent_repeated(A, [n])
         assert np.allclose(p, fac(n))
+
+
+def test_brs_HOM():
+    """HOM test"""
+
+    U = np.array([[1, 1], [1, -1]]) / np.sqrt(2)
+
+    n = [1, 1]
+    d = [1, 1]
+
+    assert np.isclose(fock_threshold_prob(n, d, U), fock_prob(n, d, U))
+
+    d = [1, 0]
+    m = [2, 0]
+
+    assert np.isclose(fock_threshold_prob(n, d, U), fock_prob(n, m, U))
+
+
+@pytest.mark.parametrize("eta", [0.2, 0.5, 0.9, 1])
+def test_brs_HOM_lossy(eta):
+    """lossy HOM dip test"""
+    T = np.sqrt(eta / 2) * np.array([[1, 1], [1, -1]])
+
+    n = [1, 1]
+    d = [1, 1]
+
+    assert np.isclose(fock_prob(n, d, T), fock_threshold_prob(n, d, T))
+
+
+def test_brs_ZTL():
+    """test 3-mode ZTL suppression"""
+
+    U = np.fft.fft(np.eye(3)) / np.sqrt(3)
+
+    n = [1, 1, 1]
+    d = [1, 1, 0]
+
+    p1 = fock_threshold_prob(n, d, U)
+    p2 = fock_prob(n, [1, 2, 0], U) + fock_prob(n, [2, 1, 0], U)
+    assert np.isclose(p1, p2)
+
+    n = [1, 1, 1]
+    d = [1, 1, 1]
+
+    p1 = fock_threshold_prob(n, d, U)
+    p2 = fock_prob(n, d, U)
+    assert np.isclose(p1, p2)
+
+    T = U[:2, :]
+    d = [1, 1]
+
+    p1 = fock_threshold_prob(n, d, T)
+    p2 = fock_prob(n, [1, 1, 1], U)
+
+    assert np.isclose(p1, p2)
+
+    d = [1, 0, 0]
+
+    p1 = fock_threshold_prob(n, d, U)
+    p2 = fock_prob(n, [3, 0, 0], U)
+
+    assert np.isclose(p1, p2)
+
+    n = [1, 2, 0]
+    d = [0, 1, 1]
+
+    p1 = fock_threshold_prob(n, d, U)
+    p2 = fock_prob(n, [0, 2, 1], U) + fock_prob(n, [0, 1, 2], U)
+
+    assert np.isclose(p1, p2)
+
+
+@pytest.mark.parametrize("eta", [0.2, 0.5, 0.9, 1])
+def test_brs_ZTL_lossy(eta):
+    """test lossy 3-mode ZTL suppression"""
+    T = np.sqrt(eta) * np.fft.fft(np.eye(3)) / np.sqrt(3)
+
+    n = [1, 1, 1]
+    d = [1, 1, 0]
+
+    p1 = eta ** 2 * (1 - eta) / 3
+    p2 = fock_threshold_prob(n, d, T)
+
+    assert np.allclose(p1, p2)
+
+
+@pytest.mark.parametrize("d", [[1, 1, 1], [1, 1, 0], [1, 0, 0]])
+def test_brs_ubrs(d):
+    """test that brs and ubrs give same results for unitary transformation"""
+
+    U = np.fft.fft(np.eye(3)) / np.sqrt(3)
+
+    n = np.array([2, 1, 0])
+    d = np.array(d)
+
+    in_modes = np.array(list(chain(*[[i] * j for i, j in enumerate(n) if j > 0])))
+    click_modes = np.where(d > 0)[0]
+
+    U_dn = U[np.ix_(click_modes, in_modes)]
+
+    b1 = ubrs(U_dn)
+
+    R = sqrtm(np.eye(U.shape[1]) - U.conj().T @ U)[:, in_modes]
+    E = R.conj().T @ R
+
+    b2 = brs(U_dn, E)
+
+    assert np.allclose(b1, b2)
+
+
+@pytest.mark.parametrize("M", range(2, 7))
+def test_brs_random(M):
+    """test that brs and per agree for random matices"""
+
+    n = np.ones(M, dtype=int)
+    n[np.random.randint(0, M)] = 0
+    d = np.ones(M, dtype=int)
+    d[np.random.randint(0, M)] = 0
+
+    loss_in = np.random.random(M)
+    loss_out = np.random.random(M)
+    U = unitary_group.rvs(M)
+    T = np.diag(loss_in) @ U @ np.diag(loss_out)
+
+    p1 = fock_threshold_prob(n, d, T)
+    p2 = fock_prob(n, d, T)
+
+    assert np.isclose(p1, p2)
+
+
+@pytest.mark.parametrize("M", range(2, 5))
+def test_brs_prob_normed(M):
+    """test that fock threshold probability is normalised"""
+
+    N = M + 1  # guarentee at least some bunching
+
+    in_modes = np.random.choice(np.arange(M), N)
+    n = np.bincount(in_modes, minlength=M)
+
+    loss_in = np.random.random(M)
+    loss_out = np.random.random(M)
+    U = unitary_group.rvs(M)
+    T = np.diag(loss_in) @ U @ np.diag(loss_out)
+
+    p_total = 0
+    for det_pattern in product([0, 1], repeat=M):
+        p = fock_threshold_prob(n, det_pattern, T)
+        p_total += p
+
+    assert np.isclose(p_total, 1)