tests

mrmustard/lab_dev/states/base.py

@@ -847,7 +847,7 @@ def expectation(self, operator: CircuitComponent):
             msg = f"Expected an observable defined for modes `{self.modes}` or a subset thereof, "
             msg += f"found one defined for modes `{operator.modes}.`"
             raise ValueError(msg)
-        
+
         result = (self @ operator @ self.dual).representation
         return result.array if isinstance(result, Fock) else result.c
 

mrmustard/physics/triples.py

@@ -345,6 +345,30 @@ def beamsplitter_gate_Abc(
     return A, b, c
 
 
+def identity_Abc(n_modes: int) -> Union[Matrix, Vector, Scalar]:
+    r"""
+    The ``(A, b, c)`` triple of a tensor product of identity gates.
+
+    Args:
+        n_modes: The number of modes.
+
+    Returns:
+        The ``(A, b, c)`` triple of the identities.
+    """
+    O_n = math.zeros((n_modes, n_modes), math.complex128)
+    I_n = (
+        math.diag(math.astensor([1.0 + 0j for _ in range(n_modes)]))
+        .reshape((n_modes, n_modes))
+        .reshape((n_modes, n_modes))
+    )
+
+    A = math.block([[O_n, I_n], [I_n, O_n]])
+    b = _vacuum_B_vector(n_modes * 2)
+    c = 1.0 + 0j
+
+    return A, b, c
+
+
 # ~~~~~~~~~~
 #  Channels
 # ~~~~~~~~~~
@@ -437,38 +461,6 @@ def amplifier_Abc(g: Union[float, Iterable[float]]) -> Union[Matrix, Vector, Sca
     return A, b, c
 
 
-def identity_Abc(n_modes: int) -> Union[Matrix, Vector, Scalar]:
-    r"""
-    The ``(A, b, c)`` triple of a tensor product of identity channels
-
-    Args:
-        n_modes: The number of modes.
-
-    Returns:
-        The ``(A, b, c)`` triple of the Fock damping channels.
-    """
-    O_n = math.zeros((n_modes, n_modes), math.complex128)
-    I_n = (
-        math.diag(math.astensor([1.0 + 0j for _ in range(n_modes)]))
-        .reshape((n_modes, n_modes))
-        .reshape((n_modes, n_modes))
-    )
-
-    A = math.block(
-        [
-            [O_n, I_n, O_n, O_n],
-            [I_n, O_n, O_n, O_n],
-            [O_n, O_n, O_n, I_n],
-            [O_n, O_n, I_n, O_n],
-        ]
-    )
-
-    b = _vacuum_B_vector(n_modes * 4)
-    c = 1.0 + 0j
-
-    return A, b, c
-
-
 def fock_damping_Abc(n_modes: int) -> Union[Matrix, Vector, Scalar]:
     r"""
     The ``(A, b, c)`` triple of a tensor product of Fock dampers.

tests/test_lab_dev/test_circuit_components_utils.py

@@ -18,13 +18,16 @@
 
 import pytest
 
+from mrmustard import settings
 from mrmustard.physics.triples import identity_Abc
 from mrmustard.physics.representations import Bargmann
 from mrmustard.lab_dev.circuit_components_utils import TraceOut
 from mrmustard.lab_dev.states import Coherent
-from mrmustard.lab_dev.transformations import Dgate, Attenuator, Unitary
 from mrmustard.lab_dev.wires import Wires
 
+# original settings
+autocutoff_max0 = settings.AUTOCUTOFF_MAX_CUTOFF
+
 
 class TestTraceOut:
     r"""
@@ -39,6 +42,16 @@ def test_init(self, modes):
         assert tr.wires == Wires(modes_in_bra=set(modes), modes_in_ket=set(modes))
         assert tr.representation == Bargmann(*identity_Abc(len(modes)))
 
-    def test_trace_out_states(self):
-        assert Coherent([0, 1, 2], x=1) >> TraceOut([0]) == Coherent([1, 2], x=1).dm()
-        assert Coherent([0, 1, 2], x=1) >> TraceOut([1, 2]) == Coherent([0], x=1).dm()
+    def test_trace_out_bargmann_states(self):
+        state = Coherent([0, 1, 2], x=1)
+        assert state >> TraceOut([0]) == Coherent([1, 2], x=1).dm()
+        assert state >> TraceOut([1, 2]) == Coherent([0], x=1).dm()
+
+    def test_trace_out_fock_states(self):
+        settings.AUTOCUTOFF_MAX_CUTOFF = 10
+
+        state = Coherent([0, 1, 2], x=1).to_fock_component()
+        assert state >> TraceOut([0]) == Coherent([1, 2], x=1).to_fock_component().dm()
+        assert state >> TraceOut([1, 2]) == Coherent([0], x=1).to_fock_component().dm()
+
+        settings.AUTOCUTOFF_MAX_CUTOFF = autocutoff_max0

tests/test_lab_dev/test_states.py

@@ -169,30 +169,30 @@ def test_expectation_bargmann(self):
         res0 = ((dm @ op0) >> TraceOut(dm.modes)).representation.c
         res1 = ((dm @ op1) >> TraceOut(dm.modes)).representation.c
         res01 = ((dm @ op01) >> TraceOut(dm.modes)).representation.c
-        
-        assert math.allclose(ket.expectation(op0), res0) 
-        assert math.allclose(ket.expectation(op1), res1) 
+
+        assert math.allclose(ket.expectation(op0), res0)
+        assert math.allclose(ket.expectation(op1), res1)
         assert math.allclose(ket.expectation(op01), res01)
 
-    # def test_expectation_fock(self):
-    #     settings.AUTOCUTOFF_MAX_CUTOFF = 10
+    def test_expectation_fock(self):
+        settings.AUTOCUTOFF_MAX_CUTOFF = 10
 
-    #     ket = Coherent([0, 1], x=1, y=[2, 3]).to_fock_component()
-    #     dm = ket.dm()
+        ket = Coherent([0, 1], x=1, y=[2, 3]).to_fock_component()
+        dm = ket.dm()
 
-    #     op0 = Dgate([1], x=0.1)
-    #     op1 = Dgate([0], x=0.2)
-    #     op01 = Dgate([0, 1], x=[0.3, 0.4])
+        op0 = Dgate([1], x=0.1)
+        op1 = Dgate([0], x=0.2)
+        op01 = Dgate([0, 1], x=[0.3, 0.4])
 
-    #     res0 = ((dm @ op0) >> TraceOut(dm.modes)).representation.array
-    #     res1 = ((dm @ op1) >> TraceOut(dm.modes)).representation.array
-    #     res01 = ((dm @ op01) >> TraceOut(dm.modes)).representation.array
-
-    #     assert math.allclose(ket.expectation(op0), res0) 
-    #     assert math.allclose(ket.expectation(op1), res1) 
-    #     assert math.allclose(ket.expectation(op01), res01)
+        res0 = ((dm @ op0) >> TraceOut(dm.modes)).representation.array
+        res1 = ((dm @ op1) >> TraceOut(dm.modes)).representation.array
+        res01 = ((dm @ op01) >> TraceOut(dm.modes)).representation.array
 
-    #     settings.AUTOCUTOFF_MAX_CUTOFF = autocutoff_max0
+        assert math.allclose(ket.expectation(op0), res0)
+        assert math.allclose(ket.expectation(op1), res1)
+        assert math.allclose(ket.expectation(op01), res01)
+
+        settings.AUTOCUTOFF_MAX_CUTOFF = autocutoff_max0
 
     def test_expectation_error(self):
         ket = Coherent([0, 1], x=1, y=[2, 3])
@@ -209,8 +209,6 @@ def test_expectation_error(self):
         with pytest.raises(ValueError, match="Expected an observable defined for modes"):
             ket.expectation(op3)
 
-
-
     def test_rshift(self):
         ket = Coherent([0, 1], 1)
         unitary = Dgate([0], 1)
@@ -360,6 +358,57 @@ def test_purity(self):
         assert math.allclose(state.purity, 1)
         assert state.is_pure
 
+    def test_expectation_bargmann(self):
+        ket = Coherent([0, 1], x=1, y=[2, 3])
+        dm = ket.dm()
+
+        op0 = Dgate([1], x=0.1)
+        op1 = Dgate([0], x=0.2)
+        op01 = Dgate([0, 1], x=[0.3, 0.4])
+
+        res0 = ((dm @ op0) >> TraceOut(dm.modes)).representation.c
+        res1 = ((dm @ op1) >> TraceOut(dm.modes)).representation.c
+        res01 = ((dm @ op01) >> TraceOut(dm.modes)).representation.c
+
+        assert math.allclose(dm.expectation(op0), res0)
+        assert math.allclose(dm.expectation(op1), res1)
+        assert math.allclose(dm.expectation(op01), res01)
+
+    def test_expectation_fock(self):
+        settings.AUTOCUTOFF_MAX_CUTOFF = 10
+
+        ket = Coherent([0, 1], x=1, y=[2, 3]).to_fock_component()
+        dm = ket.dm()
+
+        op0 = Dgate([1], x=0.1)
+        op1 = Dgate([0], x=0.2)
+        op01 = Dgate([0, 1], x=[0.3, 0.4])
+
+        res0 = ((dm @ op0) >> TraceOut(dm.modes)).representation.array
+        res1 = ((dm @ op1) >> TraceOut(dm.modes)).representation.array
+        res01 = ((dm @ op01) >> TraceOut(dm.modes)).representation.array
+
+        assert math.allclose(dm.expectation(op0), res0)
+        assert math.allclose(dm.expectation(op1), res1)
+        assert math.allclose(dm.expectation(op01), res01)
+
+        settings.AUTOCUTOFF_MAX_CUTOFF = autocutoff_max0
+
+    def test_expectation_error(self):
+        dm = Coherent([0, 1], x=1, y=[2, 3]).dm()
+
+        op1 = Attenuator([0])
+        with pytest.raises(ValueError, match="wires on the bra side"):
+            dm.expectation(op1)
+
+        op2 = CircuitComponent("", None, modes_in_ket=[0], modes_out_ket=[1])
+        with pytest.raises(ValueError, match="different modes"):
+            dm.expectation(op2)
+
+        op3 = Dgate([2])
+        with pytest.raises(ValueError, match="Expected an observable defined for modes"):
+            dm.expectation(op3)
+
     def test_rshift(self):
         ket = Coherent([0, 1], 1)
         unitary = Dgate([0], 1)

tests/test_physics/test_triples.py

@@ -222,6 +222,17 @@ def test_beamsplitter_gate_Abc(self):
         assert math.allclose(b3, np.zeros((4)))
         assert math.allclose(c3, 1)
 
+    def test_identity_Abc(self):
+        A1, b1, c1 = triples.identity_Abc(1)
+        assert math.allclose(A1, [[0, 1], [1, 0]])
+        assert math.allclose(b1, [0, 0])
+        assert math.allclose(c1, 1)
+
+        A2, b2, c2 = triples.identity_Abc(2)
+        assert math.allclose(A2, [[0, 0, 1, 0], [0, 0, 0, 1], [1, 0, 0, 0], [0, 1, 0, 0]])
+        assert math.allclose(b2, [0, 0, 0, 0])
+        assert math.allclose(c2, 1)
+
     def test_attenuator_Abc(self):
         A1, b1, c1 = triples.attenuator_Abc(0.1)
         e = 0.31622777