nodes can access masses, copy constructor done

lib/nBodySim.cpp

@@ -49,7 +49,7 @@ nBodySim::nBodySim(std::string datafile) {
     
     file.close();
 
-    tree_ = new Tree(positions_, forces_, nParticles_, 4*maxCoord, new double[3]{0, 0, 0});
+    tree_ = new Tree(masses_, positions_, forces_, nParticles_, 4*maxCoord, new double[3]{0, 0, 0});
 }
 
 nBodySim::~nBodySim() {
@@ -112,7 +112,9 @@ void nBodySim::calculateForces() {
 }
 
 void nBodySim::treeCalculateForces() {
-    // TODO
+    for (unsigned i=0; i < nParticles_; ++i) {
+        tree_->calculateForce(i);
+    }
 }
 
 double nBodySim::calculateMeanInterparticleDistance() {

lib/node.cpp

 #include "node.hpp"
 
+#include <cmath>
+
 /* 
 Node::Node() {
     root_ = nullptr;
@@ -15,15 +17,24 @@ Node::Node() {
 }
  */
 
-Node::Node(Node* parent, double* particles, unsigned nParticles, unsigned depth, double size, double center[3], unsigned* localParticles, unsigned nLocalParticles) {
+Node::Node(Node* parent, double* masses, double* particles, unsigned nParticles, unsigned depth, double size, double center[3], unsigned* localParticles, unsigned nLocalParticles) {
     parent_ = parent;
     particles_ = particles;
     nParticles_ = nParticles;
+
     depth_ = depth;
     size_ = size;
+
     for (unsigned i=0; i<3; ++i) center_[i] = center[i];
+
     localParticles_ = localParticles;
     nLocalParticles_ = nLocalParticles;
+
+    mass_ = 0;
+    for (unsigned i=0; i<nLocalParticles; ++i) {
+        mass_ += masses_[localParticles_[i]];
+    }
+
     // allocate children
     // TODO
 }
@@ -33,4 +44,35 @@ Node::~Node() {
         delete children_[i];
     }
     delete[] localParticles_;
+}
+
+Node& Node::operator=(const Node& other) {
+    root_ = other.root_;
+    parent_ = other.parent_;
+    particles_ = other.particles_;
+    nParticles_ = other.nParticles_;
+    depth_ = other.depth_;
+    size_ = other.size_;
+    for (unsigned i=0; i<3; ++i) center_[i] = other.center_[i];
+    localParticles_ = other.localParticles_;
+    nLocalParticles_ = other.nLocalParticles_;
+    mass_ = other.mass_;
+    for (unsigned i=0; i<8; ++i) children_[i] = other.children_[i];
+    return *this;
+}
+
+double* Node::calculateForce(unsigned particle) {
+    // calculate theta
+    double lambda = std::sqrt((center_[0] - particles_[3*particle + 0])*(center_[0] - particles_[3*particle + 0])
+                            + (center_[1] - particles_[3*particle + 1])*(center_[1] - particles_[3*particle + 1])
+                            + (center_[2] - particles_[3*particle + 2])*(center_[2] - particles_[3*particle + 2]));
+    double theta = size_/lambda;
+    // if angle too large, calculate force on particle from children
+    // else calculate force on particle from particles in this node
+    if (theta > theta0_) {
+        // TODO
+    } else {
+        // TODO
+    }
+    return nullptr;
 }
\ No newline at end of file

lib/node.hpp

@@ -7,12 +7,20 @@ public:
     Node() = default;
     // whoever creates the node is responsible for deciding which particles are in which octant and for allocating the localParticles array 
     // node will write to localParticles array
-    Node(Node* parent, double* particles, unsigned nParticles, unsigned depth, double size, double center[3], unsigned* localParticles, unsigned nLocalParticles);
+    Node(Node* parent, double* masses, double* particles, unsigned nParticles, unsigned depth, double size, double center[3], unsigned* localParticles, unsigned nLocalParticles);
+    // destuctor
     ~Node();
+    // copy constructor
+    Node(const Node&) = default;
+    // copy assignment
+    Node& operator=(const Node&);
+    // calculate force on a particle recursively
+    double* calculateForce(unsigned particle);
 private:
     Node* root_;
     Node* parent_;
     Node* children_[8];
+    double* masses_;
     double* particles_;
     unsigned nParticles_;
     unsigned depth_;
@@ -20,6 +28,8 @@ private:
     double center_[3];
     unsigned* localParticles_;
     unsigned nLocalParticles_;
+    double mass_;
+    const double theta0_ = 0.4;
 };
 
 #endif /* NODE_HPP */
\ No newline at end of file

lib/tree.cpp

 #include "tree.hpp"
 #include "node.hpp"
 
-Tree::Tree(double* particles, double* forces, unsigned nParticles, double size, double* center) {
+Tree::Tree(double* masses, double* particles, double* forces, unsigned nParticles, double size, double* center) {
     particles_ = particles;
     forces_ = forces;
     nParticles_ = nParticles;
@@ -13,13 +13,18 @@ Tree::Tree(double* particles, double* forces, unsigned nParticles, double size,
     for (unsigned i=0; i < nParticles_; ++i) localParticles_[i] = i;
     
     root_ = new Node();
-    *root_ = Node(root_, particles_, nParticles_, 0, size_, center_, localParticles_, nParticles_);
+    *root_ = Node(root_, masses_, particles_, nParticles_, 0, size_, center_, localParticles_, nParticles_);
 }
 
 Tree::~Tree() {
     delete root_;
 }
 
+void Tree::calculateForce(unsigned particle) {
+    double* force = root_->calculateForce(particle);
+    for (unsigned i=0; i < 3; ++i) forces_[3*particle + i] = force[i];
+}
+
 void Tree::update() {
     // TODO
 }

lib/tree.hpp

@@ -7,9 +7,11 @@ class Tree {
 public:
     // constructor
     Tree() = delete;
-    Tree(double* particles, double* forces_, unsigned nParticles, double size, double* center);
+    Tree(double* masses, double* particles, double* forces_, unsigned nParticles, double size, double* center);
     // destructor
     ~Tree();
+    // calculate force on a particle
+    void calculateForce(unsigned particle);
     // 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 positions
@@ -18,6 +20,7 @@ public:
     void kick(double dt);
 private:
     Node* root_;
+    double* masses_;
     double* particles_;
     double* forces_;
     unsigned nParticles_;