Skip to main content

qcs_sdk/qpu/
result_data.rs

1use std::collections::HashMap;
2
3use pyo3::{
4    exceptions::PyNotImplementedError,
5    pyclass,
6    pyclass::CompareOp,
7    pymethods,
8    types::{PyComplex, PyFloat, PyInt, PyList},
9    IntoPy, Py, PyResult, Python,
10};
11use qcs::qpu::{result_data::MemoryValues, QpuResultData, ReadoutValues};
12use rigetti_pyo3::{impl_repr, py_wrap_type, py_wrap_union_enum, PyTryFrom, PyWrapper, ToPython};
13
14py_wrap_union_enum! {
15    PyReadoutValues(ReadoutValues) as "ReadoutValues" {
16        integer: Integer => Vec<Py<PyInt>>,
17        real: Real => Vec<Py<PyFloat>>,
18        complex: Complex => Vec<Py<PyComplex>>
19    }
20}
21
22py_wrap_union_enum! {
23    #[derive(Debug, PartialEq)]
24    PyMemoryValues(MemoryValues) as "MemoryValues" {
25        binary: Binary => Vec<Py<PyInt>>,
26        integer: Integer => Vec<Py<PyInt>>,
27        real: Real => Vec<Py<PyFloat>>
28    }
29}
30impl_repr!(PyMemoryValues);
31
32#[pymethods]
33impl PyMemoryValues {
34    fn __richcmp__(&self, other: &PyMemoryValues, op: CompareOp) -> PyResult<bool> {
35        match op {
36            CompareOp::Eq => Ok(self == other),
37            CompareOp::Ne => Ok(self != other),
38            _ => Err(PyNotImplementedError::new_err(
39                "MemoryValues only supports equality comparisons",
40            )),
41        }
42    }
43}
44
45py_wrap_type! {
46    PyQpuResultData(QpuResultData) as "QPUResultData"
47}
48impl_repr!(PyQpuResultData);
49
50#[pymethods]
51impl PyQpuResultData {
52    #[new]
53    fn __new__(
54        py: Python<'_>,
55        mappings: HashMap<String, String>,
56        readout_values: HashMap<String, PyReadoutValues>,
57        memory_values: HashMap<String, PyMemoryValues>,
58    ) -> PyResult<Self> {
59        Ok(Self(QpuResultData::from_mappings_and_values(
60            mappings,
61            HashMap::<String, ReadoutValues>::py_try_from(py, &readout_values)?,
62            HashMap::<String, MemoryValues>::py_try_from(py, &memory_values)?,
63        )))
64    }
65
66    #[getter]
67    fn mappings(&self, py: Python<'_>) -> PyResult<HashMap<String, String>> {
68        self.as_inner().mappings().to_python(py)
69    }
70
71    #[getter]
72    fn readout_values(&self, py: Python<'_>) -> PyResult<HashMap<String, PyReadoutValues>> {
73        self.as_inner().readout_values().to_python(py)
74    }
75
76    #[getter]
77    fn memory_values(&self, py: Python<'_>) -> PyResult<HashMap<String, PyMemoryValues>> {
78        self.as_inner().memory_values().to_python(py)
79    }
80
81    pub(crate) fn to_raw_readout_data(&self, py: Python<'_>) -> RawQpuReadoutData {
82        RawQpuReadoutData {
83            mappings: self.as_inner().mappings().clone(),
84            readout_values: self
85                .as_inner()
86                .readout_values()
87                .iter()
88                .map(|(register, values)| {
89                    (
90                        register.to_string(),
91                        match values {
92                            ReadoutValues::Integer(values) => PyList::new(py, values).into_py(py),
93                            ReadoutValues::Real(values) => PyList::new(py, values).into_py(py),
94                            ReadoutValues::Complex(values) => PyList::new(py, values).into_py(py),
95                        },
96                    )
97                })
98                .collect::<HashMap<String, Py<PyList>>>(),
99            memory_values: self
100                .as_inner()
101                .memory_values()
102                .iter()
103                .map(|(register, memory_values)| {
104                    (
105                        register.to_string(),
106                        match memory_values {
107                            MemoryValues::Binary(values) => PyList::new(py, values).into_py(py),
108                            MemoryValues::Integer(values) => PyList::new(py, values).into_py(py),
109                            MemoryValues::Real(values) => PyList::new(py, values).into_py(py),
110                        },
111                    )
112                })
113                .collect::<HashMap<String, Py<PyList>>>(),
114        }
115    }
116}
117
118/// A wrapper type for data returned by the QPU in a more flat structure than
119/// [`PyQpuResultData`] offers. This makes it more convenient to work with
120/// the data if you don't care what type of number the readout values for
121/// each register contains.
122#[derive(Debug)]
123#[pyclass(name = "RawQPUReadoutData")]
124#[pyo3(get_all)]
125pub struct RawQpuReadoutData {
126    mappings: HashMap<String, String>,
127    readout_values: HashMap<String, Py<PyList>>,
128    memory_values: HashMap<String, Py<PyList>>,
129}
130
131impl RawQpuReadoutData {
132    fn __repr__(&self) -> String {
133        format!("{self:?}")
134    }
135}