arael-sym 0.5.2

Symbolic math library: expression trees, automatic differentiation, simplification, CSE, code generation
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148

use std::collections::{BTreeSet, HashMap};
use super::{Expr, E};

impl Expr {
    /// Evaluate the expression numerically given variable bindings.
    ///
    /// Returns `Err` if any symbol in the expression is not bound in `vars`.
    pub fn eval(&self, vars: &HashMap<&str, f64>) -> Result<f64, String> {
        match self {
            Expr::Sym(name) => {
                vars.get(name.as_str()).copied()
                    .ok_or_else(|| format!("unbound symbol: {name}"))
            }
            Expr::Const(v) => Ok(*v),
            Expr::NamedConst { value, .. } => Ok(*value),
            Expr::Neg(a) => Ok(-a.eval(vars)?),
            Expr::Add(a, b) => Ok(a.eval(vars)? + b.eval(vars)?),
            Expr::Sub(a, b) => Ok(a.eval(vars)? - b.eval(vars)?),
            Expr::Mul(a, b) => Ok(a.eval(vars)? * b.eval(vars)?),
            Expr::Div(a, b) => Ok(a.eval(vars)? / b.eval(vars)?),
            Expr::Pow(a, b) => Ok(a.eval(vars)?.powf(b.eval(vars)?)),
            Expr::Sin(a) => Ok(a.eval(vars)?.sin()),
            Expr::Cos(a) => Ok(a.eval(vars)?.cos()),
            Expr::Tan(a) => Ok(a.eval(vars)?.tan()),
            Expr::Asin(a) => Ok(a.eval(vars)?.asin()),
            Expr::Acos(a) => Ok(a.eval(vars)?.acos()),
            Expr::Atan(a) => Ok(a.eval(vars)?.atan()),
            Expr::Atan2(y, x) => Ok(y.eval(vars)?.atan2(x.eval(vars)?)),
            Expr::Sinh(a) => Ok(a.eval(vars)?.sinh()),
            Expr::Cosh(a) => Ok(a.eval(vars)?.cosh()),
            Expr::Tanh(a) => Ok(a.eval(vars)?.tanh()),
            Expr::Exp(a) => Ok(a.eval(vars)?.exp()),
            Expr::Ln(a) => Ok(a.eval(vars)?.ln()),
            Expr::Log2(a) => Ok(a.eval(vars)?.log2()),
            Expr::Log10(a) => Ok(a.eval(vars)?.log10()),
            Expr::Sqrt(a) => Ok(a.eval(vars)?.sqrt()),
            Expr::Abs(a) => Ok(a.eval(vars)?.abs()),
            Expr::Heaviside(a) => {
                let v = a.eval(vars)?;
                Ok(if v < 0.0 { 0.0 } else { 1.0 })
            }
            Expr::Clamp(val, lo, hi) => {
                let v = val.eval(vars)?;
                let l = lo.eval(vars)?;
                let h = hi.eval(vars)?;
                Ok(v.clamp(l, h))
            }
            Expr::Func { params, kind, args, .. } => {
                if let Some(f) = kind.eval_fn() {
                    let vals: Result<Vec<f64>, _> = args.iter().map(|a| a.eval(vars)).collect();
                    Ok(f(&vals?))
                } else {
                    let body = kind.body().expect("FuncKind must have body or eval_fn");
                    super::expand_func(params, body, args).eval(vars)
                }
            }
        }
    }

    /// Substitute all occurrences of the named variable with `replacement`.
    ///
    /// `var` can be any [`AsVarName`] -- a `&str`, a `String`, or an
    /// [`E`] handle wrapping a `Sym` node. Returns a new expression
    /// with the substitution applied throughout.
    pub fn subs(&self, var: impl crate::AsVarName, replacement: &E) -> E {
        self.subs_by_name(var.var_name(), replacement)
    }

    fn subs_by_name(&self, var: &str, replacement: &E) -> E {
        match self {
            Expr::Sym(name) if name == var => replacement.clone(),
            Expr::Sym(_) | Expr::Const(_) | Expr::NamedConst { .. } => E::new(self.clone()),
            Expr::Neg(a) => -a.subs_by_name(var, replacement),
            Expr::Add(a, b) => a.subs_by_name(var, replacement) + b.subs_by_name(var, replacement),
            Expr::Sub(a, b) => a.subs_by_name(var, replacement) - b.subs_by_name(var, replacement),
            Expr::Mul(a, b) => a.subs_by_name(var, replacement) * b.subs_by_name(var, replacement),
            Expr::Div(a, b) => a.subs_by_name(var, replacement) / b.subs_by_name(var, replacement),
            Expr::Pow(a, b) => E::new(Expr::Pow(a.subs_by_name(var, replacement), b.subs_by_name(var, replacement))),
            Expr::Sin(a) => E::new(Expr::Sin(a.subs_by_name(var, replacement))),
            Expr::Cos(a) => E::new(Expr::Cos(a.subs_by_name(var, replacement))),
            Expr::Tan(a) => E::new(Expr::Tan(a.subs_by_name(var, replacement))),
            Expr::Asin(a) => E::new(Expr::Asin(a.subs_by_name(var, replacement))),
            Expr::Acos(a) => E::new(Expr::Acos(a.subs_by_name(var, replacement))),
            Expr::Atan(a) => E::new(Expr::Atan(a.subs_by_name(var, replacement))),
            Expr::Atan2(y, x) => E::new(Expr::Atan2(y.subs_by_name(var, replacement), x.subs_by_name(var, replacement))),
            Expr::Sinh(a) => E::new(Expr::Sinh(a.subs_by_name(var, replacement))),
            Expr::Cosh(a) => E::new(Expr::Cosh(a.subs_by_name(var, replacement))),
            Expr::Tanh(a) => E::new(Expr::Tanh(a.subs_by_name(var, replacement))),
            Expr::Exp(a) => E::new(Expr::Exp(a.subs_by_name(var, replacement))),
            Expr::Ln(a) => E::new(Expr::Ln(a.subs_by_name(var, replacement))),
            Expr::Log2(a) => E::new(Expr::Log2(a.subs_by_name(var, replacement))),
            Expr::Log10(a) => E::new(Expr::Log10(a.subs_by_name(var, replacement))),
            Expr::Sqrt(a) => E::new(Expr::Sqrt(a.subs_by_name(var, replacement))),
            Expr::Abs(a) => E::new(Expr::Abs(a.subs_by_name(var, replacement))),
            Expr::Heaviside(a) => E::new(Expr::Heaviside(a.subs_by_name(var, replacement))),
            Expr::Clamp(a, lo, hi) => E::new(Expr::Clamp(a.subs_by_name(var, replacement), lo.subs_by_name(var, replacement), hi.subs_by_name(var, replacement))),
            Expr::Func { name, params, kind, args } => {
                let new_args = args.iter().map(|a| a.subs_by_name(var, replacement)).collect();
                E::new(Expr::Func { name: name.clone(), params: params.clone(), kind: kind.clone(), args: new_args })
            }
        }
    }

    /// Collect all free (unbound) variable names in the expression.
    ///
    /// Returns a sorted set of variable name strings.
    pub fn free_vars(&self) -> BTreeSet<String> {
        let mut set = BTreeSet::new();
        self.collect_vars(&mut set);
        set
    }

    fn collect_vars(&self, set: &mut BTreeSet<String>) {
        match self {
            Expr::Sym(name) => { set.insert(name.clone()); }
            Expr::Const(_) | Expr::NamedConst { .. } => {}
            Expr::Neg(a) => a.collect_vars(set),
            Expr::Add(a, b) | Expr::Sub(a, b) | Expr::Mul(a, b)
            | Expr::Div(a, b) | Expr::Pow(a, b) | Expr::Atan2(a, b) => {
                a.collect_vars(set);
                b.collect_vars(set);
            }
            Expr::Sin(a) | Expr::Cos(a) | Expr::Tan(a)
            | Expr::Asin(a) | Expr::Acos(a) | Expr::Atan(a)
            | Expr::Sinh(a) | Expr::Cosh(a) | Expr::Tanh(a)
            | Expr::Exp(a) | Expr::Ln(a) | Expr::Log2(a) | Expr::Log10(a)
            | Expr::Sqrt(a) | Expr::Abs(a)
            | Expr::Heaviside(a) => {
                a.collect_vars(set);
            }
            Expr::Clamp(a, lo, hi) => {
                a.collect_vars(set);
                lo.collect_vars(set);
                hi.collect_vars(set);
            }
            Expr::Func { args, .. } => {
                for arg in args { arg.collect_vars(set); }
            }
        }
    }

    /// Differentiate with respect to multiple variables, returning a vector.
    ///
    /// Equivalent to calling [`Expr::diff`] for each variable in order.
    pub fn diff_all(&self, vars: &[&str]) -> Vec<E> {
        vars.iter().map(|v| self.diff(v)).collect()
    }
}