Skip to main content

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