1use super::{Float, IterScan, One, Zero};
2use std::{
3 cmp::Ordering,
4 iter::{Product, Sum},
5 ops::{Add, AddAssign, Div, DivAssign, Mul, MulAssign, Neg, Sub, SubAssign},
6};
7
8#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, PartialOrd, Ord, Hash)]
9pub struct Complex<T> {
10 pub re: T,
11 pub im: T,
12}
13
14impl<T> Complex<T> {
15 pub fn new(re: T, im: T) -> Self {
16 Self { re, im }
17 }
18 pub fn transpose(self) -> Self {
19 Self::new(self.im, self.re)
20 }
21 pub fn map<U>(self, mut f: impl FnMut(T) -> U) -> Complex<U> {
22 Complex::new(f(self.re), f(self.im))
23 }
24}
25impl<T> Zero for Complex<T>
26where
27 T: Zero,
28{
29 fn zero() -> Self {
30 Self::new(T::zero(), T::zero())
31 }
32}
33impl<T> One for Complex<T>
34where
35 T: Zero + One,
36{
37 fn one() -> Self {
38 Self::new(T::one(), T::zero())
39 }
40}
41impl<T> Complex<T>
42where
43 T: Zero + One,
44{
45 pub fn i() -> Self {
46 Self::new(T::zero(), T::one())
47 }
48}
49impl<T> Complex<T>
50where
51 T: Neg<Output = T>,
52{
53 pub fn conjugate(self) -> Self {
54 Self::new(self.re, -self.im)
55 }
56}
57impl<T> Complex<T>
58where
59 T: Mul,
60 <T as Mul>::Output: Add,
61{
62 pub fn dot(self, rhs: Self) -> <<T as Mul>::Output as Add>::Output {
63 self.re * rhs.re + self.im * rhs.im
64 }
65}
66impl<T> Complex<T>
67where
68 T: Mul,
69 <T as Mul>::Output: Sub,
70{
71 pub fn cross(self, rhs: Self) -> <<T as Mul>::Output as Sub>::Output {
72 self.re * rhs.im - self.im * rhs.re
73 }
74}
75impl<T> Complex<T>
76where
77 T: Mul + Clone,
78 <T as Mul>::Output: Add,
79{
80 pub fn norm(self) -> <<T as Mul>::Output as Add>::Output {
81 self.re.clone() * self.re + self.im.clone() * self.im
82 }
83}
84impl<T> Complex<T>
85where
86 T: Zero + Ord + Mul,
87 <T as Mul>::Output: Ord,
88{
89 pub fn cmp_by_arg(self, other: Self) -> Ordering {
90 fn pos<T>(c: &Complex<T>) -> bool
91 where
92 T: Zero + Ord,
93 {
94 let zero = T::zero();
95 c.im < zero || c.im <= zero && c.re < zero
96 }
97 pos(&self)
98 .cmp(&pos(&other))
99 .then_with(|| (self.re * other.im).cmp(&(self.im * other.re)).reverse())
100 }
101}
102impl<T> Complex<T>
103where
104 T: Float,
105{
106 pub fn polar(r: T, theta: T) -> Self {
107 Self::new(r * theta.cos(), r * theta.sin())
108 }
109 pub fn primitive_nth_root_of_unity(n: T) -> Self {
110 let theta = T::TAU / n;
111 Self::new(theta.cos(), theta.sin())
112 }
113 pub fn abs(self) -> T {
114 self.re.hypot(self.im)
115 }
116 pub fn unit(self) -> Self {
117 self / self.abs()
118 }
119 pub fn angle(self) -> T {
120 self.im.atan2(self.re)
121 }
122}
123impl<T> Add for Complex<T>
124where
125 T: Add,
126{
127 type Output = Complex<<T as Add>::Output>;
128 fn add(self, rhs: Self) -> Self::Output {
129 Complex::new(self.re + rhs.re, self.im + rhs.im)
130 }
131}
132impl<T> Add<T> for Complex<T>
133where
134 T: Add<Output = T>,
135{
136 type Output = Self;
137 fn add(self, rhs: T) -> Self::Output {
138 Self::new(self.re + rhs, self.im)
139 }
140}
141impl<T> Sub for Complex<T>
142where
143 T: Sub,
144{
145 type Output = Complex<<T as Sub>::Output>;
146 fn sub(self, rhs: Self) -> Self::Output {
147 Complex::new(self.re - rhs.re, self.im - rhs.im)
148 }
149}
150impl<T> Sub<T> for Complex<T>
151where
152 T: Sub<Output = T>,
153{
154 type Output = Self;
155 fn sub(self, rhs: T) -> Self::Output {
156 Self::new(self.re - rhs, self.im)
157 }
158}
159impl<T, U> Mul for Complex<T>
160where
161 T: Clone + Mul,
162 <T as Mul>::Output: Add<Output = U> + Sub<Output = U>,
163{
164 type Output = Complex<U>;
165 fn mul(self, rhs: Self) -> Self::Output {
166 Complex::new(
167 self.re.clone() * rhs.re.clone() - self.im.clone() * rhs.im.clone(),
168 self.re * rhs.im + self.im * rhs.re,
169 )
170 }
171}
172impl<T> Mul<T> for Complex<T>
173where
174 T: Clone + Mul,
175{
176 type Output = Complex<<T as Mul>::Output>;
177 fn mul(self, rhs: T) -> Self::Output {
178 Complex::new(self.re * rhs.clone(), self.im * rhs)
179 }
180}
181impl<T> Div for Complex<T>
182where
183 T: Clone + Add<Output = T> + Sub<Output = T> + Mul<Output = T> + Div,
184{
185 type Output = Complex<<T as Div>::Output>;
186 fn div(self, rhs: Self) -> Self::Output {
187 let d = rhs.re.clone() * rhs.re.clone() + rhs.im.clone() * rhs.im.clone();
188 Complex::new(
189 (self.re.clone() * rhs.re.clone() + self.im.clone() * rhs.im.clone()) / d.clone(),
190 (self.im * rhs.re - self.re * rhs.im) / d,
191 )
192 }
193}
194impl<T> Div<T> for Complex<T>
195where
196 T: Clone + Div,
197{
198 type Output = Complex<<T as Div>::Output>;
199 fn div(self, rhs: T) -> Self::Output {
200 Complex::new(self.re / rhs.clone(), self.im / rhs)
201 }
202}
203impl<T> Neg for Complex<T>
204where
205 T: Neg,
206{
207 type Output = Complex<<T as Neg>::Output>;
208 fn neg(self) -> Self::Output {
209 Complex::new(-self.re, -self.im)
210 }
211}
212macro_rules! impl_complex_ref_binop {
213 (impl<$T:ident> $imp:ident $method:ident ($l:ty, $r:ty) where $($w:ident)*) => {
214 impl<$T> $imp<$r> for &$l
215 where
216 $T: Clone $(+ $w<Output = $T>)*,
217 {
218 type Output = <$l as $imp<$r>>::Output;
219 fn $method(self, rhs: $r) -> <$l as $imp<$r>>::Output {
220 $imp::$method(self.clone(), rhs)
221 }
222 }
223 impl<$T> $imp<&$r> for $l
224 where
225 $T: Clone $(+ $w<Output = $T>)*,
226 {
227 type Output = <$l as $imp<$r>>::Output;
228 fn $method(self, rhs: &$r) -> <$l as $imp<$r>>::Output {
229 $imp::$method(self, rhs.clone())
230 }
231 }
232 impl<$T> $imp<&$r> for &$l
233 where
234 $T: Clone $(+ $w<Output = $T>)*,
235 {
236 type Output = <$l as $imp<$r>>::Output;
237 fn $method(self, rhs: &$r) -> <$l as $imp<$r>>::Output {
238 $imp::$method(self.clone(), rhs.clone())
239 }
240 }
241 };
242}
243impl_complex_ref_binop!(impl<T> Add add (Complex<T>, Complex<T>) where Add);
244impl_complex_ref_binop!(impl<T> Add add (Complex<T>, T) where Add);
245impl_complex_ref_binop!(impl<T> Sub sub (Complex<T>, Complex<T>) where Sub);
246impl_complex_ref_binop!(impl<T> Sub sub (Complex<T>, T) where Sub);
247impl_complex_ref_binop!(impl<T> Mul mul (Complex<T>, Complex<T>) where Add Sub Mul);
248impl_complex_ref_binop!(impl<T> Mul mul (Complex<T>, T) where Mul);
249impl_complex_ref_binop!(impl<T> Div div (Complex<T>, Complex<T>) where Add Sub Mul Div);
250impl_complex_ref_binop!(impl<T> Div div (Complex<T>, T) where Div);
251macro_rules! impl_complex_ref_unop {
252 (impl<$T:ident> $imp:ident $method:ident ($t:ty) where $($w:ident)*) => {
253 impl<$T> $imp for &$t
254 where
255 $T: Clone $(+ $w<Output = $T>)*,
256 {
257 type Output = <$t as $imp>::Output;
258 fn $method(self) -> <$t as $imp>::Output {
259 $imp::$method(self.clone())
260 }
261 }
262 };
263}
264impl_complex_ref_unop!(impl<T> Neg neg (Complex<T>) where Neg);
265macro_rules! impl_complex_op_assign {
266 (impl<$T:ident> $imp:ident $method:ident ($l:ty, $r:ty) $fromimp:ident $frommethod:ident where $($w:ident)*) => {
267 impl<$T> $imp<$r> for $l
268 where
269 $T: Clone $(+ $w<Output = $T>)*,
270 {
271 fn $method(&mut self, rhs: $r) {
272 *self = $fromimp::$frommethod(self.clone(), rhs);
273 }
274 }
275 impl<$T> $imp<&$r> for $l
276 where
277 $T: Clone $(+ $w<Output = $T>)*,
278 {
279 fn $method(&mut self, rhs: &$r) {
280 $imp::$method(self, rhs.clone());
281 }
282 }
283 };
284}
285impl_complex_op_assign!(impl<T> AddAssign add_assign (Complex<T>, Complex<T>) Add add where Add);
286impl_complex_op_assign!(impl<T> AddAssign add_assign (Complex<T>, T) Add add where Add);
287impl_complex_op_assign!(impl<T> SubAssign sub_assign (Complex<T>, Complex<T>) Sub sub where Sub);
288impl_complex_op_assign!(impl<T> SubAssign sub_assign (Complex<T>, T) Sub sub where Sub);
289impl_complex_op_assign!(impl<T> MulAssign mul_assign (Complex<T>, Complex<T>) Mul mul where Add Sub Mul);
290impl_complex_op_assign!(impl<T> MulAssign mul_assign (Complex<T>, T) Mul mul where Mul);
291impl_complex_op_assign!(impl<T> DivAssign div_assign (Complex<T>, Complex<T>) Div div where Add Sub Mul Div);
292impl_complex_op_assign!(impl<T> DivAssign div_assign (Complex<T>, T) Div div where Div);
293macro_rules! impl_complex_fold {
294 (impl<$T:ident> $imp:ident $method:ident ($t:ty) $identimp:ident $identmethod:ident $fromimp:ident $frommethod:ident where $($w:ident)* $(+ $x:ident)*) => {
295 impl<$T> $imp for $t
296 where
297 $T: $identimp $(+ $w<Output = $T>)* $(+ $x)*,
298 {
299 fn $method<I: Iterator<Item = Self>>(iter: I) -> Self {
300 iter.fold(<Self as $identimp>::$identmethod(), $fromimp::$frommethod)
301 }
302 }
303 impl<'a, $T: 'a> $imp<&'a $t> for $t
304 where
305 $T: Clone + $identimp $(+ $w<Output = $T>)* $(+ $x)*,
306 {
307 fn $method<I: Iterator<Item = &'a $t>>(iter: I) -> Self {
308 iter.fold(<Self as $identimp>::$identmethod(), $fromimp::$frommethod)
309 }
310 }
311 };
312}
313impl_complex_fold!(impl<T> Sum sum (Complex<T>) Zero zero Add add where Add);
314impl_complex_fold!(impl<T> Product product (Complex<T>) One one Mul mul where Add Sub Mul + Zero + Clone);
315
316impl<T: IterScan> IterScan for Complex<T> {
317 type Output = Complex<<T as IterScan>::Output>;
318 fn scan<'a, I: Iterator<Item = &'a str>>(iter: &mut I) -> Option<Self::Output> {
319 Some(Complex::new(
320 <T as IterScan>::scan(iter)?,
321 <T as IterScan>::scan(iter)?,
322 ))
323 }
324}