euler.py

import sys
import time
import math
import numpy as np


G = 1.0


class Body:
    def __init__(self, r, v, m):
        self.r = r
        self.v = v
        self.m = m


    def __repr__(self):
        return f"{self.r}:{self.v}:{self.m}"


def load(fn):
    bodies = []
    with open(fn, "rb") as f:
        G = float(f.readline())
        nbodies = int(f.readline())
        for i in range(nbodies):
            v = f.readline().split()
            rad = np.array([
                float(v[1]),
                float(v[2]),
                float(v[3])
            ], dtype=float)
            vel = np.array([
                float(v[4]),
                float(v[5]),
                float(v[6])
            ], dtype=float)
            m = float(v[7])
            b = Body(rad, vel, m)
            bodies.append(b)
    return bodies


class Euler:
    def __init__(self, bodies, dt):
        self.bodies = bodies
        self.positions = np.array([b.r for b in self.bodies])
        self.velocities = np.array([b.v for b in self.bodies])
        self.masses = np.array([b.m for b in self.bodies])
        self.dt = dt

    def update(self):
        # Расчет разностей координат для каждой пары тел
        r_ij = self.positions[np.newaxis, :, :] - self.positions[:, np.newaxis, :]
        R_ij = np.linalg.norm(r_ij, axis=2)  # Евклидово расстояние между парами тел


        np.fill_diagonal(R_ij, np.inf)

        # eps = 0.1
        # R_ij[R_ij < eps] = eps  # Замена слишком малых расстояний на eps
        R_ij = R_ij[:, :, np.newaxis]

        # Вычисление сил
        F_ij = G * self.masses[:, np.newaxis, np.newaxis] * self.masses[np.newaxis, :, np.newaxis] * r_ij / (R_ij**3)

        # Сумма всех сил, действующих на каждое тело
        F_total = np.sum(F_ij, axis=1)  # Сумма по второй оси

        # Обновление скоростей и позиций
        accelerations = F_total / self.masses[:, np.newaxis]  # Ускорения
        self.velocities += self.dt * accelerations  # Обновление скоростей
        self.positions += self.dt * self.velocities  # Обновление позиций


def kepler(dt):
    bodies = []
    G = 1
    MM = 1e5
    bodies.append(
        Body(
            np.array([0.0, 0.0, 0.0, 0.0], dtype=float),
            np.array([0.0, 0.0, 0.0, 0.0], dtype=float),
            MM
        )
    )
    bodies.append(
        Body(
            np.array([0.0, 1.0, 0.0, 0.0], dtype=float),
            np.array([math.sqrt(G * MM), 0.0, 0.0, 0.0], dtype=float),
            1
        )
    )

    max_err = 0
    euler = Euler(bodies, dt)
    T = 0.1
    t = 0

    while t < T:
        euler.update()
        r = np.linalg.norm(euler.positions[1])
        err = math.fabs(r - 1.0)

        if max_err < err:
            max_err = err

        t += dt

    return max_err


def test():
    print("Test")
    err1 = kepler(0.001)
    err2 = kepler(0.0001)
    err3 = kepler(0.00001)
    print(f"{err1} {err2} {err3}")
    print(f"{err1 / err2} {err1 / err3}")
    print(err1 < 0.3)
    print(err2 < 0.02)
    print(err3 < 0.002)


def bench(fn):
    dt = 0.0001
    bodies = load(fn)
    euler = Euler(bodies, dt)

    t = 0.0

    t0 = time.time()
    while t < 0.01: # TODO: incr t to  0.1
        euler.update()
        t += dt
    t1 = time.time()
    print(f"Elapsed: {t1-t0}")


def main():
    test_mode = False
    bench_mode = False
    fn = ""

    i = 1
    nargs = len(sys.argv)

    while i < nargs:
        if i < nargs - 1 and sys.argv[i] == "--input":
            i = i + 1
            fn = sys.argv[i]
        elif sys.argv[i] == "--test":
            test_mode = True
        elif sys.argv[i] == "--bench":
            bench_mode = True

        i = i + 1


    if test_mode:
        test()
        return

    if bench_mode:
        bench(fn)
        return

if __name__ == "__main__":
    main()