using new overlap method

mrmustard/lab_dev/bb.py

@@ -28,7 +28,8 @@
 class GraphComponent:
     r"""A lightweight "CircuitComponent" without the actual representation.
     Basically a wrapper around Wires, so that it can emulate components in
-    a circuit. It exposes the repr, wires, shape, name and order (in the circuit).
+    a circuit. It exposes the repr, wires, shape, name and cost of obtaining
+    the component from previous contractions.
     """
 
     def __init__(self, repr: str, wires: Wires, shape: list[int], name: str = "", cost: int = 0):
@@ -207,8 +208,7 @@ def parse_components(components: list[CircuitComponent]) -> Graph:
         comp = GraphComponent.from_circuitcomponent(A)
         graph.add_node(i, component=comp.copy())
         for j, B in enumerate(components[i + 1 :]):
-            ovlp_ket = comp.wires.output.ket.modes & B.wires.input.ket.modes
-            ovlp_bra = comp.wires.output.bra.modes & B.wires.input.bra.modes
+            ovlp_bra, ovlp_ket = A.wires.overlap(B.wires)
             if ovlp_ket or ovlp_bra:
                 graph.add_edge(i, i + j + 1)
                 comp.wires = Wires(

mrmustard/lab_dev/wires.py

@@ -16,7 +16,6 @@
 
 from __future__ import annotations
 from functools import cached_property
-from typing import Optional, Sequence
 import numpy as np
 
 from IPython.display import display
@@ -194,7 +193,12 @@ def adjoint(self) -> Wires:
         New ``Wires`` object obtained by swapping ket and bra wires.
         """
         return Wires(
-            self.args[2], self.args[3], self.args[0], self.args[1], self.args[4], self.args[5]
+            self.args[2],
+            self.args[3],
+            self.args[0],
+            self.args[1],
+            self.args[4],
+            self.args[5],
         )
 
     @cached_property
@@ -221,7 +225,12 @@ def dual(self) -> Wires:
         New ``Wires`` object obtained by swapping input and output wires.
         """
         return Wires(
-            self.args[1], self.args[0], self.args[3], self.args[2], self.args[5], self.args[4]
+            self.args[1],
+            self.args[0],
+            self.args[3],
+            self.args[2],
+            self.args[5],
+            self.args[4],
         )
 
     @cached_property
@@ -343,6 +352,11 @@ def output(self) -> Wires:
         ret._original = self.original or self  # pylint: disable=protected-access
         return ret
 
+    def overlap(self, other: Wires) -> tuple[set[int], set[int]]:
+        ovlp_ket = self.output.ket.modes & other.input.ket.modes
+        ovlp_bra = self.output.bra.modes & other.input.bra.modes
+        return ovlp_bra, ovlp_ket
+
     @cached_property
     def adjoint(self) -> Wires:
         r"New ``Wires`` object obtained by swapping ket and bra wires."
@@ -492,9 +506,8 @@ def __repr__(self) -> str:
         return f"Wires{self.args}"
 
     def contracted_indices(self, other: Wires):
-        """Returns the indices being contracted on self and other"""
-        ovlp_bra = sorted(self.output.bra.modes & other.input.bra.modes)
-        ovlp_ket = sorted(self.output.ket.modes & other.input.ket.modes)
+        """Returns the indices being contracted on self and other on quantum wires."""
+        ovlp_bra, ovlp_ket = self.overlap(other)
         idxA = self.output.bra[ovlp_bra].indices + self.output.ket[ovlp_ket].indices
         idxB = other.input.bra[ovlp_bra].indices + other.input.ket[ovlp_ket].indices
         return idxA, idxB