examples/pde/swe.cpp

#include <stdio.h>
#include <string.h>

#include <arrayfire.h>
using namespace af;

static void swe_demo(bool console)
{
    double time_total = 20; // run for N seconds

    // Grid length, number and spacing
    const unsigned Lx = 512, nx = Lx + 1;
    const unsigned Ly = 512, ny = Ly + 1;

    const float dx = Lx / (nx - 1);
    const float dy = Ly / (ny - 1);

    array ZERO = constant(0,nx, ny);
    array um = ZERO, vm = ZERO;

    unsigned io = (unsigned)floor(Lx / 5.0f),
             jo = (unsigned)floor(Ly / 5.0f),
             k = 20;

    array x = tile(array(seq(nx),nx,1), 1,ny);
    array y = tile(array(seq(ny),1,ny), nx,1);
    array etam = 0.01f * exp((-((x - io) * (x - io) + (y - jo) * (y - jo))) / (k * k));
    array eta = etam;

    float dt = 0.5;

    // conv kernels
    unsigned diff_x_dims[] = {3, 1}, diff_y_dims[] = {1, 3}, lap_dims[] = {3, 3};
    float h_diff_kernel[] = {9.81f * (dt / dx), 0, -9.81f * (dt / dx)};
    float h_lap_kernel[] = {0, 1, 0, 1, -4, 1, 0, 1, 0};

    timer time_start, time_last;
    time_start = time_last = timer::start();
    int iter = 0, iter_last = 0;
    double max_rate = 0;

    while (true) {

        // compute
        array up = um + convolve(eta, 2, diff_x_dims, h_diff_kernel);
        array vp = um + convolve(eta, 2, diff_y_dims, h_diff_kernel);

        array e = convolve(eta, 2, lap_dims, h_lap_kernel);
        array etap = 2 * eta - etam + (2 * dt * dt) / (dx * dy) * e;
        etam = eta;
        eta = etap;

        if (!console) {
            // viz
            fig("sub",2,1,1);   image(eta);         fig("title","wave propogation");
            fig("sub",2,1,2);   plot3(eta, up, vp);  fig("title","gradients versus magnitude");
        }

        double elapsed = timer::stop(time_last);
        if (elapsed > 1) {
            double rate = (iter - iter_last) / elapsed;
            double total_elapsed = timer::stop(time_start);
            time_last = timer::start();
            iter_last = iter;
            max_rate = std::max(max_rate, rate);
            if (total_elapsed >= time_total) {
                break;
            }
            if (!console)
                printf("  iterations per second: %.0f   (progress %.0f%%)\n",
                       rate, 100.0f * total_elapsed / time_total);
        }
        iter++;
    }

    if (console) {
        printf(" ### shallow_water_example: %f Iterations per second\n", max_rate);
    }

}

int main(int argc, char* argv[]) {

    bool console = false;
    if (argc > 2 || (argc == 2 && strcmp(argv[1], "-"))) {
        printf("usage: swe [-]\n");
        return -1;
    } else if (argc == 2 && !strcmp(argv[1], "-")) {
        console = true;
    }

    printf("simulation of shallow water equations\n");
    try {
        swe_demo(console);
    } catch (af::exception& e) {
        fprintf(stderr, "%s\n", e.what());
    }

#ifdef WIN32 // pause in Windows
    if (!console) {
        printf("hit [enter]...");
        fflush(stdout);
        getchar();
    }
#endif
    return 0;
}