python物理引擎模拟三体_三体世界的模拟

1 '''三体问题求解及可视化,去掉了动图模块'''

2 '''Coworker：聂嘉颖，张瀚文'''

3 importnumpy as np4 from numpy importarange5 importmatplotlib.pyplot as plt6 from mpl_toolkits.mplot3d importAxes3D7 from math importsqrt8

9

10 #得到太阳对行星夹角的cot值

11 defsun_height(x0, y0, z0, x1, y1, z1):12 a0 = x1 -x013 b0 = y1 -y014 c0 = z1 -z015 return a0 / sqrt(b0 ** 2 + c0 ** 2)16

17

18 #得到两星体间的距离

19 defdistants(x0, y0, z0, x1, y1, z1):20 return sqrt((x0 - x1) ** 2 + (y0 - y1) ** 2 + (z0 - z1) ** 2)21

22

23 #计算某一瞬时某一星体对行星的辐射强度

24 defsun_heat(d, sun_temperature, x0, y0, z0, x1, y1, z1):25 theta = d/distants(x0, y0, z0, x1, y1, z1)26 earth_temperature = 0.5 * sqrt(theta) *sun_temperature27 #print(earth_temperature)

28 returnearth_temperature29

30 #定义星体质量

31 #star_4选定为行星，其余为恒星

32 star_1_mass = 3e30 #kg

33 star_2_mass = 2e30 #kg

34 star_3_mass = 3e30 #kg

35 star_4_mass = 2e30 #kg

36

37 diameter0 = 1.0392e10 #米

38 diameter1 = 1.0392e10 #米

39 diameter2 = 1.00392e11 #米

40

41 #定义恒星表面温度

42 sun_temperature0 = 60. #K

43 sun_temperature1 = 600. #K

44 sun_temperature2 = 60. #K

45

46 #引力常数

47 gravitational_constant = 6.67e-11 #m3 / kg s2

48

49 #行星运动总时长(按地球年计算)

50 #每两小时计算一次星体位置

51 end_time = 16 * 365.26 * 24. * 3600. #s

52 h = 2. * 3600 #s

53 num_steps = int(end_time /h)54 tpoints =arange(0, end_time, h)55

56

57 defthree_body_problem():58 '''主程序，计算星体位置和表面温度'''

59 positions = np.zeros([num_steps + 1, 4, 3]) #m

60 velocities = np.zeros([num_steps + 1, 4, 3]) #m / s

61

62 positions[0] = np.array([[1., 3., 2.], [6., -5., 4.], [7., 8., -7.], [7., -2., 9.]]) * 1e11 #m

63 velocities[0] = np.array([[-2., 0.5, 5.], [7., 0.5, 2.], [4., -0.5, 3.], [-10., 4., -2.]]) * 1e3 #m/s

64 positions[0] = np.array([[1., 3., 2.], [3., -5., 1.], [2., 8., -7.], [-3., -2., 9.]]) * 1e11 #m

65 velocities[0] = np.array([[0, 0., 0.], [0., 0., 0.], [0., 0., 0.], [0., 0., 0.]]) * 1e3 #m/s

66

67 defacceleration(positions):68 a = np.zeros([4, 3])69 a[0] = gravitational_constant * star_2_mass / np.linalg.norm(positions[0] - positions[1]) ** 3 *(70 positions[1] - positions[0]) + gravitational_constant * star_3_mass /np.linalg.norm(71 positions[0] - positions[2]) ** 3 * (positions[2] -positions[72 0]) + gravitational_constant * star_4_mass / np.linalg.norm(positions[0] - positions[3]) ** 3 *(73 positions[3] -positions[0])74 a[1] = gravitational_constant * star_1_mass / np.linalg.norm(positions[1] - positions[0]) ** 3 *(75 positions[0] - positions[1]) + gravitational_constant * star_3_mass /np.linalg.norm(76 positions[1] - positions[2]) ** 3 * (positions[2] -positions[77 1]) + gravitational_constant * star_4_mass / np.linalg.norm(positions[1] - positions[3]) ** 3 *(78 positions[3] - positions[1])79 a[2] = gravitational_constant * star_1_mass / np.linalg.norm(positions[2] - positions[0]) ** 3 *(80 positions[0] - positions[2]) + gravitational_constant * star_2_mass /np.linalg.norm(81 positions[2] - positions[1]) ** 3 * (positions[1] -positions[82 2]) + gravitational_constant * star_4_mass / np.linalg.norm(positions[2] - positions[3]) ** 3 *(83 positions[3] - positions[2])84 a[3] = gravitational_constant * star_1_mass / np.linalg.norm(positions[3] - positions[0]) ** 3 *(85 positions[0] - positions[3]) + gravitational_constant * star_2_mass /np.linalg.norm(86 positions[3] - positions[1]) ** 3 * (positions[1] -positions[87 3]) + gravitational_constant * star_3_mass / np.linalg.norm(positions[3] - positions[2]) ** 3 *(88 positions[2] - positions[3])89 returna90

91 defvelocity(p, t, velo):92 v = np.zeros([4, 3])93 v[0] = acceleration(p)[0] * t +velo[0]94 v[1] = acceleration(p)[1] * t + velo[1]95 v[2] = acceleration(p)[2] * t + velo[2]96 v[3] = acceleration(p)[3] * t + velo[3]97 returnvelocities[step]98

99 heat_sum, t = [0.1], [0]100

101 per_0 =0102 for step inrange(num_steps):103 #四阶龙格库塔法求星体位置

104 j1 = h *velocity(positions[step], h, velocities[step])105 j2 = h * velocity(positions[step] + 0.5 * j1, h + 0.5 *h, velocities[step])106 j3 = h * velocity(positions[step] + 0.5 * j2, h + 0.5 *h, velocities[step])107 j4 = h * velocity(positions[step] + j3, h +h, velocities[step])108 positions[step + 1] = positions[step] + (j1 + 2 * j2 + 2 * j3 + j4) / 6

109 velocities[step + 1] = velocities[step] + h * acceleration(positions[step + 1])110

111 #从 positions 中取出坐标值

112 x0, y0, z0 = positions[step, 0, 0], positions[step, 0, 1], positions[step, 0, 2]113 x1, y1, z1 = positions[step, 1, 0], positions[step, 1, 1], positions[step, 1, 2]114 x2, y2, z2 = positions[step, 2, 0], positions[step, 2, 1], positions[step, 2, 2]115 x3, y3, z3 = positions[step, 3, 0], positions[step, 3, 1], positions[step, 3, 2]116

117 #计算三个太阳对行星表面的总辐射强度

118 heat0 =sun_heat(diameter0, sun_temperature0, x0, y0, z0, x3, y3, z3)119 heat1 =sun_heat(diameter1, sun_temperature1, x1, y1, z1, x3, y3, z3)120 heat2 =sun_heat(diameter2, sun_temperature2, x2, y2, z2, x3, y3, z3)121 heat_sum.append(heat0 + heat1 +heat2)122

123 per = int(100 * step /num_steps)124 if per >per_0:125 print(per, '%')126 per_0 =per127

128 t.append(step)129

130 returnpositions, t, heat_sum131

132 positions, t, heat_sum =three_body_problem()133

134

135 empty, extinction_line, frozen_line =[], [], []136

137 for i inrange(len(t)):138 empty.append(0)139 extinction_line.append(100)140 frozen_line.append(40)141

142

143 fig, ax =plt.subplots()144 fig.set_tight_layout(True)145 plt.plot(t, heat_sum, 'b')146 plt.plot(t, empty, 'g')147 plt.plot(t, extinction_line, 'r')148 plt.plot(t, frozen_line, 'r')149

150 ax.set_xlabel('X')151 ax.set_xlim(0, len(t)+10e3)152 ax.set_ylabel('Y')153 plt.show()154

155 print("\033[1;31;47m---Begin ploting image---\033[0m")156

157 fig =plt.figure()158 ax2 =Axes3D(fig)159

160 ax2.plot(positions[:, 0, 0], positions[:, 0, 1], positions[:, 0, 2], color='b')161 ax2.plot(positions[:, 1, 0], positions[:, 1, 1], positions[:, 1, 2], color='y')162 ax2.plot(positions[:, 2, 0], positions[:, 2, 1], positions[:, 2, 2], color='b')163 ax2.plot(positions[:, 3, 0], positions[:, 3, 1], positions[:, 3, 2], color='r')164

165 zoom = 1.2e12

166 ax2.set_xlabel('X')167 #ax2.set_xlim3d(-zoom, zoom + 3.e12)

168 ax2.set_ylabel('Y')169 #ax2.set_ylim3d(-zoom, zoom)

170 ax2.set_zlabel('Z')171 #ax2.set_zlim3d(-zoom, zoom + 2.e12)

172 print("100 %")173

174 plt.show()