Skip to content

Commit

Permalink
re-work spectrogram method, so that RAM can be re-used
Browse files Browse the repository at this point in the history
  • Loading branch information
@v923z
v923z committed Jan 27, 2024
1 parent 9a1d03d commit decc601
Show file tree
Hide file tree
Showing 7 changed files with 249 additions and 75 deletions.
14 changes: 7 additions & 7 deletions code/numpy/fft/fft.c
View file
Original file line number Diff line number Diff line change
Expand Up @@ -5,7 +5,7 @@
*
* The MIT License (MIT)
*
* Copyright (c) 2019-2021 Zoltán Vörös
* Copyright (c) 2019-2024 Zoltán Vörös
* 2020 Scott Shawcroft for Adafruit Industries
* 2020 Taku Fukada
*/
Expand Down Expand Up @@ -43,16 +43,16 @@
//|
#if ULAB_SUPPORTS_COMPLEX & ULAB_FFT_IS_NUMPY_COMPATIBLE
static mp_obj_t fft_fft(mp_obj_t arg) {
return fft_fft_ifft_spectrogram(arg, FFT_FFT);
return fft_fft_ifft(arg, FFT_FFT);
}

MP_DEFINE_CONST_FUN_OBJ_1(fft_fft_obj, fft_fft);
#else
static mp_obj_t fft_fft(size_t n_args, const mp_obj_t *args) {
if(n_args == 2) {
return fft_fft_ifft_spectrogram(n_args, args[0], args[1], FFT_FFT);
return fft_fft_ifft(n_args, args[0], args[1], FFT_FFT);
} else {
return fft_fft_ifft_spectrogram(n_args, args[0], mp_const_none, FFT_FFT);
return fft_fft_ifft(n_args, args[0], mp_const_none, FFT_FFT);
}
}

Expand All @@ -71,17 +71,17 @@ MP_DEFINE_CONST_FUN_OBJ_VAR_BETWEEN(fft_fft_obj, 1, 2, fft_fft);

#if ULAB_SUPPORTS_COMPLEX & ULAB_FFT_IS_NUMPY_COMPATIBLE
static mp_obj_t fft_ifft(mp_obj_t arg) {
return fft_fft_ifft_spectrogram(arg, FFT_IFFT);
return fft_fft_ifft(arg, FFT_IFFT);
}

MP_DEFINE_CONST_FUN_OBJ_1(fft_ifft_obj, fft_ifft);
#else
static mp_obj_t fft_ifft(size_t n_args, const mp_obj_t *args) {
NOT_IMPLEMENTED_FOR_COMPLEX()
if(n_args == 2) {
return fft_fft_ifft_spectrogram(n_args, args[0], args[1], FFT_IFFT);
return fft_fft_ifft(n_args, args[0], args[1], FFT_IFFT);
} else {
return fft_fft_ifft_spectrogram(n_args, args[0], mp_const_none, FFT_IFFT);
return fft_fft_ifft(n_args, args[0], mp_const_none, FFT_IFFT);
}
}

Expand Down
47 changes: 13 additions & 34 deletions code/numpy/fft/fft_tools.c
View file
Original file line number Diff line number Diff line change
Expand Up @@ -5,7 +5,7 @@
*
* The MIT License (MIT)
*
* Copyright (c) 2019-2021 Zoltán Vörös
* Copyright (c) 2019-2024 Zoltán Vörös
*/

#include <math.h>
Expand Down Expand Up @@ -45,7 +45,7 @@
imag[i] = data[2i+1]
*/
void fft_kernel_complex(mp_float_t *data, size_t n, int isign) {
void fft_kernel(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;
Expand Down Expand Up @@ -94,9 +94,9 @@ void fft_kernel_complex(mp_float_t *data, size_t n, int isign) {
/*
* 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.
* routine can be called from utils.spectrogram.
*/
mp_obj_t fft_fft_ifft_spectrogram(mp_obj_t data_in, uint8_t type) {
mp_obj_t fft_fft_ifft(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"));
}
Expand Down Expand Up @@ -134,20 +134,10 @@ mp_obj_t fft_fft_ifft_spectrogram(mp_obj_t data_in, uint8_t type) {
}
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);
}
if(type == FFT_FFT) {
fft_kernel(data, len, 1);
} else { // inverse transform
fft_kernel_complex(data, len, -1);
fft_kernel(data, len, -1);
// TODO: numpy accepts the norm keyword argument
for(size_t i = 0; i < 2 * len; i++) {
*data++ /= len;
Expand Down Expand Up @@ -202,7 +192,7 @@ void fft_kernel(mp_float_t *real, mp_float_t *imag, size_t n, int isign) {
}
}

mp_obj_t fft_fft_ifft_spectrogram(size_t n_args, mp_obj_t arg_re, mp_obj_t arg_im, uint8_t type) {
mp_obj_t fft_fft_ifft(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"));
}
Expand Down Expand Up @@ -258,15 +248,8 @@ mp_obj_t fft_fft_ifft_spectrogram(size_t n_args, mp_obj_t arg_re, mp_obj_t arg_i
data_im -= len;
}

if((type == FFT_FFT) || (type == FFT_SPECTROGRAM)) {
if(type == FFT_FFT) {
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
Expand All @@ -275,13 +258,9 @@ mp_obj_t fft_fft_ifft_spectrogram(size_t n_args, mp_obj_t arg_re, mp_obj_t arg_i
*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);
}
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 */
5 changes: 2 additions & 3 deletions code/numpy/fft/fft_tools.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -14,15 +14,14 @@
enum FFT_TYPE {
FFT_FFT,
FFT_IFFT,
FFT_SPECTROGRAM,
};

#if ULAB_SUPPORTS_COMPLEX & ULAB_FFT_IS_NUMPY_COMPATIBLE
void fft_kernel(mp_float_t *, size_t , int );
mp_obj_t fft_fft_ifft_spectrogram(mp_obj_t , uint8_t );
mp_obj_t fft_fft_ifft(mp_obj_t , uint8_t );
#else
void fft_kernel(mp_float_t *, mp_float_t *, size_t , int );
mp_obj_t fft_fft_ifft_spectrogram(size_t , mp_obj_t , mp_obj_t , uint8_t );
mp_obj_t fft_fft_ifft(size_t , mp_obj_t , mp_obj_t , uint8_t );
#endif /* ULAB_SUPPORTS_COMPLEX & ULAB_FFT_IS_NUMPY_COMPATIBLE */

#endif /* _FFT_TOOLS_ */
2 changes: 1 addition & 1 deletion code/ulab.c
View file
Original file line number Diff line number Diff line change
Expand Up @@ -33,7 +33,7 @@
#include "user/user.h"
#include "utils/utils.h"

#define ULAB_VERSION 6.5.0
#define ULAB_VERSION 6.5.1
#define xstr(s) str(s)
#define str(s) #s

Expand Down
192 changes: 179 additions & 13 deletions code/utils/utils.c
View file
Original file line number Diff line number Diff line change
Expand Up @@ -5,7 +5,7 @@
*
* The MIT License (MIT)
*
* Copyright (c) 2020-2021 Zoltán Vörös
* Copyright (c) 2020-2024 Zoltán Vörös
*/

#include <math.h>
Expand All @@ -16,6 +16,7 @@
#include "py/misc.h"
#include "utils.h"

#include "../ulab_tools.h"
#include "../numpy/fft/fft_tools.h"

#if ULAB_HAS_UTILS_MODULE
Expand Down Expand Up @@ -203,23 +204,188 @@ MP_DEFINE_CONST_FUN_OBJ_KW(utils_from_uint32_buffer_obj, 1, utils_from_uint32_bu
//| ...
//|

mp_obj_t utils_spectrogram(size_t n_args, const mp_obj_t *args) {
#if ULAB_SUPPORTS_COMPLEX & ULAB_FFT_IS_NUMPY_COMPATIBLE
return fft_fft_ifft_spectrogram(args[0], FFT_SPECTROGRAM);
mp_obj_t utils_spectrogram(size_t n_args, const mp_obj_t *pos_args, mp_map_t *kw_args) {
static const mp_arg_t allowed_args[] = {
{ MP_QSTR_, MP_ARG_REQUIRED | MP_ARG_OBJ, { .u_rom_obj = MP_ROM_NONE } } ,
#if !ULAB_FFT_IS_NUMPY_COMPATIBLE
{ MP_QSTR_, MP_ARG_OBJ, { .u_rom_obj = MP_ROM_NONE } },
#else
{ MP_QSTR_scratchpad, MP_ARG_KW_ONLY | MP_ARG_OBJ, { .u_rom_obj = MP_ROM_NONE } },
{ MP_QSTR_out, MP_ARG_KW_ONLY | MP_ARG_OBJ, { .u_rom_obj = MP_ROM_NONE } },
{ MP_QSTR_log, MP_ARG_KW_ONLY | MP_ARG_OBJ, { .u_rom_obj = MP_ROM_FALSE } },
#endif
};

mp_arg_val_t args[MP_ARRAY_SIZE(allowed_args)];
mp_arg_parse_all(n_args, pos_args, kw_args, MP_ARRAY_SIZE(allowed_args), allowed_args, args);

if(!mp_obj_is_type(args[0].u_obj, &ulab_ndarray_type)) {
mp_raise_NotImplementedError(MP_ERROR_TEXT("spectrogram is defined for ndarrays only"));
}
ndarray_obj_t *in = MP_OBJ_TO_PTR(args[0].u_obj);

#if ULAB_MAX_DIMS > 1
if(in->ndim != 1) {
mp_raise_TypeError(MP_ERROR_TEXT("spectrogram is implemented for 1D 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 = NULL;

#if ULAB_FFT_IS_NUMPY_COMPATIBLE
mp_obj_t scratchpad_object = args[1].u_obj;
mp_obj_t out_object = args[2].u_obj;
mp_obj_t log_object = args[3].u_obj;
#else
if(n_args == 2) {
return fft_fft_ifft_spectrogram(n_args, args[0], args[1], FFT_SPECTROGRAM);
mp_obj_t scratchpad_object = args[2].u_obj;
mp_obj_t out_object = args[3].u_obj;
mp_obj_t log_object = args[4].u_obj;
#endif

if(out_object != mp_const_none) {
if(!mp_obj_is_type(out_object, &ulab_ndarray_type)) {
mp_raise_TypeError(MP_ERROR_TEXT("out must be an ndarray"));
}
out = MP_OBJ_TO_PTR(out_object);
if((out->dtype != NDARRAY_FLOAT) || (out->ndim != 1)){
mp_raise_TypeError(MP_ERROR_TEXT("out must be a 1D array of float type"));
}
if(len != out->len) {
mp_raise_ValueError(MP_ERROR_TEXT("input and out arrays are not compatible"));
}
} else {
return fft_fft_ifft_spectrogram(n_args, args[0], mp_const_none, FFT_SPECTROGRAM);
out = ndarray_new_linear_array(len, NDARRAY_FLOAT);
}
#endif

ndarray_obj_t *scratchpad = NULL;
mp_float_t *tmp = NULL;

if(scratchpad_object != mp_const_none) {
if(!mp_obj_is_type(scratchpad_object, &ulab_ndarray_type)) {
mp_raise_TypeError(MP_ERROR_TEXT("scratchpad must be an ndarray"));
}

scratchpad = MP_OBJ_TO_PTR(scratchpad_object);
if(!ndarray_is_dense(scratchpad) || (scratchpad->ndim != 1)) {
mp_raise_ValueError(MP_ERROR_TEXT("scratchpad must be a 1D dense array"));
}

#if ULAB_SUPPORTS_COMPLEX
if((scratchpad->dtype != NDARRAY_COMPLEX) && (scratchpad->dtype != NDARRAY_FLOAT)) {
mp_raise_ValueError(MP_ERROR_TEXT("scratchpad must be of complex or float type"));
}
if(((scratchpad->dtype == NDARRAY_COMPLEX) && (len != scratchpad->len)) ||
((scratchpad->dtype == NDARRAY_FLOAT) && (2 * len != scratchpad->len))) {
mp_raise_ValueError(MP_ERROR_TEXT("in and scratchpad arrays are not compatible"));
}
#else
if((scratchpad->dtype != NDARRAY_FLOAT)){
mp_raise_TypeError(MP_ERROR_TEXT("scratchpad must be of float type"));
}
if(2 * len != scratchpad->len) {
mp_raise_ValueError(MP_ERROR_TEXT("input and scratchpad arrays are not compatible"));
}
#endif

tmp = (mp_float_t *)scratchpad->array;
} else {
tmp = m_new0(mp_float_t, 2 * len);
}

uint8_t *array = (uint8_t *)in->array;

#if ULAB_FFT_IS_NUMPY_COMPATIBLE // ULAB_SUPPORTS_COMPLEX is automatically true
if(in->dtype == NDARRAY_COMPLEX) {
uint8_t sz = 2 * sizeof(mp_float_t);
for(size_t i = 0; i < len; i++) {
memcpy(tmp, array, sz);
tmp += 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++) {
*tmp++ = func(array); // real part
*tmp++ = 0; // imaginary part, clear
array += in->strides[ULAB_MAX_DIMS - 1];
}
}

tmp -= 2 * len;
fft_kernel(tmp, len, 1);
#else // we might have two real input vectors

ndarray_obj_t *in2 = NULL;

if(n_args == 2) {
if(!mp_obj_is_type(args[1].u_obj, &ulab_ndarray_type)) {
mp_raise_NotImplementedError(MP_ERROR_TEXT("spectrogram is defined for ndarrays only"));
}
in2 = MP_OBJ_TO_PTR(args[1].u_obj);

#if ULAB_MAX_DIMS > 1
if(in2->ndim != 1) {
mp_raise_TypeError(MP_ERROR_TEXT("spectrogram is implemented for 1D arrays only"));
}
#endif
if(len != in2->len) {
mp_raise_TypeError(MP_ERROR_TEXT("input arrays are not compatible"));
}
}

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
*tmp++ = func(array);
array += in->strides[ULAB_MAX_DIMS - 1];
}

if(n_args == 2) {
mp_float_t (*func2)(void *) = ndarray_get_float_function(in2->dtype);
array = (uint8_t *)in2->array;
for(size_t i = 0; i < len; i++) {
*tmp++ = func2(array);
array += in2->strides[ULAB_MAX_DIMS - 1];
}
tmp -= len;
}
tmp -= len;

fft_kernel(tmp, tmp + len, len, 1);
#endif /* ULAB_FFT_IS_NUMPY_COMPATIBLE */

mp_float_t *spectrum = (mp_float_t *)out->array;
uint8_t spectrum_sz = out->strides[ULAB_MAX_DIMS - 1] / sizeof(mp_float_t);

for(size_t i = 0; i < len; i++) {
#if ULAB_FFT_IS_NUMPY_COMPATIBLE
*spectrum = MICROPY_FLOAT_C_FUN(sqrt)(*tmp * *tmp + *(tmp + 1) * *(tmp + 1));
tmp += 2;
#else
*spectrum = MICROPY_FLOAT_C_FUN(sqrt)(*tmp * *tmp + *(tmp + len) * *(tmp + len));
tmp++;
#endif
if(log_object == mp_const_true) {
*spectrum = MICROPY_FLOAT_C_FUN(log)(*spectrum);
}
spectrum += spectrum_sz;
}

// FIXME:
// if(scratchpad_object == mp_const_none) {
// m_del(mp_float_t, tmp, 2 * len);
// }
return MP_OBJ_FROM_PTR(out);
}

#if ULAB_SUPPORTS_COMPLEX & ULAB_FFT_IS_NUMPY_COMPATIBLE
MP_DEFINE_CONST_FUN_OBJ_VAR_BETWEEN(utils_spectrogram_obj, 1, 1, utils_spectrogram);
#else
MP_DEFINE_CONST_FUN_OBJ_VAR_BETWEEN(utils_spectrogram_obj, 1, 2, utils_spectrogram);
#endif
MP_DEFINE_CONST_FUN_OBJ_KW(utils_spectrogram_obj, 1, utils_spectrogram);

#endif /* ULAB_UTILS_HAS_SPECTROGRAM */

Expand Down
Loading

0 comments on commit decc601

Please sign in to comment.