Add end to end noise aware zne test

qermit/noise_model/transpiler_backend.py

@@ -3,8 +3,8 @@
 from pytket.backends.backendresult import BackendResult
 from pytket.utils.outcomearray import OutcomeArray
 import uuid
-from pytket.passes import BasePass
-from typing import Dict, List, Optional, Iterator
+from pytket.passes import BasePass, CustomPass
+from typing import Dict, List, Optional, Iterator, Sequence, Iterable
 from pytket import Circuit, Bit
 
 
@@ -51,6 +51,20 @@ def __init__(
         self.max_batch_size = max_batch_size
         self.result_dict = result_dict
 
+    def default_compilation_pass(self, **kwargs) -> BasePass:
+
+        def transform(circuit: Circuit) -> Circuit:
+            return circuit
+
+        return CustomPass(transform=transform)
+
+    def rebase_pass(self):
+
+        def transform(circuit: Circuit) -> Circuit:
+            return circuit
+
+        return CustomPass(transform=transform)
+
     def run_circuit(
         self,
         circuit: Circuit,
@@ -70,6 +84,17 @@ def run_circuit(
         handle = self.process_circuit(circuit, n_shots, **kwargs)
         return self.get_result(handle=handle)
 
+    def process_circuits(
+        self,
+        circuits: Sequence[Circuit],
+        n_shots: Sequence[int],
+    ) -> List[uuid.UUID]:
+
+        return [
+            self.process_circuit(circuit=circuit, n_shots=n)
+            for circuit, n in zip(circuits, n_shots)
+        ]
+
     def process_circuit(
         self,
         circuit: Circuit,
@@ -104,6 +129,11 @@ def process_circuit(
 
         return handle
 
+    def get_results(self, handles: Iterable[uuid.UUID]) -> List[BackendResult]:
+        return [
+            self.get_result(handle) for handle in handles
+        ]
+
     def get_result(self, handle: uuid.UUID) -> BackendResult:
         """Retrieve result from backend.
 

tests/zne_test.py

@@ -47,6 +47,7 @@
 from qermit.taskgraph import gen_MeasurementReduction_MitEx
 from qermit.noise_model import NoiseModel, ErrorDistribution
 from qermit.zero_noise_extrapolation.zne import gen_noise_scaled_mitex
+from qermit.noise_model import TranspilerBackend, PauliErrorTranspile
 
 n_qubits = 2
 
@@ -842,6 +843,60 @@ def test_end_to_end_noise_scaled_mitex():
     assert abs(qubit_pauli_operator_list[1]._dict[qps_noisless_zero]) < 0.1
 
 
+@pytest.mark.high_compute
+def test_end_to_end_noise_aware_zne():
+
+    error_rate = 0.1
+    error_distribution = ErrorDistribution(
+        distribution={(Pauli.X, Pauli.I): error_rate}
+    )
+    noise_model = NoiseModel(
+        noise_model={OpType.CZ: error_distribution}
+    )
+    transpiler = PauliErrorTranspile(noise_model=noise_model)
+    backend = TranspilerBackend(transpiler=transpiler)
+
+    zne_mitex = gen_ZNE_MitEx(
+        backend=backend,
+        noise_scaling_list=[2, 5, 3, 1, 4],
+        show_fit=True,
+        folding_type=Folding.noise_aware,
+        fit_type=Fit.exponential,
+        noise_model=noise_model,
+        n_noisy_circuit_samples=1000,
+    )
+
+    circuit_noisy = Circuit(2).CZ(0, 1)
+
+    qps_noisy_noisy = QubitPauliString(map={Qubit(0): Pauli.Z, Qubit(1): Pauli.Z})
+
+    qubit_pauli_operator_noisy = QubitPauliOperator(
+        dictionary={qps_noisy_noisy: 1}
+    )
+
+    observable_tracker_noisy = ObservableTracker(
+        qubit_pauli_operator=qubit_pauli_operator_noisy
+    )
+
+    shots = 10000
+    ansatz_circuit_noisy = AnsatzCircuit(
+        Circuit=circuit_noisy,
+        Shots=shots,
+        SymbolsDict={}
+    )
+
+    observable_experiment_noisy = ObservableExperiment(
+        AnsatzCircuit=ansatz_circuit_noisy,
+        ObservableTracker=observable_tracker_noisy,
+    )
+
+    qubit_pauli_operator_list = zne_mitex.run(
+        mitex_wires=[observable_experiment_noisy]
+    )
+
+    assert abs(qubit_pauli_operator_list[0]._dict[qps_noisy_noisy] - 1) < 0.1
+
+
 if __name__ == "__main__":
     test_no_qubit_relabel()
     test_extrapolation_task_gen()