include/PenningTrap.hpp

@@ -36,6 +36,10 @@ private:
     double V_0;                      ///< Applied potential
     double d;                        ///< Characteristic dimension
     std::vector<Particle> particles; ///< The particles in the Penning trap
+    arma::vec::fixed<3> *k_v;
+    arma::vec::fixed<3> *k_r;
+    std::function<arma::vec(int, double)> v_funcs[4];
+    std::function<arma::vec(int, double)> r_funcs[4];
 
 public:
     /** @brief Set B_0, V_0 and d.

src/PenningTrap.cpp

@@ -26,28 +26,51 @@ PenningTrap::PenningTrap(double B_0, double V_0, double d)
     this->B_0 = B_0;
     this->V_0 = V_0;
     this->d = d;
+
+    v_funcs[0] = [this](int i, double dt) { return .5 * dt * this->k_v[i]; };
+    v_funcs[1] = [this](int i, double dt) {
+        return .5 * dt * this->k_v[this->particles.size() + i];
+    };
+    v_funcs[2] = [this](int i, double dt) {
+        return dt * this->k_v[2 * this->particles.size() + i];
+    };
+    v_funcs[3] = [this](int i, double dt) {
+        arma::vec res = this->k_v[i] + this->k_v[this->particles.size() + i] +
+                        this->k_v[2 * this->particles.size() + i] +
+                        this->k_v[3 * this->particles.size() + i];
+        return (dt / 6) * res;
+    };
+
+    r_funcs[0] = [this](int i, double dt) { return .5 * dt * this->k_r[i]; };
+    r_funcs[1] = [this](int i, double dt) {
+        return .5 * dt * this->k_r[this->particles.size() + i];
+    };
+    r_funcs[2] = [this](int i, double dt) {
+        return dt * this->k_r[2 * this->particles.size() + i];
+    };
+    r_funcs[3] = [this](int i, double dt) {
+        arma::vec res = this->k_r[i] + this->k_r[this->particles.size() + i] +
+                        this->k_r[2 * this->particles.size() + i] +
+                        this->k_r[3 * this->particles.size() + i];
+        return (dt / 6) * res;
+    };
 }
 
 PenningTrap::PenningTrap(int i, double B_0, double V_0, double d)
+    : PenningTrap::PenningTrap(B_0, V_0, d)
 {
-    this->B_0 = B_0;
-    this->V_0 = V_0;
-    this->d = d;
-
     arma::vec r, v;
     for (int j = 0; j < i; j++) {
-        r = arma::vec(3).randn() * .1 * this->d;
+        r = arma::vec::fixed<3>().randn() * .1 * this->d;
-        v = arma::vec(3).randn() * .1 * this->d;
+        v = arma::vec::fixed<3>().randn() * .1 * this->d;
         this->add_particle(Particle(1., 40., r, v));
     }
 }
 
 PenningTrap::PenningTrap(std::vector<Particle> particles, double B_0,
                          double V_0, double d)
+    : PenningTrap::PenningTrap(B_0, V_0, d)
 {
-    this->B_0 = B_0;
-    this->V_0 = V_0;
-    this->d = d;
     this->particles = particles;
 }
 
@@ -58,11 +81,8 @@ void PenningTrap::add_particle(Particle particle)
 
 arma::vec PenningTrap::external_E_field(arma::vec r)
 {
-    arma::vec::fixed<3> res;
+    arma::vec::fixed<3> res{r(0), r(1), -2. * r(2)};
     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;
 }
@@ -91,6 +111,10 @@ arma::vec PenningTrap::force_on_particle(int i, int j)
 
 arma::vec PenningTrap::total_force_external(int i)
 {
+    if (arma::norm(this->particles.at(i).r_vec) > this->d) {
+        return arma::vec::fixed<3>{0., 0., 0.};
+    }
+
     Particle p = this->particles.at(i);
 
     arma::vec::fixed<3> B = this->external_B_field(p.r_vec);
@@ -99,7 +123,8 @@ arma::vec PenningTrap::total_force_external(int i)
                                   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);
+    arma::vec::fixed<3> force =
+        p.q * (this->external_E_field(p.r_vec) + v_cross_B);
 
     return force;
 }
@@ -108,7 +133,7 @@ arma::vec PenningTrap::total_force_particles(int i)
 {
     Particle p = this->particles.at(i);
 
-    arma::vec res(3);
+    arma::vec::fixed<3> res;
 
     for (int j = 0; j < this->particles.size(); j++) {
         if (i == j) {
@@ -131,81 +156,34 @@ arma::vec PenningTrap::total_force(int i)
 void PenningTrap::evolve_RK4(double dt, bool particle_interaction)
 {
 
+    std::vector<Particle> original_particles = this->particles;
     std::vector<Particle> tmp_particles = this->particles;
 
     std::function<arma::vec(int)> force;
     if (particle_interaction) {
-        force = [this](int i) {
+        force = [this](int i) { return this->total_force(i); };
-            return this->total_force(i);
-        };
     }
     else {
-        force = [this](int i) {
+        force = [this](int i) { return this->total_force_external(i); };
-            return this->total_force_external(i);
-        };
     }
 
-    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 = new arma::vec::fixed<3>[size * 4];
-        k_v[i] = this->total_force(i) / this->particles.at(i).m;
+    k_r = new arma::vec::fixed<3>[size * 4];
-        k_r[i] = this->particles.at(i).v_vec;
-    }
 
-    for (int i = 0; i < size; i++) {
+    for (int i = 0; i < 4; i++) {
-        Particle *p = &this->particles.at(i);
+#pragma omp parallel for
+        for (int j = 0; j < size; j++) {
+            k_v[i * size + j] = this->total_force(j) / this->particles.at(j).m;
+            k_r[i * size + j] = this->particles.at(j).v_vec;
 
-        p->v_vec = tmp_particles.at(i).v_vec + (dt / 2) * k_v[i];
+            Particle *p = &tmp_particles.at(j);
-        p->r_vec = tmp_particles.at(i).r_vec + (dt / 2) * k_r[i];
-    }
 
-    for (int i = 0; i < size; i++) {
+            p->v_vec = original_particles.at(j).v_vec + this->v_funcs[i](j, dt);
-        k_v[1 * size + i] = this->total_force(i) / this->particles.at(i).m;
+            p->r_vec = original_particles.at(j).r_vec + this->r_funcs[i](j, dt);
-        k_r[1 * size + i] = this->particles.at(i).v_vec;
+        }
-    }
+        this->particles = tmp_particles;
-
-    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;
@@ -259,6 +237,7 @@ sim_arr PenningTrap::simulate(double time, int steps, std::string method,
     }
 
     for (int j = 0; j < steps; j++) {
+        #pragma omp parallel for
         for (int i = 0; i < this->particles.size(); i++) {
             res[i][j] = this->particles[i].r_vec;
         }
@@ -287,6 +266,7 @@ void PenningTrap::write_simulation_to_dir(std::string path, double time,
     std::ofstream ofile;
     sim_rows row;
 
+#pragma omp parallel for private(ofile, row)
     for (int i = 0; i < number_of_particles; i++) {
         row = res[i];
         ofile.open(path + "particle_" + std::to_string(i) + ".txt");

src/animate_100_particles.py

@@ -63,8 +63,8 @@ def animate():
         fig, update, N, fargs=(lines, arr), interval=1, blit=False
     )
 
-    ani.save("../images/100_particles.gif", writer=animation.FFMpegFileWriter(fps=50))
+    # ani.save("../images/100_particles.gif", writer=animation.FFMpegFileWriter(fps=50))
-    # plt.show()
+    plt.show()
 
 
 if __name__ == "__main__":