Repository: qiboteam/qibo

Issue:
Draw Noise Channel Gates
**Is your feature request related to a problem? Please describe.**
I have implement a customised noise model for my quantum circuit. However, I could not use the built in .draw() function to visualise my quantum circuit. It would be difficult to check whether I have placed my quantum channels correctly.

> ~/miniconda3/envs/mypsi4/lib//python3.8/site-packages/qibo/abstractions/circuit.py in draw(self, line_wrap)
>953         for gate in self.queue:
>954             if gate.name not in labels:
>--> 955                 raise_error(NotImplementedError, f"{gate.name} gate is not supported by `circuit.draw`")
>956             gate_name = labels.get(gate.name)
>957             targets = list(gate.target_qubits)
>
>~/miniconda3/envs/mypsi4/lib//python3.8/site-packages/qibo/config.py in raise_error(exception, message, args)
>50         raise exception(message, args)
>51     else:
>---> 52         raise exception(message)
>53 
>54 
> 
> NotImplementedError: PauliNoiseChannel gate is not supported by `circuit.draw`

**Describe the solution you'd like**
Draw a quantum circuit with simple special symbols that represents the single qubit noise channels gates (Thermal, Pauli, etc) for the purpose of checking gates placement quickly.

**Describe alternatives you've considered**
I have place all the gates in a list before adding them to the quantum circuit one by one. By printing out the gate type and qubit index of each gate, one can 'visualise' the quantum circuit with great difficulty.

**Additional context**
None

Discussion:

Patch:
diff --git a/setup.py b/setup.py
index 3332da87f..7a1aaa439 100644
--- a/setup.py
+++ b/setup.py
@@ -60,6 +60,7 @@ setup(
         "psutil",
         "pyyaml",
         "importlib_metadata",
+        "tabulate"
     ],
     extras_require={
         "docs": ["sphinx", "sphinx_rtd_theme", "recommonmark", "sphinxcontrib-bibtex", "sphinx_markdown_tables", "nbsphinx", "IPython"],
diff --git a/src/qibo/abstractions/circuit.py b/src/qibo/abstractions/circuit.py
index bee382e2d..3029054f3 100644
--- a/src/qibo/abstractions/circuit.py
+++ b/src/qibo/abstractions/circuit.py
@@ -937,12 +937,14 @@ class AbstractCircuit(ABC):
 
         return len(qubits), gate_list
 
-    def draw(self, line_wrap=70) -> str:
+    def draw(self, line_wrap=70, legend=False) -> str:
         """Draw text circuit using unicode symbols.
 
         Args:
             line_wrap (int): maximum number of characters per line. This option
                 split the circuit text diagram in chunks of line_wrap characters.
+            legend (bool): If ``True`` prints a legend below the circuit for
+                callbacks and channels. Default is ``False``.
 
         Return:
             String containing text circuit diagram.
@@ -955,7 +957,12 @@ class AbstractCircuit(ABC):
                   "crx": "RX", "cry": "RY", "crz": "RZ",
                   "cu1": "U1", "cu3": "U3", "ccx": "X",
                   "id": "I", "measure": "M", "fsim": "f",
-                  "generalizedfsim": "gf", "Unitary": "U", "fswap":"fx"}
+                  "generalizedfsim": "gf", "Unitary": "U", "fswap":"fx",
+                  "PauliNoiseChannel": "PN", "KrausChannel": "K",
+                  "UnitaryChannel": "U", "ThermalRelaxationChannel": "TR",
+                  "ResetChannel": "R", "PartialTrace": "PT",
+                  "EntanglementEntropy": "EE", "Norm": "N",
+                  "Overlap": "O", "Energy": "E"}
 
         # build string representation of gates
         matrix = [[] for _ in range(self.nqubits)]
@@ -965,7 +972,10 @@ class AbstractCircuit(ABC):
             if gate.name not in labels:
                 raise_error(NotImplementedError, f"{gate.name} gate is not supported by `circuit.draw`")
             gate_name = labels.get(gate.name)
-            targets = list(gate.target_qubits)
+            if isinstance(gate, gates.CallbackGate):
+                targets = list(range(self.nqubits))
+            else:
+                targets = list(gate.target_qubits)
             controls = list(gate.control_qubits)
 
             # identify boundaries
@@ -1013,6 +1023,17 @@ class AbstractCircuit(ABC):
             output += f'q{q}' + ' ' * (len(str(self.nqubits))-len(str(q))) + \
                        ': ─' + ''.join(matrix[q]) + '\n'
 
+        # legend
+        if legend:
+            from tabulate import tabulate
+            names = [i.name for i in self.queue \
+                     if isinstance(i,gates.CallbackGate) or "Channel" in i.name]
+            names = list(dict.fromkeys(names))
+            table = tabulate([[i, labels[i]] for i in names],
+                              headers=['Gate', 'Symbol'],
+                              tablefmt='orgtbl')
+            table = '\n Legend for callbacks and channels: \n' + table
+
         # line wrap
         if line_wrap:
             loutput = output.splitlines()
@@ -1040,4 +1061,7 @@ class AbstractCircuit(ABC):
             if loutput is not None:
                 output = ''.join(loutput)
 
+        if legend:
+            output += table
+
         return output.rstrip('\n')


Test patch:
diff --git a/src/qibo/tests/test_abstract_circuit.py b/src/qibo/tests/test_abstract_circuit.py
index dfa3ba1ef..05c6b567b 100644
--- a/src/qibo/tests/test_abstract_circuit.py
+++ b/src/qibo/tests/test_abstract_circuit.py
@@ -549,3 +549,49 @@ def test_circuit_draw_not_supported_gates():
     c.add(gates.Flatten(1))
     with pytest.raises(NotImplementedError):
         c.draw()
+@pytest.mark.parametrize("legend", [True, False])
+def test_circuit_draw_channels(legend):
+    """Check that channels are drawn correctly"""
+    from qibo.models import Circuit as circuit
+    c = circuit(2, density_matrix=True)
+    c.add(gates.H(0))
+    c.add(gates.PauliNoiseChannel(0, 0.1, 0.0, 0.2))
+    c.add(gates.H(1))
+    c.add(gates.PauliNoiseChannel(1, 0.0, 0.2, 0.1))
+    c.add(gates.CNOT(0, 1))
+    c.add(gates.PauliNoiseChannel(0, 0.1, 0.0, 0.2))
+    c.add(gates.PauliNoiseChannel(1, 0.0, 0.2, 0.1))
+
+    ref = 'q0: ─H─PN─o─PN─\n' \
+          'q1: ─H─PN─X─PN─'
+
+    if legend:
+        ref += '\n\n Legend for callbacks and channels: \n' \
+               '| Gate              | Symbol   |\n' \
+               '|-------------------+----------|\n' \
+               '| PauliNoiseChannel | PN       |'
+
+    assert c.draw(legend=legend) == ref
+
+@pytest.mark.parametrize("legend", [True, False])
+def test_circuit_draw_callbacks(legend):
+    """Check that callbacks are drawn correcly"""
+    from qibo.callbacks import EntanglementEntropy
+    entropy = EntanglementEntropy([0])
+    c = Circuit(2)
+    c.add(gates.CallbackGate(entropy))
+    c.add(gates.H(0))
+    c.add(gates.CallbackGate(entropy))
+    c.add(gates.CNOT(0, 1))
+    c.add(gates.CallbackGate(entropy))
+
+    ref = 'q0: ─EE─H─EE─o─EE─\n' \
+          'q1: ─EE───EE─X─EE─'
+
+    if legend:
+        ref += '\n\n Legend for callbacks and channels: \n' \
+               '| Gate                | Symbol   |\n'\
+               '|---------------------+----------|\n'\
+               '| EntanglementEntropy | EE       |'
+
+    assert c.draw(legend=legend) == ref