1from collections.abc import Iterable
2from typing import TypeVar
3
4from qiskit import QuantumCircuit, transpile
5from qiskit.primitives import BackendEstimatorV2, BackendSamplerV2
6from qiskit.providers.backend import BackendV1, BackendV2
7from qiskit_ibm_runtime import EstimatorV2, SamplerV2
8
9AUTO_TRANSPILE_SAMPLER_TYPE = TypeVar('AUTO_TRANSPILE_SAMPLER_TYPE', BackendSamplerV2, SamplerV2)
10AUTO_TRANSPILE_ESTIMATOR_TYPE = TypeVar('AUTO_TRANSPILE_ESTIMATOR_TYPE', BackendEstimatorV2, EstimatorV2)
11
12
13def _get_transpiled_sampler_pubs(
14 pubs: list[tuple | QuantumCircuit], backend: BackendV1 | BackendV2, optimization_level: int = 2
15) -> list[tuple | QuantumCircuit]:
16 new_pubs = []
17 for pub in pubs:
18 if isinstance(pub, QuantumCircuit):
19 pub = transpile(pub, backend, optimization_level=optimization_level)
20 elif isinstance(pub, tuple):
21 pub = (transpile(pub[0], backend, optimization_level=optimization_level), *pub[1:])
22 new_pubs.append(pub)
23 return new_pubs
24
25
26def _get_transpiled_estimator_pubs(
27 pubs: list[tuple], backend: BackendV1 | BackendV2, optimization_level: int = 2
28) -> list[tuple | QuantumCircuit]:
29 new_pubs = []
30 for pub in pubs:
31 circuit, operator, *args = pub
32 transp_circ: QuantumCircuit = transpile(circuit, backend, optimization_level=optimization_level)
33
34 if isinstance(operator, Iterable):
35 transp_op = [op.apply_layout(transp_circ.layout, num_qubits=transp_circ.num_qubits) for op in operator]
36 else:
37 transp_op = operator.apply_layout(transp_circ.layout, num_qubits=transp_circ.num_qubits)
38
39 pub = (transp_circ, transp_op, *args)
40 new_pubs.append(pub)
41 return new_pubs
42
43
44def _assign_sampler_pubs(pubs: list[tuple | QuantumCircuit]) -> list[QuantumCircuit]:
45 new_pubs = []
46 for pub in pubs:
47 if isinstance(pub, tuple):
48 qc, *params = pub
49 pub = qc.assign_parameters(*params)
50 new_pubs.append(pub)
51 return new_pubs
52
53
54def _assign_estimator_pubs(pubs: list[tuple]) -> list[tuple]:
55 new_pubs = []
56 for pub in pubs:
57 if len(pub) > 2:
58 qc, observable, *params = pub
59 pub = (qc.assign_parameters(*params), observable)
60 new_pubs.append(pub)
61 return new_pubs
62
63
[docs]
64def set_sampler_auto_run_behavior(
65 sampler: AUTO_TRANSPILE_SAMPLER_TYPE, auto_transpile: bool = False, auto_transpile_level: int = 2, auto_assign: bool = False
66) -> AUTO_TRANSPILE_SAMPLER_TYPE:
67 """
68 Set chosen automatic behavior on a sampler instance.
69
70 Args:
71 sampler (AUTO_TRANSPILE_SAMPLER_TYPE): Compatible sampler instance
72 auto_transpile (bool, optional): Whether to automatically transpile to the samplers backend. Defaults to False.
73 auto_transpile_level (int, optional): What level of optimization to set. Defaults to 2.
74 auto_assign (bool, optional): Whether to automatically assign parameters in case of parameterized circuits. Defaults to False.
75
76 Returns:
77 AUTO_TRANSPILE_SAMPLER_TYPE: Same instance with modified run() method.
78 """
79 func = sampler.run
80
81 def run_wrapper(pubs, *args, shots: int | None = None):
82 if auto_transpile:
83 pubs = _get_transpiled_sampler_pubs(pubs, sampler._backend, optimization_level=auto_transpile_level)
84 if auto_assign:
85 pubs = _assign_sampler_pubs(pubs)
86
87 return func(pubs, *args, shots=shots)
88
89 sampler.run = run_wrapper
90 return sampler
91
92
[docs]
93def set_estimator_auto_run_behavior(
94 estimator: AUTO_TRANSPILE_ESTIMATOR_TYPE, auto_transpile: bool = False, auto_transpile_level: int = 2, auto_assign: bool = False
95) -> AUTO_TRANSPILE_ESTIMATOR_TYPE:
96 """
97 Set chosen automatic behavior on a estimator instance.
98
99 Args:
100 estimator (AUTO_TRANSPILE_ESTIMATOR_TYPE): Compatible estimator instance
101 auto_transpile (bool, optional): Whether to automatically transpile to the estimators backend. Defaults to False.
102 auto_transpile_level (int, optional): What level of optimization to set. Defaults to 2.
103 auto_assign (bool, optional): Whether to automatically assign parameters in case of parameterized circuits. Defaults to False.
104
105 Returns:
106 AUTO_TRANSPILE_ESTIMATOR_TYPE: Same instance with modified run() method.
107 """
108 func = estimator.run
109
110 def run_wrapper(pubs, *args, precision: float | None = None):
111 if auto_transpile:
112 pubs = _get_transpiled_estimator_pubs(pubs, estimator._backend, optimization_level=auto_transpile_level)
113 if auto_assign:
114 pubs = _assign_estimator_pubs(pubs)
115 return func(pubs, *args, precision=precision)
116
117 estimator.run = run_wrapper
118 return estimator
