Implement Sp/Sn runtime

Author: david-pl

src/bloqade/pyqrack/squin/op.py

@@ -4,6 +4,8 @@
 from bloqade.pyqrack.base import PyQrackInterpreter
 
 from .runtime import (
+    SnRuntime,
+    SpRuntime,
     RotRuntime,
     KronRuntime,
     MultRuntime,
@@ -114,10 +116,10 @@ def projector(
         state = isinstance(stmt, op.stmts.P1)
         return (ProjectorRuntime(to_state=state),)
 
-    @interp.impl(op.stmts.Sn)
-    def sn(self, interp: PyQrackInterpreter, frame: interp.Frame, stmt: op.stmts.Sn):
-        raise NotImplementedError()
-
     @interp.impl(op.stmts.Sp)
     def sp(self, interp: PyQrackInterpreter, frame: interp.Frame, stmt: op.stmts.Sp):
-        raise NotImplementedError()
+        return (SpRuntime(),)
+
+    @interp.impl(op.stmts.Sn)
+    def sn(self, interp: PyQrackInterpreter, frame: interp.Frame, stmt: op.stmts.Sn):
+        return (SnRuntime(),)

src/bloqade/pyqrack/squin/runtime.py

@@ -141,18 +141,14 @@ def control_apply(self, *qubits: PyQrackQubit, adjoint: bool = False) -> None:
 
 
 @dataclass(frozen=True)
-class PhaseOpRuntime(OperatorRuntimeABC):
-    theta: float
-    global_: bool
-
-    def mat(self, adjoint: bool):
-        sign = (-1) ** (not adjoint)
-        phase = np.exp(sign * 1j * self.theta)
-        return [self.global_ * phase, 0, 0, phase]
+class MtrxOpRuntime(OperatorRuntimeABC):
+    def mat(self, adjoint: bool) -> list[complex]:
+        raise NotImplementedError("Override this method in the subclass!")
 
     def apply(self, *qubits: PyQrackQubit, adjoint: bool = False) -> None:
         target = qubits[-1]
-        target.sim_reg.mtrx(self.mat(adjoint=adjoint), target.addr)
+        m = self.mat(adjoint=adjoint)
+        target.sim_reg.mtrx(m, target.addr)
 
     def control_apply(self, *qubits: PyQrackQubit, adjoint: bool = False) -> None:
         target = qubits[-1]
@@ -163,6 +159,35 @@ def control_apply(self, *qubits: PyQrackQubit, adjoint: bool = False) -> None:
         target.sim_reg.mcmtrx(ctrls, m, target.addr)
 
 
+@dataclass(frozen=True)
+class SpRuntime(MtrxOpRuntime):
+    def mat(self, adjoint: bool) -> list[complex]:
+        if adjoint:
+            return [0, 0, 1, 0]
+        else:
+            return [0, 1, 0, 0]
+
+
+@dataclass(frozen=True)
+class SnRuntime(MtrxOpRuntime):
+    def mat(self, adjoint: bool) -> list[complex]:
+        if adjoint:
+            return [0, 1, 0, 0]
+        else:
+            return [0, 0, 1, 0]
+
+
+@dataclass(frozen=True)
+class PhaseOpRuntime(MtrxOpRuntime):
+    theta: float
+    global_: bool
+
+    def mat(self, adjoint: bool) -> list[complex]:
+        sign = (-1) ** (not adjoint)
+        phase = np.exp(sign * 1j * self.theta)
+        return [self.global_ * phase, 0, 0, phase]
+
+
 @dataclass(frozen=True)
 class RotRuntime(OperatorRuntimeABC):
     axis: OperatorRuntimeABC

test/pyqrack/test_squin.py

@@ -170,12 +170,42 @@ def main():
     assert result == [0]
 
 
+def test_sp():
+    @squin.kernel
+    def main():
+        q = squin.qubit.new(1)
+        sp = squin.op.spin_p()
+        squin.qubit.apply(sp, q)
+        return q
+
+    target = PyQrack(1)
+    result = target.run(main)
+    assert isinstance(result, ilist.IList)
+    assert isinstance(qubit := result[0], PyQrackQubit)
+
+    assert qubit.sim_reg.out_ket() == [0, 0]
+
+    @squin.kernel
+    def main2():
+        q = squin.qubit.new(1)
+        sn = squin.op.spin_n()
+        sp = squin.op.spin_p()
+        squin.qubit.apply(sn, q)
+        squin.qubit.apply(sp, q)
+        return squin.qubit.measure(q)
+
+    target = PyQrack(1)
+    result = target.run(main2)
+    assert result == [0]
+
+
 # TODO: remove
 # test_qubit()
 # test_x()
 # test_basic_ops("x")
 # test_cx()
 # test_mult()
 # test_kron()
-test_scale()
-test_phase()
+# test_scale()
+# test_phase()
+test_sp()