Merge pull request #354 from Quantum-TII/empdraw

Updating multiline circuit draw
qiboteam · Mar 16, 2021 · 34bd341 · 34bd341
2 parents f45dd66 + 452de4b
commit 34bd341
Show file tree

Hide file tree

Showing 2 changed files with 44 additions and 38 deletions.
diff --git a/src/qibo/abstractions/circuit.py b/src/qibo/abstractions/circuit.py
@@ -933,7 +933,7 @@ def read_args(args):
 
         return len(qubits), gate_list
 
-    def draw(self, line_wrap=None) -> str:
+    def draw(self, line_wrap=70) -> str:
         """Draw text circuit using unicode symbols.
 
         Args:
@@ -1010,24 +1010,29 @@ def draw(self, line_wrap=None) -> str:
 
         # line wrap
         if line_wrap:
-            output = output.splitlines()
+            loutput = output.splitlines()
             def chunkstring(string, length):
                 nchunks = range(0, len(string), length)
                 return (string[i:length+i] for i in nchunks), len(nchunks)
             for row in range(self.nqubits):
-                chunks, nchunks = chunkstring(output[row], line_wrap)
+                chunks, nchunks = chunkstring(loutput[row][3 + len(str(self.nqubits)):], line_wrap)
+                if nchunks == 1:
+                    loutput = None
+                    break
                 for i, c in enumerate(chunks):
-                    output += ['' for _ in range(self.nqubits)]
-                    suffix = ' ...\n'
+                    loutput += ['' for _ in range(self.nqubits)]
+                    suffix = f' ...\n'
+                    prefix = f'q{row}' + ' ' * (len(str(self.nqubits))-len(str(row))) + ': '
                     if i == 0:
-                        prefix = ''
+                        prefix += ' ' * 4
                     elif row == 0:
-                        prefix = '\n... '
+                        prefix = '\n' + prefix + '... '
                     else:
-                        prefix = '... '
+                        prefix += '... '
                     if i == nchunks-1:
                         suffix = '\n'
-                    output[row + i * self.nqubits] = prefix + c + suffix
-            output = ''.join(output)
+                    loutput[row + i * self.nqubits] = prefix + c + suffix
+            if loutput is not None:
+                output = ''.join(loutput)
 
         return output.rstrip('\n')
diff --git a/src/qibo/tests_new/test_abstract_circuit.py b/src/qibo/tests_new/test_abstract_circuit.py
@@ -495,36 +495,37 @@ def test_circuit_draw():
 def test_circuit_draw_line_wrap():
     """Test circuit text draw with line wrap."""
     ref_line_wrap_50 = \
-          'q0: ─H─U1─U1─U1─U1───────────────────────────x───I ...\n' \
-          'q1: ───o──|──|──|──H─U1─U1─U1────────────────|─x─I ...\n' \
-          'q2: ──────o──|──|────o──|──|──H─U1─U1────────|─|── ...\n' \
-          'q3: ─────────o──|───────o──|────o──|──H─U1───|─x── ...\n' \
-          'q4: ────────────o──────────o───────o────o──H─x──── ...\n' \
-          '\n' \
-          '... ───f─o────gf───M─\n' \
-          '... ───|─U3───|──o─M─\n' \
-          '... ───|────X─gf─o─M─\n' \
-          '... ─M─|────o────o───\n' \
-          '... ───f────o────X───'
+        'q0:     ─H─U1─U1─U1─U1───────────────────────────x───I───f ...\n' \
+        'q1:     ───o──|──|──|──H─U1─U1─U1────────────────|─x─I───| ...\n' \
+        'q2:     ──────o──|──|────o──|──|──H─U1─U1────────|─|─────| ...\n' \
+        'q3:     ─────────o──|───────o──|────o──|──H─U1───|─x───M─| ...\n' \
+        'q4:     ────────────o──────────o───────o────o──H─x───────f ...\n' \
+        '\n' \
+        'q0: ... ─o────gf───M─\n' \
+        'q1: ... ─U3───|──o─M─\n' \
+        'q2: ... ────X─gf─o─M─\n' \
+        'q3: ... ────o────o───\n' \
+        'q4: ... ────o────X───'
 
     ref_line_wrap_30 = \
-           'q0: ─H─U1─U1─U1─U1──────────── ...\n' \
-           'q1: ───o──|──|──|──H─U1─U1─U1─ ...\n' \
-           'q2: ──────o──|──|────o──|──|── ...\n' \
-           'q3: ─────────o──|───────o──|── ...\n' \
-           'q4: ────────────o──────────o── ...\n' \
-           '\n' \
-           '... ───────────────x───I───f─o──── ...\n' \
-           '... ───────────────|─x─I───|─U3─── ...\n' \
-           '... H─U1─U1────────|─|─────|────X─ ...\n' \
-           '... ──o──|──H─U1───|─x───M─|────o─ ...\n' \
-           '... ─────o────o──H─x───────f────o─ ...\n' \
-           '\n' \
-           '... gf───M─\n' \
-           '... |──o─M─\n' \
-           '... gf─o─M─\n' \
-           '... ───o───\n' \
-           '... ───X───'
+        'q0:     ─H─U1─U1─U1─U1──────────────── ...\n' \
+        'q1:     ───o──|──|──|──H─U1─U1─U1───── ...\n' \
+        'q2:     ──────o──|──|────o──|──|──H─U1 ...\n' \
+        'q3:     ─────────o──|───────o──|────o─ ...\n' \
+        'q4:     ────────────o──────────o────── ...\n' \
+        '\n' \
+        'q0: ... ───────────x───I───f─o────gf── ...\n' \
+        'q1: ... ───────────|─x─I───|─U3───|──o ...\n' \
+        'q2: ... ─U1────────|─|─────|────X─gf─o ...\n' \
+        'q3: ... ─|──H─U1───|─x───M─|────o────o ...\n' \
+        'q4: ... ─o────o──H─x───────f────o────X ...\n' \
+        '\n' \
+        'q0: ... ─M─\n' \
+        'q1: ... ─M─\n' \
+        'q2: ... ─M─\n' \
+        'q3: ... ───\n' \
+        'q4: ... ───'
+
     import numpy as np
     circuit = Circuit(5)
     for i1 in range(5):