DFT Trait

src/complex.rs

@@ -1,22 +1,22 @@
-use std::{fmt::Display, ops::{Add, Div, Mul, Neg, Sub}};
+use std::{
+    f32::consts::PI,
+    fmt::Display,
+    ops::{Add, Div, Mul, Neg, Sub},
+};
 
 #[derive(Copy, Clone, Debug)]
-pub struct Complex<T>
-{
+pub struct Complex<T> {
     pub re: T,
-    pub im: T
+    pub im: T,
 }
 
-impl<T> Complex<T>
-{
-    pub fn new(re: T, im: T) -> Self
-    {
+impl<T> Complex<T> {
+    pub fn new(re: T, im: T) -> Self {
         Complex { re, im }
     }
 }
 
-impl<T: Clone + Add<Output = T>> Add<Complex<T>> for Complex<T>
-{
+impl<T: Clone + Add<Output = T>> Add<Complex<T>> for Complex<T> {
     type Output = Complex<T>;
 
     #[inline]
@@ -25,8 +25,7 @@ impl<T: Clone + Add<Output = T>> Add<Complex<T>> for Complex<T>
     }
 }
 
-impl<T: Clone + Sub<Output = T>> Sub<Complex<T>> for Complex<T>
-{
+impl<T: Clone + Sub<Output = T>> Sub<Complex<T>> for Complex<T> {
     type Output = Complex<T>;
 
     #[inline]
@@ -35,8 +34,7 @@ impl<T: Clone + Sub<Output = T>> Sub<Complex<T>> for Complex<T>
     }
 }
 
-impl<T: Clone + Add<Output = T>> Add<T> for Complex<T>
-{
+impl<T: Clone + Add<Output = T>> Add<T> for Complex<T> {
     type Output = Complex<T>;
 
     #[inline]
@@ -45,18 +43,21 @@ impl<T: Clone + Add<Output = T>> Add<T> for Complex<T>
     }
 }
 
-impl<T: Clone + Add<Output = T> + Mul<Output = T> + Sub<Output = T>> Mul<Complex<T>> for Complex<T>
+impl<T: Clone + Add<Output = T> + Mul<Output = T> + Sub<Output = T>> Mul<Complex<T>>
+    for Complex<T>
 {
     type Output = Complex<T>;
 
     #[inline]
     fn mul(self, rhs: Complex<T>) -> Self::Output {
-        self::Complex::new(self.re.clone() * rhs.re.clone() - self.im.clone() * rhs.im.clone(), self.re * rhs.im + self.im * rhs.re)
+        self::Complex::new(
+            self.re.clone() * rhs.re.clone() - self.im.clone() * rhs.im.clone(),
+            self.re * rhs.im + self.im * rhs.re,
+        )
     }
 }
 
-impl<T: Clone + Mul<Output = T>> Mul<T> for Complex<T>
-{
+impl<T: Clone + Mul<Output = T>> Mul<T> for Complex<T> {
     type Output = Complex<T>;
 
     #[inline]
@@ -65,8 +66,7 @@ impl<T: Clone + Mul<Output = T>> Mul<T> for Complex<T>
     }
 }
 
-impl<T: Clone + Neg<Output = T>> Neg for Complex<T>
-{
+impl<T: Clone + Neg<Output = T>> Neg for Complex<T> {
     type Output = Complex<T>;
 
     fn neg(self) -> Self::Output {
@@ -74,17 +74,42 @@ impl<T: Clone + Neg<Output = T>> Neg for Complex<T>
     }
 }
 
-impl<T: Clone + Neg<Output = T>> Complex<T>
-{
-    pub fn conj(&self) -> Complex<T>
-    {
+impl<T: Clone + Neg<Output = T>> Complex<T> {
+    pub fn conj(&self) -> Complex<T> {
         Self::new(self.re.clone(), -self.im.clone())
     }
 }
 
-impl<T: Display> Display for Complex<T>
-{
+impl<T: Display> Display for Complex<T> {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
         write!(f, "{} + i{}", self.re, self.im)
     }
 }
+
+pub type Complex32 = Complex<f32>;
+impl Complex32 {
+    pub fn zero() -> Self {
+        Complex { re: 0.0, im: 0.0 }
+    }
+
+    pub fn cexp(angle: f32) -> Complex32 {
+        Complex32 {
+            re: angle.cos(),
+            im: angle.sin(),
+        }
+    }
+
+    pub fn mag(&self) -> f32 {
+        (self.re * self.re + self.im * self.im).sqrt()
+    }
+
+    pub fn arg(&self) -> f32 {
+        if self.im == 0. {
+            if self.re >= 0. { 0. } else { PI }
+        } else if self.re == 0. {
+            if self.im >= 0. { PI / 2.0 } else { -PI / 2.0 }
+        } else {
+            (self.im / self.re).atan()
+        }
+    }
+}