idsp/
complex.rs

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

idsp/
complex.rs
