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