Cython Lattice

beamsplitter	Calculates the Fock representation of the beamsplitter.
beamsplitter_batched	Calculates the Fock representation of the batched beamsplitter.
displacement	Calculates the matrix elements of the displacement gate using a recurrence relation.
displacement_batched	Calculates the matrix elements of the displacement gate using a recurrence relation.
homodyne_projector	Fast homodyne projector.
homodyne_projector_batched	Fast homodyne projector for batched inputs.
jacobian_displacement	Calculates the jacobian of the displacement gate with respect to the complex displacement alpha and its conjugate.
squeezed	Calculates the matrix elements of the single-mode squeezed state using recurrence relations.
squeezed_batched	Calculates the matrix elements of the single-mode squeezed state using recurrence relations.
squeezer	Calculates the matrix elements of the squeezing gate using a recurrence relation.
squeezer_batched	Calculates the matrix elements of the squeezing gate using a recurrence relation.
vanilla	Vanilla algorithm for calculating the fock representation of a Gaussian tensor.
vanilla_batched	The batched implementation of vanilla.
vanilla_vjp	Vanilla vjp function.
vanilla_vjp_batched	Batched implementation of vanilla_vjp.

beamsplitter

mrmustard.mathlib.cython_lattice.beamsplitter(shape, theta, phi, stable=False)

Calculates the Fock representation of the beamsplitter. It takes advantage of input-output particle conservation (m+n=p+q) to avoid one for loop. Inspired from the original implementation in the walrus by @ziofil. Here is how the parameters are used in the code (see eq. 73-75 in https://arxiv.org/abs/2004.11002):

\[ \begin{align}\begin{aligned}\begin{split}A = \begin{bmatrix} 0 & V \\ V^T & 0 \end{bmatrix} \quad \text{(BS Bargmann matrix)}\end{split}\\\begin{split}V = \begin{bmatrix} ct & -st e^{-i\phi} \\ st e^{i\phi} & ct \end{bmatrix} \quad \text{(BS unitary)}\end{split}\end{aligned}\end{align} \]

Parameters:

• shape – The Fock shape of the resulting array.

• theta – The beamsplitter angle.

• phi – The beamsplitter phase.

• stable – Whether to use the stable algorithm.

Returns:

The matrix representing the beamsplitter operation.

beamsplitter_batched

mrmustard.mathlib.cython_lattice.beamsplitter_batched(shape, theta, phi, stable=False)

Calculates the Fock representation of the batched beamsplitter. It takes advantage of input-output particle conservation (m+n=p+q) to avoid one for loop. Inspired from the original implementation in the walrus by @ziofil. Here is how the parameters are used in the code (see eq. 73-75 in https://arxiv.org/abs/2004.11002):

\[ \begin{align}\begin{aligned}\begin{split}A = \begin{bmatrix} 0 & V \\ V^T & 0 \end{bmatrix} \quad \text{(BS Bargmann matrix)}\end{split}\\\begin{split}V = \begin{bmatrix} ct & -st e^{-i\phi} \\ st e^{i\phi} & ct \end{bmatrix} \quad \text{(BS unitary)}\end{split}\end{aligned}\end{align} \]

Note: batch dimensions must be flattened.

Parameters:

• shape – The Fock shape of the resulting array.

• theta – The batched beamsplitter angles.

• phi – The batched beamsplitter phases.

• stable – Whether to use the stable algorithm.

Returns:

The batched matrices representing the beamsplitter operation.

displacement

mrmustard.mathlib.cython_lattice.displacement(cutoffs, alpha)

Calculates the matrix elements of the displacement gate using a recurrence relation.

Parameters:

• cutoffs – The Fock ladder output-input cutoffs.

• alpha – The displacement magnitude and angle.

Returns:

The matrix representing the displacement operation.

displacement_batched

mrmustard.mathlib.cython_lattice.displacement_batched(cutoffs, alpha)

Calculates the matrix elements of the displacement gate using a recurrence relation. Supports batched alpha values.

Note: batch dimensions must be flattened.

Parameters:

• cutoffs – The Fock ladder output-input cutoffs.

• alpha – The batched displacement magnitude and angles.

Returns:

The batch of matrices representing the displacement operation.

homodyne_projector

mrmustard.mathlib.cython_lattice.homodyne_projector(fock_dim, A, b, c, out=None)

Fast homodyne projector.

Parameters:

• fock_dim – The Fock dimension (cutoff+1) of the remaining wires.

• A – The Bargmann A matrix.

• b – The Bargmann b vector.

• c – The Bargmann c scalar.

• out – An optional output array to write to. Should be of shape (fock_dim, fock_dim, fock_dim).

Returns:

The resulting (fock_dim, fock_dim, fock_dim) array.

homodyne_projector_batched

mrmustard.mathlib.cython_lattice.homodyne_projector_batched(fock_dim, A, b, c, out=None)

Fast homodyne projector for batched inputs.

Note: batch dimensions must be flattened.

Parameters:

• fock_dim – The Fock dimension (cutoff+1) of the remaining wires.

• A – The batched Bargmann A matrix.

• b – The batched Bargmann b vector.

• c – The batched Bargmann c scalar.

• out – An optional output array to write to. Should be of shape (batch_size, fock_dim, fock_dim, fock_dim).

Returns:

The resulting (batch_size, fock_dim, fock_dim, fock_dim) array.

jacobian_displacement

mrmustard.mathlib.cython_lattice.jacobian_displacement(D, alpha)

Calculates the jacobian of the displacement gate with respect to the complex displacement alpha and its conjugate. Both are needed for backprop, as the displacement gate is not a holomorphic function, as the Fock amplitudes depend on both alpha and on its conjugate. Each jacobian in this case has the same shape as the array D, as the displacement is a scalar.

Parameters:

• D (ndarray) – The D(alpha) gate in Fock representation.

• alpha (ndarray) – The alpha parameter of D(alpha).

Returns:

The Jacobian of the displacement gate with respect to alpha and alpha_conj.

Return type:

ndarray

squeezed

mrmustard.mathlib.cython_lattice.squeezed(cutoff, r, theta)

Calculates the matrix elements of the single-mode squeezed state using recurrence relations.

Parameters:

• cutoff – The Fock cutoff for the ket.

• r – The squeezing magnitude.

• theta – The squeezing angle.

Returns:

The matrix representing the squeezing gate.

squeezed_batched

mrmustard.mathlib.cython_lattice.squeezed_batched(cutoff, r, theta)

Calculates the matrix elements of the single-mode squeezed state using recurrence relations. Supports batched r and theta values.

Note: batch dimensions must be flattened.

Parameters:

• cutoff – The Fock cutoff for the ket.

• r – The batched squeezing magnitudes.

• theta – The batched squeezing angles.

Returns:

The batch of matrices representing the squeezing gate.

squeezer

mrmustard.mathlib.cython_lattice.squeezer(shape, r, theta)

Calculates the matrix elements of the squeezing gate using a recurrence relation. (See eq. 50-52 in https://arxiv.org/abs/2004.11002)

Parameters:

• shape – The Fock cutoffs for the output and input indices.

• r – The squeezing magnitude.

• theta – the squeezing angle.

Returns:

The matrix representing the squeezing gate.

squeezer_batched

mrmustard.mathlib.cython_lattice.squeezer_batched(shape, r, theta)

Calculates the matrix elements of the squeezing gate using a recurrence relation. Supports batched r and theta values. (See eq. 50-52 in https://arxiv.org/abs/2004.11002)

Note: batch dimensions must be flattened.

Parameters:

• shape – The Fock cutoffs for the output and input indices.

• r – The batched squeezing magnitudes.

• theta – The batched squeezing angles.

Returns:

The batch of matrices representing the squeezing gate.

vanilla

mrmustard.mathlib.cython_lattice.vanilla(shape, A, b, c, stable=False, out=None)

Vanilla algorithm for calculating the fock representation of a Gaussian tensor. This implementation works on flattened tensors and reshapes the tensor before returning.

The vanilla algorithm implements the flattened version of the following recursion which calculates the Fock amplitude at index \(k\) using a pivot at index \(k - 1_i\) and its neighbours at indices \(k - 1_i - 1_j\):

\[G_{k} = \frac{1}{\sqrt{k_i}} \left[ b_{i} G_{k-1_i} + \sum_j A_{ij} \sqrt{k_j - \delta_{ij}} G_{k-1_i-1_j} \right ]\]

where \(1_i\) is the vector of zeros with a 1 at index \(i\), and \(\delta_{ij}\) is the Kronecker delta. In this formula \(k\) is the vector of indices indexing into the Fock lattice. In the implementation the indices are flattened into a single integer index. This simplifies the bounds check when calculating the index of the pivot \(k-1_i\), which need to be done only until the index is smaller than the maximum stride.

see https://quantum-journal.org/papers/q-2020-11-30-366/ and https://arxiv.org/abs/2209.06069 for more details.

Parameters:

• shape (tuple[int, ...]) – shape of the output tensor

• A (np.ndarray) – A matrix of the Bargmann representation

• b (np.ndarray) – b vector of the Bargmann representation

• c (complex) – vacuum amplitude

• stable (bool) – whether to run the stable implementation of vanilla.

• out (np.ndarray) – if provided, the result will be stored in this tensor.

Returns:

Fock representation of the Gaussian tensor with shape shape

Return type:

np.ndarray

vanilla_batched

mrmustard.mathlib.cython_lattice.vanilla_batched(shape, A, b, c, stable=False, out=None)

The batched implementation of vanilla. Returns the Fock representation of the Gaussian tensor by parallelizing the single-instance vanilla over the batch dimension.

Args: shape (tuple[int, …]): shape of the output tensor A (np.ndarray): A matrix of the Bargmann representation b (np.ndarray): b vector of the Bargmann representation c (np.ndarray): vacuum amplitude out (np.ndarray): if provided, the result will be stored in this tensor.

Returns:

Fock representation of the Gaussian tensor with shape shape

Return type:

np.ndarray

Parameters:

• shape (tuple[int, ...])

• stable (bool)

• out (ndarray | None)

vanilla_vjp

mrmustard.mathlib.cython_lattice.vanilla_vjp(G, c, dLdG)

Vanilla vjp function. Returns dL/dA, dL/db, dL/dc.

Parameters:

• G (np.ndarray) – Tensor result of the forward pass

• c (complex) – vacuum amplitude

• dLdG (np.ndarray) – gradient of the loss with respect to the output tensor

Returns:

dL/dA, dL/db, dL/dc

Return type:

tuple[np.ndarray, np.ndarray, complex]

vanilla_vjp_batched

mrmustard.mathlib.cython_lattice.vanilla_vjp_batched(G, c, dLdG)

Batched implementation of vanilla_vjp. Returns dL/dA, dL/db, dL/dc by parallelizing the single-instance vanilla_vjp over the batch dimension.

Parameters:

• G (np.ndarray) – Tensor result of the forward pass with shape (batch_size,) + shape.

• c (np.ndarray) – Batched vacuum amplitudes with shape (batch_size,).

• dLdG (np.ndarray) – Gradient of the loss with respect to the output tensor G,

• G. (with the same shape as)

Returns:

dL/dA, dL/db, dL/dc.

dL/dA has shape (batch_size, D, D), dL/db has shape (batch_size, D), dL/dc has shape (batch_size,), where D is the number of modes (last dimension of G).

Return type:

tuple[np.ndarray, np.ndarray, np.ndarray]

code/mathlib/cython_lattice

 

Download Python script

 

Download Notebook

 

View on GitHub