Implement simulate method

include/PenningTrap.hpp

@@ -18,6 +18,10 @@
 #include "Particle.hpp"
 #include "constants.hpp"
 
+typedef std::vector<arma::vec::fixed<3>> sim_cols;
+typedef std::vector<arma::vec::fixed<3>> sim_rows;
+typedef std::vector<sim_cols> sim_arr;
+
 #pragma omp declare reduction(+ : arma::vec : omp_out += omp_in)               \
     initializer(omp_priv = omp_orig)
 
@@ -80,6 +84,8 @@ public:
      * */
     void evolve_forward_euler(double dt);
 
+    sim_arr simulate(double time, int steps, std::string method = "rk4");
+
     arma::vec get_particle(int i);
 
     double get_d();

src/PenningTrap.cpp

@@ -16,6 +16,7 @@
 #include "PenningTrap.hpp"
 #include "constants.hpp"
 #include "utils.hpp"
+#include <cstdlib>
 
 PenningTrap::PenningTrap(double B_0, double V_0, double d)
 {
@@ -105,8 +106,10 @@ 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_v =
-    arma::vec::fixed<3> *k_r = new arma::vec::fixed<3>[this->particles.size()*4];
+        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();
 
@@ -122,7 +125,6 @@ void PenningTrap::evolve_RK4(double dt)
         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;
@@ -147,7 +149,6 @@ void PenningTrap::evolve_RK4(double dt)
         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;
@@ -156,8 +157,16 @@ void PenningTrap::evolve_RK4(double dt)
     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->v_vec = tmp_particles.at(i).v_vec +
-        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;
+                   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;
@@ -180,6 +189,39 @@ void PenningTrap::evolve_forward_euler(double dt)
     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;