Add oqpy.gate for gate definitions

oqpy/program.py

@@ -99,11 +99,13 @@ def __init__(self, version: Optional[str] = "3.0", simplify_constants: bool = Tr
             tuple[tuple[str, ...], str, tuple[str, ...]], ast.CalibrationDefinition
         ] = {}
         self.subroutines: dict[str, ast.SubroutineDefinition] = {}
+        self.gates: dict[str, ast.QuantumGateDefinition] = {}
         self.externs: dict[str, ast.ExternDeclaration] = {}
         self.declared_vars: dict[str, Var] = {}
         self.undeclared_vars: dict[str, Var] = {}
         self.simplify_constants = simplify_constants
         self.declared_subroutines: set[str] = set()
+        self.declared_gates: set[str] = set()
 
         if version is None or (
             len(version.split(".")) in [1, 2]
@@ -126,6 +128,8 @@ def __iadd__(self, other: Program) -> Program:
             self._add_subroutine(
                 name, subroutine_stmt, needs_declaration=name not in other.declared_subroutines
             )
+        for name, gate_stmt in other.gates.items():
+            self._add_gate(name, gate_stmt, needs_declaration=name not in other.declared_gates)
         self.externs.update(other.externs)
         for var in other.declared_vars.values():
             self._mark_var_declared(var)
@@ -221,6 +225,17 @@ def _add_subroutine(
         if not needs_declaration:
             self.declared_subroutines.add(name)
 
+    def _add_gate(
+        self, name: str, stmt: ast.QuantumGateDefinition, needs_declaration: bool = True
+    ) -> None:
+        """Register a gate definition which has been used.
+
+        Gate definitions are added to the top of the program upon conversion to ast.
+        """
+        self.gates[name] = stmt
+        if not needs_declaration:
+            self.declared_gates.add(name)
+
     def _add_defcal(
         self,
         qubit_names: list[str],
@@ -293,11 +308,19 @@ def to_ast(
         statements = []
         if include_externs:
             statements += mutating_prog._make_externs_statements(encal_declarations)
-        statements += [
-            mutating_prog.subroutines[subroutine_name]
-            for subroutine_name in mutating_prog.subroutines
-            if subroutine_name not in mutating_prog.declared_subroutines
-        ] + mutating_prog._state.body
+        statements += (
+            [
+                mutating_prog.subroutines[subroutine_name]
+                for subroutine_name in mutating_prog.subroutines
+                if subroutine_name not in mutating_prog.declared_subroutines
+            ]
+            + [
+                mutating_prog.gates[gate_name]
+                for gate_name in mutating_prog.gates
+                if gate_name not in mutating_prog.declared_gates
+            ]
+            + mutating_prog._state.body
+        )
         if encal:
             statements = [ast.CalibrationStatement(statements)]
         if encal_declarations:

oqpy/quantum_types.py

@@ -24,13 +24,13 @@
 from openpulse.printer import dumps
 
 from oqpy.base import AstConvertible, Var, make_annotations, to_ast
-from oqpy.classical_types import _ClassicalVar
+from oqpy.classical_types import AngleVar, _ClassicalVar
 
 if TYPE_CHECKING:
     from oqpy.program import Program
 
 
-__all__ = ["Qubit", "QubitArray", "defcal", "PhysicalQubits", "Cal"]
+__all__ = ["Qubit", "QubitArray", "defcal", "gate", "PhysicalQubits", "Cal"]
 
 
 class Qubit(Var):
@@ -73,6 +73,56 @@ class QubitArray:
     """Represents an array of qubits."""
 
 
+@contextlib.contextmanager
+def gate(
+    program: Program,
+    qubits: Union[Qubit, list[Qubit]],
+    name: str,
+    arguments: Optional[list[AstConvertible]] = None,
+    declare_here: bool = False,
+) -> Union[Iterator[None], Iterator[list[AngleVar]], Iterator[AngleVar]]:
+    """Context manager for creating a gate.
+
+    .. code-block:: python
+
+        with gate(program, q1, "HRzH", [AngleVar(name="theta")]) as theta:
+            program.gate(q1, "H")
+            program.gate(q1, "Rz", theta)
+            program.gate(q1, "H")
+    """
+    if isinstance(qubits, Qubit):
+        qubits = [qubits]
+
+    arguments_ast = []
+    variables = []
+    if arguments is not None:
+        for arg in arguments:
+            if not isinstance(arg, AngleVar):
+                raise ValueError(arg, "Gates only support args of type AngleVar.")
+            arguments_ast.append(ast.Identifier(name=arg.name))
+            arg._needs_declaration = False
+            variables.append(arg)
+
+    program._push()
+    if len(variables) > 1:
+        yield variables
+    elif len(variables) == 1:
+        yield variables[0]
+    else:
+        yield None
+    state = program._pop()
+
+    stmt = ast.QuantumGateDefinition(
+        name=ast.Identifier(name),
+        arguments=arguments_ast,
+        qubits=[ast.Identifier(q.name) for q in qubits],
+        body=state.body,
+    )
+    if declare_here:
+        program._add_statement(stmt)
+    program._add_gate(name, stmt, needs_declaration=not declare_here)
+
+
 @contextlib.contextmanager
 def defcal(
     program: Program,

tests/test_directives.py

@@ -2195,3 +2195,57 @@ def g() -> int[32] {
     ).strip()
 
     assert prog.to_qasm() == expected
+
+
+def test_invalid_gates():
+    # missing qubits argument
+    prog = oqpy.Program()
+    with pytest.raises(TypeError):
+        with oqpy.gate(prog, None, "u"):
+            pass
+
+    # invalid argument type
+    prog = oqpy.Program()
+    with pytest.raises(ValueError):
+        q = oqpy.Qubit("q", needs_declaration=False)
+        with oqpy.gate(prog, q, "u", [oqpy.FloatVar(name="a")]) as a:
+            pass
+
+
+def test_gate_declarations():
+    prog = oqpy.Program()
+    q = oqpy.Qubit("q", needs_declaration=False)
+    with oqpy.gate(prog, q, "u", [oqpy.AngleVar(name="alpha"), oqpy.AngleVar(name="beta"), oqpy.AngleVar(name="gamma")]) as (alpha, beta, gamma):
+        prog.gate(q, "a", alpha)
+        prog.gate(q, "b", beta)
+        prog.gate(q, "c", gamma)
+        prog.gate(q, "d")
+    with oqpy.gate(prog, q, "rz", [oqpy.AngleVar(name="theta")], declare_here=True) as theta:
+        prog.gate(q, "u", theta, 0, 0)
+    with oqpy.gate(prog, q, "t"):
+        prog.gate(q, "rz", oqpy.pi / 4)
+
+    prog.gate(oqpy.PhysicalQubits[1], "t")
+    prog.gate(oqpy.PhysicalQubits[2], "t")
+
+    expected = textwrap.dedent(
+        """
+        OPENQASM 3.0;
+        gate u(alpha, beta, gamma) q {
+            a(alpha) q;
+            b(beta) q;
+            c(gamma) q;
+            d q;
+        }
+        gate t q {
+            rz(pi / 4) q;
+        }
+        gate rz(theta) q {
+            u(theta, 0, 0) q;
+        }
+        t $1;
+        t $2;
+        """
+    ).strip()
+
+    assert prog.to_qasm() == expected