refined_float/
complex.rs

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}