the Bristolian (#316)

* pasting in the functions, writing some docstrings@

* added tests

* running black

* update changelog

* pragma no cover on jitted functions

* attempt to make codefactor happy

* run black

* fix bug nico found and add new test

* remove unused variable

* remove sqrtm when input pattern is a bitstring

* fix codefactor

* new test for bunched inputs and normalised probability distributions

* add to autosummary

* change assert to raise value error and add tests

* add test docstring

* Adds contributors name

* remove permanent of empty matrix from bristolian loop

* simply E construction for bristolian

* Update thewalrus/_permanent.py

Co-authored-by: Nicolas Quesada <991946+nquesada@users.noreply.github.com>

* Passes black

Co-authored-by: Nicolas Quesada <nquesada@pop-os.localdomain>
Co-authored-by: Nicolas Quesada <991946+nquesada@users.noreply.github.com>

.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
 
 * Recursive Torontonian added for faster computation based on paper ["Polynomial speedup in Torontonian calculation by a scalable recursive algorithm" by Ágoston Kaposi, Zoltán Kolarovszki, Tamás Kozsik, Zoltán Zimborás, and Péter Rakyta](https://arxiv.org/pdf/2109.04528.pdf). [#321](https://github.com/XanaduAI/thewalrus/pull/321)
@@ -16,7 +18,7 @@
 
 This release contains contributions from (in alphabetical order):
 
-Gregory Morse
+Jake Bulmer, Gregory Morse
 
 ---
 

thewalrus/__init__.py

@@ -82,6 +82,8 @@
     tor
     perm
     permanent_repeated
+    brs
+    ubrs
     hermite_multidimensional
     hafnian_banded
     reduction
@@ -118,7 +120,7 @@
     interferometer,
     grad_hermite_multidimensional,
 )
-from ._permanent import perm, permanent_repeated
+from ._permanent import perm, permanent_repeated, brs, ubrs
 
 from ._torontonian import (
     tor,
@@ -141,6 +143,8 @@
     "tor",
     "perm",
     "permanent_repeated",
+    "brs",
+    "ubrs",
     "reduction",
     "hermite_multidimensional",
     "grad_hermite_multidimensional",

thewalrus/_permanent.py

@@ -21,9 +21,14 @@
   (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 ._hafnian import hafnian_repeated, find_kept_edges
 
 
 def perm(A, method="bbfg"):
@@ -188,3 +193,129 @@ 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(1, 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
+    """
+    if sum(n) != sum(m):
+        raise ValueError("number of input photons must equal number of output photons")
+
+    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)
+
+    if len(n) != T.shape[1]:
+        raise ValueError("length of n must matrix number of input modes of T")
+    if len(d) != T.shape[0]:
+        raise ValueError("length of d must match number of output modes of T")
+    if T.shape[0] > T.shape[1]:
+        raise ValueError("number of output modes cannot be larger than number of input modes")
+
+    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
+    E = np.eye(T.shape[1]) - T.conj().T @ T
+    if np.allclose(E, np.zeros((T.shape[1], T.shape[1]))):
+        U_dn = T[np.ix_(C, in_modes)]
+        return ubrs(U_dn).real / fac_prod
+
+    E_n = E[np.ix_(in_modes, in_modes)]
+
+    return brs(A, E_n).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
@@ -181,3 +188,185 @@ 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)
+
+
+def test_fock_thresh_valueerror():
+    """test that input checks are raised"""
+    with pytest.raises(ValueError):
+        n = [1, 1, 1]
+        T = np.ones((2, 2))
+        d = [1, 1]
+        fock_threshold_prob(n, d, T)
+
+    with pytest.raises(ValueError):
+        n = [1, 1]
+        d = [1, 1, 1]
+        T = np.ones((2, 2))
+        fock_threshold_prob(n, d, T)
+
+    with pytest.raises(ValueError):
+        n = [1, 1]
+        d = [1, 1, 1]
+        T = np.ones((3, 2))
+        fock_threshold_prob(n, d, T)
+
+
+def test_fock_prob_valueerror():
+    """test that input checks are raised"""
+    with pytest.raises(ValueError):
+        n = [1, 1, 2, 1]
+        m = [1, 1, 1, 3]
+
+        U = np.eye((4))
+
+        fock_prob(n, m, U)