more progress

mrmustard/lab_dev/states/__init__.py

@@ -23,7 +23,7 @@
 from .bargmann_eigenstate import BargmannEigenstate
 from .coherent import Coherent
 from .displaced_squeezed import DisplacedSqueezed
-from .gstate import GKet, GDM
+from .gaussian_state import GKet, GDM
 from .number import Number
 from .quadrature_eigenstate import QuadratureEigenstate
 from .squeezed_vacuum import SqueezedVacuum

mrmustard/lab_dev/transformations/rgate.py

@@ -40,7 +40,7 @@ class Rgate(Unitary):
         >>> import numpy as np
         >>> from mrmustard.lab_dev import Rgate
 
-        >>> unitary = Rgate(modes=[1, 2], phi=0.1)
+        >>> unitary = Rgate(modes=[1, 2], theta=0.1)
         >>> assert unitary.modes == [1, 2]
 
     Args:
@@ -55,15 +55,15 @@ class Rgate(Unitary):
     def __init__(
         self,
         modes: Sequence[int],
-        phi: float | Sequence[float] = 0.0,
-        phi_trainable: bool = False,
-        phi_bounds: tuple[float | None, float | None] = (0.0, None),
+        theta: float | Sequence[float] = 0.0,
+        theta_trainable: bool = False,
+        theta_bounds: tuple[float | None, float | None] = (0.0, None),
     ):
         super().__init__(name="Rgate")
-        (phis,) = list(reshape_params(len(modes), phi=phi))
-        self._add_parameter(make_parameter(phi_trainable, phis, "phi", phi_bounds))
+        (thetas,) = list(reshape_params(len(modes), theta=theta))
+        self._add_parameter(make_parameter(theta_trainable, thetas, "theta", theta_bounds))
         self._representation = self.from_ansatz(
             modes_in=modes,
             modes_out=modes,
-            ansatz=PolyExpAnsatz.from_function(fn=triples.rotation_gate_Abc, theta=self.phi),
+            ansatz=PolyExpAnsatz.from_function(fn=triples.rotation_gate_Abc, theta=self.theta),
         ).representation

tests/test_lab_dev/test_transformations/test_rgate.py

@@ -29,21 +29,21 @@ class TestRgate:
     """
 
     modes = [[0], [1, 2], [9, 7]]
-    phis = [[1], 1, [1, 2]]
+    thetas = [[1], 1, [1, 2]]
 
-    @pytest.mark.parametrize("modes,phi", zip(modes, phis))
-    def test_init(self, modes, phi):
-        gate = Rgate(modes, phi)
+    @pytest.mark.parametrize("modes,theta", zip(modes, thetas))
+    def test_init(self, modes, theta):
+        gate = Rgate(modes, theta)
 
         assert gate.name == "Rgate"
         assert gate.modes == [modes] if not isinstance(modes, list) else sorted(modes)
 
     def test_init_error(self):
-        with pytest.raises(ValueError, match="phi"):
-            Rgate(modes=[0, 1], phi=[2, 3, 4])
+        with pytest.raises(ValueError, match="theta"):
+            Rgate(modes=[0, 1], theta=[2, 3, 4])
 
     def test_representation(self):
-        rep1 = Rgate(modes=[0], phi=0.1).ansatz
+        rep1 = Rgate(modes=[0], theta=0.1).ansatz
         assert math.allclose(
             rep1.A,
             [
@@ -56,7 +56,7 @@ def test_representation(self):
         assert math.allclose(rep1.b, np.zeros((1, 2)))
         assert math.allclose(rep1.c, [1.0 + 0.0j])
 
-        rep2 = Rgate(modes=[0, 1], phi=[0.1, 0.3]).ansatz
+        rep2 = Rgate(modes=[0, 1], theta=[0.1, 0.3]).ansatz
         assert math.allclose(
             rep2.A,
             [
@@ -71,7 +71,7 @@ def test_representation(self):
         assert math.allclose(rep2.b, np.zeros((1, 4)))
         assert math.allclose(rep2.c, [1.0 + 0.0j])
 
-        rep3 = Rgate(modes=[1], phi=0.1).ansatz
+        rep3 = Rgate(modes=[1], theta=0.1).ansatz
         assert math.allclose(
             rep3.A,
             [
@@ -89,11 +89,11 @@ def test_trainable_parameters(self):
         gate2 = Rgate([0], 1, True, (-2, 2))
 
         with pytest.raises(AttributeError):
-            gate1.phi.value = 3
+            gate1.theta.value = 3
 
-        gate2.phi.value = 2
-        assert gate2.phi.value == 2
+        gate2.theta.value = 2
+        assert gate2.theta.value == 2
 
     def test_representation_error(self):
         with pytest.raises(ValueError):
-            Rgate(modes=[0], phi=[0.1, 0.2]).ansatz
+            Rgate(modes=[0], theta=[0.1, 0.2]).ansatz

tests/test_training/test_opt.py

@@ -21,15 +21,23 @@
 from thewalrus.symplectic import two_mode_squeezing
 
 from mrmustard import math, settings
-from mrmustard.lab_dev import Circuit, Sgate, S2gate, Vacuum, BSgate, Ggate, Interferometer
-from mrmustard.lab.gates import (
+from mrmustard.lab_dev import (
+    Circuit,
+    Sgate,
+    S2gate,
+    Vacuum,
+    BSgate,
+    Ggate,
+    Interferometer,
+    Rgate,
     Dgate,
-    # Ggate,
-    # Interferometer,
     RealInterferometer,
-    Rgate,
+    DisplacedSqueezed,
+    SqueezedVacuum,
+    GKet,
+    Number,
 )
-from mrmustard.lab.states import DisplacedSqueezed, Fock, Gaussian, SqueezedVacuum
+from mrmustard.lab.states import Fock, Gaussian
 from mrmustard.math.parameters import Variable, update_euclidean
 from mrmustard.physics import fidelity
 from mrmustard.physics.gaussian import trace, von_neumann_entropy
@@ -368,31 +376,31 @@ def cost_fn():
         opt.minimize(cost_fn, by_optimizing=[circ], max_steps=300)
         assert np.allclose(np.sinh(S_12.r.value) ** 2, 1, atol=1e-2)
 
-    # def test_parameter_passthrough(self):
-    #     """Same as the test above, but with param passthrough"""
-    #     skip_np()
+    def test_parameter_passthrough(self):
+        """Same as the test above, but with param passthrough"""
+        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)
 
-    #     r = np.arcsinh(1.0)
-    #     r_var = Variable(r, "r", (0.0, None))
-    #     phi_var = Variable(settings.rng.normal(), "phi", (None, None))
+        r = np.arcsinh(1.0)
+        r_var = Variable(r, "r", (0.0, None))
+        phi_var = Variable(settings.rng.normal(), "phi", (None, None))
 
-    #     ops = [
-    #         S2gate(r=r, phi=0.0, phi_trainable=True)[0, 1],
-    #         S2gate(r=r, phi=0.0, phi_trainable=True)[2, 3],
-    #         S2gate(r=r_var, phi=phi_var)[1, 2],
-    #     ]
-    #     circ = Circuit(ops)
+        state_in = Vacuum((0, 1, 2, 3))
+        s2_gate0 = S2gate((0, 1), r=r, phi=0.0, phi_trainable=True)
+        s2_gate1 = S2gate((2, 3), r=r, phi=0.0, phi_trainable=True)
+        s2_gate2 = S2gate((1, 2), r=r_var, phi=phi_var)
 
-    #     def cost_fn():
-    #         return math.abs((Vacuum(4) >> circ).ket(cutoffs=[2, 2, 2, 2])[1, 1, 1, 1]) ** 2
+        circ = Circuit([state_in, s2_gate0, s2_gate1, s2_gate2])
+
+        def cost_fn():
+            return math.abs(circ.contract().fock_array((2, 2, 2, 2))[1, 1, 1, 1]) ** 2
 
-    #     opt = Optimizer(euclidean_lr=0.001)
-    #     opt.minimize(cost_fn, by_optimizing=[r_var, phi_var], max_steps=300)
-    #     assert np.allclose(np.sinh(r_var.value) ** 2, 1, atol=1e-2)
+        opt = Optimizer(euclidean_lr=0.001)
+        opt.minimize(cost_fn, by_optimizing=[r_var, phi_var], max_steps=300)
+        assert np.allclose(np.sinh(r_var.value) ** 2, 1, atol=1e-2)
 
     # def test_making_thermal_state_as_one_half_two_mode_squeezed_vacuum(self):
     #     """Optimizes a Ggate on two modes so as to prepare a state with the same entropy
@@ -437,97 +445,100 @@ def cost_fn():
     #     rng.reset_from_seed(settings.SEED)
 
     #     rotation_angle = np.pi / 2
-    #     target_state = SqueezedVacuum(r=1.0, phi=rotation_angle)
+    #     target_state = SqueezedVacuum((0,), r=1.0, phi=rotation_angle)
 
     #     # angle of rotation gate
     #     r_angle = math.new_variable(0, bounds=(0, np.pi), name="r_angle")
     #     # trainable squeezing
-    #     S = Sgate(r=0.1, phi=0, r_trainable=True, phi_trainable=False)
+    #     S = Sgate((0,), r=0.1, phi=0, r_trainable=True, phi_trainable=False)
 
     #     def cost_fn_sympl():
-    #         state_out = Vacuum(1) >> S >> Rgate(angle=r_angle)
+    #         state_out = Vacuum((0,)) >> S >> Rgate((0,), theta=r_angle)
+    #         # TODO: fidelity
     #         return 1 - fidelity(state_out, target_state)
 
     #     opt = Optimizer(symplectic_lr=0.1, euclidean_lr=0.05)
     #     opt.minimize(cost_fn_sympl, by_optimizing=[S, r_angle])
 
     #     assert np.allclose(math.asnumpy(r_angle), rotation_angle / 2, atol=1e-4)
 
-    # def test_dgate_optimization(self):
-    #     """Test that Dgate is optimized correctly."""
-    #     skip_np()
+    def test_dgate_optimization(self):
+        """Test that Dgate is optimized correctly."""
+        skip_np()
 
-    #     settings.SEED = 24
-    #     rng = tf.random.get_global_generator()
-    #     rng.reset_from_seed(settings.SEED)
+        settings.SEED = 24
+        rng = tf.random.get_global_generator()
+        rng.reset_from_seed(settings.SEED)
 
-    #     dgate = Dgate(x_trainable=True, y_trainable=True)
-    #     target_state = DisplacedSqueezed(r=0.0, x=0.1, y=0.2).ket(cutoffs=[40])
+        dgate = Dgate((0,), x_trainable=True, y_trainable=True)
+        target_state = DisplacedSqueezed((0,), r=0.0, x=0.1, y=0.2).fock_array((40,))
 
-    #     def cost_fn():
-    #         state_out = Vacuum(1) >> dgate
-    #         return -math.abs(math.sum(math.conj(state_out.ket([40])) * target_state)) ** 2
+        def cost_fn():
+            state_out = Vacuum((0,)) >> dgate
+            return -math.abs(math.sum(math.conj(state_out.fock_array((40,))) * target_state)) ** 2
 
-    #     opt = Optimizer()
-    #     opt.minimize(cost_fn, by_optimizing=[dgate])
+        opt = Optimizer()
+        opt.minimize(cost_fn, by_optimizing=[dgate])
 
-    #     assert np.allclose(dgate.x.value, 0.1, atol=0.01)
-    #     assert np.allclose(dgate.y.value, 0.2, atol=0.01)
+        assert np.allclose(dgate.x.value, 0.1, atol=0.01)
+        assert np.allclose(dgate.y.value, 0.2, atol=0.01)
 
-    # def test_sgate_optimization(self):
-    #     """Test that Sgate is optimized correctly."""
-    #     skip_np()
+    def test_sgate_optimization(self):
+        """Test that Sgate is optimized correctly."""
+        skip_np()
 
-    #     settings.SEED = 25
-    #     rng = tf.random.get_global_generator()
-    #     rng.reset_from_seed(settings.SEED)
+        settings.SEED = 25
+        rng = tf.random.get_global_generator()
+        rng.reset_from_seed(settings.SEED)
 
-    #     sgate = Sgate(r=0.2, phi=0.1, r_trainable=True, phi_trainable=True)
-    #     target_state = SqueezedVacuum(r=0.1, phi=0.2).ket(cutoffs=[40])
+        sgate = Sgate((0,), r=0.2, phi=0.1, r_trainable=True, phi_trainable=True)
+        target_state = SqueezedVacuum((0,), r=0.1, phi=0.2).fock_array((40,))
 
-    #     def cost_fn():
-    #         state_out = Vacuum(1) >> sgate
+        def cost_fn():
+            state_out = Vacuum((0,)) >> sgate
 
-    #         return -math.abs(math.sum(math.conj(state_out.ket([40])) * target_state)) ** 2
+            return -math.abs(math.sum(math.conj(state_out.fock_array((40,))) * target_state)) ** 2
 
-    #     opt = Optimizer()
-    #     opt.minimize(cost_fn, by_optimizing=[sgate])
+        opt = Optimizer()
+        opt.minimize(cost_fn, by_optimizing=[sgate])
 
-    #     assert np.allclose(sgate.r.value, 0.1, atol=0.01)
-    #     assert np.allclose(sgate.phi.value, 0.2, atol=0.01)
+        assert np.allclose(sgate.r.value, 0.1, atol=0.01)
+        assert np.allclose(sgate.phi.value, 0.2, atol=0.01)
 
-    # def test_bsgate_optimization(self):
-    #     """Test that Sgate is optimized correctly."""
-    #     skip_np()
+    def test_bsgate_optimization(self):
+        """Test that Sgate is optimized correctly."""
+        skip_np()
 
-    #     settings.SEED = 25
-    #     rng = tf.random.get_global_generator()
-    #     rng.reset_from_seed(settings.SEED)
+        settings.SEED = 25
+        rng = tf.random.get_global_generator()
+        rng.reset_from_seed(settings.SEED)
 
-    #     G = Gaussian(2)
+        G = GKet((0, 1))
 
-    #     bsgate = BSgate(0.05, 0.1, theta_trainable=True, phi_trainable=True)
-    #     target_state = (G >> BSgate(0.1, 0.2)).ket(cutoffs=[40, 40])
+        bsgate = BSgate((0, 1), 0.05, 0.1, theta_trainable=True, phi_trainable=True)
+        target_state = (G >> BSgate((0, 1), 0.1, 0.2)).fock_array((40, 40))
 
-    #     def cost_fn():
-    #         state_out = G >> bsgate
+        def cost_fn():
+            state_out = G >> bsgate
 
-    #         return -math.abs(math.sum(math.conj(state_out.ket([40, 40])) * target_state)) ** 2
+            return (
+                -math.abs(math.sum(math.conj(state_out.fock_array((40, 40))) * target_state)) ** 2
+            )
 
-    #     opt = Optimizer()
-    #     opt.minimize(cost_fn, by_optimizing=[bsgate])
+        opt = Optimizer()
+        opt.minimize(cost_fn, by_optimizing=[bsgate])
 
-    #     assert np.allclose(bsgate.theta.value, 0.1, atol=0.01)
-    #     assert np.allclose(bsgate.phi.value, 0.2, atol=0.01)
+        assert np.allclose(bsgate.theta.value, 0.1, atol=0.01)
+        assert np.allclose(bsgate.phi.value, 0.2, atol=0.01)
 
     # def test_squeezing_grad_from_fock(self):
     #     """Test that the gradient of a squeezing gate is computed from the fock representation."""
     #     skip_np()
 
-    #     squeezing = Sgate(r=1, r_trainable=True)
+    #     squeezing = Sgate((0,), r=1, r_trainable=True)
 
     #     def cost_fn():
-    #         return -(Fock(2) >> squeezing << Vacuum(1))
+    #         return -(Number((0,), 2) >> squeezing >> Vacuum((0,)).dual)
 
     #     opt = Optimizer(euclidean_lr=0.05)
     #     opt.minimize(cost_fn, by_optimizing=[squeezing], max_steps=100)

tests/test_utils/test_serialize.py

@@ -180,7 +180,7 @@ def test_all_components_serializable(self):
                 BSgate([1, 2], theta=0.1, theta_trainable=True, theta_bounds=(-0.5, 0.5)),
                 Dgate([0], x=1.1, y=2.2),
                 Identity([1, 2]),
-                Rgate([1, 2], phi=0.1),
+                Rgate([1, 2], theta=0.1),
                 S2gate([0, 1], 1, 1),
                 Sgate([0], 0.1, 0.2, r_trainable=True),
                 FockDamping([0], damping=0.1),