Repository: PennyLaneAI/pennylane
Pull request #3760 (2023-02-14 15:58:20)

Update qml.draw() to support quantum functions as well
**Context:**
Sometimes, users want to draw a sub-function in a QNode, or some non-complete quantum circuit. This enables them to do that without digging through pennylane source code.

**Description of the Change:**
Update `qml.draw(_mpl)` to support the drawing of quantum functions, and not just QNodes.

**Benefits:**
Users can visualize parts of a circuit without creating a whole QNode.

Example:
```pycon
>>> import pennylane as qml
>>> def qfunc(x):
...     qml.RX(x, [0])
...     qml.CNOT([0, 1])
... 
>>> print(qml.draw(qfunc)(1.1))
0: ──RX─╭●─┤  
1: ─────╰X─┤
```

**Possible Drawbacks:**
The code is a bit more messy? I had originally put them into their own functions (like `qml.draw_qfunc`) but I think that might be awkward/unexpected/unnecessary for users (ty @KetpuntoG for the nudge 😉). I addressed this concern in the in-line comment, lmk what you (reviewers) think.

**Related GitHub Issues:**
#3666 

github-actions[bot]: Hello. You may have forgotten to update the changelog!
Please edit [`doc/releases/changelog-dev.md`](/PennyLaneAI/pennylane/blob/master/doc/releases/changelog-dev.md) with:
* A one-to-two sentence description of the change. You may include a small working example for new features.
* A link back to this PR.
* Your name (or GitHub username) in the contributors section.
codecov[bot]: # [Codecov](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3760?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) Report
> Merging [#3760](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3760?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) (85b0506) into [master](https://codecov.io/gh/PennyLaneAI/pennylane/commit/023e001cd437e5283a4ffa1b2f73bda38ebf4bc9?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) (023e001) will **decrease** coverage by `0.02%`.
> The diff coverage is `55.55%`.

```diff
@@            Coverage Diff             @@
##           master    #3760      +/-   ##
==========================================
- Coverage   99.81%   99.79%   -0.02%     
==========================================
  Files         331      331              
  Lines       28982    28989       +7     
==========================================
+ Hits        28927    28930       +3     
- Misses         55       59       +4     
```


| [Impacted Files](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3760?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) | Coverage Δ | |
|---|---|---|
| [pennylane/drawer/draw.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3760?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL2RyYXdlci9kcmF3LnB5) | `90.47% <42.85%> (-9.53%)` | :arrow_down: |
| [pennylane/\_\_init\_\_.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3760?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL19faW5pdF9fLnB5) | `96.77% <100.00%> (ø)` | |
| [pennylane/drawer/\_\_init\_\_.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3760?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL2RyYXdlci9fX2luaXRfXy5weQ==) | `100.00% <100.00%> (ø)` | |

Help us with your feedback. Take ten seconds to tell us [how you rate us](https://about.codecov.io/nps?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI). Have a feature suggestion? [Share it here.](https://app.codecov.io/gh/feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI)

timmysilv: update: i separated the two different use-cases into their own functions internally, while preserving the user interface of `qml.draw(QNode or Callable)`. Makes the code more readable and hopefully easier to maintain when we get to things such as the batch transform cleanup mentioned above. The only common thing was wrapping `return tape_text(...)` so the split felt fine. I'd use `singledispatch`, but of course, batch-transformed QNodes would go to the wrong place!

Patch:
diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md
index 1063830d2..2eea1e931 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -626,6 +626,10 @@
 * `transforms.fold_global` now supports non-integer values.
   [(#3781)](https://github.com/PennyLaneAI/pennylane/pull/3781)
 
+* Updated `qml.draw` and `qml.draw_mpl` to draw any quantum function,
+  which allows for visualizing only part of a complete circuit/QNode.
+  [(#3760)](https://github.com/PennyLaneAI/pennylane/pull/3760)
+
 <h3>Breaking changes</h3>
 
 * The argument `mode` in execution is replaced by the boolean `grad_on_execution` in the new execution pipeline.
diff --git a/pennylane/drawer/draw.py b/pennylane/drawer/draw.py
index cf4d62fc5..7a5f0f8c5 100644
--- a/pennylane/drawer/draw.py
+++ b/pennylane/drawer/draw.py
@@ -17,7 +17,9 @@
 Contains the drawing function.
 """
 from functools import wraps
+import warnings
 
+import pennylane as qml
 from .tape_mpl import tape_mpl
 from .tape_text import tape_text
 
@@ -31,10 +33,10 @@ def draw(
     show_matrices=False,
     expansion_strategy=None,
 ):
-    """Create a function that draws the given qnode.
+    """Create a function that draws the given qnode or quantum function.
 
     Args:
-        qnode (.QNode): the input QNode that is to be drawn.
+        qnode (.QNode or Callable): the input QNode or quantum function that is to be drawn.
         wire_order (Sequence[Any]): the order (from top to bottom) to print the wires of the circuit
         show_all_wires (bool): If True, all wires, including empty wires, are printed.
         decimals (int): How many decimal points to include when formatting operation parameters.
@@ -42,7 +44,7 @@ def draw(
         max_length (int): Maximum string width (columns) when printing the circuit
         show_matrices=False (bool): show matrix valued parameters below all circuit diagrams
         expansion_strategy (str): The strategy to use when circuit expansions or decompositions
-            are required.
+            are required. Note that this is ignored if the input is not a QNode.
 
             - ``gradient``: The QNode will attempt to decompose
               the internal circuit such that all circuit operations are supported by the gradient
@@ -54,7 +56,7 @@ def draw(
 
     Returns:
         A function that has the same argument signature as ``qnode``. When called,
-        the function will draw the QNode.
+        the function will draw the QNode/qfunc.
 
     **Example**
 
@@ -178,8 +180,61 @@ def draw(
 
         0: ──RX(0.90)─┤  <Z>
 
+    The function also accepts quantum functions rather than QNodes. This can be especially
+    helpful if you want to visualize only a part of a circuit that may not be convertible into
+    a QNode, such as a sub-function that does not return any measurements.
+
+    >>> def qfunc(x):
+    ...     qml.RX(x, wires=[0])
+    ...     qml.CNOT(wires=[0, 1])
+    >>> print(qml.draw(qfunc)(1.1))
+    0: ──RX(1.10)─╭●─┤
+    1: ───────────╰X─┤
+
     """
+    if hasattr(qnode, "construct"):
+        return _draw_qnode(
+            qnode,
+            wire_order=wire_order,
+            show_all_wires=show_all_wires,
+            decimals=decimals,
+            max_length=max_length,
+            show_matrices=show_matrices,
+            expansion_strategy=expansion_strategy,
+        )
+
+    if expansion_strategy is not None:
+        warnings.warn(
+            "When the input to qml.draw is not a QNode, the expansion_strategy argument is ignored.",
+            UserWarning,
+        )
 
+    @wraps(qnode)
+    def wrapper(*args, **kwargs):
+        tape = qml.tape.make_qscript(qnode)(*args, **kwargs)
+        _wire_order = wire_order or tape.wires
+
+        return tape_text(
+            tape,
+            wire_order=_wire_order,
+            show_all_wires=show_all_wires,
+            decimals=decimals,
+            show_matrices=show_matrices,
+            max_length=max_length,
+        )
+
+    return wrapper
+
+
+def _draw_qnode(
+    qnode,
+    wire_order=None,
+    show_all_wires=False,
+    decimals=2,
+    max_length=100,
+    show_matrices=False,
+    expansion_strategy=None,
+):
     @wraps(qnode)
     def wrapper(*args, **kwargs):
         original_expansion_strategy = getattr(qnode, "expansion_strategy", None)
@@ -237,7 +292,7 @@ def draw_mpl(
     """Draw a qnode with matplotlib
 
     Args:
-        qnode (.QNode): the input QNode that is to be drawn.
+        qnode (.QNode or Callable): the input QNode/quantum function that is to be drawn.
 
     Keyword Args:
         wire_order (Sequence[Any]): the order (from top to bottom) to print the wires of the circuit
@@ -429,7 +484,49 @@ def draw_mpl(
                 :target: javascript:void(0);
 
     """
+    if hasattr(qnode, "construct"):
+        return _draw_mpl_qnode(
+            qnode,
+            wire_order=wire_order,
+            show_all_wires=show_all_wires,
+            decimals=decimals,
+            expansion_strategy=expansion_strategy,
+            style=style,
+            **kwargs,
+        )
+
+    if expansion_strategy is not None:
+        warnings.warn(
+            "When the input to qml.draw is not a QNode, the expansion_strategy argument is ignored.",
+            UserWarning,
+        )
+
+    @wraps(qnode)
+    def wrapper(*args, **kwargs):
+        tape = qml.tape.make_qscript(qnode)(*args, **kwargs)
+        _wire_order = wire_order or tape.wires
+
+        return tape_mpl(
+            tape,
+            wire_order=_wire_order,
+            show_all_wires=show_all_wires,
+            decimals=decimals,
+            style=style,
+            **kwargs,
+        )
 
+    return wrapper
+
+
+def _draw_mpl_qnode(
+    qnode,
+    wire_order=None,
+    show_all_wires=False,
+    decimals=None,
+    expansion_strategy=None,
+    style="black_white",
+    **kwargs,
+):
     @wraps(qnode)
     def wrapper(*args, **kwargs_qnode):
         original_expansion_strategy = getattr(qnode, "expansion_strategy", None)


Test patch:
diff --git a/tests/drawer/test_draw.py b/tests/drawer/test_draw.py
index 5ae4292a8..f6c0b43e1 100644
--- a/tests/drawer/test_draw.py
+++ b/tests/drawer/test_draw.py
@@ -38,17 +38,17 @@ class TestLabelling:
         """Test wire labels with different kinds of objects."""
 
         split_str = draw(circuit)(1.2, 2.3, 3.4).split("\n")
-        assert split_str[0][0:6] == "    0:"
-        assert split_str[1][0:6] == "    a:"
-        assert split_str[2][0:6] == "1.234:"
+        assert split_str[0][:6] == "    0:"
+        assert split_str[1][:6] == "    a:"
+        assert split_str[2][:6] == "1.234:"
 
     def test_wire_order(self):
         """Test wire_order keyword changes order of the wires."""
 
         split_str = draw(circuit, wire_order=[1.234, "a", 0, "b"])(1.2, 2.3, 3.4).split("\n")
-        assert split_str[0][0:6] == "1.234:"
-        assert split_str[1][0:6] == "    a:"
-        assert split_str[2][0:6] == "    0:"
+        assert split_str[0][:6] == "1.234:"
+        assert split_str[1][:6] == "    a:"
+        assert split_str[2][:6] == "    0:"
 
     def test_show_all_wires(self):
         """Test show_all_wires=True forces empty wires to display."""
@@ -58,8 +58,8 @@ class TestLabelling:
             return qml.expval(qml.PauliZ(0))
 
         split_str = draw(circuit, show_all_wires=True)().split("\n")
-        assert split_str[0][0:2] == "0:"
-        assert split_str[1][0:2] == "1:"
+        assert split_str[0][:2] == "0:"
+        assert split_str[1][:2] == "1:"
 
     def test_show_all_wires_and_wire_order(self):
         """Test show_all_wires forces empty wires to display when empty wire is in wire order."""
@@ -69,8 +69,8 @@ class TestLabelling:
             return qml.expval(qml.PauliZ(0))
 
         split_str = draw(circuit, wire_order=[0, "a"], show_all_wires=True)().split("\n")
-        assert split_str[0][0:2] == "0:"
-        assert split_str[1][0:2] == "a:"
+        assert split_str[0][:2] == "0:"
+        assert split_str[1][:2] == "a:"
 
 
 class TestDecimals:
@@ -341,10 +341,38 @@ def test_expansion_strategy():
         return qml.probs(wires=0)
 
     expected_gradient = "0: ─╭ApproxTimeEvolution─┤  Probs\n1: ─╰ApproxTimeEvolution─┤       "
-    assert draw(circuit, expansion_strategy="gradient", decimals=None)(0.5)
+    assert draw(circuit, expansion_strategy="gradient", decimals=None)(0.5) == expected_gradient
 
-    expected_device = (
-        "0: ──H────────MultiRZ──H──H─╭MultiRZ──H──H────────MultiRZ──H──H─╭MultiRZ──H─┤  Probs\n"
-        "1: ──MultiRZ──H─────────────╰MultiRZ──H──MultiRZ──H─────────────╰MultiRZ──H─┤       "
-    )
-    assert draw(circuit, expansion_strategy="device", decimals=None)(0.5)
+    expected_device = "0: ──RX─╭RXX──RX─╭RXX─┤  Probs\n1: ──RZ─╰RXX──RZ─╰RXX─┤       "
+    assert draw(circuit, expansion_strategy="device", decimals=None)(0.5) == expected_device
+
+
+def test_draw_with_qfunc():
+    """Test a non-qnode qfunc can be drawn."""
+
+    def qfunc(x):
+        qml.RX(x, wires=[0])
+        qml.PauliZ(1)
+
+    assert qml.draw(qfunc)(1.1) == "0: ──RX(1.10)─┤  \n1: ──Z────────┤  "
+
+
+def test_draw_with_qfunc_with_measurements():
+    """Test a non-qnode qfunc with measurements can be drawn."""
+
+    def qfunc(x):
+        qml.RX(x, wires=[0])
+        qml.CNOT([0, 1])
+        return qml.expval(qml.PauliZ(1))
+
+    assert qml.draw(qfunc)(1.1) == "0: ──RX(1.10)─╭●─┤     \n1: ───────────╰X─┤  <Z>"
+
+
+def test_draw_with_qfunc_warns_with_expansion_strategy():
+    """Test that draw warns the user about expansion_strategy being ignored."""
+
+    def qfunc():
+        qml.PauliZ(0)
+
+    with pytest.warns(UserWarning, match="the expansion_strategy argument is ignored"):
+        _ = qml.draw(qfunc, expansion_strategy="gradient")
diff --git a/tests/drawer/test_draw_mpl.py b/tests/drawer/test_draw_mpl.py
index c1ee555ec..507b1d0ff 100644
--- a/tests/drawer/test_draw_mpl.py
+++ b/tests/drawer/test_draw_mpl.py
@@ -49,7 +49,7 @@ def test_standard_use():
     fig, ax = qml.draw_mpl(circuit1)(1.23, 2.34)
 
     assert isinstance(fig, mpl.figure.Figure)
-    assert isinstance(ax, mpl.axes._axes.Axes)
+    assert isinstance(ax, mpl.axes._axes.Axes)  # pylint:disable=protected-access
 
     # proxy for whether correct things were drawn
     assert len(ax.patches) == 7  # two boxes, 2 circles for CNOT, 3 patches for measure
@@ -79,7 +79,7 @@ def test_expansion_strategy(strategy, initial_strategy, n_lines):
         qml.Permute([2, 0, 1], wires=(0, 1, 2))
         return qml.expval(qml.PauliZ(0))
 
-    fig, ax = qml.draw_mpl(circuit, expansion_strategy=strategy)()
+    _, ax = qml.draw_mpl(circuit, expansion_strategy=strategy)()
 
     assert len(ax.lines) == n_lines
     assert circuit.expansion_strategy == initial_strategy
@@ -114,7 +114,7 @@ class TestKwargs:
             1.23, 2.34
         )
 
-        for l in ax.texts[0:3]:  # three labels
+        for l in ax.texts[:3]:  # three labels
             assert l.get_color() == "purple"
             assert l.get_fontsize() == 20
         plt.close()
@@ -230,7 +230,7 @@ class TestWireBehaviour:
 
         _, ax = qml.draw_mpl(circuit1, wire_options={"color": "black", "linewidth": 4})(1.23, 2.34)
 
-        for w in ax.lines[0:3]:  # three wires
+        for w in ax.lines[:3]:  # three wires
             assert w.get_color() == "black"
             assert w.get_linewidth() == 4
 
@@ -253,7 +253,7 @@ class TestMPLIntegration:
         assert ax.patches[0].get_facecolor() == rgba_red
         assert ax.patches[1].get_facecolor() == rgba_red
 
-        for l in ax.lines[0:-1]:  # final is fancy arrow, has different styling
+        for l in ax.lines[:-1]:  # final is fancy arrow, has different styling
             assert l.get_color() == rgba_green
 
         plt.style.use("default")
@@ -271,7 +271,7 @@ class TestMPLIntegration:
         assert ax.patches[1].get_facecolor() == expected_facecolor
 
         expected_linecolor = mpl.colors.to_rgba(plt.rcParams["lines.color"])
-        for l in ax.lines[0:-1]:  # final is fancy arrow, has different styling
+        for l in ax.lines[:-1]:  # final is fancy arrow, has different styling
             assert mpl.colors.to_rgba(l.get_color()) == expected_linecolor
 
         plt.style.use("default")
@@ -305,3 +305,31 @@ class TestMPLIntegration:
         assert mpl.rcParams["axes.facecolor"] == initial_facecolor
         assert mpl.rcParams["patch.facecolor"] == initial_patch_facecolor
         assert mpl.rcParams["patch.edgecolor"] == initial_patch_edgecolor
+
+
+def test_draw_mpl_supports_qfuncs():
+    """Test that draw_mpl works with non-QNode quantum functions."""
+
+    def qfunc(x):
+        qml.RX(x, 0)
+
+    fig, ax = qml.draw_mpl(qfunc)(1.1)
+
+    assert isinstance(fig, mpl.figure.Figure)
+    assert isinstance(ax, mpl.axes._axes.Axes)  # pylint:disable=protected-access
+    assert len(ax.patches) == 1
+    assert len(ax.lines) == 1
+    assert len(ax.texts) == 2
+    assert ax.texts[0].get_text() == "0"
+    assert ax.texts[1].get_text() == "RX"
+    plt.close()
+
+
+def test_draw_mpl_with_qfunc_warns_with_expansion_strategy():
+    """Test that draw warns the user about expansion_strategy being ignored."""
+
+    def qfunc():
+        qml.PauliZ(0)
+
+    with pytest.warns(UserWarning, match="the expansion_strategy argument is ignored"):
+        _ = qml.draw_mpl(qfunc, expansion_strategy="gradient")
diff --git a/tests/drawer/test_drawer_utils.py b/tests/drawer/test_drawer_utils.py
index 6fe0123e6..dbb661b0f 100644
--- a/tests/drawer/test_drawer_utils.py
+++ b/tests/drawer/test_drawer_utils.py
@@ -99,6 +99,8 @@ class TestConvertWireOrder:
 class TestUnwrapControls:
     """Tests the ``unwrap_controls`` utility function."""
 
+    # pylint:disable=too-few-public-methods
+
     @pytest.mark.parametrize(
         "op,expected_control_wires,expected_control_values",
         [
diff --git a/tests/tests_passing_pylint.txt b/tests/tests_passing_pylint.txt
index cc09f4674..18d8e29be 100644
--- a/tests/tests_passing_pylint.txt
+++ b/tests/tests_passing_pylint.txt
@@ -2,6 +2,8 @@ tests/data/**
 tests/devices/qubit/test_apply_operation.py
 tests/devices/test_execution_config.py
 tests/drawer/test_tape_mpl.py
+tests/drawer/test_draw_mpl.py
+tests/drawer/test_drawer_utils.py
 tests/measurements/test_probs.py
 tests/qchem/test_givens_rotations.py
 tests/templates/test_subroutines/test_basis_rotation.py