Add Steane-type and teleportation-based QEC for the Steane code

python/quantum-pecos/src/pecos/qeclib/steane/steane_class.py

@@ -14,6 +14,10 @@
 from typing import TYPE_CHECKING
 from warnings import warn
 
+from pecos.qeclib.steane.decoders.lookup import (
+    FlagLookupQASMActiveCorrectionX,
+    FlagLookupQASMActiveCorrectionZ,
+)
 from pecos.qeclib.steane.gates_sq import paulis, sqrt_paulis
 from pecos.qeclib.steane.gates_sq.hadamards import H
 from pecos.qeclib.steane.gates_tq import transversal_tq
@@ -162,9 +166,7 @@ def prep_t_plus_state(
                 d=self.d,
                 a=self.a,
                 out=self.scratch,
-                reject=self.scratch[
-                    2
-                ],  # the first two bits of self.scratch are used by "out"
+                reject=self.scratch[2],  # the first two bits are used by "out"
                 flag_x=self.flag_x,
                 flag_z=self.flag_z,
                 flags=self.flags,
@@ -400,7 +402,7 @@ def mz(self, log: Bit | None = None):
         """Logical destructive measurement of the logical Z operator."""
         return self.m("Z", log=log)
 
-    def qec(self, flag_bit: Bit | None = None):
+    def qec(self, flag: Bit | None = None):
         block = ParallelFlagQECActiveCorrection(
             q=self.d,
             a=self.a,
@@ -416,6 +418,156 @@ def qec(self, flag_bit: Bit | None = None):
             pf_z=self.pf_z,
             scratch=self.scratch,
         )
-        if flag_bit is not None:
-            block.extend(If(self.flags != 0).Then(flag_bit.set(1)))
+        if flag is not None:
+            block.extend(If(self.flags != 0).Then(flag.set(1)))
+        return block
+
+    def qec_steane(
+        self,
+        aux: Steane,
+        reject_x: Bit | None = None,
+        reject_z: Bit | None = None,
+        flag_x: Bit | None = None,
+        flag_z: Bit | None = None,
+        rus_limit: int | None = None,
+    ) -> Block:
+        """Run a Steane-type error-correction cycle of this code."""
+        return Block(
+            self.qec_steane_z(
+                aux,
+                reject=reject_z,
+                flag=flag_z,
+                rus_limit=rus_limit,
+            ),
+            self.qec_steane_x(
+                aux,
+                reject=reject_x,
+                flag=flag_x,
+                rus_limit=rus_limit,
+            ),
+        )
+
+    def qec_steane_z(
+        self,
+        aux: Steane,
+        reject: Bit | None = None,
+        flag: Bit | None = None,
+        rus_limit: int | None = None,
+    ) -> Block:
+        """Run a Steane-type error-correction cycle for Z stabilizers (X errors)."""
+        warn("Using experimental feature: qec_steane_z", stacklevel=2)
+        flag = flag or self.scratch.elems[7]
+        return Block(
+            aux.px(reject=reject, rus_limit=rus_limit),
+            self.cx(aux),
+            aux.mz(),
+            If(aux.syn_meas != 0).Then(flag.set(1)),
+            self.last_raw_syn_z.set(0),
+            self.pf_x.set(0),
+            FlagLookupQASMActiveCorrectionZ(
+                self.d,
+                aux.syn_meas,
+                self.syndromes,
+                self.last_raw_syn_z,
+                self.pf_x,
+                flag,
+                aux.syn_meas,
+                self.scratch,
+            ),
+        )
+
+    def qec_steane_x(
+        self,
+        aux: Steane,
+        reject: Bit | None = None,
+        flag: Bit | None = None,
+        rus_limit: int | None = None,
+    ) -> Block:
+        """Run a Steane-type error-correction cycle for X stabilizers (Z errors)."""
+        warn("Using experimental feature: qec_steane_x", stacklevel=2)
+        flag = flag or self.scratch.elems[7]
+        return Block(
+            aux.pz(reject=reject, rus_limit=rus_limit),
+            aux.cx(self),
+            aux.mx(),
+            If(aux.syn_meas != 0).Then(flag.set(1)),
+            self.last_raw_syn_x.set(0),
+            self.pf_z.set(0),
+            FlagLookupQASMActiveCorrectionX(
+                self.d,
+                aux.syn_meas,
+                self.syndromes,
+                self.last_raw_syn_x,
+                self.pf_z,
+                flag,
+                aux.syn_meas,
+                self.scratch,
+            ),
+        )
+
+    def qec_tel(
+        self,
+        aux: Steane,
+        reject_x: Bit | None = None,
+        reject_z: Bit | None = None,
+        flag_x: Bit | None = None,
+        flag_z: Bit | None = None,
+        rus_limit: int | None = None,
+    ) -> Block:
+        """Run a teleportation-based error correction cycle."""
+        return Block(
+            self.qec_tel_x(aux, reject_x, flag_x, rus_limit),
+            self.qec_tel_z(aux, reject_z, flag_z, rus_limit),
+        )
+
+    def qec_tel_x(
+        self,
+        aux: Steane,
+        reject: Bit | None = None,
+        flag: Bit | None = None,
+        rus_limit: int | None = None,
+    ) -> Block:
+        """Run a teleportation-based error correction cycle for X errors."""
+        warn("Using experimental feature: qec_tel_x", stacklevel=2)
+        block = Block(
+            # teleport
+            aux.px(reject=reject, rus_limit=rus_limit),
+            aux.cx(self),
+            self.mz(),
+            If(self.log).Then(aux.x()),
+            Permute(self.d, aux.d),
+            # update syndromes and pauli frame
+            self.last_raw_syn_x.set(0),
+            self.last_raw_syn_z.set(0),
+            self.syn_z.set(self.syn_meas),
+            self.pf_x.set(0),
+        )
+        if flag is not None:
+            block.extend(If(self.syn_meas != 0).Then(flag.set(1)))
+        return block
+
+    def qec_tel_z(
+        self,
+        aux: Steane,
+        reject: Bit | None = None,
+        flag: Bit | None = None,
+        rus_limit: int | None = None,
+    ) -> Block:
+        """Run a teleportation-based error correction cycle for Z errors."""
+        warn("Using experimental feature: qec_tel_z", stacklevel=2)
+        block = Block(
+            # teleport
+            aux.pz(reject=reject, rus_limit=rus_limit),
+            self.cx(aux),
+            self.mx(),
+            If(self.log).Then(aux.z()),
+            Permute(self.d, aux.d),
+            # update syndromes and pauli frame
+            self.last_raw_syn_x.set(0),
+            self.last_raw_syn_z.set(0),
+            self.syn_x.set(self.syn_meas),
+            self.pf_z.set(0),
+        )
+        if flag is not None:
+            block.extend(If(self.syn_meas != 0).Then(flag.set(1)))
         return block