qudit_expr/expressions/
base.rs

1use num::FromPrimitive;
2use num::ToPrimitive;
3use num::bigint::BigInt;
4use num::rational::Ratio;
5use qudit_core::RealScalar;
6use serde::{Deserialize, Serialize};
7use std::collections::HashMap;
8use std::collections::HashSet;
9
10use crate::analysis::simplify;
11
12pub type Rational = Ratio<BigInt>;
13pub type Constant = Rational;
14
15#[derive(Clone, Serialize, Deserialize)]
16pub enum Expression {
17    Pi,
18    Variable(String),
19    Constant(Constant),
20    Neg(Box<Expression>),
21    Add(Box<Expression>, Box<Expression>),
22    Sub(Box<Expression>, Box<Expression>),
23    Mul(Box<Expression>, Box<Expression>),
24    Div(Box<Expression>, Box<Expression>),
25    Pow(Box<Expression>, Box<Expression>),
26    Sqrt(Box<Expression>),
27    Sin(Box<Expression>),
28    Cos(Box<Expression>),
29}
30
31impl Expression {
32    pub fn zero() -> Self {
33        Expression::Constant(Constant::new(BigInt::from(0), BigInt::from(1)))
34    }
35
36    pub fn one() -> Self {
37        Expression::Constant(Constant::new(BigInt::from(1), BigInt::from(1)))
38    }
39
40    pub fn from_int(n: i64) -> Self {
41        Expression::Constant(Constant::new(BigInt::from(n), BigInt::from(1)))
42    }
43
44    pub fn from_float(f: f64) -> Self {
45        Self::from_float_64(f)
46    }
47
48    pub fn from_float_32(f: f32) -> Self {
49        Expression::Constant(Constant::from_f32(f).unwrap())
50    }
51
52    pub fn from_float_64(f: f64) -> Self {
53        Expression::Constant(Constant::from_f64(f).unwrap())
54    }
55
56    pub fn to_float(&self) -> f64 {
57        match self {
58            Expression::Constant(c) => c.to_f64().unwrap(),
59            Expression::Variable(_) => panic!("Cannot convert variable to float"),
60            Expression::Pi => std::f64::consts::PI,
61            Expression::Neg(expr) => -expr.to_float(),
62            Expression::Add(lhs, rhs) => lhs.to_float() + rhs.to_float(),
63            Expression::Sub(lhs, rhs) => lhs.to_float() - rhs.to_float(),
64            Expression::Mul(lhs, rhs) => lhs.to_float() * rhs.to_float(),
65            Expression::Div(lhs, rhs) => lhs.to_float() / rhs.to_float(),
66            Expression::Pow(lhs, rhs) => lhs.to_float().powf(rhs.to_float()),
67            Expression::Sqrt(expr) => expr.to_float().sqrt(),
68            Expression::Sin(expr) => expr.to_float().sin(),
69            Expression::Cos(expr) => expr.to_float().cos(),
70        }
71    }
72
73    pub fn to_constant(&self) -> Constant {
74        // TODO: Figure out how to maintain precision by doing math over Constants
75        Constant::from_float(self.to_float()).unwrap()
76    }
77
78    pub fn gather_context(&self) -> HashSet<String> {
79        let mut context = HashSet::new();
80        context.insert(self.to_string());
81        match self {
82            Expression::Pi => {
83                context.insert("pi".to_string());
84            }
85            Expression::Variable(var) => {
86                context.insert(var.clone());
87            }
88            Expression::Constant(_) => {
89                context.insert(self.to_string());
90                context.insert(self.to_float().to_string());
91            }
92            Expression::Neg(expr) => {
93                context.extend(expr.gather_context());
94            }
95            Expression::Add(lhs, rhs) => {
96                context.extend(lhs.gather_context());
97                context.extend(rhs.gather_context());
98            }
99            Expression::Sub(lhs, rhs) => {
100                context.extend(lhs.gather_context());
101                context.extend(rhs.gather_context());
102            }
103            Expression::Mul(lhs, rhs) => {
104                context.extend(lhs.gather_context());
105                context.extend(rhs.gather_context());
106            }
107            Expression::Div(lhs, rhs) => {
108                context.extend(lhs.gather_context());
109                context.extend(rhs.gather_context());
110            }
111            Expression::Pow(lhs, rhs) => {
112                context.extend(lhs.gather_context());
113                context.extend(rhs.gather_context());
114            }
115            Expression::Sqrt(expr) => {
116                context.extend(expr.gather_context());
117            }
118            Expression::Sin(expr) => {
119                context.extend(expr.gather_context());
120            }
121            Expression::Cos(expr) => {
122                context.extend(expr.gather_context());
123            }
124        }
125        context
126    }
127
128    pub fn is_zero(&self) -> bool {
129        match self {
130            Expression::Constant(c) => *c.numer() == BigInt::from(0),
131            Expression::Neg(expr) => expr.is_zero(),
132            Expression::Add(lhs, rhs) => lhs.is_zero() && rhs.is_zero(),
133            Expression::Sub(lhs, rhs) => (lhs.is_zero() && rhs.is_zero()) || lhs == rhs,
134            Expression::Mul(lhs, rhs) => lhs.is_zero() || rhs.is_zero(),
135            Expression::Div(lhs, _) => lhs.is_zero(),
136            Expression::Pow(lhs, rhs) => lhs.is_zero() && !rhs.is_zero(),
137            Expression::Sqrt(expr) => expr.is_zero(),
138            Expression::Sin(expr) => expr.is_zero(),
139            Expression::Cos(expr) => {
140                !expr.is_parameterized()
141                    && (expr.eval::<f64>(&HashMap::new()) - std::f64::consts::PI / 2.0) < 1e-6
142            }
143            Expression::Pi => false,
144            Expression::Variable(_) => false,
145        }
146    }
147
148    /// Conservative check for zero. This is faster than the exact check.
149    pub fn is_zero_fast(&self) -> bool {
150        match self {
151            Expression::Constant(c) => *c.numer() == BigInt::from(0),
152            Expression::Neg(expr) => expr.is_zero_fast(),
153            Expression::Add(lhs, rhs) => lhs.is_zero_fast() && rhs.is_zero_fast(),
154            Expression::Sub(lhs, rhs) => lhs.is_zero_fast() && rhs.is_zero_fast(),
155            Expression::Mul(lhs, rhs) => lhs.is_zero_fast() || rhs.is_zero_fast(),
156            Expression::Div(lhs, _) => lhs.is_zero_fast(),
157            Expression::Pow(lhs, rhs) => lhs.is_zero_fast() && !rhs.is_zero_fast(),
158            Expression::Sqrt(expr) => expr.is_zero_fast(),
159            Expression::Sin(expr) => expr.is_zero_fast(),
160            Expression::Cos(_expr) => false,
161            Expression::Pi => false,
162            Expression::Variable(_) => false,
163        }
164    }
165
166    pub fn is_one(&self) -> bool {
167        match self {
168            Expression::Constant(c) => *c.numer() == *c.denom(),
169            Expression::Neg(expr) => {
170                !expr.is_parameterized() && expr.eval::<f64>(&HashMap::new()) == -1.0
171            }
172            Expression::Add(lhs, rhs) => {
173                lhs.is_one() && rhs.is_zero() || lhs.is_zero() && rhs.is_one()
174            }
175            Expression::Sub(lhs, rhs) => lhs.is_one() && rhs.is_zero(),
176            Expression::Mul(lhs, rhs) => lhs.is_one() && rhs.is_one(),
177            Expression::Div(lhs, rhs) => lhs == rhs && !rhs.is_zero(),
178            Expression::Pow(lhs, _rhs) => lhs.is_one(),
179            Expression::Sqrt(expr) => expr.is_one(),
180            Expression::Sin(expr) => {
181                !expr.is_parameterized()
182                    && (expr.eval::<f64>(&HashMap::new()) - std::f64::consts::PI / 2.0) < 1e-6
183            }
184            Expression::Cos(expr) => expr.is_zero(),
185            Expression::Pi => false,
186            Expression::Variable(_) => false,
187        }
188    }
189
190    pub fn is_one_fast(&self) -> bool {
191        match self {
192            Expression::Constant(c) => *c.numer() == *c.denom(),
193            Expression::Neg(_expr) => false,
194            Expression::Add(lhs, rhs) => {
195                lhs.is_one_fast() && rhs.is_zero_fast() || lhs.is_zero_fast() && rhs.is_one_fast()
196            }
197            Expression::Sub(lhs, rhs) => lhs.is_one_fast() && rhs.is_zero_fast(),
198            Expression::Mul(lhs, rhs) => lhs.is_one_fast() && rhs.is_one_fast(),
199            Expression::Div(lhs, rhs) => lhs.is_one_fast() && rhs.is_one_fast(),
200            Expression::Pow(lhs, _rhs) => lhs.is_one_fast(),
201            Expression::Sqrt(expr) => expr.is_one_fast(),
202            Expression::Sin(_expr) => false,
203            Expression::Cos(expr) => expr.is_zero_fast(),
204            Expression::Pi => false,
205            Expression::Variable(_) => false,
206        }
207    }
208
209    pub fn contains_variable<T: AsRef<str>>(&self, var: T) -> bool {
210        let var = var.as_ref();
211        match self {
212            Expression::Pi => false,
213            Expression::Variable(v) => v == var,
214            Expression::Constant(_) => false,
215            Expression::Neg(expr) => expr.contains_variable(var),
216            Expression::Add(lhs, rhs) => lhs.contains_variable(var) || rhs.contains_variable(var),
217            Expression::Sub(lhs, rhs) => lhs.contains_variable(var) || rhs.contains_variable(var),
218            Expression::Mul(lhs, rhs) => lhs.contains_variable(var) || rhs.contains_variable(var),
219            Expression::Div(lhs, rhs) => lhs.contains_variable(var) || rhs.contains_variable(var),
220            Expression::Pow(lhs, rhs) => lhs.contains_variable(var) || rhs.contains_variable(var),
221            Expression::Sqrt(expr) => expr.contains_variable(var),
222            Expression::Sin(expr) => expr.contains_variable(var),
223            Expression::Cos(expr) => expr.contains_variable(var),
224        }
225    }
226
227    pub fn is_parameterized(&self) -> bool {
228        match self {
229            Expression::Pi => false,
230            Expression::Variable(_) => true,
231            Expression::Constant(_) => false,
232            Expression::Neg(expr) => expr.is_parameterized(),
233            Expression::Add(lhs, rhs) => lhs.is_parameterized() || rhs.is_parameterized(),
234            Expression::Sub(lhs, rhs) => lhs.is_parameterized() || rhs.is_parameterized(),
235            Expression::Mul(lhs, rhs) => lhs.is_parameterized() || rhs.is_parameterized(),
236            Expression::Div(lhs, rhs) => lhs.is_parameterized() || rhs.is_parameterized(),
237            Expression::Pow(lhs, rhs) => lhs.is_parameterized() || rhs.is_parameterized(),
238            Expression::Sqrt(expr) => expr.is_parameterized(),
239            Expression::Sin(expr) => expr.is_parameterized(),
240            Expression::Cos(expr) => expr.is_parameterized(),
241        }
242    }
243
244    pub fn eval<R: RealScalar>(&self, args: &HashMap<&str, R>) -> R {
245        match self {
246            Expression::Pi => R::PI(),
247            Expression::Variable(var) => {
248                if let Some(val) = args.get(var.as_str()) {
249                    *val
250                } else {
251                    panic!("Variable {} not found in arguments", var)
252                }
253            }
254            Expression::Constant(c) => R::from_ratio(c.clone()).unwrap(),
255            Expression::Neg(expr) => -expr.eval(args),
256            Expression::Add(lhs, rhs) => lhs.eval(args) + rhs.eval(args),
257            Expression::Sub(lhs, rhs) => lhs.eval(args) - rhs.eval(args),
258            Expression::Mul(lhs, rhs) => lhs.eval(args) * rhs.eval(args),
259            Expression::Div(lhs, rhs) => lhs.eval(args) / rhs.eval(args),
260            Expression::Pow(lhs, rhs) => lhs.eval(args).powf(rhs.eval(args)),
261            Expression::Sqrt(expr) => expr.eval(args).sqrt(),
262            Expression::Sin(expr) => expr.eval(args).sin(),
263            Expression::Cos(expr) => expr.eval(args).cos(),
264        }
265    }
266
267    /// Uses a magic value to evaluate the expression. This is useful for hashing expressions.
268    pub fn hash_eval(&self) -> f64 {
269        let val = match self {
270            Expression::Pi => self.to_float(),
271            Expression::Variable(_) => 1.7,
272            Expression::Constant(_) => self.to_float(),
273            Expression::Neg(expr) => -expr.hash_eval(),
274            Expression::Add(lhs, rhs) => lhs.hash_eval() + rhs.hash_eval(),
275            Expression::Sub(lhs, rhs) => lhs.hash_eval() - rhs.hash_eval(),
276            Expression::Mul(lhs, rhs) => lhs.hash_eval() * rhs.hash_eval(),
277            Expression::Div(lhs, rhs) => lhs.hash_eval() / rhs.hash_eval(),
278            Expression::Pow(lhs, rhs) => lhs.hash_eval().powf(rhs.hash_eval()),
279            Expression::Sqrt(expr) => expr.hash_eval().sqrt(),
280            Expression::Sin(expr) => expr.hash_eval().sin(),
281            Expression::Cos(expr) => expr.hash_eval().cos(),
282        };
283
284        if val.is_nan() || val.is_subnormal() {
285            0.0
286        } else {
287            val
288        }
289    }
290
291    pub fn map_var_names(&self, var_map: &HashMap<String, String>) -> Self {
292        match self {
293            Expression::Pi => Expression::Pi,
294            Expression::Variable(var) => {
295                if let Some(new_var) = var_map.get(var.as_str()) {
296                    Expression::Variable(new_var.to_string())
297                } else {
298                    Expression::Variable(var.clone())
299                }
300            }
301            Expression::Constant(c) => Expression::Constant(c.clone()),
302            Expression::Neg(expr) => Expression::Neg(Box::new(expr.map_var_names(var_map))),
303            Expression::Add(lhs, rhs) => Expression::Add(
304                Box::new(lhs.map_var_names(var_map)),
305                Box::new(rhs.map_var_names(var_map)),
306            ),
307            Expression::Sub(lhs, rhs) => Expression::Sub(
308                Box::new(lhs.map_var_names(var_map)),
309                Box::new(rhs.map_var_names(var_map)),
310            ),
311            Expression::Mul(lhs, rhs) => Expression::Mul(
312                Box::new(lhs.map_var_names(var_map)),
313                Box::new(rhs.map_var_names(var_map)),
314            ),
315            Expression::Div(lhs, rhs) => Expression::Div(
316                Box::new(lhs.map_var_names(var_map)),
317                Box::new(rhs.map_var_names(var_map)),
318            ),
319            Expression::Pow(lhs, rhs) => Expression::Pow(
320                Box::new(lhs.map_var_names(var_map)),
321                Box::new(rhs.map_var_names(var_map)),
322            ),
323            Expression::Sqrt(expr) => Expression::Sqrt(Box::new(expr.map_var_names(var_map))),
324            Expression::Sin(expr) => Expression::Sin(Box::new(expr.map_var_names(var_map))),
325            Expression::Cos(expr) => Expression::Cos(Box::new(expr.map_var_names(var_map))),
326        }
327    }
328
329    pub fn rename_variable<S: AsRef<str>, T: AsRef<str>>(&self, original: S, new: T) -> Self {
330        let original = original.as_ref();
331        let new = new.as_ref();
332        match self {
333            Expression::Pi => Expression::Pi,
334            Expression::Variable(var) => {
335                if var == original {
336                    Expression::Variable(new.to_string())
337                } else {
338                    Expression::Variable(var.clone())
339                }
340            }
341            Expression::Constant(c) => Expression::Constant(c.clone()),
342            Expression::Neg(expr) => Expression::Neg(Box::new(expr.rename_variable(original, new))),
343            Expression::Add(lhs, rhs) => Expression::Add(
344                Box::new(lhs.rename_variable(original, new)),
345                Box::new(rhs.rename_variable(original, new)),
346            ),
347            Expression::Sub(lhs, rhs) => Expression::Sub(
348                Box::new(lhs.rename_variable(original, new)),
349                Box::new(rhs.rename_variable(original, new)),
350            ),
351            Expression::Mul(lhs, rhs) => Expression::Mul(
352                Box::new(lhs.rename_variable(original, new)),
353                Box::new(rhs.rename_variable(original, new)),
354            ),
355            Expression::Div(lhs, rhs) => Expression::Div(
356                Box::new(lhs.rename_variable(original, new)),
357                Box::new(rhs.rename_variable(original, new)),
358            ),
359            Expression::Pow(lhs, rhs) => Expression::Pow(
360                Box::new(lhs.rename_variable(original, new)),
361                Box::new(rhs.rename_variable(original, new)),
362            ),
363            Expression::Sqrt(expr) => {
364                Expression::Sqrt(Box::new(expr.rename_variable(original, new)))
365            }
366            Expression::Sin(expr) => Expression::Sin(Box::new(expr.rename_variable(original, new))),
367            Expression::Cos(expr) => Expression::Cos(Box::new(expr.rename_variable(original, new))),
368        }
369    }
370
371    pub fn differentiate<S: AsRef<str>>(&self, wrt: S) -> Self {
372        let wrt = wrt.as_ref();
373        match self {
374            Expression::Pi => Expression::zero(),
375            Expression::Variable(var) => {
376                if var == wrt {
377                    Expression::one()
378                } else {
379                    Expression::zero()
380                }
381            }
382            Expression::Constant(_) => Expression::zero(),
383            Expression::Neg(expr) => Expression::Neg(Box::new(expr.differentiate(wrt))),
384            Expression::Add(lhs, rhs) => Expression::Add(
385                Box::new(lhs.differentiate(wrt)),
386                Box::new(rhs.differentiate(wrt)),
387            ),
388            Expression::Sub(lhs, rhs) => Expression::Sub(
389                Box::new(lhs.differentiate(wrt)),
390                Box::new(rhs.differentiate(wrt)),
391            ),
392            Expression::Mul(lhs, rhs) => {
393                lhs.differentiate(wrt) * *rhs.clone() + *lhs.clone() * rhs.differentiate(wrt)
394            }
395            Expression::Div(lhs, rhs) => {
396                (lhs.differentiate(wrt) * *rhs.clone() - *lhs.clone() * rhs.differentiate(wrt))
397                    / (*rhs.clone() * *rhs.clone())
398            }
399            Expression::Pow(lhs, rhs) => {
400                let base_fn_x = lhs.contains_variable(wrt);
401                let exponent_fn_x = rhs.contains_variable(wrt);
402
403                if !base_fn_x && !exponent_fn_x {
404                    Expression::zero()
405                } else if !base_fn_x && exponent_fn_x {
406                    if lhs.is_parameterized() {
407                        todo!(
408                            "Cannot differentiate with respect to a parameterized power base until ln is implemented"
409                        )
410                    } else {
411                        self.clone()
412                            * rhs.differentiate(wrt)
413                            * Expression::from_float(lhs.eval::<f64>(&HashMap::new()).ln())
414                    }
415                } else if base_fn_x && !exponent_fn_x {
416                    *rhs.clone()
417                        * Expression::Pow(
418                            Box::new(*lhs.clone()),
419                            Box::new(*rhs.clone() - Expression::one()),
420                        )
421                        * lhs.differentiate(wrt)
422                } else {
423                    todo!(
424                        "Cannot differentiate with respect to a parameterized base and exponent until ln is implemented"
425                    )
426                }
427            }
428            Expression::Sqrt(expr) => {
429                let two = Expression::from_int(2);
430                (Expression::one() / (two * self.clone())) * expr.differentiate(wrt)
431            }
432            Expression::Sin(expr) => {
433                Expression::Cos(Box::new(*expr.clone())) * expr.differentiate(wrt)
434            }
435            Expression::Cos(expr) => {
436                Expression::Neg(Box::new(Expression::Sin(Box::new(*expr.clone()))))
437                    * expr.differentiate(wrt)
438            }
439        }
440    }
441
442    pub fn get_ancestors<S: AsRef<str>>(&self, variable: S) -> Vec<Expression> {
443        let variable = variable.as_ref();
444        let mut ancestors = Vec::new();
445        match self {
446            Expression::Pi => {}
447            Expression::Variable(var) => {
448                if var == variable {
449                    ancestors.push(self.clone());
450                }
451            }
452            Expression::Constant(_) => {}
453            Expression::Neg(expr) => {
454                let node_ancsestors = expr.get_ancestors(variable);
455                let is_empty = node_ancsestors.is_empty();
456                ancestors.extend(node_ancsestors);
457                if !is_empty {
458                    ancestors.push(self.clone());
459                }
460            }
461            Expression::Add(lhs, rhs) => {
462                let lhs_ancestors = lhs.get_ancestors(variable);
463                let rhs_ancestors = rhs.get_ancestors(variable);
464                let is_empty = lhs_ancestors.is_empty() && rhs_ancestors.is_empty();
465                ancestors.extend(lhs_ancestors);
466                ancestors.extend(rhs_ancestors);
467                if !is_empty {
468                    ancestors.push(self.clone());
469                }
470            }
471            Expression::Sub(lhs, rhs) => {
472                let lhs_ancestors = lhs.get_ancestors(variable);
473                let rhs_ancestors = rhs.get_ancestors(variable);
474                let is_empty = lhs_ancestors.is_empty() && rhs_ancestors.is_empty();
475                ancestors.extend(lhs_ancestors);
476                ancestors.extend(rhs_ancestors);
477                if !is_empty {
478                    ancestors.push(self.clone());
479                }
480            }
481            Expression::Mul(lhs, rhs) => {
482                let lhs_ancestors = lhs.get_ancestors(variable);
483                let rhs_ancestors = rhs.get_ancestors(variable);
484                let is_empty = lhs_ancestors.is_empty() && rhs_ancestors.is_empty();
485                ancestors.extend(lhs_ancestors);
486                ancestors.extend(rhs_ancestors);
487                if !is_empty {
488                    ancestors.push(self.clone());
489                }
490            }
491            Expression::Div(lhs, rhs) => {
492                let lhs_ancestors = lhs.get_ancestors(variable);
493                let rhs_ancestors = rhs.get_ancestors(variable);
494                let is_empty = lhs_ancestors.is_empty() && rhs_ancestors.is_empty();
495                ancestors.extend(lhs_ancestors);
496                ancestors.extend(rhs_ancestors);
497                if !is_empty {
498                    ancestors.push(self.clone());
499                }
500            }
501            Expression::Pow(lhs, rhs) => {
502                let lhs_ancestors = lhs.get_ancestors(variable);
503                let rhs_ancestors = rhs.get_ancestors(variable);
504                let is_empty = lhs_ancestors.is_empty() && rhs_ancestors.is_empty();
505                ancestors.extend(lhs_ancestors);
506                ancestors.extend(rhs_ancestors);
507                if !is_empty {
508                    ancestors.push(self.clone());
509                }
510            }
511            Expression::Sqrt(expr) => {
512                let node_ancsestors = expr.get_ancestors(variable);
513                let is_empty = node_ancsestors.is_empty();
514                ancestors.extend(node_ancsestors);
515                if !is_empty {
516                    ancestors.push(self.clone());
517                }
518            }
519            Expression::Sin(expr) => {
520                let node_ancsestors = expr.get_ancestors(variable);
521                let is_empty = node_ancsestors.is_empty();
522                ancestors.extend(node_ancsestors);
523                if !is_empty {
524                    ancestors.push(self.clone());
525                }
526            }
527            Expression::Cos(expr) => {
528                let node_ancsestors = expr.get_ancestors(variable);
529                let is_empty = node_ancsestors.is_empty();
530                ancestors.extend(node_ancsestors);
531                if !is_empty {
532                    ancestors.push(self.clone());
533                }
534            }
535        }
536        ancestors
537    }
538
539    pub fn fast_eq(&self, other: &Expression) -> bool {
540        match (self, other) {
541            (Expression::Pi, Expression::Pi) => true,
542            (Expression::Variable(var1), Expression::Variable(var2)) => var1 == var2,
543            (Expression::Constant(c1), Expression::Constant(c2)) => c1 == c2,
544            (Expression::Neg(expr1), Expression::Neg(expr2)) => expr1.fast_eq(expr2),
545            (Expression::Add(lhs1, rhs1), Expression::Add(lhs2, rhs2)) => {
546                (lhs1.fast_eq(lhs2) && rhs1.fast_eq(rhs2))
547                    || (lhs1.fast_eq(rhs2) && rhs1.fast_eq(lhs2))
548            }
549            (Expression::Sub(lhs1, rhs1), Expression::Sub(lhs2, rhs2)) => {
550                (lhs1.fast_eq(lhs2) && rhs1.fast_eq(rhs2))
551                    || (lhs1.fast_eq(rhs2) && rhs1.fast_eq(lhs2))
552            }
553            (Expression::Mul(lhs1, rhs1), Expression::Mul(lhs2, rhs2)) => {
554                (lhs1.fast_eq(lhs2) && rhs1.fast_eq(rhs2))
555                    || (lhs1.fast_eq(rhs2) && rhs1.fast_eq(lhs2))
556            }
557            (Expression::Div(lhs1, rhs1), Expression::Div(lhs2, rhs2)) => {
558                (lhs1.fast_eq(lhs2) && rhs1.fast_eq(rhs2))
559                    || (lhs1.fast_eq(rhs2) && rhs1.fast_eq(lhs2))
560            }
561            (Expression::Pow(lhs1, rhs1), Expression::Pow(lhs2, rhs2)) => {
562                (lhs1.fast_eq(lhs2) && rhs1.fast_eq(rhs2))
563                    || (lhs1.fast_eq(rhs2) && rhs1.fast_eq(lhs2))
564            }
565            (Expression::Sqrt(expr1), Expression::Sqrt(expr2)) => expr1.fast_eq(expr2),
566            (Expression::Sin(expr1), Expression::Sin(expr2)) => expr1.fast_eq(expr2),
567            (Expression::Cos(expr1), Expression::Cos(expr2)) => expr1.fast_eq(expr2),
568            _ => false,
569        }
570    }
571
572    pub fn substitute<S: AsRef<Expression>, T: AsRef<Expression>>(
573        &self,
574        original: S,
575        substitution: T,
576    ) -> Self {
577        let original = original.as_ref();
578        let substitution = substitution.as_ref();
579        if self.fast_eq(original) {
580            return substitution.clone();
581        }
582        match self {
583            Expression::Pi => self.clone(),
584            Expression::Variable(_) => self.clone(),
585            Expression::Constant(_) => self.clone(),
586            Expression::Neg(expr) => {
587                Expression::Neg(Box::new(expr.substitute(original, substitution)))
588            }
589            Expression::Add(lhs, rhs) => Expression::Add(
590                Box::new(lhs.substitute(original, substitution)),
591                Box::new(rhs.substitute(original, substitution)),
592            ),
593            Expression::Sub(lhs, rhs) => Expression::Sub(
594                Box::new(lhs.substitute(original, substitution)),
595                Box::new(rhs.substitute(original, substitution)),
596            ),
597            Expression::Mul(lhs, rhs) => Expression::Mul(
598                Box::new(lhs.substitute(original, substitution)),
599                Box::new(rhs.substitute(original, substitution)),
600            ),
601            Expression::Div(lhs, rhs) => Expression::Div(
602                Box::new(lhs.substitute(original, substitution)),
603                Box::new(rhs.substitute(original, substitution)),
604            ),
605            Expression::Pow(lhs, rhs) => Expression::Pow(
606                Box::new(lhs.substitute(original, substitution)),
607                Box::new(rhs.substitute(original, substitution)),
608            ),
609            Expression::Sqrt(expr) => {
610                Expression::Sqrt(Box::new(expr.substitute(original, substitution)))
611            }
612            Expression::Sin(expr) => {
613                Expression::Sin(Box::new(expr.substitute(original, substitution)))
614            }
615            Expression::Cos(expr) => {
616                Expression::Cos(Box::new(expr.substitute(original, substitution)))
617            }
618        }
619    }
620
621    pub fn simplify(&self) -> Self {
622        simplify(self)
623    }
624
625    pub fn get_unique_variables(&self) -> Vec<String> {
626        match self {
627            Expression::Pi => {
628                vec![]
629            }
630            Expression::Variable(s) => {
631                vec![s.clone()]
632            }
633            Expression::Constant(_) => {
634                vec![]
635            }
636            Expression::Neg(expr) => expr.get_unique_variables(),
637            Expression::Add(lhs, rhs) => {
638                let mut l = lhs.get_unique_variables();
639                for r in rhs.get_unique_variables().into_iter() {
640                    if !l.contains(&r) {
641                        l.push(r)
642                    }
643                }
644                l
645            }
646            Expression::Sub(lhs, rhs) => {
647                let mut l = lhs.get_unique_variables();
648                for r in rhs.get_unique_variables().into_iter() {
649                    if !l.contains(&r) {
650                        l.push(r)
651                    }
652                }
653                l
654            }
655            Expression::Mul(lhs, rhs) => {
656                let mut l = lhs.get_unique_variables();
657                for r in rhs.get_unique_variables().into_iter() {
658                    if !l.contains(&r) {
659                        l.push(r)
660                    }
661                }
662                l
663            }
664            Expression::Div(lhs, rhs) => {
665                let mut l = lhs.get_unique_variables();
666                for r in rhs.get_unique_variables().into_iter() {
667                    if !l.contains(&r) {
668                        l.push(r)
669                    }
670                }
671                l
672            }
673            Expression::Pow(lhs, rhs) => {
674                let mut l = lhs.get_unique_variables();
675                for r in rhs.get_unique_variables().into_iter() {
676                    if !l.contains(&r) {
677                        l.push(r)
678                    }
679                }
680                l
681            }
682            Expression::Sqrt(expr) => expr.get_unique_variables(),
683            Expression::Sin(expr) => expr.get_unique_variables(),
684            Expression::Cos(expr) => expr.get_unique_variables(),
685        }
686    }
687}
688
689impl std::ops::Add<Expression> for Expression {
690    type Output = Self;
691
692    fn add(self, other: Self) -> Self {
693        &self + &other
694    }
695}
696
697impl std::ops::Add<&Expression> for Expression {
698    type Output = Expression;
699
700    fn add(self, other: &Expression) -> Expression {
701        &self + other
702    }
703}
704
705impl std::ops::Add<Expression> for &Expression {
706    type Output = Expression;
707
708    fn add(self, other: Expression) -> Expression {
709        self + &other
710    }
711}
712
713impl std::ops::Add<&Expression> for &Expression {
714    type Output = Expression;
715
716    fn add(self, other: &Expression) -> Expression {
717        if let Expression::Constant(c1) = self
718            && let Expression::Constant(c2) = other
719        {
720            return Expression::Constant(c1 + c2);
721        }
722        if other.is_zero_fast() {
723            self.clone()
724        } else if self.is_zero_fast() {
725            other.clone()
726        } else {
727            Expression::Add(Box::new(self.clone()), Box::new(other.clone()))
728        }
729    }
730}
731
732impl std::ops::Sub<Expression> for Expression {
733    type Output = Self;
734
735    fn sub(self, other: Self) -> Self {
736        &self - &other
737    }
738}
739
740impl std::ops::Sub<&Expression> for Expression {
741    type Output = Expression;
742
743    fn sub(self, other: &Expression) -> Expression {
744        &self - other
745    }
746}
747
748impl std::ops::Sub<Expression> for &Expression {
749    type Output = Expression;
750
751    fn sub(self, other: Expression) -> Expression {
752        self - &other
753    }
754}
755
756impl std::ops::Sub<&Expression> for &Expression {
757    type Output = Expression;
758
759    fn sub(self, other: &Expression) -> Expression {
760        if let Expression::Constant(c1) = self
761            && let Expression::Constant(c2) = other
762        {
763            return Expression::Constant(c1 - c2);
764        }
765        if other.is_zero_fast() {
766            self.clone()
767        } else if self.is_zero_fast() {
768            -other.clone()
769        } else {
770            Expression::Sub(Box::new(self.clone()), Box::new(other.clone()))
771        }
772    }
773}
774
775impl std::ops::Mul<Expression> for Expression {
776    type Output = Self;
777
778    fn mul(self, other: Self) -> Self {
779        &self * &other
780    }
781}
782
783impl std::ops::Mul<&Expression> for Expression {
784    type Output = Expression;
785
786    fn mul(self, other: &Expression) -> Expression {
787        &self * other
788    }
789}
790
791impl std::ops::Mul<Expression> for &Expression {
792    type Output = Expression;
793
794    fn mul(self, other: Expression) -> Expression {
795        self * &other
796    }
797}
798
799impl std::ops::Mul<&Expression> for &Expression {
800    type Output = Expression;
801
802    fn mul(self, other: &Expression) -> Expression {
803        if let Expression::Constant(c1) = self
804            && let Expression::Constant(c2) = other
805        {
806            return Expression::Constant(c1 * c2);
807        }
808        if other.is_zero_fast() || self.is_zero_fast() {
809            Expression::zero()
810        } else if other.is_one_fast() {
811            self.clone()
812        } else if self.is_one_fast() {
813            other.clone()
814        } else {
815            Expression::Mul(Box::new(self.clone()), Box::new(other.clone()))
816        }
817    }
818}
819
820impl std::ops::Div<Expression> for Expression {
821    type Output = Self;
822
823    fn div(self, other: Self) -> Self {
824        &self / &other
825    }
826}
827
828impl std::ops::Div<&Expression> for Expression {
829    type Output = Expression;
830
831    fn div(self, other: &Expression) -> Expression {
832        &self / other
833    }
834}
835
836impl std::ops::Div<Expression> for &Expression {
837    type Output = Expression;
838
839    fn div(self, other: Expression) -> Expression {
840        self / &other
841    }
842}
843
844impl std::ops::Div<&Expression> for &Expression {
845    type Output = Expression;
846
847    fn div(self, other: &Expression) -> Expression {
848        if other.is_zero_fast() {
849            panic!("Cannot divide by zero")
850        } else if let (Expression::Constant(c1), Expression::Constant(c2)) = (self, other) {
851            Expression::Constant(c1 / c2)
852        } else if self.is_zero_fast() {
853            Expression::zero()
854        } else {
855            Expression::Div(Box::new(self.clone()), Box::new(other.clone()))
856        }
857    }
858}
859
860impl std::ops::Neg for Expression {
861    type Output = Self;
862
863    fn neg(self) -> Self {
864        -&self
865    }
866}
867
868impl std::ops::Neg for &Expression {
869    type Output = Expression;
870
871    fn neg(self) -> Expression {
872        if self.is_zero_fast() {
873            self.clone()
874        } else {
875            Expression::Neg(Box::new(self.clone()))
876        }
877    }
878}
879
880impl std::fmt::Debug for Expression {
881    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
882        write!(f, "{}", self)
883    }
884}
885
886impl PartialEq for Expression {
887    fn eq(&self, other: &Self) -> bool {
888        self.fast_eq(other)
889        // if self.fast_eq(other) {
890        //     return true;
891        // }
892        // check_equality(self, other)
893    }
894}
895
896impl Eq for Expression {}
897
898impl std::hash::Hash for Expression {
899    fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
900        let val = self.hash_eval();
901        (val * 1e5_f64).round().to_bits().hash(state);
902    }
903}
904
905impl AsRef<Expression> for Expression {
906    fn as_ref(&self) -> &Expression {
907        self
908    }
909}
910
911impl std::fmt::Display for Expression {
912    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
913        let inner = match self {
914            Expression::Pi => "pi".to_string(),
915            Expression::Variable(var) => var.clone(),
916            Expression::Constant(_c) => self.to_float().to_string(),
917            Expression::Neg(expr) => format!("~ {}", expr),
918            Expression::Add(lhs, rhs) => format!("+ {} {}", lhs, rhs),
919            Expression::Sub(lhs, rhs) => format!("- {} {}", lhs, rhs),
920            Expression::Mul(lhs, rhs) => format!("* {} {}", lhs, rhs),
921            Expression::Div(lhs, rhs) => format!("/ {} {}", lhs, rhs),
922            Expression::Pow(lhs, rhs) => format!("pow {} {}", lhs, rhs),
923            Expression::Sqrt(expr) => format!("sqrt {}", expr),
924            Expression::Sin(expr) => format!("sin {}", expr),
925            Expression::Cos(expr) => format!("cos {}", expr),
926        };
927        write!(f, "({})", inner)
928    }
929}
930
931impl<R: RealScalar> From<R> for Expression {
932    fn from(value: R) -> Self {
933        Expression::from_float(value.to64())
934    }
935}
qudit_expr/expressions/base.rs

qudit_expr/expressions/
base.rs
