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}