sim_lib_numbers_cas_diff/implementation/
diff.rs

1//! Symbolic differentiation of `CasExpr` trees: the recursive `diff` rules for
2//! the builtin operators, falling through to the extensible rule registry.
3
4use sim_kernel::{Cx, Error, Result, Symbol, Value};
5use sim_lib_numbers_cas::{CasExpr, simplify_expr};
6use sim_lib_numbers_core::domains;
7
8use super::registry::apply_registered_rule;
9
10/// The `diff` symbol: the symbolic differentiation entry point.
11pub fn diff_symbol() -> Symbol {
12    Symbol::new("diff")
13}
14
15/// Differentiate a [`CasExpr`] with respect to `var`, returning a simplified
16/// derivative tree.
17///
18/// Builtin operators (`+`, `-`, `*`, `/`, `^`, and the standard trig functions)
19/// use their hard-coded rules; unrecognized operators fall through to the
20/// extensible differentiation-rule registry.
21pub fn diff_cas(cx: &mut Cx, expr: &CasExpr, var: &Symbol) -> Result<CasExpr> {
22    let derivative = match expr {
23        CasExpr::Num(_) => zero(cx)?,
24        CasExpr::Var(symbol) if symbol == var => one(cx)?,
25        CasExpr::Var(_) => zero(cx)?,
26        CasExpr::Op(operator, args) if *operator == math("add") => {
27            op(math("add"), diff_all(cx, args, var)?)
28        }
29        CasExpr::Op(operator, args) if *operator == math("sub") => diff_sub(cx, args, var)?,
30        CasExpr::Op(operator, args) if *operator == math("mul") => diff_mul(cx, args, var)?,
31        CasExpr::Op(operator, args) if *operator == math("div") => diff_div(cx, args, var)?,
32        CasExpr::Op(operator, args) if *operator == math("pow") => diff_pow(cx, args, var)?,
33        CasExpr::Op(operator, args) if *operator == Symbol::new("sin") => {
34            chain_rule(cx, Symbol::new("cos"), args, var)?
35        }
36        CasExpr::Op(operator, args) if *operator == Symbol::new("cos") => {
37            let [arg] = one_arg(args, operator)?;
38            op(
39                math("mul"),
40                vec![
41                    neg_one(cx)?,
42                    diff_cas(cx, arg, var)?,
43                    op(Symbol::new("sin"), vec![arg.clone()]),
44                ],
45            )
46        }
47        CasExpr::Op(operator, args) if *operator == Symbol::new("tan") => {
48            let [arg] = one_arg(args, operator)?;
49            op(
50                math("div"),
51                vec![
52                    diff_cas(cx, arg, var)?,
53                    op(
54                        math("pow"),
55                        vec![
56                            op(Symbol::new("cos"), vec![arg.clone()]),
57                            num_constant(cx, "2")?,
58                        ],
59                    ),
60                ],
61            )
62        }
63        CasExpr::Op(operator, args) if *operator == Symbol::new("ln") => {
64            let [arg] = one_arg(args, operator)?;
65            op(math("div"), vec![diff_cas(cx, arg, var)?, arg.clone()])
66        }
67        CasExpr::Op(operator, args) if *operator == Symbol::new("exp") => {
68            chain_rule(cx, Symbol::new("exp"), args, var)?
69        }
70        CasExpr::Op(operator, args) => {
71            if let Some(custom) = apply_registered_rule(operator, args, var) {
72                custom
73            } else {
74                op(
75                    diff_symbol(),
76                    vec![
77                        CasExpr::Op(operator.clone(), args.clone()),
78                        CasExpr::Var(var.clone()),
79                    ],
80                )
81            }
82        }
83    };
84    simplify_expr(cx, derivative)
85}
86
87fn diff_all(cx: &mut Cx, args: &[CasExpr], var: &Symbol) -> Result<Vec<CasExpr>> {
88    args.iter().map(|arg| diff_cas(cx, arg, var)).collect()
89}
90
91fn diff_sub(cx: &mut Cx, args: &[CasExpr], var: &Symbol) -> Result<CasExpr> {
92    match args {
93        [] => Err(Error::Eval(
94            "cannot differentiate an empty subtraction".to_owned(),
95        )),
96        [arg] => Ok(op(math("mul"), vec![neg_one(cx)?, diff_cas(cx, arg, var)?])),
97        _ => Ok(op(math("sub"), diff_all(cx, args, var)?)),
98    }
99}
100
101fn diff_mul(cx: &mut Cx, args: &[CasExpr], var: &Symbol) -> Result<CasExpr> {
102    if args.is_empty() {
103        return Err(Error::Eval(
104            "cannot differentiate an empty multiplication".to_owned(),
105        ));
106    }
107    let mut terms = Vec::with_capacity(args.len());
108    for (index, _) in args.iter().enumerate() {
109        let mut factors = Vec::with_capacity(args.len());
110        for (offset, arg) in args.iter().enumerate() {
111            if index == offset {
112                factors.push(diff_cas(cx, arg, var)?);
113            } else {
114                factors.push(arg.clone());
115            }
116        }
117        terms.push(op(math("mul"), factors));
118    }
119    Ok(op(math("add"), terms))
120}
121
122fn diff_div(cx: &mut Cx, args: &[CasExpr], var: &Symbol) -> Result<CasExpr> {
123    match args {
124        [] => Err(Error::Eval(
125            "cannot differentiate an empty division".to_owned(),
126        )),
127        [arg] => Ok(op(
128            math("div"),
129            vec![
130                op(math("mul"), vec![neg_one(cx)?, diff_cas(cx, arg, var)?]),
131                op(math("pow"), vec![arg.clone(), num_constant(cx, "2")?]),
132            ],
133        )),
134        [left, right] => {
135            let left_diff = diff_cas(cx, left, var)?;
136            let right_diff = diff_cas(cx, right, var)?;
137            Ok(op(
138                math("div"),
139                vec![
140                    op(
141                        math("sub"),
142                        vec![
143                            op(math("mul"), vec![left_diff, right.clone()]),
144                            op(math("mul"), vec![left.clone(), right_diff]),
145                        ],
146                    ),
147                    op(math("pow"), vec![right.clone(), num_constant(cx, "2")?]),
148                ],
149            ))
150        }
151        [head, tail @ ..] => diff_div(cx, &[head.clone(), op(math("mul"), tail.to_vec())], var),
152    }
153}
154
155fn diff_pow(cx: &mut Cx, args: &[CasExpr], var: &Symbol) -> Result<CasExpr> {
156    let [base, exponent] = two_args(args, &math("pow"))?;
157    let base_diff = diff_cas(cx, base, var)?;
158    if let CasExpr::Num(value) = exponent
159        && let Some(decremented) = decrement_value(cx, value)?
160    {
161        return Ok(op(
162            math("mul"),
163            vec![
164                CasExpr::num(cx, value.clone())?,
165                op(
166                    math("pow"),
167                    vec![base.clone(), CasExpr::num(cx, decremented)?],
168                ),
169                base_diff,
170            ],
171        ));
172    }
173    let exponent_diff = diff_cas(cx, exponent, var)?;
174    Ok(op(
175        math("mul"),
176        vec![
177            op(math("pow"), vec![base.clone(), exponent.clone()]),
178            op(
179                math("add"),
180                vec![
181                    op(
182                        math("mul"),
183                        vec![exponent_diff, op(Symbol::new("ln"), vec![base.clone()])],
184                    ),
185                    op(
186                        math("div"),
187                        vec![
188                            op(math("mul"), vec![exponent.clone(), base_diff]),
189                            base.clone(),
190                        ],
191                    ),
192                ],
193            ),
194        ],
195    ))
196}
197
198fn chain_rule(cx: &mut Cx, outer: Symbol, args: &[CasExpr], var: &Symbol) -> Result<CasExpr> {
199    let [arg] = one_arg(args, &outer)?;
200    Ok(op(
201        math("mul"),
202        vec![diff_cas(cx, arg, var)?, op(outer, vec![arg.clone()])],
203    ))
204}
205
206fn one_arg<'a>(args: &'a [CasExpr], operator: &Symbol) -> Result<[&'a CasExpr; 1]> {
207    let [arg] = args else {
208        return Err(Error::Eval(format!(
209            "{operator} expects exactly one CAS operand"
210        )));
211    };
212    Ok([arg])
213}
214
215fn two_args<'a>(args: &'a [CasExpr], operator: &Symbol) -> Result<[&'a CasExpr; 2]> {
216    let [left, right] = args else {
217        return Err(Error::Eval(format!(
218            "{operator} expects exactly two CAS operands"
219        )));
220    };
221    Ok([left, right])
222}
223
224fn zero(cx: &mut Cx) -> Result<CasExpr> {
225    num_constant(cx, "0")
226}
227
228fn one(cx: &mut Cx) -> Result<CasExpr> {
229    num_constant(cx, "1")
230}
231
232fn neg_one(cx: &mut Cx) -> Result<CasExpr> {
233    num_constant(cx, "-1")
234}
235
236fn num_constant(cx: &mut Cx, canonical: &str) -> Result<CasExpr> {
237    let value = number_constant(cx, canonical)?;
238    CasExpr::num(cx, value)
239}
240
241fn number_constant(cx: &mut Cx, canonical: &str) -> Result<Value> {
242    if cx
243        .registry()
244        .number_domain_by_symbol(&domains::i64())
245        .is_some()
246    {
247        return cx
248            .factory()
249            .number_literal(domains::i64(), canonical.to_owned());
250    }
251    if cx
252        .registry()
253        .number_domain_by_symbol(&domains::f64())
254        .is_some()
255    {
256        let canonical = if canonical == "-1" {
257            "-1.0".to_owned()
258        } else {
259            format!("{canonical}.0")
260        };
261        return cx.factory().number_literal(domains::f64(), canonical);
262    }
263    Err(Error::Eval(
264        "CAS differentiation requires a loaded integer or f64 number domain".to_owned(),
265    ))
266}
267
268fn decrement_value(cx: &mut Cx, value: &Value) -> Result<Option<Value>> {
269    if literal_number(cx, value)?.is_none() {
270        return Ok(None);
271    }
272    let one = number_constant(cx, "1")?;
273    let decremented = cx.apply_value_number_binary_op(&math("sub"), value.clone(), one)?;
274    Ok(cx
275        .number_value_ref(decremented.clone())?
276        .and_then(|number| number.literal)
277        .map(|_| decremented))
278}
279
280use sim_lib_numbers_cas::literal_number;
281
282pub(crate) fn math(name: &str) -> Symbol {
283    Symbol::qualified("math", name)
284}
285
286pub(crate) fn op(operator: Symbol, args: Vec<CasExpr>) -> CasExpr {
287    CasExpr::Op(operator, args)
288}
sim_lib_numbers_cas_diff/implementation/diff.rs

sim_lib_numbers_cas_diff/implementation/
diff.rs
