Skip to main content

qudit_expr/expressions/
ket.rs

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