Properly check for atom loss in apply methods

Author: david-pl

src/bloqade/pyqrack/squin/qubit.py

@@ -27,11 +27,6 @@ def new(self, interp: PyQrackInterpreter, frame: interp.Frame, stmt: qubit.New):
     @interp.impl(qubit.Apply)
     def apply(self, interp: PyQrackInterpreter, frame: interp.Frame, stmt: qubit.Apply):
         qubits: ilist.IList[PyQrackQubit, Any] = frame.get(stmt.qubits)
-
-        for qbit in qubits:
-            if not qbit.is_active():
-                return ()
-
         operator: OperatorRuntimeABC = frame.get(stmt.operator)
         operator.apply(*qubits)
 

src/bloqade/pyqrack/squin/runtime.py

@@ -55,8 +55,14 @@ def apply(self, qubit: PyQrackQubit, adjoint: bool = False) -> None:
         getattr(qubit.sim_reg, method_name)(qubit.addr)
 
     def control_apply(self, *qubits: PyQrackQubit, adjoint: bool = False) -> None:
-        ctrls = [qbit.addr for qbit in qubits[:-1]]
         target = qubits[-1]
+        if not target.is_active():
+            return
+
+        ctrls = [qbit.addr for qbit in qubits[:-1] if qbit.is_active()]
+        if len(ctrls) == 0:
+            return
+
         method_name = self.get_method_name(adjoint=adjoint, control=True)
         getattr(target.sim_reg, method_name)(ctrls, target.addr)
 
@@ -77,13 +83,19 @@ def apply(self, *qubits: PyQrackQubit, adjoint: bool = False) -> None:
 class ProjectorRuntime(OperatorRuntimeABC):
     to_state: bool
 
-    def apply(self, *qubits: PyQrackQubit, adjoint: bool = False) -> None:
-        qubits[-1].sim_reg.force_m(qubits[-1].addr, self.to_state)
+    def apply(self, qubit: PyQrackQubit, adjoint: bool = False) -> None:
+        if not qubit.is_active():
+            return
+        qubit.sim_reg.force_m(qubit.addr, self.to_state)
 
     def control_apply(self, *qubits: PyQrackQubit, adjoint: bool = False) -> None:
-        m = [not self.to_state, 0, 0, self.to_state]
         target = qubits[-1]
+        if not target.is_active():
+            return
+
         ctrls = [qbit.addr for qbit in qubits[:-1]]
+
+        m = [not self.to_state, 0, 0, self.to_state]
         target.sim_reg.mcmtrx(ctrls, m, target.addr)
 
 
@@ -155,9 +167,11 @@ def mat(self, adjoint: bool):
             return [self.factor, 0, 0, self.factor]
 
     def apply(self, *qubits: PyQrackQubit, adjoint: bool = False) -> None:
-        self.op.apply(*qubits, adjoint=adjoint)
-
         target = qubits[-1]
+        if not target.is_active():
+            return
+
+        self.op.apply(*qubits, adjoint=adjoint)
 
         # NOTE: just factor * eye(2)
         m = self.mat(adjoint)
@@ -166,13 +180,16 @@ def apply(self, *qubits: PyQrackQubit, adjoint: bool = False) -> None:
         target.sim_reg.mtrx(m, target.addr)
 
     def control_apply(self, *qubits: PyQrackQubit, adjoint: bool = False) -> None:
-        self.op.control_apply(*qubits, adjoint=adjoint)
-
         target = qubits[-1]
-        ctrls = [qbit.addr for qbit in qubits[:-1]]
+        if not target.is_active():
+            return
 
-        m = self.mat(adjoint=adjoint)
+        ctrls = [qbit.addr for qbit in qubits[:-1] if qbit.is_active()]
+        if len(ctrls) == 0:
+            return
 
+        self.op.control_apply(*qubits, adjoint=adjoint)
+        m = self.mat(adjoint=adjoint)
         target.sim_reg.mcmtrx(ctrls, m, target.addr)
 
 
@@ -183,15 +200,22 @@ def mat(self, adjoint: bool) -> list[complex]:
 
     def apply(self, *qubits: PyQrackQubit, adjoint: bool = False) -> None:
         target = qubits[-1]
+        if not target.is_active():
+            return
+
         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]
-        ctrls = [qbit.addr for qbit in qubits[:-1]]
+        if not target.is_active():
+            return
 
-        m = self.mat(adjoint=adjoint)
+        ctrls = [qbit.addr for qbit in qubits[:-1] if qbit.is_active()]
+        if len(ctrls) == 0:
+            return
 
+        m = self.mat(adjoint=adjoint)
         target.sim_reg.mcmtrx(ctrls, m, target.addr)
 
 
@@ -251,19 +275,25 @@ def __post_init__(self):
         object.__setattr__(self, "pyqrack_axis", axis)
 
     def apply(self, *qubits: PyQrackQubit, adjoint: bool = False) -> None:
-        sign = (-1) ** adjoint
-        angle = sign * self.angle
         target = qubits[-1]
+        if not target.is_active():
+            return
 
+        sign = (-1) ** adjoint
+        angle = sign * self.angle
         target.sim_reg.r(self.pyqrack_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 target.is_active():
+            return
+
+        ctrls = [qbit.addr for qbit in qubits[:-1] if qbit.is_active()]
+        if len(ctrls) == 0:
+            return
 
+        sign = (-1) ** (not adjoint)
+        angle = sign * self.angle
         target.sim_reg.mcr(self.pyqrack_axis, angle, ctrls, target.addr)
 
 
@@ -293,12 +323,21 @@ def angles(self, adjoint: bool) -> tuple[float, float, float]:
 
     def apply(self, *qubits: PyQrackQubit, adjoint: bool = False) -> None:
         target = qubits[-1]
+        if not target.is_active():
+            return
+
         angles = self.angles(adjoint=adjoint)
         target.sim_reg.u(target.addr, *angles)
 
     def control_apply(self, *qubits: PyQrackQubit, adjoint: bool = False) -> None:
         target = qubits[-1]
-        ctrls = [qbit.addr for qbit in qubits[:-1]]
+        if not target.is_active():
+            return
+
+        ctrls = [qbit.addr for qbit in qubits[:-1] if qbit.is_active()]
+        if len(ctrls) == 0:
+            return
+
         angles = self.angles(adjoint=adjoint)
         target.sim_reg.mcu(ctrls, target.addr, *angles)