use num_traits::Num;
use num_traits::NumCast;
use std::ops::Div;
#[derive(Copy, Clone)]
pub struct ComplexNum<T: Num + Copy> {
pub real: T,
pub imaginary: T,
}
impl<T> ComplexNum<T>
where
T: Copy + Clone + NumCast + Num,
{
pub fn new(real: T, imaginary: T) -> Self {
ComplexNum { real, imaginary }
}
pub fn conjugate(self) -> Self {
ComplexNum::new(self.real, T::zero() - self.imaginary)
}
pub fn as_tuple(self) -> (T, T) {
(self.real, self.imaginary)
}
pub fn from_polar(polar: ComplexPolar<T>) -> Option<Self> {
let radius = polar.radius.to_f64()?;
let angle = polar.angle.to_f64()?;
let real_f64 = radius * angle.cos();
let imaginary_f64 = radius * angle.sin();
Some(ComplexNum {
real: NumCast::from(real_f64)?,
imaginary: NumCast::from(imaginary_f64)?,
})
}
pub fn as_polar(&self) -> Option<ComplexPolar<T>> {
let a = self.real.to_f64()?;
let b = self.imaginary.to_f64()?;
let radius = (a * a + b * b).sqrt();
let angle = b.atan2(a);
Some(ComplexPolar {
radius: NumCast::from(radius)?,
angle: NumCast::from(angle)?,
})
}
pub fn safe_div(self, rhs: Self) -> Self {
let denom_mag_f64 = {
let real_f64 = rhs.real.to_f64().unwrap_or(0.0);
let imag_f64 = rhs.imaginary.to_f64().unwrap_or(0.0);
real_f64 * real_f64 + imag_f64 * imag_f64
};
if denom_mag_f64.abs() < 1e-12 {
let inf = T::from(f64::INFINITY).unwrap_or(rhs.real); ComplexNum::new(inf, inf)
} else {
let numer = self * rhs.conjugate();
ComplexNum::new(
numer.real / rhs.real * rhs.real + rhs.imaginary * rhs.imaginary,
numer.imaginary / rhs.real * rhs.real + rhs.imaginary * rhs.imaginary,
)
}
}
}
#[derive(Copy, Clone)]
pub struct ComplexPolar<T: Num + Copy> {
radius: T,
angle: T,
}
impl<T> ComplexPolar<T>
where
T: Num + Copy,
{
pub fn get_radius(&self) -> T {
self.radius
}
pub fn get_angle(&self) -> T {
self.angle
}
}
impl<T> std::ops::Add for ComplexNum<T>
where
T: Num + Copy + Clone + NumCast,
{
type Output = ComplexNum<T>;
fn add(self, rhs: Self) -> Self::Output {
ComplexNum::new(self.real + rhs.real, self.imaginary + rhs.imaginary)
}
}
impl<T> std::ops::Sub for ComplexNum<T>
where
T: Num + Copy + Clone + NumCast,
{
type Output = ComplexNum<T>;
fn sub(self, rhs: Self) -> Self::Output {
ComplexNum::new(self.real - rhs.real, self.imaginary - rhs.imaginary)
}
}
impl<T> std::ops::Mul for ComplexNum<T>
where
T: Num + Copy + Clone + NumCast,
{
type Output = ComplexNum<T>;
fn mul(self, rhs: Self) -> Self::Output {
ComplexNum::new(
self.real * rhs.real - self.imaginary * rhs.imaginary,
self.real * rhs.imaginary + self.imaginary * rhs.real,
)
}
}
impl<T> std::ops::Div for ComplexNum<T>
where
T: Num + Copy + Div<Output = T> + Clone + NumCast,
{
type Output = ComplexNum<T>;
fn div(self, rhs: Self) -> Self::Output {
let denom = rhs.real * rhs.real + rhs.imaginary * rhs.imaginary;
let numer = self * rhs.conjugate();
ComplexNum::new(numer.real / denom, numer.imaginary / denom)
}
}