qudit_expr/expressions/
complex.rs

1use std::collections::HashMap;
2
3use qudit_core::ComplexScalar;
4use serde::Deserialize;
5use serde::Serialize;
6
7use crate::expressions::Constant;
8use crate::expressions::Expression;
9use crate::qgl::Expression as CiscExpression;
10use crate::qgl::parse_scalar;
11
12#[derive(Clone, PartialEq, Eq, Hash, Deserialize, Serialize)]
13pub struct ComplexExpression {
14    pub real: Expression,
15    pub imag: Expression,
16}
17
18impl ComplexExpression {
19    pub fn from_string(input: impl AsRef<str>) -> Self {
20        ComplexExpression::new(
21            parse_scalar(input.as_ref()).unwrap_or_else(|e| panic!("Invalid input string: {e}")),
22        )
23    }
24
25    pub fn from_real_64(input: f64) -> Self {
26        ComplexExpression {
27            real: Expression::from_float_64(input),
28            imag: Expression::zero(),
29        }
30    }
31
32    pub fn from_real_32(input: f32) -> Self {
33        ComplexExpression {
34            real: Expression::from_float_32(input),
35            imag: Expression::zero(),
36        }
37    }
38
39    pub fn new(cisc_expr: CiscExpression) -> Self {
40        match cisc_expr {
41            CiscExpression::Number(num) => ComplexExpression {
42                real: Expression::Constant(
43                    Constant::from_float(num.parse::<f64>().unwrap()).unwrap(),
44                ),
45                imag: Expression::zero(),
46            },
47            CiscExpression::Variable(var) => {
48                if var == "i" {
49                    ComplexExpression {
50                        real: Expression::zero(),
51                        imag: Expression::one(),
52                    }
53                } else if var == "π" || var == "pi" {
54                    ComplexExpression {
55                        real: Expression::Pi,
56                        imag: Expression::zero(),
57                    }
58                } else {
59                    ComplexExpression {
60                        real: Expression::Variable(var),
61                        imag: Expression::zero(),
62                    }
63                }
64            }
65            CiscExpression::Unary { op, expr } => {
66                let risc_expr = ComplexExpression::new(*expr);
67                match op {
68                    '~' => ComplexExpression {
69                        real: Expression::Neg(Box::new(risc_expr.real)),
70                        imag: Expression::Neg(Box::new(risc_expr.imag)),
71                    },
72                    _ => panic!("Invalid unary operator: {}", op),
73                }
74            }
75            CiscExpression::Binary { op, lhs, rhs } => {
76                let risc_lhs = ComplexExpression::new(*lhs);
77                let risc_rhs = ComplexExpression::new(*rhs);
78                match op {
79                    '+' => risc_lhs + risc_rhs,
80                    '-' => risc_lhs - risc_rhs,
81                    '*' => risc_lhs * risc_rhs,
82                    '/' => risc_lhs / risc_rhs,
83                    '^' => {
84                        if risc_lhs.is_e() {
85                            // checking `is_imag()` requires a non-zero imaginary part
86                            // and therefore rejects an exponent that simplifies to zero
87                            // (e.g. 'i*0')
88                            assert!(
89                                risc_rhs.real.is_zero(),
90                                "Exponential power must be imaginary",
91                            );
92                            ComplexExpression {
93                                real: Expression::Cos(Box::new(risc_rhs.imag.clone())),
94                                imag: Expression::Sin(Box::new(risc_rhs.imag)),
95                            }
96                        } else {
97                            assert!(risc_lhs.is_real(), "Power base must be real");
98                            assert!(risc_rhs.is_real(), "Power exponent must be real");
99                            ComplexExpression {
100                                real: Expression::Pow(
101                                    Box::new(risc_lhs.real),
102                                    Box::new(risc_rhs.real),
103                                ),
104                                imag: Expression::zero(),
105                            }
106                        }
107                    }
108                    _ => panic!("Invalid binary operator: {}", op),
109                }
110            }
111            CiscExpression::Call { fn_name, args } => match fn_name.as_str() {
112                "sqrt" => {
113                    let risc_arg = ComplexExpression::new(args[0].clone());
114                    assert!(args.len() == 1, "sqrt function takes exactly one argument");
115                    assert!(
116                        risc_arg.is_real(),
117                        "sqrt function is only supported for real numbers"
118                    );
119                    ComplexExpression {
120                        real: Expression::Sqrt(Box::new(risc_arg.real)),
121                        imag: Expression::zero(),
122                    }
123                }
124                "sin" => {
125                    let risc_arg = ComplexExpression::new(args[0].clone());
126                    assert!(args.len() == 1, "sin function takes exactly one argument");
127                    assert!(
128                        risc_arg.is_real(),
129                        "sin function is only supported for real numbers"
130                    );
131                    ComplexExpression {
132                        real: Expression::Sin(Box::new(risc_arg.real)),
133                        imag: Expression::zero(),
134                    }
135                }
136                "cos" => {
137                    let risc_arg = ComplexExpression::new(args[0].clone());
138                    assert!(args.len() == 1, "cos function takes exactly one argument");
139                    assert!(
140                        risc_arg.is_real(),
141                        "cos function is only supported for real numbers"
142                    );
143                    ComplexExpression {
144                        real: Expression::Cos(Box::new(risc_arg.real)),
145                        imag: Expression::zero(),
146                    }
147                }
148                "tan" => {
149                    let risc_arg = ComplexExpression::new(args[0].clone());
150                    assert!(args.len() == 1, "tan function takes exactly one argument");
151                    assert!(
152                        risc_arg.is_real(),
153                        "tan function is only supported for real numbers"
154                    );
155                    ComplexExpression {
156                        real: Expression::Div(
157                            Box::new(Expression::Sin(Box::new(risc_arg.real.clone()))),
158                            Box::new(Expression::Cos(Box::new(risc_arg.real))),
159                        ),
160                        imag: Expression::zero(),
161                    }
162                }
163                "csc" => {
164                    let risc_arg = ComplexExpression::new(args[0].clone());
165                    assert!(args.len() == 1, "csc function takes exactly one argument");
166                    assert!(
167                        risc_arg.is_real(),
168                        "csc function is only supported for real numbers"
169                    );
170                    ComplexExpression {
171                        real: Expression::Div(
172                            Box::new(Expression::one()),
173                            Box::new(Expression::Sin(Box::new(risc_arg.real))),
174                        ),
175                        imag: Expression::zero(),
176                    }
177                }
178                "sec" => {
179                    let risc_arg = ComplexExpression::new(args[0].clone());
180                    assert!(args.len() == 1, "sec function takes exactly one argument");
181                    assert!(
182                        risc_arg.is_real(),
183                        "sec function is only supported for real numbers"
184                    );
185                    ComplexExpression {
186                        real: Expression::Div(
187                            Box::new(Expression::one()),
188                            Box::new(Expression::Cos(Box::new(risc_arg.real))),
189                        ),
190                        imag: Expression::zero(),
191                    }
192                }
193                "cot" => {
194                    let risc_arg = ComplexExpression::new(args[0].clone());
195                    assert!(args.len() == 1, "cot function takes exactly one argument");
196                    assert!(
197                        risc_arg.is_real(),
198                        "cot function is only supported for real numbers"
199                    );
200                    ComplexExpression {
201                        real: Expression::Div(
202                            Box::new(Expression::Cos(Box::new(risc_arg.real.clone()))),
203                            Box::new(Expression::Sin(Box::new(risc_arg.real))),
204                        ),
205                        imag: Expression::zero(),
206                    }
207                }
208                _ => panic!("Invalid function name: {}", fn_name),
209            },
210            _ => panic!("Unexpected expression during complex conversion"),
211        }
212    }
213
214    pub fn one() -> Self {
215        ComplexExpression {
216            real: Expression::one(),
217            imag: Expression::zero(),
218        }
219    }
220
221    pub fn zero() -> Self {
222        ComplexExpression {
223            real: Expression::zero(),
224            imag: Expression::zero(),
225        }
226    }
227
228    pub fn i() -> Self {
229        ComplexExpression {
230            real: Expression::zero(),
231            imag: Expression::one(),
232        }
233    }
234
235    pub fn pi() -> Self {
236        ComplexExpression {
237            real: Expression::Pi,
238            imag: Expression::zero(),
239        }
240    }
241
242    pub fn is_real(&self) -> bool {
243        self.imag.is_zero()
244    }
245
246    pub fn is_imag(&self) -> bool {
247        self.real.is_zero() && !self.imag.is_zero()
248    }
249
250    pub fn is_cplx(&self) -> bool {
251        !self.real.is_zero() && !self.imag.is_zero()
252    }
253
254    pub fn is_real_fast(&self) -> bool {
255        self.imag.is_zero_fast()
256    }
257
258    pub fn is_imag_fast(&self) -> bool {
259        self.real.is_zero_fast() && !self.imag.is_zero_fast()
260    }
261
262    pub fn is_cplx_fast(&self) -> bool {
263        !self.real.is_zero_fast() && !self.imag.is_zero_fast()
264    }
265
266    pub fn is_e(&self) -> bool {
267        match &self.real {
268            Expression::Variable(var) => var == "e",
269            _ => false,
270        }
271    }
272
273    pub fn is_zero(&self) -> bool {
274        self.real.is_zero() && self.imag.is_zero()
275    }
276
277    pub fn is_one(&self) -> bool {
278        self.real.is_one() && self.imag.is_zero()
279    }
280
281    pub fn is_zero_fast(&self) -> bool {
282        self.real.is_zero_fast() && self.imag.is_zero_fast()
283    }
284
285    pub fn is_one_fast(&self) -> bool {
286        self.real.is_one_fast() && self.imag.is_zero_fast()
287    }
288
289    pub fn eval<C: ComplexScalar>(&self, args: &HashMap<&str, C::R>) -> C {
290        C::new(self.real.eval(args), self.imag.eval(args))
291    }
292
293    pub fn map_var_names(&self, var_map: &HashMap<String, String>) -> Self {
294        ComplexExpression {
295            real: self.real.map_var_names(var_map),
296            imag: self.imag.map_var_names(var_map),
297        }
298    }
299
300    pub fn conjugate(&self) -> Self {
301        ComplexExpression {
302            real: self.real.clone(),
303            imag: Expression::Neg(Box::new(self.imag.clone())),
304        }
305    }
306
307    pub fn conjugate_in_place(&mut self) {
308        self.imag = Expression::Neg(Box::new(self.imag.clone()));
309    }
310
311    pub fn differentiate(&self, wrt: &str) -> Self {
312        ComplexExpression {
313            real: self.real.differentiate(wrt),
314            imag: self.imag.differentiate(wrt),
315        }
316    }
317
318    pub fn simplify(&self) -> Self {
319        ComplexExpression {
320            real: self.real.simplify(),
321            imag: self.imag.simplify(),
322        }
323    }
324
325    pub fn get_ancestors(&self, variable: &str) -> Vec<Expression> {
326        let mut ancestors = self.real.get_ancestors(variable);
327        let im_ancestors = self.imag.get_ancestors(variable);
328        if ancestors.is_empty() {
329            return im_ancestors;
330        }
331        if !im_ancestors.is_empty() {
332            ancestors.retain(|x| im_ancestors.contains(x));
333        }
334        ancestors
335    }
336
337    pub fn substitute<S: AsRef<Expression>, T: AsRef<Expression>>(
338        &self,
339        original: S,
340        substitution: T,
341    ) -> Self {
342        // TODO: Allow ComplexExpression substitution
343        ComplexExpression {
344            real: self
345                .real
346                .substitute(original.as_ref(), substitution.as_ref()),
347            imag: self
348                .imag
349                .substitute(original.as_ref(), substitution.as_ref()),
350        }
351    }
352
353    pub fn rename_variable<S: AsRef<str>, T: AsRef<str>>(&self, original: S, new: T) -> Self {
354        ComplexExpression {
355            real: self.real.rename_variable(original.as_ref(), new.as_ref()),
356            imag: self.imag.rename_variable(original.as_ref(), new.as_ref()),
357        }
358    }
359
360    pub fn get_unique_variables(&self) -> Vec<String> {
361        let mut real = self.real.get_unique_variables();
362        for i in self.imag.get_unique_variables().into_iter() {
363            if !real.contains(&i) {
364                real.push(i)
365            }
366        }
367        real
368    }
369
370    pub fn is_parameterized(&self) -> bool {
371        self.real.is_parameterized() || self.imag.is_parameterized()
372    }
373}
374
375impl std::ops::Mul<ComplexExpression> for ComplexExpression {
376    type Output = Self;
377
378    fn mul(self, rhs: Self) -> Self {
379        &self * &rhs
380    }
381}
382
383impl std::ops::Mul<&ComplexExpression> for ComplexExpression {
384    type Output = ComplexExpression;
385
386    fn mul(self, rhs: &ComplexExpression) -> ComplexExpression {
387        &self * rhs
388    }
389}
390
391impl std::ops::Mul<ComplexExpression> for &ComplexExpression {
392    type Output = ComplexExpression;
393
394    fn mul(self, rhs: ComplexExpression) -> ComplexExpression {
395        self * &rhs
396    }
397}
398
399impl std::ops::Mul<&ComplexExpression> for &ComplexExpression {
400    type Output = ComplexExpression;
401
402    fn mul(self, rhs: &ComplexExpression) -> ComplexExpression {
403        ComplexExpression {
404            real: &self.real * &rhs.real - &self.imag * &rhs.imag,
405            imag: &self.real * &rhs.imag + &self.imag * &rhs.real,
406        }
407    }
408}
409
410impl std::ops::Add<ComplexExpression> for ComplexExpression {
411    type Output = Self;
412
413    fn add(self, rhs: Self) -> Self {
414        &self + &rhs
415    }
416}
417
418impl std::ops::Add<&ComplexExpression> for ComplexExpression {
419    type Output = ComplexExpression;
420
421    fn add(self, rhs: &ComplexExpression) -> ComplexExpression {
422        &self + rhs
423    }
424}
425
426impl std::ops::Add<ComplexExpression> for &ComplexExpression {
427    type Output = ComplexExpression;
428
429    fn add(self, rhs: ComplexExpression) -> ComplexExpression {
430        self + &rhs
431    }
432}
433
434impl std::ops::Add<&ComplexExpression> for &ComplexExpression {
435    type Output = ComplexExpression;
436
437    fn add(self, rhs: &ComplexExpression) -> ComplexExpression {
438        ComplexExpression {
439            real: &self.real + &rhs.real,
440            imag: &self.imag + &rhs.imag,
441        }
442    }
443}
444
445impl std::ops::Sub<ComplexExpression> for ComplexExpression {
446    type Output = Self;
447
448    fn sub(self, rhs: Self) -> Self {
449        &self - &rhs
450    }
451}
452
453impl std::ops::Sub<&ComplexExpression> for ComplexExpression {
454    type Output = ComplexExpression;
455
456    fn sub(self, rhs: &ComplexExpression) -> ComplexExpression {
457        &self - rhs
458    }
459}
460
461impl std::ops::Sub<ComplexExpression> for &ComplexExpression {
462    type Output = ComplexExpression;
463
464    fn sub(self, rhs: ComplexExpression) -> ComplexExpression {
465        self - &rhs
466    }
467}
468
469impl std::ops::Sub<&ComplexExpression> for &ComplexExpression {
470    type Output = ComplexExpression;
471
472    fn sub(self, rhs: &ComplexExpression) -> ComplexExpression {
473        ComplexExpression {
474            real: &self.real - &rhs.real,
475            imag: &self.imag - &rhs.imag,
476        }
477    }
478}
479
480impl std::ops::Neg for ComplexExpression {
481    type Output = Self;
482
483    fn neg(self) -> Self {
484        -&self
485    }
486}
487
488impl std::ops::Neg for &ComplexExpression {
489    type Output = ComplexExpression;
490
491    fn neg(self) -> ComplexExpression {
492        ComplexExpression {
493            real: -&self.real,
494            imag: -&self.imag,
495        }
496    }
497}
498
499impl std::ops::Div<ComplexExpression> for ComplexExpression {
500    type Output = Self;
501
502    fn div(self, rhs: Self) -> Self {
503        &self / &rhs
504    }
505}
506
507impl std::ops::Div<&ComplexExpression> for ComplexExpression {
508    type Output = ComplexExpression;
509
510    fn div(self, rhs: &ComplexExpression) -> ComplexExpression {
511        &self / rhs
512    }
513}
514
515impl std::ops::Div<ComplexExpression> for &ComplexExpression {
516    type Output = ComplexExpression;
517
518    fn div(self, rhs: ComplexExpression) -> ComplexExpression {
519        self / &rhs
520    }
521}
522
523impl std::ops::Div<&ComplexExpression> for &ComplexExpression {
524    type Output = ComplexExpression;
525
526    fn div(self, rhs: &ComplexExpression) -> ComplexExpression {
527        let dem = &rhs.real * &rhs.real + &rhs.imag * &rhs.imag;
528        ComplexExpression {
529            real: (&self.real * &rhs.real + &self.imag * &rhs.imag) / &dem,
530            imag: (&self.imag * &rhs.real - &self.real * &rhs.imag) / &dem,
531        }
532    }
533}
534
535impl std::ops::AddAssign for ComplexExpression {
536    fn add_assign(&mut self, rhs: Self) {
537        *self = &*self + &rhs;
538    }
539}
540
541impl std::ops::AddAssign<&ComplexExpression> for ComplexExpression {
542    fn add_assign(&mut self, rhs: &Self) {
543        *self = &*self + rhs;
544    }
545}
546
547impl std::ops::SubAssign for ComplexExpression {
548    fn sub_assign(&mut self, rhs: Self) {
549        *self = &*self - &rhs;
550    }
551}
552
553impl std::ops::SubAssign<&ComplexExpression> for ComplexExpression {
554    fn sub_assign(&mut self, rhs: &Self) {
555        *self = &*self - rhs;
556    }
557}
558
559impl std::ops::MulAssign for ComplexExpression {
560    fn mul_assign(&mut self, rhs: Self) {
561        *self = &*self * &rhs;
562    }
563}
564
565impl std::ops::MulAssign<&ComplexExpression> for ComplexExpression {
566    fn mul_assign(&mut self, rhs: &Self) {
567        *self = &*self * rhs;
568    }
569}
570
571impl std::ops::DivAssign for ComplexExpression {
572    fn div_assign(&mut self, rhs: Self) {
573        *self = &*self / &rhs;
574    }
575}
576
577impl std::ops::DivAssign<&ComplexExpression> for ComplexExpression {
578    fn div_assign(&mut self, rhs: &Self) {
579        *self = &*self / rhs;
580    }
581}
582
583impl std::fmt::Debug for ComplexExpression {
584    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
585        f.debug_struct("ComplexExpression")
586            .field("real", &self.real)
587            .field("imag", &self.imag)
588            .finish()
589    }
590}
591
592impl<C: ComplexScalar> From<C> for ComplexExpression {
593    fn from(value: C) -> Self {
594        ComplexExpression {
595            real: value.re().into(),
596            imag: value.im().into(),
597        }
598    }
599}
qudit_expr/expressions/complex.rs

qudit_expr/expressions/
complex.rs
