# Source code for thewalrus._hermite_multidimensional

# 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
"""
Hermite Multidimensional Python interface
"""
from typing import Tuple, Generator, Iterable
from numba import jit
from numba.cpython.unsafe.tuple import tuple_setitem
import numpy as np

from ._hafnian import input_validation

# pylint: disable=too-many-arguments
[docs]def hermite_multidimensional(
R, cutoff, y=None, C=1, renorm=False, make_tensor=True, modified=False, rtol=1e-05, atol=1e-08
):
r"""Returns photon number statistics of a Gaussian state for a given
covariance matrix as described in *Multidimensional Hermite polynomials
and photon distribution for polymode mixed light*
arxiv:9308033 <https://arxiv.org/abs/hep-th/9308033>_.

Here :math:R is an :math:n \times n square matrix, and
:math:y is an :math:n dimensional vector. The polynomials :math:H_k^{(R)}(y) are
parametrized by the multi-index :math:k=(k_0,k_1,\ldots,k_{n-1}),
and are calculated for all values :math:0 \leq k_j < \text{cutoff},
thus a tensor of dimensions :math:\text{cutoff}^n is returned.

This tensor can either be flattened into a vector or returned as an actual
tensor with :math:n indices.

This implementation is based on the MATLAB code available at github
clementsw/gaussian-optics <https://github.com/clementsw/gaussian-optics>_.

.. note::

Note that if :math:R = (1) then :math:H_k^{(R)}(y)
are precisely the well known **probabilists' Hermite polynomials** :math:He_k(y),
and if :math:R = (2) then :math:H_k^{(R)}(y) are precisely the well known
**physicists' Hermite polynomials** :math:H_k(y).

Args:
R (array): square matrix parametrizing the Hermite polynomial family
cutoff (int): maximum size of the subindices in the Hermite polynomial
y (array): vector argument of the Hermite polynomial
C (complex): first value of the Hermite polynomials, the default value is 1
renorm (bool): If True, normalizes the returned multidimensional Hermite
polynomials such that :math:H_k^{(R)}(y)/\prod_i k_i!
make_tensor (bool): If False, returns a flattened one dimensional array
containing the values of the polynomial
modified (bool): whether to return the modified multidimensional Hermite polynomials or the standard ones
rtol (float): the relative tolerance parameter used in np.allclose
atol (float): the absolute tolerance parameter used in np.allclose
Returns:
(array): the multidimensional Hermite polynomials
"""

input_validation(R, atol=atol, rtol=rtol)
n, _ = R.shape

if (modified is False) and (y is not None):
m = y.shape
if m == n:
ym = R @ y
return hermite_multidimensional(
R, cutoff, y=ym, C=C, renorm=renorm, make_tensor=make_tensor, modified=True
)

if y is None:
y = np.zeros([n], dtype=complex)

m = y.shape
if m != n:
raise ValueError("The matrix R and vector y have incompatible dimensions")

num_indices = len(y)
# we want to catch np.ndarray(int) of ndim=0 which cannot be cast to tuple
if isinstance(cutoff, np.ndarray) and (cutoff.ndim == 0 or len(cutoff) == 1):
cutoff = int(cutoff)
if isinstance(cutoff, Iterable):
cutoffs = tuple(cutoff)
else:
cutoffs = tuple([cutoff]) * num_indices

Rt = np.real_if_close(R)
yt = np.real_if_close(y)

dtype = np.find_common_type([Rt.dtype.name, yt.dtype.name], [np.array(C).dtype.name])
array = np.zeros(cutoffs, dtype=dtype)
array[(0,) * num_indices] = C

if renorm:
values = np.array(_hermite_multidimensional_renorm(Rt, yt, array))
else:
values = np.array(_hermite_multidimensional(Rt, yt, array))

if not make_tensor:
values = values.flatten()

return values

def interferometer(R, cutoff, C=1, renorm=True, make_tensor=True, rtol=1e-05, atol=1e-08):
r"""Returns the matrix elements of an interferometer parametrized in terms of its R matrix.

Here :math:R is an :math:n \times n square matrix. The polynomials are
parametrized by the multi-index :math:k=(k_0,k_1,\ldots,k_{n-1}),
and are calculated for all values :math:0 \leq k_j < \text{cutoff},
thus a tensor of dimensions :math:\text{cutoff}^n is returned.

This tensor can either be flattened into a vector or returned as an actual
tensor with :math:n indices.

.. note::

Note that interferometer uses the normalized multidimensional Hermite polynomials.

Args:
R (array): square matrix parametrizing the Hermite polynomial family
cutoff (int): maximum size of the subindices in the Hermite polynomial
C (complex): first value of the Hermite polynomials, the default value is 1
renorm (bool): If True, normalizes the returned multidimensional Hermite
polynomials such that :math:H_k^{(R)}(y)/\prod_i k_i!
make_tensor (bool): If False, returns a flattened one dimensional array
containing the values of the polynomial
rtol (float): the relative tolerance parameter used in np.allclose
atol (float): the absolute tolerance parameter used in np.allclose
Returns:
(array): the multidimensional Hermite polynomials
"""

input_validation(R, atol=atol, rtol=rtol)
n, num_indices = R.shape

# we want to catch np.ndarray(int) of ndim=0 which cannot be cast to tuple
if isinstance(cutoff, np.ndarray) and (cutoff.ndim == 0 or len(cutoff) == 1):
cutoff = int(cutoff)
if isinstance(cutoff, Iterable):
cutoffs = tuple(cutoff)
else:
cutoffs = tuple([cutoff]) * num_indices

Rt = np.real_if_close(R)

dtype = np.find_common_type([Rt.dtype.name], [np.array(C).dtype.name])
array = np.zeros(cutoffs, dtype=dtype)
array[(0,) * num_indices] = C

if renorm:
values = np.array(_interferometer_renorm(Rt, array))
else:
values = np.array(_interferometer(Rt, array))

if make_tensor:
shape = cutoff * np.ones([n], dtype=int)
values = np.reshape(values, shape)

return values

# pylint: disable=too-many-arguments
[docs]def hafnian_batched(A, cutoff, mu=None, rtol=1e-05, atol=1e-08, renorm=False, make_tensor=True):
r"""Calculates the hafnian of :func:reduction(A, k) <hafnian.reduction>
for all possible values of vector k below the specified cutoff.

Here,

* :math:A is am :math:n\times n square matrix
* :math:k is a vector of (non-negative) integers with the same dimensions as :math:A,
i.e., :math:k = (k_0,k_1,\ldots,k_{n-1}), and where :math:0 \leq k_j < \texttt{cutoff}.

The function :func:~.hafnian_repeated can be used to calculate the reduced hafnian
for a *specific* value of :math:k; see the documentation for more information.

.. note::

If mu is not None, this function instead returns
hafnian(np.fill_diagonal(reduction(A, k), reduction(mu, k)), loop=True).
This calculation can only be performed if the matrix :math:A is invertible.

Args:
A (array): a square, symmetric :math:N\times N array.
cutoff (int): maximum size of the subindices in the Hermite polynomial
mu (array): a vector of length :math:N representing the vector of means/displacement
renorm (bool): If True, the returned hafnians are *normalized*, that is,
:math:haf(reduction(A, k))/\ \sqrt{prod_i k_i!}
(or :math:lhaf(fill\_diagonal(reduction(A, k),reduction(mu, k))) if
mu is not None)
make_tensor: If False, returns a flattened one dimensional array instead
of a tensor with the values of the hafnians.
rtol (float): the relative tolerance parameter used in np.allclose.
atol (float): the absolute tolerance parameter used in np.allclose.
Returns:
(array): the values of the hafnians for each value of :math:k up to the cutoff
"""
input_validation(A, atol=atol, rtol=rtol)
n, _ = A.shape

if mu is None:
mu = np.zeros([n], dtype=complex)
# The minus signs are intentional and are due to the fact that
# Dodonov et al. in PRA 50, 813 (1994) use (p,q) ordering instead of (q,p) ordering
return hermite_multidimensional(
-A, cutoff, y=mu, renorm=renorm, make_tensor=make_tensor, modified=True
)

@jit(nopython=True)
def dec(tup: Tuple[int], i: int) -> Tuple[int, ...]:  # pragma: no cover
r"""returns a copy of the given tuple of integers where the ith element has been decreased by 1
Args:
tup (Tuple[int]): the given tuple
i (int): the position of the element to be decreased
Returns:
Tuple[int,...]: the new tuple with the decrease on i-th element by 1
"""
copy = tup[:]
return tuple_setitem(copy, i, tup[i] - 1)

@jit(nopython=True)
def remove(
pattern: Tuple[int, ...]
) -> Generator[Tuple[int, Tuple[int, ...]], None, None]:  # pragma: no cover
r"""returns a generator for all the possible ways to decrease elements of the given tuple by 1
Args:
pattern (Tuple[int, ...]): the pattern given to be decreased
Returns:
Generator[Tuple[int, Tuple[int, ...]], None, None]: the generator
"""
for p, n in enumerate(pattern):
if n > 0:
yield p, dec(pattern, p)

SQRT = np.sqrt(np.arange(1000))  # saving the time to recompute square roots

@jit(nopython=True)
def _hermite_multidimensional_renorm(R, y, G):  # pragma: no cover
r"""Numba-compiled function to fill an array with the Hermite polynomials. It expects an array
initialized with zeros everywhere except at index (0,...,0) (i.e. the seed value).

Args:
R (array[complex]): square matrix parametrizing the Hermite polynomial
y (vector[complex]): vector argument of the Hermite polynomial
G (array[complex]): array to be filled with the Hermite polynomials

Returns:
array[complex]: the multidimensional Hermite polynomials
"""
indices = np.ndindex(G.shape)
next(indices)  # skip the first index (0,...,0)
for idx in indices:
i = 0
for i, val in enumerate(idx):
if val > 0:
break
ki = dec(idx, i)
u = y[i] * G[ki]
for l, kl in remove(ki):
u -= SQRT[ki[l]] * R[i, l] * G[kl]
G[idx] = u / SQRT[idx[i]]
return G

@jit(nopython=True)
def _hermite_multidimensional(R, y, G):  # pragma: no cover
r"""Numba-compiled function to fill an array with the Hermite polynomials. It expects an array
initialized with zeros everywhere except at index (0,...,0) (i.e. the seed value).

Args:
R (array[complex]): square matrix parametrizing the Hermite polynomial
y (vector[complex]): vector argument of the Hermite polynomial
G (array[complex]): array to be filled with the Hermite polynomials

Returns:
array[complex]: the multidimensional Hermite polynomials
"""
indices = np.ndindex(G.shape)
next(indices)  # skip the first index (0,...,0)
for idx in indices:
i = 0
for i, val in enumerate(idx):
if val > 0:
break
ki = dec(idx, i)
u = y[i] * G[ki]
for l, kl in remove(ki):
u -= ki[l] * R[i, l] * G[kl]
G[idx] = u
return G

@jit(nopython=True)
def _interferometer_renorm(R, G):  # pragma: no cover
r"""Numba-compiled function returning the matrix elements of an interferometer
parametrized in terms of its R matrix

Args:
R (array[complex]): square matrix parametrizing the Hermite polynomial
array (array[complex]): array to be filled with the Hermite polynomials

Returns:
array[complex]: the multidimensional Hermite polynomials
"""
dim, _ = R.shape
num_modes = dim / 2

indices = np.ndindex(G.shape)
next(indices)  # skip the first index (0,...,0)
for idx in indices:
bran = 0
for ii in range(0, num_modes):
bran += idx[ii]

ketn = 0
for ii in range(num_modes, dim):
ketn += idx[ii]

if bran == ketn:
i = 0
for i, val in enumerate(idx):
if val > 0:
break
ki = dec(idx, i)
u = 0
for l, kl in remove(ki):
u -= SQRT[ki[l]] * R[i, l] * G[kl]
G[idx] = u / SQRT[idx[i]]

return G

@jit(nopython=True)
def _interferometer(R, G):  # pragma: no cover
r"""Numba-compiled function returning the matrix elements of an interferometer
parametrized in terms of its R matrix

Args:
R (array[complex]): square matrix parametrizing the Hermite polynomial
G (array[complex]): array to be filled with the Hermite polynomials

Returns:
array[complex]: the multidimensional Hermite polynomials
"""
dim, _ = R.shape
num_modes = dim / 2

indices = np.ndindex(G.shape)
next(indices)  # skip the first index (0,...,0)
for idx in indices:
bran = 0
for ii in range(0, num_modes):
bran += idx[ii]

ketn = 0
for ii in range(num_modes, dim):
ketn += idx[ii]

if bran == ketn:
i = 0
for i, val in enumerate(idx):
if val > 0:
break
ki = dec(idx, i)
u = 0
for l, kl in remove(ki):
u -= ki[l] * R[i, l] * G[kl]
G[idx] = u

return G

[docs]def grad_hermite_multidimensional(G, R, y, C=1, renorm=True, dtype=None):
# pylint: disable=too-many-arguments
r"""Calculates the gradients of the renormalized multidimensional Hermite polynomials :math:C*H_k^{(R)}(y) with respect to its parameters :math:C, :math:y and :math:R.

Args:
G (array): the multidimensional Hermite polynomials
R (array[complex]): square matrix parametrizing the Hermite polynomial
y (vector[complex]): vector argument of the Hermite polynomial
C (complex): first value of the Hermite polynomials
renorm (bool): If True, uses the normalized multidimensional Hermite
polynomials such that :math:H_k^{(R)}(y)/\prod_i k_i!
dtype (data type): Specifies the data type used for the calculation

Returns:
array[data type], array[data type], array[data type]: the gradients of the multidimensional Hermite polynomials with respect to C, R and y
"""
if dtype is None:
dtype = np.find_common_type(
[G.dtype.name, R.dtype.name, y.dtype.name], [np.array(C).dtype.name]
)
n, _ = R.shape
if y.shape != n:
raise ValueError(
f"The matrix R and vector y have incompatible dimensions ({R.shape} vs {y.shape})"
)
dG_dC = np.array(G / C).astype(dtype)
dG_dR = np.zeros(G.shape + R.shape, dtype=dtype)
dG_dy = np.zeros(G.shape + y.shape, dtype=dtype)
if renorm:
dG_dR, dG_dy = _grad_hermite_multidimensional_renorm(R, y, G, dG_dR, dG_dy)
else:
dG_dR, dG_dy = _grad_hermite_multidimensional(R, y, G, dG_dR, dG_dy)

return dG_dC, dG_dR, dG_dy

@jit(nopython=True)
def _grad_hermite_multidimensional_renorm(R, y, G, dG_dR, dG_dy):  # pragma: no cover
r"""
Numba-compiled function to fill two arrays (dG_dR, dG_dy) with the gradients of the renormalized multidimensional Hermite polynomials
with respect to its parameters :math:R and :math:y. It needs the array of the multidimensional Hermite polynomials.

Args:
R (array[complex]): square matrix parametrizing the Hermite polynomial
y (vector[complex]): vector argument of the Hermite polynomial
G (array[complex]): array of the multidimensional Hermite polynomials
dG_dR (array[complex]): array to be filled with the gradients of the renormalized multidimensional Hermite polynomials with respect to R
dG_dy (array[complex]): array to be filled with the gradients of the renormalized multidimensional Hermite polynomials with respect to y

Returns:
dG_dR[complex], dG_dy[complex]: the gradients of the renormalized multidimensional Hermite polynomials with respect to R and y
"""
indices = np.ndindex(G.shape)
next(indices)  # skip the first index (0,...,0)
for idx in indices:
i = 0
for i, val in enumerate(idx):
if val > 0:
break
ki = dec(idx, i)
dy = y[i] * dG_dy[ki]
dy[i] += G[ki]
dR = y[i] * dG_dR[ki]
for l, kl in remove(ki):
dy -= SQRT[ki[l]] * dG_dy[kl] * R[i, l]
dR -= SQRT[ki[l]] * R[i, l] * dG_dR[kl]
dR[i, l] -= SQRT[ki[l]] * G[kl]
dG_dR[idx] = dR / SQRT[idx[i]]
dG_dy[idx] = dy / SQRT[idx[i]]
return dG_dR, dG_dy

@jit(nopython=True)
def _grad_hermite_multidimensional(R, y, G, dG_dR, dG_dy):  # pragma: no cover
r"""
Numba-compiled function to fill two arrays (dG_dR, dG_dy) with the gradients of the renormalized multidimensional Hermite polynomials
with respect to its parameters :math:R and :math:y. It needs the array of the multidimensional Hermite polynomials.

Args:
R (array[complex]): square matrix parametrizing the Hermite polynomial
y (vector[complex]): vector argument of the Hermite polynomial
G (array[complex]): array of the multidimensional Hermite polynomials
dG_dR (array[complex]): array to be filled with the gradients of the renormalized multidimensional Hermite polynomials with respect to R
dG_dy (array[complex]): array to be filled with the gradients of the renormalized multidimensional Hermite polynomials with respect to y

Returns:
dG_dR[complex], dG_dy[complex]: the gradients of the renormalized multidimensional Hermite polynomials with respect to R and y
"""
indices = np.ndindex(G.shape)
next(indices)  # skip the first index (0,...,0)
for idx in indices:
i = 0
for i, val in enumerate(idx):
if val > 0:
break
ki = dec(idx, i)
dy = y[i] * dG_dy[ki]
dy[i] += G[ki]
dR = y[i] * dG_dR[ki]
for l, kl in remove(ki):
dy -= ki[l] * dG_dy[kl] * R[i, l]
dR -= ki[l] * R[i, l] * dG_dR[kl]
dR[i, l] -= ki[l] * G[kl]
dG_dR[idx] = dR
dG_dy[idx] = dy
return dG_dR, dG_dy


Getting started

Background

The Walrus API