# Source code for thewalrus._permanent

# Copyright 2019 Xanadu Quantum Technologies Inc.

# 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
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
r"""
This submodule provides access to tools for finding the permanent of a matrix. The algorithms implemented
here was first derived in

* Ryser, Herbert John (1963).
Combinatorial Mathematics, The Carus Mathematical Monographs, Vol. 14, Mathematical Association of America.
* Glynn, David G.
(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, prange

from scipy.special import factorial

from ._hafnian import hafnian_repeated, find_kept_edges

[docs]def perm(A, method="bbfg"):
"""Returns the permanent of a matrix using various methods.

Args:
A (array[float or complex]): a square array.
method (string): Set this to "ryser" to use the
Ryser formula
<https://en.wikipedia.org/wiki/Computing_the_permanent#Ryser_formula>_,
or "bbfg" to use the
BBFG formula
<https://en.wikipedia.org/wiki/Computing_the_permanent#Balasubramanian%E2%80%93Bax%E2%80%93Franklin%E2%80%93Glynn_formula>_.

Returns:
float or complex: the permanent of matrix A
"""

if not isinstance(A, np.ndarray):
raise TypeError("Input matrix must be a NumPy array.")

matshape = A.shape

if matshape != matshape:
raise ValueError("Input matrix must be square.")

if np.isnan(A).any():
raise ValueError("Input matrix must not contain NaNs.")

if matshape == 0:
return A.dtype.type(1.0)

if matshape == 1:
return A[0, 0]

if matshape == 2:
return A[0, 0] * A[1, 1] + A[0, 1] * A[1, 0]

if matshape == 3:
return (
A[0, 2] * A[1, 1] * A[2, 0]
+ A[0, 1] * A[1, 2] * A[2, 0]
+ A[0, 2] * A[1, 0] * A[2, 1]
+ A[0, 0] * A[1, 2] * A[2, 1]
+ A[0, 1] * A[1, 0] * A[2, 2]
+ A[0, 0] * A[1, 1] * A[2, 2]
)

isRyser = bool(method != "bbfg")

return perm_ryser(A) if isRyser else perm_bbfg(A)

@jit(nopython=True)
def perm_ryser(M):  # pragma: no cover
"""
Returns the permanent of a matrix using the Ryser formula in Gray ordering.

The code is an re-implementation from a Python 2 code found in
Permanent code golf
<https://codegolf.stackexchange.com/questions/97060/calculate-the-permanent-as-quickly-as-possible>_
using Numba.

Args:
M (array) : a square array.

Returns:
float or complex: the permanent of matrix M
"""
n = len(M)
if n == 0:
return M.dtype.type(1.0)
# row_comb keeps the sum of previous subsets.
# Every iteration, it removes a term and/or adds a new term
# to give the term to add for the next subset
row_comb = np.zeros((n), dtype=M.dtype)
total = 0
old_grey = 0
sign = +1
binary_power_dict = [2**i for i in range(n)]
num_loops = 2**n
for k in range(0, num_loops):
bin_index = (k + 1) % num_loops
reduced = np.prod(row_comb)
total += sign * reduced
new_grey = bin_index ^ (bin_index // 2)
grey_diff = old_grey ^ new_grey
grey_diff_index = binary_power_dict.index(grey_diff)
new_vector = M[grey_diff_index]
direction = (old_grey > new_grey) - (old_grey < new_grey)
for i in range(n):
row_comb[i] += new_vector[i] * direction
sign = -sign
old_grey = new_grey

@jit(nopython=True)
def perm_bbfg(M):  # pragma: no cover
"""
Returns the permanent of a matrix using the bbfg formula in Gray ordering

The code is a re-implementation from a Python 2 code found in
Permanent code golf
<https://codegolf.stackexchange.com/questions/97060/calculate-the-permanent-as-quickly-as-possible>_
using Numba.

Args:
M (array) : a square array.

Returns:
float or complex: the permanent of a matrix M
"""

n = len(M)
if n == 0:
return M.dtype.type(1.0)
row_comb = np.sum(M, 0)
total = 0
old_gray = 0
sign = +1
binary_power_dict = [2**i for i in range(n)]
num_loops = 2 ** (n - 1)
for bin_index in range(1, num_loops + 1):
reduced = np.prod(row_comb)
total += sign * reduced
new_gray = bin_index ^ (bin_index // 2)
gray_diff = old_gray ^ new_gray
gray_diff_index = binary_power_dict.index(gray_diff)
new_vector = M[gray_diff_index]
direction = 2 * ((old_gray > new_gray) - (old_gray < new_gray))
for i in range(n):
row_comb[i] += new_vector[i] * direction
sign = -sign
old_gray = new_gray

[docs]def permanent_repeated(A, rpt):
r"""Calculates the permanent of matrix :math:A, where the ith row/column
of :math:A is repeated :math:rpt_i times.

This function constructs the matrix

.. math:: B = \begin{bmatrix} 0 & A\\ A^T & 0 \end{bmatrix},

and then calculates :math:perm(A)=haf(B), by calling

>>> hafnian_repeated(B, rpt*2, loop=False)

Args:
A (array): matrix of size [N, N]
rpt (Sequence): sequence of N positive integers indicating the corresponding rows/columns
of A to be repeated.

Returns:
int or float or complex: the permanent of matrix A
"""
n = A.shape
O = np.zeros([n, n])
B = np.vstack([np.hstack([O, A]), np.hstack([A.T, O])])

return hafnian_repeated(B, rpt * 2, loop=False)

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

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)
Ay = A[kept_rows, :]
plusminus = (-1) ** ((m - len(kept_rows)) % 2)
total += plusminus * perm_bbfg(Ay.conj().T @ Ay + E)

[docs]@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
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)
Az = A[kept_rows, :]
plusminus = (-1) ** ((m - len(kept_rows)) % 2)
total += plusminus * perm_bbfg(Az.conj().T @ Az)

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:
raise ValueError("length of n must matrix number of input modes of T")
if len(d) != T.shape:
raise ValueError("length of d must match number of output modes of T")
if T.shape > T.shape:
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)

A = T[np.ix_(C, in_modes)]

# if matrix is unitary, use the Unitary Bristolian
E = np.eye(T.shape) - T.conj().T @ T
if np.allclose(E, np.zeros((T.shape, T.shape))):
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
`