[CQT-421] Implement QGymMapper mapper pass

CHANGELOG.md

@@ -14,6 +14,7 @@ This project adheres to [Semantic Versioning](http://semver.org/).
 
 ### Added
 
+- `QGymMapper` mapper pass
 - `U`(theta, phi, lambda) gate to default single-qubit gates
 - `Z90` and `mZ90` pi-half rotation gates (equivalent to `S` and `Sdag` gates)
 

docs/compilation-passes/mapping/index.md

@@ -14,3 +14,9 @@ thereby reducing the overall error rate and improving the fidelity of the quantu
 Efficient qubit mapping can significantly enhance the performance of the quantum circuit.
 By strategically placing qubits, the compiler can reduce the number of additional operations required,
 leading to faster and more reliable quantum computations.
+
+The following mapping passes are available in Opensquirrel:
+- [Hardcoded Mapper](hardcoded-mapper.md) (`HardcodedMapper`)
+- [Identity Mapper](identity-mapper.md) (`IdentitiyMapper`)
+- [Random Mapper](random-mapper.md) (`RandomMapper`)
+- [Qgym Mapper](qgym-mapper.md) (`QGymMapper`)

docs/compilation-passes/mapping/qgym-mapper.md

@@ -0,0 +1,69 @@
+The [Qgym](https://github.com/QuTech-Delft/qgym) package functions in a manner similar 
+to the well known gym package, in the  sense that it provides a number of environments 
+on which reinforcement learning (RL) agents can be applied. The main purpose of qgym is 
+to develop reinforcement learning environments which represent various passes of the 
+[OpenQL framework](https://arxiv.org/abs/2005.13283).
+
+The package offers RL-based environments resembling quantum compilation steps, namely 
+for initial mapping, qubit routing, and gate scheduling. The environments offer all the 
+relevant components needed to train agents, including states and action spaces, and 
+(customizable) reward functions (basically all the components required by a Markov 
+Decision Process). Furthermore, the actual training of the agents is handled by the 
+[StableBaselines3](https://github.com/DLR-RM/stable-baselines3) python package, which 
+offers reliable, customizable, out of the box Pytorch implementations of DRL agents.
+
+The initial mapping problem is translated to a RL context within Qgym in the following 
+manner. The setup begins with a fixed connection graph (an undirected graph 
+representation of the hardware connectivity scheme), static across all episodes. Each 
+episode introduces a novel, randomly generated interaction graph (undirected graph 
+representation of the qubit interactions within the circuit) for the agent to observe, 
+alongside an initially empty mapping. At every step, the agent can map a virtual 
+(logical) qubit to a physical qubit until the mapping is fully established. In theory, 
+this process enables the training of agents that are capable of managing various 
+interaction graphs on a predetermined hardware layout. Both the interaction and 
+connection graphs are easily represented via [Networkx](https://networkx.org/) graphs.
+
+At the moment, the following DRL agents can be used to map circuits in Opensquirrel:
+
+- Proximal Policy Optimization (PPO)
+- Advantage Actor-Critic (A2C)
+- Trust Region Policy Optimization (TRPO)
+- Recurrent PPO
+- PPO with illegal action masking
+
+The last three agents in the above list can be imported from the extension/experimental 
+package of StableBaselines3, namely [sb3-contrib](https://github.com/Stable-Baselines-Team/stable-baselines3-contrib). 
+
+The following code snippet demonstrates the usage of the `QGymMapper`. Assume that you
+have a `connectivities.json` file containing some hardware connectivity schemes, as 
+well as a `TRPO.zip` file containing the weights of a trained agent in your working 
+directory. 
+
+```python
+from opensquirrel.passes.mapper import QGymMapper
+from opensquirrel import CircuitBuilder
+import networkx as nx
+import json
+
+with open('connectivities.json', 'r') as file:
+    connectivities = json.load(file)
+
+hardware_connectivitiy = connectivities["tuna-5"]
+
+connection_graph = nx.Graph()
+connection_graph.add_edges_from(hardware_connectivity)
+
+qgym_mapper = QGymMapper(agent_class = "TRPO", agent_path = "TRPO.zip", 
+                        connection_graph=connection_graph)
+
+builder = CircuitBuilder(5)
+builder.H(0)
+builder.CNOT(0, 1)
+builder.H(2)
+builder.CNOT(1, 2)
+builder.CNOT(2, 4)
+builder.CNOT(3, 4)
+circuit = builder.to_circuit()
+
+circuit.map(mapper = qgym_mapper)
+```

opensquirrel/passes/mapper/__init__.py

@@ -1,9 +1,15 @@
 from opensquirrel.passes.mapper.mip_mapper import MIPMapper
 from opensquirrel.passes.mapper.simple_mappers import HardcodedMapper, IdentityMapper, RandomMapper
 
+try:
+    from opensquirrel.passes.mapper.qgym_mapper import QGymMapper
+except ImportError:
+    QGymMapper: type | None = None  # type: ignore[no-redef]
+
 __all__ = [
     "HardcodedMapper",
     "IdentityMapper",
     "MIPMapper",
+    "QGymMapper",
     "RandomMapper",
 ]

opensquirrel/passes/mapper/connectivities.json

@@ -0,0 +1 @@
+{"tuna-5": [[0, 2], [1, 2], [2, 3], [2, 4]], "starmon-7": [[0, 2], [0, 3], [2, 5], [3, 5], [3, 1], [3, 6], [6, 4], [1, 4]]}

opensquirrel/passes/mapper/qgym_mapper.py

@@ -0,0 +1,126 @@
+from __future__ import annotations
+
+import importlib
+from itertools import combinations
+from typing import TYPE_CHECKING, Any
+
+import networkx as nx
+from qgym.envs import InitialMapping
+
+from opensquirrel.ir import IR, Instruction
+from opensquirrel.passes.mapper.general_mapper import Mapper
+from opensquirrel.passes.mapper.mapping import Mapping
+
+if TYPE_CHECKING:
+    from stable_baselines3.common.base_class import BaseAlgorithm
+
+
+class QGymMapper(Mapper):
+    """
+    QGym-based Mapper using a Stable-Baselines3 agent.
+    - Builds a qubit interaction graph from the IR.
+    - Runs the InitialMapping environment with the trained agent.
+    - Returns a Mapping compatible with OpenSquirrel.
+    """
+
+    def __init__(
+        self,
+        agent_class: str,
+        agent_path: str,
+        hardware_connectivity: nx.Graph,
+        env_kwargs: dict[str, Any] | None = None,
+        **kwargs: Any,
+    ) -> None:
+        super().__init__(**kwargs)
+        self.env = InitialMapping(connection_graph=hardware_connectivity, **(env_kwargs or {}))
+        self.hardware_connectivity = hardware_connectivity
+        self.agent = self._load_agent(agent_class, agent_path)
+
+    def _load_agent(self, agent_class: str, agent_path: str) -> BaseAlgorithm:
+        """Load a trained Stable-Baselines3 agent from a file."""
+        if agent_class in ["PPO", "A2C"]:
+            sb3 = importlib.import_module("stable_baselines3")
+        else:
+            sb3 = importlib.import_module("sb3_contrib")
+        agent_cls = getattr(sb3, agent_class)
+        return agent_cls.load(agent_path)
+
+    def map(self, ir: IR, qubit_register_size: int) -> Mapping:
+        """
+        Compute an initial logical-to-physical qubit mapping using a trained
+        Stable-Baselines3 agent acting in the QGym InitialMapping environment.
+
+        Args:
+            ir (IR): Intermediate representation of the quantum circuit to be mapped.
+            qubit_register_size (int): Number of logical (virtual) qubits in the circuit.
+
+        Returns:
+            Mapping: Mapping from virtual to physical qubits.
+
+        Raises:
+            ValueError: If the number of logical qubits differs from the number of physical qubits.
+            ValueError: If the agent produces an incomplete or invalid mapping.
+            RuntimeError: If no 'mapping' key is found in the final observation.
+        """
+        num_physical = self.hardware_connectivity.number_of_nodes()
+        if qubit_register_size != num_physical:
+            error_msg = f"QGym requires equal logical and physical qubits: logical={qubit_register_size}, physical={num_physical}"  # noqa: E501
+            raise ValueError(error_msg)
+
+        circuit_graph = self._ir_to_interaction_graph(ir)
+
+        obs, _ = self.env.reset(options={"interaction_graph": circuit_graph})
+
+        done = False
+        last_obs: Any = obs
+        while not done:
+            action, _ = self.agent.predict(obs, deterministic=True)
+            obs, _, terminated, truncated, _ = self.env.step(action)
+            done = terminated or truncated
+            last_obs = obs
+
+        mapping_data = self._extract_mapping_data(last_obs)
+        return self._get_mapping(mapping_data, qubit_register_size)
+
+    def _ir_to_interaction_graph(self, ir: IR) -> nx.Graph:
+        """Build an undirected interaction graph representation of the IR."""
+        interaction_graph = nx.Graph()
+        for instr in ir.statements:
+            if not isinstance(instr, Instruction):
+                continue
+            qubit_indices = [q.index for q in instr.get_qubit_operands()]
+            for q_index in qubit_indices:
+                interaction_graph.add_node(q_index)
+            if len(qubit_indices) >= 2:
+                for q_i, q_j in combinations(qubit_indices, 2):
+                    if interaction_graph.has_edge(q_i, q_j):
+                        interaction_graph[q_i][q_j]["weight"] = interaction_graph[q_i][q_j].get("weight", 1) + 1
+                    else:
+                        interaction_graph.add_edge(q_i, q_j, weight=1)
+        return interaction_graph
+
+    def _get_mapping(self, mapping_data: Any, qubit_register_size: int) -> Mapping:
+        """Convert QGym's physical-to-logical mapping to OpenSquirrel's logical-to-physical mapping."""
+        physical_to_logical = mapping_data.tolist()
+
+        if len(physical_to_logical) != qubit_register_size:
+            error_msg = f"Mapping length {len(physical_to_logical)} != qubit_register_size {qubit_register_size}."
+            raise ValueError(error_msg)
+
+        logical_to_physical = [-1] * qubit_register_size
+        for physical_qubit, logical_qubit in enumerate(physical_to_logical):
+            if logical_qubit < qubit_register_size:
+                logical_to_physical[logical_qubit] = physical_qubit
+
+        if -1 in logical_to_physical:
+            error_msg = f"Incomplete mapping. Physical-to-logical: {physical_to_logical}"
+            raise ValueError(error_msg)
+
+        return Mapping(logical_to_physical)
+
+    def _extract_mapping_data(self, last_obs: Any) -> Any:
+        """Extract mapping from the observation dict only."""
+        if isinstance(last_obs, dict) and last_obs.get("mapping") is not None:
+            return last_obs["mapping"]
+        error_msg = "QGym env did not provide 'mapping' in observation."
+        raise RuntimeError(error_msg)

pyproject.toml

@@ -17,7 +17,7 @@ dependencies = [
     "libqasm==1.2.1",
     "networkx>=3.4.2",
     "numpy>=2.2.6",
-    "scipy>=1.15.3"
+    "scipy>=1.15.3",
 ]
 classifiers = [
     "Development Status :: 3 - Alpha",
@@ -48,6 +48,11 @@ export = [
     "pyqt5-qt5==5.15.2; sys_platform != 'darwin'",
     "quantify-scheduler==0.26.0; sys_platform != 'darwin'",
 ]
+qgym_mapper = [
+    "qgym==0.3.1",
+    "stable-baselines3==2.7.0",
+    "sb3-contrib==2.7.0",
+]
 docs = [
     "mike>=2.1.3",
     "mkdocs>=1.6.1,<2",
@@ -176,3 +181,5 @@ flake8-type-checking.exempt-modules = ["typing", "typing_extensions", "numpy", "
 [tool.mypy]
 ignore_missing_imports = true
 strict = true
+
+

tests/passes/mapper/test_qgym_mapper.py

@@ -0,0 +1,99 @@
+# Tests for the QGymMapper class
+import importlib.util
+import json
+from importlib.resources import files
+
+import networkx as nx
+import pytest
+
+from opensquirrel import CircuitBuilder
+from opensquirrel.circuit import Circuit
+from opensquirrel.passes.mapper import QGymMapper
+from opensquirrel.passes.mapper.mapping import Mapping
+
+if importlib.util.find_spec("qgym") is None:
+    pytest.skip("qgym not installed; skipping QGym mapper tests", allow_module_level=True)
+
+if importlib.util.find_spec("stable_baselines3") is None and importlib.util.find_spec("sb3_contrib") is None:
+    pytest.skip("stable-baselines3 and sb3_contrib not installed; skipping QGym mapper tests", allow_module_level=True)
+
+CONNECTIVITY_SCHEMES = json.loads(
+    (files("opensquirrel.passes.mapper") / "connectivities.json").read_text(encoding="utf-8")
+)
+AGENT1 = "opensquirrel/passes/mapper/TRPO_tuna5_2e5.zip"
+AGENT2 = "opensquirrel/passes/mapper/TRPO_starmon7_5e5.zip"
+AGENT_CLASS = "TRPO"
+
+
+@pytest.fixture
+def mapper1() -> QGymMapper:
+    agent_class = AGENT_CLASS
+    agent_path = AGENT1
+    connectivity = CONNECTIVITY_SCHEMES["tuna-5"]
+    connection_graph = nx.Graph()
+    connection_graph.add_edges_from(connectivity)
+    return QGymMapper(agent_class, agent_path, connection_graph)
+
+
+@pytest.fixture
+def mapper2() -> QGymMapper:
+    agent_class = AGENT_CLASS
+    agent_path = AGENT2
+    connectivity = CONNECTIVITY_SCHEMES["starmon-7"]
+    connection_graph = nx.Graph()
+    connection_graph.add_edges_from(connectivity)
+    return QGymMapper(agent_class, agent_path, connection_graph)
+
+
+@pytest.fixture
+def circuit1() -> Circuit:
+    builder = CircuitBuilder(5)
+    builder.H(0)
+    builder.CNOT(0, 1)
+    builder.H(2)
+    builder.CNOT(1, 2)
+    builder.CNOT(2, 4)
+    builder.CNOT(3, 4)
+    return builder.to_circuit()
+
+
+@pytest.fixture
+def circuit2() -> Circuit:
+    builder = CircuitBuilder(7)
+    builder.H(0)
+    builder.CNOT(0, 6)
+    builder.H(2)
+    builder.CNOT(1, 5)
+    builder.CNOT(2, 4)
+    builder.CNOT(3, 6)
+    builder.H(5)
+    builder.CNOT(0, 2)
+    builder.CNOT(1, 3)
+    builder.CNOT(4, 5)
+    builder.CNOT(5, 6)
+    return builder.to_circuit()
+
+
+@pytest.mark.parametrize(
+    "mapper, circuit, expected_mapping_length",  # noqa: PT006
+    [("mapper1", "circuit1", 5), ("mapper2", "circuit2", 7)],
+)
+def test_mapping(
+    mapper: QGymMapper, circuit: Circuit, expected_mapping_length: int, request: pytest.FixtureRequest
+) -> None:
+    circuit = request.getfixturevalue(circuit)  # type: ignore[arg-type]
+    mapper = request.getfixturevalue(mapper)  # type: ignore[arg-type]
+    mapping = mapper.map(circuit.ir, circuit.qubit_register_size)
+    assert isinstance(mapping, Mapping)
+    assert len(mapping) == expected_mapping_length
+
+
+def test_map_on_circuit(mapper1: QGymMapper, circuit1: Circuit) -> None:
+    initial_circuit = str(circuit1)
+    circuit1.map(mapper=mapper1)
+    assert str(circuit1) != initial_circuit
+
+
+def test_check_not_many_logical_as_physical_qubits(mapper1: QGymMapper, circuit2: Circuit) -> None:
+    with pytest.raises(ValueError, match=r"QGym requires equal logical and physical qubits: logical=7, physical=5"):
+        circuit2.map(mapper1)