qudit_expr/expressions/
named.rs

1use itertools::Itertools;
2use qudit_core::RealScalar;
3use serde::{Deserialize, Serialize};
4use std::{
5    collections::HashMap,
6    ops::{Deref, DerefMut},
7};
8
9use crate::ComplexExpression;
10
11#[derive(PartialEq, Eq, Debug, Clone, Serialize, Deserialize)]
12pub struct ExpressionBody {
13    body: Vec<ComplexExpression>,
14}
15
16impl ExpressionBody {
17    pub fn new(body: Vec<ComplexExpression>) -> Self {
18        Self { body }
19    }
20
21    pub fn num_elements(&self) -> usize {
22        self.body.len()
23    }
24
25    pub fn elements(&self) -> &[ComplexExpression] {
26        &self.body
27    }
28
29    pub fn elements_mut(&mut self) -> &mut [ComplexExpression] {
30        &mut self.body
31    }
32
33    pub fn conjugate(&mut self) {
34        for expr in self.body.iter_mut() {
35            expr.conjugate_in_place()
36        }
37    }
38
39    pub fn apply_element_permutation(&mut self, elem_perm: &[usize]) {
40        // TODO: do physical element permutation in place via transpositions
41        let mut swap_vec = vec![];
42        std::mem::swap(&mut swap_vec, &mut self.body);
43        self.body = swap_vec
44            .into_iter()
45            .enumerate()
46            .sorted_by(|(old_idx_a, _), (old_idx_b, _)| {
47                elem_perm[*old_idx_a].cmp(&elem_perm[*old_idx_b])
48            })
49            .map(|(_, expr)| expr)
50            .collect();
51    }
52
53    pub fn rename(&mut self, var_map: HashMap<String, String>) {
54        for elem in self.body.iter_mut() {
55            *elem = elem.map_var_names(&var_map);
56        }
57    }
58}
59
60impl From<Vec<ComplexExpression>> for ExpressionBody {
61    fn from(value: Vec<ComplexExpression>) -> Self {
62        ExpressionBody::new(value)
63    }
64}
65
66impl From<ExpressionBody> for Vec<ComplexExpression> {
67    fn from(value: ExpressionBody) -> Self {
68        value.body
69    }
70}
71
72impl AsRef<[ComplexExpression]> for ExpressionBody {
73    fn as_ref(&self) -> &[ComplexExpression] {
74        self.elements()
75    }
76}
77
78impl Deref for ExpressionBody {
79    type Target = [ComplexExpression];
80
81    fn deref(&self) -> &Self::Target {
82        &self.body
83    }
84}
85
86impl DerefMut for ExpressionBody {
87    fn deref_mut(&mut self) -> &mut Self::Target {
88        &mut self.body
89    }
90}
91
92#[derive(PartialEq, Eq, Debug, Clone, Serialize, Deserialize)]
93pub struct BoundExpressionBody {
94    variables: Vec<String>,
95    body: ExpressionBody,
96}
97
98impl BoundExpressionBody {
99    pub fn new<B: Into<ExpressionBody>>(variables: Vec<String>, body: B) -> Self {
100        Self {
101            variables,
102            body: body.into(),
103        }
104    }
105
106    pub fn num_params(&self) -> usize {
107        self.variables.len()
108    }
109
110    pub fn variables(&self) -> &[String] {
111        &self.variables
112    }
113
114    pub fn set_variables(&mut self, new_variables: Vec<String>) {
115        self.variables = new_variables;
116    }
117
118    pub fn num_elements(&self) -> usize {
119        self.body.num_elements()
120    }
121
122    pub fn elements(&self) -> &[ComplexExpression] {
123        self.body.elements()
124    }
125
126    pub fn elements_mut(&mut self) -> &mut [ComplexExpression] {
127        self.body.elements_mut()
128    }
129
130    pub fn conjugate(&mut self) {
131        self.body.conjugate()
132    }
133
134    pub fn destruct(self) -> (Vec<String>, Vec<ComplexExpression>) {
135        let Self { variables, body } = self;
136        (variables, body.into())
137    }
138
139    pub fn apply_element_permutation(&mut self, elem_perm: &[usize]) {
140        self.body.apply_element_permutation(elem_perm);
141    }
142
143    pub fn alpha_rename(&mut self, starting_number: Option<usize>) {
144        let start = starting_number.unwrap_or_default();
145
146        let mut var_map = HashMap::new();
147        let mut new_vars = Vec::new();
148
149        for (var_id, var) in (start..).zip(self.variables()) {
150            let new_var_name = format!("alpha_{}", var_id);
151            new_vars.push(new_var_name.clone());
152            var_map.insert(var.clone(), new_var_name);
153        }
154
155        self.body.rename(var_map);
156        self.variables = new_vars;
157    }
158
159    pub fn get_arg_map<R: RealScalar>(&self, args: &[R]) -> HashMap<&str, R> {
160        self.variables()
161            .iter()
162            .zip(args.iter())
163            .map(|(a, b)| (a.as_str(), *b))
164            .collect()
165    }
166}
167
168impl AsRef<[ComplexExpression]> for BoundExpressionBody {
169    fn as_ref(&self) -> &[ComplexExpression] {
170        self.elements()
171    }
172}
173
174impl Deref for BoundExpressionBody {
175    type Target = ExpressionBody;
176
177    fn deref(&self) -> &Self::Target {
178        &self.body
179    }
180}
181
182impl DerefMut for BoundExpressionBody {
183    fn deref_mut(&mut self) -> &mut Self::Target {
184        &mut self.body
185    }
186}
187
188#[derive(Debug, Clone, Serialize, Deserialize)]
189pub struct NamedExpression {
190    name: String,
191    body: BoundExpressionBody,
192}
193
194impl NamedExpression {
195    pub fn new<S: Into<String>, B: Into<ExpressionBody>>(
196        name: S,
197        variables: Vec<String>,
198        body: B,
199    ) -> Self {
200        Self {
201            name: name.into(),
202            body: BoundExpressionBody::new(variables, body),
203        }
204    }
205
206    pub fn from_body_with_name(name: String, body: BoundExpressionBody) -> Self {
207        Self { name, body }
208    }
209
210    pub fn name(&self) -> &str {
211        &self.name
212    }
213
214    pub fn set_name(&mut self, new_name: impl Into<String>) {
215        self.name = new_name.into();
216    }
217
218    pub fn set_variables(&mut self, new_variables: Vec<String>) {
219        self.body.set_variables(new_variables);
220    }
221
222    pub fn destruct(self) -> (String, Vec<String>, Vec<ComplexExpression>) {
223        let Self { name, body } = self;
224        let (variables, body) = body.destruct();
225        (name, variables, body)
226    }
227
228    pub fn apply_element_permutation(&mut self, elem_perm: &[usize]) {
229        self.body.apply_element_permutation(elem_perm);
230    }
231}
232
233impl AsRef<[ComplexExpression]> for NamedExpression {
234    fn as_ref(&self) -> &[ComplexExpression] {
235        self.elements()
236    }
237}
238
239impl AsRef<BoundExpressionBody> for NamedExpression {
240    fn as_ref(&self) -> &BoundExpressionBody {
241        &self.body
242    }
243}
244
245impl<B: AsRef<[ComplexExpression]>> PartialEq<B> for NamedExpression {
246    fn eq(&self, other: &B) -> bool {
247        self.elements() == other.as_ref()
248    }
249}
250
251impl Eq for NamedExpression {}
252
253impl Deref for NamedExpression {
254    type Target = BoundExpressionBody;
255
256    fn deref(&self) -> &Self::Target {
257        &self.body
258    }
259}
260
261impl DerefMut for NamedExpression {
262    fn deref_mut(&mut self) -> &mut Self::Target {
263        &mut self.body
264    }
265}
qudit_expr/expressions/named.rs

qudit_expr/expressions/
named.rs
