1use core::fmt::{Debug, Display};
2use core::iter::{Product, Sum};
3use core::ops::{Add, AddAssign, Div, DivAssign, Mul, MulAssign, Neg, Sub, SubAssign};
4
5#[cfg(feature = "approx")]
6use approx::AbsDiffEq;
7
8use crate::{Float32, Float64};
9
10pub trait FloatTraitsForComplex:
11 Sized
12 + Copy
13 + Clone
14 + Add<Output = Self>
15 + AddAssign
16 + Sub<Output = Self>
17 + SubAssign
18 + Mul<Output = Self>
19 + MulAssign
20 + Div<Output = Self>
21 + DivAssign
22 + Neg<Output = Self>
23 + Sum
24 + Product
25{
26 fn sin(&self) -> Self;
27 fn cos(&self) -> Self;
28 fn atan(&self) -> Self;
29 fn sqrt(&self) -> Self;
30
31 const ZERO: Self;
32 const ONE: Self;
33
34 #[inline]
35 fn sin_cos(&self) -> (Self, Self) {
36 (self.sin(), self.cos())
37 }
38 #[inline]
39 fn square(&self) -> Self {
40 self.mul(*self)
41 }
42}
43
44impl FloatTraitsForComplex for Float32 {
45 #[inline]
46 fn cos(&self) -> Self {
47 Float32::cos(*self)
48 }
49
50 #[inline]
51 fn sin(&self) -> Self {
52 Float32::sin(*self)
53 }
54
55 #[inline]
56 fn sqrt(&self) -> Self {
57 Float32::sqrt(*self)
58 }
59
60 #[inline]
61 fn atan(&self) -> Self {
62 Float32::atan(*self)
63 }
64
65 const ZERO: Self = Float32::ZERO;
66 const ONE: Self = Float32::ONE;
67}
68
69impl FloatTraitsForComplex for Float64 {
70 #[inline]
71 fn cos(&self) -> Self {
72 Float64::cos(*self)
73 }
74
75 #[inline]
76 fn sin(&self) -> Self {
77 Float64::sin(*self)
78 }
79
80 #[inline]
81 fn sqrt(&self) -> Self {
82 Float64::sqrt(*self)
83 }
84
85 #[inline]
86 fn atan(&self) -> Self {
87 Float64::atan(*self)
88 }
89
90 const ZERO: Self = Float64::ZERO;
91 const ONE: Self = Float64::ONE;
92}
93
94#[derive(Clone, Copy, PartialEq)]
95pub struct Complex<F: FloatTraitsForComplex> {
96 pub real: F,
97 pub imag: F,
98}
99
100impl<F: FloatTraitsForComplex> Complex<F> {
101 #[inline]
102 pub fn new(real: F, imag: F) -> Self {
103 Self { imag, real }
104 }
105
106 #[inline]
107 pub fn from_len_angle(len: F, angle: F) -> Self {
108 let (sin, cos) = angle.sin_cos();
109 Self::new(len * sin, len * cos)
110 }
111
112 #[inline]
113 pub fn length(&self) -> F {
114 (self.imag.square() + self.real.square()).sqrt()
115 }
116
117 #[inline]
118 pub fn angle(&self) -> F {
119 (self.imag / self.real).atan()
120 }
121
122 #[inline]
123 pub fn to_len_angle(&self) -> (F, F) {
124 let angle = (self.imag / self.real).atan();
125 let len = (self.imag.square() + self.real.square()).sqrt();
126 (len, angle)
127 }
128
129 #[inline]
130 pub fn conjugate(&self) -> Self {
131 Self::new(self.real, -self.imag)
132 }
133}
134
135impl<F: FloatTraitsForComplex> From<(F, F)> for Complex<F> {
136 #[inline]
137 fn from(value: (F, F)) -> Self {
138 Self::new(value.0, value.1)
139 }
140}
141
142impl<F: FloatTraitsForComplex> Add for Complex<F> {
143 type Output = Self;
144
145 #[inline]
146 fn add(self, rhs: Self) -> Self::Output {
147 Self::new(self.real + rhs.real, self.imag + rhs.imag)
148 }
149}
150
151impl<F: FloatTraitsForComplex> AddAssign for Complex<F> {
152 #[inline]
153 fn add_assign(&mut self, rhs: Self) {
154 self.real.add_assign(rhs.real);
155 self.imag.add_assign(rhs.imag);
156 }
157}
158
159impl<F: FloatTraitsForComplex> Sub for Complex<F> {
160 type Output = Self;
161
162 #[inline]
163 fn sub(self, rhs: Self) -> Self::Output {
164 Self::new(self.real - rhs.real, self.imag - rhs.imag)
165 }
166}
167
168impl<F: FloatTraitsForComplex> SubAssign for Complex<F> {
169 #[inline]
170 fn sub_assign(&mut self, rhs: Self) {
171 self.real.sub_assign(rhs.real);
172 self.imag.sub_assign(rhs.imag);
173 }
174}
175
176impl<F: FloatTraitsForComplex> Mul for Complex<F> {
177 type Output = Self;
178
179 #[inline]
180 fn mul(self, rhs: Self) -> Self::Output {
181 Self::new(
182 self.real * rhs.real - self.imag * rhs.imag,
183 self.real * rhs.imag + self.imag * rhs.real,
184 )
185 }
186}
187
188impl<F: FloatTraitsForComplex> MulAssign for Complex<F> {
189 #[inline]
190 fn mul_assign(&mut self, rhs: Self) {
191 *self = self.mul(rhs);
192 }
193}
194
195impl<F: FloatTraitsForComplex> Div for Complex<F> {
196 type Output = Self;
197
198 #[inline]
199 fn div(self, rhs: Self) -> Self::Output {
200 let x = rhs.real.square() + rhs.imag.square();
201 Self::new(
202 (self.real * rhs.real + self.imag * rhs.imag) / x,
203 (self.imag * rhs.real - self.real * rhs.imag) / x,
204 )
205 }
206}
207
208impl<F: FloatTraitsForComplex> DivAssign for Complex<F> {
209 #[inline]
210 fn div_assign(&mut self, rhs: Self) {
211 *self = self.div(rhs);
212 }
213}
214
215impl<F: FloatTraitsForComplex> Neg for Complex<F> {
216 type Output = Self;
217
218 fn neg(self) -> Self::Output {
219 Self::new(-self.real, -self.imag)
220 }
221}
222
223impl<F: FloatTraitsForComplex> Sum for Complex<F> {
224 fn sum<I: Iterator<Item = Self>>(iter: I) -> Self {
225 iter.fold(Self::new(F::ZERO, F::ZERO), |a, b| a + b)
226 }
227}
228
229impl<F: FloatTraitsForComplex> core::iter::Product for Complex<F> {
230 fn product<I: Iterator<Item = Self>>(iter: I) -> Self {
231 iter.fold(Self::new(F::ONE, F::ZERO), |a, b| a * b)
232 }
233}
234
235impl<F: FloatTraitsForComplex + Debug + PartialOrd> Debug for Complex<F> {
236 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
237 write!(
238 f,
239 "{}({:#?} {} {:#?} I)",
240 stringify!(Complex<F>),
241 self.real,
242 if self.imag > F::ZERO { '+' } else { '-' },
243 if self.imag > F::ZERO {
244 self.imag
245 } else {
246 -self.imag
247 },
248 )
249 }
250}
251
252impl<F: FloatTraitsForComplex + Display + PartialOrd> Display for Complex<F> {
253 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
254 write!(
255 f,
256 "{} {} {} I",
257 self.real,
258 if self.imag > F::ZERO { '+' } else { '-' },
259 if self.imag > F::ZERO {
260 self.imag
261 } else {
262 -self.imag
263 },
264 )
265 }
266}
267
268#[cfg(feature = "approx")]
269impl<F: FloatTraitsForComplex + AbsDiffEq<Epsilon = F> + PartialEq> approx::AbsDiffEq
270 for Complex<F>
271{
272 type Epsilon = F;
273
274 fn default_epsilon() -> Self::Epsilon {
275 F::default_epsilon()
276 }
277
278 fn abs_diff_eq(&self, other: &Self, epsilon: Self::Epsilon) -> bool {
279 self.real.abs_diff_eq(&other.real, epsilon) && self.imag.abs_diff_eq(&other.imag, epsilon)
280 }
281}