corrected leapfrog integration

lib/nBodySim.cpp

@@ -30,7 +30,6 @@ nBodySim::nBodySim(std::string datafile) {
     for (unsigned i=0; i < 3; ++i) {
         for (unsigned j=0; j < nParticles_; ++j) {
             file >> velocities_[i+3*j];
-            velocities_[i+3*j] *= 0.0001;
         }
     }
 
@@ -44,14 +43,11 @@ nBodySim::nBodySim(std::string datafile) {
     }
 
     forces_ = new double[3*nParticles_];
-    #pragma omp parallel for
-    for (unsigned i=0; i < 3*nParticles_; ++i) {
-        forces_[i] = 0.0;
-    }
+    updateForces();
     
     file.close();
 
-    softening_ = 0.1;// * calculateMeanInterparticleDistance();
+    softening_ = 0.01;
 
     Vector center{0.0, 0.0, 0.0};
     tree_ = new Tree(mass_, positions_, forces_, softening_, nParticles_, maxCoord, center);
@@ -66,25 +62,31 @@ nBodySim::~nBodySim() {
 }
 
 void nBodySim::runSimulation(double dt, unsigned nSteps) {
+    const unsigned skip = 10;
     std::ofstream file("../out/positions.dat");
-    double avgtime = 0;
+    file << nParticles_ << " " << skip << " " << dt << std::endl;
     for (unsigned i=0; i < nSteps; ++i) {
 
+        write2file(positions_, file, 3, skip);
+
+        // Integrate with leapfrog
         auto start = std::chrono::high_resolution_clock::now();
-        treeUpdateForces();
+        kick(dt/2.0);
+        drift(dt);
+        updateForces();
+        kick(dt/2.0);
         auto stop = std::chrono::high_resolution_clock::now();
         auto duration = std::chrono::duration_cast<std::chrono::milliseconds>(stop - start);
-        avgtime += duration.count();
-        std::cout << std::right << "Step "    << std::setw(5) << i 
+
+        std::cout << std::right << "Step "    << std::setw(5)  << i 
                                 << ": MFM = " << std::setw(12) << calculateMeanForceMagnitude() 
-                                << ", DEV = " << std::setw(12) << verifyForces() 
-                                << ", MID = " << std::setw(8) << calculateMeanInterparticleDistance() 
-                                << ", T = "   << std::setw(8) << duration.count()/1000.0 << " s"
+                                // << ", DEV = " << std::setw(12) << verifyForces() 
+                                << ", MCD = " << std::setw(8)  << calculateMeanCenterDistance()
+                                // << ", MID = " << std::setw(8)  << calculateMeanCenterDistance() 
+                                << ", E = "   << std::setw(12) << calculateTotalEnergy()
+                                << ", T = "   << std::setw(8)  << duration.count()/1000.0 << " s"
                                 << std::endl;
-        updateForces();
-        std::cout << std::right << "    correct MFM = " << std::setw(12) << calculateMeanForceMagnitude() << std::endl;
-        doTimeStep(dt);
-        updateTree();
+        // updateTree();
     }
     file.close();
 }
@@ -139,10 +141,6 @@ void nBodySim::updateForces() {
 }
 
 void nBodySim::treeUpdateForces() {
-    #pragma omp parallel for
-    for (unsigned i=0; i < 3*nParticles_; ++i) {
-        forces_[i] = 0.0;
-    }
     #pragma omp parallel for
     for (unsigned i=0; i < nParticles_; ++i) {
         tree_->updateForce(i);
@@ -187,6 +185,28 @@ double nBodySim::verifyForces() const {
     return avgdev;
 }
 
+void nBodySim::kick(double dt) {
+    #pragma omp parallel for
+    for (unsigned i=0; i < nParticles_; ++i) {
+        for (unsigned j=0; j < 3; ++j) {
+            velocities_[3*i+j] += forces_[3*i+j] * dt / mass_;
+        }
+    }
+}
+
+void nBodySim::drift(double dt) {
+    #pragma omp parallel for
+    for (unsigned i=0; i < nParticles_; ++i) {
+        for (unsigned j=0; j < 3; ++j) {
+            positions_[3*i+j] += velocities_[3*i+j] * dt;
+        }
+    }
+}
+
+void nBodySim::updateTree() {
+    tree_->update();
+}
+
 double nBodySim::calculateMeanInterparticleDistance() const {
     double meanInterparticleDistance = 0.0;
     // loop over each pair of particles and calculate the distance between them
@@ -226,24 +246,52 @@ double nBodySim::calculateMeanForceMagnitude() const {
     return forceMagnitude;
 }
 
-void nBodySim::doTimeStep(double dt) {
-    // update velocities, then positions
-    #pragma omp parallel for
-    for (unsigned i=0; i < nParticles_; ++i) {
-        for (unsigned j=0; j < 3; ++j) {
-            velocities_[3*i+j] += 0.5 * forces_[3*i+j] * dt / mass_;
-            positions_[3*i+j] += velocities_[3*i+j] * dt;
-        }
+double nBodySim::calculateMeanCenterDistance() const {
+    double avgDist = 0.0;
+    #pragma omp parallel for reduction(+:avgDist)
+    for (unsigned i = 0; i < nParticles_; i++) {
+        double x = positions_[3*i];
+        double y = positions_[3*i+1];
+        double z = positions_[3*i+2];
+        double d2 = x*x + y*y + z*z;
+        double d = std::sqrt(d2);
+        avgDist += d;
     }
+    avgDist /= nParticles_;
+    return avgDist;
 }
 
-void nBodySim::updateTree() {
-    tree_->update();
+double nBodySim::calculateTotalEnergy() const {
+    double totalEnergy = 0.0;
+    #pragma omp parallel for reduction(+:totalEnergy)
+    for (unsigned i = 0; i < nParticles_; i++) {
+        double px = positions_[3*i];
+        double py = positions_[3*i+1];
+        double pz = positions_[3*i+2];
+        double vx = velocities_[3*i];
+        double vy = velocities_[3*i+1];
+        double vz = velocities_[3*i+2];
+        double v2 = vx*vx + vy*vy + vz*vz;
+        double kineticEnergy = 0.5 * mass_ * v2;
+        double potentialEnergy = 0.0;
+        for (unsigned j = 0; j < nParticles_; ++j) {
+            if (j == i) continue;
+            // calculate distance between particles
+            double dx = positions_[3*j]   - px;
+            double dy = positions_[3*j+1] - py;
+            double dz = positions_[3*j+2] - pz;
+            double r2 = dx*dx + dy*dy + dz*dz;
+            double r = std::sqrt(r2);
+            potentialEnergy -= mass_*mass_ / r;
+        }
+        totalEnergy += kineticEnergy + potentialEnergy;
+    }
+    return totalEnergy;
 }
 
-void nBodySim::write2file(const double* array, std::ofstream& file, unsigned dim) const {
+void nBodySim::write2file(const double* array, std::ofstream& file, unsigned dim, unsigned skip) const {
     // writes array to file, with N rows and dim columns
-    for (unsigned i = 0; i < nParticles_; i++) {
+    for (unsigned i = 0; i < nParticles_; i+=skip) {
         for (unsigned j = 0; j < dim; j++) {
             file << std::right << std::setw(16) << array[dim*i+j] << " ";
         }

lib/nBodySim.hpp

@@ -21,18 +21,24 @@ public:
     // Loops over all pairs of particles and calculates forces, checks if forces are correct
     double verifyForces() const;
 
+    // Leapfrog: update velocities
+    void kick(double dt);
+    // Leapfrog: update positions
+    void drift(double dt);
+    // Update tree
+    void updateTree();
+
     // Loop over all pairs of particles and calculate mean interparticle distance
     double calculateMeanInterparticleDistance() const;
     // Loop over all force vectors and calculate mean force magnitude
     double calculateMeanForceMagnitude() const;
-
-    // Integrate using leapfrog algorithm. doesn't update forces, doesn't update tree
-    void doTimeStep(double dt);
-    // Update tree
-    void updateTree();
+    // loop over every particle and calculate the mean distance from the center
+    double calculateMeanCenterDistance() const;
+    // Loop over all particles and calculate the total energy
+    double calculateTotalEnergy() const;
 
     // Writes data to file, every row is a particle
-    void write2file(const double* array, std::ofstream& file, unsigned dim) const;
+    void write2file(const double* array, std::ofstream& file, unsigned dim, unsigned skip) const;
     // Write current state of simulation to file
     void saveState2file(unsigned step, std::ofstream& file) const;