n-body-2.c

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

struct model;

struct model* model_new(int capacity, double dt, double G);
void model_free(struct model*);
int model_body_add(struct model* model, double* r, double* v, double m, int fixed);
void model_step_naive(struct model* model);
void model_step_runge4(struct model* model);
void model_print(struct model* model);

typedef void (*model_step_t)(struct model* );
struct body
{
    double m;
    int fixed;
};

struct model
{
    int capacity;
    int n;
    int dim;

    double dt;
    double G;

    double* r;
    double* v;

    double* k1r;
    double* k2r;
    double* k3r;
    double* k4r;

    double* k1v;
    double* k2v;
    double* k3v;
    double* k4v;

    double* tmp1;
    double* tmp2;

    struct body* bodies;

    FILE* output;
};

struct model* model_new(int capacity, double dt, double G)
{
    int dim = 3;
    struct model* m = malloc(sizeof(struct model));
    m->capacity = capacity;
    m->n = 0;
    m->dim = dim;
    m->dt = dt;
    m->G = G;
    m->r = malloc(dim * capacity * sizeof(double));
    m->v = malloc(dim * capacity * sizeof(double));
    m->bodies = malloc(capacity * sizeof(struct body));

    m->k1r = malloc(dim * capacity * sizeof(double));
    m->k2r = malloc(dim * capacity * sizeof(double));
    m->k3r = malloc(dim * capacity * sizeof(double));
    m->k4r = malloc(dim * capacity * sizeof(double));

    m->k1v = malloc(dim * capacity * sizeof(double));
    m->k2v = malloc(dim * capacity * sizeof(double));
    m->k3v = malloc(dim * capacity * sizeof(double));
    m->k4v = malloc(dim * capacity * sizeof(double));

    m->tmp1 = malloc(dim * capacity * sizeof(double));
    m->tmp2 = malloc(dim * capacity * sizeof(double));

    m->output = stdout; //fopen("output.txt", "wb");

    return m;
}

void model_free(struct model* m) {
    free(m->r);
    free(m->v);
    free(m->bodies);

    free(m->k1r);
    free(m->k2r);
    free(m->k3r);
    free(m->k4r);

    free(m->k1v);
    free(m->k2v);
    free(m->k3v);
    free(m->k4v);

    free(m->tmp1);
    free(m->tmp2);

    fclose(m->output);

    free(m);
}

void model_print(struct model* model) {
    int dim = model->dim;

    for (int i = 0; i < model->n; ++i) {
        double x = model->r[i * dim];
        double y = model->r[i * dim + 1];

        double vx = model->v[i * dim];
        double vy = model->v[i * dim + 1];

        fprintf(model->output, "%lf %lf %lf ", x, y, sqrt(vx*vx + vy*vy));
    }
    fprintf(model->output, "\n");
}

int model_body_add(struct model* m, double* r, double* v, double mass, int fixed)
{
    if (m->n >= m->capacity) {
        return -1;
    } else {
        int i = m->n++;
        int dim = m->dim;
        char buf[256];
        memcpy(&m->r[i * dim], r, dim * sizeof(double));
        memcpy(&m->v[i * dim], v, dim * sizeof(double));
        m->bodies[i].m = mass;
        m->bodies[i].fixed = fixed;
        return 0;
    }
}

static void fr(struct model* model, double* f, const double* r)
{
    memcpy(f, r, model->n * model->dim * sizeof(double));
}

static void fv(struct model* model, double* f, const double* r)
{
    double R;
    double* ff;
    const double* r1;
    const double* r2;
    int dim = model->dim;
    double G= model->G;

    for (int i = 0; i < model->n; ++i) {
        ff = &f[i * dim];

        if (model->bodies[i].fixed) {
            memcpy(ff, r, dim * sizeof(double));
            continue;
        }

        for (int k = 0; k < dim; ++k) {
            ff[k] = 0;
        }

        for (int j = 0; j < model->n; ++j) {
            if (i == j) {
                continue;
            }

            r1 = &r[i * dim];
            r2 = &r[j * dim];

            R = 0.0;
            for (int k = 0; k < dim; ++k) {
                R += (r1[k] - r2[k]) * (r1[k] - r2[k]);
            }
            R = sqrt(R);

            for (int k = 0; k < dim; ++k) {
                ff[k] += G*model->bodies[j].m * (r2[k] - r1[k]) / R / R / R;
            }
        }
    }
}

void model_step_naive(struct model* model)
{
    int dim = model->dim;
    double dt = model->dt;
    double* tmp = model->k1v;

    fv(model, tmp, model->r);

    for (int i = 0; i < model->n; ++i) {
        if (model->bodies[i].fixed) {
            continue;
        }

        for (int k = 0; k < model->dim; ++k) {
            model->v[i * dim + k] += dt * tmp[i * dim + k];
        }
    }

    for (int i = 0; i < model->n; ++i) {
        if (model->bodies[i].fixed) {
            continue;
        }

        for (int k = 0; k < model->dim; ++k) {
            model->r[i * dim + k] += dt * model->v[i * dim + k];
        }
    }
}

void model_step_runge4(struct model* model)
{
    int dim = model->dim;
    double dt = model->dt;

    fv(model, model->k1v, model->r);
    fr(model, model->k1r, model->v);

    for (int i = 0; i < model->n * dim; ++i) {
        model->tmp1[i] = model->v[i] + 0.5 * dt * model->k1v[i];
        model->tmp2[i] = model->r[i] + 0.5 * dt * model->k1r[i];
    }
    fv(model, model->k2v, model->tmp2);
    fr(model, model->k2r, model->tmp1);

    for (int i = 0; i < model->n * dim; ++i) {
        model->tmp1[i] = model->v[i] + 0.5 * dt * model->k2v[i];
        model->tmp2[i] = model->r[i] + 0.5 * dt * model->k2r[i];
    }
    fv(model, model->k3v, model->tmp2);
    fr(model, model->k3r, model->tmp1);

    for (int i = 0; i < model->n * dim; ++i) {
        model->tmp1[i] = model->v[i] + dt * model->k3v[i];
        model->tmp2[i] = model->r[i] + dt * model->k3r[i];
    }
    fv(model, model->k4v, model->tmp2);
    fr(model, model->k4r, model->tmp1);

    for (int i = 0; i < model->n; ++i) {
        if (model->bodies[i].fixed) {
            continue;
        }
        for (int k = 0; k < dim; ++k) {
            int j = i * 3 + k;
            model->v[j] += 1.0 / 6.0 * dt * (model->k1v[j] + 2.0 * model->k2v[j] + 2.0 * model->k3v[j] + model->k4v[j]);
            model->r[j] += 1.0 / 6.0 * dt * (model->k1r[j] + 2.0 * model->k2r[j] + 2.0 * model->k3r[j] + model->k4r[j]);
        }
    }
}

/*
runge4
0.01
8
2.92e-6
10
  0  0       0  0    0  0  333333  1
  0  0.39    0  1.58 0  0  0.038   0
  0  0.72    0  1.17 0  0  0.82    0
  0  1       0  1    0  0  1       0
  0  1.00256 0  1.03 0  0  0.012   0
  0  1.51    0  0.8  0  0  0.1     0
  0  5.2     0  0.43 0  0  317     0
  0  9.3     0  0.32 0  0  95      0
  0  19.3    0  0.23 0  0  14.5    0
  0  30      0  0.18 0  0  16.7    0
# x  y       z  vx   vy vz mass    fixed
*/
int main(int argc, char** argv)
{
    int n;
    double dt = 0.01;
    struct model* m;
    char method[256];
    model_step_t step;
    FILE * f;
    double t, T;
    double G;

    if (argc < 2) {
        fprintf(stderr, "config file requered\n");
        return -1;
    }

    f = fopen(argv[1], "rb");
    if (!f) {
        fprintf(stderr, "cannot open config file '%s'\n", argv[1]);
    }

    if (fscanf(f, "%255s", method) != 1) {
        fprintf(stderr, "cannot fscanf method\n");
        return -1;
    }

    if (!strcmp(method, "runge4")) {
        step = model_step_runge4;
    } else if (!strcmp(method, "naive")) {
        step = model_step_naive;
    } else {
        fprintf(stderr, "unknown method '%s'\n", method);
        return -1;
    }

    if (fscanf(f, "%lf", &dt) != 1) {
        fprintf(stderr, "cannot fscanf dt\n");
        return -1;
    };

    if (fscanf(f, "%lf", &T) != 1) {
        fprintf(stderr, "cannot fscanf T\n");
        return -1;
    };

    if (fscanf(f, "%lf", &G) != 1) {
        fprintf(stderr, "cannot fscanf G\n");
        return -1;
    };

    if (fscanf(f, "%d", &n) != 1) {
        fprintf(stderr, "cannot fscanf n\n");
        return -1;
    }

    m = model_new(n, dt, G);

    for (int i = 0; i < n; ++i) {
        double r[3];
        double v[3];
        double mass;
        int fixed;

        if (fscanf(f, "%lf%lf%lf%lf%lf%lf%lf%d", &r[0], &r[1], &r[2], &v[0], &v[1], &v[2], &mass, &fixed) != 8) {
            fprintf(stderr, "bad config format\n");
            return -1;
        }
        printf("add body: %lf %lf %lf %lf %lf %lf %lf %d\n", r[0], r[1], r[2], v[0], v[1], v[2], mass, fixed);

        if (model_body_add(m, r, v, mass, fixed) != 0) {
            fprintf(stderr, "cannot add body\n");
            return -1;
        }
    }

    while (t < T) {
        model_print(m);
        step(m);
        t += dt;
    }

    model_print(m);
    model_free(m);
}