ggate fix

mrmustard/lab_dev/transformations/ggate.py

@@ -20,6 +20,7 @@
 
 from typing import Sequence
 from mrmustard import math
+from mrmustard.math.parameters import update_symplectic
 from mrmustard.utils.typing import RealMatrix
 
 from .base import Unitary
@@ -58,17 +59,20 @@ def __init__(
         super().__init__(name="Ggate")
 
         symplectic = symplectic if symplectic is not None else math.random_symplectic(len(modes))
-        S = make_parameter(symplectic_trainable, symplectic, "symplectic", (None, None))
-        self.parameter_set.add_parameter(S)
+        self._add_parameter(
+            make_parameter(
+                symplectic_trainable, symplectic, "symp", (None, None), update_symplectic
+            )
+        )
         self._representation = self.from_ansatz(
             modes_in=modes,
             modes_out=modes,
             ansatz=PolyExpAnsatz.from_function(
                 fn=lambda s: Unitary.from_symplectic(modes, s).bargmann_triple(),
-                s=self.parameter_set.symplectic,
+                s=self.symp,
             ),
         ).representation
 
     @property
     def symplectic(self):
-        return self.parameter_set.symplectic.value
+        return self.symp.value

tests/test_training/test_opt.py

@@ -163,24 +163,24 @@ def cost_fn():
         opt.minimize(cost_fn, by_optimizing=[circ], max_steps=300)
         assert np.allclose(-cost_fn(), 0.25, atol=1e-5)
 
-    # def test_learning_two_mode_Ggate(self):
-    #     """Finding the optimal Ggate to make a pair of single photons"""
-    #     skip_np()
+    def test_learning_two_mode_Ggate(self):
+        """Finding the optimal Ggate to make a pair of single photons"""
+        skip_np()
 
-    #     settings.SEED = 42
-    #     rng = tf.random.get_global_generator()
-    #     rng.reset_from_seed(settings.SEED)
+        settings.SEED = 42
+        rng = tf.random.get_global_generator()
+        rng.reset_from_seed(settings.SEED)
 
-    #     G = Ggate((0,1), symplectic_trainable=True)
+        G = Ggate((0, 1), symplectic_trainable=True)
 
-    #     def cost_fn():
-    #         amps = (Vacuum((0,1)) >> G).fock_array((2,2))
-    #         return -math.abs(amps[1, 1]) ** 2 + math.abs(amps[0, 1]) ** 2
+        def cost_fn():
+            amps = (Vacuum((0, 1)) >> G).fock_array((2, 2))
+            return -math.abs(amps[1, 1]) ** 2 + math.abs(amps[0, 1]) ** 2
 
-    #     opt = Optimizer(symplectic_lr=0.5, euclidean_lr=0.01)
+        opt = Optimizer(symplectic_lr=0.5, euclidean_lr=0.01)
 
-    #     opt.minimize(cost_fn, by_optimizing=[G], max_steps=500)
-    #     assert np.allclose(-cost_fn(), 0.25, atol=1e-4)
+        opt.minimize(cost_fn, by_optimizing=[G], max_steps=500)
+        assert np.allclose(-cost_fn(), 0.25, atol=1e-4)
 
     # def test_learning_two_mode_Interferometer(self):
     #     """Finding the optimal Interferometer to make a pair of single photons"""
@@ -418,10 +418,11 @@ def cost_fn():
     #     def thermal_entropy(nbar):
     #         return -(nbar * np.log((nbar) / (1 + nbar)) - np.log(1 + nbar))
 
-    #     G = Ggate(num_modes=2, symplectic_trainable=True, symplectic=S_init)
+    #     G = Ggate((0,1), symplectic=S_init, symplectic_trainable=True)
 
     #     def cost_fn():
-    #         state = Vacuum(2) >> G
+    #         state = Vacuum((0,1)) >> G
+    #         # TODO: cov and means
     #         cov1, _ = trace(state.cov, state.means, [0])
     #         mean1 = state.number_means[0]
     #         mean2 = state.number_means[1]