/
universe.py
83 lines (67 loc) · 2.31 KB
/
universe.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
#-----------------------------------------------------------------------
# universe.py
# python universe.py 20000 < 2bodyTiny.txt
# python universe.py 20000 < 2body.txt
# python universe.py 20000 < 3body.txt
# python universe.py 20000 < 4body.txt
#-----------------------------------------------------------------------
import stdio
import vector
import sys
import stddraw
import body
#-----------------------------------------------------------------------
class Universe:
# read universe from standard input
def __init__(self):
self._N = stdio.readInt() # Number of bodies
self._radius = stdio.readFloat() # Radius of universe
# the set scale for drawing on screen
stddraw.setXscale(-self._radius, +self._radius)
stddraw.setYscale(-self._radius, +self._radius)
# read in the _N bodies
self.orbs = [] # Array of bodies
for i in range(self._N):
rx = stdio.readFloat()
ry = stdio.readFloat()
vx = stdio.readFloat()
vy = stdio.readFloat()
mass = stdio.readFloat()
position = [ rx, ry ]
velocity = [ vx, vy ]
r = vector.Vector(position)
v = vector.Vector(velocity)
self.orbs += [body.Body(r, v, mass)]
# increment time by dt units, assume forces are constant
# in given interval
def increaseTime(self, dt):
# initialize the forces to zero
f = []
for i in range(self._N):
f += [vector.Vector([0.0, 0.0])]
# compute the forces
for i in range(self._N):
for j in range(self._N):
if i != j:
f[i] = f[i] + self.orbs[j].forceTo(self.orbs[i])
# move the bodies
for i in range(self._N):
self.orbs[i].move(f[i], dt)
# draw the _N bodies
def draw(self):
for i in range(self._N):
self.orbs[i].draw()
#-----------------------------------------------------------------------
# Simulate a universe
def main(argv):
stddraw.createWindow()
newton = Universe()
dt = float(argv[1])
while True:
newton.increaseTime(dt)
newton.draw()
stddraw.sleep(10)
stddraw.show()
stddraw.clear()
if __name__ == '__main__':
main(sys.argv)