diff --git a/src/main.cpp b/src/main.cpp
index 19ef8fd..0644c61 100644
--- a/src/main.cpp
+++ b/src/main.cpp
@@ -13,10 +13,8 @@
 #include <cmath>
 #include <complex>
 #include <fstream>
-#include <functional>
 #include <omp.h>
 #include <string>
-#include <sys/stat.h>
 #include <vector>
 
 #include "PenningTrap.hpp"
@@ -24,9 +22,10 @@
 
 #define PARTICLES 100
 #define N 40000
-#define CHARGE 1.
-#define MASS 40. // unit: amu
+#define CHARGE 1. // unit: e
+#define MASS 40.  // unit: amu
 
+// Particles used for testing
 Particle p1(CHARGE, MASS, vec_3d{20., 0., 20.}, vec_3d{0., 25., 0.});
 Particle p2(CHARGE, MASS, vec_3d{25., 25., 0.}, vec_3d{0., 40., 5.});
 
@@ -38,31 +37,34 @@ vec_3d analytical_solution_particle_1(double t)
     double w_n = (w_0 - std::sqrt(w_0 * w_0 - 2. * w_z2)) / 2.;
     double A_p = (25. + w_n * 20.) / (w_n - w_p);
     double A_n = -(25. + w_p * 20.) / (w_n - w_p);
-    std::complex<double> f = A_p * std::exp(std::complex<double>(0., -w_p * t)) +
-                             A_n * std::exp(std::complex<double>(0., -w_n * t));
+    std::complex<double> f =
+        A_p * std::exp(std::complex<double>(0., -w_p * t)) +
+        A_n * std::exp(std::complex<double>(0., -w_n * t));
     vec_3d res{std::real(f), std::imag(f), 20. * std::cos(std::sqrt(w_z2) * t)};
     return res;
 }
 
 void simulate_single_particle()
 {
-    DEBUG("Inside single particle sim");
+    // Initialize trap with particle 1
     PenningTrap trap(std::vector<Particle>{p1});
 
     double time = 50.; // microseconds
 
-    DEBUG("Write to dir");
+    // Simulate and write results to file
     trap.write_simulation_to_dir("output/simulate_single_particle", time, N,
                                  "rk4", false);
 }
 
 void simulate_two_particles()
 {
+    // Initialize traps with particles
     PenningTrap trap_no_interaction(std::vector<Particle>{p1, p2});
     PenningTrap trap_with_interaction(std::vector<Particle>{p1, p2});
 
     double time = 50.; // microseconds
 
+    // Simulate and write results to files
     trap_no_interaction.write_simulation_to_dir(
         "output/simulate_2_particles/no_interaction", time, N, "rk4", false);
     trap_with_interaction.write_simulation_to_dir(
@@ -71,38 +73,39 @@ void simulate_two_particles()
 
 void simulate_single_particle_with_different_steps()
 {
-
     double time = 50.; // microseconds
 
     std::ofstream ofile;
 
+    // Calculate relative error for RK4
+    std::string path = "output/relative_error/RK4/";
+    mkpath(path);
     for (int i = 0; i < 4; i++) {
         int steps = 4000 * std::pow(2, i);
         double dt = time / (double)steps;
-        std::string path = "output/relative_error/RK4/";
-        mkpath(path);
         ofile.open(path + std::to_string(steps) + "_steps.txt");
         PenningTrap trap(std::vector<Particle>{p1});
+        simulation_t res = trap.simulate(time, steps, "rk4", false);
         for (int i = 0; i < steps; i++) {
-            trap.evolve_RK4(dt);
-            ofile << arma::norm(trap.get_r(0) -
-                         analytical_solution_particle_1(trap.get_t()))
+            ofile << arma::norm(res.r_vecs[0][i] -
+                                analytical_solution_particle_1(dt*i))
                   << "\n";
         }
         ofile.close();
     }
 
+    // Calculate relative error for forward Euler
+    path = "output/relative_error/euler/";
+    mkpath(path);
     for (int i = 0; i < 4; i++) {
         int steps = 4000 * std::pow(2, i);
         double dt = time / (double)steps;
-        std::string path = "output/relative_error/euler/";
-        mkpath(path);
         ofile.open(path + std::to_string(steps) + "_steps.txt");
         PenningTrap trap(std::vector<Particle>{p1});
+        simulation_t res = trap.simulate(time, steps, "euler", false);
         for (int i = 0; i < steps; i++) {
-            trap.evolve_forward_euler(dt);
-            ofile << arma::norm(trap.get_r(0) -
-                         analytical_solution_particle_1(trap.get_t()))
+            ofile << arma::norm(res.r_vecs[0][i] -
+                                analytical_solution_particle_1(dt*i))
                   << "\n";
         }
         ofile.close();
@@ -111,19 +114,18 @@ void simulate_single_particle_with_different_steps()
 
 void simulate_100_particles()
 {
-    PenningTrap trap((unsigned)100, T,
-                     [](double t) {
-                         return 25. * V / 1000. * (1. + .4 * std::cos(1.5 * t));
-                     },
-                     500., 0);
+    PenningTrap trap((unsigned)100);
 
-    double time = 500.; // microseconds
+    double time = 50.; // microseconds
 
-    trap.write_simulation_to_dir("output/simulate_100_particles", time, N * 4);
+    trap.write_simulation_to_dir("output/simulate_100_particles", time, N, "rk4", false);
 }
 
+// Wide sweep
 void simulate_100_particles_with_time_potential()
 {
+    double time = 500.;
+
     double amplitudes[]{.1, .4, .7};
 
     double freq_start = .2;
@@ -156,7 +158,7 @@ void simulate_100_particles_with_time_potential()
                     amplitudes[i], res[0][j], std::placeholders::_1),
                 500., 0.);
             res[i + 1][j] =
-                trap.fraction_of_particles_left(500., 40000, "rk4", false);
+                trap.fraction_of_particles_left(time, N, "rk4", false);
         }
     }
 
@@ -170,21 +172,23 @@ void simulate_100_particles_with_time_potential()
 
 int main()
 {
-    double start = omp_get_wtime();
+    double t0 = omp_get_wtime();
 
-    simulate_single_particle();
+    // simulate_single_particle();
 
-    simulate_two_particles();
+    // simulate_two_particles();
 
      simulate_single_particle_with_different_steps();
 
-    // simulate_100_particles();
+    double t1 = omp_get_wtime();
 
-    // simulate_100_particles_with_time_potential();
+    simulate_100_particles();
+
+    //simulate_100_particles_with_time_potential();
 
     double end = omp_get_wtime();
 
-    std::cout << "Time: " << (end - start) << " seconds" << std::endl;
+    std::cout << "Time: " << (end - t1) << " seconds" << std::endl;
 
     return 0;
 }