三体运行模拟器

scenes/func.py

@@ -6,6 +6,7 @@
 # link            :https://gitcode.net/pythoncr/
 # python_version  :3.8
 # ==============================================================================
+import matplotlib.pyplot as plt
 from common.consts import SECONDS_PER_WEEK
 from common.system import System
 
@@ -66,6 +67,22 @@ def mpl_run(bodies, dt=SECONDS_PER_WEEK, gif_file_name=None, gif_max_frame=200):
     simulator.run(dt, gif_file_name=gif_file_name, gif_max_frame=gif_max_frame)
 
 
+COSMIC_BG_COLOR = "#002563"
+COSMIC_FORE_COLOR = "white"
+
+
+def create_fig_ax(styles={}):
+    bg_color = styles["bg_color"] if "bg_color" in styles else COSMIC_BG_COLOR
+    fore_color = styles["fore_color"] if "fore_color" in styles else COSMIC_FORE_COLOR
+    if bg_color is None:
+        fig = plt.figure('天体模拟运行效果', figsize=(20, 12))
+    else:
+        fig = plt.figure('天体模拟运行效果', figsize=(20, 12), facecolor=bg_color)
+    ax = fig.gca(projection="3d")
+
+    return fig, ax
+
+
 if __name__ == '__main__':
     from bodies import Sun, Earth
 

scenes/three_body_chatgpt_01.py

+# -*- coding:utf-8 -*-
+# title           :三体场景模拟02
+# description     :三体场景模拟（3个太阳、1个地球）
+# author          :Python超人
+# date            :2023-02-11
+# link            :https://gitcode.net/pythoncr/
+# python_version  :3.8
+# ==============================================================================
+from mayavi import mlab
+from bodies import Sun, Earth
+from common.consts import SECONDS_PER_WEEK
+from scenes.func import mayavi_run, mpl_run
+
+if __name__ == '__main__':
+    """
+    注释： 3个太阳（ChatGPT生成的效果）
+    可以修改影响效果的参数为： 
+    1、三个方向的初始位置 init_position[x, y, z]
+    2、三个方向的初始速度 init_velocity[x, y, z]
+    3、天体质量 mass    
+    """
+    # 代码案例如下：
+    SIZE_SCALE = 5e1  # 生成的太阳放大 50 倍
+    RUN_DIAMETER = 1.392e6  # 真实太阳的直径
+    bodies = [
+        Sun(name="太阳1", mass=2.5e30, init_position=[100 * RUN_DIAMETER, 200 * RUN_DIAMETER, 300 * RUN_DIAMETER],
+            init_velocity=[-12.5, -12.0, 11.5],
+            size_scale=SIZE_SCALE, texture="sun1.jpg"),
+        Sun(name="太阳2", mass=2e30, init_position=[-150 * RUN_DIAMETER, 250 * RUN_DIAMETER, 350 * RUN_DIAMETER],
+            init_velocity=[11.5, 12.0, 12.5],
+            size_scale=SIZE_SCALE, texture="sun2.jpg"),
+        Sun(name="太阳3", mass=2.8e30, init_position=[200 * RUN_DIAMETER, -300 * RUN_DIAMETER, 400 * RUN_DIAMETER],
+            init_velocity=[-11.5, -12.5, -12.0],
+            size_scale=SIZE_SCALE, texture="sun2.jpg"),
+    ]
+
+    # 按照以上代码案例格式生成代码，要求 init_position 、init_velocity、mass 生成后，要保证在万有引力的作用下，能正常运行，不会碰撞
+    """
+太阳1：
+
+    初始位置：(100000, 200000, 300000) km
+    初始速度：(-1.0, -2.0, 3.0) km/s
+    质量：2.5 x 10^30 kg
+
+太阳2：
+
+    初始位置：(-150000, 250000, 350000) km
+    初始速度：(2.0, -3.0, -1.0) km/s
+    质量：2.0 x 10^30 kg
+
+太阳3：
+
+    初始位置：(200000, -300000, 400000) km
+    初始速度：(-3.0, 1.0, -2.0) km/s
+    质量：2.8 x 10^30 kg
+    """
+    # mayavi_run(bodies, SECONDS_PER_WEEK, view_azimuth=0)
+    mpl_run(bodies, SECONDS_PER_WEEK)

simulators/func.py

+# -*- coding:utf-8 -*-
+# title           :模拟器用功能库
+# description     :模拟器用功能库
+# author          :Python超人
+# date            :2023-02-11
+# link            :https://gitcode.net/pythoncr/
+# python_version  :3.8
+# ==============================================================================
+import matplotlib.pyplot as plt
+import matplotlib.colors as mcolors
+from matplotlib import ticker
+
+from common.consts import SECONDS_PER_WEEK
+from common.system import System
+
+COSMIC_BG_COLOR = "#002563"
+COSMIC_FORE_COLOR = "white"
+
+
+def get_default_colors(styles={}):
+    bg_color = styles["bg_color"] if "bg_color" in styles else "white"  # COSMIC_BG_COLOR
+    fore_color = styles["fore_color"] if "fore_color" in styles else "black"  # COSMIC_FORE_COLOR
+    # bg_color = styles["bg_color"] if "bg_color" in styles else COSMIC_BG_COLOR
+    # fore_color = styles["fore_color"] if "fore_color" in styles else COSMIC_FORE_COLOR
+    styles["bg_color"] = bg_color
+    styles["fore_color"] = fore_color
+    return bg_color, fore_color
+
+
+def create_fig_ax(styles={}):
+    bg_color, fore_color = get_default_colors(styles)
+
+    plt.rcParams['patch.facecolor'] = bg_color
+    plt.rcParams['xtick.color'] = fore_color
+    plt.rcParams['ytick.color'] = fore_color
+
+    plt.rcParams['axes.labelcolor'] = fore_color
+    plt.rcParams['axes.edgecolor'] = fore_color
+
+    plt.rcParams['axes.facecolor'] = fore_color
+    plt.rcParams['figure.facecolor'] = fore_color
+
+    if bg_color is None:
+        fig = plt.figure('天体模拟运行效果', figsize=(20, 12))
+    else:
+        fig = plt.figure('天体模拟运行效果', figsize=(20, 12), facecolor=bg_color)
+    ax = fig.gca(projection="3d")
+
+    # 调整子图区域
+    plt.subplots_adjust(left=-0.8, right=1.8, bottom=0.0, top=0.95)
+
+    # # 设置窗口大小
+    # fig.set_size_inches(20, 12)
+    # ax.set_box_aspect([2, 1, 1])  # 将三个轴设置为等比例缩放
+
+    return fig, ax
+
+
+def update_ax_styles(ax, styles={}):
+    """
+
+    :param ax:
+    :param styles:
+    :return:
+    """
+    plt.cla()
+    bg_color, fore_color = get_default_colors(styles)
+
+    # # 设置 x 轴刻度
+    # ax.xaxis.set_major_locator(ticker.MultipleLocator(base=5))
+    # ax.xaxis.set_minor_locator(ticker.MultipleLocator(base=0.01))
+    #
+    # # 设置 y 轴刻度
+    # ax.yaxis.set_major_locator(ticker.MultipleLocator(base=5))
+    # ax.yaxis.set_minor_locator(ticker.MultipleLocator(base=0.01))
+    #
+    # # 设置 z 轴刻度
+    # ax.zaxis.set_major_locator(ticker.MultipleLocator(base=5))
+    # ax.zaxis.set_minor_locator(ticker.MultipleLocator(base=0.01))
+
+    # ax.w_xaxis.line.set_color(fore_color)
+    # ax.w_yaxis.line.set_color(fore_color)
+    # ax.w_zaxis.line.set_color(fore_color)
+    # ax.xaxis.label.set_color(fore_color)
+    # ax.yaxis.label.set_color(fore_color)
+    # ax.zaxis.label.set_color(fore_color)
+
+    # ax.tick_params(axis='both', colors=fore_color)
+    # ax.tick_params(axis='x', colors=fore_color)  # only affects
+    # ax.tick_params(axis='y', colors=fore_color)  # tick labels
+    # ax.tick_params(axis='z', colors=fore_color)  # not tick marks
+    # ax.margins(x=1e20, y=None, z=None, tight=False)
+    ax.patch.set_color(bg_color)  # 设置 ax 区域背景颜色
+    bg_color = mcolors.to_rgba(bg_color)
+    fore_color_alpha = mcolors.to_rgba(fore_color, alpha=0.3)
+    # bg_color = (0.2, 0.2, 1.0, 1.0)
+    # 设置网格背景颜色
+    # ax.grid(color=fore_color)
+    # ax.grid(color=fore_color, linestyle='dashed', linewidth=1, alpha=0.1)
+    ax.xaxis._axinfo["grid"]['color'] = fore_color_alpha
+    ax.yaxis._axinfo["grid"]['color'] = fore_color_alpha
+    ax.zaxis._axinfo["grid"]['color'] = fore_color_alpha
+    ax.xaxis._axinfo["grid"]['linestyle'] = 'dashed'
+    ax.yaxis._axinfo["grid"]['linestyle'] = 'dashed'
+    ax.zaxis._axinfo["grid"]['linestyle'] = 'dashed'
+    ax.xaxis._axinfo["grid"]['linewidth'] = 0.5
+    ax.yaxis._axinfo["grid"]['linewidth'] = 0.5
+    ax.zaxis._axinfo["grid"]['linewidth'] = 0.5
+    # ax.xaxis._axinfo["grid"]['alpha'] = 0.1
+    # ax.yaxis._axinfo["grid"]['alpha'] = 0.1
+    # ax.zaxis._axinfo["grid"]['alpha'] = 0.1
+
+    ax.w_xaxis.set_pane_color(bg_color)
+    ax.w_yaxis.set_pane_color(bg_color)
+    ax.w_zaxis.set_pane_color(bg_color)
+    # ax.axes.yaxis.grid(color='red', linestyle='--', linewidth=1, alpha=1)
+    # ax.axes.yaxis.gridlines
+    # ax.set_bgcolor('red')
+    # ax.xaxis.grid(color='r', linestyle='--', linewidth=1, alpha=1)
+    # ax.axes.yaxis.grid(color='r', linestyle='--', linewidth=1, alpha=1)
+    # ax.spines['right'].set_color('blue')
+    # ax.spines['top'].set_color('none')
+    # axs0=plt.subplot(221,axisbg='#FFDAB9')
+
+    ax.set_title('天体模拟运行效果', loc='center', pad=0, fontsize="24", color=fore_color)
+    # ax.set_title('Title', )
+
+    ax.set_xlabel('X(天文单位:AU)', fontsize="18", color=fore_color)
+    ax.set_ylabel('Y(天文单位:AU)', fontsize="18", color=fore_color)
+    ax.set_zlabel('Z(天文单位:AU)', fontsize="18", color=fore_color)

simulators/mpl_simulator.py

@@ -8,9 +8,12 @@
 # ==============================================================================
 import matplotlib.pyplot as plt
 import matplotlib.animation as animation
+import matplotlib.colors as mcolors
 from simulators.simulator import Simulator
 from common.system import System
+from common.consts import AU
 from simulators.views.mpl_view import MplView
+from simulators.func import create_fig_ax, update_ax_styles
 import numpy as np
 import copy
 
@@ -26,7 +29,7 @@ class MplSimulator(Simulator):
     def __init__(self, bodies_sys: System):
         super().__init__(bodies_sys, MplView)
 
-    def save_as_gif(self, dt, gif_max_frame=200, gif_file_name='bodies_run.gif'):
+    def save_as_gif(self, dt, gif_max_frame=200, gif_file_name='bodies_run.gif', styles={}):
         """
         保存 GIF 文件
         :param dt: 单位：秒，按时间差进行演变，值越小越精确，但演变速度会慢。
@@ -34,8 +37,9 @@ class MplSimulator(Simulator):
         :param gif_file_name: 导出的 gif 文件名
         :return:
         """
-        fig = plt.figure('天体模拟运行效果', figsize=(16, 12))
-        ax = fig.gca(projection="3d")
+
+        fig, ax = create_fig_ax()
+
         views_frames = []
         for i in range(gif_max_frame):
             self.evolve(dt)
@@ -45,7 +49,7 @@ class MplSimulator(Simulator):
         def update(num):
             body_views = views_frames[num]
             print("\rGIF 生成进度：%d/%d %.2f" % (num + 1, gif_max_frame, ((num + 1) / gif_max_frame) * 100) + "%", end='')
-            return self.show_figure(ax, body_views, pause=0)
+            return self.show_figure(ax, body_views, pause=0, update_ax=update_ax_styles, styles=styles)
 
         ani = animation.FuncAnimation(fig=fig, func=update, frames=np.arange(0, gif_max_frame), interval=1)
         ani.save(gif_file_name)
@@ -61,22 +65,23 @@ class MplSimulator(Simulator):
         """
         gif_file_name = kwargs["gif_file_name"] if "gif_file_name" in kwargs else None
         gif_max_frame = kwargs["gif_max_frame"] if "gif_max_frame" in kwargs else None
+        styles = kwargs["styles"] if "styles" in kwargs else {}
 
         if gif_file_name is not None:
-            self.save_as_gif(dt, gif_max_frame=gif_max_frame, gif_file_name=gif_file_name)
+            self.save_as_gif(dt, gif_max_frame=gif_max_frame, gif_file_name=gif_file_name, styles=styles)
             return
 
-        fig = plt.figure('天体模拟运行效果', figsize=(16, 12))
-        ax = fig.gca(projection="3d")
+        fig, ax = create_fig_ax()
+
         # TODO: 注意：显示动态图，需先进行以下设置：
         # Pycharm：：File –> Settings –> Tools –> Python Scientific –> Show plots in tool window(取消打勾)
 
         while True:
             self.evolve(dt)
             body_views = copy.deepcopy(self.body_views)
-            self.show_figure(ax, body_views, pause=0.1)
+            self.show_figure(ax, body_views, pause=0.1, update_ax=update_ax_styles, styles=styles)
 
-    def show_figure(self, ax, bodies, pause=0.1):
+    def show_figure(self, ax, bodies, pause=0.1, update_ax=None, styles={}):
         """
 
         :param ax:
@@ -84,12 +89,9 @@ class MplSimulator(Simulator):
         :param pause:
         :return:
         """
-        plt.cla()
-
-        ax.set_title('天体模拟运行效果')
-        ax.set_xlabel('X')
-        ax.set_ylabel('Y')
-        ax.set_zlabel('Z')
+        if update_ax is not None:
+            # 更新 ax
+            update_ax(ax, styles)
 
         for idx, body in enumerate(bodies):
             if body.is_fixed_star:
@@ -98,21 +100,40 @@ class MplSimulator(Simulator):
                 color = 'blue'
             # size = 800 if str(body.name).lower().startswith("sun") else 500
             size = body.raduis / 80000
-            pos = body.position
-            # 天体名称
-            ax.text(pos[0], pos[1], pos[2] + 0.006, s=body.name, color=color, fontsize=12)
+            # size = pow(body.raduis / AU * body.size_scale,3)
+            pos = body.position / AU
+
             # 天体
             ax.scatter(pos[0], pos[1], pos[2], color=color, s=size, alpha=0.8)
             # for _his_pos in body.his_position():
             #     ax.scatter3D(_his_pos[0], _his_pos[1], _his_pos[2], color=color, alpha=0.5)
             # ax.scatter(his_pos[0], his_pos[1], his_pos[2], color=colors[idx], s=10)
-            his_pos = body.his_position
+            his_pos = np.array(body.his_position) / AU
             tail_len = len(his_pos)
             if tail_len > 1:
                 _his_pos = list(zip(*his_pos))
                 # 历史轨迹线
                 ax.plot3D(_his_pos[0], _his_pos[1], _his_pos[2], color=color, alpha=0.5)
 
+            z_range = ax.get_zlim()[1] - ax.get_zlim()[0]
+            # z_range = 0.11   1.52
+            ax.text(pos[0], pos[1], pos[2] + size*(z_range/5000), s=body.name, color=color, fontsize=12)
+            # 计算每个坐标轴的数据范围
+            # x_range = ax.get_xlim()[1] - ax.get_xlim()[0]
+            # y_range = ax.get_ylim()[1] - ax.get_ylim()[0]
+            # z_range = ax.get_zlim()[1] - ax.get_zlim()[0]
+            # r = (body.raduis/AU)
+            # k = r / (max(x_range, y_range, z_range))
+            # k = r / np.array([x_range, y_range, z_range])
+            # 计算文本的偏移量
+            # 偏移量与刻度、球体大小相关
+            # text_offset = (max(x_range, y_range, z_range) + (size)) / 80
+            # text_offset = *(z_range/3)
+            # ax.text(pos[0] + k[0] * r, pos[1] + k[1] * r, pos[2] + k[2] * r,
+            #         s=body.name,
+            #         color=color,
+            #         fontsize=12)
+
         if pause > 0:
             plt.pause(pause)