fixing/simplifying tests

Author: weinbe58

test/analysis/fidelity/test_fidelity.py

@@ -8,64 +8,41 @@
 from bloqade.qasm2.passes.noise import NoisePass
 
 
-class NoiseTestModel(native.MoveNoiseModelABC):
-    def parallel_cz_errors(self, ctrls, qargs, rest):
-        return {(0.01, 0.01, 0.01, 0.01): ctrls + qargs + rest}
-
+def test_atom_loss_analysis():
 
-def test_basic_noise():
+    p_loss = 0.01
 
     @qasm2.extended
     def main():
         q = qasm2.qreg(2)
-        qasm2.x(q[0])
+        native.atom_loss_channel([q[0]], prob=p_loss)
         return q
 
-    main.print()
-
     fid_analysis = FidelityAnalysis(main.dialects)
     fid_analysis.run_analysis(main, no_raise=False)
 
     assert fid_analysis.gate_fidelity == fid_analysis._current_gate_fidelity == 1
-
-    px = 0.01
-    py = 0.01
-    pz = 0.01
-    p_loss = 0.01
-
-    noise_params = native.GateNoiseParams(
-        global_loss_prob=p_loss,
-        global_px=px,
-        global_py=py,
-        global_pz=pz,
-        local_px=0.002,
-    )
-
-    model = NoiseTestModel()
-
-    NoisePass(main.dialects, noise_model=model, gate_noise_params=noise_params)(main)
-
-    main.print()
-
-    fid_analysis = FidelityAnalysis(main.dialects)
-    fid_analysis.run_analysis(main, no_raise=False)
-
-    p_noise = noise_params.local_px + noise_params.local_py + noise_params.local_pz
-    assert (
-        fid_analysis.gate_fidelity
-        == fid_analysis._current_gate_fidelity
-        == (1 - p_noise)
-    )
-
-    assert 0.9 < fid_analysis.atom_survival_probability[0] < 1
-    assert fid_analysis.atom_survival_probability[0] == 1 - noise_params.local_loss_prob
+    assert fid_analysis.atom_survival_probability[0] == 1 - p_loss
     assert fid_analysis.atom_survival_probability[1] == 1
 
 
-def test_c_noise():
+def test_cz_noise():
+    p_ch = 0.01 / 3.0
+
     @qasm2.extended
     def main():
         q = qasm2.qreg(2)
+        native.cz_pauli_channel(
+            [q[0]],
+            [q[1]],
+            px_ctrl=p_ch,
+            py_ctrl=p_ch,
+            pz_ctrl=p_ch,
+            px_qarg=p_ch,
+            py_qarg=p_ch,
+            pz_qarg=p_ch,
+            paired=True,
+        )
         qasm2.cz(q[0], q[1])
         return q
 
@@ -74,43 +51,35 @@ def main():
     fid_analysis = FidelityAnalysis(main.dialects)
     fid_analysis.run_analysis(main, no_raise=False)
 
-    assert fid_analysis.gate_fidelity == fid_analysis._current_gate_fidelity == 1
+    expected_fidelity = (1 - 3 * p_ch) ** 2
 
-    px = 0.01
-    py = 0.01
-    pz = 0.01
-    p_loss = 0.01
+    assert fid_analysis.gate_fidelity == fid_analysis._current_gate_fidelity
+    assert math.isclose(fid_analysis.gate_fidelity, expected_fidelity)
 
-    noise_params = native.GateNoiseParams(
-        global_loss_prob=p_loss,
-        global_px=px,
-        global_py=py,
-        global_pz=pz,
-        local_px=0.002,
-    )
 
-    model = NoiseTestModel()
+def test_single_qubit_noise():
+    p_ch = 0.01 / 3.0
 
-    NoisePass(main.dialects, noise_model=model, gate_noise_params=noise_params)(main)
+    @qasm2.extended
+    def main():
+        q = qasm2.qreg(2)
+        native.pauli_channel([q[0]], px=p_ch, py=p_ch, pz=p_ch)
+        return q
 
     main.print()
 
     fid_analysis = FidelityAnalysis(main.dialects)
     fid_analysis.run_analysis(main, no_raise=False)
 
-    # two cz channels (**2 for each one since we look at both control & target)
-    fid_cz = (1 - 3 * noise_params.cz_paired_gate_px) ** 4
-
-    # one pauli channel
-    fid_cz *= 1 - noise_params.global_px * 3
+    expected_fidelity = 1 - 3 * p_ch
 
     assert fid_analysis.gate_fidelity == fid_analysis._current_gate_fidelity
-    assert math.isclose(fid_cz, fid_analysis.gate_fidelity, abs_tol=1e-14)
+    assert math.isclose(fid_analysis.gate_fidelity, expected_fidelity)
 
-    assert 0.9 < fid_analysis.atom_survival_probability[0] < 1
-    assert fid_analysis.atom_survival_probability[0] == (
-        1 - noise_params.cz_gate_loss_prob
-    ) * (1 - p_loss)
+
+class NoiseTestModel(native.MoveNoiseModelABC):
+    def parallel_cz_errors(self, ctrls, qargs, rest):
+        return {(0.01, 0.01, 0.01, 0.01): ctrls + qargs + rest}
 
 
 @pytest.mark.xfail
@@ -145,23 +114,19 @@ def main_if():
     pz = 0.01
     p_loss = 0.01
 
-    noise_params = native.GateNoiseParams(
+    model = NoiseTestModel(
         global_loss_prob=p_loss,
         global_px=px,
         global_py=py,
         global_pz=pz,
         local_px=0.002,
     )
-
-    model = NoiseTestModel()
-    NoisePass(main.dialects, noise_model=model, gate_noise_params=noise_params)(main)
+    NoisePass(main.dialects, noise_model=model)(main)
     fid_analysis = FidelityAnalysis(main.dialects)
     fid_analysis.run_analysis(main, no_raise=False)
 
     model = NoiseTestModel()
-    NoisePass(main_if.dialects, noise_model=model, gate_noise_params=noise_params)(
-        main_if
-    )
+    NoisePass(main_if.dialects, noise_model=model)(main_if)
     fid_if_analysis = FidelityAnalysis(main_if.dialects)
     fid_if_analysis.run_analysis(main_if, no_raise=False)
 
@@ -211,24 +176,20 @@ def main_for():
     pz = 0.01
     p_loss = 0.01
 
-    noise_params = native.GateNoiseParams(
+    model = NoiseTestModel(
         global_loss_prob=p_loss,
         global_px=px,
         global_py=py,
         global_pz=pz,
         local_px=0.002,
         local_loss_prob=0.03,
     )
-
-    model = NoiseTestModel()
-    NoisePass(main.dialects, noise_model=model, gate_noise_params=noise_params)(main)
+    NoisePass(main.dialects, noise_model=model)(main)
     fid_analysis = FidelityAnalysis(main.dialects)
     fid_analysis.run_analysis(main, no_raise=False)
 
     model = NoiseTestModel()
-    NoisePass(main_for.dialects, noise_model=model, gate_noise_params=noise_params)(
-        main_for
-    )
+    NoisePass(main_for.dialects, noise_model=model)(main_for)
 
     main_for.print()
 

test/qasm2/passes/test_heuristic_noise.py

@@ -22,15 +22,14 @@ def test_single_qubit_noise():
     pz = 0.01
     p_loss = 0.01
 
-    noise_params = native.GateNoiseParams(
+    model = NoiseTestModel(
         local_px=px, local_py=py, local_pz=pz, local_loss_prob=p_loss
     )
-    model = NoiseTestModel()
 
     test_qubit = ir.TestValue(type=qasm2.QubitType)
     address_analysis = {test_qubit: address.AddressQubit(0)}
     qubit_ssa_value = {0: test_qubit}
-    rule = NoiseRewriteRule(address_analysis, qubit_ssa_value, noise_params, model)
+    rule = NoiseRewriteRule(address_analysis, qubit_ssa_value, model)
     rule.qubit_ssa_value[0] = test_qubit
     block = ir.Block(
         [
@@ -58,10 +57,9 @@ def test_parallel_qubit_noise():
     pz = 0.01
     p_loss = 0.01
 
-    noise_params = native.GateNoiseParams(
+    model = NoiseTestModel(
         local_px=px, local_py=py, local_pz=pz, local_loss_prob=p_loss
     )
-    model = NoiseTestModel()
 
     test_qubit = ir.TestValue(type=qasm2.QubitType)
     qubit_list = ilist.New(values=[test_qubit])
@@ -73,7 +71,7 @@ def test_parallel_qubit_noise():
     }
     qubit_ssa_value = {0: test_qubit}
 
-    rule = NoiseRewriteRule(address_analysis, qubit_ssa_value, noise_params, model)
+    rule = NoiseRewriteRule(address_analysis, qubit_ssa_value, model)
     rule.qubit_ssa_value[0] = test_qubit
     block = ir.Block(
         [
@@ -112,7 +110,7 @@ def test_cz_gate_noise():
     pz = 0.01
     p_loss = 0.01
 
-    noise_params = native.GateNoiseParams(
+    model = NoiseTestModel(
         cz_paired_gate_px=px,
         cz_paired_gate_py=py,
         cz_paired_gate_pz=pz,
@@ -123,16 +121,14 @@ def test_cz_gate_noise():
         cz_unpaired_loss_prob=p_loss,
     )
 
-    model = NoiseTestModel()
-
     ctrl_qubit = ir.TestValue(type=qasm2.QubitType)
     qarg_qubit = ir.TestValue(type=qasm2.QubitType)
     address_analysis = {
         ctrl_qubit: address.AddressQubit(0),
         qarg_qubit: address.AddressQubit(1),
     }
     qubit_ssa_value = {0: ctrl_qubit, 1: qarg_qubit}
-    rule = NoiseRewriteRule(address_analysis, qubit_ssa_value, noise_params, model)
+    rule = NoiseRewriteRule(address_analysis, qubit_ssa_value, model)
     rule.qubit_ssa_value[0] = ctrl_qubit
     rule.qubit_ssa_value[1] = qarg_qubit
     block = ir.Block(
@@ -187,7 +183,7 @@ def test_parallel_cz_gate_noise():
     pz = 0.01
     p_loss = 0.01
 
-    noise_params = native.GateNoiseParams(
+    model = NoiseTestModel(
         cz_paired_gate_px=px,
         cz_paired_gate_py=py,
         cz_paired_gate_pz=pz,
@@ -198,8 +194,6 @@ def test_parallel_cz_gate_noise():
         cz_unpaired_loss_prob=p_loss,
     )
 
-    model = NoiseTestModel()
-
     ctrl_qubit = ir.TestValue(type=qasm2.QubitType)
     qarg_qubit = ir.TestValue(type=qasm2.QubitType)
     ctrl_list = ilist.New(values=[ctrl_qubit])
@@ -211,7 +205,7 @@ def test_parallel_cz_gate_noise():
         qarg_list.result: address.AddressTuple([address.AddressQubit(1)]),
     }
     qubit_ssa_value = {0: ctrl_qubit, 1: qarg_qubit}
-    rule = NoiseRewriteRule(address_analysis, qubit_ssa_value, noise_params, model)
+    rule = NoiseRewriteRule(address_analysis, qubit_ssa_value, model)
     rule.qubit_ssa_value[0] = ctrl_qubit
     rule.qubit_ssa_value[1] = qarg_qubit
     block = ir.Block(
@@ -275,17 +269,13 @@ def test_method():
     pz = 0.01
     p_loss = 0.01
 
-    noise_params = native.GateNoiseParams(
+    model = NoiseTestModel(
         global_loss_prob=p_loss, global_px=px, global_py=py, global_pz=pz
     )
 
-    model = NoiseTestModel()
-
     test_method.print()
 
-    NoisePass(test_method.dialects, noise_model=model, gate_noise_params=noise_params)(
-        test_method
-    )
+    NoisePass(test_method.dialects, noise_model=model)(test_method)
 
     expected_block = ir.Block(
         [