mod float;
mod ops;
pub use float::Float;
use core::fmt;
#[derive(Clone, Copy, Debug, Default, PartialEq)]
#[repr(C)]
pub struct Complex<T> {
pub re: T,
pub im: T,
}
impl<T> Complex<T> {
#[inline]
pub const fn new(re: T, im: T) -> Self {
Complex { re, im }
}
}
impl<T: Float> Complex<T> {
pub const ZERO: Self = Complex::new(T::ZERO, T::ZERO);
pub const ONE: Self = Complex::new(T::ONE, T::ZERO);
pub const I: Self = Complex::new(T::ZERO, T::ONE);
#[inline]
pub fn from_polar(r: T, theta: T) -> Self {
let (sin, cos) = theta.sin_cos();
Complex::new(r * cos, r * sin)
}
#[inline]
pub fn to_polar(self) -> (T, T) {
(self.abs(), self.arg())
}
#[inline]
pub fn conj(self) -> Self {
Complex::new(self.re, -self.im)
}
#[inline]
pub fn abs(self) -> T {
self.re.hypot(self.im)
}
#[inline]
pub fn norm_sqr(self) -> T {
self.re * self.re + self.im * self.im
}
#[inline]
pub fn arg(self) -> T {
self.im.atan2(self.re)
}
#[inline]
pub fn recip(self) -> Self {
Self::ONE / self
}
#[inline]
pub fn is_finite(self) -> bool {
self.re.is_finite() && self.im.is_finite()
}
#[inline]
pub fn is_nan(self) -> bool {
self.re.is_nan() || self.im.is_nan()
}
pub fn exp(self) -> Self {
let r = self.re.exp();
if self.im == T::ZERO {
return Complex::new(r, self.im);
}
let (sin, cos) = self.im.sin_cos();
if r.is_infinite() && self.re.is_finite() {
let h = (self.re * T::HALF).exp();
return Complex::new(h * cos * h, h * sin * h);
}
Complex::new(r * cos, r * sin)
}
pub fn ln(self) -> Self {
let (a, b) = (self.re.abs(), self.im.abs());
let (big, small) = if a >= b { (a, b) } else { (b, a) };
let re = if big > T::HALF && big < T::TWO {
((big - T::ONE) * (big + T::ONE) + small * small).ln_1p() * T::HALF
} else {
let m = self.abs();
if m.is_finite() && m >= T::MIN_POSITIVE {
m.ln()
} else {
Self::ln_of_extreme(m, big, small)
}
};
Complex::new(re, self.arg())
}
#[cold]
#[inline(never)]
fn ln_of_extreme(m: T, big: T, small: T) -> T {
if big == T::ZERO || !big.is_finite() || !small.is_finite() {
return m.ln();
}
let r = small / big;
big.ln() + (r * r).ln_1p() * T::HALF
}
pub fn sqrt(self) -> Self {
let (re, im) = (self.re, self.im);
if re == T::ZERO && im == T::ZERO {
return Complex::new(T::ZERO, im);
}
if im.is_infinite() {
return Complex::new(T::INFINITY, im);
}
if re.is_infinite() {
let zero_or_nan = if im.is_nan() {
im
} else {
T::ZERO.copysign(im)
};
return if re > T::ZERO {
Complex::new(re, zero_or_nan)
} else {
Complex::new(zero_or_nan.abs(), T::INFINITY.copysign(im))
};
}
if re.is_nan() || im.is_nan() {
return Complex::new(T::NAN, T::NAN);
}
let big = re.abs().max(im.abs());
if big < T::from_f64(4.0) * T::MIN_POSITIVE {
let up = T::from_f64(18_446_744_073_709_551_616.0);
return (self * up).sqrt() / T::from_f64(4_294_967_296.0);
}
if big > T::MAX / T::from_f64(4.0) {
return (self * T::from_f64(0.25)).sqrt() * T::TWO;
}
let t = ((re.abs() + self.abs()) * T::HALF).sqrt();
if re >= T::ZERO {
Complex::new(t, im / (T::TWO * t))
} else {
Complex::new(im.abs() / (T::TWO * t), t.copysign(im))
}
}
pub fn powi(self, n: i32) -> Self {
let mut base = if n < 0 { self.recip() } else { self };
let mut n = n.unsigned_abs();
let mut r = Self::ONE;
while n > 0 {
if n & 1 == 1 {
r *= base;
}
n >>= 1;
if n > 0 {
base *= base;
}
}
r
}
pub fn powf(self, x: T) -> Self {
if self.re == T::ZERO && self.im == T::ZERO {
return if x == T::ZERO { Self::ONE } else { Self::ZERO };
}
(self.ln() * x).exp()
}
pub fn powc(self, w: Self) -> Self {
if self.re == T::ZERO && self.im == T::ZERO {
return if w == Self::ZERO {
Self::ONE
} else {
Self::ZERO
};
}
(self.ln() * w).exp()
}
}
impl<T: Float> From<T> for Complex<T> {
#[inline]
fn from(re: T) -> Self {
Complex::new(re, T::ZERO)
}
}
impl<T: Float> fmt::Display for Complex<T> {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
let sign = if self.im.is_sign_negative() { '-' } else { '+' };
match f.precision() {
Some(p) => write!(f, "{:.*}{}{:.*}i", p, self.re, sign, p, self.im.abs()),
None => write!(f, "{}{}{}i", self.re, sign, self.im.abs()),
}
}
}