diff --git a/lib/nBodySim.cpp b/lib/nBodySim.cpp
index a92c3f11f98fe15f33b74c247fffcb9a136ef286..7ed7323932773ca233de9a0230d8d98cb1494c88 100644
--- a/lib/nBodySim.cpp
+++ b/lib/nBodySim.cpp
@@ -62,7 +62,6 @@ nBodySim::~nBodySim() {
 
 void nBodySim::runSimulation(double dt, unsigned nSteps) {
     for (unsigned i=0; i < nSteps; ++i) {
-        calculateForces();
         doTimeStep(dt);
     }
 }
diff --git a/lib/tree.hpp b/lib/tree.hpp
index 9578237d61370acd0d2d12c1d2fae069bba66430..38cca7a79d8dc9f5f0818f86db781137299b5c49 100644
--- a/lib/tree.hpp
+++ b/lib/tree.hpp
@@ -12,9 +12,9 @@ public:
     ~Tree();
     // update tree: visit each node, check if it needs to be split or merged (for example if a particle has left the region)
     void update();
-    // drift tree: visit each node, update center of mass and center of mass velocity
+    // drift tree: visit each node, update positions
     void drift(double dt);
-    // kick tree: visit each node, update center of mass velocity
+    // kick tree: visit each node, update velocities
     void kick(double dt);
 private:
     Node* root_;
diff --git a/main.cpp b/main.cpp
index d723a1fb08966befaf135b1af618003513347c3a..5224b2965021278f4c274d5ad3f8d7fbdc1621a1 100644
--- a/main.cpp
+++ b/main.cpp
@@ -26,36 +26,42 @@ int main(int argc, char** argv) {
     std::cout << "N = " << N << std::endl;
 
     /*
+    // print positions
     for (unsigned i=0; i<N; ++i) {
         for (unsigned j=0; j<3; ++j) {
             std::cout << sim.getPositions()[3*i+j] << " ";
         }
         std::cout << std::endl;
     }
-    */
 
-    try {
-        // calculate mean interparticle distance
-        double meanInterparticleDistance = sim.calculateMeanInterparticleDistance();
-        std::cout << "Mean interparticle distance: " << meanInterparticleDistance << std::endl;
+    // calculate mean interparticle distance
+    double meanInterparticleDistance = sim.calculateMeanInterparticleDistance();
+    std::cout << "Mean interparticle distance: " << meanInterparticleDistance << std::endl;
+
+    // calculate forces and measure computation time
+    auto start = std::chrono::high_resolution_clock::now();
+    sim.calculateForces();
+    auto stop = std::chrono::high_resolution_clock::now();
 
-        // calculate forces and measure computation time
-        auto start = std::chrono::high_resolution_clock::now();
-        sim.calculateForces();
-        auto stop = std::chrono::high_resolution_clock::now();
+    auto duration = std::chrono::duration_cast<std::chrono::microseconds>(stop - start);
+    std::cout << "Computation Time: " << duration.count()/1000000.0 << " seconds" << std::endl;
 
-        auto duration = std::chrono::duration_cast<std::chrono::microseconds>(stop - start);
-        std::cout << "Computation Time: " << duration.count()/1000000.0 << " seconds" << std::endl;
+    // calculate mean force magnitude
+    double meanForceMagnitude = sim.calculateMeanForceMagnitude();
+    std::cout << "Mean force magnitude: " << meanForceMagnitude << std::endl;
 
-        // calculate mean force magnitude
-        double meanForceMagnitude = sim.calculateMeanForceMagnitude();
-        std::cout << "Mean force magnitude: " << meanForceMagnitude << std::endl;
+    sim.write2file(sim.getForces(), "../out/forces.txt", 3);
+    */
 
-        sim.write2file(sim.getForces(), "../out/forces.txt", 3);
+    try {
+        sim.runSimulation(dt, nSteps);
     } catch (std::exception& e) {
         std::cerr << e.what() << std::endl;
         return 1;
     }
 
+    // save state of simulation
+    sim.saveState2file(nSteps, "../out/state.txt");
+
     return 0;
 }
\ No newline at end of file