Improved text drawing for dynamically allocated wires

pennylane/allocation.py

@@ -158,10 +158,10 @@ def circuit():
             return qp.expval(qp.Z(0))
 
     >>> print(qp.draw(circuit)())
-                0: ──H────────╭●────╭●──────────────────────┤  <Z>
-    <DynamicWire>: ──Allocate─╰X────╰X───────────Deallocate─┤
-    <DynamicWire>: ─╭Allocate─╭SWAP─╭Deallocate─────────────┤
-    <DynamicWire>: ─╰Allocate─╰SWAP─╰Deallocate─────────────┤
+    0: ──────────────H────╭●─╭●────┤  <Z>
+         |0>├─╭SWAP──┤    │  │
+         |0>├─╰SWAP──┤    │  │
+                     |0>├─╰X─╰X──┤
 
     Here, three dynamic wires were allocated in the circuit originally. When PennyLane determines
     which concrete values to use for dynamic wires to send to the device for execution, we can see
@@ -276,10 +276,10 @@ def circuit():
             return qp.expval(qp.Z(0))
 
     >>> print(qp.draw(circuit)())
-                0: ──H───────────────────────┤  <Z>
-                1: ──H───────────────────────┤
-    <DynamicWire>: ─╭Allocate──H─╭Deallocate─┤
-    <DynamicWire>: ─╰Allocate──H─╰Deallocate─┤
+    0: ───────────H─┤  <Z>
+    1: ───────────H─┤
+         |0>├──H──┤
+         |0>├──H──┤
 
     Equivalenty, ``allocate`` can be used in-line along with :func:`~.deallocate` for manual
     handling:
@@ -318,12 +318,12 @@ def circuit():
                 return qp.expval(qp.Z(0))
 
         >>> print(qp.draw(circuit)())
-                    0: ──H─────────────────────╭●───────────────────────╭●─────────────┤  <Z>
-        <DynamicWire>: ──Allocate──┤↗│  │0⟩────│──────────Deallocate────│──────────────┤
-        <DynamicWire>: ──Allocate───║────────Z─╰X─────────Deallocate────│──────────────┤
-        <DynamicWire>: ─────────────║────────║──Allocate──┤↗│  │0⟩──────│───Deallocate─┤
-        <DynamicWire>: ─────────────║────────║──Allocate───║──────────Z─╰X──Deallocate─┤
-                                    ╚════════╝             ╚══════════╝
+        0: ──────────────────H─────╭●───────────╭●────┤  <Z>
+            |0>├──┤↗│  │0⟩────────│──────────┤ │
+            ├──────║────────Z─────╰X─────────┤ │
+                   ║        ║|0>├──┤↗│  │0⟩────│───┤
+                   ║        ║├──────║────────Z─╰X──┤
+                   ╚════════╝       ╚════════╝
 
         The user-level circuit drawing shows four separate allocations and deallocations (two per
         loop iteration). However, the circuit that the device receives gets automatically compiled
@@ -356,9 +356,9 @@ def circuit():
                 return qp.expval(qp.Z(0))
 
         >>> print(qp.draw(circuit, level="user")())
-        <DynamicWire>: ──Allocate──H──Deallocate─┤
-        <DynamicWire>: ──Allocate──X──Deallocate─┤
-                    0: ──────────────────────────┤  <Z>
+        0: ──────────────────────┤  <Z>
+            |0>├──X──┤
+                    |0>├──H──┤
         >>> print(qp.draw(circuit, level="device")())
         0: ─────────────────┤  <Z>
         1: ──H──┤↗│  │0⟩──X─┤

pennylane/decomposition/decomposition_rule.py

@@ -375,11 +375,11 @@ def circuit():
               return qp.probs(wires=[0, 1, 2, 3])
 
        >>> print(qp.draw(circuit)())
-       <DynamicWire>: ──Allocate─╭X─╭●───────────────────╭X──Deallocate─┤
-                   0: ───────────├●─│────────────────────├●─────────────┤ ╭Probs
-                   1: ───────────├●─│────────────────────├●─────────────┤ ├Probs
-                   2: ───────────╰●─│────────────────────╰●─────────────┤ ├Probs
-                   3: ──────────────╰Rot(0.10,0.20,0.30)────────────────┤ ╰Probs
+        0: ───────╭●──────────────────────╭●────┤ ╭Probs
+        1: ───────├●──────────────────────├●────┤ ├Probs
+        2: ───────├●──────────────────────├●────┤ ├Probs
+        3: ───────│──╭Rot(0.10,0.20,0.30)─│─────┤ ╰Probs
+             |0>├─╰X─╰●───────────────────╰X──┤
 
     """
 

pennylane/drawer/_add_obj.py

@@ -27,6 +27,7 @@
 
 from functools import singledispatch
 
+from pennylane.allocation import Allocate, Deallocate
 from pennylane.measurements import (
     CountsMP,
     DensityMatrixMP,
@@ -69,8 +70,8 @@ def _add_cond_grouping_symbols(op, layer_str, config):
         ctrl_symbol = "╝"
     layer_str[max_b + n_wires] = f"═{ctrl_symbol}"
 
-    for w in range(max_w + 1, max(config.wire_map.values()) + 1):
-        layer_str[w] = "─║"
+    for row in range(max_w + 1, config.n_wires):
+        layer_str[row] = config.wire_filler(row) + "║"
 
     for b in range(max_b):
         if b in mapped_bits:
@@ -121,8 +122,8 @@ def _add_mid_measure_grouping_symbols(op, layer_str, config):
     bit = config.bit_map[op] + n_wires
     layer_str[bit] += " ╚"
 
-    for w in range(mapped_wire + 1, n_wires):
-        layer_str[w] += "─║"
+    for row in range(mapped_wire + 1, config.n_wires):
+        layer_str[row] += config.wire_filler(row) + "║"
 
     for b in range(n_wires, bit):
         filler = " " if layer_str[b][-1] == " " else "═"
@@ -164,6 +165,21 @@ def _add_obj(
     raise NotImplementedError(f"unable to draw object {obj}")
 
 
+@_add_obj.register(Allocate)
+def _(obj, layer_str, config, tape_cache=None, skip_grouping_symbols=False):
+    label = "|0>├" if obj.state.value == "zero" else "├"
+    for w in obj.wires:
+        layer_str[config.wire_map[w]] += label
+    return layer_str
+
+
+@_add_obj.register(Deallocate)
+def _(obj, layer_str, config, tape_cache=None, skip_grouping_symbols=False):
+    for w in obj.wires:
+        layer_str[config.wire_map[w]] += "┤"
+    return layer_str
+
+
 @_add_obj.register
 def _add_cond(obj: Conditional, layer_str, config, tape_cache=None, skip_grouping_symbols=False):
     layer_str = _add_cond_grouping_symbols(obj, layer_str, config)

pennylane/drawer/draw.py

@@ -23,6 +23,7 @@
 from typing import TYPE_CHECKING, Literal
 
 from pennylane import math
+from pennylane.allocation import DynamicWire
 from pennylane.tape import make_qscript
 from pennylane.workflow import construct_batch
 
@@ -317,9 +318,9 @@ def wrapper(*args, **kwargs):
             _wire_order = wire_order
         else:
             try:
-                _wire_order = sorted(tape.wires)
+                _wire_order = sorted(w for w in tape.wires if not isinstance(w, DynamicWire))
             except TypeError:
-                _wire_order = tape.wires
+                _wire_order = [w for w in tape.wires if not isinstance(w, DynamicWire)]
 
         return tape_text(
             tape,
@@ -356,9 +357,9 @@ def wrapper(*args, **kwargs):
             _wire_order = qnode.device.wires
         else:
             try:
-                _wire_order = sorted(tapes[0].wires)
+                _wire_order = sorted(w for w in tapes[0].wires if not isinstance(w, DynamicWire))
             except TypeError:
-                _wire_order = tapes[0].wires
+                _wire_order = [w for w in tapes[0].wires if not isinstance(w, DynamicWire)]
 
         cache = {"tape_offset": 0, "matrices": []}
         res = [

pennylane/drawer/drawable_layers.py

@@ -15,8 +15,10 @@
 This module contains a helper function to sort operations into layers.
 """
 
+from dataclasses import dataclass, field
 from functools import singledispatch
 
+from pennylane.allocation import Allocate, Deallocate, DynamicWire
 from pennylane.measurements import MeasurementProcess
 from pennylane.ops import (
     Conditional,
@@ -147,7 +149,71 @@ def _handle_cond(op: Conditional, wire_map, bit_map):
     return set(range(min_wire, max_wire + 1))
 
 
-def drawable_layers(operations, wire_map=None, bit_map=None):
+@dataclass
+class LayersData:
+    """Data for putting operations into layers."""
+
+    ops_in_layer: list[list] = field(default_factory=lambda: [[]])
+    """Lists of which operators should be placed in each layer."""
+
+    occupied_wires_per_layer: list[set[int]] = field(default_factory=lambda: [set()])
+    """The mapped wires that will be occupied in the drawing at each layer."""
+
+    used_cwires_per_layer: list[set[int]] = field(default_factory=lambda: [set()])
+    """The clasical wires that will be occupied in each layer."""
+
+    waiting_dynamic_wires: list[DynamicWire] = field(default_factory=list)
+    """DynamicWires that are waiting for the first interaction between
+    the dynamic wires and an algorithmic wire. 
+
+    This allows us to push the allocation and initial setup to the right and conserve
+    drawer space, keeping things together.  See ``insert_waiting_ops``.
+    """
+
+    waiting_dynamic_ops: list = field(default_factory=list)
+    """The Allocate instructions and operators that are waiting on the first interaction between
+    the waiting_dynamic_wires and an algorithmic wire.
+    """
+
+    @property
+    def max_layer(self) -> int:
+        """The maximum number of layers present."""
+        return len(self.ops_in_layer) - 1
+
+    def add_to_layer(self, op, op_layer, occupied_wires, used_cwires):
+        """Adds an operation to a layer."""
+        if op_layer > self.max_layer:
+            self.occupied_wires_per_layer.append(set())
+            self.ops_in_layer.append([])
+            self.used_cwires_per_layer.append(set())
+
+        self.ops_in_layer[op_layer].append(op)
+        self.occupied_wires_per_layer[op_layer].update(occupied_wires)
+        self.used_cwires_per_layer[op_layer].update(used_cwires)
+
+    def insert_waiting_ops(self, op_layer, wire_map, bit_map):
+        """Insert ops in ``waiting_dynamic_ops`` before ``op_layer``."""
+        inner_layers = drawable_layers(
+            self.waiting_dynamic_ops,
+            wire_map=wire_map,
+            bit_map=bit_map,
+            _dynamic_wires=False,
+        )
+        for i, layer in enumerate(reversed(inner_layers)):
+            insert_layer = op_layer - i - 1
+            if insert_layer < 0:
+                self.ops_in_layer.insert(0, [])
+                op_layer += 1
+                insert_layer = op_layer - i - 1
+                self.occupied_wires_per_layer.insert(0, set())
+                self.used_cwires_per_layer.insert(0, set())
+            self.ops_in_layer[insert_layer].extend(layer)
+        self.waiting_dynamic_wires = []
+        self.waiting_dynamic_ops = []
+        return op_layer
+
+
+def drawable_layers(operations, wire_map=None, bit_map=None, _dynamic_wires=True):
     """Determine non-overlapping yet dense placement of operations into layers for drawing.
 
     Args:
@@ -157,6 +223,8 @@ def drawable_layers(operations, wire_map=None, bit_map=None):
         wire_map (dict): A map from wire label to non-negative integers. Defaults to None.
         bit_map (dict): A map containing mid-circuit measurements used for classical conditions
             or collecting statistics as keys. Defaults to None.
+        _dynamic_wires (bool): **Internal**.  Whether or not try pushing allocations and ops
+            on only dynamic wires to the right next to the first time the dynamic wires are used.
 
     Returns:
         (list[set[~.Operator]], list[set[~.MeasurementProcess]]) : Each index is a set of operations
@@ -188,14 +256,22 @@ def drawable_layers(operations, wire_map=None, bit_map=None):
     bit_map = bit_map or {}
 
     # initialize for operation layers
-    max_layer = 0
-    occupied_wires_per_layer = [set()]
-    ops_in_layer = [[]]
-    used_cwires_per_layer = [set()]
+    data = LayersData()
 
     # loop over operations
     for op in operations:
-        if isinstance(op, MeasurementProcess) and op.mv is not None:
+        if _dynamic_wires and isinstance(op, Allocate):
+            data.waiting_dynamic_ops.append(op)
+            data.waiting_dynamic_wires.extend(op.wires)
+
+        elif (
+            _dynamic_wires
+            and all(w in data.waiting_dynamic_wires for w in op.wires)
+            and op.wires  # if no wires (GlobalPhase) then do not put into waiting_dynamic_ops
+            and not isinstance(op, Deallocate)  # deallocate should force putting into circuit
+        ):
+            data.waiting_dynamic_ops.append(op)
+        elif isinstance(op, MeasurementProcess) and op.mv is not None:
             # Only terminal measurements that collect mid-circuit measurement statistics have
             # op.mv != None.
             # Get the occupied classical wires of the measurement process and find which layer
@@ -208,34 +284,29 @@ def drawable_layers(operations, wire_map=None, bit_map=None):
             )
             op_occupied_cwires = set(range(min(mapped_cwires), max(mapped_cwires) + 1))
             op_layer = _recursive_find_mcm_stats_layer(
-                max_layer, op_occupied_cwires, used_cwires_per_layer
+                data.max_layer, op_occupied_cwires, data.used_cwires_per_layer
             )
 
+            data.add_to_layer(op, op_layer, op_occupied_wires, op_occupied_cwires)
+
         else:
             # Find occupied wires of the operator/measurement process and find which layer to
             # put it in.
             op_occupied_wires = _get_op_occupied_wires(op, wire_map, bit_map)
             try:
                 op_layer = _recursive_find_layer(
-                    max_layer, op_occupied_wires, occupied_wires_per_layer
+                    data.max_layer, op_occupied_wires, data.occupied_wires_per_layer
                 )
             except RecursionError as e:
                 raise RecursionError(
-                    f"Drawer is currently at depth {max_layer}, which is too deep to handle. "
+                    f"Drawer is currently at depth {data.max_layer}, which is too deep to handle. "
                     "Try drawing a smaller subset of your circuit instead."
                 ) from e
             op_occupied_cwires = set()
 
-        # see if need to add new layer
-        if op_layer > max_layer:
-            max_layer += 1
-            occupied_wires_per_layer.append(set())
-            ops_in_layer.append([])
-            used_cwires_per_layer.append(set())
+            if _dynamic_wires and any(w in data.waiting_dynamic_wires for w in op.wires):
+                op_layer = data.insert_waiting_ops(op_layer, wire_map, bit_map)
 
-        # add to op_layer
-        ops_in_layer[op_layer].append(op)
-        occupied_wires_per_layer[op_layer].update(op_occupied_wires)
-        used_cwires_per_layer[op_layer].update(op_occupied_cwires)
+            data.add_to_layer(op, op_layer, op_occupied_wires, op_occupied_cwires)
 
-    return list(filter(None, ops_in_layer[:-1])) + ops_in_layer[-1:]
+    return list(filter(None, data.ops_in_layer[:-1])) + data.ops_in_layer[-1:]