7 changed files with 373 additions and 41 deletions
@ -1,22 +1,103 @@
|
||||
use ndarray::Array; |
||||
use numpy::{IntoPyArray, PyArray1, PyReadonlyArray2}; |
||||
use ndarray::prelude::*; |
||||
|
||||
use ndarray::{Array, Array2, ArrayView2}; |
||||
use ndrustfft::{ndfft_r2c, ndifft_r2c, Complex, R2cFftHandler}; |
||||
use numpy::{IntoPyArray, PyArray2, PyReadonlyArray2}; |
||||
use pyo3::prelude::*; |
||||
|
||||
fn triangular_taper(taper_size: usize, plateau: usize) -> Array<f32, _> { |
||||
assert!(plateau < taper_size, "Plateau cannot be larger than size."); |
||||
assert!(taper_size as f32 % 2.0 != 0., "Size has to be even"); |
||||
use rayon::prelude::*; |
||||
|
||||
use std::cmp; |
||||
|
||||
fn triangular_taper(window_size: usize, plateau: usize) -> Array2<f32> { |
||||
assert!(plateau < window_size, "Plateau cannot be larger than size."); |
||||
assert!(window_size as f32 % 2.0 == 0., "Size has to be even"); |
||||
|
||||
let ramp_size = (window_size - plateau) / 2; |
||||
let ramp = Array::<f32, _>::linspace(0., 1., ramp_size); |
||||
let mut taper = Array::<f32, _>::ones(window_size); |
||||
|
||||
taper.slice_mut(s![..ramp_size]).assign(&ramp); |
||||
taper |
||||
.slice_mut(s![window_size - ramp_size..]) |
||||
.assign(&ramp.slice(s![..;-1])); |
||||
let window = &taper * &taper.slice(s![.., NewAxis]); |
||||
window |
||||
} |
||||
|
||||
fn goldstein_filter( |
||||
data: ArrayView2<f32>, |
||||
window_size: usize, |
||||
overlap: usize, |
||||
exponent: f32, |
||||
) -> Array2<f32> { |
||||
let window_stride = window_size - overlap; |
||||
let window_non_overlap = window_size - 2 * overlap; |
||||
let npx_x = data.nrows(); |
||||
let npx_y = data.ncols(); |
||||
|
||||
let nx = npx_x / window_stride; |
||||
let ny = npx_y / window_stride; |
||||
|
||||
let mut filtered_data = Array2::<f32>::zeros(data.raw_dim()); |
||||
let taper = triangular_taper(window_size, window_non_overlap); |
||||
|
||||
for ix in 0..nx { |
||||
let px_x_beg = cmp::min(ix * window_stride, npx_x); |
||||
let px_x_end = cmp::min(px_x_beg + window_size, npx_x); |
||||
for ny in 0..ny { |
||||
let px_y_beg = cmp::min(ny * window_stride, npx_y); |
||||
let px_y_end = cmp::min(px_y_beg + window_size, npx_y); |
||||
|
||||
let window_data = data.slice(s![px_x_beg..px_x_end, px_y_beg..px_y_end]); |
||||
let window_shape = [window_data.shape()[0], window_data.shape()[1]]; |
||||
let taper_window = taper.slice(s![..window_shape[0], ..window_shape[1]]); |
||||
|
||||
let mut handler = R2cFftHandler::<f32>::new(window_shape[0]); |
||||
let mut window_fft = |
||||
Array2::<Complex<f32>>::default((window_shape[0] / 2 + 1, window_shape[1])); |
||||
let mut window_filtered = Array2::<f32>::default(window_data.raw_dim()); |
||||
|
||||
ndfft_r2c(&window_data, &mut window_fft, &mut handler, 0); |
||||
window_fft.mapv_inplace(|px| px.scale(px.norm().powf(exponent))); |
||||
ndifft_r2c(&window_fft, &mut window_filtered, &mut handler, 0); |
||||
|
||||
window_filtered *= &taper_window; |
||||
|
||||
let mut filtered_slice = |
||||
filtered_data.slice_mut(s![px_x_beg..px_x_end, px_y_beg..px_y_end]); |
||||
filtered_slice += &window_filtered; |
||||
} |
||||
} |
||||
filtered_data |
||||
} |
||||
|
||||
/// A Python module implemented in Rust.
|
||||
#[pymodule] |
||||
fn rust(_py: Python<'_>, m: &PyModule) -> PyResult<()> { |
||||
#[pyfn(m)] |
||||
fn rust_func<'py>( |
||||
#[pyo3(name = "triangular_taper")] |
||||
fn triangular_taper_wrapper<'py>( |
||||
py: Python<'py>, |
||||
window_size: usize, |
||||
plateau: usize, |
||||
) -> PyResult<&'py PyArray2<f32>> { |
||||
let taper = triangular_taper(window_size, plateau).into_pyarray(py); |
||||
Ok(taper) |
||||
} |
||||
|
||||
#[pyfn(m)] |
||||
#[pyo3(name = "goldstein_filter")] |
||||
fn goldstein_filter_wrapper<'py>( |
||||
py: Python<'py>, |
||||
data: PyReadonlyArray2<'py, f64>, |
||||
) -> PyResult<&'py PyArray1<f64>> { |
||||
let x = vec![1., 2., 3.].into_pyarray(py); |
||||
Ok(x) |
||||
data: PyReadonlyArray2<'py, f32>, |
||||
window_size: usize, |
||||
overlap: usize, |
||||
exponent: f32, |
||||
) -> PyResult<&'py PyArray2<f32>> { |
||||
let data_array = data.as_array(); |
||||
let result = goldstein_filter(data_array, window_size, overlap, exponent).into_pyarray(py); |
||||
Ok(result) |
||||
} |
||||
Ok(()) |
||||
} |
||||
|
||||
Loading…
Reference in new issue