Mandelbrot set explorer

#include <iostream>
#include "timer.h"
#include <stdio.h>
#include <opencv2/core/core.hpp>
#include <opencv2/highgui/highgui.hpp>
#include <opencv2/opencv.hpp>
#include <cuda.h>
#include <cuda_runtime.h>
#include <string>

using namespace std;

/*
* Global variables
*/
int numCols = 800, numRows = 600; // Image size
double cx = 0, cy = 0, zoom = 100.0f; // Center of the image and zoom
uchar *h_image; // The image
int iterationMax = 100; // Maximum number of iterations in the Mandelbrot loop
int framesPerZoom = 10; // Number of frames in the zoom animation



/*
* Operations on complex numbers
*/
__device__ void cMult(double *ar, double *ai, double *br, double *bi, double  *cr, double  *ci)
{
    *cr = *ar*(*br) - *ai*(*bi);
    *ci = *ar*(*bi) + *ai*(*br);
}
__device__ void cAdd(double *ar, double *ai, double *br, double *bi, double  *cr, double  *ci)
{
    *cr = *ar + *br;
    *ci = *ai + *bi;
}
__device__ void cMin(double *ar, double *ai, double *br, double *bi, double  *cr, double  *ci)
{
    *cr = *ar - *br;
    *ci = *ai - *bi;
}

/*
* Fractal functions
*/
// Z^2 + C
__device__ void fractal1(double *x, double *y, double a, double  b)
{
    double auxX, auxY;
    cMult(x, y, x, y, &auxX, &auxY);
    *x = auxX + a;
    *y = auxY + b;
}
// Z^4 + C
__device__ void fractal2(double *x, double *y, double  a, double  b)
{
    double auxX, auxY;
    cMult(x, y, x, y, &auxX, &auxY);
    *x = auxX;
    *y = auxY;
    cMult(x, y, x, y, &auxX, &auxY);
    *x = auxX + a;
    *y = auxY + b;
}


/*
* Kernel function
*/
__global__ void mandelbrot(uchar* const d_image, int numRows, int numCols, double cx, double cy, double zoom, int iterationMax)
{
    // Init
    const int2 thread_2D_pos = make_int2( blockIdx.x * blockDim.x + threadIdx.x,
        blockIdx.y * blockDim.y + threadIdx.y);

    const int thread_1D_pos = (thread_2D_pos.y * numCols + thread_2D_pos.x)*3;

    // Exit if we are out of the image
    if(thread_2D_pos.x >= numCols || thread_2D_pos.y >= numRows) return;

    // Mandelbrot loop
    /*
        U0 = 0
        Un+1 = Un*Un + c
        where c = a + b*i
        We stop when |Un|>2 or after iterationMax iterations

        You can find many similar formulas here : http://www.lifesmith.com/formulas.html
    */
    int i = 0;
    double a = cx + (thread_2D_pos.x-numCols/2)/zoom, b = cy + (thread_2D_pos.y- numRows/2)/zoom, x = 0, y = 0;
    for(i = 0 ; i < iterationMax ; i++)
    {
        fractal1(&x, &y, a, b);

        if(x*x + y*y > 4) break;
    }

    // Color conversion
    i = (1.*i)/iterationMax*255.0;
    if(i/64 == 0)
    {
        d_image[thread_1D_pos] = 0;
        d_image[thread_1D_pos+1] = i*4;
        d_image[thread_1D_pos+2] = 255;
    }
    else if(i/64 == 1)
    {
        d_image[thread_1D_pos] = 0;
        d_image[thread_1D_pos+1] = 255;
        d_image[thread_1D_pos+2] = 255-(i%64)*4;
    }
    else if(i/64 == 2)
    {
        d_image[thread_1D_pos] = (i%64)*4;
        d_image[thread_1D_pos+1] = 255;
        d_image[thread_1D_pos+2] = 0;
    } else {
        d_image[thread_1D_pos] = 255;
        d_image[thread_1D_pos+1] = 255-(i%64)*4;
        d_image[thread_1D_pos+2] = 0;
    }
}


/*
* Initialization of cuda
*/
void operations(uchar* h_image, const int numRows, const int numCols, double cx, double cy, double zoom)
{
    int BLOCK_SIZE_X = 20;
    int BLOCK_SIZE_Y = 20;
    const dim3 blockSize(BLOCK_SIZE_X, BLOCK_SIZE_Y, 1);
    const dim3 gridSize(numCols/BLOCK_SIZE_X+1, numRows/BLOCK_SIZE_Y+1, 1);

    uchar* d_image;
    cudaMalloc(&d_image, sizeof(uchar) * 3 * numRows * numCols);

    // Menderlbrot
    mandelbrot<<<gridSize, blockSize>>>(d_image, numRows, numCols, cx, cy, zoom, iterationMax);
    cudaDeviceSynchronize();

    cudaMemcpy(h_image, d_image, sizeof(uchar) * 3 * numRows * numCols, cudaMemcpyDeviceToHost);
    cudaFree(d_image);
}







/*
* Mouse callback
*/
void mouseEvent(int evt, int x, int y, int flags, void* param){
    if(evt==CV_EVENT_LBUTTONDOWN || evt == CV_EVENT_RBUTTONDOWN){

        for(int i = 0 ; i < framesPerZoom ; i++) {
            // New center
            cx += (x-numCols/2)/zoom/framesPerZoom;
            cy += (y-numRows/2)/zoom/framesPerZoom;
            //printf("x=%d y=%d\ncx=%f, cy=%f zoom=%f\n", x, y, cx, cy, zoom);


            // New zoom
            zoom *= (evt==CV_EVENT_LBUTTONDOWN) ? std::pow(1.5, 1./framesPerZoom) : (evt==CV_EVENT_RBUTTONDOWN) ? std::pow(0.66, 1./framesPerZoom) : 1.0;

            // New image
            //GpuTimer timer;
            //timer.Start();
            operations(h_image, numRows, numCols, cx, cy, zoom);
            //timer.Stop();
            //printf("Total time : %f ms\n", timer.Elapsed());

            cv::imshow("Mandelbrot", cv::Mat(numRows, numCols, CV_8UC3, h_image));
            cv::waitKey(10);
        }
    }
}

/*
* Main
*/
int main(int argc, char **argv) {


    h_image = (uchar*) malloc(sizeof(uchar) * 3 * numRows * numCols);

    GpuTimer timer;
    timer.Start();
    operations(h_image, numRows, numCols, cx, cy, zoom);
    timer.Stop();
    printf("Total time : %f ms\n", timer.Elapsed());

    cv::imshow("Mandelbrot", cv::Mat(numRows, numCols, CV_8UC3, h_image));
    cvSetMouseCallback("Mandelbrot", mouseEvent, 0);

    cv::waitKey();
    return 0;
}