Simplified implementation of gate modifiers (#64)

jcjaskula-aws · PhilReinhold · web-flow · commit 4bb47dd59e07 · 2023-09-25T09:09:42.000-07:00
* add gate modifier classes

* clean up

* fix isort

* implement to_ast for OQpyGateModifier

* add negctrl

* add number of controls to ctrl and negctrl

* test that only ints are passed to PhysicalQubits

* test modifier in gate context

* add gate modifier arguments

* add qubit argument to gate modifiers

* reorganize xfail tests

* add test for second method

* test overlapping qubit sets

* add overlapping set test for the first method

* clean code

* remove OQPyGateModifier

* clean tests

* change pow argument to exp

* change to singular

* check types

* update qubit array test

* encapsulate modifier building code in a method

* move function

---------

Co-authored-by: Phil Reinhold &lt;pcrein@amazon.com&gt;
diff --git a/oqpy/program.py b/oqpy/program.py
@@ -34,6 +34,7 @@
 from oqpy import classical_types, quantum_types
 from oqpy.base import (
     AstConvertible,
+    OQPyExpression,
     Var,
     expr_matches,
     map_to_ast,
@@ -495,19 +496,107 @@ def returns(self, expression: AstConvertible) -> Program:
         self._add_statement(ast.ReturnStatement(to_ast(self, expression)))
         return self
 
+    def _create_modifiers_ast(
+        self,
+        control: quantum_types.Qubit | Iterable[quantum_types.Qubit] | None,
+        neg_control: quantum_types.Qubit | Iterable[quantum_types.Qubit] | None,
+        inv: bool,
+        exp: AstConvertible,
+    ) -> tuple[list[ast.QuantumGateModifier], list[AstConvertible]]:
+        """Create the AST for the gate modifiers."""
+        used_qubits: list[AstConvertible] = []
+        modifiers: list[ast.QuantumGateModifier] = []
+
+        control = control if control is not None else []
+        control = {control} if isinstance(control, quantum_types.Qubit) else set(control)
+        if control:
+            modifiers.append(
+                ast.QuantumGateModifier(
+                    modifier=ast.GateModifierName.ctrl,
+                    argument=to_ast(self, len(control)) if len(control) > 1 else None,
+                )
+            )
+            used_qubits.extend(sorted(control))
+
+        neg_control = neg_control if neg_control is not None else []
+        neg_control = (
+            {neg_control} if isinstance(neg_control, quantum_types.Qubit) else set(neg_control)
+        )
+        if neg_control:
+            modifiers.append(
+                ast.QuantumGateModifier(
+                    modifier=ast.GateModifierName.negctrl,
+                    argument=to_ast(self, len(neg_control)) if len(neg_control) > 1 else None,
+                )
+            )
+            for qubit in sorted(neg_control):
+                if qubit in used_qubits:
+                    raise ValueError(f"Qubit {qubit} has already been defined as a control qubit.")
+                else:
+                    used_qubits.append(qubit)
+
+        if inv:
+            modifiers.append(
+                ast.QuantumGateModifier(
+                    modifier=ast.GateModifierName.inv,
+                )
+            )
+
+        if isinstance(exp, OQPyExpression) or (isinstance(exp, float) and exp != 1.0):
+            modifiers.append(
+                ast.QuantumGateModifier(
+                    modifier=ast.GateModifierName.pow, argument=to_ast(self, exp)
+                )
+            )
+        return modifiers, used_qubits
+
     def gate(
-        self, qubits: AstConvertible | Iterable[AstConvertible], name: str, *args: Any
+        self,
+        qubits: AstConvertible | Iterable[AstConvertible],
+        name: str,
+        *args: Any,
+        control: quantum_types.Qubit | Iterable[quantum_types.Qubit] | None = None,
+        neg_control: quantum_types.Qubit | Iterable[quantum_types.Qubit] | None = None,
+        inv: bool = False,
+        exp: AstConvertible = 1,
     ) -> Program:
-        """Apply a gate to a qubit or set of qubits."""
+        """Apply a gate with its modifiers to a qubit or set of qubits.
+
+        Args:
+            qubits (AstConvertible | Iterable[AstConvertible]): The qubit or list of qubits
+                to which the gate will be applied
+            name (str): The gate name
+            *args (Any): A list of parameters passed to the gate
+            control (quantum_types.Qubit | Iterable[quantum_types.Qubit] | None): The list
+                of control qubits (default: None)
+            neg_control: (quantum_types.Qubit | Iterable[quantum_types.Qubit] | None): The list
+                of negative control qubits (default: None)
+            inv (bool): Flag to use the inverse gate (default: False)
+            exp (AstConvertible): The exponent used with `pow` gate modifier
+
+        Returns:
+            Program: The OQpy program to which the gate is added
+        """
+        modifiers, used_qubits = self._create_modifiers_ast(control, neg_control, inv, exp)
+
         if isinstance(qubits, (quantum_types.Qubit, quantum_types.IndexedQubitArray)):
             qubits = [qubits]
         assert isinstance(qubits, Iterable)
+
+        for qubit in qubits:
+            if qubit in used_qubits:
+                raise ValueError(
+                    f"Qubit {qubit} has already been defined as a control qubit or a negative control qubit."
+                )
+            else:
+                used_qubits.append(qubit)
+
         self._add_statement(
             ast.QuantumGate(
-                [],
+                modifiers,
                 ast.Identifier(name),
                 map_to_ast(self, args),
-                map_to_ast(self, qubits),
+                map_to_ast(self, used_qubits),
             )
         )
         return self
diff --git a/oqpy/quantum_types.py b/oqpy/quantum_types.py
@@ -48,6 +48,15 @@ def __init__(
         self.size = size
         self.annotations = annotations
 
+    def __hash__(self) -> int:
+        return hash(self.name)
+
+    def __eq__(self, other: object) -> bool:
+        return isinstance(other, Qubit) and self.name == other.name
+
+    def __lt__(self, other: Qubit) -> bool:
+        return self.name < other.name
+
     def to_ast(self, prog: Program) -> ast.Expression:
         """Converts the OQpy variable into an ast node."""
         prog._add_var(self)
@@ -77,6 +86,7 @@ class PhysicalQubits:
     """
 
     def __class_getitem__(cls, item: int) -> Qubit:
+        assert isinstance(item, int)
         return Qubit(f"${item}", needs_declaration=False)
 
 
diff --git a/tests/test_directives.py b/tests/test_directives.py
@@ -496,6 +496,8 @@ def test_add_incomptible_type():
 def test_measure_reset_pragma():
     prog = Program()
     q = PhysicalQubits[0]
+    with pytest.raises(AssertionError):
+        reg = PhysicalQubits[0:1]
     c = BitVar(name="c")
     prog.reset(q)
     prog.pragma("CLASSIFIER linear")
@@ -2246,6 +2248,58 @@ def g() -> int[32] {
     ).strip()
 
     assert prog.to_qasm() == expected
+    _check_respects_type_hints(prog)
+
+
+def test_gate_modifiers():
+    prog = oqpy.Program()
+    qreg = [oqpy.PhysicalQubits[i] for i in range(0, 8)]
+    six_qubit_reg = [qreg[i] for i in [1, 4, 7, 2, 3, 5]]
+
+    prog.gate(qreg[2], "t", control=qreg[1])
+    prog.gate(qreg[2], "x", neg_control=qreg[1])
+    prog.gate(qreg[3], "rz", inv=True)
+    prog.gate(qreg[2], "t", exp=0.5)
+    prog.gate(qreg[0], "x", inv=True, exp=oqpy.IntVar(5, "i") / 2)
+
+    prog.gate(
+        six_qubit_reg[-1],
+        "x",
+        control=six_qubit_reg[0:2],
+        neg_control=six_qubit_reg[2:5],
+        inv=True,
+        exp=1 / 2,
+    )
+
+    prog.gate(
+        qreg[6],
+        "rz1",
+        control=[qreg[2], qreg[4], qreg[5]],
+        neg_control=[qreg[0], qreg[1], qreg[0]],
+    )
+
+    with pytest.raises(ValueError):
+        prog.gate(qreg[2], "t", control=qreg[2])
+    with pytest.raises(ValueError):
+        prog.gate(qreg[2], "x", neg_control=qreg[2])
+    with pytest.raises(ValueError):
+        prog.gate(qreg[1], "x", control=qreg[2], neg_control=qreg[2])
+
+    expected = textwrap.dedent(
+        """
+        OPENQASM 3.0;
+        int[32] i = 5;
+        ctrl @ t $1, $2;
+        negctrl @ x $1, $2;
+        inv @ rz $3;
+        pow(0.5) @ t $2;
+        inv @ pow(i / 2) @ x $0;
+        ctrl(2) @ negctrl(3) @ inv @ pow(0.5) @ x $1, $4, $2, $3, $7, $5;
+        ctrl(3) @ negctrl(2) @ rz1 $2, $4, $5, $0, $1, $6;
+        """
+    ).strip()
+    assert prog.to_qasm() == expected
+    _check_respects_type_hints(prog)
 
 
 def test_invalid_gates():
@@ -2266,6 +2320,7 @@ def test_invalid_gates():
 def test_gate_declarations():
     prog = oqpy.Program()
     q = oqpy.Qubit("q", needs_declaration=False)
+    r = oqpy.Qubit("r", needs_declaration=False)
     with oqpy.gate(
         prog,
         q,
@@ -2280,6 +2335,10 @@ def test_gate_declarations():
         prog.gate(q, "u", theta, 0, 0)
     with oqpy.gate(prog, q, "t"):
         prog.gate(q, "rz", oqpy.pi / 4)
+    with oqpy.gate(prog, [q, r], "cnot"):
+        prog.gate(r, "x", control=q)
+    with oqpy.gate(prog, [q, r], "ncphaseshift", [oqpy.AngleVar(name="theta")]) as theta:
+        prog.gate(r, "phase", theta, neg_control=[q])
 
     prog.gate(oqpy.PhysicalQubits[1], "t")
     prog.gate(oqpy.PhysicalQubits[2], "t")
@@ -2296,6 +2355,12 @@ def test_gate_declarations():
         gate t q {
             rz(pi / 4) q;
         }
+        gate cnot q, r {
+            ctrl @ x q, r;
+        }
+        gate ncphaseshift(theta) q, r {
+            negctrl @ phase(theta) q, r;
+        }
         gate rz(theta) q {
             u(theta, 0, 0) q;
         }
@@ -2305,6 +2370,7 @@ def test_gate_declarations():
     ).strip()
 
     assert prog.to_qasm() == expected
+    _check_respects_type_hints(prog)
 
 
 def test_include():
@@ -2321,6 +2387,7 @@ def test_include():
     ).strip()
 
     assert prog.to_qasm() == expected
+    _check_respects_type_hints(prog)
 
 
 def test_qubit_array():
