Ticket #307: testFFT.c

#include <stdio.h>
#include <math.h>
#include "pslib.h"

double getTime(void)
/* Gets the current time.  Got this from Nick Kaiser's fetchpix.c */
{
    struct timeval tv;
    gettimeofday(&tv, NULL);
    return(tv.tv_sec + 1.e-6 * tv.tv_usec);
}
 

int main(int argc, char **argv)
{
    if (argc != 2) {
        fprintf(stderr, "Generates an image and FFTs it, outputting timing statistics.\n"
                "Usage: %s SIZE\n"
                "where SIZE is the size (in pixels) of the square image generated.\n"
                "\n", argv[0]);
        exit(EXIT_SUCCESS);
    }
    int size = atoi(argv[1]);           // Size of image
    double startTime = getTime();       // Start time

    psRandom *rng = psRandomAlloc(PS_RANDOM_TAUS, 0); // Seed RNG automagically.
    float yPeriod = psRandomUniform(rng) * size / 2.0;
    float xPeriod = psRandomUniform(rng) * size / 2.0;
    float xPhase = psRandomUniform(rng) * 2.0 * M_PI;
    float yPhase = psRandomUniform(rng) * 2.0 * M_PI;
    psFree(rng);

    // Generate the image
    psImage *image = psImageAlloc(size, size, PS_TYPE_F32);
    for (int j = 0; j < size; j++) {
        float yValue = 2.0 * M_PI / yPeriod * (float)j + yPhase;
        for (int i = 0; i < size; i++) {
            float xValue = 2.0 * M_PI / xPeriod * (float)i + xPhase;
            image->data.F32[j][i] = cos(xValue) * cos(yValue);
        }
    }

    // Write image out
    psFits *imageFile = psFitsAlloc("!image.fits");
    if (!psFitsWriteImage(imageFile, NULL, image, 0, NULL)) {
        psErrorStackPrint(stderr, "Unable to write image: input.fits\n");
    }
    psFree(imageFile);
    printf("%dx%d image (image.fits) constructed after %f sec.\n", size, size, getTime() - startTime);

    // Forward FFT
    psImage *fft = psImageFFT(NULL, image, PS_FFT_FORWARD);
    printf("%dx%d FFT image constructed after %f sec.\n", fft->numCols, fft->numRows, getTime() - startTime);

    // Inverse FFT
    psImage *back = psImageFFT(NULL, fft, PS_FFT_REVERSE | PS_FFT_REAL_RESULT);
    printf("%dx%d inverse FFT image constructed after %f sec.\n", back->numCols, back->numRows,
           getTime() - startTime);
    if (back->numRows != image->numRows || back->numCols != image->numCols) {
        printf("Sizes of original image (%dx%d) and end image (%dx%d) do not match!\n", image->numCols,
               image->numRows, back->numCols, back->numRows);
    }

    // Power spectrum
    psImage *power = psImagePowerSpectrum(NULL, fft);
    printf("%dx%d power spectrum constructed after %f sec.\n", power->numCols, power->numRows, getTime() - startTime);

    // Write out power spectrum
    psFits *powerFile = psFitsAlloc("!power.fits");
    if (!psFitsWriteImage(powerFile, NULL, power, 0, NULL)) {
        psErrorStackPrint(stderr, "Unable to write image: power.fits\n");
    }
    psFree(powerFile);
    printf("Power spectrum image written to power.fits after %f sec.\n", getTime() - startTime);

    // Subtract input and inverse FFT image and compare
    (void)psBinaryOp(back, back, "/", psScalarAlloc(size*size, PS_TYPE_F32));
    psImage *sub = (psImage*)psBinaryOp(NULL, back, "-", image);
    psFits *subFile = psFitsAlloc("!sub.fits");
    if (!psFitsWriteImage(subFile, NULL, sub, 0, NULL)) {
        psErrorStackPrint(stderr, "Unable to write image: sub.fits\n");
    }
    psFree(subFile);

    printf("%dx%d subtracted image constructed and written to sub.fits at %f sec.\n", size, size,
           getTime() - startTime);

    psStats *stats = psStatsAlloc(PS_STAT_SAMPLE_MEAN | PS_STAT_SAMPLE_STDEV);
    psImageStats(stats, sub, NULL, 0);
    printf("Subtracted image: mean = %g, stdev = %g\n", stats->sampleMean, stats->sampleStdev);
    psFree(stats);

    psFree(fft);
    psFree(power);
    psFree(back);
    psFree(sub);
    psFree(image);
}