view src/may/transform/fft.yeti @ 449:ea95ee1cb3a6

Add complex-complex forward/inverse

module may.transform.fft;

import edu.emory.mathcs.jtransforms.fft: DoubleFFT_1D;

vec = load may.vector;
complex = load may.complex;

load may.complex.type;

packedToComplex len p is number -> ~double[] -> array<cplx> =
   (n = len / 2;
    array
       (map do i:
            re = if i == n then p[1] else p[i*2] fi;
            im = if i == 0 or i == n then 0 else p[i*2+1] fi;
            complex.complex re im;
        done [0..n]));

complexToPacked arr =
   (n = length arr;
    v = new double[n*2-2];
    for [0..(n-1)*2-1] do i:
        ix = int (i/2);
        v[i] :=
            if i == ix*2 then
                complex.real arr[ix]
            else 
                complex.imaginary arr[ix] 
            fi;
    done;
    v[1] := complex.real arr[n-1];
    v);

unpackedToComplex len unp is number -> ~double[] -> array<cplx> =
    array
       (map2 complex.complex
           (vec.list (vec.slice unp 0 len))
           (vec.list (vec.slice unp len (len*2))));

complexToUnpacked cc is array<cplx> -> ~double[] =
    vec.primitive
       (vec.fromList (map complex.real cc ++ map complex.imaginary cc));

//!!! doc: n separately as below
forward n =
   (d = new DoubleFFT_1D(n);
    do cc:
        arr = complexToUnpacked cc;
        d#complexForward(arr);
        unpackedToComplex n arr;
    done);

//!!! doc: n separately as below
inverse n =
   (d = new DoubleFFT_1D(n);
    do cc:
        arr = complexToUnpacked cc;
        d#complexInverse(arr, true);
        unpackedToComplex n arr;
    done);

//!!! doc: n is supplied separately from the input vector to support partial evaluation
//!!! doc: output has n/2+1 complex values
//!!! doc: powers of two only? check with jtransforms
realForward n = 
   (d = new DoubleFFT_1D(n);
    do bl:
        v = vec.primitive (vec.resizedTo n bl);
        d#realForward(v);
        packedToComplex n v;
    done);

//!!! doc: n separately as above
realForwardMagnitude n =
    complex.magnitudes . (realForward n);

//!!! doc: input requires n/2+1 complex values (or should test and throw?)
//!!! doc: powers of two only? check with jtransforms
realInverse n = 
   (d = new DoubleFFT_1D(n);
    do cplx:
        v = complexToPacked (array cplx);
        d#realInverse(v, true);
        vec.vector v;
    done);

{
forward,
inverse,
realForward,
realForwardMagnitude,
realInverse,
}