use std::collections::HashMap;
use super::{eval_complex_f64, ComplexF64};
use crate::kernel::{ExprId, ExprPool};
const MAX_DEGREE: usize = 24;
const ROOT_RESIDUAL_TOL: f64 = 1e-6;
const MAX_ITERATIONS: usize = 500;
const STEP_TOL: f64 = 1e-15;
const REAL_TOL: f64 = 1e-6;
pub(crate) fn eval_root_sum_f64(
poly: ExprId,
rvar: ExprId,
body: ExprId,
env: &HashMap<ExprId, f64>,
pool: &ExprPool,
) -> Option<f64> {
let roots = minimal_polynomial_roots(poly, rvar, pool)?;
let mut bindings: HashMap<ExprId, ComplexF64> = env
.iter()
.map(|(&k, &v)| (k, ComplexF64::new(v, 0.0)))
.collect();
let (mut sum_re, mut sum_im, mut scale) = (0.0f64, 0.0f64, 0.0f64);
for root in roots {
bindings.insert(rvar, root);
let term = eval_complex_f64(body, pool, &bindings).ok()?;
sum_re += term.re;
sum_im += term.im;
scale += term.re.abs() + term.im.abs();
}
if !sum_re.is_finite() || !sum_im.is_finite() {
return None;
}
if sum_im.abs() > REAL_TOL * scale.max(1.0) {
return None;
}
Some(sum_re)
}
fn minimal_polynomial_roots(
poly: ExprId,
rvar: ExprId,
pool: &ExprPool,
) -> Option<Vec<ComplexF64>> {
let up = crate::poly::UniPoly::from_symbolic_clear_denoms(poly, rvar, pool).ok()?;
let degree = usize::try_from(up.degree()).ok()?;
if degree == 0 || degree > MAX_DEGREE {
return None;
}
let coeffs: Vec<f64> = up.coefficients().iter().map(rug::Integer::to_f64).collect();
if coeffs.iter().any(|c| !c.is_finite()) {
return None;
}
let roots = polynomial_roots(&coeffs)?;
if roots.len() != degree {
return None;
}
for &z in &roots {
let (value, magnitude) = horner_with_magnitude(&coeffs, z);
if value.re.hypot(value.im) > ROOT_RESIDUAL_TOL * magnitude.max(1.0) {
return None;
}
}
Some(roots)
}
fn polynomial_roots(coeffs: &[f64]) -> Option<Vec<ComplexF64>> {
let mut c = coeffs.to_vec();
while c.len() > 1 && c.last() == Some(&0.0) {
c.pop();
}
if c.len() <= 1 {
return Some(Vec::new());
}
let at_origin = c.iter().take_while(|v| **v == 0.0).count();
let tail = &c[at_origin..];
let degree = tail.len() - 1;
let leading = *tail.last()?;
if leading == 0.0 || !leading.is_finite() {
return None;
}
let mut roots = vec![ComplexF64::new(0.0, 0.0); at_origin];
if degree == 0 {
return Some(roots);
}
let monic: Vec<f64> = tail.iter().map(|v| v / leading).collect();
if monic.iter().any(|v| !v.is_finite()) {
return None;
}
let scale = {
let s = monic[0].abs().powf(1.0 / degree as f64);
if s.is_finite() && s > 0.0 {
s
} else {
1.0
}
};
let scaled: Vec<f64> = monic
.iter()
.enumerate()
.map(|(k, &v)| v / scale.powi((degree - k) as i32))
.collect();
if scaled.iter().any(|v| !v.is_finite()) {
return None;
}
let seed = ComplexF64::new(0.4, 0.9);
let mut z = Vec::with_capacity(degree);
let mut point = ComplexF64::new(1.0, 0.0);
for _ in 0..degree {
z.push(point);
point = cmul(point, seed);
}
for _ in 0..MAX_ITERATIONS {
let mut worst = 0.0f64;
for i in 0..degree {
let mut denom = ComplexF64::new(1.0, 0.0);
for j in 0..degree {
if i != j {
denom = cmul(denom, ComplexF64::new(z[i].re - z[j].re, z[i].im - z[j].im));
}
}
let denom_abs = denom.re.hypot(denom.im);
if denom_abs == 0.0 || !denom_abs.is_finite() {
return None;
}
let (value, _) = horner_with_magnitude(&scaled, z[i]);
let step = cdiv(value, denom);
z[i] = ComplexF64::new(z[i].re - step.re, z[i].im - step.im);
if !z[i].re.is_finite() || !z[i].im.is_finite() {
return None;
}
worst = worst.max(step.re.hypot(step.im) / (1.0 + z[i].re.hypot(z[i].im)));
}
if worst < STEP_TOL {
break;
}
}
roots.extend(
z.into_iter()
.map(|w| ComplexF64::new(w.re * scale, w.im * scale)),
);
Some(roots)
}
fn cmul(a: ComplexF64, b: ComplexF64) -> ComplexF64 {
ComplexF64::new(a.re * b.re - a.im * b.im, a.re * b.im + a.im * b.re)
}
fn cdiv(a: ComplexF64, b: ComplexF64) -> ComplexF64 {
let d = b.re * b.re + b.im * b.im;
ComplexF64::new(
(a.re * b.re + a.im * b.im) / d,
(a.im * b.re - a.re * b.im) / d,
)
}
fn horner_with_magnitude(coeffs: &[f64], z: ComplexF64) -> (ComplexF64, f64) {
let az = z.re.hypot(z.im);
let mut acc = ComplexF64::new(0.0, 0.0);
let mut magnitude = 0.0f64;
for &c in coeffs.iter().rev() {
acc = ComplexF64::new(
acc.re * z.re - acc.im * z.im + c,
acc.re * z.im + acc.im * z.re,
);
magnitude = magnitude * az + c.abs();
}
(acc, magnitude)
}
#[cfg(test)]
mod tests {
use super::*;
use crate::kernel::{Domain, ExprPool};
#[test]
fn power_sums_of_a_quadratic() {
let pool = ExprPool::new();
let r = pool.symbol("r", Domain::Real);
let m = pool.add(vec![pool.pow(r, pool.integer(2_i32)), pool.integer(-2_i32)]);
let env = HashMap::new();
let s1 = eval_root_sum_f64(m, r, r, &env, &pool).unwrap();
assert!(s1.abs() < 1e-10, "Σc = {s1}");
let s2 = eval_root_sum_f64(m, r, pool.pow(r, pool.integer(2_i32)), &env, &pool).unwrap();
assert!((s2 - 4.0).abs() < 1e-10, "Σc² = {s2}");
}
#[test]
fn a_complex_root_pair_still_sums_to_a_real() {
let pool = ExprPool::new();
let r = pool.symbol("r", Domain::Real);
let m = pool.add(vec![pool.pow(r, pool.integer(2_i32)), pool.integer(1_i32)]);
let got = eval_root_sum_f64(
m,
r,
pool.pow(r, pool.integer(2_i32)),
&HashMap::new(),
&pool,
)
.unwrap();
assert!((got + 2.0).abs() < 1e-10, "Σc² over ±i = {got}");
}
#[test]
fn a_non_real_sum_is_declined() {
let pool = ExprPool::new();
let r = pool.symbol("r", Domain::Real);
let m = pool.add(vec![pool.pow(r, pool.integer(2_i32)), pool.integer(1_i32)]);
let two_i = pool.mul(vec![pool.integer(-2_i32), pool.imaginary_unit()]);
let body = pool.pow(pool.add(vec![r, two_i]), pool.integer(-1_i32));
assert_eq!(
eval_root_sum_f64(m, r, body, &HashMap::new(), &pool),
None,
"a sum with a real imaginary part is not a real number"
);
}
#[test]
fn a_non_polynomial_minimal_polynomial_is_declined() {
let pool = ExprPool::new();
let r = pool.symbol("r", Domain::Real);
let m = pool.func("sin", vec![r]);
assert_eq!(eval_root_sum_f64(m, r, r, &HashMap::new(), &pool), None);
}
#[test]
fn the_rothstein_trager_shape_evaluates_through_eval_interp() {
let pool = ExprPool::new();
let x = pool.symbol("x", Domain::Real);
let r = pool.symbol("$root$", Domain::Real);
let m = pool.add(vec![
pool.pow(r, pool.integer(3_i32)),
pool.rational(-1, 108),
]);
let arg = pool.add(vec![x, pool.mul(vec![pool.integer(-6_i32), r])]);
let dbody = pool.mul(vec![r, pool.pow(arg, pool.integer(-1_i32))]);
let sum = pool.root_sum(m, r, dbody);
for &xv in &[0.3719_f64, 1.4231, -2.8123] {
let env: HashMap<ExprId, f64> = std::iter::once((x, xv)).collect();
let got = crate::jit::eval_interp(sum, &env, &pool)
.unwrap_or_else(|| panic!("no value at x = {xv}"));
let want = 1.0 / (xv * xv * xv - 2.0);
assert!(
(got - want).abs() <= 1e-9 * (1.0 + want.abs()),
"at x = {xv}: Σ r/(x−6r) = {got}, 1/(x³−2) = {want}"
);
}
}
#[test]
fn an_oversized_minimal_polynomial_is_declined() {
let pool = ExprPool::new();
let r = pool.symbol("r", Domain::Real);
let m = pool.add(vec![
pool.pow(r, pool.integer((MAX_DEGREE + 1) as i32)),
pool.integer(-1_i32),
]);
assert_eq!(eval_root_sum_f64(m, r, r, &HashMap::new(), &pool), None);
}
}