src/libnum/complex.rs

   1 // Copyright 2013 The Rust Project Developers. See the COPYRIGHT
   2 // file at the top-level directory of this distribution and at
   3 // http://rust-lang.org/COPYRIGHT.
   4 //
   5 // Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
   6 // http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
   7 // <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
   8 // option. This file may not be copied, modified, or distributed
   9 // except according to those terms.
  10
  11
  12 //! Complex numbers.
  13
  14 use std::fmt;
  15 use std::num::{Zero,One,ToStrRadix};
  16
  17 // FIXME #1284: handle complex NaN & infinity etc. This
  18 // probably doesn't map to C's _Complex correctly.
  19
  20 /// A complex number in Cartesian form.
  21 #[deriving(PartialEq,Clone)]
  22 pub struct Complex<T> {
  23     /// Real portion of the complex number
  24     pub re: T,
  25     /// Imaginary portion of the complex number
  26     pub im: T
  27 }
  28
  29 pub type Complex32 = Complex<f32>;
  30 pub type Complex64 = Complex<f64>;
  31
  32 impl<T: Clone + Num> Complex<T> {
  33     /// Create a new Complex
  34     #[inline]
  35     pub fn new(re: T, im: T) -> Complex<T> {
  36         Complex { re: re, im: im }
  37     }
  38
  39     /**
  40     Returns the square of the norm (since `T` doesn't necessarily
  41     have a sqrt function), i.e. `re^2 + im^2`.
  42     */
  43     #[inline]
  44     pub fn norm_sqr(&self) -> T {
  45         self.re * self.re + self.im * self.im
  46     }
  47
  48
  49     /// Returns the complex conjugate. i.e. `re - i im`
  50     #[inline]
  51     pub fn conj(&self) -> Complex<T> {
  52         Complex::new(self.re.clone(), -self.im)
  53     }
  54
  55
  56     /// Multiplies `self` by the scalar `t`.
  57     #[inline]
  58     pub fn scale(&self, t: T) -> Complex<T> {
  59         Complex::new(self.re * t, self.im * t)
  60     }
  61
  62     /// Divides `self` by the scalar `t`.
  63     #[inline]
  64     pub fn unscale(&self, t: T) -> Complex<T> {
  65         Complex::new(self.re / t, self.im / t)
  66     }
  67
  68     /// Returns `1/self`
  69     #[inline]
  70     pub fn inv(&self) -> Complex<T> {
  71         let norm_sqr = self.norm_sqr();
  72         Complex::new(self.re / norm_sqr,
  73                     -self.im / norm_sqr)
  74     }
  75 }
  76
  77 impl<T: Clone + FloatMath> Complex<T> {
  78     /// Calculate |self|
  79     #[inline]
  80     pub fn norm(&self) -> T {
  81         self.re.hypot(self.im)
  82     }
  83 }
  84
  85 impl<T: Clone + FloatMath> Complex<T> {
  86     /// Calculate the principal Arg of self.
  87     #[inline]
  88     pub fn arg(&self) -> T {
  89         self.im.atan2(self.re)
  90     }
  91     /// Convert to polar form (r, theta), such that `self = r * exp(i
  92     /// * theta)`
  93     #[inline]
  94     pub fn to_polar(&self) -> (T, T) {
  95         (self.norm(), self.arg())
  96     }
  97     /// Convert a polar representation into a complex number.
  98     #[inline]
  99     pub fn from_polar(r: &T, theta: &T) -> Complex<T> {
 100         Complex::new(*r * theta.cos(), *r * theta.sin())
 101     }
 102 }
 103
 104 /* arithmetic */
 105 // (a + i b) + (c + i d) == (a + c) + i (b + d)
 106 impl<T: Clone + Num> Add<Complex<T>, Complex<T>> for Complex<T> {
 107     #[inline]
 108     fn add(&self, other: &Complex<T>) -> Complex<T> {
 109         Complex::new(self.re + other.re, self.im + other.im)
 110     }
 111 }
 112 // (a + i b) - (c + i d) == (a - c) + i (b - d)
 113 impl<T: Clone + Num> Sub<Complex<T>, Complex<T>> for Complex<T> {
 114     #[inline]
 115     fn sub(&self, other: &Complex<T>) -> Complex<T> {
 116         Complex::new(self.re - other.re, self.im - other.im)
 117     }
 118 }
 119 // (a + i b) * (c + i d) == (a*c - b*d) + i (a*d + b*c)
 120 impl<T: Clone + Num> Mul<Complex<T>, Complex<T>> for Complex<T> {
 121     #[inline]
 122     fn mul(&self, other: &Complex<T>) -> Complex<T> {
 123         Complex::new(self.re*other.re - self.im*other.im,
 124                    self.re*other.im + self.im*other.re)
 125     }
 126 }
 127
 128 // (a + i b) / (c + i d) == [(a + i b) * (c - i d)] / (c*c + d*d)
 129 //   == [(a*c + b*d) / (c*c + d*d)] + i [(b*c - a*d) / (c*c + d*d)]
 130 impl<T: Clone + Num> Div<Complex<T>, Complex<T>> for Complex<T> {
 131     #[inline]
 132     fn div(&self, other: &Complex<T>) -> Complex<T> {
 133         let norm_sqr = other.norm_sqr();
 134         Complex::new((self.re*other.re + self.im*other.im) / norm_sqr,
 135                    (self.im*other.re - self.re*other.im) / norm_sqr)
 136     }
 137 }
 138
 139 impl<T: Clone + Num> Neg<Complex<T>> for Complex<T> {
 140     #[inline]
 141     fn neg(&self) -> Complex<T> {
 142         Complex::new(-self.re, -self.im)
 143     }
 144 }
 145
 146 /* constants */
 147 impl<T: Clone + Num> Zero for Complex<T> {
 148     #[inline]
 149     fn zero() -> Complex<T> {
 150         Complex::new(Zero::zero(), Zero::zero())
 151     }
 152
 153     #[inline]
 154     fn is_zero(&self) -> bool {
 155         self.re.is_zero() && self.im.is_zero()
 156     }
 157 }
 158
 159 impl<T: Clone + Num> One for Complex<T> {
 160     #[inline]
 161     fn one() -> Complex<T> {
 162         Complex::new(One::one(), Zero::zero())
 163     }
 164 }
 165
 166 /* string conversions */
 167 impl<T: fmt::Show + Num + PartialOrd> fmt::Show for Complex<T> {
 168     fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
 169         if self.im < Zero::zero() {
 170             write!(f, "{}-{}i", self.re, -self.im)
 171         } else {
 172             write!(f, "{}+{}i", self.re, self.im)
 173         }
 174     }
 175 }
 176
 177 impl<T: ToStrRadix + Num + PartialOrd> ToStrRadix for Complex<T> {
 178     fn to_str_radix(&self, radix: uint) -> String {
 179         if self.im < Zero::zero() {
 180             format!("{}-{}i",
 181                     self.re.to_str_radix(radix),
 182                     (-self.im).to_str_radix(radix))
 183         } else {
 184             format!("{}+{}i",
 185                     self.re.to_str_radix(radix),
 186                     self.im.to_str_radix(radix))
 187         }
 188     }
 189 }
 190
 191 #[cfg(test)]
 192 mod test {
 193     #![allow(non_uppercase_statics)]
 194
 195     use super::{Complex64, Complex};
 196     use std::num::{Zero,One,Float};
 197
 198     pub static _0_0i : Complex64 = Complex { re: 0.0, im: 0.0 };
 199     pub static _1_0i : Complex64 = Complex { re: 1.0, im: 0.0 };
 200     pub static _1_1i : Complex64 = Complex { re: 1.0, im: 1.0 };
 201     pub static _0_1i : Complex64 = Complex { re: 0.0, im: 1.0 };
 202     pub static _neg1_1i : Complex64 = Complex { re: -1.0, im: 1.0 };
 203     pub static _05_05i : Complex64 = Complex { re: 0.5, im: 0.5 };
 204     pub static all_consts : [Complex64, .. 5] = [_0_0i, _1_0i, _1_1i, _neg1_1i, _05_05i];
 205
 206     #[test]
 207     fn test_consts() {
 208         // check our constants are what Complex::new creates
 209         fn test(c : Complex64, r : f64, i: f64) {
 210             assert_eq!(c, Complex::new(r,i));
 211         }
 212         test(_0_0i, 0.0, 0.0);
 213         test(_1_0i, 1.0, 0.0);
 214         test(_1_1i, 1.0, 1.0);
 215         test(_neg1_1i, -1.0, 1.0);
 216         test(_05_05i, 0.5, 0.5);
 217
 218         assert_eq!(_0_0i, Zero::zero());
 219         assert_eq!(_1_0i, One::one());
 220     }
 221
 222     #[test]
 223     #[ignore(cfg(target_arch = "x86"))]
 224     // FIXME #7158: (maybe?) currently failing on x86.
 225     fn test_norm() {
 226         fn test(c: Complex64, ns: f64) {
 227             assert_eq!(c.norm_sqr(), ns);
 228             assert_eq!(c.norm(), ns.sqrt())
 229         }
 230         test(_0_0i, 0.0);
 231         test(_1_0i, 1.0);
 232         test(_1_1i, 2.0);
 233         test(_neg1_1i, 2.0);
 234         test(_05_05i, 0.5);
 235     }
 236
 237     #[test]
 238     fn test_scale_unscale() {
 239         assert_eq!(_05_05i.scale(2.0), _1_1i);
 240         assert_eq!(_1_1i.unscale(2.0), _05_05i);
 241         for &c in all_consts.iter() {
 242             assert_eq!(c.scale(2.0).unscale(2.0), c);
 243         }
 244     }
 245
 246     #[test]
 247     fn test_conj() {
 248         for &c in all_consts.iter() {
 249             assert_eq!(c.conj(), Complex::new(c.re, -c.im));
 250             assert_eq!(c.conj().conj(), c);
 251         }
 252     }
 253
 254     #[test]
 255     fn test_inv() {
 256         assert_eq!(_1_1i.inv(), _05_05i.conj());
 257         assert_eq!(_1_0i.inv(), _1_0i.inv());
 258     }
 259
 260     #[test]
 261     #[should_fail]
 262     #[ignore]
 263     fn test_inv_zero() {
 264         // FIXME #5736: should this really fail, or just NaN?
 265         _0_0i.inv();
 266     }
 267
 268     #[test]
 269     fn test_arg() {
 270         fn test(c: Complex64, arg: f64) {
 271             assert!((c.arg() - arg).abs() < 1.0e-6)
 272         }
 273         test(_1_0i, 0.0);
 274         test(_1_1i, 0.25 * Float::pi());
 275         test(_neg1_1i, 0.75 * Float::pi());
 276         test(_05_05i, 0.25 * Float::pi());
 277     }
 278
 279     #[test]
 280     fn test_polar_conv() {
 281         fn test(c: Complex64) {
 282             let (r, theta) = c.to_polar();
 283             assert!((c - Complex::from_polar(&r, &theta)).norm() < 1e-6);
 284         }
 285         for &c in all_consts.iter() { test(c); }
 286     }
 287
 288     mod arith {
 289         use super::{_0_0i, _1_0i, _1_1i, _0_1i, _neg1_1i, _05_05i, all_consts};
 290         use std::num::Zero;
 291
 292         #[test]
 293         fn test_add() {
 294             assert_eq!(_05_05i + _05_05i, _1_1i);
 295             assert_eq!(_0_1i + _1_0i, _1_1i);
 296             assert_eq!(_1_0i + _neg1_1i, _0_1i);
 297
 298             for &c in all_consts.iter() {
 299                 assert_eq!(_0_0i + c, c);
 300                 assert_eq!(c + _0_0i, c);
 301             }
 302         }
 303
 304         #[test]
 305         fn test_sub() {
 306             assert_eq!(_05_05i - _05_05i, _0_0i);
 307             assert_eq!(_0_1i - _1_0i, _neg1_1i);
 308             assert_eq!(_0_1i - _neg1_1i, _1_0i);
 309
 310             for &c in all_consts.iter() {
 311                 assert_eq!(c - _0_0i, c);
 312                 assert_eq!(c - c, _0_0i);
 313             }
 314         }
 315
 316         #[test]
 317         fn test_mul() {
 318             assert_eq!(_05_05i * _05_05i, _0_1i.unscale(2.0));
 319             assert_eq!(_1_1i * _0_1i, _neg1_1i);
 320
 321             // i^2 & i^4
 322             assert_eq!(_0_1i * _0_1i, -_1_0i);
 323             assert_eq!(_0_1i * _0_1i * _0_1i * _0_1i, _1_0i);
 324
 325             for &c in all_consts.iter() {
 326                 assert_eq!(c * _1_0i, c);
 327                 assert_eq!(_1_0i * c, c);
 328             }
 329         }
 330         #[test]
 331         fn test_div() {
 332             assert_eq!(_neg1_1i / _0_1i, _1_1i);
 333             for &c in all_consts.iter() {
 334                 if c != Zero::zero() {
 335                     assert_eq!(c / c, _1_0i);
 336                 }
 337             }
 338         }
 339         #[test]
 340         fn test_neg() {
 341             assert_eq!(-_1_0i + _0_1i, _neg1_1i);
 342             assert_eq!((-_0_1i) * _0_1i, _1_0i);
 343             for &c in all_consts.iter() {
 344                 assert_eq!(-(-c), c);
 345             }
 346         }
 347     }
 348
 349     #[test]
 350     fn test_to_str() {
 351         fn test(c : Complex64, s: String) {
 352             assert_eq!(c.to_str(), s);
 353         }
 354         test(_0_0i, "0+0i".to_string());
 355         test(_1_0i, "1+0i".to_string());
 356         test(_0_1i, "0+1i".to_string());
 357         test(_1_1i, "1+1i".to_string());
 358         test(_neg1_1i, "-1+1i".to_string());
 359         test(-_neg1_1i, "1-1i".to_string());
 360         test(_05_05i, "0.5+0.5i".to_string());
 361     }
 362 }