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

qudit_expr/expressions/
utrysys.rs
