Project-3/src/PenningTrap.cpp

/** @file PenningTrap.cpp
 *
 *  @author Cory Alexander Balaton (coryab)
 *  @author Janita Ovidie Sandtrøen Willumsen (janitaws)
 *
 *  @version 0.1
 *
 *  @brief The implementation of the PenningTrap class.
 *
 *  @bug No known bugs
 *
 *  @todo Implement evolve_RK4
 *  @todo Implement evolve_forward_euler
 * */

#include "PenningTrap.hpp"
#include "constants.hpp"
#include "utils.hpp"
#include <cstdlib>

PenningTrap::PenningTrap(double B_0, double V_0, double d)
{
    this->B_0 = B_0;
    this->V_0 = V_0;
    this->d = d;
}

void PenningTrap::add_particle(Particle particle)
{
    this->particles.push_back(particle);
}

arma::vec PenningTrap::external_E_field(arma::vec r)
{
    arma::vec::fixed<3> res;
    double f = this->V_0 / (this->d * this->d);
    res(0) = r(0);
    res(1) = r(1);
    res(2) = -2. * r(2);

    return f * res;
}

arma::vec PenningTrap::external_B_field(arma::vec r)
{
    arma::vec::fixed<3> res{0., 0., this->B_0};

    return res;
}

arma::vec PenningTrap::force_on_particle(int i, int j)
{
    // Calculate the difference between the particles' position
    arma::vec::fixed<3> res =
        this->particles.at(i).r_vec - this->particles.at(j).r_vec;

    // Get the distance between the particles
    double norm = arma::norm(res);

    // Multiply res with p_j's charge divided by the norm cubed
    res *= this->particles.at(j).q / (norm * norm * norm);

    return res;
}

arma::vec PenningTrap::total_force_external(int i)
{
    Particle p = this->particles.at(i);

    arma::vec::fixed<3> B = this->external_B_field(p.r_vec);

    arma::vec::fixed<3> v_cross_B{p.v_vec(1) * B(2) - p.v_vec(2) * B(1),
                                  p.v_vec(2) * B(0) - p.v_vec(0) * B(2),
                                  p.v_vec(0) * B(1) - p.v_vec(1) * B(0)};

    arma::vec force = p.q * (this->external_E_field(p.r_vec) + v_cross_B);

    return force;
}

arma::vec PenningTrap::total_force_particles(int i)
{
    Particle p = this->particles.at(i);

    arma::vec res(3);

    for (int j = 0; j < this->particles.size(); j++) {
        if (i == j) {
            continue;
        }

        res += this->force_on_particle(i, j);
    }

    res *= K_E * (p.q / p.m);

    return res;
}

arma::vec PenningTrap::total_force(int i)
{
    return this->total_force_external(i) - this->total_force_particles(i);
}

void PenningTrap::evolve_RK4(double dt)
{
    std::vector<Particle> tmp_particles = this->particles;

    arma::vec::fixed<3> *k_v =
        new arma::vec::fixed<3>[this->particles.size() * 4];
    arma::vec::fixed<3> *k_r =
        new arma::vec::fixed<3>[this->particles.size() * 4];

    int size = this->particles.size();

    for (int i = 0; i < size; i++) {
        k_v[i] = this->total_force(i) / this->particles.at(i).m;
        k_r[i] = this->particles.at(i).v_vec;
    }

    for (int i = 0; i < size; i++) {
        Particle *p = &this->particles.at(i);

        p->v_vec = tmp_particles.at(i).v_vec + (dt / 2) * k_v[i];
        p->r_vec = tmp_particles.at(i).r_vec + (dt / 2) * k_r[i];
    }

    for (int i = 0; i < size; i++) {
        k_v[1 * size + i] = this->total_force(i) / this->particles.at(i).m;
        k_r[1 * size + i] = this->particles.at(i).v_vec;
    }

    for (int i = 0; i < size; i++) {
        Particle *p = &this->particles.at(i);

        p->v_vec = tmp_particles.at(i).v_vec + (dt / 2) * k_v[1 * size + i];
        p->r_vec = tmp_particles.at(i).r_vec + (dt / 2) * k_r[1 * size + i];
    }

    for (int i = 0; i < size; i++) {
        k_v[2 * size + i] = this->total_force(i) / this->particles.at(i).m;
        k_r[2 * size + i] = this->particles.at(i).v_vec;
    }

    for (int i = 0; i < size; i++) {
        Particle *p = &this->particles.at(i);

        p->v_vec = tmp_particles.at(i).v_vec + dt * k_v[2 * size + i];
        p->r_vec = tmp_particles.at(i).r_vec + dt * k_r[2 * size + i];
    }

    for (int i = 0; i < size; i++) {
        k_v[3 * size + i] = this->total_force(i) / this->particles.at(i).m;
        k_r[3 * size + i] = this->particles.at(i).v_vec;
    }

    for (int i = 0; i < size; i++) {
        Particle *p = &this->particles.at(i);

        p->v_vec = tmp_particles.at(i).v_vec +
                   dt *
                       (k_v[i] + k_v[size + i] + k_v[2 * size + i] +
                        k_v[3 * size + i]) /
                       6;
        p->r_vec = tmp_particles.at(i).r_vec +
                   dt *
                       (k_r[i] + k_r[size + i] + k_r[2 * size + i] +
                        k_r[3 * size + i]) /
                       6;
    }

    delete[] k_v;
    delete[] k_r;
}

void PenningTrap::evolve_forward_euler(double dt)
{
    std::vector<Particle> new_state = this->particles;

    Particle *p;

#pragma omp parallel for private(p)
    for (int i = 0; i < this->particles.size(); i++) {
        p = &new_state.at(i);
        p->v_vec += dt * this->total_force(i) / new_state.at(i).m;
        p->r_vec += dt * this->particles.at(i).v_vec;
    }

    this->particles = new_state;
}

sim_arr PenningTrap::simulate(double time, int steps, std::string method)
{
    double dt = time / (double)steps;
    sim_arr res(this->particles.size(), sim_cols(steps));

    std::function<void(double)> func;
    if (method == "rk4") {
        func = [this](double dt) {
            this->evolve_RK4(dt);
        };
    }
    else if(method == "euler") {
        func = [this](double dt) {
            this->evolve_forward_euler(dt);
        };
    }
    else {
        std::cout << "Not a valid method!" << std::endl;
        abort();
    }

    int size = this->particles.size();

    for (int j=0; j<steps; j++) {
        for (int i=0;i<size;i++) {
            res[i][j] = this->particles[i].r_vec;
        }
        func(dt); 
    }

    return res;
}

arma::vec PenningTrap::get_particle(int i)
{
    return this->particles.at(i).r_vec;
}

double PenningTrap::get_d()
{
    return this->d;
}