diff --git a/qermit/noise_model/transpiler_backend.py b/qermit/noise_model/transpiler_backend.py index 5178d52d..8c05ff52 100644 --- a/qermit/noise_model/transpiler_backend.py +++ b/qermit/noise_model/transpiler_backend.py @@ -37,7 +37,7 @@ def __init__( transpiler: BasePass, result_dict: Dict[ResultHandle, BackendResult] = {}, max_batch_size: int = 1000, - n_cores: Optional[int] = None, + n_cores: int = 1, ): """Initialisation method. @@ -52,9 +52,9 @@ def __init__( results within backend, defaults to {} :type result_dict: Dict[ResultHandle, BackendResult], optional :param n_cores: Shots will be taken in parallel. This parameter - specifies the number of cores to use. - :type n_cores: Optional[int]. Defaults to None, using all available - cores. + specifies the number of cores to use. The default is to use + one core. + :type n_cores: Optional[int]. Defaults to 1. """ self.transpiler = transpiler @@ -261,16 +261,28 @@ def get_counts( :rtype: Iterator[Counter] """ - if cbits is not None: - cbits_list = [cbit.to_list() for cbit in cbits] - else: - cbits_list = None + if self.n_cores > 1: - with multiprocessing.Pool(self.n_cores) as pool: - processes = [ - pool.apply_async(self._get_batch_counts, args=(circuit_list, cbits_list)) - for circuit_list in self._gen_batches(circuit, n_shots) - ] - counter_list = [p.get() for p in processes] + if cbits is not None: + cbits_list = [cbit.to_list() for cbit in cbits] + else: + cbits_list = None - return sum(counter_list, Counter()) + with multiprocessing.Pool(self.n_cores) as pool: + processes = [ + pool.apply_async(self._get_batch_counts, args=(circuit_list, cbits_list)) + for circuit_list in self._gen_batches(circuit, n_shots) + ] + counter_list = [p.get() for p in processes] + + return sum(counter_list, Counter()) + + else: + + counter: Counter = Counter() + for circuit_list in self._gen_batches(circuit, n_shots): + result_list = self.backend.run_circuits(circuit_list, n_shots=1) + counter += sum((result.get_counts(cbits=cbits) + for result in result_list), Counter()) + + return counter diff --git a/tests/zne_test.py b/tests/zne_test.py index e3f0b6e6..41837611 100644 --- a/tests/zne_test.py +++ b/tests/zne_test.py @@ -863,7 +863,10 @@ def test_end_to_end_noise_aware_zne_mitex_starting_from_ptm() -> None: } ) transpiler = PauliErrorTranspile(noise_model=noise_model) - backend = TranspilerBackend(transpiler=transpiler) + backend = TranspilerBackend( + transpiler=transpiler, + n_cores=1, + ) # Here we perform ZNE with some unevenly spaced # noise scaling values. @@ -920,7 +923,10 @@ def test_end_to_end_noise_aware_zne_mitex(): noise_model={OpType.CZ: error_distribution} ) transpiler = PauliErrorTranspile(noise_model=noise_model) - backend = TranspilerBackend(transpiler=transpiler) + backend = TranspilerBackend( + transpiler=transpiler, + n_cores=1, + ) zne_mitex = gen_ZNE_MitEx( backend=backend,