1use pyo3::types::PyList;
2use qcs::{
3 qvm::{self, http, QvmOptions, QvmResultData},
4 RegisterData,
5};
6use rigetti_pyo3::{
7 create_init_submodule, impl_as_mut_for_wrapper, impl_repr, py_wrap_error, py_wrap_type,
8 pyo3::{exceptions::PyRuntimeError, prelude::*, Python},
9 wrap_error, PyTryFrom, PyWrapper, PyWrapperMut, ToPython, ToPythonError,
10};
11use std::num::NonZeroU16;
12use std::{collections::HashMap, time::Duration};
13
14use crate::register_data::PyRegisterData;
15
16mod api;
17
18use api::PyAddressRequest;
19
20create_init_submodule! {
21 classes: [PyQvmResultData, PyQvmOptions, RawQvmReadoutData, PyQvmClient],
22 errors: [QVMError],
23 funcs: [py_run, py_run_async],
24 submodules: [
25 "api": api::init_submodule
26 ],
27}
28
29wrap_error!(RustQvmError(qcs::qvm::Error));
30py_wrap_error!(api, RustQvmError, QVMError, PyRuntimeError);
31
32#[derive(Clone)]
33pub enum QvmClient {
34 Http(qvm::http::HttpClient),
35 #[cfg(feature = "libquil")]
36 Libquil(qvm::libquil::Client),
37}
38
39#[pyclass(name = "QVMClient")]
40#[derive(Clone)]
41pub struct PyQvmClient {
42 inner: QvmClient,
43}
44
45impl PyQvmClient {
46 pub fn as_client(&self) -> &(dyn qvm::Client + Send + Sync) {
47 match &self.inner {
48 QvmClient::Http(client) => client,
49 #[cfg(feature = "libquil")]
50 QvmClient::Libquil(client) => client,
51 }
52 }
53}
54
55#[pymethods]
56impl PyQvmClient {
57 #[new]
58 fn new() -> PyResult<Self> {
59 Err(PyRuntimeError::new_err("QVMClient cannot be instantiated directly. See the static methods: QVMClient.new_http() and QVMClient.new_libquil()."))
60 }
61
62 #[staticmethod]
63 fn new_http(endpoint: &str) -> PyResult<Self> {
64 let http_client = qvm::http::HttpClient::new(endpoint.to_string());
65 Ok(Self {
66 inner: QvmClient::Http(http_client),
67 })
68 }
69
70 #[cfg(feature = "libquil")]
71 #[staticmethod]
72 fn new_libquil() -> PyResult<Self> {
73 Ok(Self {
74 inner: QvmClient::Libquil(qvm::libquil::Client {}),
75 })
76 }
77
78 #[cfg(not(feature = "libquil"))]
79 #[staticmethod]
80 fn new_libquil() -> PyResult<Self> {
81 Err(PyRuntimeError::new_err(
82 "Cannot create a libquil QVM client as feature is not enabled.",
83 ))
84 }
85
86 #[getter]
87 fn qvm_url(&self) -> PyResult<String> {
88 match &self.inner {
89 QvmClient::Http(client) => Ok(client.qvm_url.to_string()),
90 #[cfg(feature = "libquil")]
91 QvmClient::Libquil(_) => Ok("".into()),
92 }
93 }
94}
95
96#[async_trait::async_trait]
97impl qvm::Client for PyQvmClient {
98 async fn get_version_info(&self, options: &QvmOptions) -> Result<String, qvm::Error> {
100 self.as_client().get_version_info(options).await
101 }
102 async fn run(
104 &self,
105 request: &http::MultishotRequest,
106 options: &QvmOptions,
107 ) -> Result<http::MultishotResponse, qvm::Error> {
108 self.as_client().run(request, options).await
109 }
110 async fn run_and_measure(
120 &self,
121 request: &http::MultishotMeasureRequest,
122 options: &QvmOptions,
123 ) -> Result<Vec<Vec<i64>>, qvm::Error> {
124 self.as_client().run_and_measure(request, options).await
125 }
126 async fn measure_expectation(
128 &self,
129 request: &http::ExpectationRequest,
130 options: &QvmOptions,
131 ) -> Result<Vec<f64>, qvm::Error> {
132 self.as_client().measure_expectation(request, options).await
133 }
134 async fn get_wavefunction(
136 &self,
137 request: &http::WavefunctionRequest,
138 options: &QvmOptions,
139 ) -> Result<Vec<u8>, qvm::Error> {
140 self.as_client().get_wavefunction(request, options).await
141 }
142}
143
144py_wrap_type! {
145 PyQvmResultData(QvmResultData) as "QVMResultData"
146}
147
148#[pymethods]
149impl PyQvmResultData {
150 #[new]
151 fn new(memory: HashMap<String, PyRegisterData>) -> Self {
152 let memory = memory
153 .into_iter()
154 .map(|(key, value)| (key, value.into_inner()))
155 .collect();
156 Self::from(QvmResultData::from_memory_map(memory))
157 }
158
159 #[staticmethod]
160 fn from_memory_map(py: Python<'_>, memory: HashMap<String, PyRegisterData>) -> PyResult<Self> {
161 Ok(Self(QvmResultData::from_memory_map(HashMap::<
162 String,
163 RegisterData,
164 >::py_try_from(
165 py, &memory
166 )?)))
167 }
168
169 #[getter]
170 fn memory(&self, py: Python<'_>) -> PyResult<HashMap<String, PyRegisterData>> {
171 self.as_inner().memory().to_python(py)
172 }
173
174 pub(crate) fn to_raw_readout_data(&self, py: Python<'_>) -> RawQvmReadoutData {
175 RawQvmReadoutData {
176 memory: self
177 .as_inner()
178 .memory()
179 .iter()
180 .map(|(register, matrix)| {
181 (
182 register.to_string(),
183 match matrix {
184 RegisterData::I8(matrix) => PyList::new(py, matrix).into_py(py),
185 RegisterData::F64(matrix) => PyList::new(py, matrix).into_py(py),
186 RegisterData::I16(matrix) => PyList::new(py, matrix).into_py(py),
187 RegisterData::Complex32(matrix) => PyList::new(py, matrix).into_py(py),
188 },
189 )
190 })
191 .collect::<HashMap<String, Py<PyList>>>(),
192 }
193 }
194}
195
196#[pyclass(name = "RawQVMReadoutData")]
197#[derive(Debug)]
198pub(crate) struct RawQvmReadoutData {
199 #[pyo3(get)]
200 memory: HashMap<String, Py<PyList>>,
201}
202
203impl RawQvmReadoutData {
204 fn __repr__(&self) -> String {
205 format!("{self:?}")
206 }
207}
208
209py_wrap_type! {
210 #[derive(Default)]
211 PyQvmOptions(QvmOptions) as "QVMOptions"
212}
213impl_as_mut_for_wrapper!(PyQvmOptions);
214impl_repr!(PyQvmOptions);
215
216#[pymethods]
217impl PyQvmOptions {
218 #[new]
219 #[pyo3(signature = (timeout_seconds = None))]
220 pub fn new(timeout_seconds: Option<f64>) -> Self {
221 Self(QvmOptions {
222 timeout: timeout_seconds.map(Duration::from_secs_f64),
223 })
224 }
225
226 #[getter]
227 pub fn timeout(&self) -> Option<f32> {
228 self.as_inner()
229 .timeout
230 .map(|duration| duration.as_secs_f32())
231 }
232
233 #[setter]
234 pub fn set_timeout(&mut self, timeout_seconds: Option<f64>) {
235 self.as_inner_mut().timeout = timeout_seconds.map(Duration::from_secs_f64);
236 }
237
238 #[staticmethod]
239 #[pyo3(name = "default")]
240 pub fn py_default() -> Self {
241 <Self as Default>::default()
242 }
243}
244
245crate::py_sync::py_function_sync_async! {
246 #[pyfunction]
247 #[tracing::instrument(skip_all)]
248 #[pyo3_opentelemetry::pypropagate(on_context_extraction_failure="ignore")]
249 async fn run(
250 quil: String,
251 #[pyo3(from_py_with = "crate::from_py::non_zero_u16")]
252 shots: NonZeroU16,
253 addresses: HashMap<String, PyAddressRequest>,
254 params: HashMap<String, Vec<f64>>,
255 client: PyQvmClient,
256 measurement_noise: Option<(f64, f64, f64)>,
257 gate_noise: Option<(f64, f64, f64)>,
258 rng_seed: Option<i64>,
259 options: Option<PyQvmOptions>,
260 ) -> PyResult<PyQvmResultData> {
261 let params = params.into_iter().map(|(key, value)| (key.into_boxed_str(), value)).collect();
262 let addresses = addresses.into_iter().map(|(address, request)| (address, request.as_inner().clone())).collect();
263 let options = options.unwrap_or_default();
264 Ok(
265 PyQvmResultData(
266 qcs::qvm::run(
267 &quil,
268 shots,
269 addresses,
270 ¶ms,
271 measurement_noise,
272 gate_noise,
273 rng_seed,
274 &client,
275 options.as_inner()
276 )
277 .await
278 .map_err(RustQvmError::from)
279 .map_err(RustQvmError::to_py_err)?))
280 }
281}