import sys
import math
import pygame
import random
from collections import defaultdict
# Ablak mérete
WIDTH, HEIGHT = 900, 600
WIDTHD2, HEIGHTD2 = WIDTH / 2., HEIGHT / 2.
# Bolygók száma
PLANETS = 30
# Sűrűség, tömeg számításához
DENSITY = 0.001
# Gravitáció erőssége
GRAVITYSTRENGTH = 1.e4
# Globális lista
g_listOfPlanets = []
class State:
def __init__(self, x, y, vx, vy):
self._x, self._y, self._vx, self._vy = x, y, vx, vy
def __repr__(self):
return f'x:{self._x} y:{self._y} vx:{self._vx} vy:{self._vy}'
class Derivative:
def __init__(self, dx, dy, dvx, dvy):
self._dx, self._dy, self._dvx, self._dvy = dx, dy, dvx, dvy
def __repr__(self):
return f'dx:{self._dx} dy:{self._dy} dvx:{self._dvx} dvy:{self._dvy}'
class Planet:
def __init__(self):
if PLANETS == 1:
self._st = State(150, 300, 0, 2)
else:
self._st = State(
float(random.randint(0, WIDTH)),
float(random.randint(0, HEIGHT)),
float(random.randint(0, 300) / 100.) - 1.5,
float(random.randint(0, 300) / 100.) - 1.5
)
self._r = 1.5
self.setMassFromRadius()
self._merged = False
def __repr__(self):
return repr(self._st)
def acceleration(self, state, unused_t):
ax = ay = 0.0
for p in g_listOfPlanets:
if p is self or p._merged:
continue
dx = p._st._x - state._x
dy = p._st._y - state._y
dsq = dx * dx + dy * dy
dr = math.sqrt(dsq)
force = GRAVITYSTRENGTH * self._m * p._m / dsq if dsq > 1e-10 else 0.
ax += force * dx / dr
ay += force * dy / dr
return ax, ay
def initialDerivative(self, state, t):
ax, ay = self.acceleration(state, t)
return Derivative(state._vx, state._vy, ax, ay)
def nextDerivative(self, initialState, derivative, t, dt):
state = State(
initialState._x + derivative._dx * dt,
initialState._y + derivative._dy * dt,
initialState._vx + derivative._dvx * dt,
initialState._vy + derivative._dvy * dt
)
ax, ay = self.acceleration(state, t + dt)
return Derivative(state._vx, state._vy, ax, ay)
def updatePlanet(self, t, dt):
a = self.initialDerivative(self._st, t)
b = self.nextDerivative(self._st, a, t, dt * 0.5)
c = self.nextDerivative(self._st, b, t, dt * 0.5)
d = self.nextDerivative(self._st, c, t, dt)
dxdt = (a._dx + 2 * (b._dx + c._dx) + d._dx) / 6.0
dydt = (a._dy + 2 * (b._dy + c._dy) + d._dy) / 6.0
dvxdt = (a._dvx + 2 * (b._dvx + c._dvx) + d._dvx) / 6.0
dvydt = (a._dvy + 2 * (b._dvy + c._dvy) + d._dvy) / 6.0
self._st._x += dxdt * dt
self._st._y += dydt * dt
self._st._vx += dvxdt * dt
self._st._vy += dvydt * dt
def setMassFromRadius(self):
self._m = DENSITY * 4. * math.pi * (self._r ** 3.) / 3.
def setRadiusFromMass(self):
self._r = (3. * self._m / (DENSITY * 4. * math.pi)) ** (0.3333)
def main():
pygame.init()
win = pygame.display.set_mode((WIDTH, HEIGHT))
pygame.display.set_caption('Gravity simulation (SPACE: show/hide, +/-: zoom)')
keysPressed = defaultdict(bool)
def ScanKeyboard():
for evt in pygame.event.get():
if evt.type in [pygame.KEYDOWN, pygame.KEYUP]:
keysPressed[evt.key] = evt.type == pygame.KEYDOWN
elif evt.type == pygame.QUIT:
pygame.quit()
sys.exit()
global g_listOfPlanets, PLANETS
if len(sys.argv) == 2:
PLANETS = int(sys.argv[1])
g_listOfPlanets = [Planet() for _ in range(PLANETS)]
def planetsTouch(p1, p2):
dx = p1._st._x - p2._st._x
dy = p1._st._y - p2._st._y
return math.hypot(dx, dy) <= (p1._r + p2._r)
# Nap-szerű objektum
sun = Planet()
sun._st._x, sun._st._y = WIDTHD2, HEIGHTD2
sun._st._vx = sun._st._vy = 0.
sun._m *= 1000
sun.setRadiusFromMass()
g_listOfPlanets.append(sun)
for p in g_listOfPlanets:
if p is not sun and planetsTouch(p, sun):
p._merged = True
zoom = 1.0
t, dt = 0., 1.
bClearScreen = True
while True:
t += dt
ScanKeyboard()
if bClearScreen:
win.fill((0, 0, 0))
for p in g_listOfPlanets:
if not p._merged:
pygame.draw.circle(
win, (255, 255, 255),
(int(WIDTHD2 + zoom * (p._st._x - WIDTHD2)),
int(HEIGHTD2 + zoom * (p._st._y - HEIGHTD2))),
int(p._r * zoom)
)
pygame.display.flip()
for p in g_listOfPlanets:
if not p._merged and p is not sun:
p.updatePlanet(t, dt)
# Merge logic
for i, p1 in enumerate(g_listOfPlanets):
if p1._merged:
continue
for j, p2 in enumerate(g_listOfPlanets):
if i >= j or p2._merged:
continue
if planetsTouch(p1, p2):
if p1._m < p2._m:
p1, p2 = p2, p1
p2._merged = True
if p1 is sun:
continue
newvx = (p1._st._vx * p1._m + p2._st._vx * p2._m) / (p1._m + p2._m)
newvy = (p1._st._vy * p1._m + p2._st._vy * p2._m) / (p1._m + p2._m)
p1._m += p2._m
p1.setRadiusFromMass()
p1._st._vx, p1._st._vy = newvx, newvy
# Zoom
if keysPressed[pygame.K_KP_PLUS]:
zoom /= 0.99
if keysPressed[pygame.K_KP_MINUS]:
zoom /= 1.01
if keysPressed[pygame.K_ESCAPE]:
pygame.quit()
sys.exit()
if keysPressed[pygame.K_SPACE]:
bClearScreen = not bClearScreen
if __name__ == "__main__":
main()
--------------------
Nincsenek megjegyzések:
Megjegyzés küldése