qudit_expr/expressions/
kraus.rs

1use std::ops::{Deref, DerefMut};
2
3use crate::{
4    GenerationShape, TensorExpression,
5    expressions::JittableExpression,
6    index::{IndexDirection, TensorIndex},
7};
8
9use super::NamedExpression;
10use qudit_core::QuditSystem;
11use qudit_core::Radices;
12
13#[derive(PartialEq, Eq, Debug, Clone)]
14pub struct KrausOperatorsExpression {
15    inner: NamedExpression,
16    input_radices: Radices,
17    output_radices: Radices,
18    num_operators: usize,
19}
20
21impl KrausOperatorsExpression {
22    pub fn new<T: AsRef<str>>(input: T) -> Self {
23        TensorExpression::new(input).try_into().unwrap()
24    }
25
26    pub fn num_qudits(&self) -> usize {
27        if self.input_radices == self.output_radices {
28            self.input_radices.num_qudits()
29        } else {
30            panic!("Input and output number of qudits are different for kraus operator.")
31        }
32    }
33}
34
35impl JittableExpression for KrausOperatorsExpression {
36    fn generation_shape(&self) -> GenerationShape {
37        GenerationShape::Tensor3D(
38            self.num_operators,
39            self.output_radices.dimension(),
40            self.input_radices.dimension(),
41        )
42    }
43}
44
45impl AsRef<NamedExpression> for KrausOperatorsExpression {
46    fn as_ref(&self) -> &NamedExpression {
47        &self.inner
48    }
49}
50
51impl From<KrausOperatorsExpression> for NamedExpression {
52    fn from(value: KrausOperatorsExpression) -> Self {
53        value.inner
54    }
55}
56
57impl Deref for KrausOperatorsExpression {
58    type Target = NamedExpression;
59
60    fn deref(&self) -> &Self::Target {
61        &self.inner
62    }
63}
64
65impl DerefMut for KrausOperatorsExpression {
66    fn deref_mut(&mut self) -> &mut Self::Target {
67        &mut self.inner
68    }
69}
70
71impl From<KrausOperatorsExpression> for TensorExpression {
72    fn from(value: KrausOperatorsExpression) -> Self {
73        let KrausOperatorsExpression {
74            inner,
75            input_radices,
76            output_radices,
77            num_operators,
78        } = value;
79        // TODO: add a proper implementation of into_iter for QuditRadices
80        let indices = [num_operators]
81            .into_iter()
82            .map(|r| (IndexDirection::Batch, r))
83            .chain(
84                output_radices
85                    .iter()
86                    .map(|r| (IndexDirection::Output, usize::from(*r))),
87            )
88            .chain(
89                input_radices
90                    .iter()
91                    .map(|r| (IndexDirection::Input, usize::from(*r))),
92            )
93            .enumerate()
94            .map(|(i, (d, r))| TensorIndex::new(d, i, r))
95            .collect();
96        TensorExpression::from_raw(indices, inner)
97    }
98}
99
100impl TryFrom<TensorExpression> for KrausOperatorsExpression {
101    // TODO: Come up with proper error handling
102    type Error = String;
103
104    fn try_from(value: TensorExpression) -> Result<Self, Self::Error> {
105        let mut num_operators = None;
106        let mut input_radices = vec![];
107        let mut output_radices = vec![];
108        for idx in value.indices() {
109            match idx.direction() {
110                IndexDirection::Batch => match num_operators {
111                    Some(n) => num_operators = Some(n * idx.index_size()),
112                    None => num_operators = Some(idx.index_size()),
113                },
114                IndexDirection::Input => {
115                    input_radices.push(idx.index_size());
116                }
117                IndexDirection::Output => {
118                    output_radices.push(idx.index_size());
119                }
120                _ => unreachable!(),
121            }
122        }
123
124        Ok(KrausOperatorsExpression {
125            inner: value.into(),
126            input_radices: input_radices.into(),
127            output_radices: output_radices.into(),
128            num_operators: num_operators.unwrap_or(1),
129        })
130    }
131}
132
133#[cfg(feature = "python")]
134mod python {
135    use super::*;
136    use crate::python::PyExpressionRegistrar;
137    use pyo3::prelude::*;
138    use pyo3_stub_gen::derive::*;
139    use qudit_core::Radix;
140
141    /// A symbolic expression representing a set of Kraus operators.
142    #[gen_stub_pyclass]
143    #[pyclass(name = "KrausOperatorsExpression", module = "openqudit.expressions")]
144    pub struct PyKrausOperatorsExpression {
145        expr: KrausOperatorsExpression,
146    }
147
148    #[gen_stub_pymethods]
149    #[pymethods]
150    impl PyKrausOperatorsExpression {
151        /// Parses a Kraus operators expression from its string representation.
152        ///
153        /// # Arguments
154        ///
155        /// * `expr` - The textual definition of the Kraus operators expression.
156        #[new]
157        fn new(expr: String) -> Self {
158            Self {
159                expr: KrausOperatorsExpression::new(expr),
160            }
161        }
162
163        /// Returns the number of free (unbound) parameters in this expression.
164        fn num_params(&self) -> usize {
165            self.expr.num_params()
166        }
167
168        /// Returns the name assigned to this expression.
169        fn name(&self) -> String {
170            self.expr.name().to_string()
171        }
172
173        /// Returns the radix of each qudit acted on by these Kraus operators.
174        fn radices(&self) -> Vec<Radix> {
175            self.expr.input_radices.to_vec()
176        }
177
178        /// Returns the number of qudits acted on. Panics if the input and output
179        /// radices of the Kraus operators differ.
180        fn num_qudits(&self) -> usize {
181            self.expr.num_qudits()
182        }
183
184        /// Returns the number of Kraus operators in this set.
185        fn num_operators(&self) -> usize {
186            self.expr.num_operators
187        }
188
189        /// Returns the Hilbert space dimension each Kraus operator acts on.
190        fn dimension(&self) -> usize {
191            self.expr.input_radices.dimension()
192        }
193
194        fn __repr__(&self) -> String {
195            format!(
196                "KrausOperatorsExpression(name='{}', radices={:?}, num_operators={}, params={})",
197                self.expr.name(),
198                self.expr.input_radices.to_vec(),
199                self.expr.num_operators,
200                self.expr.num_params()
201            )
202        }
203    }
204
205    impl From<KrausOperatorsExpression> for PyKrausOperatorsExpression {
206        fn from(value: KrausOperatorsExpression) -> Self {
207            PyKrausOperatorsExpression { expr: value }
208        }
209    }
210
211    impl From<PyKrausOperatorsExpression> for KrausOperatorsExpression {
212        fn from(value: PyKrausOperatorsExpression) -> Self {
213            value.expr
214        }
215    }
216
217    impl<'py> IntoPyObject<'py> for KrausOperatorsExpression {
218        type Target = <PyKrausOperatorsExpression as IntoPyObject<'py>>::Target;
219        type Output = Bound<'py, Self::Target>;
220        type Error = PyErr;
221
222        fn into_pyobject(self, py: Python<'py>) -> Result<Self::Output, Self::Error> {
223            let py_expr = PyKrausOperatorsExpression::from(self);
224            Bound::new(py, py_expr)
225        }
226    }
227
228    impl<'a, 'py> FromPyObject<'a, 'py> for KrausOperatorsExpression {
229        type Error = PyErr;
230
231        fn extract(ob: Borrowed<'a, 'py, PyAny>) -> PyResult<Self> {
232            let py_expr: PyRef<PyKrausOperatorsExpression> = ob.extract()?;
233            Ok(py_expr.expr.clone())
234        }
235    }
236
237    /// Registers the KrausOperatorsExpression class with the Python module.
238    fn register(parent_module: &Bound<'_, PyModule>) -> PyResult<()> {
239        parent_module.add_class::<PyKrausOperatorsExpression>()?;
240        Ok(())
241    }
242    inventory::submit!(PyExpressionRegistrar { func: register });
243}
qudit_expr/expressions/kraus.rs

qudit_expr/expressions/
kraus.rs
