sim_lib_numbers_cas_diff/implementation/
integrate.rs1use sim_kernel::{Cx, Error, Result, Symbol, Value};
5use sim_lib_numbers_cas::{CasExpr, simplify_expr};
6use sim_lib_numbers_core::domains;
7
8use super::diff::math;
9
10pub fn integrate_sym_symbol() -> Symbol {
12 Symbol::new("integrate-sym")
13}
14
15pub fn integrate_cas(cx: &mut Cx, expr: &CasExpr, var: &Symbol) -> Result<CasExpr> {
21 let integral = match expr {
22 CasExpr::Num(value) => op(
23 math("mul"),
24 vec![CasExpr::num(cx, value.clone())?, CasExpr::Var(var.clone())],
25 ),
26 CasExpr::Var(symbol) if symbol == var => integrate_power_of_var(cx, 1, var)?,
27 CasExpr::Var(symbol) => op(
28 math("mul"),
29 vec![CasExpr::Var(symbol.clone()), CasExpr::Var(var.clone())],
30 ),
31 CasExpr::Op(operator, args) if *operator == math("add") => {
32 op(math("add"), integrate_all(cx, args, var)?)
33 }
34 CasExpr::Op(operator, args) if *operator == math("sub") => {
35 op(math("sub"), integrate_all(cx, args, var)?)
36 }
37 CasExpr::Op(operator, args) if *operator == math("mul") => integrate_mul(cx, args, var)?,
38 CasExpr::Op(operator, args) if *operator == math("pow") => integrate_pow(cx, args, var)?,
39 _ => {
40 return Err(Error::Eval(format!(
41 "{} only supports constants, sums, scalar products, and powers of the integration variable",
42 integrate_sym_symbol()
43 )));
44 }
45 };
46 simplify_expr(cx, integral)
47}
48
49fn integrate_all(cx: &mut Cx, args: &[CasExpr], var: &Symbol) -> Result<Vec<CasExpr>> {
50 args.iter().map(|arg| integrate_cas(cx, arg, var)).collect()
51}
52
53fn integrate_mul(cx: &mut Cx, args: &[CasExpr], var: &Symbol) -> Result<CasExpr> {
54 let mut constant = Vec::new();
55 let mut variable = Vec::new();
56 for arg in args {
57 if depends_on(arg, var) {
58 variable.push(arg.clone());
59 } else {
60 constant.push(arg.clone());
61 }
62 }
63 if variable.len() != 1 {
64 return Err(Error::Eval(format!(
65 "{} only handles products with exactly one variable-dependent factor",
66 integrate_sym_symbol()
67 )));
68 }
69 let mut out = constant;
70 out.push(integrate_cas(cx, &variable[0], var)?);
71 Ok(op(math("mul"), out))
72}
73
74fn integrate_pow(cx: &mut Cx, args: &[CasExpr], var: &Symbol) -> Result<CasExpr> {
75 let [base, exponent] = two_args(args)?;
76 if !matches!(base, CasExpr::Var(symbol) if symbol == var) {
77 return Err(Error::Eval(format!(
78 "{} only supports powers of the integration variable",
79 integrate_sym_symbol()
80 )));
81 }
82 let exponent = literal_i64(cx, exponent)?.ok_or_else(|| {
83 Error::Eval(format!(
84 "{} only supports integer exponents for symbolic powers",
85 integrate_sym_symbol()
86 ))
87 })?;
88 integrate_power_of_var(cx, exponent, var)
89}
90
91fn integrate_power_of_var(cx: &mut Cx, exponent: i64, var: &Symbol) -> Result<CasExpr> {
92 if exponent == -1 {
93 return Ok(op(Symbol::new("ln"), vec![CasExpr::Var(var.clone())]));
94 }
95 let next = exponent.checked_add(1).ok_or_else(|| {
96 Error::Eval(format!(
97 "{} exponent {exponent} overflows when raised for integration",
98 integrate_sym_symbol()
99 ))
100 })?;
101 let coefficient = rational_constant(cx, 1, next)?;
102 let exponent = integer_constant(cx, next)?;
103 Ok(op(
104 math("mul"),
105 vec![
106 CasExpr::num(cx, coefficient)?,
107 op(
108 math("pow"),
109 vec![CasExpr::Var(var.clone()), CasExpr::num(cx, exponent)?],
110 ),
111 ],
112 ))
113}
114
115fn depends_on(expr: &CasExpr, var: &Symbol) -> bool {
116 match expr {
117 CasExpr::Num(_) => false,
118 CasExpr::Var(symbol) => symbol == var,
119 CasExpr::Op(_, args) => args.iter().any(|arg| depends_on(arg, var)),
120 }
121}
122
123fn literal_i64(cx: &mut Cx, expr: &CasExpr) -> Result<Option<i64>> {
124 let CasExpr::Num(value) = expr else {
125 return Ok(None);
126 };
127 let display = value.object().display(cx)?;
128 Ok(display.parse::<i64>().ok())
129}
130
131fn integer_constant(cx: &mut Cx, value: i64) -> Result<Value> {
132 if cx
133 .registry()
134 .number_domain_by_symbol(&domains::i64())
135 .is_some()
136 {
137 return cx
138 .factory()
139 .number_literal(domains::i64(), value.to_string());
140 }
141 cx.factory()
142 .number_literal(domains::f64(), format!("{value}.0"))
143}
144
145fn rational_constant(cx: &mut Cx, num: i64, den: i64) -> Result<Value> {
146 if cx
147 .registry()
148 .number_domain_by_symbol(&domains::rational())
149 .is_some()
150 {
151 return cx
152 .factory()
153 .number_literal(domains::rational(), format!("{num}/{den}"));
154 }
155 let value = num as f64 / den as f64;
156 cx.factory()
157 .number_literal(domains::f64(), value.to_string())
158}
159
160fn two_args(args: &[CasExpr]) -> Result<[&CasExpr; 2]> {
161 let [left, right] = args else {
162 return Err(Error::Eval(format!(
163 "{} expects exactly two operands",
164 math("pow")
165 )));
166 };
167 Ok([left, right])
168}
169
170use super::diff::op;