jepler/micropython-ulab
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity

factored out fft_fft_ifft_spectrogram

Browse source
This commit is contained in:
Zoltán Vörös 2020-11-27 07:42:20 +01:00
parent 941d36f635
commit aa649b28bb
6 changed files with 104 additions and 100 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

93
code/fft/fft.c
View file

@ -26,87 +26,6 @@
//|
mp_obj_t fft_fft_ifft_spectrum(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(translate("FFT is defined for ndarrays only"));
}
if(n_args == 2) {
if(!MP_OBJ_IS_TYPE(arg_im, &ulab_ndarray_type)) {
mp_raise_NotImplementedError(translate("FFT is defined for ndarrays only"));
}
}
// Check if input is of length of power of 2
ndarray_obj_t *re = MP_OBJ_TO_PTR(arg_re);
#if ULAB_MAX_DIMS > 1
if(re->ndim != 1) {
mp_raise_TypeError(translate("FFT is implemented for linear arrays only"));
}
#endif
size_t len = re->len;
if((len & (len-1)) != 0) {
mp_raise_ValueError(translate("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) {
mp_raise_TypeError(translate("FFT is implemented for linear arrays only"));
}
#endif
if (re->len != im->len) {
mp_raise_ValueError(translate("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_TO_PTR(out_re);
} else {
mp_obj_t tuple[2];
tuple[0] = out_re;
tuple[1] = out_im;
return mp_obj_new_tuple(2, tuple);
}
}
//| def fft(r: ulab.array, c: Optional[ulab.array] = None) -> Tuple[ulab.array, ulab.array]:
//| """
//| :param ulab.array r: A 1-dimension array of values whose size is a power of 2
@ -121,9 +40,9 @@ mp_obj_t fft_fft_ifft_spectrum(size_t n_args, mp_obj_t arg_re, mp_obj_t arg_im,
//|
mp_obj_t fft_fft(size_t n_args, const mp_obj_t *args) {
if(n_args == 2) {
return fft_fft_ifft_spectrum(n_args, args[0], args[1], FFT_FFT);
return fft_fft_ifft_spectrogram(n_args, args[0], args[1], FFT_FFT);
} else {
return fft_fft_ifft_spectrum(n_args, args[0], mp_const_none, FFT_FFT);
return fft_fft_ifft_spectrogram(n_args, args[0], mp_const_none, FFT_FFT);
}
}
@ -141,9 +60,9 @@ MP_DEFINE_CONST_FUN_OBJ_VAR_BETWEEN(fft_fft_obj, 1, 2, fft_fft);
mp_obj_t fft_ifft(size_t n_args, const mp_obj_t *args) {
if(n_args == 2) {
return fft_fft_ifft_spectrum(n_args, args[0], args[1], FFT_IFFT);
return fft_fft_ifft_spectrogram(n_args, args[0], args[1], FFT_IFFT);
} else {
return fft_fft_ifft_spectrum(n_args, args[0], mp_const_none, FFT_IFFT);
return fft_fft_ifft_spectrogram(n_args, args[0], mp_const_none, FFT_IFFT);
}
}
@ -160,9 +79,9 @@ MP_DEFINE_CONST_FUN_OBJ_VAR_BETWEEN(fft_ifft_obj, 1, 2, fft_ifft);
mp_obj_t fft_spectrogram(size_t n_args, const mp_obj_t *args) {
if(n_args == 2) {
return fft_fft_ifft_spectrum(n_args, args[0], args[1], FFT_SPECTROGRAM);
return fft_fft_ifft_spectrogram(n_args, args[0], args[1], FFT_SPECTROGRAM);
} else {
return fft_fft_ifft_spectrum(n_args, args[0], mp_const_none, FFT_SPECTROGRAM);
return fft_fft_ifft_spectrogram(n_args, args[0], mp_const_none, FFT_SPECTROGRAM);
}
}

10
code/fft/fft.h
View file

@ -17,20 +17,12 @@
#include "../ndarray.h"
#include "fft_tools.h"
#define SWAP(t, a, b) { t tmp = a; a = b; b = tmp; }
enum FFT_TYPE {
FFT_FFT,
FFT_IFFT,
FFT_SPECTROGRAM,
};
extern mp_obj_module_t ulab_fft_module;
MP_DECLARE_CONST_FUN_OBJ_VAR_BETWEEN(fft_fft_obj);
MP_DECLARE_CONST_FUN_OBJ_VAR_BETWEEN(fft_ifft_obj);
MP_DECLARE_CONST_FUN_OBJ_VAR_BETWEEN(fft_spectrogram_obj);
mp_obj_t fft_fft_ifft_spectrum(size_t , mp_obj_t , mp_obj_t , uint8_t );
mp_obj_t fft_fft_ifft_spectrogram(size_t , mp_obj_t , mp_obj_t , uint8_t );
#endif

91
code/fft/fft_tools.c
View file

@ -11,6 +11,8 @@
#include <math.h>
#include "py/runtime.h"
#include "../ndarray.h"
#include "../ulab_tools.h"
#include "fft_tools.h"
#ifndef MP_PI
@ -20,8 +22,6 @@
#define MP_E MICROPY_FLOAT_CONST(2.71828182845904523536)
#endif
#define SWAP(t, a, b) { t tmp = a; a = b; b = tmp; }
/*
* The following function takes two arrays, namely, the real and imaginary
* parts of a complex array, and calculates the Fourier transform in place.
@ -76,3 +76,90 @@ void fft_kernel(mp_float_t *real, mp_float_t *imag, size_t n, int isign) {
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(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(translate("FFT is defined for ndarrays only"));
}
if(n_args == 2) {
if(!MP_OBJ_IS_TYPE(arg_im, &ulab_ndarray_type)) {
mp_raise_NotImplementedError(translate("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) {
mp_raise_TypeError(translate("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(translate("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) {
mp_raise_TypeError(translate("FFT is implemented for linear arrays only"));
}
#endif
if (re->len != im->len) {
mp_raise_ValueError(translate("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_TO_PTR(out_re);
} else {
mp_obj_t tuple[2];
tuple[0] = out_re;
tuple[1] = out_im;
return mp_obj_new_tuple(2, tuple);
}
}

6
code/fft/fft_tools.h
View file

@ -11,6 +11,12 @@
#ifndef _FFT_TOOLS_
#define _FFT_TOOLS_
enum FFT_TYPE {
FFT_FFT,
FFT_IFFT,
FFT_SPECTROGRAM,
};
void fft_kernel(mp_float_t *, mp_float_t *, size_t , int );
#endif /* _FFT_TOOLS_ */

2
code/ndarray.h
View file

@ -45,8 +45,6 @@ typedef struct _mp_obj_float_t {
#define translate(x) MP_ERROR_TEXT(x)
#endif
#define SWAP(t, a, b) { t tmp = a; a = b; b = tmp; }
#define NDARRAY_NUMERIC 0
#define NDARRAY_BOOLEAN 1

2
code/ulab_tools.h
View file

@ -11,6 +11,8 @@
#ifndef _TOOLS_
#define _TOOLS_
#define SWAP(t, a, b) { t tmp = a; a = b; b = tmp; }
mp_float_t ndarray_get_float_uint8(void *);
mp_float_t ndarray_get_float_int8(void *);
mp_float_t ndarray_get_float_uint16(void *);