Adding __matmul__ to CircuitComponents

mrmustard/lab_dev/__init__.py

@@ -11,3 +11,10 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+
+from .circuit_components import CircuitComponent
+from .circuits import Circuit
+from .simulator import Simulator
+from .states import *
+from .transformations import *
+from .wires import Wires

mrmustard/lab_dev/circuit_components.py

@@ -20,7 +20,7 @@
 
 from typing import Optional, Sequence, Union
 
-from ..physics.representations import Bargmann, Representation
+from ..physics.representations import Bargmann, Fock, Representation
 from ..math.parameter_set import ParameterSet
 from ..math.parameters import Constant, Variable
 from ..utils.typing import Batch, ComplexMatrix, ComplexTensor, ComplexVector, Mode
@@ -140,6 +140,7 @@ def wires(self) -> Wires:
         """
         return self._wires
 
+    @property
     def adjoint(self) -> CircuitComponent:
         r"""
         Light-copies this component, then returns the adjoint of it, obtained by taking the
@@ -151,6 +152,7 @@ def adjoint(self) -> CircuitComponent:
         representation = ret.representation.conj()
         return CircuitComponent.from_attributes(name, wires, representation)
 
+    @property
     def dual(self) -> CircuitComponent:
         r"""
         Light-copies this component, then returns the dual of it, obtained by taking the
@@ -172,6 +174,14 @@ def light_copy(self) -> CircuitComponent:
         instance.__dict__["_wires"] = self.wires.copy()
         return instance
 
+    def __eq__(self, other) -> bool:
+        r"""
+        Whether this component is equal to another component.
+
+        Compares representations and wires, but not the other attributes (including name and parameter set).
+        """
+        return self.representation == other.representation and self.wires == other.wires
+
     def __getitem__(self, idx: Union[Mode, Sequence[Mode]]):
         r"""
         Returns a slice of this component for the given modes.
@@ -181,6 +191,48 @@ def __getitem__(self, idx: Union[Mode, Sequence[Mode]]):
         ret._parameter_set = self.parameter_set
         return ret
 
+    def __matmul__(self, other: CircuitComponent) -> CircuitComponent:
+        r"""
+        Contracts ``self`` and ``other``, without adding adjoints.
+        """
+        # initialized the ``Wires`` of the returned component
+        wires_ret = self.wires @ other.wires
+
+        # find the indices of the wires being contracted on the bra side
+        bra_modes = set(self.wires.bra.output.modes).intersection(other.wires.bra.input.modes)
+        idx_z = self.wires[bra_modes].bra.output.indices
+        idx_zconj = other.wires[bra_modes].bra.input.indices
+
+        # find the indices of the wires being contracted on the ket side
+        ket_modes = set(self.wires.ket.output.modes).intersection(other.wires.ket.input.modes)
+        idx_z += self.wires[ket_modes].ket.output.indices
+        idx_zconj += other.wires[ket_modes].ket.input.indices
+
+        # convert Bargmann -> Fock if needed
+        LEFT = self.representation
+        RIGHT = other.representation
+        if isinstance(LEFT, Bargmann) and isinstance(RIGHT, Fock):
+            raise ValueError("Cannot contract objects with different representations.")
+            # shape = [s if i in idx_z else None for i, s in enumerate(other.representation.shape)]
+            # LEFT = Fock(self.fock(shape=shape), batched=False)
+        elif isinstance(LEFT, Fock) and isinstance(RIGHT, Bargmann):
+            raise ValueError("Cannot contract objects with different representations.")
+            # shape = [s if i in idx_zconj else None for i, s in enumerate(self.representation.shape)]
+            # RIGHT = Fock(other.fock(shape=shape), batched=False)
+
+        # calculate the representation of the returned component
+        representation_ret = LEFT[idx_z] @ RIGHT[idx_zconj]
+
+        # reorder the representation
+        contracted_idx = [self.wires.ids[i] for i in range(len(self.wires.ids)) if i not in idx_z]
+        contracted_idx += [
+            other.wires.ids[i] for i in range(len(other.wires.ids)) if i not in idx_zconj
+        ]
+        order = [contracted_idx.index(id) for id in wires_ret.ids]
+        representation_ret = representation_ret.reorder(order)
+
+        return CircuitComponent.from_attributes("", wires_ret, representation_ret)
+
 
 def connect(components: Sequence[CircuitComponent]) -> Sequence[CircuitComponent]:
     r"""
@@ -227,6 +279,6 @@ def add_bra(components: Sequence[CircuitComponent]) -> Sequence[CircuitComponent
     for component in components:
         ret.append(component.light_copy())
         if not component.wires.bra:
-            ret.append(component.adjoint())
+            ret.append(component.adjoint)
 
     return ret

mrmustard/lab_dev/simulator.py

@@ -92,7 +92,7 @@ def _contract(self, components: Sequence[CircuitComponent]) -> CircuitComponent:
         # calculate the ``Wires`` of the returned component, alongside its substrings
         wires_out = components[0].wires
         for c in components[1:]:
-            wires_out >>= c.wires
+            wires_out @= c.wires
         subs_out = "".join([ids_to_subs[id] for id in wires_out.ids])
 
         # grab the representation that remains in ``subs_to_rep``

mrmustard/lab_dev/transformations.py

@@ -185,15 +185,15 @@ def representation(self) -> Bargmann:
         A = math.cast(
             np.array(
                 [
-                    [0, np.sqrt(e), 1 - e, 0],
-                    [np.sqrt(e), 0, 0, 0],
-                    [1 - e, 0, 0, np.sqrt(e)],
-                    [0, 0, np.sqrt(e), 0],
+                    [0, np.sqrt(e), 0, 0],
+                    [np.sqrt(e), 0, 0, 1 - e],
+                    [0, 0, 0, np.sqrt(e)],
+                    [0, 1 - e, np.sqrt(e), 0],
                 ],
             ),
             math.complex128,
         )
         B = math.cast([0.0, 0.0, 0.0, 0.0], math.complex128)
-        C = math.cast(np.sqrt(e), math.complex128)
+        C = 1
 
         return Bargmann(A, B, C)

mrmustard/lab_dev/wires.py

@@ -300,6 +300,26 @@ def __bool__(self) -> bool:
         """
         return len(self.ids) > 0
 
+    def __eq__(self, other) -> bool:
+        r"""
+        Returns ``True`` if this ``Wires`` acts on the same modes as ``other``, ``False`` otherwise.
+        """
+        if self.output.bra.modes != other.output.bra.modes:
+            return False
+        if self.input.bra.modes != other.input.bra.modes:
+            return False
+        if self.output.ket.modes != other.output.ket.modes:
+            return False
+        if self.input.ket.modes != other.input.ket.modes:
+            return False
+        return True
+
+    def __neq__(self, other) -> bool:
+        r"""
+        Returns ``False`` if this ``Wires`` is equal to ``other``, ``True`` otherwise.
+        """
+        return not self == other
+
     def __getitem__(self, modes: Iterable[int] | int) -> Wires:
         r"""
         A view of this Wires object with wires only on the given modes.
@@ -308,9 +328,6 @@ def __getitem__(self, modes: Iterable[int] | int) -> Wires:
         idxs = tuple(list(self._modes).index(m) for m in set(self._modes).difference(modes))
         return self._view(masked_rows=idxs)
 
-    def __lshift__(self, other: Wires) -> Wires:
-        return (other.dual >> self.dual).dual  # how cool is this
-
     @staticmethod
     def _outin(self_in: int, self_out: int, other_in: int, other_out: int) -> np.ndarray:
         r"""
@@ -331,12 +348,17 @@ def _outin(self_in: int, self_out: int, other_in: int, other_out: int) -> np.nda
         else:  # no wires on other
             return np.array([self_out, self_in], dtype=np.int64)
 
-    def __rshift__(self, other: Wires) -> Wires:
+    def __matmul__(self, other: Wires) -> Wires:
         r"""
-        A new Wires object with the wires of ``self`` and ``other`` combined as two
-        components in a circuit: the output of self connects to the input of other wherever
-        they match. All surviving wires are arranged in the standard order.
-        A ValueError is raised if there are any surviving wires that overlap.
+        A new ``Wires`` object with the wires of ``self`` and ``other`` combined.
+
+        The output of ``self`` connects to the input of ``other`` wherever they match. All
+        surviving wires are arranged in the standard order.
+
+        This function does not add missing adjoints.
+
+        Raises:
+            ValueError: If there are any surviving wires that overlap.
         """
         all_modes = sorted(set(self.modes) | set(other.modes))
         new_id_array = np.zeros((len(all_modes), 4), dtype=np.int64)

mrmustard/physics/triples.py

@@ -392,7 +392,7 @@ def attenuator_Abc(eta: Union[Scalar, Iterable]) -> Union[Matrix, Vector, Scalar
 
     A = A[reshape_list, :][:, reshape_list]
     b = _vacuum_B_vector(n_modes * 2)
-    c = np.prod(eta)
+    c = 1.0
 
     return A, b, c
 

tests/test_lab_dev/test_circuit_components.py

@@ -16,9 +16,13 @@
 Tests for circuit components.
 """
 
+import numpy as np
+
+from mrmustard.physics.representations import Bargmann
 from mrmustard.lab_dev.circuit_components import connect, add_bra, CircuitComponent
 from mrmustard.lab_dev.states import Vacuum
 from mrmustard.lab_dev.transformations import Dgate, Attenuator
+from mrmustard.lab_dev.wires import Wires
 
 
 class TestCircuitComponent:
@@ -37,6 +41,79 @@ def test_light_copy(self):
         assert d.y is d_copy.y
         assert d.wires is not d_copy.wires
 
+    def test_matmul_one_mode(self):
+        r"""
+        Tests that ``__matmul__`` produces the correct outputs for one-mode components.
+        """
+        vac0 = Vacuum(1)
+        d0 = Dgate(1, modes=[0])
+        a0 = Attenuator(0.9, modes=[0])
+
+        result1 = vac0 @ d0
+        result1 = (result1 @ result1.adjoint) @ a0
+
+        assert result1.wires == Wires(modes_out_bra=[0], modes_out_ket=[0])
+        assert np.allclose(result1.representation.A, 0)
+        assert np.allclose(result1.representation.b, [0.9486833, 0.9486833])
+        assert np.allclose(result1.representation.c, 0.40656966)
+
+        result2 = result1 @ vac0.dual @ vac0.dual.adjoint
+        assert not result2.wires
+        assert np.allclose(result2.representation.A, 0)
+        assert np.allclose(result2.representation.b, 0)
+        assert np.allclose(result2.representation.c, 0.40656966)
+
+    def test_matmul_multi_modes(self):
+        r"""
+        Tests that ``__matmul__`` produces the correct outputs for multi-mode components.
+        """
+        vac012 = Vacuum(3)
+        d0 = Dgate(0.1, 0.1, modes=[0])
+        d1 = Dgate(0.1, 0.1, modes=[1])
+        d2 = Dgate(0.1, 0.1, modes=[2])
+        a0 = Attenuator(0.8, modes=[0])
+        a1 = Attenuator(0.8, modes=[1])
+        a2 = Attenuator(0.7, modes=[2])
+
+        result = vac012 @ d0 @ d1 @ d2
+        result = result @ result.adjoint @ a0 @ a1 @ a2
+
+        assert result.wires == Wires(modes_out_bra=[0, 1, 2], modes_out_ket=[0, 1, 2])
+        assert np.allclose(result.representation.A, 0)
+        assert np.allclose(
+            result.representation.b,
+            [
+                0.08944272 - 0.08944272j,
+                0.08944272 - 0.08944272j,
+                0.083666 - 0.083666j,
+                0.08944272 + 0.08944272j,
+                0.08944272 + 0.08944272j,
+                0.083666 + 0.083666j,
+            ],
+        )
+        assert np.allclose(result.representation.c, 0.95504196)
+
+    def test_matmul_is_associative(self):
+        r"""
+        Tests that ``__matmul__`` is associative, meaning ``a @ (b @ c) == (a @ b) @ c``.
+        """
+        vac012 = Vacuum(3)
+        d0 = Dgate(0.1, 0.1, modes=[0])
+        d1 = Dgate(0.1, 0.1, modes=[1])
+        d2 = Dgate(0.1, 0.1, modes=[2])
+        a0 = Attenuator(0.8, modes=[0])
+        a1 = Attenuator(0.8, modes=[1])
+        a2 = Attenuator(0.7, modes=[2])
+
+        result1 = vac012 @ d0 @ d1 @ a0 @ a1 @ a2 @ d2
+        result2 = (vac012 @ d0) @ (d1 @ a0) @ a1 @ (a2 @ d2)
+        result3 = vac012 @ (d0 @ (d1 @ a0 @ a1) @ a2 @ d2)
+        result4 = vac012 @ (d0 @ (d1 @ (a0 @ (a1 @ (a2 @ d2)))))
+
+        assert result1 == result2
+        assert result1 == result3
+        assert result1 == result4
+
 
 class TestConnect:
     r"""
@@ -82,7 +159,7 @@ def test_ket_and_bra(self):
         Tests the ``connect`` function with components with kets and bras.
         """
         d1 = Dgate(1, modes=[0, 8, 9])
-        d1_adj = d1.adjoint()
+        d1_adj = d1.adjoint
         a1 = Attenuator(0.1, modes=[8])
 
         components = connect([d1, d1_adj, a1])

tests/test_lab_dev/test_wires.py

@@ -167,15 +167,30 @@ def test_getitem(self):
         assert w1.modes == [1]
         assert w1.ids == [w.ids[1], w.ids[3]]
 
-    def test_rshift(self):
+    def test_eq_neq(self):
+        w1 = Wires([0, 1], [2, 3], [4, 5], [6, 7])
+        w2 = Wires([0, 1], [2, 3], [4, 5], [6, 7])
+        w3 = Wires([], [2, 3], [4, 5], [6, 7])
+        w4 = Wires([0, 1], [], [4, 5], [6, 7])
+        w5 = Wires([0, 1], [2, 3], [], [6, 7])
+        w6 = Wires([0, 1], [2, 3], [4, 5], [])
+
+        assert w1 == w1
+        assert w1 == w2
+        assert w1 != w3
+        assert w1 != w4
+        assert w1 != w5
+        assert w1 != w6
+
+    def test_matmul(self):
         # contracts 1,1 on bra side
         # contracts 3,3 and 13,13 on ket side (note order doesn't matter)
         u = Wires([1, 5], [2, 6, 15], [3, 7, 13], [4, 8])
         v = Wires([0, 9, 14], [1, 10], [2, 11], [13, 3, 12])
-        assert (u >> v)._args() == ((0, 5, 9, 14), (2, 6, 10, 15), (2, 7, 11), (4, 8, 12))
+        assert (u @ v)._args() == ((0, 5, 9, 14), (2, 6, 10, 15), (2, 7, 11), (4, 8, 12))
 
-    def test_rshift_error(self):
+    def test_matmul_error(self):
         u = Wires([], [], [0], [])  # only output wire
         v = Wires([], [], [0], [])  # only output wire
         with pytest.raises(ValueError):
-            u >> v  # pylint: disable=pointless-statement
+            u @ v  # pylint: disable=pointless-statement

tests/test_physics/test_triples.py

@@ -220,7 +220,7 @@ def test_attenuator_Abc(self):
         e = 0.31622777
         assert np.allclose(A1, [[0, e, 0, 0], [e, 0, 0, 0.9], [0, 0, 0, e], [0, 0.9, e, 0]])
         assert np.allclose(b1, 0)
-        assert np.allclose(c1, 0.1)
+        assert np.allclose(c1, 1.0)
 
         A2, b2, c2 = triples.attenuator_Abc([0.1, 1])
         e = 0.31622777
@@ -238,7 +238,7 @@ def test_attenuator_Abc(self):
             ],
         )
         assert np.allclose(b2, 0)
-        assert np.allclose(c2, 0.1)
+        assert np.allclose(c2, 1.0)
 
     def test_attenuator_Abc_error(self):
         with pytest.raises(ValueError, match="in the interval"):