forked from ltt1598/--Galaxy
-
Notifications
You must be signed in to change notification settings - Fork 0
/
galaxy_profiling.py
56 lines (47 loc) · 1.59 KB
/
galaxy_profiling.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
import taichi as ti
from celestial_objects import Star, Planet
if __name__ == "__main__":
ti.init(kernel_profiler=True, arch=ti.cuda)
# control
paused = False
export_images = False
# stars and planets
stars = Star(N=2, mass=1000)
stars.initialize(0.5, 0.5, 0.2, 10)
planets = Planet(N=1000, mass=1)
planets.initialize(0.5, 0.5, 0.4, 10)
# GUI
my_gui = ti.GUI("Galaxy", (800, 800))
h = 5e-5 # time-step size
i = 0
stars.computeForce()
planets.computeForce(stars)
for celestial_obj in (stars, planets):
celestial_obj.update(h)
ti.clear_kernel_profile_info()
for i in range(100): # while my_gui.running:
for e in my_gui.get_events(ti.GUI.PRESS):
if e.key == ti.GUI.ESCAPE:
exit()
elif e.key == ti.GUI.SPACE:
paused = not paused
print("paused =", paused)
elif e.key == 'r':
stars.initialize(0.5, 0.5, 0.2, 10)
planets.initialize(0.5, 0.5, 0.4, 10)
i = 0
elif e.key == 'i':
export_images = not export_images
if not paused:
stars.computeForce()
planets.computeForce(stars)
for celestial_obj in (stars, planets):
celestial_obj.update(h)
i += 1
stars.display(my_gui, radius=10, color=0xffd500)
planets.display(my_gui)
if export_images:
my_gui.show(f"images\output_{i:05}.png")
else:
my_gui.show()
ti.print_kernel_profile_info('count')