idsp/
complex.rs

1use super::{atan2, cossin};
2use core::num::Wrapping;
3use core::ops::{Add, Deref, DerefMut, Mul, Sub};
4use dsp_fixedpoint::{P32, Q32};
5use num_traits::AsPrimitive;
6
7/// A complex number in cartesian coordinates
8#[derive(
9    Clone,
10    Copy,
11    Debug,
12    Default,
13    PartialEq,
14    Eq,
15    PartialOrd,
16    Ord,
17    serde::Serialize,
18    serde::Deserialize,
19    bytemuck::Zeroable,
20    bytemuck::Pod,
21)]
22#[repr(transparent)]
23#[serde(transparent)]
24pub struct Complex<T>(
25    /// Real and imaginary parts
26    pub [T; 2],
27);
28
29impl<T> Deref for Complex<T> {
30    type Target = [T; 2];
31
32    fn deref(&self) -> &Self::Target {
33        &self.0
34    }
35}
36
37impl<T> DerefMut for Complex<T> {
38    fn deref_mut(&mut self) -> &mut Self::Target {
39        &mut self.0
40    }
41}
42
43impl<T: Copy> Complex<T> {
44    /// Create a new `Complex<T>`
45    pub const fn new(re: T, im: T) -> Self {
46        Self([re, im])
47    }
48
49    /// The real part
50    pub fn re(&self) -> T {
51        self.0[0]
52    }
53
54    /// The imaginary part
55    pub fn im(&self) -> T {
56        self.0[1]
57    }
58}
59
60impl<T: Copy + core::ops::Neg<Output = T>> Complex<T> {
61    /// Conjugate
62    pub fn conj(self) -> Self {
63        Self([self.0[0], -self.0[1]])
64    }
65}
66
67macro_rules! fwd_binop {
68    ($tr:ident::$meth:ident) => {
69        impl<T: Copy + core::ops::$tr<Output = T>> core::ops::$tr for Complex<T> {
70            type Output = Self;
71            fn $meth(self, rhs: Self) -> Self {
72                Self([self.0[0].$meth(rhs.0[0]), self.0[1].$meth(rhs.0[1])])
73            }
74        }
75    };
76}
77fwd_binop!(Add::add);
78fwd_binop!(Sub::sub);
79fwd_binop!(BitAnd::bitand);
80fwd_binop!(BitOr::bitor);
81fwd_binop!(BitXor::bitxor);
82
83macro_rules! fwd_binop_inner {
84    ($tr:ident::$meth:ident) => {
85        impl<T: Copy + core::ops::$tr<Output = T>> core::ops::$tr<T> for Complex<T> {
86            type Output = Self;
87            fn $meth(self, rhs: T) -> Self {
88                Self([self.0[0].$meth(rhs), self.0[1].$meth(rhs)])
89            }
90        }
91    };
92}
93fwd_binop_inner!(Mul::mul);
94fwd_binop_inner!(Div::div);
95fwd_binop_inner!(Rem::rem);
96fwd_binop_inner!(BitAnd::bitand);
97fwd_binop_inner!(BitOr::bitor);
98fwd_binop_inner!(BitXor::bitxor);
99
100macro_rules! fwd_unop {
101    ($tr:ident::$meth:ident) => {
102        impl<T: Copy + core::ops::$tr<Output = T>> core::ops::$tr for Complex<T> {
103            type Output = Self;
104            fn $meth(self) -> Self {
105                Self([self.0[0].$meth(), self.0[1].$meth()])
106            }
107        }
108    };
109}
110fwd_unop!(Not::not);
111fwd_unop!(Neg::neg);
112
113impl<T: 'static + Copy + Mul<Output = A>, A: Add<Output = A> + Sub<Output = A> + AsPrimitive<T>> Mul
114    for Complex<T>
115{
116    type Output = Self;
117    fn mul(self, rhs: Self) -> Self {
118        Self([
119            (self.0[0] * rhs.0[0] - self.0[1] * rhs.0[1]).as_(),
120            (self.0[0] * rhs.0[1] + self.0[1] * rhs.0[0]).as_(),
121        ])
122    }
123}
124
125impl<T> core::iter::Sum for Complex<T>
126where
127    Self: Default + Add<Output = Self>,
128{
129    fn sum<I: Iterator<Item = Self>>(iter: I) -> Self {
130        iter.fold(Default::default(), |c, i| c + i)
131    }
132}
133
134impl<T> core::iter::Product for Complex<T>
135where
136    Self: Default + Mul<Output = Self>,
137{
138    fn product<I: Iterator<Item = Self>>(iter: I) -> Self {
139        iter.fold(Default::default(), |c, i| c * i)
140    }
141}
142
143impl Complex<Q32<31>> {
144    /// Return a Complex on the unit circle given an angle.
145    ///
146    /// Example:
147    ///
148    /// ```
149    /// use core::num::Wrapping as W;
150    /// use idsp::Complex;
151    /// Complex::<_>::from_angle(W(0));
152    /// Complex::<_>::from_angle(W(1 << 30)); // pi/2
153    /// Complex::<_>::from_angle(W(-1 << 30)); // -pi/2
154    /// ```
155    pub fn from_angle(angle: Wrapping<i32>) -> Self {
156        let (c, s) = cossin(angle.0);
157        Self::new(Q32::new(c), Q32::new(s))
158    }
159}
160
161impl Complex<i32> {
162    /// Return the absolute square (the squared magnitude).
163    ///
164    /// Note(panic): This will panic for `Complex(i32::MIN, i32::MIN)`
165    ///
166    /// Example:
167    ///
168    /// ```
169    /// use dsp_fixedpoint::{P32, Q32};
170    /// use idsp::Complex;
171    /// assert_eq!(Complex::new(i32::MIN, 0).norm_sqr(), P32::new(1 << 31));
172    /// assert_eq!(
173    ///     Complex::new(i32::MAX, i32::MAX).norm_sqr(),
174    ///     P32::new(u32::MAX - 3)
175    /// );
176    /// assert_eq!(
177    ///     Complex::new(i32::MIN, i32::MAX).norm_sqr(),
178    ///     P32::new(u32::MAX - 1)
179    /// );
180    /// ```
181    pub fn norm_sqr(&self) -> P32<31> {
182        let [x, y] = self.0.map(|x| x as i64 * x as i64);
183        P32::new(((x + y) >> 31) as _)
184    }
185
186    /// trunc(log2(power)) re full scale (approximation)
187    ///
188    /// TODO: scale up, interpolate
189    ///
190    /// Example:
191    ///
192    /// ```
193    /// use idsp::Complex;
194    /// assert_eq!(Complex::new(i32::MIN, i32::MIN).log2(), 0);
195    /// assert_eq!(Complex::new(i32::MAX, i32::MAX).log2(), -1);
196    /// assert_eq!(Complex::new(i32::MIN, 0).log2(), -1);
197    /// assert_eq!(Complex::new(i32::MAX, 0).log2(), -2);
198    /// assert_eq!(Complex::new(-1, 0).log2(), -63);
199    /// assert_eq!(Complex::new(1, 0).log2(), -63);
200    /// assert_eq!(Complex::new(0, 0).log2(), -64);
201    /// ```
202    pub fn log2(&self) -> i32 {
203        let [x, y] = self.0.map(|x| x as i64 * x as i64);
204        -(x.wrapping_add(y).leading_zeros() as i32)
205    }
206
207    /// Return the angle.
208    ///
209    /// Note: Normalization is `1 << 31 == pi`.
210    ///
211    /// Example:
212    ///
213    /// ```
214    /// use core::num::Wrapping as W;
215    /// use idsp::Complex;
216    /// assert_eq!(Complex::new(0, 0).arg(), W(0));
217    /// assert_eq!(Complex::new(0, 1).arg(), W((1 << 30) - 1));
218    /// ```
219    pub fn arg(&self) -> Wrapping<i32> {
220        Wrapping(atan2(self.im(), self.re()))
221    }
222}
idsp/complex.rs

idsp/
complex.rs
