added uniform interface to fft, ifft, spectrum, and significantly ext…

…ended manual
v923z · Oct 17, 2019 · 31e40da · 31e40da
1 parent d4acce7
commit 31e40da
Show file tree

Hide file tree

Showing 7 changed files with 775 additions and 195 deletions.
diff --git a/code/fft.c b/code/fft.c
@@ -19,6 +19,13 @@
 #include "ndarray.h"
 #include "fft.h"
 
+
+enum FFT_TYPE {
+    FFT_FFT,
+    FFT_IFFT,
+    FFT_SPECTRUM,
+};
+
 void fft_kernel(float *real, float *imag, int n, int isign) {
     // This is basically a modification of four1 from Numerical Recipes
     // The main difference is that this function takes two arrays, one 
@@ -68,19 +75,17 @@ void fft_kernel(float *real, float *imag, int n, int isign) {
     }
 }
 
-mp_obj_t fft_fft(size_t n_args, const mp_obj_t *args) {
-    // TODO: return the absolute value, if keyword argument is specified
-    // TODO: transform the data in place, if keyword argument is specified
-    if(!MP_OBJ_IS_TYPE(args[0], &ulab_ndarray_type)) {
+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("FFT is defined for ndarrays only");
     } 
     if(n_args == 2) {
-        if(!MP_OBJ_IS_TYPE(args[1], &ulab_ndarray_type)) {
+        if(!MP_OBJ_IS_TYPE(arg_im, &ulab_ndarray_type)) {
             mp_raise_NotImplementedError("FFT is defined for ndarrays only");
         }
     }
     // Check if input is of length of power of 2
-    ndarray_obj_t *re = MP_OBJ_TO_PTR(args[0]);
+    ndarray_obj_t *re = MP_OBJ_TO_PTR(arg_re);
     uint16_t len = re->array->len;
     if((len & (len-1)) != 0) {
         mp_raise_ValueError("input array length must be power of 2");
@@ -89,7 +94,9 @@ mp_obj_t fft_fft(size_t n_args, const mp_obj_t *args) {
     ndarray_obj_t *out_re = create_new_ndarray(1, len, NDARRAY_FLOAT);
     float *data_re = (float *)out_re->array->items;
 
-    if(re->array->typecode == NDARRAY_FLOAT) {
+    if(re->array->typecode == NDARRAY_FLOAT) { 
+        // By treating this case separately, we can save a bit of time.
+        // I don't know if it is worthwhile, though...
         memcpy((float *)out_re->array->items, (float *)re->array->items, re->bytes);
     } else {
         for(size_t i=0; i < len; i++) {
@@ -100,7 +107,7 @@ mp_obj_t fft_fft(size_t n_args, const mp_obj_t *args) {
     float *data_im = (float *)out_im->array->items;
 
     if(n_args == 2) {
-        ndarray_obj_t *im = MP_OBJ_TO_PTR(args[1]);
+        ndarray_obj_t *im = MP_OBJ_TO_PTR(arg_im);
         if (re->array->len != im->array->len) {
             mp_raise_ValueError("real and imaginary parts must be of equal length");
         }
@@ -111,36 +118,51 @@ mp_obj_t fft_fft(size_t n_args, const mp_obj_t *args) {
                 data_im[i] = ndarray_get_float_value(im->array->items, im->array->typecode, i);
             }
         }
-    }    
-    fft_kernel(data_re, data_im, len, 1);
-    mp_obj_t tuple[2];
-    tuple[0] = out_re;
-    tuple[1] = out_im;
-    return mp_obj_new_tuple(2, tuple);
-}
-
-mp_obj_t fft_spectrum(mp_obj_t oin) {
-    // calculates the the spectrum of a single real ndarray in place
-    if(!MP_OBJ_IS_TYPE(oin, &ulab_ndarray_type)) {
-        mp_raise_NotImplementedError("FFT is defined for ndarrays only");
     }
-    ndarray_obj_t *re = MP_OBJ_TO_PTR(oin);
-    uint16_t len = re->array->len;
-    if((re->m > 1) && (re->n > 1)) {
-        mp_raise_ValueError("input data must be an array");
+    if((type == FFT_FFT) || (type == FFT_SPECTRUM)) {
+        fft_kernel(data_re, data_im, len, 1);
+        if(type == FFT_SPECTRUM) {
+            for(size_t i=0; i < len; i++) {
+                data_re[i] = sqrtf(data_re[i]*data_re[i] + data_im[i]*data_im[i]);
+            }
+        }
+    } else { // inverse transform
+        fft_kernel(data_re, data_im, len, -1);
+        for(size_t i=0; i < len; i++) {
+            data_re[i] /= len;
+            data_im[i] /= len;
+        }
     }
-    if((len & (len-1)) != 0) {
-        mp_raise_ValueError("input array length must be power of 2");
+    if(type == FFT_SPECTRUM) {
+        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);
     }
-    if(re->array->typecode != NDARRAY_FLOAT) {
-        mp_raise_TypeError("input array must be of type float");
+}
+
+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);
+    } else {
+        return fft_fft_ifft_spectrum(n_args, args[0], mp_const_none, 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);
+    } else {
+        return fft_fft_ifft_spectrum(n_args, args[0], mp_const_none, FFT_IFFT);
     }
-    float *data_re = (float *)re->array->items;
-    float *data_im = m_new(float, len);
-    fft_kernel(data_re, data_im, len, 1);
-    for(size_t i=0; i < len; i++) {
-        data_re[i] = sqrtf(data_re[i]*data_re[i] + data_im[i]*data_im[i]);
+}
+
+mp_obj_t fft_spectrum(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_SPECTRUM);
+    } else {
+        return fft_fft_ifft_spectrum(n_args, args[0], mp_const_none, FFT_SPECTRUM);
     }
-    m_del(float, data_im, len);
-    return mp_const_none;
 }
diff --git a/code/fft.h b/code/fft.h
@@ -18,5 +18,6 @@
 #define SWAP(t, a, b) { t tmp = a; a = b; b = tmp; }
 
 mp_obj_t fft_fft(size_t , const mp_obj_t *);
-mp_obj_t fft_spectrum(mp_obj_t );
+mp_obj_t fft_ifft(size_t , const mp_obj_t *);
+mp_obj_t fft_spectrum(size_t , const mp_obj_t *);
 #endif
diff --git a/code/ulab.c b/code/ulab.c
@@ -24,7 +24,7 @@
 #include "fft.h"
 #include "numerical.h"
 
-#define ULAB_VERSION 0.201
+#define ULAB_VERSION 0.21
 
 typedef struct _mp_obj_float_t {
     mp_obj_base_t base;
@@ -87,7 +87,10 @@ STATIC MP_DEFINE_CONST_FUN_OBJ_2(poly_polyval_obj, poly_polyval);
 STATIC MP_DEFINE_CONST_FUN_OBJ_VAR_BETWEEN(poly_polyfit_obj, 2, 3, poly_polyfit);
 
 STATIC MP_DEFINE_CONST_FUN_OBJ_VAR_BETWEEN(fft_fft_obj, 1, 2, fft_fft);
-STATIC MP_DEFINE_CONST_FUN_OBJ_1(fft_spectrum_obj, fft_spectrum);
+STATIC MP_DEFINE_CONST_FUN_OBJ_VAR_BETWEEN(fft_ifft_obj, 1, 2, fft_ifft);
+STATIC MP_DEFINE_CONST_FUN_OBJ_VAR_BETWEEN(fft_spectrum_obj, 1, 2, fft_spectrum);
+
+//STATIC MP_DEFINE_CONST_FUN_OBJ_1(fft_spectrum_obj, fft_spectrum);
 
 STATIC const mp_rom_map_elem_t ulab_ndarray_locals_dict_table[] = {
     { MP_ROM_QSTR(MP_QSTR_shape), MP_ROM_PTR(&ndarray_shape_obj) },
@@ -159,6 +162,7 @@ STATIC const mp_map_elem_t ulab_globals_table[] = {
     { MP_OBJ_NEW_QSTR(MP_QSTR_polyval), (mp_obj_t)&poly_polyval_obj },
     { MP_OBJ_NEW_QSTR(MP_QSTR_polyfit), (mp_obj_t)&poly_polyfit_obj },
     { MP_OBJ_NEW_QSTR(MP_QSTR_fft), (mp_obj_t)&fft_fft_obj },
+    { MP_OBJ_NEW_QSTR(MP_QSTR_ifft), (mp_obj_t)&fft_ifft_obj },    
     { MP_OBJ_NEW_QSTR(MP_QSTR_spectrum), (mp_obj_t)&fft_spectrum_obj },
     // class constants
     { MP_ROM_QSTR(MP_QSTR_uint8), MP_ROM_INT(NDARRAY_UINT8) },

diff --git a/code/vectorise.c b/code/vectorise.c
@@ -85,7 +85,7 @@ MATH_FUN_1(acos, acos);
 MATH_FUN_1(acosh, acosh);
 MATH_FUN_1(asin, asin);
 MATH_FUN_1(asinh, asinh);
-MATH_FUN_1(atan, atan);	
+MATH_FUN_1(atan, atan);
 MATH_FUN_1(atanh, atanh);
 MATH_FUN_1(ceil, ceil);
 MATH_FUN_1(cos, cos);

diff --git a/docs/ulab-change-log.md b/docs/ulab-change-log.md
@@ -1,4 +1,10 @@
 
+Thu, 17 Oct 2019
+
+version 0.21
+
+    implemented uniform interface for fft, and spectrum, and added ifft.
+
 Wed, 16 Oct 2019
 
 version 0.20