qcs_sdk/
execution_data.rs

1use std::collections::HashMap;
2use std::time::Duration;
3
4use numpy::{Complex64, PyArray2};
5use pyo3::exceptions::{PyKeyError, PyRuntimeError};
6use pyo3::pyclass::CompareOp;
7use pyo3::{
8    exceptions::PyValueError,
9    pyclass, pymethods,
10    types::{PyBytes, PyDelta},
11    IntoPy, Py, PyObject, PyRef, PyRefMut, PyResult, Python, ToPyObject,
12};
13use qcs::qvm::QvmResultData;
14use qcs::{ExecutionData, RegisterMap, RegisterMatrix, ResultData};
15use rigetti_pyo3::{
16    impl_repr, py_wrap_data_struct, py_wrap_error, py_wrap_type, py_wrap_union_enum, wrap_error,
17    PyTryFrom, PyWrapper, ToPython, ToPythonError,
18};
19
20use crate::qpu::PyQpuResultData;
21use crate::qvm::PyQvmResultData;
22
23py_wrap_union_enum! {
24    #[pyo3(module = "qcs_sdk")]
25    PyResultData(ResultData) as "ResultData" {
26        qpu: Qpu => PyQpuResultData,
27        qvm: Qvm => PyQvmResultData
28    }
29}
30impl_repr!(PyResultData);
31
32impl Default for PyResultData {
33    fn default() -> Self {
34        PyResultData(ResultData::Qvm(QvmResultData::from_memory_map(
35            HashMap::new(),
36        )))
37    }
38}
39
40wrap_error!(RustRegisterMatrixConversionError(
41    qcs::RegisterMatrixConversionError
42));
43py_wrap_error!(
44    execution_data,
45    RustRegisterMatrixConversionError,
46    RegisterMatrixConversionError,
47    PyValueError
48);
49
50#[pymethods]
51impl PyResultData {
52    pub fn to_register_map(&self, py: Python) -> PyResult<PyRegisterMap> {
53        self.as_inner()
54            .to_register_map()
55            .map_err(RustRegisterMatrixConversionError)
56            .map_err(ToPythonError::to_py_err)?
57            .to_python(py)
58    }
59
60    pub fn to_raw_readout_data(&self, py: Python) -> PyResult<PyObject> {
61        match self.as_inner() {
62            ResultData::Qpu(_) => self
63                .to_qpu(py)
64                .map(|data| data.to_raw_readout_data(py).into_py(py)),
65            ResultData::Qvm(_) => self
66                .to_qvm(py)
67                .map(|data| data.to_raw_readout_data(py).into_py(py)),
68        }
69    }
70}
71
72py_wrap_data_struct! {
73    #[pyo3(module = "qcs_sdk")]
74    #[derive(Debug, PartialEq)]
75    PyExecutionData(ExecutionData) as "ExecutionData" {
76        result_data: ResultData => PyResultData,
77        duration: Option<Duration> => Option<Py<PyDelta>>
78    }
79}
80impl_repr!(PyExecutionData);
81
82#[pymethods]
83impl PyExecutionData {
84    #[new]
85    fn __new__(
86        py: Python<'_>,
87        // result_data is optional here because pickling an object requires calling __new__ without
88        // any arguments.
89        result_data: Option<PyResultData>,
90        duration: Option<Py<PyDelta>>,
91    ) -> PyResult<Self> {
92        Ok(Self(ExecutionData {
93            result_data: ResultData::py_try_from(py, &result_data.unwrap_or_default())?,
94            duration: duration
95                .map(|delta| {
96                    delta
97                        .as_ref(py)
98                        .call_method0("total_seconds")
99                        .map(|result| result.extract::<f64>())?
100                        .map(Duration::from_secs_f64)
101                })
102                .transpose()?,
103        }))
104    }
105
106    pub fn __getstate__<'a>(&self, py: Python<'a>) -> PyResult<&'a PyBytes> {
107        Ok(PyBytes::new(
108            py,
109            &serde_json::to_vec(self.as_inner())
110                .map_err(|e| PyRuntimeError::new_err(format!("failed to serialize: {e}")))?,
111        ))
112    }
113
114    pub fn __setstate__(&mut self, state: &PyBytes) -> PyResult<()> {
115        let execution_data: ExecutionData = serde_json::from_slice(state.as_bytes())
116            .map_err(|e| PyRuntimeError::new_err(format!("failed to deserialize: {e}")))?;
117        *self = PyExecutionData(execution_data);
118        Ok(())
119    }
120
121    fn __richcmp__(&self, other: &Self, op: CompareOp, py: Python<'_>) -> PyObject {
122        match op {
123            CompareOp::Eq => (self == other).into_py(py),
124            CompareOp::Ne => (self != other).into_py(py),
125            _ => py.NotImplemented(),
126        }
127    }
128}
129
130py_wrap_type! {
131    #[pyo3(mapping)]
132    PyRegisterMap(RegisterMap) as "RegisterMap";
133}
134impl_repr!(PyRegisterMap);
135
136py_wrap_type! {
137    PyRegisterMatrix(RegisterMatrix) as "RegisterMatrix"
138}
139impl_repr!(PyRegisterMatrix);
140
141#[pymethods]
142impl PyRegisterMatrix {
143    fn to_ndarray(&self, py: Python<'_>) -> PyResult<PyObject> {
144        match self.as_inner() {
145            RegisterMatrix::Integer(_) => self.to_integer(py).map(|matrix| matrix.to_object(py)),
146            RegisterMatrix::Real(_) => self.to_real(py).map(|matrix| matrix.to_object(py)),
147            RegisterMatrix::Complex(_) => self.to_complex(py).map(|matrix| matrix.to_object(py)),
148        }
149    }
150
151    #[staticmethod]
152    fn from_integer(inner: &PyArray2<i64>) -> PyRegisterMatrix {
153        Self(RegisterMatrix::Integer(inner.to_owned_array()))
154    }
155
156    fn to_integer<'a>(&self, py: Python<'a>) -> PyResult<&'a PyArray2<i64>> {
157        if let Some(matrix) = self.as_inner().as_integer() {
158            Ok(PyArray2::from_array(py, matrix))
159        } else {
160            Err(PyValueError::new_err("not a integer register"))
161        }
162    }
163
164    fn as_integer<'a>(&self, py: Python<'a>) -> Option<&'a PyArray2<i64>> {
165        if let Some(matrix) = self.as_inner().as_integer() {
166            Some(PyArray2::from_array(py, matrix))
167        } else {
168            None
169        }
170    }
171
172    fn is_integer(&self) -> bool {
173        matches!(self.as_inner(), RegisterMatrix::Integer(_))
174    }
175
176    #[staticmethod]
177    fn from_real(inner: &PyArray2<f64>) -> PyRegisterMatrix {
178        Self(RegisterMatrix::Real(inner.to_owned_array()))
179    }
180
181    fn to_real<'a>(&self, py: Python<'a>) -> PyResult<&'a PyArray2<f64>> {
182        if let Some(matrix) = self.as_inner().as_real() {
183            Ok(PyArray2::from_array(py, matrix))
184        } else {
185            Err(PyValueError::new_err("not a real numbered register"))
186        }
187    }
188
189    fn as_real<'a>(&self, py: Python<'a>) -> Option<&'a PyArray2<f64>> {
190        if let Some(matrix) = self.as_inner().as_real() {
191            Some(PyArray2::from_array(py, matrix))
192        } else {
193            None
194        }
195    }
196
197    fn is_real(&self) -> bool {
198        matches!(self.as_inner(), RegisterMatrix::Real(_))
199    }
200
201    #[staticmethod]
202    fn from_complex(inner: &PyArray2<Complex64>) -> PyRegisterMatrix {
203        Self(RegisterMatrix::Complex(inner.to_owned_array()))
204    }
205
206    fn to_complex<'a>(&self, py: Python<'a>) -> PyResult<&'a PyArray2<Complex64>> {
207        if let Some(matrix) = self.as_inner().as_complex() {
208            Ok(PyArray2::from_array(py, matrix))
209        } else {
210            Err(PyValueError::new_err("not a complex numbered register"))
211        }
212    }
213
214    fn as_complex<'a>(&self, py: Python<'a>) -> Option<&'a PyArray2<Complex64>> {
215        if let Some(matrix) = self.as_inner().as_complex() {
216            Some(PyArray2::from_array(py, matrix))
217        } else {
218            None
219        }
220    }
221
222    fn is_complex(&self) -> bool {
223        matches!(self.as_inner(), RegisterMatrix::Complex(_))
224    }
225}
226
227#[pymethods]
228impl PyRegisterMap {
229    pub fn get_register_matrix(&self, register_name: &str) -> Option<PyRegisterMatrix> {
230        self.as_inner()
231            .get_register_matrix(register_name)
232            .map(PyRegisterMatrix::from)
233    }
234
235    pub fn __len__(&self) -> usize {
236        self.as_inner().0.len()
237    }
238
239    pub fn __contains__(&self, key: &str) -> bool {
240        self.as_inner().0.contains_key(key)
241    }
242
243    pub fn __getitem__(&self, item: &str) -> PyResult<PyRegisterMatrix> {
244        self.get_register_matrix(item)
245            .ok_or(PyKeyError::new_err(format!(
246                "Key {item} not found in RegisterMap"
247            )))
248    }
249
250    pub fn __iter__(&self, py: Python<'_>) -> PyResult<Py<PyRegisterMapKeysIter>> {
251        Py::new(
252            py,
253            PyRegisterMapKeysIter {
254                inner: self.as_inner().0.clone().into_iter(),
255            },
256        )
257    }
258
259    pub fn keys(&self, py: Python<'_>) -> PyResult<Py<PyRegisterMapKeysIter>> {
260        self.__iter__(py)
261    }
262
263    pub fn values(&self, py: Python<'_>) -> PyResult<Py<PyRegisterMapValuesIter>> {
264        Py::new(
265            py,
266            PyRegisterMapValuesIter {
267                inner: self.as_inner().0.clone().into_iter(),
268            },
269        )
270    }
271
272    pub fn items(&self, py: Python<'_>) -> PyResult<Py<PyRegisterMapItemsIter>> {
273        Py::new(
274            py,
275            PyRegisterMapItemsIter {
276                inner: self.as_inner().0.clone().into_iter(),
277            },
278        )
279    }
280
281    pub fn get(&self, key: &str, default: Option<PyRegisterMatrix>) -> Option<PyRegisterMatrix> {
282        self.__getitem__(key).ok().or(default)
283    }
284}
285
286#[pyclass]
287pub struct PyRegisterMapItemsIter {
288    inner: std::collections::hash_map::IntoIter<String, RegisterMatrix>,
289}
290
291#[pymethods]
292impl PyRegisterMapItemsIter {
293    fn __iter__(slf: PyRef<'_, Self>) -> PyRef<'_, Self> {
294        slf
295    }
296
297    fn __next__(mut slf: PyRefMut<'_, Self>) -> Option<(String, PyRegisterMatrix)> {
298        slf.inner
299            .next()
300            .map(|(register, matrix)| (register, PyRegisterMatrix(matrix)))
301    }
302}
303
304// The keys and values iterators are built on the iterator of the full
305// `HashMap`, because the iterators returned by `keys()` and `values()`
306// return an iterator with a _reference_ to the underlying `HashMap`.
307// The reference would require these structs to specify a lifetime,
308// which is incompatible with `#[pyclass]`.
309#[pyclass]
310pub struct PyRegisterMapKeysIter {
311    inner: std::collections::hash_map::IntoIter<String, RegisterMatrix>,
312}
313
314#[pymethods]
315impl PyRegisterMapKeysIter {
316    fn __iter__(slf: PyRef<'_, Self>) -> PyRef<'_, Self> {
317        slf
318    }
319
320    fn __next__(mut slf: PyRefMut<'_, Self>) -> Option<String> {
321        slf.inner.next().map(|(register, _)| register)
322    }
323}
324
325#[pyclass]
326pub struct PyRegisterMapValuesIter {
327    inner: std::collections::hash_map::IntoIter<String, RegisterMatrix>,
328}
329
330#[pymethods]
331impl PyRegisterMapValuesIter {
332    fn __iter__(slf: PyRef<'_, Self>) -> PyRef<'_, Self> {
333        slf
334    }
335
336    fn __next__(mut slf: PyRefMut<'_, Self>) -> Option<PyRegisterMatrix> {
337        slf.inner.next().map(|(_, matrix)| PyRegisterMatrix(matrix))
338    }
339}
qcs_sdk/execution_data.rs

qcs_sdk/
execution_data.rs
