diff --git a/src/numerical/fft.rs b/src/numerical/fft.rs new file mode 100644 index 00000000..19d08dff --- /dev/null +++ b/src/numerical/fft.rs @@ -0,0 +1,243 @@ +use std::ops::Range; + +use num_complex::Complex64 as C64; + +/// Get the smallest power of two larger than or equal to n, and the log base two +/// of that number. +const fn next_power(mut n: usize) -> (usize, usize) { + let power_of_two = n.is_power_of_two(); + let mut count = 0; + + while n > 0 { + n >>= 1; + count += 1; + } + + if power_of_two { + count -= 1; + } + + (1 << count, count) +} + +/// Flip the bits of a number for the FFT algorithm until the logn'th bit. +const fn reverse_bits_around(num: usize, logn: usize) -> usize { + if logn > 0 { + (num << (usize::BITS as usize - logn)).reverse_bits() + } else { + 0 + } +} + +fn bit_reverse_copy(source: &Vec, target_length: Option) -> (Vec, usize, usize) +where + T: Into + Copy, +{ + let len = usize::max(source.len(), target_length.unwrap_or(0)); + let (power, logn) = next_power(len); + + let mut res = vec![(0.).into(); power]; + + for (i, val) in source.iter().enumerate() { + res[reverse_bits_around(i, logn)] = (*val).into(); + } + + (res, power, logn) +} + +/// Perform the Cooley Tukey variant of the Fast Fourier Transformation in place. +/// If `target_length > values.len()`, make the returned array of at least length +/// `target_length`, +pub fn fft(values: &Vec, target_length: Option) -> Vec { + let (mut res, power, logn) = bit_reverse_copy(values, target_length); + + for s in 1..=logn { + let m = 1 << s; + let m_div_2 = m >> 1; + + let ωm = C64::from_polar(1., -std::f64::consts::TAU / (m as f64)); + + for k in (0..power).step_by(m) { + let mut ω = C64::new(1., 0.); + + for j in 0..m_div_2 { + let t = ω * res[k + j + m_div_2]; + let u = res[k + j]; + + res[k + j] = u + t; + res[k + j + m_div_2] = u - t; + + ω *= ωm; + } + } + } + + res +} + +/// Perform the Cooley Tukey variant of the inverse Fast Fourier Transformation +/// in place. +pub fn inverse_fft(values: &Vec) -> Vec { + let (mut res, power, logn) = bit_reverse_copy(values, None); + + for s in 1..=logn { + let m = 1 << s; + let m_div_2 = m >> 1; + + let ωm = C64::from_polar(1., std::f64::consts::TAU / (m as f64)); + + for k in (0..power).step_by(m) { + let mut ω = C64::new(1., 0.); + + for j in 0..m_div_2 { + let t = ω * res[k + j + m_div_2]; + let u = res[k + j]; + + res[k + j] = u + t; + res[k + j + m_div_2] = u - t; + + ω *= ωm; + } + } + } + + res.into_iter().map(|x| x.norm() / (power as f64)).collect() +} + +#[cfg(test)] +mod fft_tests { + use super::*; + use crate::structure::complex::C64; + use float_cmp::approx_eq; + + #[test] + fn test_next_power() { + let vals = vec![1, 2, 4, 3, 9, 1023, 11, 0]; + let powers = vec![1, 2, 4, 4, 16, 1024, 16, 1]; + + for (val, power) in vals.into_iter().zip(powers) { + let (np, log) = next_power(val); + + assert_eq!(np, power); + assert_eq!(1 << log, np); + } + } + + #[test] + // Test if the values are permuted correctly + fn test_bitwise_copy() { + let vals = vec![1., 2., 3., 4.]; + + // This should be the resulting permutation + let target = vec![ + C64::new(1., 0.), + C64::new(3., 0.), + C64::new(2., 0.), + C64::new(4., 0.), + ]; + + let (swapped, power, logn) = bit_reverse_copy(&vals, None); + + assert_eq!(swapped.len(), vals.len()); + assert_eq!(swapped.len(), power); + assert_eq!(1 << logn, power); + + // Assert the swapped vec and what it is supposed to be is the same are the same + for (swapped, goal) in swapped.into_iter().zip(target.into_iter()) { + assert_eq!(swapped, goal); + } + } + + #[test] + // Test if the bitwise copy extends an array properly before copying the values. + fn test_bitwise_copy_with_extend() { + let vals = vec![1., 2., 3.]; + + let target = vec![ + C64::new(1., 0.), + C64::new(0., 0.), + C64::new(3., 0.), + C64::new(0., 0.), + C64::new(2., 0.), + C64::new(0., 0.), + C64::new(0., 0.), + C64::new(0., 0.), + ]; + + let (swapped, power, logn) = bit_reverse_copy(&vals, Some(8)); + + assert_eq!(swapped.len(), 8); + assert_eq!(swapped.len(), power); + assert_eq!(1 << logn, power); + + for (swapped, goal) in swapped.into_iter().zip(target.into_iter()) { + assert_eq!(swapped, goal); + } + } + + #[test] + fn test_simple_fft() { + // The Discrete Fourier Transform of f(x)=1 is just one. + let coeff1 = vec![1.]; + + let transformed = fft(&coeff1, None); + + let transformed_target = vec![C64::new(1., 0.)]; + + for (val1, val2) in transformed.into_iter().zip(transformed_target) { + assert!(approx_eq!(f64, (val1 - val2).norm(), 0.)); + } + } + + #[test] + fn test_simple_fft_and_inverse() { + let coeffs = vec![0., 1., 0.]; + + let transformed = fft(&coeffs, None); + + // The resulting DFT should have length 4 because it is + // the smallest larger power of two. + let transformed_target = vec![ + C64::new(1., 0.), + C64::new(0., -1.), + C64::new(-1., 0.), + C64::new(0., 1.), + ]; + + for (val1, val2) in transformed.iter().zip(transformed_target) { + assert!(approx_eq!(f64, (val1 - val2).norm(), 0.)); + } + + let coeffs_recovered = inverse_fft(&transformed); + + for (x, y) in coeffs.into_iter().zip(coeffs_recovered.into_iter()) { + assert!(approx_eq!(f64, x, y)) + } + } + + #[test] + fn test_complex_fft() { + // These coefficients represent a polynomial with the fourth root of + // unity (and all its powers except for 1.0) as one of its roots + let coeffs = vec![1., 1., 1., 1.]; + + let transformed = fft(&coeffs, None); + + let transformed_target = vec![ + C64::new(4., 0.), + C64::new(0., 0.), + C64::new(0., 0.), + C64::new(0., 0.), + ]; + + for (val1, val2) in transformed.iter().zip(transformed_target) { + assert!(approx_eq!(f64, (val1 - val2).norm(), 0.)); + } + + let coeffs_recovered = inverse_fft(&transformed); + + for (x, y) in coeffs.into_iter().zip(coeffs_recovered.into_iter()) { + assert!(approx_eq!(f64, x, y)) + } + } +} diff --git a/src/numerical/mod.rs b/src/numerical/mod.rs index 20de60cc..67cde9a5 100644 --- a/src/numerical/mod.rs +++ b/src/numerical/mod.rs @@ -3,6 +3,7 @@ //pub mod bdf; //pub mod gauss_legendre; pub mod eigen; +pub mod fft; pub mod integral; pub mod interp; pub mod newton; diff --git a/src/structure/mod.rs b/src/structure/mod.rs index 666bbeee..4aed6238 100644 --- a/src/structure/mod.rs +++ b/src/structure/mod.rs @@ -7,8 +7,8 @@ //! * DataFrame //! * Multinomial (not yet implemented) -//pub mod complex; pub mod ad; +pub mod complex; pub mod dataframe; pub mod matrix; pub mod multinomial; diff --git a/src/structure/polynomial.rs b/src/structure/polynomial.rs index 9c1f00ed..174ccad6 100644 --- a/src/structure/polynomial.rs +++ b/src/structure/polynomial.rs @@ -251,6 +251,26 @@ impl Polynomial { } } +fn polynomial_fft_mul(poly1: &Polynomial, poly2: &Polynomial) -> Polynomial { + use crate::numerical::fft; + let target_length = Some(2 * usize::max(poly1.coef.len(), poly2.coef.len())); + + let prod_length = poly1.coef.len() + poly2.coef.len() - 1; + + let mut poly1_transform = fft::fft(&poly1.coef, target_length); + let poly2_transform = fft::fft(&poly2.coef, target_length); + + for (val1, val2) in poly1_transform.iter_mut().zip(poly2_transform.iter()) { + *val1 *= val2; + } + + let mut res = fft::inverse_fft(&poly1_transform); + + res.resize(prod_length, 0.); + + poly(res) +} + /// Convenient to declare polynomial pub fn poly(coef: Vec) -> Polynomial { Polynomial::new(coef) @@ -338,30 +358,36 @@ where } } +const SWITCH_T0_FFT_FOR_MUL: usize = 1000; + impl Mul for Polynomial { type Output = Self; fn mul(self, other: Self) -> Self { - let (l1, l2) = (self.coef.len(), other.coef.len()); - let (n1, n2) = (l1 - 1, l2 - 1); - let n = n1 + n2; - let mut result = vec![0f64; n + 1]; + if usize::max(self.coef.len(), other.coef.len()) < SWITCH_T0_FFT_FOR_MUL { + let (l1, l2) = (self.coef.len(), other.coef.len()); + let (n1, n2) = (l1 - 1, l2 - 1); + let n = n1 + n2; + let mut result = vec![0f64; n + 1]; - for i in 0..l1 { - let fixed_val = self.coef[i]; - let fixed_exp = n1 - i; + for i in 0..l1 { + let fixed_val = self.coef[i]; + let fixed_exp = n1 - i; - for j in 0..l2 { - let target_val = other.coef[j]; - let target_exp = n2 - j; + for j in 0..l2 { + let target_val = other.coef[j]; + let target_exp = n2 - j; - let result_val = fixed_val * target_val; - let result_exp = fixed_exp + target_exp; + let result_val = fixed_val * target_val; + let result_exp = fixed_exp + target_exp; - result[n - result_exp] += result_val; + result[n - result_exp] += result_val; + } } - } - Self::new(result) + Self::new(result) + } else { + polynomial_fft_mul(&self, &other) + } } } @@ -585,7 +611,10 @@ pub fn chebyshev_polynomial(n: usize, kind: SpecialKind) -> Polynomial { #[cfg(test)] mod tests { + use std::task::Poll; + use super::*; + use float_cmp::approx_eq; #[test] fn test_honor_division() { @@ -606,4 +635,61 @@ mod tests { assert_eq!(a.eval(i), b.eval(i - 6)); } } + + #[test] + fn test_low_degree_mul() { + let a = Polynomial::new(vec![1., 1.]); + let b = Polynomial::new(vec![-1., 1.]); + + let target = Polynomial::new(vec![-1., 0., 1.]); + let prod = a * b; + + for (prod_coef, target_coef) in target.coef.iter().zip(prod.coef.iter()) { + assert!(approx_eq!(f64, *prod_coef, *target_coef)); + } + } + + #[test] + fn test_constant_poly_mul() { + let a = Polynomial::new(vec![3.]); + let b = Polynomial::new(vec![1., 2., 3., 4., 5.]); + + let prod = a * b.clone(); + + for (prod_coef, orig_coef) in prod.coef.iter().zip(b.coef.iter()) { + assert!(approx_eq!(f64, *prod_coef, 3. * *orig_coef)); + } + } + + #[test] + fn test_large_poly_mul() { + let a = Polynomial::new( + (1..=SWITCH_T0_FFT_FOR_MUL) + .into_iter() + .map(|i| i as f64) + .collect(), + ); + let b = Polynomial::new(vec![1.; SWITCH_T0_FFT_FOR_MUL]); + + let target_coefs = (1..2 * (SWITCH_T0_FFT_FOR_MUL - 1)) + .into_iter() + .map(|i| { + if i <= 1000 { + (i * (i + 1) / 2) as f64 + } else { + (1000 * 1001 / 2 - (i - 1000) * (i - 999) / 2) as f64 + } + }) + .collect::>(); + + for (i, (coef, target_coef)) in (a * b).coef.iter().zip(target_coefs.iter()).enumerate() { + println!("{} {} {}", i, coef, target_coef); + assert!(approx_eq!( + f64, + *coef, + *target_coef, + epsilon = (10.).powi(-8) + )); + } + } }