alkahest_cas/eval/
complex_f64.rs

1//! IEEE-754 complex evaluation for `re` / `im` / `conjugate` / `arg`.
2
3use crate::kernel::{ExprData, ExprId, ExprPool};
4use std::collections::HashMap;
5
6use super::{error, EvalError, UnsupportedReason};
7
8#[derive(Clone, Copy, Debug, PartialEq)]
9pub struct ComplexF64 {
10    pub re: f64,
11    pub im: f64,
12}
13
14impl ComplexF64 {
15    pub const ZERO: Self = Self { re: 0.0, im: 0.0 };
16    pub const ONE: Self = Self { re: 1.0, im: 0.0 };
17    pub fn new(re: f64, im: f64) -> Self {
18        Self { re, im }
19    }
20    fn add(self, o: Self) -> Self {
21        Self::new(self.re + o.re, self.im + o.im)
22    }
23    fn mul(self, o: Self) -> Self {
24        Self::new(
25            self.re * o.re - self.im * o.im,
26            self.re * o.im + self.im * o.re,
27        )
28    }
29    fn powi(self, n: i64) -> Result<Self, EvalError> {
30        if n == 0 {
31            return Ok(Self::ONE);
32        }
33        if n < 0 {
34            if self.re == 0.0 && self.im == 0.0 {
35                return Err(error(UnsupportedReason::ZeroToNegativePower));
36            }
37            let p = self.powi(-n)?;
38            let d = p.re * p.re + p.im * p.im;
39            if d == 0.0 || !d.is_finite() {
40                return Err(error(UnsupportedReason::NonFiniteResult));
41            }
42            return Ok(Self::new(p.re / d, -p.im / d));
43        }
44        let mut acc = Self::ONE;
45        let mut base = self;
46        let mut e = n;
47        while e != 0 {
48            if e & 1 == 1 {
49                acc = acc.mul(base);
50            }
51            e >>= 1;
52            if e != 0 {
53                base = base.mul(base);
54            }
55        }
56        Ok(acc)
57    }
58
59    /// Principal-branch power `z^w = exp(w · Log z)` for `z ≠ 0`.
60    fn powc(self, exp: Self) -> Result<Self, EvalError> {
61        if self.re == 0.0 && self.im == 0.0 {
62            // 0^w: only non-negative real exponents are defined in the
63            // principal sense we support here.
64            if exp.im == 0.0 && exp.re > 0.0 {
65                return Ok(Self::ZERO);
66            }
67            if exp.re == 0.0 && exp.im == 0.0 {
68                return Err(error(UnsupportedReason::UnsupportedExpression {
69                    kind: "branch_cut",
70                }));
71            }
72            return Err(error(UnsupportedReason::ZeroToNegativePower));
73        }
74        let ln = self.ln()?;
75        Ok(exp.mul(ln).exp())
76    }
77
78    fn sqrt(self) -> Result<Self, EvalError> {
79        // Use the principal logarithm rather than the textbook geometric
80        // formula: `((r±re)/2)^½` suffers catastrophic cancellation for
81        // arguments near the negative-real cut (e.g. -100 + 1e-6·i).
82        if self.re == 0.0 && self.im == 0.0 {
83            return Ok(Self::ZERO);
84        }
85        self.powc(Self::new(0.5, 0.0))
86    }
87    fn exp(self) -> Self {
88        let s = self.re.exp();
89        Self::new(s * self.im.cos(), s * self.im.sin())
90    }
91    fn ln(self) -> Result<Self, EvalError> {
92        if self.re == 0.0 && self.im == 0.0 {
93            return Err(error(UnsupportedReason::UnsupportedExpression {
94                kind: "branch_cut",
95            }));
96        }
97        let r = (self.re * self.re + self.im * self.im).sqrt();
98        Ok(Self::new(r.ln(), self.im.atan2(self.re)))
99    }
100    fn sin(self) -> Self {
101        Self::new(
102            self.re.sin() * self.im.cosh(),
103            self.re.cos() * self.im.sinh(),
104        )
105    }
106    fn cos(self) -> Self {
107        Self::new(
108            self.re.cos() * self.im.cosh(),
109            -self.re.sin() * self.im.sinh(),
110        )
111    }
112    fn principal_arg(self) -> Result<f64, EvalError> {
113        // Principal arg is undefined at 0 and discontinuous on the negative real axis.
114        if self.im == 0.0 && self.re <= 0.0 {
115            return Err(error(UnsupportedReason::UnsupportedExpression {
116                kind: "branch_cut",
117            }));
118        }
119        Ok(self.im.atan2(self.re))
120    }
121}
122
123pub fn eval_complex_f64(
124    expr: ExprId,
125    pool: &ExprPool,
126    bindings: &HashMap<ExprId, ComplexF64>,
127) -> Result<ComplexF64, EvalError> {
128    let v = eval_node(expr, pool, bindings)?;
129    if v.re.is_finite() && v.im.is_finite() {
130        Ok(v)
131    } else {
132        Err(error(UnsupportedReason::NonFiniteResult))
133    }
134}
135
136fn eval_node(
137    expr: ExprId,
138    pool: &ExprPool,
139    bindings: &HashMap<ExprId, ComplexF64>,
140) -> Result<ComplexF64, EvalError> {
141    match pool.get(expr) {
142        ExprData::Integer(n) => Ok(ComplexF64::new(n.0.to_f64(), 0.0)),
143        ExprData::Rational(r) => Ok(ComplexF64::new(r.0.to_f64(), 0.0)),
144        ExprData::Float(f) => Ok(ComplexF64::new(f.inner.to_f64(), 0.0)),
145        ExprData::Symbol { .. } => {
146            if let Some(&v) = bindings.get(&expr) {
147                Ok(v)
148            } else if pool.is_imaginary_unit(expr) {
149                // Canonical I evaluates to 0+1j in complex mode without an
150                // explicit binding (matches symbolic `i² → −1` folding).
151                Ok(ComplexF64::new(0.0, 1.0))
152            } else {
153                Err(error(UnsupportedReason::UnboundSymbol { symbol: expr }))
154            }
155        }
156        ExprData::Add(args) => args.iter().try_fold(ComplexF64::ZERO, |a, &x| {
157            Ok(a.add(eval_node(x, pool, bindings)?))
158        }),
159        ExprData::Mul(args) => args.iter().try_fold(ComplexF64::ONE, |a, &x| {
160            Ok(a.mul(eval_node(x, pool, bindings)?))
161        }),
162        ExprData::Pow { base, exp } => {
163            let b = eval_node(base, pool, bindings)?;
164            match pool.get(exp) {
165                ExprData::Integer(n) => b.powi(n.0.to_i64().unwrap_or(0)),
166                ExprData::Rational(r) if *r.0.denom() == 1 => {
167                    b.powi(r.0.numer().to_i64().unwrap_or(0))
168                }
169                // Principal branch: z^w = exp(w · Log z). Covers float and
170                // non-integer rational exponents (e.g. (-1)^(1/2) → i).
171                _ => {
172                    let e = eval_node(exp, pool, bindings)?;
173                    // Fast path: pure integer-valued real exponent.
174                    if e.im == 0.0 && e.re.fract() == 0.0 && e.re.abs() < (i64::MAX as f64) {
175                        b.powi(e.re as i64)
176                    } else {
177                        b.powc(e)
178                    }
179                }
180            }
181        }
182        ExprData::Func { name, args } if args.len() == 1 => {
183            let x = eval_node(args[0], pool, bindings)?;
184            match name.as_str() {
185                "sin" => Ok(x.sin()),
186                "cos" => Ok(x.cos()),
187                "exp" => Ok(x.exp()),
188                "log" => x.ln(),
189                "sqrt" => x.sqrt(),
190                "re" => Ok(ComplexF64::new(x.re, 0.0)),
191                "im" => Ok(ComplexF64::new(x.im, 0.0)),
192                "conjugate" => Ok(ComplexF64::new(x.re, -x.im)),
193                "arg" => Ok(ComplexF64::new(x.principal_arg()?, 0.0)),
194                _ => Err(error(UnsupportedReason::UnsupportedFunction {
195                    name: name.clone(),
196                })),
197            }
198        }
199        ExprData::Func { name, .. } => Err(error(UnsupportedReason::UnsupportedFunction {
200            name: name.clone(),
201        })),
202        other => Err(error(UnsupportedReason::UnsupportedExpression {
203            kind: super::expr_kind(&other),
204        })),
205    }
206}
207
208#[cfg(test)]
209mod tests {
210    use super::*;
211
212    #[test]
213    fn arg_declines_on_branch_cut() {
214        let pool = crate::kernel::ExprPool::new();
215        let expr = pool.func("arg", vec![pool.integer(-1_i32)]);
216        assert_eq!(
217            eval_complex_f64(expr, &pool, &HashMap::new())
218                .unwrap_err()
219                .reason,
220            UnsupportedReason::UnsupportedExpression { kind: "branch_cut" }
221        );
222    }
223
224    #[test]
225    fn imaginary_unit_auto_binds() {
226        let pool = crate::kernel::ExprPool::new();
227        let i = pool.imaginary_unit();
228        let v = eval_complex_f64(i, &pool, &HashMap::new()).unwrap();
229        assert_eq!(v, ComplexF64::new(0.0, 1.0));
230        let i2 = pool.mul(vec![i, i]);
231        let v2 = eval_complex_f64(i2, &pool, &HashMap::new()).unwrap();
232        assert!((v2.re + 1.0).abs() < 1e-12 && v2.im.abs() < 1e-12);
233    }
234
235    #[test]
236    fn principal_sqrt_of_negative_one() {
237        let pool = crate::kernel::ExprPool::new();
238        let expr = pool.func("sqrt", vec![pool.integer(-1_i32)]);
239        let v = eval_complex_f64(expr, &pool, &HashMap::new()).unwrap();
240        assert!((v.re).abs() < 1e-12 && (v.im - 1.0).abs() < 1e-12);
241    }
242
243    #[test]
244    fn principal_half_power_of_negative_one() {
245        let pool = crate::kernel::ExprPool::new();
246        let expr = pool.pow(pool.integer(-1_i32), pool.rational(1, 2));
247        let v = eval_complex_f64(expr, &pool, &HashMap::new()).unwrap();
248        assert!((v.re).abs() < 1e-12 && (v.im - 1.0).abs() < 1e-12);
249    }
250
251    #[test]
252    fn log_of_negative_one_is_i_pi() {
253        let pool = crate::kernel::ExprPool::new();
254        let expr = pool.func("log", vec![pool.integer(-1_i32)]);
255        let v = eval_complex_f64(expr, &pool, &HashMap::new()).unwrap();
256        assert!(v.re.abs() < 1e-12 && (v.im - std::f64::consts::PI).abs() < 1e-12);
257    }
258}
alkahest_cas/eval/complex_f64.rs

alkahest_cas/eval/
complex_f64.rs
