|
| 1 | +//! Complex math functions matching Python's cmath module behavior. |
| 2 | +//! |
| 3 | +//! These implementations follow the algorithms from cmathmodule.c |
| 4 | +//! to ensure numerical precision and correct handling of edge cases. |
| 5 | +
|
| 6 | +mod exponential; |
| 7 | +mod misc; |
| 8 | +mod trigonometric; |
| 9 | + |
| 10 | +pub use exponential::{exp, log, log10, sqrt}; |
| 11 | +pub use misc::{abs, isclose, isfinite, isinf, isnan, phase, polar, rect}; |
| 12 | +pub use trigonometric::{acos, acosh, asin, asinh, atan, atanh, cos, cosh, sin, sinh, tan, tanh}; |
| 13 | + |
| 14 | +#[cfg(test)] |
| 15 | +use crate::Result; |
| 16 | +use crate::m; |
| 17 | +#[cfg(test)] |
| 18 | +use num_complex::Complex64; |
| 19 | + |
| 20 | +// Shared constants |
| 21 | + |
| 22 | +const M_LN2: f64 = core::f64::consts::LN_2; |
| 23 | + |
| 24 | +/// Used to avoid spurious overflow in sqrt, log, inverse trig/hyperbolic functions. |
| 25 | +const CM_LARGE_DOUBLE: f64 = f64::MAX / 4.0; |
| 26 | +const CM_LOG_LARGE_DOUBLE: f64 = 709.0895657128241; // log(CM_LARGE_DOUBLE) |
| 27 | + |
| 28 | +const INF: f64 = f64::INFINITY; |
| 29 | + |
| 30 | +// Special value table constants |
| 31 | +const P: f64 = core::f64::consts::PI; |
| 32 | +const P14: f64 = 0.25 * core::f64::consts::PI; |
| 33 | +const P12: f64 = 0.5 * core::f64::consts::PI; |
| 34 | +const P34: f64 = 0.75 * core::f64::consts::PI; |
| 35 | +const N: f64 = f64::NAN; |
| 36 | +#[allow(clippy::excessive_precision)] |
| 37 | +const U: f64 = -9.5426319407711027e33; // unlikely value, used as placeholder |
| 38 | + |
| 39 | +/// Helper to create Complex64 in const context (for special value tables) |
| 40 | +#[inline] |
| 41 | +const fn c(re: f64, im: f64) -> num_complex::Complex64 { |
| 42 | + num_complex::Complex64::new(re, im) |
| 43 | +} |
| 44 | + |
| 45 | +/// Special value types for classifying doubles. |
| 46 | +#[derive(Clone, Copy, Debug, PartialEq, Eq)] |
| 47 | +#[repr(usize)] |
| 48 | +enum SpecialType { |
| 49 | + NInf = 0, // negative infinity |
| 50 | + Neg = 1, // negative finite (nonzero) |
| 51 | + NZero = 2, // -0. |
| 52 | + PZero = 3, // +0. |
| 53 | + Pos = 4, // positive finite (nonzero) |
| 54 | + PInf = 5, // positive infinity |
| 55 | + Nan = 6, // NaN |
| 56 | +} |
| 57 | + |
| 58 | +/// Return special value from table if input is non-finite. |
| 59 | +macro_rules! special_value { |
| 60 | + ($z:expr, $table:expr) => { |
| 61 | + if !$z.re.is_finite() || !$z.im.is_finite() { |
| 62 | + return Ok($table[special_type($z.re) as usize][special_type($z.im) as usize]); |
| 63 | + } |
| 64 | + }; |
| 65 | +} |
| 66 | +pub(crate) use special_value; |
| 67 | + |
| 68 | +/// Classify a double into one of seven special types. |
| 69 | +#[inline] |
| 70 | +fn special_type(d: f64) -> SpecialType { |
| 71 | + if d.is_finite() { |
| 72 | + if d != 0.0 { |
| 73 | + if m::copysign(1.0, d) == 1.0 { |
| 74 | + SpecialType::Pos |
| 75 | + } else { |
| 76 | + SpecialType::Neg |
| 77 | + } |
| 78 | + } else if m::copysign(1.0, d) == 1.0 { |
| 79 | + SpecialType::PZero |
| 80 | + } else { |
| 81 | + SpecialType::NZero |
| 82 | + } |
| 83 | + } else if d.is_nan() { |
| 84 | + SpecialType::Nan |
| 85 | + } else if m::copysign(1.0, d) == 1.0 { |
| 86 | + SpecialType::PInf |
| 87 | + } else { |
| 88 | + SpecialType::NInf |
| 89 | + } |
| 90 | +} |
| 91 | + |
| 92 | +#[cfg(test)] |
| 93 | +pub(crate) mod tests { |
| 94 | + use super::*; |
| 95 | + |
| 96 | + /// Compare complex result with CPython, allowing small ULP differences for finite values. |
| 97 | + pub fn assert_complex_eq(py_re: f64, py_im: f64, rs: Complex64, func: &str, re: f64, im: f64) { |
| 98 | + let check_component = |py: f64, rs: f64, component: &str| { |
| 99 | + if py.is_nan() && rs.is_nan() { |
| 100 | + // Both NaN - OK |
| 101 | + } else if py.is_nan() || rs.is_nan() { |
| 102 | + panic!("{func}({re}, {im}).{component}: py={py} vs rs={rs} (one is NaN)",); |
| 103 | + } else if py.is_infinite() && rs.is_infinite() { |
| 104 | + // Check sign matches |
| 105 | + if py.is_sign_positive() != rs.is_sign_positive() { |
| 106 | + panic!("{func}({re}, {im}).{component}: py={py} vs rs={rs} (sign mismatch)",); |
| 107 | + } |
| 108 | + } else if py.is_infinite() || rs.is_infinite() { |
| 109 | + panic!("{func}({re}, {im}).{component}: py={py} vs rs={rs} (one is infinite)",); |
| 110 | + } else { |
| 111 | + // Both finite - allow small ULP difference |
| 112 | + let py_bits = py.to_bits() as i64; |
| 113 | + let rs_bits = rs.to_bits() as i64; |
| 114 | + let ulp_diff = (py_bits - rs_bits).abs(); |
| 115 | + if ulp_diff != 0 { |
| 116 | + panic!( |
| 117 | + "{func}({re}, {im}).{component}: py={py} (bits={:#x}) vs rs={rs} (bits={:#x}), ULP diff={ulp_diff}", |
| 118 | + py.to_bits(), |
| 119 | + rs.to_bits() |
| 120 | + ); |
| 121 | + } |
| 122 | + } |
| 123 | + }; |
| 124 | + check_component(py_re, rs.re, "re"); |
| 125 | + check_component(py_im, rs.im, "im"); |
| 126 | + } |
| 127 | + |
| 128 | + pub fn test_cmath_func<F>(func_name: &str, rs_func: F, re: f64, im: f64) |
| 129 | + where |
| 130 | + F: Fn(Complex64) -> Result<Complex64>, |
| 131 | + { |
| 132 | + use pyo3::prelude::*; |
| 133 | + |
| 134 | + let rs_result = rs_func(Complex64::new(re, im)); |
| 135 | + |
| 136 | + pyo3::Python::attach(|py| { |
| 137 | + let cmath = pyo3::types::PyModule::import(py, "cmath").unwrap(); |
| 138 | + let py_func = cmath.getattr(func_name).unwrap(); |
| 139 | + let py_result = py_func.call1((pyo3::types::PyComplex::from_doubles(py, re, im),)); |
| 140 | + |
| 141 | + match py_result { |
| 142 | + Ok(result) => { |
| 143 | + use pyo3::types::PyComplexMethods; |
| 144 | + let c = result.cast::<pyo3::types::PyComplex>().unwrap(); |
| 145 | + let py_re = c.real(); |
| 146 | + let py_im = c.imag(); |
| 147 | + match rs_result { |
| 148 | + Ok(rs) => { |
| 149 | + assert_complex_eq(py_re, py_im, rs, func_name, re, im); |
| 150 | + } |
| 151 | + Err(e) => { |
| 152 | + panic!( |
| 153 | + "{func_name}({re}, {im}): py=({py_re}, {py_im}) but rs returned error {e:?}" |
| 154 | + ); |
| 155 | + } |
| 156 | + } |
| 157 | + } |
| 158 | + Err(e) => { |
| 159 | + // CPython raised an exception - check we got an error too |
| 160 | + if rs_result.is_ok() { |
| 161 | + let rs = rs_result.unwrap(); |
| 162 | + // Some special cases may return values for domain errors in Python |
| 163 | + // Check if it's a domain error |
| 164 | + if e.is_instance_of::<pyo3::exceptions::PyValueError>(py) { |
| 165 | + panic!( |
| 166 | + "{func_name}({re}, {im}): py raised ValueError but rs=({}, {})", |
| 167 | + rs.re, rs.im |
| 168 | + ); |
| 169 | + } else if e.is_instance_of::<pyo3::exceptions::PyOverflowError>(py) { |
| 170 | + panic!( |
| 171 | + "{func_name}({re}, {im}): py raised OverflowError but rs=({}, {})", |
| 172 | + rs.re, rs.im |
| 173 | + ); |
| 174 | + } |
| 175 | + } |
| 176 | + // Both raised errors - OK |
| 177 | + } |
| 178 | + } |
| 179 | + }); |
| 180 | + } |
| 181 | +} |
0 commit comments