main.cpp

#include <cstdio>
#include <vapor/vapor.hpp>


using namespace vapor;



using particle_t = vec_t<double, 6>;
using particle_array_t = memory_backed_array_t<1, particle_t, ref_counted_ptr_t>;
auto G = 1.0;
auto M = 1.0;




struct State
{
    double time;
    double iter;
    particle_array_t part;
};

static auto average(const State& a, const State& b, double x) -> State
{
    return x == 1.0 ? b : State{
        (a.time * (1.0 - x) + b.time * x),
        (a.iter * (1.0 - x) + b.iter * x),
        (a.part * (1.0 - x) + b.part * x).cache(),
    };
}





auto pos(const particle_t& particle)
{
    return vec(particle[0], particle[1], particle[2]);
}

auto vel(const particle_t& particle)
{
    return vec(particle[3], particle[4], particle[5]);
}

auto norm(const vec_t<double, 3>& v)
{
    return sqrt(v[0] * v[0] + v[1] * v[1] + v[2] * v[2]);
}

auto initial_particles()
{
    auto num_particles = 100;
    auto r0 = 1.0;
    auto r1 = 2.0;
    auto dr = (r1 - r0) / (num_particles - 1);
    return range(num_particles).map([=] (int i) {
        auto r = r0 + i * dr;
        auto v = sqrt(G * M / r);
        auto p = particle_t{r, 0.0, 0.0, 0.0, v, 0.0};
        return p;
    }).cache();
}

auto initial_state()
{
    return State{
        0.0,
        0.0,
        initial_particles(),
    };
}

auto gravity(vec_t<double, 3> x)
{
    auto rhat = x / norm(x);
    auto r2 = dot(x, x);
    auto a = -G * M / r2 * rhat;
    return a;
}

auto min_dt(const State& state)
{
    return min(state.part.map([] (auto p)
    {
        auto x = pos(p);
        auto v = vel(p);
        auto a = gravity(x);
        return norm(v) / norm(a);
    }));
}

auto time_derivative(particle_array_t particles)
{
    return particles.map([] (auto p)
    {
        auto x = pos(p);
        auto v = vel(p);
        auto a = gravity(x);
        return concat(v, a);
    });
}

auto next(State state, double dt)
{
    return State{
        state.time + dt,
        state.iter + 1,
        (state.part + time_derivative(state.part) * dt).cache()
    };
}

static void update_state(State& state)
{
    auto rk_order = 3;
    auto s0 = state;
    auto dt = min_dt(state) * 0.1;

    switch (rk_order)
    {
        case 1: {
            state = next(s0, dt);
            break;
        }
        case 2: {
            auto s1 = average(s0, next(s0, dt), 1./1);
            auto s2 = average(s0, next(s1, dt), 1./2);
            state = s2;
            break;
        }
        case 3: {
            auto s1 = average(s0, next(s0, dt), 1./1);
            auto s2 = average(s0, next(s1, dt), 1./4);
            auto s3 = average(s0, next(s2, dt), 2./3);
            state = s3;
            break;
        }
    }
}

void output(particle_array_t particles, const char* filename)
{
    FILE* file = fopen(filename, "w");
    if (!file)
    {
        perror("Failed to open file");
        return;
    }
    for (int n = 0; n < particles.size(); ++n)
    {
        auto p = particles[n];
        fprintf(
            file, "%+12.10e %+12.10e %+12.10e %+12.10e %+12.10e %+12.10e\n",
            p[0], p[1], p[2], p[3], p[4], p[5]
        );
    }
    fclose(file);
}

void output_bin(particle_array_t particles, const char* filename)
{
    FILE* file = fopen(filename, "wb");
    if (!file)
    {
        perror("Failed to open file");
        return;
    }
    for (int n = 0; n < particles.size(); ++n)
    {
        auto p = particles[n];
        fwrite(&p, sizeof(particle_t), 1, file);
    }
    fclose(file);
}

int main()
{
    auto state = initial_state();
    auto num_outputs = 0;

    while (state.time < 2 * M_PI)
    {
        update_state(state);
        if (int(state.iter) % 1 == 0) {
            auto fname = format("dust.%04d.bin", num_outputs);
            output_bin(state.part, fname.data);
            num_outputs++;
        }
        printf("[%04d] t=%.3f\n", int(state.iter), state.time);
    }
    return 0;
}