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

qudit_expr/expressions/
brasys.rs
