idsp/
complex.rs

1pub use num_complex::Complex;
2
3use super::{atan2, cossin};
4
5/// Complex extension trait offering DSP (fast, good accuracy) functionality.
6pub trait ComplexExt<T, U> {
7    /// Unit magnitude from angle
8    fn from_angle(angle: T) -> Self;
9    /// Square of magnitude
10    fn abs_sqr(&self) -> U;
11    /// Log2 approximation
12    fn log2(&self) -> T;
13    /// Angle
14    fn arg(&self) -> T;
15    /// Staturating addition
16    fn saturating_add(&self, other: Self) -> Self;
17    /// Saturating subtraction
18    fn saturating_sub(&self, other: Self) -> Self;
19}
20
21impl ComplexExt<i32, u32> for Complex<i32> {
22    /// Return a Complex on the unit circle given an angle.
23    ///
24    /// Example:
25    ///
26    /// ```
27    /// use idsp::{Complex, ComplexExt};
28    /// Complex::<i32>::from_angle(0);
29    /// Complex::<i32>::from_angle(1 << 30); // pi/2
30    /// Complex::<i32>::from_angle(-1 << 30); // -pi/2
31    /// ```
32    fn from_angle(angle: i32) -> Self {
33        let (c, s) = cossin(angle);
34        Self::new(c, s)
35    }
36
37    /// Return the absolute square (the squared magnitude).
38    ///
39    /// Note: Normalization is `1 << 32`, i.e. U0.32.
40    ///
41    /// Note(panic): This will panic for `Complex(i32::MIN, i32::MIN)`
42    ///
43    /// Example:
44    ///
45    /// ```
46    /// use idsp::{Complex, ComplexExt};
47    /// assert_eq!(Complex::new(i32::MIN, 0).abs_sqr(), 1 << 31);
48    /// assert_eq!(Complex::new(i32::MAX, i32::MAX).abs_sqr(), u32::MAX - 3);
49    /// ```
50    fn abs_sqr(&self) -> u32 {
51        (((self.re as i64) * (self.re as i64) + (self.im as i64) * (self.im as i64)) >> 31) as u32
52    }
53
54    /// log2(power) re full scale approximation
55    ///
56    /// TODO: scale up, interpolate
57    ///
58    /// Panic:
59    /// This will panic for `Complex(i32::MIN, i32::MIN)`
60    ///
61    /// Example:
62    ///
63    /// ```
64    /// use idsp::{Complex, ComplexExt};
65    /// assert_eq!(Complex::new(i32::MAX, i32::MAX).log2(), -1);
66    /// assert_eq!(Complex::new(i32::MAX, 0).log2(), -2);
67    /// assert_eq!(Complex::new(1, 0).log2(), -63);
68    /// assert_eq!(Complex::new(0, 0).log2(), -64);
69    /// ```
70    fn log2(&self) -> i32 {
71        let a = (self.re as i64) * (self.re as i64) + (self.im as i64) * (self.im as i64);
72        -(a.leading_zeros() as i32)
73    }
74
75    /// Return the angle.
76    ///
77    /// Note: Normalization is `1 << 31 == pi`.
78    ///
79    /// Example:
80    ///
81    /// ```
82    /// use idsp::{Complex, ComplexExt};
83    /// assert_eq!(Complex::new(0, 0).arg(), 0);
84    /// ```
85    fn arg(&self) -> i32 {
86        atan2(self.im, self.re)
87    }
88
89    fn saturating_add(&self, other: Self) -> Self {
90        Self::new(
91            self.re.saturating_add(other.re),
92            self.im.saturating_add(other.im),
93        )
94    }
95
96    fn saturating_sub(&self, other: Self) -> Self {
97        Self::new(
98            self.re.saturating_sub(other.re),
99            self.im.saturating_sub(other.im),
100        )
101    }
102}
103
104/// Full scale fixed point multiplication.
105pub trait MulScaled<T> {
106    /// Scaled multiplication for fixed point
107    fn mul_scaled(self, other: T) -> Self;
108}
109
110impl MulScaled<Complex<i32>> for Complex<i32> {
111    fn mul_scaled(self, other: Self) -> Self {
112        let a = self.re as i64;
113        let b = self.im as i64;
114        let c = other.re as i64;
115        let d = other.im as i64;
116        Complex {
117            re: ((a * c - b * d) >> 31) as i32,
118            im: ((b * c + a * d) >> 31) as i32,
119        }
120    }
121}
122
123impl MulScaled<i32> for Complex<i32> {
124    fn mul_scaled(self, other: i32) -> Self {
125        Complex {
126            re: ((other as i64 * self.re as i64) >> 31) as i32,
127            im: ((other as i64 * self.im as i64) >> 31) as i32,
128        }
129    }
130}
131
132impl MulScaled<i16> for Complex<i32> {
133    fn mul_scaled(self, other: i16) -> Self {
134        Complex {
135            re: (other as i32 * (self.re >> 16) + (1 << 14)) >> 15,
136            im: (other as i32 * (self.im >> 16) + (1 << 14)) >> 15,
137        }
138    }
139}