verlet.c

#include <stdio.h>
#include <string.h>
#include <stdlib.h>
#include <math.h>
#include <time.h>

struct body {
    char name[16];
    char color[16];
    double rad;
    double r[4];
    double v[4];
    double a[4];
    double a_next[4];
    double m;
    double max_rad;
    double min_rad;
    int fixed;
};

struct data {
    int nbodies;
    struct body* bodies;
    double G;
    double dt;
};

void verlet_init(struct data* data) {
    int n = data->nbodies;
    double G = data->G;
    double rij[4];

    // new acc
#pragma omp parallel for
    for (int i = 0; i < n; i++) {
        struct body* b1 = &data->bodies[i];
        if (b1->fixed) continue;

        for (int k = 0; k < 4; k++) {
            b1->a[k] = 0;
        }

        for (int j = 0; j < n; j++) {
            if (i == j) continue;

            struct body* b2 = &data->bodies[j];

            for (int k = 0; k < 4; k++) {
                rij[k] = (b2->r[k] - b1->r[k]);
            }

            double R = 0;
            for (int k = 0; k < 4; k++) {
                R += rij[k] * rij[k];
            }
            R = sqrt(R);
            double Koef = G * b2->m / (R*R*R);

            for (int k = 0; k < 4; k++) {
                b1->a[k] += Koef * rij[k];
            }
        }
    }
}

void verlet_next(struct data* data) {
    int n = data->nbodies;
    double G = data->G;
    double dt = data->dt;

#pragma omp parallel for
    for (int i = 0; i < n; i++) {
        struct body* b = &data->bodies[i];

        for (int k = 0; k < 4; k++) {
            // new pos
            b->r[k] = b->r[k] + b->v[k] * dt + b->a[k] * dt * dt * 0.5;
        }

        double R = 0;
        if (b->min_rad >= 0 || b->max_rad >= 0) {
            for (int k = 0; k < 4; k++) {
                R += b->r[k] * b->r[k];
            }
            R = sqrt(R);
        }

        if (b->min_rad > 0 && R < b->min_rad) {
            for (int k = 0; k < 4; k++) {
                b->r[k] = b->min_rad * b->r[k] / R;
            }
        }
        if (b->max_rad > 0 && R > b->max_rad) {
            for (int k = 0; k < 4; k++) {
                b->r[k] = b->max_rad * b->r[k] / R;
            }
        }
    }

    // new acc
#pragma omp parallel for
    for (int i = 0; i < n; i++) {
        double rij[4]; 
        struct body* b1 = &data->bodies[i];
        if (b1->fixed) continue;

        for (int k = 0; k < 4; k++) {
            b1->a_next[k] = 0;
        }

        for (int j = 0; j < n; j++) {
            if (i == j) continue;

            struct body* b2 = &data->bodies[j];
            for (int k = 0; k < 4; k++) {
                rij[k] = (b2->r[k] - b1->r[k]);
            }

            double R = 0;
            for (int k = 0; k < 3; k++) {
                R += rij[k] * rij[k];
            }
            R = sqrt(R);
            double Koef = G * b2->m / (R*R*R);

            for (int k = 0; k < 3; k++) {
                b1->a_next[k] += Koef * rij[k];
            }
        }
    }

#pragma omp parallel for
    for (int i = 0; i < n; i++) {
        struct body* b = &data->bodies[i];

        for (int k = 0; k < 4; k++) {
            // new vel
            b->v[k] = b->v[k] + 0.5 * dt * (b->a[k] + b->a_next[k]);
            // a = new acc
            b->a[k] = b->a_next[k];
        }
    }
}

double kepler(double dt) {
    double G = 1;
    double MM = 1e5;

    struct body bodies[] = {
        {
            .r = {0, 0, 0, 0},
            .v = {0, 0, 0, 0},
            .a = {0, 0, 0, 0},
            .m = MM,
            .fixed = 1
        },
        {
            .r = {0, 1, 0, 0},
            .v = {sqrt(G * MM), 0, 0, 0},
            .a = {0, 0, 0, 0},
            .m = 1,
            .fixed = 0
        },
    };

    struct data data = {
        .nbodies = 2,
        .bodies = &bodies[0],
        .G = G,
        .dt = dt
    };

    double max_err = 0;

    verlet_init(&data);

    double T = 0.1;
    double t = 0;
    while (t < T) {
        verlet_next(&data);

        double r = 0;
        for (int k = 0; k < 3; k++) {
            r += data.bodies[1].r[k] * data.bodies[1].r[k];
        }
        r = sqrt(r);
        double err = fabs(r - 1.0);
        if (max_err < err) {
            max_err = err;
        }
	t += dt;
    }
    return max_err;
}

void run_test() {
    double err1 = kepler(0.001);
    double err2 = kepler(0.0001);
    double err3 = kepler(0.00001);
    printf("%f %f %f\n", err1, err2, err3);
    if (err1 / 97 < err2) {
        printf("Error1 %f\n", err1/err2);
        exit(1);
    }
    if (err1 / 9700 < err3) {
        printf("Error2 %f\n", err1/err3);
        exit(2);
    }
    printf("Ok\n");
    exit(0);
}

/*
  file format:
  G
  N
  Body1 r0 r1 r2 v0 v1 v2 Mass
  Body2 r0 r1 r2 v0 v1 v2 Mass
  ...
  BodyN r0 r1 r2 v0 v1 v2 Mass
 */

void load(struct data* data, const char* fn) {
    FILE* f = fopen(fn, "rb");
    if (!f) { goto err; }

    if (fscanf(f, "%lf %d", &data->G, &data->nbodies) != 2) { goto err; }
    data->bodies = calloc(data->nbodies, sizeof(struct body));
    for (int i = 0; i < data->nbodies; i++) {
        if (fscanf(
                f, "%15s %lf %lf %lf %lf %lf %lf %lf",
                data->bodies[i].name,
                &data->bodies[i].r[0], &data->bodies[i].r[1], &data->bodies[i].r[2],
                &data->bodies[i].v[0], &data->bodies[i].v[1], &data->bodies[i].v[2],
                &data->bodies[i].m) != 8)
        {
            goto err;
        }
        strcpy(data->bodies[i].color, "000000");
        data->bodies[i].min_rad = -1;
        data->bodies[i].max_rad = -1;
        data->bodies[i].rad = 1;
    }

    // properties
    int i;
    char color[10];
    double min_radius, max_radius;
    double rad;
    while (fscanf(f, "%d %10s %lf %lf %lf", &i, color, &min_radius, &max_radius, &rad) == 5) {
        strcpy(data->bodies[i].color, color);
        data->bodies[i].min_rad = min_radius;
        data->bodies[i].max_rad = max_radius;
        data->bodies[i].rad = rad;
    }

    fclose(f);

    return;

err:
    fprintf(stderr, "Cannot open or parse file: '%s'\n", fn);
    exit(1);
}

void print_header(struct data* data) {
    // column names
    printf("t ");
    for (int i = 0; i < data->nbodies; i++) {
        for (int j = 0; j < 3; j++) {
            printf("r%d,%d ", i, j);
        }
        for (int j = 0; j < 3; j++) {
            printf("v%d,%d ", i, j);
        }
    }
    printf("\n");
    // comment
    for (int i = 0; i < data->nbodies; i++) {
        printf("# %s %le %s %lf\n", data->bodies[i].name, data->bodies[i].m, data->bodies[i].color, data->bodies[i].rad);
    }
}

void print(struct data* data, double t) {
    printf("%e ", t);
    for (int i = 0; i < data->nbodies; i++) {
        printf(
            "%e %e %e %e %e %e ",
            data->bodies[i].r[0], data->bodies[i].r[1], data->bodies[i].r[2],
            data->bodies[i].v[0], data->bodies[i].v[1], data->bodies[i].v[2]);
    }
    printf("\n");
}

void solve(struct data* data, double T) {
    double t = 0;
    print_header(data);
    print(data, t);
    verlet_init(data);
    while (t < T) {
        verlet_next(data);
        t += data->dt;
        print(data, t);
    }
}

double get_seconds_elapsed(struct timespec start, struct timespec end) {
    return (double)(end.tv_sec - start.tv_sec) + (double)(end.tv_nsec - start.tv_nsec) / 1.0e9;
}

void bench(struct data* data) {
    double t = 0;
    struct timespec start_time, end_time;
    double elapsed_time;

    clock_gettime(CLOCK_MONOTONIC, &start_time);
    verlet_init(data);
    while (t < 0.1) {
        verlet_next(data);
        t += data->dt;
    }
    clock_gettime(CLOCK_MONOTONIC, &end_time);
    elapsed_time = get_seconds_elapsed(start_time, end_time);

    printf("Elapsed: %fs\n", elapsed_time);
}

void usage(const char* name) {
    fprintf(stderr, "%s --input file.txt [--dt 0.001] [--T 10] [--test]\n", name);
    exit(0);
}

int main(int argc, char** argv) {
    const char* fn = NULL;
    double dt = 0.0001;
    double T = 10.0;
    int test_mode = 0;
    int bench_mode = 0;
    for (int i = 1; i < argc; i++) {
        if (i < argc - 1 && !strcmp(argv[i], "--input")) {
            fn = argv[++i];
        } else if (i < argc - 1 && !strcmp(argv[i], "--dt")) {
            dt = atof(argv[++i]);
        } else if (i < argc - 1 && !strcmp(argv[i], "--T")) {
            T = atof(argv[++i]);
        } else if (!strcmp(argv[i], "--test")) {
            test_mode = 1;
        } else if (!strcmp(argv[i], "--bench")) {
            bench_mode = 1;
        } else {
            usage(argv[0]);
        }
    }
    if (test_mode) {
        run_test(); return 0;
    }
    if (!fn) {
        usage(argv[0]);
    }

    struct data data = {.dt = dt};
    load(&data, fn);
    if (bench_mode) {
        bench(&data);
    } else {
        solve(&data, T);
    }
    free(data.bodies);

    return 0;
}