Skip to content

Commit

Permalink
Add Complex::powi and assorted Pow impls
Browse files Browse the repository at this point in the history
  • Loading branch information
cuviper committed Jun 4, 2019
1 parent 813e021 commit 7af251e
Show file tree
Hide file tree
Showing 2 changed files with 140 additions and 2 deletions.
41 changes: 39 additions & 2 deletions src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -38,13 +38,15 @@ use core::str::FromStr;
#[cfg(feature = "std")]
use std::error::Error;

use traits::{Inv, MulAdd, Num, One, Signed, Zero};
use traits::{Inv, MulAdd, Num, One, Pow, Signed, Zero};

#[cfg(feature = "std")]
use traits::float::Float;
use traits::float::FloatCore;

mod cast;
mod pow;

#[cfg(feature = "rand")]
mod crand;
#[cfg(feature = "rand")]
Expand Down Expand Up @@ -121,6 +123,12 @@ impl<T: Clone + Num> Complex<T> {
pub fn unscale(&self, t: T) -> Self {
Self::new(self.re.clone() / t.clone(), self.im.clone() / t)
}

/// Raises `self` to an unsigned integer power.
#[inline]
pub fn powu(&self, exp: u32) -> Self {
Pow::pow(self, exp)
}
}

impl<T: Clone + Num + Neg<Output = T>> Complex<T> {
Expand All @@ -139,6 +147,12 @@ impl<T: Clone + Num + Neg<Output = T>> Complex<T> {
-self.im.clone() / norm_sqr,
)
}

/// Raises `self` to a signed integer power.
#[inline]
pub fn powi(&self, exp: i32) -> Self {
Pow::pow(self, exp)
}
}

impl<T: Clone + Signed> Complex<T> {
Expand Down Expand Up @@ -1487,6 +1501,25 @@ mod test {
assert_eq!(_4_2i.l1_norm(), 6.0);
}

#[test]
fn test_pow() {
for c in all_consts.iter() {
assert_eq!(c.powi(0), _1_0i);
let mut pos = _1_0i;
let mut neg = _1_0i;
for i in 1i32..20 {
pos *= c;
assert_eq!(pos, c.powi(i));
if c.is_zero() {
assert!(c.powi(-i).is_nan());
} else {
neg /= c;
assert_eq!(neg, c.powi(-i));
}
}
}
}

#[cfg(feature = "std")]
mod float {
use super::*;
Expand Down Expand Up @@ -1535,7 +1568,11 @@ mod test {

fn close_to_tol(a: Complex64, b: Complex64, tol: f64) -> bool {
// returns true if a and b are reasonably close
(a == b) || (a - b).norm() < tol
let close = (a == b) || (a - b).norm() < tol;
if !close {
println!("{:?} != {:?}", a, b);
}
close
}

#[test]
Expand Down
101 changes: 101 additions & 0 deletions src/pow.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,101 @@
use super::Complex;

use std::ops::Neg;
use traits::{Float, Num, One, Pow};

macro_rules! pow_impl {
($U:ty, $S:ty) => {
impl<'a, T: Clone + Num> Pow<$U> for &'a Complex<T> {
type Output = Complex<T>;

#[inline]
fn pow(self, mut exp: $U) -> Self::Output {
if exp == 0 {
return Complex::one();
}
let mut base = self.clone();

while exp & 1 == 0 {
base = base.clone() * base;
exp >>= 1;
}

if exp == 1 {
return base;
}

let mut acc = base.clone();
while exp > 1 {
exp >>= 1;
base = base.clone() * base;
if exp & 1 == 1 {
acc = acc * base.clone();
}
}
acc
}
}

impl<'a, 'b, T: Clone + Num> Pow<&'b $U> for &'a Complex<T> {
type Output = Complex<T>;

#[inline]
fn pow(self, exp: &$U) -> Self::Output {
self.pow(*exp)
}
}

impl<'a, T: Clone + Num + Neg<Output = T>> Pow<$S> for &'a Complex<T> {
type Output = Complex<T>;

#[inline]
fn pow(self, exp: $S) -> Self::Output {
if exp < 0 {
Pow::pow(&self.inv(), exp.wrapping_neg() as $U)
} else {
Pow::pow(self, exp as $U)
}
}
}

impl<'a, 'b, T: Clone + Num + Neg<Output = T>> Pow<&'b $S> for &'a Complex<T> {
type Output = Complex<T>;

#[inline]
fn pow(self, exp: &$S) -> Self::Output {
self.pow(*exp)
}
}
};
}

pow_impl!(u8, i8);
pow_impl!(u16, i16);
pow_impl!(u32, i32);
pow_impl!(u64, i64);
pow_impl!(usize, isize);
#[cfg(has_i128)]
pow_impl!(u128, i128);

// Note: the impls above are for `&Complex<T>`, while those below are for `Complex<T>`. This is
// fine since `Float: Copy` anyway, but it's also necessary to avoid conflicting implementations.
// Otherwise rustc would insist that those `Pow<{integer}>` impls could overlap with `Pow<T>` if
// integers ever implement `Float`, though of course we know they won't...

impl<T: Float> Pow<T> for Complex<T> {
type Output = Complex<T>;

#[inline]
fn pow(self, exp: T) -> Self::Output {
self.powf(exp)
}
}

impl<T: Float> Pow<Complex<T>> for Complex<T> {
type Output = Complex<T>;

#[inline]
fn pow(self, exp: Complex<T>) -> Self::Output {
self.powc(exp)
}
}

0 comments on commit 7af251e

Please sign in to comment.