610 mzgate add native support (#617)

* interface [backends] - mzgate

* implementation [backends/fockbackend] - mzgate

* implementation [compilers/fock] - set mzgate as primitive gate

* implementation [ops] - mzgate as primitive _apply method

* implementation [backends/tfbackend] - mzgate

* test [frontend] - test decomposition of MZgate

* requirement [requirements.txt] - require dev version of The Walrus

* fix [backends] - fix CodeFactor issues

* fix [tests/frontend] - fix comment text

* requirement [requirements.txt] - download dev version of The Walrus from wheel

* add changes to CHANGELOG

* test [backend] - test MZgate operation

* fix [changelog] - typo

Co-authored-by: Josh Izaac <josh146@gmail.com>

* fix [docs/requirements] - use dev version of teh walrus from wheel

* refactor [fockbackend/circuit] - remove duplicated code for MZgate and BSgate

Both MZgate and BSgate use the same selections rules. This commit implement the same method for both.

* fix [doc/requirements] - re-commit: use dev version of the walrus with docs

* refactor [fockbackend/circuit] - CodeFactor

* Apply suggestions from code review

* fix [tfbackend/ops] - fix docstrings

fix misleading docstrings for mzgate and two_mode_squeeze

* remove breakpoint

* test [test_tf_integration] - test_MZ_state_gradients

* refactor [test_mzgate_operation] - all tests

use initial fock |1,1> state instead of a coherent state

* fix [tests/backend] - test_mzgate_operation

get the test to catch and fail the malfunctioning mzgate on the tf backend

* fix [tfbackend] - circuit: wrong parameter handling

fix wrong handling of parameters when batching

* impl [integration] - test_MZ_state_gradients

implement assertions and refactor test

* impl [backend] - test_mzgate_operation

implement new test with initial |1,1> fock state checking P_11 and P_02

* fix [tests/backend] - test_mzgate_operation: parameters as list insted of np.array

use a list for parametrized tests instead of numpy arrays

* Revert "fix [tests/backend] - test_mzgate_operation: parameters as list insted of np.array"

This reverts commit aa8a823.
The commit does not solve the problem

* fix [tests/backend] - test_mzgate_operation: use all_fock_prob

use `all_fock_probabilities` instead of `ket` state to do assertions

* fix [tests] - mzgate: bypass batched mode

Co-authored-by: Josh Izaac <josh146@gmail.com>
Co-authored-by: Nicolas Quesada <nicolas@xanadu.ai>
Co-authored-by: Nicolas Quesada <zeitus@gmail.com>

.github/CHANGELOG.md

@@ -144,9 +144,14 @@
 * Speed improvements to ``gaussian_unitary`` compiler
   [(#603)](https://github.com/XanaduAI/strawberryfields/pull/603)
 
+
+* Added native support in the Fock backend for the MZgate.
+  [#610](https://github.com/XanaduAI/strawberryfields/issues/610)
+
 * `measure_threshold` is now supported in the `bosonic` backend.
   [(#618)](https://github.com/XanaduAI/strawberryfields/pull/618)
 
+
 <h3>Breaking Changes</h3>
 
 <h3>Bug fixes</h3>
@@ -183,7 +188,9 @@
 
 This release contains contributions from (in alphabetical order):
 
-J. Eli Bourassa, Jake Bulmer, Theodor Isacsson, Aaron Robertson, Jeremy Swinarton, Antal Száva, Federico Rueda, Yuan Yao.
+
+J. Eli Bourassa, Jake Bulmer, Sebastian Duque, Theodor Isacsson, Aaron Robertson, Jeremy Swinarton, Antal Száva, Federico Rueda, Yuan Yao.
+
 
 # Release 0.18.0 (current release)
 

doc/requirements.txt

@@ -15,5 +15,5 @@ sphinx-copybutton
 sphinx-automodapi
 sphinxcontrib-bibtex==0.4.2
 tensorflow==2.5.0
-thewalrus>=0.15.0
-toml
+thewalrus>=0.16.0
+toml

strawberryfields/backends/base.py

@@ -98,7 +98,7 @@ def delete(self, modes):
         if self.valid(modes):
             new_map = []
             ctr = 0
-            for m in range(len(self._map)):
+            for m, _ in enumerate(self._map):
                 if m in modes or self._map[m] is None:
                     new_map.append(None)
                 else:
@@ -338,6 +338,17 @@ def beamsplitter(self, theta, phi, mode1, mode2):
         """
         raise NotImplementedError
 
+    def mzgate(self, phi_in, phi_ex, mode1, mode2):
+        """Apply the Mach-Zehnder interferometer operation to the specified modes.
+
+        Args:
+            phi_in (float): internal phase
+            phi_ex (float): external phase
+            mode1 (int): first mode that MZ interferometer acts on
+            mode2 (int): second mode that MZ interferometer acts on
+        """
+        raise NotImplementedError
+
     def loss(self, T, mode):
         r"""Perform a loss channel operation on the specified mode.
 

strawberryfields/backends/fockbackend/backend.py

@@ -178,6 +178,9 @@ def two_mode_squeeze(self, r, phi, mode1, mode2):
     def beamsplitter(self, theta, phi, mode1, mode2):
         self.circuit.beamsplitter(theta, phi, self._remap_modes(mode1), self._remap_modes(mode2))
 
+    def mzgate(self, phi_in, phi_ex, mode1, mode2):
+        self.circuit.mzgate(phi_in, phi_ex, self._remap_modes(mode1), self._remap_modes(mode2))
+
     def measure_homodyne(self, phi, mode, shots=1, select=None, **kwargs):
         """Perform a homodyne measurement on the specified mode.
 

strawberryfields/backends/fockbackend/circuit.py

@@ -24,7 +24,6 @@
 from numpy import sqrt, pi
 from scipy.special import factorial as bang
 from numba import jit
-from numba.typed import List
 
 from . import ops
 
@@ -221,14 +220,14 @@ def apply_twomode_gate(self, mat, modes, gate="BSgate"):
         """Applies a two-mode gate to the 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.
+        the Numba compiler for faster application. Currently, only the beamsplitter, the
+        Mach-Zehnder and two-mode squeeze gate are supported.
 
         Args:
             mat (array[complex]): The numeric operator to be applied to the state, of shape `[trunc]*(2*n)`
             modes (list[int]): The list of modes to which the operator is applied on
             gate (str): The gate that is being applied. This argument determines the selection rules that
-                are used. Options are ``"BSgate"`` and ``"S2gate"``.
+                are used. Options are ``"BSgate"``, ``"MZgate`` and ``"S2gate"``.
 
         Returns:
             array[complex]: The state after application of the two-mode operation
@@ -248,14 +247,14 @@ def apply_twomode_gate(self, mat, modes, gate="BSgate"):
             self._state = self._state.transpose(switch_list_1)
             self._state = self._state.transpose(switch_list_2)
 
-            if gate == "BSgate":
-                self._state = self._apply_BS(mat, self._state, self._trunc)
+            if gate in ("BSgate", "MZgate"):
+                self._state = self._apply_two_mode_passive(mat, self._state, self._trunc)
             elif gate == "S2gate":
                 self._state = self._apply_S2(mat, self._state, self._trunc)
             else:
                 raise NotImplementedError(
-                    "Currently, selection rules are only implemented for the BSgate "
-                    "and the S2gate. The {} gate is not supported".format(gate)
+                    "Currently, selection rules are only implemented for the BSgate, "
+                    "the MZgate and the S2gate. The {} gate is not supported".format(gate)
                 )
 
             self._state = self._state.transpose(switch_list_2)
@@ -282,12 +281,12 @@ def apply_twomode_gate(self, mat, modes, gate="BSgate"):
             self._state = self._state.transpose(transpose_list)
             self._state = self._state.transpose(switch_list_1)
 
-            if gate == "BSgate":
-                self._state = self._apply_BS(mat, self._state, self._trunc)
+            if gate in ("BSgate", "MZgate"):
+                self._state = self._apply_two_mode_passive(mat, self._state, self._trunc)
                 self._state = self._state.transpose(switch_list_1)
 
                 self._state = self._state.transpose(switch_list_2)
-                self._state = self._apply_BS(mat.conj(), self._state, self._trunc)
+                self._state = self._apply_two_mode_passive(mat.conj(), self._state, self._trunc)
 
             elif gate == "S2gate":
                 self._state = self._apply_S2(mat, self._state, self._trunc)
@@ -307,7 +306,7 @@ def apply_twomode_gate(self, mat, modes, gate="BSgate"):
     # ignored in covtest (doesn't work well with the jit decorator)
     @staticmethod
     @jit(nopython=True)
-    def _apply_BS(mat, state, trunc):  # pragma: no cover
+    def _apply_two_mode_passive(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
@@ -555,6 +554,13 @@ def beamsplitter(self, theta, phi, mode1, mode2):
         mat = ops.beamsplitter(theta, phi, self._trunc)
         self._state = self.apply_twomode_gate(mat, [mode1, mode2], gate="BSgate")
 
+    def mzgate(self, phi_in, phi_ex, mode1, mode2):
+        """
+        Applies a MZ gate.
+        """
+        mat = ops.mzgate(phi_in, phi_ex, self._trunc)
+        self._state = self.apply_twomode_gate(mat, [mode1, mode2], gate="MZgate")
+
     def squeeze(self, r, theta, mode):
         """
         Applies a squeezing gate.

strawberryfields/backends/fockbackend/ops.py

@@ -32,6 +32,7 @@
     squeezing as squeezing_tw,
     two_mode_squeezing as two_mode_squeezing_tw,
     beamsplitter as beamsplitter_tw,
+    mzgate as mzgate_tw,
 )
 
 def_type = np.complex128
@@ -185,7 +186,7 @@ def project_reset(modes, x, state, pure, n, trunc):
 
     def intersperse(lst):
         # pylint: disable=missing-docstring
-        return tuple([lst[i // 2] for i in range(len(lst) * 2)])
+        return tuple(lst[i // 2] for i in range(len(lst) * 2))
 
     if pure:
         ret = np.zeros([trunc for i in range(n)], dtype=def_type)
@@ -328,6 +329,20 @@ def beamsplitter(theta, phi, trunc):
     return BS_tw.transpose((0, 2, 1, 3))
 
 
+@functools.lru_cache()
+def mzgate(phi_in, phi_ex, cutoff):
+    r"""The MZ gate :math:`\mathrm{MZ}(\phi_{in}, \phi_{ex})`.
+
+    Uses the `MZ gate operation from The Walrus`_ to calculate the Fock representation of the Mach-Zehnder interferometer.
+
+    .. _`MZ gate operation from The Walrus`: https://the-walrus.readthedocs.io/en/latest/code/api/thewalrus.fock_gradients.mzgate.html
+    """
+    ret = mzgate_tw(phi_in, phi_ex, cutoff)
+
+    # Transpose needed because of different conventions in SF and The Walrus.
+    return ret.transpose((0, 2, 1, 3))
+
+
 @functools.lru_cache()
 def proj(i, j, trunc):
     r"""

strawberryfields/backends/tfbackend/backend.py

@@ -198,6 +198,11 @@ def beamsplitter(self, theta, phi, mode1, mode2):
             remapped_modes = self._remap_modes([mode1, mode2])
             self.circuit.beamsplitter(theta, phi, remapped_modes[0], remapped_modes[1])
 
+    def mzgate(self, phi_in, phi_ex, mode1, mode2):
+        with tf.name_scope("MZgate"):
+            remapped_modes = self._remap_modes([mode1, mode2])
+            self.circuit.mzgate(phi_in, phi_ex, remapped_modes[0], remapped_modes[1])
+
     def two_mode_squeeze(self, r, phi, mode1, mode2):
         with tf.name_scope("Two-mode_squeezing"):
             remapped_modes = self._remap_modes([mode1, mode2])

strawberryfields/backends/tfbackend/circuit.py

@@ -494,6 +494,26 @@ def beamsplitter(self, theta, phi, mode1, mode2):
         )
         self._update_state(new_state)
 
+    def mzgate(self, phi_in, phi_ex, mode1, mode2):
+        """
+        Apply a MZ-gate operator to the two specified modes.
+        """
+        phi_in = self._maybe_batch(phi_in)
+        phi_ex = self._maybe_batch(phi_ex)
+        self._check_incompatible_batches(phi_in, phi_ex)
+        new_state = ops.mzgate(
+            phi_in,
+            phi_ex,
+            mode1,
+            mode2,
+            self._state,
+            self._cutoff_dim,
+            self._state_is_pure,
+            self._batched,
+            dtype=self._dtype,
+        )
+        self._update_state(new_state)
+
     def two_mode_squeeze(self, r, phi, mode1, mode2):
         """
         Apply a two-mode squeezing operator to the two specified modes.

strawberryfields/backends/tfbackend/ops.py

@@ -41,6 +41,8 @@
 from thewalrus.fock_gradients import grad_squeezing as grad_squeezing_tw
 from thewalrus.fock_gradients import beamsplitter as beamsplitter_tw
 from thewalrus.fock_gradients import grad_beamsplitter as grad_beamsplitter_tw
+from thewalrus.fock_gradients import mzgate as mzgate_tw
+from thewalrus.fock_gradients import grad_mzgate as grad_mzgate_tw
 from thewalrus.fock_gradients import two_mode_squeezing as two_mode_squeezing_tw
 from thewalrus.fock_gradients import grad_two_mode_squeezing as grad_two_mode_squeezing_tw
 
@@ -353,6 +355,42 @@ def beamsplitter_matrix(theta, phi, cutoff, batched=False, dtype=tf.complex64):
     )
 
 
+@tf.custom_gradient
+def single_mzgate_matrix(phi_in, phi_ex, cutoff, dtype=tf.complex64.as_numpy_dtype):
+    """creates a single mode mzgate matrix"""
+    phi_in = phi_in.numpy()
+    phi_ex = phi_ex.numpy()
+
+    gate = mzgate_tw(phi_in, phi_ex, cutoff, dtype)
+    gate = np.transpose(gate, [0, 2, 1, 3])
+
+    def grad(dy):
+        Dtheta, Dphi = grad_mzgate_tw(np.transpose(gate, [0, 2, 1, 3]), phi_in, phi_ex)
+        Dtheta = np.transpose(Dtheta, [0, 2, 1, 3])
+        Dphi = np.transpose(Dphi, [0, 2, 1, 3])
+        grad_theta = tf.math.real(tf.reduce_sum(dy * tf.math.conj(Dtheta)))
+        grad_phi = tf.math.real(tf.reduce_sum(dy * tf.math.conj(Dphi)))
+        return grad_theta, grad_phi, None
+
+    return gate, grad
+
+
+def mzgate_matrix(phi_in, phi_ex, cutoff, batched=False, dtype=tf.complex64):
+    """creates a single mode mzgate matrix accounting for batching"""
+    phi_in = tf.cast(phi_in, dtype)
+    phi_ex = tf.cast(phi_ex, dtype)
+    if batched:
+        return tf.stack(
+            [
+                single_mzgate_matrix(phi_in_, phi_ex_, cutoff, dtype=dtype.as_numpy_dtype)
+                for phi_in_, phi_ex_ in tf.transpose([phi_in, phi_ex])
+            ]
+        )
+    return tf.convert_to_tensor(
+        single_mzgate_matrix(phi_in, phi_ex, cutoff, dtype=dtype.as_numpy_dtype)
+    )
+
+
 @tf.custom_gradient
 def single_two_mode_squeezing_matrix(theta, phi, cutoff, dtype=tf.complex64.as_numpy_dtype):
     """creates a single mode two-mode squeezing matrix"""
@@ -795,10 +833,21 @@ def beamsplitter(
     return output
 
 
+def mzgate(
+    phi_in, phi_ex, mode1, mode2, in_modes, cutoff, pure=True, batched=False, dtype=tf.complex64
+):
+    """returns mzgate unitary matrix on specified input modes"""
+    phi_in = tf.cast(phi_in, dtype)
+    phi_ex = tf.cast(phi_ex, dtype)
+    matrix = mzgate_matrix(phi_in, phi_ex, cutoff, batched, dtype)
+    output = two_mode_gate(matrix, mode1, mode2, in_modes, pure, batched)
+    return output
+
+
 def two_mode_squeeze(
     r, theta, mode1, mode2, in_modes, cutoff, pure=True, batched=False, dtype=tf.complex64
 ):
-    """returns beamsplitter unitary matrix on specified input modes"""
+    """returns two mode squeeze unitary matrix on specified input modes"""
     matrix = two_mode_squeezer_matrix(r, theta, cutoff, batched, dtype)
     output = two_mode_gate(matrix, mode1, mode2, in_modes, pure, batched)
     return output

strawberryfields/compilers/fock.py

@@ -50,6 +50,7 @@ class Fock(Compiler):
         "Kgate",
         # two mode gates
         "BSgate",
+        "MZgate",
         "CKgate",
         "S2gate",
     }
@@ -63,7 +64,6 @@ class Fock(Compiler):
         "Pgate": {},
         "CXgate": {},
         "CZgate": {},
-        "MZgate": {},
         "sMZgate": {},
         "Xgate": {},
         "Zgate": {},

strawberryfields/ops.py

@@ -1953,6 +1953,10 @@ class MZgate(Gate):
     def __init__(self, phi_in, phi_ex):
         super().__init__([phi_in, phi_ex])
 
+    def _apply(self, reg, backend, **kwargs):
+        phi_in, phi_ex = par_evaluate(self.p)
+        backend.mzgate(phi_in, phi_ex, *reg)
+
     def _decompose(self, reg, **kwargs):
         # into local phase shifts and two 50-50 beamsplitters
         return [