Task: N-body problem


C++: #include <algorithm>
#include <iomanip>
#include <iostream>
#include <fstream>
#include <string>
#include <vector>

class Vector {
private:
    double px, py, pz;

public:
    Vector() : px(0.0), py(0.0), pz(0.0) {
        
    }

    Vector(double x, double y, double z) : px(x), py(y), pz(z) {
        
    }

    double mod() const {
        return sqrt(px*px + py * py + pz * pz);
    }

    Vector operator+(const Vector& rhs) const {
        return Vector(px + rhs.px, py + rhs.py, pz + rhs.pz);
    }

    Vector operator-(const Vector& rhs) const {
        return Vector(px - rhs.px, py - rhs.py, pz - rhs.pz);
    }

    Vector operator*(double s) const {
        return Vector(px*s, py*s, pz*s);
    }

    bool operator==(const Vector& rhs) const {
        return px == rhs.px
            && py == rhs.py
            && pz == rhs.pz;
    }

    friend std::istream& operator>>(std::istream&, Vector&);
    friend std::ostream& operator<<(std::ostream&, Vector&);
};

std::istream& operator>>(std::istream& in, Vector& v) {
    return in >> v.px >> v.py >> v.pz;
}

std::ostream& operator<<(std::ostream& out, Vector& v) {
    auto precision = out.precision();
    auto width = out.width();
    out << std::fixed << std::setw(width) << std::setprecision(precision) << v.px << "  ";
    out << std::fixed << std::setw(width) << std::setprecision(precision) << v.py << "  ";
    out << std::fixed << std::setw(width) << std::setprecision(precision) << v.pz;
    return out;
}

const Vector ORIGIN{ 0.0, 0.0, 0.0 };

class NBody {
private:
    double gc;
    int bodies;
    int timeSteps;
    std::vector<double> masses;
    std::vector<Vector> positions;
    std::vector<Vector> velocities;
    std::vector<Vector> accelerations;

    void resolveCollisions() {
        for (int i = 0; i < bodies; ++i) {
            for (int j = i + 1; j < bodies; ++j) {
                if (positions[i] == positions[j]) {
                    std::swap(velocities[i], velocities[j]);
                }
            }
        }
    }

    void computeAccelerations() {
        for (int i = 0; i < bodies; ++i) {
            accelerations[i] = ORIGIN;
            for (int j = 0; j < bodies; ++j) {
                if (i != j) {
                    double temp = gc * masses[j] / pow((positions[i] - positions[j]).mod(), 3);
                    accelerations[i] = accelerations[i] + (positions[j] - positions[i]) * temp;
                }
            }
        }
    }

    void computeVelocities() {
        for (int i = 0; i < bodies; ++i) {
            velocities[i] = velocities[i] + accelerations[i];
        }
    }

    void computePositions() {
        for (int i = 0; i < bodies; ++i) {
            positions[i] = positions[i] + velocities[i] + accelerations[i] * 0.5;
        }
    }

public:
    NBody(std::string& fileName) {
        using namespace std;

        ifstream ifs(fileName);
        if (!ifs.is_open()) {
            throw runtime_error("Could not open file.");
        }

        ifs >> gc >> bodies >> timeSteps;

        masses.resize(bodies);
        positions.resize(bodies);
        fill(positions.begin(), positions.end(), ORIGIN);
        velocities.resize(bodies);
        fill(velocities.begin(), velocities.end(), ORIGIN);
        accelerations.resize(bodies);
        fill(accelerations.begin(), accelerations.end(), ORIGIN);

        for (int i = 0; i < bodies; ++i) {
            ifs >> masses[i] >> positions[i] >> velocities[i];
        }

        cout << "Contents of " << fileName << '\n';
        cout << gc << ' ' << bodies << ' ' << timeSteps << '\n';
        for (int i = 0; i < bodies; ++i) {
            cout << masses[i] << '\n';
            cout << positions[i] << '\n';
            cout << velocities[i] << '\n';
        }
        cout << "\nBody   :      x          y          z    |     vx         vy         vz\n";
    }

    int getTimeSteps() {
        return timeSteps;
    }

    void simulate() {
        computeAccelerations();
        computePositions();
        computeVelocities();
        resolveCollisions();
    }

    friend std::ostream& operator<<(std::ostream&, NBody&);
};

std::ostream& operator<<(std::ostream& out, NBody& nb) {
    for (int i = 0; i < nb.bodies; ++i) {
        out << "Body " << i + 1 << " : ";
        out << std::setprecision(6) << std::setw(9) << nb.positions[i];
        out << " | ";
        out << std::setprecision(6) << std::setw(9) << nb.velocities[i];
        out << '\n';
    }
    return out;
}

int main() {
    std::string fileName = "nbody.txt";
    NBody nb(fileName);

    for (int i = 0; i < nb.getTimeSteps(); ++i) {
        std::cout << "\nCycle " << i + 1 << '\n';
        nb.simulate();
        std::cout << nb;
    }

    return 0;
}






Python: import math

class Vector:
    def __init__(self, x, y, z):
        self.x = x
        self.y = y
        self.z = z

    def __add__(self, other):
        return Vector(self.x + other.x, self.y + other.y, self.z + other.z)

    def __sub__(self, other):
        return Vector(self.x - other.x, self.y - other.y, self.z - other.z)

    def __mul__(self, other):
        return Vector(self.x * other, self.y * other, self.z * other)

    def __div__(self, other):
        return Vector(self.x / other, self.y / other, self.z / other)

    def __eq__(self, other):
        if isinstance(other, Vector):
            return self.x == other.x and self.y == other.y and self.z == other.z
        return False

    def __ne__(self, other):
        return not self.__eq__(other)

    def __str__(self):
        return '({x}, {y}, {z})'.format(x=self.x, y=self.y, z=self.z)

    def abs(self):
        return math.sqrt(self.x*self.x + self.y*self.y + self.z*self.z)

origin = Vector(0, 0, 0)

class NBody:
    def __init__(self, fileName):
        with open(fileName, "r") as fh:
            lines = fh.readlines()
            gbt = lines[0].split()
            self.gc = float(gbt[0])
            self.bodies = int(gbt[1])
            self.timeSteps = int(gbt[2])
            self.masses = [0.0 for i in range(self.bodies)]
            self.positions = [origin for i in range(self.bodies)]
            self.velocities = [origin for i in range(self.bodies)]
            self.accelerations = [origin for i in range(self.bodies)]
            for i in range(self.bodies):
                self.masses[i] = float(lines[i*3 + 1])
                self.positions[i] = self.__decompose(lines[i*3 + 2])
                self.velocities[i] = self.__decompose(lines[i*3 + 3])

            print "Contents of", fileName
            for line in lines:
                print line.rstrip()
            print
            print "Body   :      x          y          z    |",
            print "     vx         vy         vz"

    def __decompose(self, line):
        xyz = line.split()
        x = float(xyz[0])
        y = float(xyz[1])
        z = float(xyz[2])
        return Vector(x, y, z)

    def __computeAccelerations(self):
        for i in xrange(self.bodies):
            self.accelerations[i] = origin
            for j in xrange(self.bodies):
                if i != j:
                    temp = self.gc * self.masses[j] / math.pow((self.positions[i] - self.positions[j]).abs(), 3)
                    self.accelerations[i] += (self.positions[j] - self.positions[i]) * temp
        return None

    def __computePositions(self):
        for i in xrange(self.bodies):
            self.positions[i] += self.velocities[i] + self.accelerations[i] * 0.5
        return None

    def __computeVelocities(self):
        for i in xrange(self.bodies):
            self.velocities[i] += self.accelerations[i]
        return None

    def __resolveCollisions(self):
        for i in xrange(self.bodies):
            for j in xrange(self.bodies):
                if self.positions[i] == self.positions[j]:
                    (self.velocities[i], self.velocities[j]) = (self.velocities[j], self.velocities[i])
        return None

    def simulate(self):
        self.__computeAccelerations()
        self.__computePositions()
        self.__computeVelocities()
        self.__resolveCollisions()
        return None

    def printResults(self):
        fmt = "Body %d : % 8.6f  % 8.6f  % 8.6f | % 8.6f  % 8.6f  % 8.6f"
        for i in xrange(self.bodies):
            print fmt % (i+1, self.positions[i].x, self.positions[i].y, self.positions[i].z, self.velocities[i].x, self.velocities[i].y, self.velocities[i].z)
        return None

nb = NBody("nbody.txt")
for i in xrange(nb.timeSteps):
    print "\nCycle %d" % (i + 1)
    nb.simulate()
    nb.printResults()