diff --git a/src/main.cpp b/src/main.cpp
index 1c55279..b6d83c7 100644
--- a/src/main.cpp
+++ b/src/main.cpp
@@ -21,6 +21,89 @@
 #define CHARGE 1.
 #define MASS 40. // unit: amu
 
+Particle p1(1., 40., arma::vec{20., 0., 20.}, arma::vec{0., 25., 0.});
+Particle p2(1., 40., arma::vec{25., 25., 0.}, arma::vec{0., 40., 5.});
+
+void simulate_single_particle()
+{
+    PenningTrap trap;
+
+    trap.add_particle(p1);
+
+    double time = 50.; // microseconds
+    double dt = time / (double)N;
+    
+    auto res = new arma::vec::fixed<3>[N];
+
+    for (int i=0; i<N;i++) {
+        res[i] = trap.get_particle(0);
+        trap.evolve_RK4(dt);
+    }
+
+    mkdir("output", 0777);
+    mkdir("output/simulate_single_particle", 0777);
+
+    std::ofstream ofile;
+
+    ofile.open("output/simulate_single_particle/out.txt");
+    for (int i=0; i<N;i++) {
+        arma::vec p = res[i];
+        ofile << p(0) << "," << p(1) << "," << p(3) << "\n";
+    }
+    ofile.close();
+    delete [] res;
+}
+
+void simulate_two_particles()
+{
+    PenningTrap trap;
+
+    trap.add_particle(p1);
+    trap.add_particle(p2);
+
+    double time = 50.; // microseconds
+    double dt = time / (double)N;
+
+    auto res = new arma::vec::fixed<3>[2][N];
+
+
+    int counter = 0;
+
+    // Get the path of all particles
+    for (int j = 0; j < N; j++) {
+#pragma omp parallel for
+        for (int i = 0; i < 2; i++) {
+            res[i][j] = trap.get_particle(i);
+        }
+        trap.evolve_RK4(dt);
+    }
+
+    std::cout << counter << std::endl;
+
+    arma::vec::fixed<3> *cur_row;
+    arma::vec::fixed<3> cur_elem;
+
+    mkdir("output", 0777);
+    mkdir("output/simulate_2_particles", 0777);
+
+    std::ofstream ofile;
+
+// Write particle paths to file
+#pragma omp parallel for private(cur_row, cur_elem, ofile)
+    for (int i = 0; i < 2; i++) {
+        cur_row = res[i];
+        ofile.open("output/simulate_100_particles/p" + std::to_string(i) +
+                   ".txt");
+        for (int j = 0; j < N; j++) {
+            cur_elem = cur_row[j];
+            ofile << cur_elem(0) << "," << cur_elem(1) << "," << cur_elem(2)
+                  << "\n";
+        }
+        ofile.close();
+    }
+
+}
+
 void simulate_100_particles()
 {
     PenningTrap trap;
@@ -37,22 +120,9 @@ void simulate_100_particles()
     double time = 50.; // microseconds
     double dt = time / (double)N;
 
-    auto res = new arma::vec::fixed<3>[PARTICLES][N];
+    sim_arr res = trap.simulate(time, N);
 
-    int counter = 0;
-
-    // Get the path of all particles
-    for (int j = 0; j < N; j++) {
-#pragma omp parallel for
-        for (int i = 0; i < PARTICLES; i++) {
-            res[i][j] = trap.get_particle(i);
-        }
-        trap.evolve_RK4(dt);
-    }
-
-    std::cout << counter << std::endl;
-
-    arma::vec::fixed<3> *cur_row;
+    sim_rows cur_row;
     arma::vec::fixed<3> cur_elem;
 
     mkdir("output", 0777);
@@ -64,7 +134,8 @@ void simulate_100_particles()
 #pragma omp parallel for private(cur_row, cur_elem, ofile)
     for (int i = 0; i < PARTICLES; i++) {
         cur_row = res[i];
-        ofile.open("output/simulate_100_particles/p" + std::to_string(i) + ".txt");
+        ofile.open("output/simulate_100_particles/p" + std::to_string(i) +
+                   ".txt");
         for (int j = 0; j < N; j++) {
             cur_elem = cur_row[j];
             ofile << cur_elem(0) << "," << cur_elem(1) << "," << cur_elem(2)