✨ Add missing TKET RL actions

CHANGELOG.md

@@ -12,6 +12,8 @@ releases may include breaking changes.
 
 ### Added
 
+- ✨ Add TKET's `KAKDecomposition`, `GraphPlacement`, and `NoiseAwarePlacement`
+  passes to the RL actions ([#796]) ([**@flowerthrower**])
 - ✨ Add BQSKit's `QSDPass` to the RL synthesis actions ([#795])
   ([**@flowerthrower**])
 - ✨ Add Qiskit's `TrivialLayout`, `ElidePermutations`, `SabreSwap`,
@@ -30,6 +32,8 @@ releases may include breaking changes.
 
 ### Changed
 
+- 🐛 Mask TKET layout and routing actions for circuits containing operations
+  wider than two qubits ([#796]) ([**@flowerthrower**])
 - 🐛 Make the `OptimizeCliffords` RL action collect standard Clifford gates
   before optimizing them ([#794]) ([**@flowerthrower**])
 - 🔥 Drop support for Python 3.10 ([#773]) ([**@denialhaag**])
@@ -102,6 +106,7 @@ for previous changelogs._
 
 [#773]: https://github.com/munich-quantum-toolkit/predictor/pull/771
 [#769]: https://github.com/munich-quantum-toolkit/predictor/pull/769
+[#796]: https://github.com/munich-quantum-toolkit/predictor/pull/796
 [#795]: https://github.com/munich-quantum-toolkit/predictor/pull/795
 [#794]: https://github.com/munich-quantum-toolkit/predictor/pull/794
 [#758]: https://github.com/munich-quantum-toolkit/predictor/pull/758

UPGRADING.md

@@ -6,6 +6,17 @@ of changes including minor and patch releases, please refer to the
 
 ## [Unreleased]
 
+### Expanded TKET action set
+
+The RL action space now includes TKET's `KAKDecomposition` optimization action
+and the `GraphPlacement` and `NoiseAwarePlacement` layout actions. TKET layout
+and routing actions are masked while a circuit contains operations on more than
+two qubits; those operations must be decomposed first.
+
+Existing RL models must be retrained because the action-space size and the
+indices of later actions have changed. Code that persists or selects actions by
+numeric index must be updated.
+
 ### Expanded Qiskit action set
 
 The RL action space now includes the following Qiskit passes:

src/mqt/predictor/rl/actions/registry.py

@@ -59,6 +59,7 @@ def get_actions_by_pass_type() -> dict[PassType, list[Action]]:
     qiskit_actions.qiskit_o3_action(),
     *qiskit_actions.qiskit_optimization_actions(),
     qiskit_actions.qiskit_final_optimization_action(),
+    *tket_actions.tket_layout_actions(),
     tket_actions.tket_routing_action(),
     *tket_actions.tket_optimization_actions(),
     *bqskit_actions.bqskit_layout_actions(),

src/mqt/predictor/rl/actions/tket_actions.py

@@ -10,17 +10,28 @@
 
 from __future__ import annotations
 
+import logging
 import operator
+from collections import defaultdict
+from functools import cache
 from typing import TYPE_CHECKING, cast
 
 from pytket import Qubit
 from pytket._tket.passes import BasePass as TketBasePass  # ruff:ignore[import-private-name]
 from pytket.architecture import Architecture
 from pytket.circuit import Node
 from pytket.extensions.qiskit import qiskit_to_tk, tk_to_qiskit
-from pytket.passes import CliffordSimp, FullPeepholeOptimise, PeepholeOptimise2Q, RemoveRedundancies, RoutingPass
-from pytket.placement import place_with_map
-from qiskit.transpiler import Layout
+from pytket.passes import (
+    CliffordSimp,
+    FullPeepholeOptimise,
+    KAKDecomposition,
+    PeepholeOptimise2Q,
+    RemoveRedundancies,
+    RoutingPass,
+)
+from pytket.placement import GraphPlacement, NoiseAwarePlacement, Placement, place_with_map
+from qiskit.transpiler import CouplingMap, Layout, PassManager, TranspileLayout
+from qiskit.transpiler.passes import ApplyLayout, EnlargeWithAncilla, FullAncillaAllocation, SetLayout
 
 from mqt.predictor.rl.actions.base import CompilationOrigin, DeferredDeviceAction, DeviceIndependentAction, PassType
 
@@ -29,11 +40,13 @@
 
     from pytket import Circuit
     from qiskit import QuantumCircuit
-    from qiskit.passmanager.base_tasks import Task
-    from qiskit.transpiler import Target, TranspileLayout
+    from qiskit.circuit import Qubit as QiskitQubit
+    from qiskit.transpiler import Target
 
     from mqt.predictor.rl.actions.base import Action
 
+logger = logging.getLogger("mqt-predictor")
+
 
 class PreProcessTKETRoutingAfterQiskitLayout:
     """Pre-process TKET routing for circuits that already carry a Qiskit layout.
@@ -49,6 +62,90 @@ def apply(self, circuit: Circuit) -> None:
         place_with_map(circuit=circuit, qmap=mapping)
 
 
+@cache
+def _prepare_noise_data(device: Target) -> tuple[dict[Node, float], dict[tuple[Node, Node], float], dict[Node, float]]:
+    """Extract calibration errors for TKET's noise-aware placement."""
+    node_error_samples: defaultdict[Node, list[float]] = defaultdict(list)
+    link_error_samples: defaultdict[tuple[Node, Node], list[float]] = defaultdict(list)
+    readout_errors: dict[Node, float] = {}
+
+    for operation_name in device.operation_names:
+        if operation_name == "measure":
+            continue
+        for qubits, properties in device[operation_name].items():
+            if qubits is None or properties is None or properties.error is None:
+                continue
+            if len(qubits) == 1:
+                node_error_samples[Node(qubits[0])].append(properties.error)
+            elif len(qubits) == 2:
+                link_error_samples[Node(qubits[0]), Node(qubits[1])].append(properties.error)
+
+    if "measure" in device:
+        for qubits, properties in device["measure"].items():
+            if qubits is not None and len(qubits) == 1 and properties is not None and properties.error is not None:
+                readout_errors[Node(qubits[0])] = properties.error
+
+    node_errors = {node: sum(errors) / len(errors) for node, errors in node_error_samples.items()}
+    link_errors = {link: sum(errors) / len(errors) for link, errors in link_error_samples.items()}
+    return node_errors, link_errors, readout_errors
+
+
+def _noise_aware_placement(device: Target) -> list[Placement]:
+    node_errors, link_errors, readout_errors = _prepare_noise_data(device)
+    return [
+        NoiseAwarePlacement(
+            Architecture(list(device.build_coupling_map())),
+            node_errors=node_errors,
+            link_errors=link_errors,
+            readout_errors=readout_errors,
+            timeout=5000,
+            maximum_matches=5000,
+        )
+    ]
+
+
+def _translate_placement(
+    circuit: QuantumCircuit,
+    placement: dict[Qubit, Node],
+    action_name: str,
+    num_device_qubits: int,
+) -> Layout | None:
+    qiskit_qubits: dict[tuple[str, tuple[int, ...]], QiskitQubit] = {}
+    for qubit in circuit.qubits:
+        location = circuit.find_bit(qubit)
+        if location.registers:
+            register, register_index = location.registers[0]
+            qiskit_qubits[register.name, (register_index,)] = qubit
+        else:
+            qiskit_qubits["q", (location.index,)] = qubit
+
+    qiskit_mapping: dict[QiskitQubit, int] = {}
+    unassigned_qubits: list[QiskitQubit] = []
+    used_physical_indices: set[int] = set()
+    for tket_qubit, target_node in placement.items():
+        qiskit_qubit = qiskit_qubits.get((str(tket_qubit.reg_name), tuple(int(i) for i in tket_qubit.index)))
+        if qiskit_qubit is None:
+            logger.warning("Placement failed (%s): unknown logical qubit %s.", action_name, tket_qubit)
+            return None
+
+        if target_node.reg_name == "node" and target_node.index:
+            physical_index = int(target_node.index[0])
+            qiskit_mapping[qiskit_qubit] = physical_index
+            used_physical_indices.add(physical_index)
+        else:
+            unassigned_qubits.append(qiskit_qubit)
+
+    unassigned_qubits.extend(qubit for qubit in circuit.qubits if qubit not in qiskit_mapping)
+    unassigned_qubits = list(dict.fromkeys(unassigned_qubits))
+    remaining_indices = [index for index in range(num_device_qubits) if index not in used_physical_indices]
+    if len(remaining_indices) < len(unassigned_qubits):
+        logger.warning("Placement failed (%s): insufficient free physical qubits.", action_name)
+        return None
+
+    qiskit_mapping.update(zip(unassigned_qubits, remaining_indices, strict=False))
+    return Layout(qiskit_mapping)
+
+
 def tket_optimization_actions() -> list[Action]:
     """Returns the TKET optimization actions."""
     return [
@@ -70,6 +167,15 @@ def tket_optimization_actions() -> list[Action]:
             preserves_routing=False,
             preserves_synthesis=False,
         ),
+        DeviceIndependentAction(
+            "KAKDecomposition",
+            CompilationOrigin.TKET,
+            PassType.OPT,
+            [KAKDecomposition(allow_swaps=False)],
+            preserves_layout=True,
+            preserves_routing=True,
+            preserves_synthesis=False,
+        ),
         DeviceIndependentAction(
             "FullPeepholeOptimiseCX",
             CompilationOrigin.TKET,
@@ -91,19 +197,40 @@ def tket_optimization_actions() -> list[Action]:
     ]
 
 
+def tket_layout_actions() -> list[Action]:
+    """Return the TKET layout actions."""
+    return [
+        DeferredDeviceAction(
+            "GraphPlacement",
+            CompilationOrigin.TKET,
+            PassType.LAYOUT,
+            transpile_pass=lambda device: [
+                GraphPlacement(
+                    Architecture(list(device.build_coupling_map())),
+                    timeout=5000,
+                    maximum_matches=5000,
+                )
+            ],
+        ),
+        DeferredDeviceAction(
+            "NoiseAwarePlacement",
+            CompilationOrigin.TKET,
+            PassType.LAYOUT,
+            transpile_pass=_noise_aware_placement,
+        ),
+    ]
+
+
 def tket_routing_action() -> Action:
     """Returns the TKET routing action."""
     return DeferredDeviceAction(
         "RoutingPass",
         CompilationOrigin.TKET,
         PassType.ROUTING,
-        transpile_pass=lambda device: cast(
-            "list[Task]",
-            [
-                PreProcessTKETRoutingAfterQiskitLayout(),
-                RoutingPass(Architecture(list(device.build_coupling_map()))),
-            ],
-        ),
+        transpile_pass=lambda device: [
+            PreProcessTKETRoutingAfterQiskitLayout(),
+            RoutingPass(Architecture(list(device.build_coupling_map()))),
+        ],
     )
 
 
@@ -135,10 +262,43 @@ def run_tket_action(
     """Apply a TKET action and return the updated circuit and layout metadata."""
     tket_qc = qiskit_to_tk(circuit, preserve_param_uuid=True)
     if callable(action.transpile_pass):
-        factory = cast("Callable[[Target], list[Task]]", action.transpile_pass)
+        factory = cast(
+            "Callable[[Target], list[TketBasePass | PreProcessTKETRoutingAfterQiskitLayout | Placement]]",
+            action.transpile_pass,
+        )
         passes = factory(device)
     else:
-        passes = cast("list[Task]", action.transpile_pass)
+        passes = cast("list[TketBasePass | PreProcessTKETRoutingAfterQiskitLayout | Placement]", action.transpile_pass)
+
+    if action.pass_type == PassType.LAYOUT:
+        if len(passes) != 1 or not isinstance(passes[0], Placement):
+            msg = f"TKET layout action {action.name} must provide exactly one placement."
+            raise TypeError(msg)
+        try:
+            placement = passes[0].get_placement_map(tket_qc)
+        except (RuntimeError, TypeError, ValueError) as error:
+            logger.warning("Placement failed (%s): %s.", action.name, error)
+            return circuit, layout
+
+        qiskit_layout = _translate_placement(circuit, placement, action.name, device.num_qubits)
+        if qiskit_layout is None:
+            return circuit, layout
+        pass_manager = PassManager([
+            SetLayout(qiskit_layout),
+            FullAncillaAllocation(coupling_map=CouplingMap(device.build_coupling_map())),
+            EnlargeWithAncilla(),
+            ApplyLayout(),
+        ])
+        altered_qc = pass_manager.run(circuit)
+        applied_layout = cast("Layout", pass_manager.property_set["layout"])
+        return altered_qc, TranspileLayout(
+            initial_layout=applied_layout,
+            input_qubit_mapping=pass_manager.property_set["original_qubit_indices"],
+            final_layout=pass_manager.property_set["final_layout"],
+            _output_qubit_list=altered_qc.qubits,
+            _input_qubit_count=circuit.num_qubits,
+        )
+
     for pass_ in passes:
         assert isinstance(pass_, TketBasePass | PreProcessTKETRoutingAfterQiskitLayout)
         pass_.apply(tket_qc)
@@ -154,8 +314,10 @@ def run_tket_action(
     return altered_qc, layout
 
 
-def is_tket_action_available(*, action: Action, has_layout: bool) -> bool:
+def is_tket_action_available(*, action: Action, has_layout: bool, has_wide_operations: bool) -> bool:
     """Return whether a TKET action is available for the current layout state."""
+    if has_wide_operations and action.pass_type in {PassType.LAYOUT, PassType.ROUTING}:
+        return False
     # TKET layout/optimization actions must not run after a Qiskit layout has been set
     # (it is not clear how tket will handle the layout). TKET routing actions, however, are
     #  designed to work after a Qiskit layout via PreProcessTKETRoutingAfterQiskitLayout.

src/mqt/predictor/rl/predictorenv.py

@@ -493,6 +493,7 @@ def action_masks(self) -> list[bool]:
             A dense boolean mask ordered like ``self.action_set``.
         """
         has_layout = self.layout is not None
+        has_wide_operations = any(len(instruction.qubits) > 2 for instruction in self.state.data)
         valid_action_indices = set(self.valid_actions)
         action_mask: list[bool] = []
 
@@ -508,7 +509,13 @@ def action_masks(self) -> list[bool]:
             if action.origin == CompilationOrigin.QISKIT:
                 action_mask.append(is_qiskit_action_available(action, self.device))
             elif action.origin == CompilationOrigin.TKET:
-                action_mask.append(is_tket_action_available(action=action, has_layout=has_layout))
+                action_mask.append(
+                    is_tket_action_available(
+                        action=action,
+                        has_layout=has_layout,
+                        has_wide_operations=has_wide_operations,
+                    )
+                )
             elif action.origin == CompilationOrigin.BQSKIT:
                 action_mask.append(
                     is_bqskit_action_available(