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 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}