Repository: XanaduAI/strawberryfields

Issue:
Custom contraction for beamsplitter and two-mode squeezing
The beamsplitter  and two-mode squeezing operation have special symmetries, namely SU(2) and SU(1,1). This implies that their matrix elements B_{nm}^{kl}, T_{ij}^{kl} satisfy the selection rules

![image](https://user-images.githubusercontent.com/991946/73875187-d5b94c00-4822-11ea-98a8-26fb47331197.png)

This implies that one can contract pure states (without loss of generality we assume only two modes) using one less loop as follows
![image](https://user-images.githubusercontent.com/991946/73875247-f681a180-4822-11ea-9bbf-8dec1819a546.png)
and 
![image](https://user-images.githubusercontent.com/991946/73875263-01d4cd00-4823-11ea-8de1-2effa308411e.png)
where c_{ij} is the tensor representing the two mode state and n is the cutoff in Fock space.
Similarly for a mixed state one should be able to do the update of the state using two less summations.

These observations should result in significant savings when acting beamsplitters and two-mode squeezing in the Fock backend.




Discussion:
Note that the Rgate and the Kgate can also be contracted in a more clever way.

Resolving patch:
diff --git a/strawberryfields/backends/fockbackend/circuit.py b/strawberryfields/backends/fockbackend/circuit.py
--- a/strawberryfields/backends/fockbackend/circuit.py
+++ b/strawberryfields/backends/fockbackend/circuit.py
@@ -23,6 +23,7 @@
 import numpy as np
 from numpy import sqrt, pi
 from scipy.special import factorial as bang
+from numba.typed import List
 
 from . import ops
 
@@ -314,7 +315,14 @@ def beamsplitter(self, t, r, phi, mode1, mode2):
         """
         Applies a beamsplitter.
         """
-        self._apply_gate(ops.beamsplitter(t, r, phi, self._trunc), [mode1, mode2])
+        mat = ops.beamsplitter(t, r, phi, self._trunc)
+
+        modes = List()
+        modes.append(mode1)
+        modes.append(mode2)
+
+        args = [mat, self._state, self._pure, modes, self._num_modes, self._trunc]
+        self._state = ops.apply_twomode_gate(*args, gate="BSgate")
 
     def squeeze(self, r, theta, mode):
         """
@@ -326,7 +334,15 @@ def two_mode_squeeze(self, r, theta, mode1, mode2):
         """
         Applies a two-mode squeezing gate.
         """
-        self._apply_gate(ops.two_mode_squeezing(r, theta, self._trunc), [mode1, mode2])
+        # self._apply_gate(ops.two_mode_squeezing(r, theta, self._trunc), [mode1, mode2])
+        mat = ops.two_mode_squeezing(r, theta, self._trunc)
+
+        modes = List()
+        modes.append(mode1)
+        modes.append(mode2)
+
+        args = [mat, self._state, self._pure, modes, self._num_modes, self._trunc]
+        self._state = ops.apply_twomode_gate(*args, gate="S2gate")
 
     def kerr_interaction(self, kappa, mode):
         """
diff --git a/strawberryfields/backends/fockbackend/ops.py b/strawberryfields/backends/fockbackend/ops.py
--- a/strawberryfields/backends/fockbackend/ops.py
+++ b/strawberryfields/backends/fockbackend/ops.py
@@ -27,6 +27,8 @@
 from scipy.special import factorial as fac
 from scipy.linalg import expm as matrixExp
 
+from numba import jit
+
 from thewalrus.fock_gradients import Dgate, Sgate, S2gate, BSgate
 
 
@@ -221,6 +223,7 @@ def apply_gate_BLAS(mat, state, pure, modes, n, trunc):
     einsum doesn't actually use BLAS but rather a c implementation. In theory
     if reshaping is efficient this should be faster.
     """
+    # pylint: disable=too-many-branches
 
     size = len(modes)
     dim = trunc**size
@@ -231,8 +234,14 @@ def apply_gate_BLAS(mat, state, pure, modes, n, trunc):
     transpose_list = [2*i for i in range(size)] + [2*i + 1 for i in range(size)]
     matview = np.transpose(mat, transpose_list).reshape((dim, dim))
 
+    # checks if matview is diagonal and, if so, use faster contractions
+    diag = np.all(matview == np.diag(np.diagonal(matview)))
+
     if pure:
         if n == 1:
+            if diag:
+                mat_diag = mat.diagonal()
+                return np.multiply(mat_diag, state)
             return np.dot(mat, state)
 
         # Transpose the state into the following form:
@@ -243,7 +252,11 @@ def apply_gate_BLAS(mat, state, pure, modes, n, trunc):
         # Apply matrix to each substate
         ret = np.zeros([trunc for i in range(n)], dtype=def_type)
         for i in product(*([range(trunc) for j in range(n - size)])):
-            ret[i] = np.dot(matview, view[i].ravel()).reshape(stshape)
+            if diag:
+                mat_diag = matview.diagonal()
+                ret[i] = np.multiply(mat_diag, view[i].ravel()).reshape(stshape)
+            else:
+                ret[i] = np.dot(matview, view[i].ravel()).reshape(stshape)
 
         # "untranspose" the return matrix ret
         untranspose_list = [0] * len(transpose_list)
@@ -254,6 +267,9 @@ def apply_gate_BLAS(mat, state, pure, modes, n, trunc):
 
     # otherwise, the state is mixed
     if n == 1:
+        if diag:
+            mat_diag = mat.diagonal().reshape(-1, 1)
+            return np.multiply(mat_diag, np.multiply(state, dagger(mat_diag)))
         return np.dot(mat, np.dot(state, dagger(mat)))
 
     # Transpose the state into the following form:
@@ -265,7 +281,11 @@ def apply_gate_BLAS(mat, state, pure, modes, n, trunc):
     # Apply matrix to each substate
     ret = np.zeros([trunc for i in range(n*2)], dtype=def_type)
     for i in product(*([range(trunc) for j in range((n - size)*2)])):
-        ret[i] = np.dot(matview, np.dot(view[i].reshape((dim, dim)), dagger(matview))).reshape(stshape + stshape)
+        if diag:
+            mat_diag = matview.diagonal().reshape(-1, 1)
+            ret[i] = np.multiply(mat_diag, np.multiply(view[i].reshape((dim, dim)), dagger(mat_diag))).reshape(stshape + stshape)
+        else:
+            ret[i] = np.dot(matview, np.dot(view[i].reshape((dim, dim)), dagger(matview))).reshape(stshape + stshape)
 
     # "untranspose" the return matrix ret
     untranspose_list = [0] * len(transpose_list)
@@ -318,6 +338,161 @@ def apply_gate_einsum(mat, state, pure, modes, n, trunc):
     return np.einsum(einstring, mat, state, mat.conj())
 
 
+def apply_twomode_gate(mat, state, pure, modes, n, trunc, gate="BSgate"):
+    """Applies a two-mode gate to a state.
+
+    Applies the specified two-mode gate to the state using custom tensor contractions and
+    the Numba compiler for faster application. Currently, only the beamsplitter and the
+    two-mode squeeze gate are supported.
+
+    Args:
+        mat (array[complex]): The numeric operator to be applied to the state, of shape `[trunc]*(2*n)`
+        state (array[complex]): The state that the operator is applied to
+        pure (bool): whether the state is pure or mixed
+        modes (list[int]): the list of modes to which the operator is applied on
+        n (int): The total number of modes
+        trunc (int): The Hilbert space truncation/cutoff
+        gate (str): The gate that is being applied. This argument determines the selection rules that
+            are used. Options are ``"BSgate"`` and ``"S2gate"``.
+
+    Returns:
+        array[complex]: state after application of the two-mode operation
+    """
+    # the transpositions below are necessary for using the jitted apply-gate
+    # functions by moving and grouping the modes on which the gate is applied
+    # to the front indices, since the other modes can then simply be ignored
+    if pure:
+        t1 = modes[0]
+        t2 = modes[1]
+
+        # moves the modes on which the gate should be applied to the
+        # front indices of the state, for mode 1 and 2 respectively
+        # e.g. from [i1, i2, i3] --> [i2, i3, i1] if the gate is applied to modes 2 and 3
+        switch_list_1 = np.arange(n)
+        switch_list_2 = np.arange(n)
+        switch_list_1[[0, t1]] = switch_list_1[[t1, 0]]
+        switch_list_2[[1, t2]] = switch_list_2[[t2, 1]]
+
+        state = state.transpose(switch_list_1)
+        state = state.transpose(switch_list_2)
+
+        if gate == "BSgate":
+            state = _apply_BS(mat, state, trunc)
+        elif gate == "S2gate":
+            state = _apply_S2(mat, state, trunc)
+        else:
+            raise NotImplementedError("Currently, selection rules are only implemented for the BSgate "
+                                      "and the S2gate. The {} gate is not supported".format(gate))
+
+        state = state.transpose(switch_list_2)
+        ret = state.transpose(switch_list_1)
+    else:
+        t1 = 2 * modes[0]
+        t2 = 2 * modes[1]
+
+        # moves the modes on which the gate should be applied to the
+        # front indices of the state, for mode 1 and 2 respectively
+        # e.g. from [i1, j1, i2, j2, i3, j3] --> [i2, j2, i3, j3, i1, j1]
+        # if the gate is applied to modes 2 and 3
+        switch_list_1 = np.arange(2 * n)
+        switch_list_2 = np.arange(2 * n)
+        switch_list_1[[0, 1, t1, t1+1]] = switch_list_1[[t1, t1+1, 0, 1]]
+        switch_list_2[[0, 1, t2, t2+1]] = switch_list_2[[t2, t2+1, 0, 1]]
+
+        # puts the modes on which the gate should be applied together
+        # e.g. from [i2, j2, i3, j3, i1, j1] --> [i2, i3, j2, j3, i1, j1]
+        # if the gate is applied to modes 2 and 3 (which are already switched to the first two modes)
+        transpose_list = np.arange(2 * n)
+        transpose_list[[t1+1, t2]] = transpose_list[[t2, t1+1]]
+
+        state = state.transpose(transpose_list)
+        state = state.transpose(switch_list_1)
+
+        if gate == "BSgate":
+            state = _apply_BS(mat, state, trunc)
+            state = state.transpose(switch_list_1)
+
+            state = state.transpose(switch_list_2)
+            state = _apply_BS(mat.conj(), state, trunc)
+
+        elif gate == "S2gate":
+            state = _apply_S2(mat, state, trunc)
+            state = state.transpose(switch_list_1)
+
+            state = state.transpose(switch_list_2)
+            state = _apply_S2(mat.conj(), state, trunc)
+
+        else:
+            raise NotImplementedError("Currently, selection rules are only implemented for the BSgate "
+                                      "and the S2gate. The {} gate is not supported".format(gate))
+
+        state = state.transpose(switch_list_2)
+        ret = state.transpose(transpose_list)
+
+    return ret
+
+
+# ignored in covtest (doesn't work well with the jit decorator)
+@jit(nopython=True)
+def _apply_BS(mat, state, trunc):  # pragma: no cover
+    r"""Applies the BS gate to the first bra in state.
+
+    The beamsplitter matrix elements :math:`B_{ij}^{kl}` satisfy the selection
+    rules
+
+    .. math::
+        B_{ij}^{kl} \propto \delta_{i+j}^{k+l}
+
+    This implies that one can contract pure states (without loss of generality
+    we assume only two modes) using one less loop as follows:
+
+    .. math::
+        c'_{i, j} = \sum\limits_{k = \text{max}(1+i+j-n, 0)}
+                   ^{\text{min}(i+j, n-1) + 1}
+        B_{ik}^{j, i+j-k} c_{k, i+j-k}
+
+    where :math:`c_{kl}` is the tensor representing the two mode state and
+    :math:`n` is the cutoff in Fock space. Similarly for a mixed state it is
+    possible to do the update of the state using two less summations.
+    """
+    ret = np.zeros_like(state, dtype=np.complex128)
+    for i in range(trunc):
+        for j in range(trunc):
+            for k in range(max(1+i+j-trunc, 0), min(i+j, trunc-1) + 1):
+                ret[i, j] += mat[i, k, j, i+j-k] * state[k, i+j-k]
+    return ret
+
+
+# ignored in covtest (doesn't work well with the jit decorator)
+@jit(nopython=True)
+def _apply_S2(mat, state, trunc):  # pragma: no cover
+    r"""Applies the S2 gate to the first bra in state.
+
+    The two-mode squeeze matrix elements :math:`T_{ij}^{kl}` satisfy the selection
+    rules
+
+    .. math::
+        T_{ij}^{kl} \propto \delta_{i-j}^{k-l}
+
+    This implies that one can contract pure states (without loss of generality
+    we assume only two modes) using one less loop as follows:
+
+    .. math::
+        c'_{i, k} = \sum\limits_{k = \text{max}(i-j, 0)}^{\text{min}(i-j, 0) + n}
+        T_{ij}^{k, k+j-i} c_{j, k+j-i}
+
+    where :math:`c_{jl}` is the tensor representing the two mode state and
+    :math:`n` is the cutoff in Fock space. Similarly for a mixed state it is
+    possible to do the update of the state using two less summations.
+    """
+    ret = np.zeros_like(state, dtype=np.complex128)
+    for i in range(trunc):
+        for j in range(trunc):
+            for k in range(max(i-j, 0), trunc + min(i-j, 0)):
+                ret[i, k] += mat[i, j, k, k + j-i] * state[j, k + j-i]
+    return ret
+
+
 # ============================================
 #
 # Gates


Test patch:
diff --git a/tests/backend/test_beamsplitter_operation.py b/tests/backend/test_beamsplitter_operation.py
--- a/tests/backend/test_beamsplitter_operation.py
+++ b/tests/backend/test_beamsplitter_operation.py
@@ -21,6 +21,8 @@
 Equivalently, we have t:=\cos(\theta) (t assumed real) and r:=\exp{i\phi}\sin(\theta)
 """
 
+import itertools as it
+
 import pytest
 
 import numpy as np
@@ -31,6 +33,7 @@
 PHASE_R = np.linspace(0, 2 * np.pi, 3, endpoint=False)
 ALPHA = 0.1
 MAG_ALPHAS = np.linspace(0.0, ALPHA, 3)
+MODES = list(it.combinations(range(4), 2))
 
 
 class TestRepresentationIndependent:
@@ -102,7 +105,7 @@ def test_normalized_beamsplitter_output(self, setup_backend, t, r_phi, tol):
     @pytest.mark.parametrize("mag_alpha", MAG_ALPHAS[1:])
     @pytest.mark.parametrize("r_phi", PHASE_R)
     def test_coherent_vacuum_interfered_fock_elements(
-        self, setup_backend, mag_alpha, t, r_phi, cutoff, pure, tol
+            self, setup_backend, mag_alpha, t, r_phi, cutoff, pure, tol
     ):
         r"""Tests if a range of beamsplitter output states (formed from a coherent state interfering with vacuum)
             have the correct Fock basis elements.
@@ -151,3 +154,54 @@ def test_coherent_vacuum_interfered_fock_elements(
             )
 
         assert np.allclose(numer_state, ref_state, atol=tol, rtol=0)
+
+
+@pytest.mark.backends("fock")
+class TestModeSubsets:
+    """Tests on mode subsets using the Fock basis representation."""
+
+    @pytest.mark.parametrize("t", T_VALUES)
+    @pytest.mark.parametrize("mag_alpha", MAG_ALPHAS[1:])
+    @pytest.mark.parametrize("r_phi", PHASE_R)
+    @pytest.mark.parametrize("modes", MODES)
+    def test_beamsplitter_on_mode_subset(
+            self, setup_backend, mag_alpha, t, r_phi, modes, cutoff, tol
+    ):
+        """Tests applying the beamsplitter on different mode subsets."""
+
+        phase_alpha = np.pi / 5
+        alpha = mag_alpha * np.exp(1j * phase_alpha)
+        r = np.exp(1j * r_phi) * np.sqrt(1.0 - np.abs(t) ** 2)
+
+        backend = setup_backend(4)
+
+        backend.displacement(alpha, modes[0])
+        backend.beamsplitter(t, r, *modes)
+        state = backend.state()
+
+        alpha_outA = t * alpha
+        alpha_outB = r * alpha
+
+        n = np.arange(cutoff)
+        ref_stateA = (
+            np.exp(-0.5 * np.abs(alpha_outA) ** 2)
+            * alpha_outA ** n
+            / np.sqrt(factorial(n))
+        )
+        ref_stateB = (
+            np.exp(-0.5 * np.abs(alpha_outB) ** 2)
+            * alpha_outB ** n
+            / np.sqrt(factorial(n))
+        )
+
+        numer_state = state.reduced_dm(list(modes))
+
+        ref_state = np.einsum(
+            "i,j,k,l->ijkl",
+            ref_stateA,
+            np.conj(ref_stateA),
+            ref_stateB,
+            np.conj(ref_stateB),
+        )
+
+        assert np.allclose(numer_state, ref_state, atol=tol, rtol=0)
diff --git a/tests/backend/test_twomode_squeezing_operation.py b/tests/backend/test_twomode_squeezing_operation.py
--- a/tests/backend/test_twomode_squeezing_operation.py
+++ b/tests/backend/test_twomode_squeezing_operation.py
@@ -26,6 +26,7 @@
 
 PHASE = np.linspace(0, 2 * np.pi, 3, endpoint=False)
 MAG = np.linspace(0.0, 0.2, 5, endpoint=False)
+MODES = list(it.combinations(range(4), 2))
 
 
 def get_amplitude(k, r, p):
@@ -74,3 +75,28 @@ def test_two_mode_squeezing(self, setup_backend, r, p, cutoff, pure, tol):
                     tmsv2 = get_amplitude(k, r, p) * np.conj(get_amplitude(l, r, p))
 
                     assert np.allclose(state.data[t], tmsv2, atol=tol, rtol=0)
+
+
+    @pytest.mark.parametrize("r", MAG)
+    @pytest.mark.parametrize("p", PHASE)
+    @pytest.mark.parametrize("modes", MODES)
+    def test_squeezing_on_mode_subset(self, setup_backend, r, p, modes, cutoff, pure, tol):
+        r""" Test two-mode squeezing on vacuum-state for both pure states and
+        mixed states on different mode subsets with the amplitude given by
+	        :math:`\delta_{kl} \frac{e^{in\phi} \tanh^n{r}}{\cosh{r}}`
+        """
+        z = r * np.exp(1j * p)
+
+        backend = setup_backend(4)
+        backend.two_mode_squeeze(z, *modes)
+
+        state = backend.state()
+
+        state_data = state.reduced_dm(list(modes))
+
+        for k in it.product(range(cutoff), repeat=2):
+            for l in it.product(range(cutoff), repeat=2):
+                t = (k[0], l[0], k[1], l[1])
+                tmsv2 = get_amplitude(k, r, p) * np.conj(get_amplitude(l, r, p))
+
+                assert np.allclose(state_data[t], tmsv2, atol=tol, rtol=0)