int
main(int argc, char **argv)
{
mfft *mtmh;
double *sig, *psd, *specgram;
double sigpow;
printf("* Testing TFR library:\n");
printf("* N = %d\n", N);
printf("* shift = %d\n", step);
printf("* MTM NW = %3.2f\n", NW);
printf("* TFR Np = %d\n", Np);
printf("* TFR k = %d\n", k);
printf("* TFR tm = %3.2f\n", tm);
sig = (double*)malloc(17590 * sizeof(double));
fmsin(sig, npoints, 0.15, 0.45, 1024, 256./4, 0.3, -1);
printf("* Input signal to tfr_in.dat\n");
write_file("tfr_in.dat", sig, npoints, 1);
mtmh = mtm_init_dpss(N, NW, (int)(NW*2-1));
printf("* MTM PSD to tfr_out_psd\n");
psd = (double*)malloc(N * sizeof(double));
sigpow = mtfft(mtmh, sig+8300, N);
mtpower(mtmh, psd, sigpow);
write_file("tfr_out_psd.dat", psd, N/2 + 1, 1);
free(psd);
const int l = (npoints - Np + 1) / step;
printf("* MTM spectrogram to tfr_out_mtm\n");
specgram = (double*)calloc(l * (N/2+1), sizeof(double));
mtm_spec(mtmh, specgram, sig, npoints, step, 1);
write_file("tfr_out_mtm.dat", specgram, l, (N/2+1));
free(specgram);
mtm_destroy(mtmh);
printf("* TFR spectrogram to tfr_out_tfr\n");
mtmh = mtm_init_herm(N, Np, k, tm);
specgram = (double*)calloc(l * (N/2+1), sizeof(double));
tfr_spec(mtmh, specgram, sig, npoints, -1, step, 0.01, 5, 0, NULL);
write_file("tfr_out_tfr.dat", specgram, l, (N/2+1));
free(specgram);
mtm_destroy(mtmh);
free(sig);
}
static PyObject*
libtfr_stft(PyObject *self, PyObject *args)
{
/* arguments */
PyObject *o = NULL;
PyArrayObject *signal = NULL;
PyArrayObject *signal_cast = NULL;
PyObject *o2 = NULL;
PyArrayObject *window = NULL;
int step;
int N = 0;
int Npoints, Ntapers;
/* output data */
npy_intp out_shape[3];
PyArrayObject *outdata = NULL;
int do_complex = 0;
double *spec;
complex double *zspec_tmp;
npy_cdouble *zspec;
/* internal stuff */
mfft *mtmh;
double *samples = NULL;
double *windowp;
/* parse arguments */
if (!PyArg_ParseTuple(args, "OOi|ii", &o, &o2, &step, &N, &do_complex))
return NULL;
signal = (PyArrayObject*) PyArray_FromAny(o, NULL, 1, 1, NPY_CONTIGUOUS, NULL);
if (signal==NULL) {
PyErr_SetString(PyExc_TypeError, "Input signal must be an ndarray");
return NULL;
}
window = (PyArrayObject*) PyArray_FromAny(o2, NULL, 1, 2, NPY_CONTIGUOUS, NULL);
if (window==NULL) {
PyErr_SetString(PyExc_TypeError, "Window must be a 1D or 2D ndarray");
goto fail;
}
/* determine dimensions of window */
if (PyArray_NDIM(window)==1) {
Ntapers = 1;
Npoints = PyArray_DIM(window, 0);
}
else {
Ntapers = PyArray_DIM(window, 0);
Npoints = PyArray_DIM(window, 1);
}
/* coerce data to proper type */
samples = coerce_ndarray_double(signal, &signal_cast);
if (samples==NULL) {
PyErr_SetString(PyExc_TypeError, "Unable to cast signal to supported data type");
goto fail;
}
/* need to copy the window function b/c mtm_destroy() will demalloc it */
if (PyArray_TYPE(window)!=NPY_DOUBLE) {
PyErr_SetString(PyExc_TypeError, "Window function must be double precision float");
goto fail;
}
windowp = malloc(Npoints * Ntapers * sizeof(double));
memcpy(windowp, PyArray_DATA(window), Npoints * Ntapers * sizeof(double));
/* allocate outputs and do the transform */
if (N < 1)
N = Npoints;
mtmh = mtm_init(N, Npoints, Ntapers, windowp, NULL);
out_shape[0] = N/2+1;
if (do_complex) {
out_shape[1] = Ntapers;
out_shape[2] = SPEC_NFRAMES(mtmh, PyArray_SIZE(signal), step);
outdata = (PyArrayObject*) PyArray_ZEROS(3,out_shape,NPY_CDOUBLE,1); // fortran-order
zspec = (npy_cdouble*) PyArray_DATA(outdata);
//printf("output dimensions: %d, %d, %d\n", out_shape[0], out_shape[1], out_shape[2]);
zspec_tmp = (complex double*)malloc(PyArray_SIZE(outdata) * sizeof(complex double));
mtm_zspec(mtmh, zspec_tmp, samples, PyArray_SIZE(signal), step);
cmplx_c99tonpy(zspec_tmp, zspec, PyArray_SIZE(outdata));
free(zspec_tmp);
}
else {
out_shape[1] = SPEC_NFRAMES(mtmh, PyArray_SIZE(signal), step);
outdata = (PyArrayObject*) PyArray_ZEROS(2,out_shape,NPY_DOUBLE,1); // fortran-order
spec = (double*) PyArray_DATA(outdata);
mtm_spec(mtmh, spec, samples, PyArray_SIZE(signal), step, 0);
}
mtm_destroy(mtmh);
Py_DECREF(signal);
Py_DECREF(window);
Py_XDECREF(signal_cast);
return PyArray_Return(outdata);
fail:
Py_XDECREF(signal_cast);
Py_XDECREF(signal);
Py_XDECREF(window);
Py_XDECREF(outdata);
return NULL;
}
