jepler/micropython-ulab
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity
micropython-ulab/code/numpy/fft/fft_tools.c

287 lines
9.4 KiB
C
Raw Blame History

/*
* This file is part of the micropython-ulab project,
*
* https://github.com/v923z/micropython-ulab
*
* The MIT License (MIT)
*
* Copyright (c) 2019-2021 Zoltán Vörös
*/
#include <math.h>
#include <string.h>
#include "py/runtime.h"
#include "../../ndarray.h"
#include "../../ulab_tools.h"
#include "../carray/carray_tools.h"
#include "fft_tools.h"
#ifndef MP_PI
#define MP_PI MICROPY_FLOAT_CONST(3.14159265358979323846)
#endif
#ifndef MP_E
#define MP_E MICROPY_FLOAT_CONST(2.71828182845904523536)
#endif
/* Kernel implementation for the case, when ulab has no complex support
* The following function takes two arrays, namely, the real and imaginary
* parts of a complex array, and calculates the Fourier transform in place.
*
* The function is basically a modification of four1 from Numerical Recipes,
* has no dependencies beyond micropython itself (for the definition of mp_float_t),
* and can be used independent of ulab.
*/
#if ULAB_SUPPORTS_COMPLEX & ULAB_FFT_IS_NUMPY_COMPATIBLE
/* Kernel implementation for the complex case. Data are contained in data as
data[0], data[1], data[2], data[3], .... , data[2n - 2], data[2n-1]
real[0], imag[0], real[1], imag[1], .... , real[n-1], imag[n-1]
In general
real[i] = data[2i]
imag[i] = data[2i+1]
*/
void fft_kernel_complex(mp_float_t *data, size_t n, int isign) {
size_t j, m, mmax, istep;
mp_float_t tempr, tempi;
mp_float_t wtemp, wr, wpr, wpi, wi, theta;
j = 0;
for(size_t i = 0; i < n; i++) {
if (j > i) {
SWAP(mp_float_t, data[2*i], data[2*j]);
SWAP(mp_float_t, data[2*i+1], data[2*j+1]);
}
m = n >> 1;
while (j >= m && m > 0) {
j -= m;
m >>= 1;
}
j += m;
}
mmax = 1;
while (n > mmax) {
istep = mmax << 1;
theta = MICROPY_FLOAT_CONST(-2.0)*isign*MP_PI/istep;
wtemp = MICROPY_FLOAT_C_FUN(sin)(MICROPY_FLOAT_CONST(0.5) * theta);
wpr = MICROPY_FLOAT_CONST(-2.0) * wtemp * wtemp;
wpi = MICROPY_FLOAT_C_FUN(sin)(theta);
wr = MICROPY_FLOAT_CONST(1.0);
wi = MICROPY_FLOAT_CONST(0.0);
for(m = 0; m < mmax; m++) {
for(size_t i = m; i < n; i += istep) {
j = i + mmax;
tempr = wr * data[2*j] - wi * data[2*j+1];
tempi = wr * data[2*j+1] + wi * data[2*j];
data[2*j] = data[2*i] - tempr;
data[2*j+1] = data[2*i+1] - tempi;
data[2*i] += tempr;
data[2*i+1] += tempi;
}
wtemp = wr;
wr = wr*wpr - wi*wpi + wr;
wi = wi*wpr + wtemp*wpi + wi;
}
mmax = istep;
}
}
/*
* The following function is a helper interface to the python side.
* It has been factored out from fft.c, so that the same argument parsing
* routine can be called from scipy.signal.spectrogram.
*/
mp_obj_t fft_fft_ifft_spectrogram(mp_obj_t data_in, uint8_t type) {
if(!mp_obj_is_type(data_in, &ulab_ndarray_type)) {
mp_raise_NotImplementedError(MP_ERROR_TEXT("FFT is defined for ndarrays only"));
}
ndarray_obj_t *in = MP_OBJ_TO_PTR(data_in);
#if ULAB_MAX_DIMS > 1
if(in->ndim != 1) {
mp_raise_TypeError(MP_ERROR_TEXT("FFT is implemented for linear arrays only"));
}
#endif
size_t len = in->len;
// Check if input is of length of power of 2
if((len & (len-1)) != 0) {
mp_raise_ValueError(MP_ERROR_TEXT("input array length must be power of 2"));
}
ndarray_obj_t *out = ndarray_new_linear_array(len, NDARRAY_COMPLEX);
mp_float_t *data = (mp_float_t *)out->array;
uint8_t *array = (uint8_t *)in->array;
if(in->dtype == NDARRAY_COMPLEX) {
uint8_t sz = 2 * sizeof(mp_float_t);
for(size_t i = 0; i < len; i++) {
memcpy(data, array, sz);
data += 2;
array += in->strides[ULAB_MAX_DIMS - 1];
}
} else {
mp_float_t (*func)(void *) = ndarray_get_float_function(in->dtype);
for(size_t i = 0; i < len; i++) {
// real part; the imaginary part is 0, no need to assign
*data = func(array);
data += 2;
array += in->strides[ULAB_MAX_DIMS - 1];
}
}
data -= 2 * len;
if((type == FFT_FFT) || (type == FFT_SPECTROGRAM)) {
fft_kernel_complex(data, len, 1);
if(type == FFT_SPECTROGRAM) {
ndarray_obj_t *spectrum = ndarray_new_linear_array(len, NDARRAY_FLOAT);
mp_float_t *sarray = (mp_float_t *)spectrum->array;
for(size_t i = 0; i < len; i++) {
*sarray++ = MICROPY_FLOAT_C_FUN(sqrt)(data[0] * data[0] + data[1] * data[1]);
data += 2;
}
m_del(mp_float_t, data, 2 * len);
return MP_OBJ_FROM_PTR(spectrum);
}
} else { // inverse transform
fft_kernel_complex(data, len, -1);
// TODO: numpy accepts the norm keyword argument
for(size_t i = 0; i < 2 * len; i++) {
*data++ /= len;
}
}
return MP_OBJ_FROM_PTR(out);
}
#else /* ULAB_SUPPORTS_COMPLEX & ULAB_FFT_IS_NUMPY_COMPATIBLE */
void fft_kernel(mp_float_t *real, mp_float_t *imag, size_t n, int isign) {
size_t j, m, mmax, istep;
mp_float_t tempr, tempi;
mp_float_t wtemp, wr, wpr, wpi, wi, theta;
j = 0;
for(size_t i = 0; i < n; i++) {
if (j > i) {
SWAP(mp_float_t, real[i], real[j]);
SWAP(mp_float_t, imag[i], imag[j]);
}
m = n >> 1;
while (j >= m && m > 0) {
j -= m;
m >>= 1;
}
j += m;
}
mmax = 1;
while (n > mmax) {
istep = mmax << 1;
theta = MICROPY_FLOAT_CONST(-2.0)*isign*MP_PI/istep;
wtemp = MICROPY_FLOAT_C_FUN(sin)(MICROPY_FLOAT_CONST(0.5) * theta);
wpr = MICROPY_FLOAT_CONST(-2.0) * wtemp * wtemp;
wpi = MICROPY_FLOAT_C_FUN(sin)(theta);
wr = MICROPY_FLOAT_CONST(1.0);
wi = MICROPY_FLOAT_CONST(0.0);
for(m = 0; m < mmax; m++) {
for(size_t i = m; i < n; i += istep) {
j = i + mmax;
tempr = wr * real[j] - wi * imag[j];
tempi = wr * imag[j] + wi * real[j];
real[j] = real[i] - tempr;
imag[j] = imag[i] - tempi;
real[i] += tempr;
imag[i] += tempi;
}
wtemp = wr;
wr = wr*wpr - wi*wpi + wr;
wi = wi*wpr + wtemp*wpi + wi;
}
mmax = istep;
}
}
mp_obj_t fft_fft_ifft_spectrogram(size_t n_args, mp_obj_t arg_re, mp_obj_t arg_im, uint8_t type) {
if(!mp_obj_is_type(arg_re, &ulab_ndarray_type)) {
mp_raise_NotImplementedError(MP_ERROR_TEXT("FFT is defined for ndarrays only"));
}
if(n_args == 2) {
if(!mp_obj_is_type(arg_im, &ulab_ndarray_type)) {
mp_raise_NotImplementedError(MP_ERROR_TEXT("FFT is defined for ndarrays only"));
}
}
ndarray_obj_t *re = MP_OBJ_TO_PTR(arg_re);
#if ULAB_MAX_DIMS > 1
if(re->ndim != 1) {
COMPLEX_DTYPE_NOT_IMPLEMENTED(re->dtype)
mp_raise_TypeError(MP_ERROR_TEXT("FFT is implemented for linear arrays only"));
}
#endif
size_t len = re->len;
// Check if input is of length of power of 2
if((len & (len-1)) != 0) {
mp_raise_ValueError(MP_ERROR_TEXT("input array length must be power of 2"));
}
ndarray_obj_t *out_re = ndarray_new_linear_array(len, NDARRAY_FLOAT);
mp_float_t *data_re = (mp_float_t *)out_re->array;
uint8_t *array = (uint8_t *)re->array;
mp_float_t (*func)(void *) = ndarray_get_float_function(re->dtype);
for(size_t i=0; i < len; i++) {
*data_re++ = func(array);
array += re->strides[ULAB_MAX_DIMS - 1];
}
data_re -= len;
ndarray_obj_t *out_im = ndarray_new_linear_array(len, NDARRAY_FLOAT);
mp_float_t *data_im = (mp_float_t *)out_im->array;
if(n_args == 2) {
ndarray_obj_t *im = MP_OBJ_TO_PTR(arg_im);
#if ULAB_MAX_DIMS > 1
if(im->ndim != 1) {
COMPLEX_DTYPE_NOT_IMPLEMENTED(im->dtype)
mp_raise_TypeError(MP_ERROR_TEXT("FFT is implemented for linear arrays only"));
}
#endif
if (re->len != im->len) {
mp_raise_ValueError(MP_ERROR_TEXT("real and imaginary parts must be of equal length"));
}
array = (uint8_t *)im->array;
func = ndarray_get_float_function(im->dtype);
for(size_t i=0; i < len; i++) {
*data_im++ = func(array);
array += im->strides[ULAB_MAX_DIMS - 1];
}
data_im -= len;
}
if((type == FFT_FFT) || (type == FFT_SPECTROGRAM)) {
fft_kernel(data_re, data_im, len, 1);
if(type == FFT_SPECTROGRAM) {
for(size_t i=0; i < len; i++) {
*data_re = MICROPY_FLOAT_C_FUN(sqrt)(*data_re * *data_re + *data_im * *data_im);
data_re++;
data_im++;
}
}
} else { // inverse transform
fft_kernel(data_re, data_im, len, -1);
// TODO: numpy accepts the norm keyword argument
for(size_t i=0; i < len; i++) {
*data_re++ /= len;
*data_im++ /= len;
}
}
if(type == FFT_SPECTROGRAM) {
return MP_OBJ_FROM_PTR(out_re);
} else {
mp_obj_t tuple[2];
tuple[0] = MP_OBJ_FROM_PTR(out_re);
tuple[1] = MP_OBJ_FROM_PTR(out_im);
return mp_obj_new_tuple(2, tuple);
}
}
#endif /* ULAB_SUPPORTS_COMPLEX & ULAB_FFT_IS_NUMPY_COMPATIBLE */
Reference in a new issue View git blame Copy permalink