Implement simple rotations about X, Y, Z

Author: david-pl

src/bloqade/pyqrack/squin/runtime.py

@@ -195,9 +195,53 @@ def mat(self, adjoint: bool) -> list[complex]:
 
 @dataclass(frozen=True)
 class RotRuntime(OperatorRuntimeABC):
+    AXIS_MAP = {
+        "x": 1,
+        "y": 2,
+        "z": 3,
+    }
     axis: OperatorRuntimeABC
     angle: float
-    # TODO: how does this work?
+
+    def apply(self, *qubits: PyQrackQubit, adjoint: bool = False) -> None:
+        sign = (-1) ** adjoint
+        angle = sign * self.angle
+        target = qubits[-1]
+
+        if not isinstance(self.axis, OperatorRuntime):
+            raise RuntimeError(
+                f"Rotation only supported for Pauli operators! Got {self.axis}"
+            )
+
+        try:
+            axis = self.AXIS_MAP[self.axis.method_name]
+        except KeyError:
+            raise RuntimeError(
+                f"Rotation only supported for Pauli operators! Got {self.axis}"
+            )
+
+        target.sim_reg.r(axis, angle, target.addr)
+
+    def control_apply(self, *qubits: PyQrackQubit, adjoint: bool = False) -> None:
+        sign = (-1) ** (not adjoint)
+        angle = sign * self.angle
+
+        ctrls = [qbit.addr for qbit in qubits[:-1]]
+        target = qubits[-1]
+
+        if not isinstance(self.axis, OperatorRuntime):
+            raise RuntimeError(
+                f"Rotation only supported for Pauli operators! Got {self.axis}"
+            )
+
+        try:
+            axis = self.AXIS_MAP[self.axis.method_name]
+        except KeyError:
+            raise RuntimeError(
+                f"Rotation only supported for Pauli operators! Got {self.axis}"
+            )
+
+        target.sim_reg.mcr(axis, angle, ctrls, target.addr)
 
 
 @dataclass(frozen=True)

test/pyqrack/test_squin.py

@@ -213,6 +213,20 @@ def main():
     assert result == [1]
 
 
+def test_rot():
+    @squin.kernel
+    def main():
+        q = squin.qubit.new(1)
+        x = squin.op.x()
+        r = squin.op.rot(x, math.pi / 2)
+        squin.qubit.apply(r, q)
+        return squin.qubit.measure(q)
+
+    target = PyQrack(1)
+    result = target.run(main)
+    assert result == [1]
+
+
 # TODO: remove
 # test_qubit()
 # test_x()
@@ -224,3 +238,4 @@ def main():
 # test_phase()
 # test_sp()
 # test_adjoint()
+# test_rot()