From 8df35f90f5d15b3bf88639569a22846fbfab99a8 Mon Sep 17 00:00:00 2001
From: Cory <cory@balaton.dev>
Date: Sun, 8 Oct 2023 20:29:41 +0200
Subject: [PATCH] Add particle_interaction parameter

---
 include/PenningTrap.hpp | 11 ++++--
 src/PenningTrap.cpp     | 75 +++++++++++++++++++++++++++++++++++------
 2 files changed, 72 insertions(+), 14 deletions(-)

diff --git a/include/PenningTrap.hpp b/include/PenningTrap.hpp
index 3831eb4..bd88c7f 100644
--- a/include/PenningTrap.hpp
+++ b/include/PenningTrap.hpp
@@ -84,13 +84,18 @@ public:
 
     /** @brief Go forward one timestep using the RK4 method
      * */
-    void evolve_RK4(double dt);
+    void evolve_RK4(double dt, bool particle_interaction = true);
 
     /** @brief Go forward one timestep using the forward Euler method
      * */
-    void evolve_forward_euler(double dt);
+    void evolve_forward_euler(double dt, bool particle_interaction = true);
 
-    sim_arr simulate(double time, int steps, std::string method = "rk4");
+    sim_arr simulate(double time, int steps, std::string method = "rk4",
+                     bool particle_interaction = true);
+
+    void write_simulation_to_dir(std::string path, double time, int steps,
+                                 std::string method = "rk4",
+                                 bool particle_interaction = true);
 };
 
 #endif
diff --git a/src/PenningTrap.cpp b/src/PenningTrap.cpp
index 694b2fd..62c608a 100644
--- a/src/PenningTrap.cpp
+++ b/src/PenningTrap.cpp
@@ -17,6 +17,8 @@
 #include "constants.hpp"
 #include "utils.hpp"
 #include <cstdlib>
+#include <functional>
+#include <string>
 #include <vector>
 
 PenningTrap::PenningTrap(double B_0, double V_0, double d)
@@ -126,10 +128,23 @@ arma::vec PenningTrap::total_force(int i)
     return this->total_force_external(i) - this->total_force_particles(i);
 }
 
-void PenningTrap::evolve_RK4(double dt)
+void PenningTrap::evolve_RK4(double dt, bool particle_interaction)
 {
+
     std::vector<Particle> tmp_particles = this->particles;
 
+    std::function<arma::vec(int)> force;
+    if (particle_interaction) {
+        force = [this](int i) {
+            return this->total_force(i);
+        };
+    }
+    else {
+        force = [this](int 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 =
@@ -197,49 +212,87 @@ void PenningTrap::evolve_RK4(double dt)
     delete[] k_r;
 }
 
-void PenningTrap::evolve_forward_euler(double dt)
+void PenningTrap::evolve_forward_euler(double dt, bool particle_interaction)
 {
     std::vector<Particle> new_state = this->particles;
 
     Particle *p;
 
+    std::function<arma::vec(int)> force;
+    if (particle_interaction) {
+        force = [this](int i) { return this->total_force(i); };
+    }
+    else {
+        force = [this](int i) { return this->total_force_external(i); };
+    }
+
 #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->v_vec += dt * 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)
+sim_arr PenningTrap::simulate(double time, int steps, std::string method,
+                              bool particle_interaction)
 {
     double dt = time / (double)steps;
     sim_arr res(this->particles.size(), sim_cols(steps));
 
-    std::function<void(double)> func;
+    std::function<void(double, bool)> func;
     if (method == "rk4") {
-        func = [this](double dt) { this->evolve_RK4(dt); };
+        func = [this](double dt, bool particle_interaction) {
+            this->evolve_RK4(dt, particle_interaction);
+        };
     }
     else if (method == "euler") {
-        func = [this](double dt) { this->evolve_forward_euler(dt); };
+        func = [this](double dt, bool particle_interaction) {
+            this->evolve_forward_euler(dt, particle_interaction);
+        };
     }
     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++) {
+        for (int i = 0; i < this->particles.size(); i++) {
             res[i][j] = this->particles[i].r_vec;
         }
-        func(dt);
+        func(dt, particle_interaction);
     }
 
     return res;
 }
 
+void PenningTrap::write_simulation_to_dir(std::string path, double time,
+                                          int steps, std::string method,
+                                          bool particle_interaction)
+{
+    if (path.back() != '/') {
+        path += '/';
+    }
+    if (mkpath(path, 0777) != 0) {
+        std::cout << "Hello" << std::endl;
+        return;
+    }
 
+    sim_arr res = this->simulate(time, steps, method, particle_interaction);
+
+    int number_of_particles = this->particles.size();
+
+    std::ofstream ofile;
+    sim_rows row;
+
+    for (int i = 0; i < number_of_particles; i++) {
+        row = res[i];
+        ofile.open(path + "particle_" + std::to_string(i) + ".txt");
+        for (arma::vec vector : row) {
+            ofile << vector(0) << "," << vector(1) << "," << vector(2) << "\n";
+        }
+        ofile.close();
+    }
+}