aboutsummaryrefslogtreecommitdiff
path: root/circuitpython/extmod/ulab/code/numpy/fft/fft_tools.c
diff options
context:
space:
mode:
Diffstat (limited to 'circuitpython/extmod/ulab/code/numpy/fft/fft_tools.c')
-rw-r--r--circuitpython/extmod/ulab/code/numpy/fft/fft_tools.c287
1 files changed, 287 insertions, 0 deletions
diff --git a/circuitpython/extmod/ulab/code/numpy/fft/fft_tools.c b/circuitpython/extmod/ulab/code/numpy/fft/fft_tools.c
new file mode 100644
index 0000000..8a55927
--- /dev/null
+++ b/circuitpython/extmod/ulab/code/numpy/fft/fft_tools.c
@@ -0,0 +1,287 @@
+/*
+ * 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(translate("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(translate("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(translate("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);
+ uint8_t *data_ = (uint8_t *)out->array;
+ for(size_t i = 0; i < len; i++) {
+ memcpy(data_, array, sz);
+ 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 < 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(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) {
+ COMPLEX_DTYPE_NOT_IMPLEMENTED(re->dtype)
+ 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) {
+ COMPLEX_DTYPE_NOT_IMPLEMENTED(im->dtype)
+ 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);
+ }
+}
+#endif /* ULAB_SUPPORTS_COMPLEX & ULAB_FFT_IS_NUMPY_COMPATIBLE */
generated by cgit v1.2.3 (git 2.25.1) at 2025-10-23 13:03:25 +0000