Orbit GUI.py

import numpy as np
from matplotlib.figure import Figure
from PyQt5 import QtCore, QtWidgets, QtGui, Qt
from matplotlib.backends.backend_qt5agg import FigureCanvasQTAgg as FigureCanvas
from matplotlib.backends.backend_qt5agg import NavigationToolbar2QT as NavigationToolbar
import sys


class ApplicationWindow(QtWidgets.QMainWindow):
    def __init__(self, universe, plotting_objects):
        super().__init__()

        self._plotting_objects = plotting_objects
        self._universe = universe
        self._main = QtWidgets.QWidget()
        self.setCentralWidget(self._main)
        layout = QtWidgets.QVBoxLayout(self._main)
        canvas = FigureCanvas(Figure(figsize=(6, 6)))
        layout.addWidget(canvas)
        self.setLayout(layout)
        self.addToolBar(QtCore.Qt.TopToolBarArea, NavigationToolbar(canvas, self))
        self._ax = canvas.figure.subplots()
        self._timer = canvas.new_timer(16.6, [(self._update_canvas, (), {})])
        self._timer.start()
        canvas.show()

    def _update_canvas(self):
        # print(self._plotting_objects[0].velocity, self._plotting_objects[1].velocity)

        # dP = self._plotting_objects[0].position - self._plotting_objects[1].position

        # dR = (dP**2).sum()**.5

        # print(dR)

        for _ in range(1000):
            self._universe.update_objects()

        self._ax.clear()
        self._ax.set_aspect("equal")
        self._ax.set_xlim([-1.2 * universe.size, 1.2 * universe.size])
        self._ax.set_ylim([-1.2 * universe.size, 1.2 * universe.size])
        for p_object in self._plotting_objects:
            self._ax.plot(p_object.position[0], p_object.position[1], "bo")
        self._ax.figure.canvas.draw()
        # input("Press Enter to continue...")

    def closeEvent(self, a0: QtGui.QCloseEvent):
        super().closeEvent(a0)
        self._timer.stop()
        QtWidgets.QApplication.quit()


class PhysicalObject:

    def __init__(self, mass: float, initial_position: np.array, initial_velocity: np.array, fixed=False):
        self._mass = mass
        self._position = initial_position.astype(float) * 1
        self._velocity = initial_velocity.astype(float) * 1
        self._fixed = fixed
        self._time = 0

    @property
    def position(self):
        return self._position

    @property
    def mass(self):
        return self._mass

    @property
    def velocity(self):
        return self._velocity

    def calculate_distance(self, other):
        dx = other.position[0] - self._position[0]
        dy = other.position[1] - self._position[1]

        dx_squared = np.square(dx)
        dy_squared = np.square(dy)

        distance = np.sqrt(dx_squared + dy_squared)

        return distance

    def apply_force(self, force: np.array, time_step: float):
        if not self._fixed:
            acceleration = force / self._mass

            # print(__name__, acceleration)

            velocity_step = acceleration * time_step
            position_step = self._velocity * time_step

            self._velocity += velocity_step
            self._position += position_step

        self._time += time_step


class Universe:
    def __init__(self, objects: [PhysicalObject], time_step: float, size: float, gravitational_constant=6.6743015e-11):
        self._objects = objects
        self._time_step = time_step
        self._time = 0
        self._size = size
        self._gravitational_constant = gravitational_constant

    @property
    def time(self):
        return self._time

    @property
    def size(self):
        return self._size

    def center_of_mass(self):
        weighted_positions = np.array([0, 0], dtype=float)
        total_mass = 0

        for p_object in self._objects:
            weighted_positions += p_object.position * p_object.mass
            total_mass += p_object.mass

        return weighted_positions / total_mass, total_mass

    def update_objects(self):
        # cm, cm_mass = self.__center_of_mass()

        for p_object in self._objects:
            total_force = np.array([0, 0], dtype=float)

            for o_p_object in self._objects:
                if o_p_object == p_object:
                    continue

                r_vector = p_object.position - o_p_object.position
                r_hat = r_vector / np.linalg.norm(r_vector)

                force_magnitude = self._gravitational_constant * o_p_object.mass * p_object.mass / r_vector.dot(r_vector)
                force = -force_magnitude * r_hat
                total_force += force

            p_object.apply_force(total_force, self._time_step)

        self._time += self._time_step


if __name__ == "__main__":
    x_positions = []
    y_positions = []

    gravitational_constant = 6.6743015e-11
    mass_earth = 5.972e24
    mass_moon = 7.348e22

    earth_moon_radius = 384.4e6

    initial_y_velocity = np.sqrt(gravitational_constant * mass_earth / earth_moon_radius)

    earth = PhysicalObject(mass_earth, np.array([0, 0]), np.array([0, 0]), fixed=False)
    moon = PhysicalObject(mass_moon, np.array([earth_moon_radius, 0]), np.array([0, initial_y_velocity]))

    universe = Universe([earth, moon], 1, 1.2 * earth_moon_radius)

    qapp = QtWidgets.QApplication(sys.argv)
    app = ApplicationWindow(universe, [earth, moon])
    app.show()
    qapp.exec_()