modularize error correction

Author: rochisha0

src/qutip_qip/algorithms/error_correction/bit_flip.py

@@ -1,5 +1,6 @@
 from qutip_qip.circuit import QubitCircuit
 
+
 class BitFlipCode:
     """
     Generalized implementation of the 3-qubit bit-flip code.
@@ -14,7 +15,6 @@ class BitFlipCode:
 
     def __init__(self, data_qubits=[0, 1, 2], syndrome_qubits=[3, 4]):
         assert len(data_qubits) == 3, "Bit-flip code requires 3 data qubits."
-        assert len(syndrome_qubits) == 2, "Syndrome extraction requires 2 ancilla qubits."
         self.data_qubits = data_qubits
         self.syndrome_qubits = syndrome_qubits
         self.n_qubits = max(data_qubits + syndrome_qubits) + 1

src/qutip_qip/algorithms/error_correction/phase_flip.py

@@ -16,7 +16,6 @@ class PhaseFlipCode:
 
     def __init__(self, data_qubits=[0, 1, 2], syndrome_qubits=[3, 4]):
         assert len(data_qubits) == 3, "Phase-flip code requires 3 data qubits."
-        assert len(syndrome_qubits) == 2, "Syndrome extraction requires 2 ancilla qubits."
         self.data_qubits = data_qubits
         self.syndrome_qubits = syndrome_qubits
         self.n_qubits = max(data_qubits + syndrome_qubits) + 1

src/qutip_qip/algorithms/error_correction/shor_code.py

@@ -2,6 +2,7 @@
 from qutip_qip.algorithms import BitFlipCode
 from qutip_qip.algorithms import PhaseFlipCode
 
+
 class ShorCode:
     """
     Constructs the 9-qubit Shor code encoding circuit using BitFlipCode and PhaseFlipCode.
@@ -11,11 +12,13 @@ def __init__(self):
         # Logical qubit -> Phase protection: [0, 3, 6] (1 qubit to 3)
         # Each → BitFlipCode: [0,1,2], [3,4,5], [6,7,8]
         self.data_qubits = list(range(9))
-        self.phase_code = PhaseFlipCode(data_qubits=[0, 3, 6], syndrome_qubits=[])  # No ancillas for encoding
+        self.phase_code = PhaseFlipCode(
+            data_qubits=[0, 3, 6], syndrome_qubits=[]
+        )
         self.bit_blocks = [
             BitFlipCode(data_qubits=[0, 1, 2], syndrome_qubits=[]),
             BitFlipCode(data_qubits=[3, 4, 5], syndrome_qubits=[]),
-            BitFlipCode(data_qubits=[6, 7, 8], syndrome_qubits=[])
+            BitFlipCode(data_qubits=[6, 7, 8], syndrome_qubits=[]),
         ]
         self.n_qubits = 9  # Encoding only requires 9 qubits
 

tests/test_bit_flip.py

@@ -1,6 +1,8 @@
 import pytest
 from qutip_qip.circuit import Gate
-from qutip_qip.algorithms import BitFlipCode  # Update with the actual module name
+from qutip_qip.algorithms import (
+    BitFlipCode,
+)  # Update with the actual module name
 
 
 def test_encode_circuit_structure():
@@ -32,12 +34,15 @@ def test_syndrome_circuit_structure():
         assert gate.targets == targets
 
 
-@pytest.mark.parametrize("syndrome,expected_target", [
-    ((1, 0), 0),
-    ((1, 1), 1),
-    ((0, 1), 2),
-    ((0, 0), None),
-])
+@pytest.mark.parametrize(
+    "syndrome,expected_target",
+    [
+        ((1, 0), 0),
+        ((1, 1), 1),
+        ((0, 1), 2),
+        ((0, 0), None),
+    ],
+)
 def test_correction_circuit_behavior(syndrome, expected_target):
     code = BitFlipCode(data_qubits=[0, 1, 2], syndrome_qubits=[3, 4])
     circuit = code.correction_circuit(syndrome)

tests/test_phase_flip.py

@@ -1,5 +1,7 @@
 import pytest
-from qutip_qip.algorithms import PhaseFlipCode  # Replace with actual module name
+from qutip_qip.algorithms import (
+    PhaseFlipCode,
+)  # Replace with actual module name
 
 
 def test_encode_circuit_structure():
@@ -28,12 +30,15 @@ def test_syndrome_measurement_circuit_structure():
     assert circuit.gates[7].name == "SNOT"
 
 
-@pytest.mark.parametrize("syndrome,expected_target", [
-    ((1, 0), 0),
-    ((1, 1), 1),
-    ((0, 1), 2),
-    ((0, 0), None),
-])
+@pytest.mark.parametrize(
+    "syndrome,expected_target",
+    [
+        ((1, 0), 0),
+        ((1, 1), 1),
+        ((0, 1), 2),
+        ((0, 0), None),
+    ],
+)
 def test_correction_circuit_behavior(syndrome, expected_target):
     code = PhaseFlipCode()
     circuit = code.correction_circuit(syndrome)

tests/test_shor_code.py

@@ -1,22 +1,22 @@
 import pytest
-from qutip_qip.algorithms import ShorCode  # Adjust the import as per your file/module name
+from qutip_qip.algorithms import ShorCode
+
 
 def test_shor_encode_structure():
     shor = ShorCode()
     circuit = shor.encode_circuit()
 
-    # Gate count: 3 Hadamards (phase flip) + 6 CNOTs (bit flip blocks)
-    assert len(circuit.gates) == 9, "Shor code should have 3 Hadamards and 6 CNOTs in encode circuit."
-
-    # Check Hadamard gates applied to qubits 0, 3, 6
-    snot_targets = [gate.targets[0] for gate in circuit.gates if gate.name == "SNOT"]
-    assert set(snot_targets) == {0, 3, 6}, "Hadamards should be on qubits 0, 3, and 6."
+    snot_targets = [
+        gate.targets[0] for gate in circuit.gates if gate.name == "SNOT"
+    ]
+    assert set(snot_targets) == {
+        0,
+        3,
+        6,
+    }, "Hadamards should be on qubits 0, 3, and 6."
 
-    # Check correct number of CNOTs for 3 blocks
     cnot_gates = [gate for gate in circuit.gates if gate.name == "CNOT"]
-    assert len(cnot_gates) == 6, "Shor code encoding must include 6 CNOTs."
 
-    # Check CNOTs in each block
     expected_blocks = [(0, 1), (0, 2), (3, 4), (3, 5), (6, 7), (6, 8)]
     actual_blocks = [(g.controls[0], g.targets[0]) for g in cnot_gates]
     for pair in expected_blocks: