Adds custom barrier gate with conversion to/from qasm

mitiq/interface/mitiq_qiskit/conversions.py

@@ -16,7 +16,10 @@
 from cirq.contrib.qasm_import import circuit_from_qasm
 import qiskit
 
-from mitiq.utils import _simplify_circuit_exponents
+from mitiq.utils import (
+    _simplify_circuit_exponents_and_remove_barriers,
+    Barrier,
+)
 
 
 QASMType = str
@@ -45,6 +48,65 @@ def _remove_qasm_barriers(qasm: QASMType) -> QASMType:
     return "".join(lines)
 
 
+def _extract_qasm_barriers(
+    qasm: QASMType,
+) -> Tuple[QASMType, List[Tuple[int, List[int]]]]:
+    """Returns a copy of the input QASM with all barriers removed and a list
+    of tuples where each tuple contains the line number and qubit indices of a
+    barrier.
+
+    Args:
+        qasm: QASM to extract barriers from.
+    """
+    # Split the QASM into lines
+    lines = qasm.split("\n")
+
+    barrier_info = []
+
+    for i, line in enumerate(lines):
+        match = re.match(r"^\s*barrier ((?:q\[\d+\],? ?)+);", line)
+        if match is not None:
+            qubits_str = match.group(1)
+            qubits = [
+                int(qubit_index)
+                for qubit_index in re.findall(r"q\[(\d+)\]", qubits_str)
+            ]
+            barrier_info.append((i, qubits))
+
+    # Remove the barrier lines
+    lines = [line for line in lines if not line.strip().startswith("barrier")]
+    qasm_without_barriers = "\n".join(lines)
+
+    return qasm_without_barriers, barrier_info
+
+
+def _add_qasm_barriers(qasm: str, barriers: List[Tuple[int, Barrier]]) -> str:
+    """Returns a copy of the input QASM with barriers added at the specified indices."""
+    # Split the QASM into lines
+    lines = qasm.split("\n")
+
+    # Reverse the barriers list to insert from the end
+    barriers = list(reversed(barriers))
+
+    # For each barrier, create a QASM barrier string and insert it into the lines
+    for index, barrier in barriers:
+        qubits = barrier.get_qubits()
+        if qubits is None:
+            continue
+        if not all(isinstance(qubit, cirq.LineQubit) for qubit in qubits):
+            raise TypeError(
+                "All qubits must be LineQubits for QASM conversion."
+            )
+
+        qubit_strs = [f"q[{qubit._comparison_key}]" for qubit in qubits]
+        barrier_str = f"barrier {', '.join(qubit_strs)};"
+        lines.insert(index, barrier_str)
+
+    # Join the lines back together
+    qasm_with_barriers = "\n".join(lines)
+    return qasm_with_barriers
+
+
 def _map_bit_index(
     bit_index: int, new_register_sizes: List[int]
 ) -> Tuple[int, int]:
@@ -249,17 +311,29 @@ def _transform_registers(
 
 
 def to_qasm(circuit: cirq.Circuit) -> QASMType:
-    """Returns a QASM string representing the input Mitiq circuit.
+    """Converts a Cirq circuit with custom barriers to QASM and preserves
+    the barrier positions.
 
     Args:
-        circuit: Mitiq circuit to convert to a QASM string.
+        circuit: The Cirq circuit to convert.
 
     Returns:
-        QASMType: QASM string equivalent to the input Mitiq circuit.
+        The QASM string with custom barriers.
     """
-    # Simplify exponents of gates. For example, H**-1 is simplified to H.
-    _simplify_circuit_exponents(circuit)
-    return circuit.to_qasm()
+    barrier_indices = _simplify_circuit_exponents_and_remove_barriers(
+        circuit, return_barriers=True
+    )
+    qasm_without_barriers = circuit.to_qasm()
+
+    # Only add barriers if there are any
+    if barrier_indices is not None:
+        qasm_with_barriers = _add_qasm_barriers(
+            qasm_without_barriers, barrier_indices
+        )
+    else:
+        qasm_with_barriers = qasm_without_barriers
+
+    return qasm_with_barriers
 
 
 def to_qiskit(circuit: cirq.Circuit) -> qiskit.QuantumCircuit:
@@ -289,13 +363,33 @@ def from_qiskit(circuit: qiskit.QuantumCircuit) -> cirq.Circuit:
 
 
 def from_qasm(qasm: QASMType) -> cirq.Circuit:
-    """Returns a Mitiq circuit equivalent to the input QASM string.
+    """Returns a cirq.Circuit equivalent to the input QASM string.
 
     Args:
-        qasm: QASM string to convert to a Mitiq circuit.
+        qasm: QASM string to convert to a cirq.Circuit.
 
     Returns:
-        Mitiq circuit representation equivalent to the input QASM string.
+        cirq.Circuit representation equivalent to the input QASM string.
     """
-    qasm = _remove_qasm_barriers(qasm)
-    return circuit_from_qasm(qasm)
+    # Remove barriers from QASM and get barrier info
+    qasm_without_barriers, barrier_info = _extract_qasm_barriers(qasm)
+
+    # Convert QASM to Cirq circuit
+    circuit = circuit_from_qasm(qasm_without_barriers)
+
+    # Add barriers back into circuit
+    for line_number, qubits in barrier_info:
+        # Get the moment for this line number
+        moment = circuit[line_number]
+
+        # Create a barrier gate for each qubit
+        for qubit in qubits:
+            barrier = Barrier().set_qubits([list(circuit.all_qubits())[qubit]])
+            moment = moment.with_operation(
+                barrier.on(list(circuit.all_qubits())[qubit])
+            )
+
+        # Replace the moment in the circuit
+        circuit[line_number] = moment
+
+    return circuit

mitiq/tests/test_utils.py

@@ -31,7 +31,7 @@
     _equal,
     _is_measurement,
     _simplify_gate_exponent,
-    _simplify_circuit_exponents,
+    _simplify_circuit_exponents_and_remove_barriers,
     _max_ent_state_circuit,
     _circuit_to_choi,
     _operation_to_choi,
@@ -266,7 +266,7 @@ def test_simplify_circuit_exponents_controlled_gate():
     )
     copy = circuit.copy()
 
-    _simplify_circuit_exponents(circuit)
+    _simplify_circuit_exponents_and_remove_barriers(circuit)
     assert _equal(circuit, copy)
 
 
@@ -291,7 +291,7 @@ def test_simplify_circuit_exponents():
     assert inverse_qasm != expected_qasm
 
     # Simplify the circuit
-    _simplify_circuit_exponents(inverse_circuit)
+    _simplify_circuit_exponents_and_remove_barriers(inverse_circuit)
 
     # Check inverse_circuit has the expected simplified representation
     simplified_repr = inverse_circuit.__repr__()
@@ -311,7 +311,7 @@ def test_simplify_circuit_exponents_with_non_self_inverse_gates():
     inverse_qasm = inverse_circuit._to_qasm_output().__str__()
 
     # Simplify the circuit (it should not change this circuit)
-    _simplify_circuit_exponents(inverse_circuit)
+    _simplify_circuit_exponents_and_remove_barriers(inverse_circuit)
 
     # Check inverse_circuit did not change
     simplified_repr = inverse_circuit.__repr__()

mitiq/utils.py

@@ -5,27 +5,51 @@
 
 """Utility functions."""
 from copy import deepcopy
-from typing import Any, Dict, List, Tuple
+from typing import Any, Dict, List, Tuple, Optional
 
 import numpy as np
 import numpy.typing as npt
 
 from cirq import (
     LineQubit,
     Circuit,
-    EigenGate,
-    Gate,
-    GateOperation,
-    Moment,
     CNOT,
     H,
     DensityMatrixSimulator,
     ops,
     OP_TREE,
+    Gate,
+    Qid,
+    CircuitDiagramInfoArgs,
+    EigenGate,
+    GateOperation,
+    Moment,
 )
 from cirq.ops.measurement_gate import MeasurementGate
 
 
+class Barrier(Gate):
+    def __init__(self) -> None:
+        super().__init__()
+        self._qubits: Optional[List[Qid]] = None
+
+    def num_qubits(self) -> int:
+        return 1
+
+    def _decompose_(self, qubits: List[Qid]) -> List[None]:
+        return []
+
+    def _circuit_diagram_info_(self, args: CircuitDiagramInfoArgs) -> str:
+        return "B"
+
+    def set_qubits(self, qubits: List[Qid]) -> "Barrier":
+        self._qubits = qubits
+        return self
+
+    def get_qubits(self) -> Optional[List[Qid]]:
+        return self._qubits
+
+
 def _simplify_gate_exponent(gate: EigenGate) -> EigenGate:
     """Returns the input gate with a simplified exponent if possible,
     otherwise the input gate is returned without any change.
@@ -43,18 +67,28 @@ def _simplify_gate_exponent(gate: EigenGate) -> EigenGate:
     return gate
 
 
-def _simplify_circuit_exponents(circuit: Circuit) -> None:
+def _simplify_circuit_exponents_and_remove_barriers(
+    circuit: Circuit, return_barriers: bool = False
+) -> Optional[List[Tuple[int, Barrier]]]:
     """Simplifies the gate exponents of the input circuit if possible,
-    mutating the input circuit.
+    mutating the input circuit and optionally returning barrier indices.
 
     Args:
         circuit: The circuit to simplify.
+        return_barriers: Whether to return barrier indices.
     """
+    barrier_indices = []
+
     # Iterate over moments
     for moment_idx, moment in enumerate(circuit):
         simplified_operations = []
+
         # Iterate over operations in moment
         for op in moment:
+            if isinstance(op.gate, Barrier):
+                if return_barriers:
+                    barrier_indices.append((moment_idx, op.gate))
+                continue
 
             if not isinstance(op, GateOperation):
                 simplified_operations.append(op)
@@ -67,9 +101,12 @@ def _simplify_circuit_exponents(circuit: Circuit) -> None:
 
             simplified_operation = op.with_gate(simplified_gate)
             simplified_operations.append(simplified_operation)
+
         # Mutate the input circuit
         circuit[moment_idx] = Moment(simplified_operations)
 
+    return barrier_indices if return_barriers else None
+
 
 def _is_measurement(op: ops.Operation) -> bool:
     """Returns true if the operation's gate is a measurement, else False.