Add n-mode gaussian transformation gate in TFbackend (#599)

* 1,2,3,4 in tfbackend codes * test for the new git config * changes in detail * modify the gradient part * move choi trick here to be differentiable * add the expression for C from S,d,t and put all codes in TF * new modifications for the parameters * add n_mode_gate code * fix problem from codefactor * fix * modify the n_mode_gate and add the tests for it * black the files * fix n_mode_gate nad gaussian_gate_matrix * add a Symplectic optimizer * delete optimizer for now and open a new PR for that * modify the functions from thewalrus * fix * balck files * fix * fix ns proble * minor addition on fock.py compiler * fix the problem about the *modes and register * fix trailing whitespace * modify the requirement for TheWalrus * requirements * https -> http * tests file modified * test 1 * parameter * modify tests * test 2 * add new test for gaussian gate * modify the dy inside the gradient of Ggate * modify the test * add the warning when initialize the Ggate to use the specific Symplectic Optimizer * black the file * move the warning from ops into the backend * fix warning * fix gradient problem: unconnected gradientgit add strawberryfields/backends/tfbackend/ops.py * once warning * fix import issue * change the requirement file for the version of thewalrus * blacked * fix tests * transpose the gate inside gaussian gate * blacked * changelog * test coverage changes or not * coverage upup * fix coverage * fix * ahhhh * last one * add tests * test * changelog + test file changed * change d->2N * blacked * new tests * Ggate test * test * test passed locally * 1 * test passed * remove the test of symplectic -> to be discussed * mayve_batch error fix * complex64 -> 128 * fix gradient of d proble, * comments fix * batch error * black ? * test * try to fix the test error * add cutoff param * try to fix test error * changes for req * req test * dicrease the cutoff dim in the test * make the test simple to see how is it going * git add tests/integration/test_ops_integration.py * fix error * optimize choi_trick func * try to change TF version * independence conflit * remove the tests * add one test * modify tests * modifyyyy * test 2 * error * modify the transp * hopefully fixes the last bugs * making codefactor happy * blacked * cover more * removes unused functions * put optimizer function inside ops (TOBE DISCUSSED) * add tests for optimizer * blacked * import update_symplectic in sf namespace * added blank lines for clarity * blacked * blacked tests * marked backend for tests * blacked * changes pytest mark * adds setup_backend fixture * removed tqdm * fixes Ggate test * blacked tests * trigger tests * fix [tfbackend/circuit] - gaussian_gate: pass dtype argument to gate implementation * refactor [tfbackend/ops] - gaussian_gate: use default tf dtype in class `FockStateTF` * fix [ops/gaussian_gate] - is_symplectic: lower atol for `tf.complex64` * fixed docstring * updated changelog * updated changelog * updated changelog * Update CHANGELOG.md * Update .github/CHANGELOG.md * Update requirements.txt * Update CHANGELOG.md * Update .github/CHANGELOG.md Co-authored-by: Theodor <theodor.isacsson@gmail.com> * def update_symplectic if tensorflow is installed * removed ns=None * fixes import * improved example * removed unnecessary return * improved warning message * blacked * improved import of update_symplectic * blacked Co-authored-by: Yuan Yao <yuanyao@YuandeMacBook-Pro.local> Co-authored-by: Nicolas Quesada <zeitus@gmail.com> Co-authored-by: ziofil <miatto@gmail.com> Co-authored-by: Theodor <theodor@xanadu.ai> Co-authored-by: ziofil <ziofil@users.noreply.github.com> Co-authored-by: Sebastián Duque Mesa <675763+sduquemesa@users.noreply.github.com> Co-authored-by: Josh Izaac <josh146@gmail.com> Co-authored-by: Theodor <theodor.isacsson@gmail.com>
XanaduAI · Nov 3, 2021 · b446075 · b446075
1 parent 688ac9f
commit b446075
Show file tree

Hide file tree

Showing 9 changed files with 674 additions and 3 deletions.
diff --git a/.github/CHANGELOG.md b/.github/CHANGELOG.md
@@ -2,6 +2,73 @@
 
 <h3>New features since last release</h3>
 
+* The generic multimode Gaussian gate ``Ggate`` is now available in the ``sf.ops``
+  module with the backend choice of ``tf``. The N mode ``Ggate`` can be parametrized by a real
+  symplectic matrix `S` (size `2N * 2N`) and a diplacement vector `d` (size `N`). You can also
+  obtain the gradients of the Ggate gate via TensorFlow's ``tape.gradient``
+  [(#599)](https://github.com/XanaduAI/strawberryfields/pull/599)
+  [(#606)](https://github.com/XanaduAI/strawberryfields/pull/606)
+
+  ```python
+  from thewalrus.random import random_symplectic
+
+  num_mode = 2
+  cutoff = 10
+  S = tf.Variable(random_symplectic(num_mode))
+  d = tf.Variable(np.random.random(2 * num_mode))
+
+  eng = sf.Engine("tf", backend_options={"cutoff_dim": cutoff})
+  prog = sf.Program(2)
+
+  with prog.context as q:
+      sf.ops.Ggate(S, d) | (q[0], q[1])
+
+  state_out = eng.run(prog).state.ket()
+  ```
+
+  Note that in order to update the parameter `S` by using its gradient, you cannot use gradient 
+  descent directly (as the unitary would not be symplectic after the update). Please use the
+  function `sf.backends.tfbackend.update_symplectic` which is designed specifically for this purpose.
+
+  ```python
+  def overlap_loss(state, objective):
+      return -tf.abs(tf.reduce_sum(tf.math.conj(state) * objective)) ** 2
+
+  def norm_loss(state):
+      return -tf.abs(tf.linalg.norm(state)) ** 2
+
+  def loss(state, objective):
+      return overlap_loss(state, objective) + norm_loss(state)
+
+  num_mode = 1
+  cutoff = 10
+
+  S = tf.Variable(random_symplectic(num_mode))
+  d = tf.Variable(np.random.random(2 * num_mode))
+  kappa = tf.Variable(0.3)
+  objective = tf.Variable(np.eye(cutoff)[1], dtype=tf.complex64)
+
+  adam = tf.keras.optimizers.Adam(learning_rate=0.01)
+  eng = sf.Engine("tf", backend_options={"cutoff_dim": cutoff})
+  prog = sf.Program(1)
+
+  with prog.context as q:
+      sf.ops.Ggate(S, d) | q
+      sf.ops.Kgate(kappa) | q
+
+  loss_vals = []
+  for _ in range(200):
+      with tf.GradientTape() as tape:
+          state_out = eng.run(prog).state.ket()
+          loss_val = loss(state_out, objective)
+
+      eng.reset()
+      grad_S, gradients_d, gradients_kappa = tape.gradient(loss_val, [S, d, kappa])
+      adam.apply_gradients(zip([gradients_d, gradients_kappa], [d, kappa]))
+      update_symplectic(S, grad_S, lr=0.1)  # update S here
+      loss_vals.append(loss_val)
+  ```
+
 <h3>Breaking Changes</h3>
 
 <h3>Bug fixes</h3>
@@ -12,6 +79,8 @@
 
 This release contains contributions from (in alphabetical order):
 
+Sebastián Duque Mesa, Filippo Miatto, Nicolás Quesada, Yuan Yao.
+
 # Release 0.19.0 (current release)
 
 <h3>New features since last release</h3>

diff --git a/strawberryfields/backends/tfbackend/__init__.py b/strawberryfields/backends/tfbackend/__init__.py
@@ -150,7 +150,6 @@
     tf_available = True
     tf_version = tensorflow.__version__
 
-
 tf_info = """\
 To use Strawberry Fields with TensorFlow support, version 2.x of TensorFlow is required.
 This can be installed as follows:
@@ -169,3 +168,6 @@ def excepthook(type, value, traceback):
     sys.excepthook = excepthook
 
     raise ImportError(tf_info)
+
+
+from .ops import update_symplectic
diff --git a/strawberryfields/backends/tfbackend/backend.py b/strawberryfields/backends/tfbackend/backend.py
@@ -228,6 +228,12 @@ def cross_kerr_interaction(self, kappa, mode1, mode2):
             remapped_modes = self._remap_modes([mode1, mode2])
             self.circuit.cross_kerr_interaction(kappa, remapped_modes[0], remapped_modes[1])
 
+    def gaussian_gate(self, S, d, *modes):
+        """Multimode gaussian gate parametrized by its symplectic matrix and displacement vector, the number of modes to be applied is not fixed"""
+        with tf.name_scope("Gaussian_gate"):
+            remapped_modes = self._remap_modes(modes)
+            self.circuit.gaussian_gate(S, d, remapped_modes)
+
     def state(self, modes=None, **kwargs):
         r"""Returns the state of the quantum simulation, restricted to the subsystems defined by `modes`.
 

diff --git a/strawberryfields/backends/tfbackend/circuit.py b/strawberryfields/backends/tfbackend/circuit.py
@@ -574,6 +574,24 @@ def cubic_phase(self, gamma, mode):
         )
         self._update_state(new_state)
 
+    def gaussian_gate(self, S, d, modes):
+        """
+        Apply the N-mode gaussian gates to the specified mode.
+        """
+        self._batched = len(d.shape) == 2
+        self._batch_size = d.shape[0] if self._batched else None
+        new_state = ops.gaussian_gate(
+            S,
+            d,
+            modes,
+            in_modes=self._state,
+            cutoff=self._cutoff_dim,
+            pure=self._state_is_pure,
+            batched=self._batched,
+            dtype=self._dtype,
+        )
+        self._update_state(new_state)
+
     def loss(self, T, mode):
         """
         Apply a loss channel  to the specified mode.