Fix depth and ascii diagrams for circuits with overlapping multi-qubi…

src/braket/circuits/ascii_circuit_diagram.py

@@ -11,7 +11,7 @@
 # ANY KIND, either express or implied. See the License for the specific
 # language governing permissions and limitations under the License.
 
-from typing import List
+from typing import List, Tuple
 
 from braket.circuits.circuit_diagram import CircuitDiagram
 from braket.circuits.gate import Gate
@@ -28,7 +28,7 @@ def build_diagram(circuit) -> str:
         Build an ASCII string circuit diagram.
 
         Args:
-            circuit (Circuit): Circuit to build a diagram of.
+            circuit (Circuit): Circuit to build a diagram.
 
         Returns:
             str: ASCII string circuit diagram.
@@ -68,23 +68,106 @@ def build_diagram(circuit) -> str:
 
         return "\n".join(lines)
 
+    @staticmethod
+    def _ascii_moment_group_instructions(
+        instructions: List[Instruction],
+    ) -> List[Tuple[QubitSet, List[Instruction]]]:
+        """
+        Group instructions in a moment for ASCII diagram
+
+        Args:
+            instructions (List[Instruction]): list of instructions
+
+        Returns:
+            List[(QubitSet, List[Instruction])]: list of grouped instructions
+        """
+        groupings = []
+        for instr in instructions:
+            # Can only print Gate operators at the moment
+            if not isinstance(instr.operator, Gate):
+                continue
+
+            qubit_range = QubitSet(range(min(instr.target), max(instr.target) + 1))
+
+            found_grouping = False
+            for group in groupings:
+                qubits_added = group[0]
+                instr_group = group[1]
+                # Take into account overlapping multi-qubit gates
+                if not qubits_added.intersection(set(qubit_range)):
+                    instr_group.append(instr)
+                    qubits_added.update(qubit_range)
+                    found_grouping = True
+                    break
+
+            if not found_grouping:
+                groupings.append((qubit_range, [instr]))
+
+        return groupings
+
     @staticmethod
     def _ascii_diagram_moment(
         time: int, circuit_qubits: QubitSet, instructions: List[Instruction]
     ) -> str:
         """
         Return an ASCII string diagram of the circuit at a particular moment in time.
 
+        Args:
+            time (int): time of moment
+            circuit_qubits (QubitSet): qubits in circuit
+            instructions (List[Instruction]): list of instructions
+
         Returns:
             str: An ASCII string diagram for the specified moment in time.
         """
+
+        # Group instructions to separate out overlapping multi-qubit gates
+        groupings = AsciiCircuitDiagram._ascii_moment_group_instructions(instructions)
+
+        column_strs = [
+            AsciiCircuitDiagram._ascii_diagram_moment_column(circuit_qubits, grouping[1])
+            for grouping in groupings
+        ]
+
+        # Unite column strings
+        lines = column_strs[0].split("\n")
+        for column_str in column_strs[1:]:
+            for i, moment_line in enumerate(column_str.split("\n")):
+                lines[i] += moment_line
+
+        # Adjust for time width
+        time_width = len(str(time))
+        symbols_width = len(lines[0]) - 1
+        if symbols_width < time_width:
+            diff = time_width - symbols_width
+            for i in range(len(lines) - 1):
+                if lines[i].endswith("-"):
+                    lines[i] += "-" * diff
+                else:
+                    lines[i] += " "
+
+        first_line = "{:^{width}}|\n".format(str(time), width=len(lines[0]) - 1)
+
+        return first_line + "\n".join(lines)
+
+    @staticmethod
+    def _ascii_diagram_moment_column(
+        circuit_qubits: QubitSet, instructions: List[Instruction]
+    ) -> str:
+        """
+        Return an ASCII string diagram of the circuit at a particular moment in time for a column.
+
+        Args:
+            circuit_qubits (QubitSet): qubits in circuit
+            instructions (List[Instruction]): list of instructions
+
+        Returns:
+            str: An ASCII string diagram for the specified moment in time for a column.
+        """
         symbols = {qubit: "-" for qubit in circuit_qubits}
         margins = {qubit: " " for qubit in circuit_qubits}
 
         for instr in instructions:
-            # Can only print Gate operators at the moment
-            if not isinstance(instr.operator, Gate):
-                continue
 
             qubits = circuit_qubits.intersection(
                 set(range(min(instr.target), max(instr.target) + 1))
@@ -100,13 +183,13 @@ def _ascii_diagram_moment(
                 else:
                     symbols[qubit] = "|"
 
-                # set the margin to be a connector if not on the first qubit
+                # Set the margin to be a connector if not on the first qubit
                 if qubit != min(instr.target):
                     margins[qubit] = "|"
 
-        symbols_width = max([len(symbol) for symbol in symbols.values()] + [len(str(time))])
+        symbols_width = max([len(symbol) for symbol in symbols.values()])
 
-        output = "{0:{width}}|\n".format(str(time), width=symbols_width)
+        output = ""
         for qubit in circuit_qubits:
             output += "{0:{width}}\n".format(margins[qubit], width=symbols_width + 1)
             output += "{0:{fill}{align}{width}}\n".format(

src/braket/circuits/moments.py

@@ -138,12 +138,10 @@ def add(self, instructions: Iterable[Instruction]) -> None:
             self._add(instruction)
 
     def _add(self, instruction: Instruction) -> None:
-        qubit_range = range(min(instruction.target), max(instruction.target) + 1)
+        qubit_range = instruction.target
         time = max([self._max_time_for_qubit(qubit) for qubit in qubit_range]) + 1
 
-        # Mark all qubits in the range to avoid another gate being placed in the overlap.
-        # For example CNOT(0, 5) would draw a line from 0 to 5 and therefore should prevent
-        # another instruction using those qubits in that time moment.
+        # Mark all qubits in qubit_range with max_time
         for qubit in qubit_range:
             self._max_times[qubit] = max(time, self._max_time_for_qubit(qubit))
 

test/unit_tests/braket/circuits/test_ascii_circuit_diagram.py

@@ -18,6 +18,13 @@ def test_empty_circuit():
     assert AsciiCircuitDiagram.build_diagram(Circuit()) == ""
 
 
+def test_one_gate_one_qubit():
+    circ = Circuit().h(0)
+    expected = ("T  : |0|", "        ", "q0 : -H-", "", "T  : |0|")
+    expected = "\n".join(expected)
+    assert AsciiCircuitDiagram.build_diagram(circ) == expected
+
+
 def test_qubit_width():
     circ = Circuit().h(0).h(100)
     expected = (
@@ -43,13 +50,13 @@ def to_ir(self, target):
 
     circ = Circuit().h(0).h(1).add_instruction(Instruction(Foo(), 0))
     expected = (
-        "T  : |0|1  |",
+        "T  : |0| 1 |",
         "            ",
         "q0 : -H-FOO-",
         "            ",
         "q1 : -H-----",
         "",
-        "T  : |0|1  |",
+        "T  : |0| 1 |",
     )
     expected = "\n".join(expected)
     assert AsciiCircuitDiagram.build_diagram(circ) == expected
@@ -120,10 +127,48 @@ def test_connector_across_two_qubits():
     assert AsciiCircuitDiagram.build_diagram(circ) == expected
 
 
+def test_overlapping_qubits():
+    circ = Circuit().cnot(0, 2).cnot(1, 3).h(0)
+    expected = (
+        "T  : | 0 |1|",
+        "            ",
+        "q0 : -C---H-",
+        "      |     ",
+        "q1 : -|-C---",
+        "      | |   ",
+        "q2 : -X-|---",
+        "        |   ",
+        "q3 : ---X---",
+        "",
+        "T  : | 0 |1|",
+    )
+    expected = "\n".join(expected)
+    assert AsciiCircuitDiagram.build_diagram(circ) == expected
+
+
+def test_overlapping_qubits_angled_gates():
+    circ = Circuit().zz([0, 2], 0.15).cnot(1, 3).h(0)
+    expected = (
+        "T  : |    0     |1|",
+        "                   ",
+        "q0 : -ZZ(0.15)---H-",
+        "      |            ",
+        "q1 : -|--------C---",
+        "      |        |   ",
+        "q2 : -ZZ(0.15)-|---",
+        "               |   ",
+        "q3 : ----------X---",
+        "",
+        "T  : |    0     |1|",
+    )
+    expected = "\n".join(expected)
+    assert AsciiCircuitDiagram.build_diagram(circ) == expected
+
+
 def test_connector_across_gt_two_qubits():
     circ = Circuit().h(4).cnot(3, 5).h(4).h(2)
     expected = (
-        "T  : |0|1|2|",
+        "T  : | 0 |1|",
         "            ",
         "q2 : -H-----",
         "            ",
@@ -133,7 +178,7 @@ def test_connector_across_gt_two_qubits():
         "        |   ",
         "q5 : ---X---",
         "",
-        "T  : |0|1|2|",
+        "T  : | 0 |1|",
     )
     expected = "\n".join(expected)
     assert AsciiCircuitDiagram.build_diagram(circ) == expected
@@ -142,7 +187,7 @@ def test_connector_across_gt_two_qubits():
 def test_connector_across_non_used_qubits():
     circ = Circuit().h(4).cnot(3, 100).h(4).h(101)
     expected = (
-        "T    : |0|1|2|",
+        "T    : | 0 |1|",
         "              ",
         "q3   : ---C---",
         "          |   ",
@@ -152,7 +197,7 @@ def test_connector_across_non_used_qubits():
         "              ",
         "q101 : -H-----",
         "",
-        "T    : |0|1|2|",
+        "T    : | 0 |1|",
     )
     expected = "\n".join(expected)
     assert AsciiCircuitDiagram.build_diagram(circ) == expected

test/unit_tests/braket/circuits/test_moments.py

@@ -73,8 +73,11 @@ def test_overlaping_qubits():
     moments = Moments([h(0), h(0)])
     assert moments.depth == 2
 
-    moments.add([cnot(0, 2), h(1)])
-    assert moments.depth == 4
+    moments.add([cnot(0, 3), h(1)])
+    assert moments.depth == 3
+
+    moments.add([cnot(2, 4)])
+    assert moments.depth == 3
 
 
 def test_qubits():