三体运行模拟器

bodies/body.py

@@ -14,6 +14,9 @@ from common.consts import AU
 
 
 class Body(metaclass=ABCMeta):
+    """
+    天体信息基类
+    """
     def __init__(self, name, mass, init_position, init_velocity,
                  density=5e3, color=(125 / 255, 125 / 255, 125 / 255),
                  texture=None, size_scale=1.0, distance_scale=1.0):

bodies/saturn.py

@@ -46,6 +46,14 @@ class Saturn(Body):
         """
         return True
 
+    @property
+    def rings_color(self):
+        """
+        土星光环的颜色
+        :return:
+        """
+        return 173, 121, 92
+
 
 if __name__ == '__main__':
     saturn = Saturn()

simulators/mayavi_simulator.py

@@ -9,7 +9,7 @@
 from mayavi import mlab
 from simulators.simulator import Simulator
 from common.system import System
-from simulators.viewers.mayavi_view import MayaviView
+from simulators.views.mayavi_view import MayaviView
 
 
 class MayaviSimulator(Simulator):

simulators/mpl_simulator.py

@@ -10,7 +10,7 @@ import matplotlib.pyplot as plt
 import matplotlib.animation as animation
 from simulators.simulator import Simulator
 from common.system import System
-from simulators.viewers.mpl_view import MplView
+from simulators.views.mpl_view import MplView
 import numpy as np
 import time
 

simulators/simulator.py

@@ -54,6 +54,8 @@ class Simulator(metaclass=ABCMeta):
             view.his_position = body.his_position()
             view.is_fixed_star = body.is_fixed_star
             view.has_rings = body.has_rings
+            view.size_scale = body.size_scale
+            view.distance_scale = body.distance_scale
 
             view.update()
 

simulators/viewers/body_view.py → simulators/views/body_view.py

@@ -29,7 +29,7 @@ class BodyView(metaclass=ABCMeta):
                 self.color = tuple(np.array(body.color) / 255)
             else:
                 self.color = self.__get_texture_main_color(self.texture)
-        self.sphere = self.appear()
+        self.appear()
         self.position = [None, None, None]
         self.name = None
         self.mass = None

simulators/viewers/mayavi_view.py → simulators/views/mayavi_view.py

@@ -12,7 +12,7 @@ import os
 import matplotlib.pyplot as plt
 from common.func import get_dominant_colors
 
-from simulators.viewers.body_view import BodyView
+from simulators.views.body_view import BodyView
 import numpy as np
 
 
@@ -23,48 +23,59 @@ class MayaviView(BodyView):
 
     def update(self):
         """
-
+        更新天体信息和数据，比如：更新天体的位置
         :return:
         """
-        x_offset = self.body.position[0] - self.sphere.mlab_source.x
-        y_offset = self.body.position[1] - self.sphere.mlab_source.y
-        z_offset = self.body.position[2] - self.sphere.mlab_source.z
-
-        self.sphere.mlab_source.set(x=self.position[0], y=self.position[1], z=self.position[2])
-        return x_offset[0], y_offset[0], z_offset[0]
-
-    # def __find_texture(self, texture):
-    #     """
-    #     在多路径下寻找纹理图片
-    #     :param texture: 纹理图片
-    #     :return: 纹理图片的路径
-    #     """
-    #     paths = ['./textures', '../textures']
-    #     for path in paths:
-    #         p = path + "/" + texture
-    #         if os.path.exists(p):
-    #             return p
-    #
-    #     return None
-    #
-    # def __texture_to_color(self, texture):
-    #     """
-    #     根据纹理图片获取颜色
-    #     :param texture:
-    #     :return:
-    #     """
-    #     colors = get_dominant_colors(texture)
-    #     first_color = colors[0]
-    #     # print(self.name, first_color)
-    #     return tuple(np.array(first_color) / 255)
+        if hasattr(self.sphere, "mlab_source"):
+            x_offset = self.body.position[0] - self.sphere.mlab_source.x
+            y_offset = self.body.position[1] - self.sphere.mlab_source.y
+            z_offset = self.body.position[2] - self.sphere.mlab_source.z
+            self.sphere.mlab_source.set(x=self.position[0], y=self.position[1], z=self.position[2])
+        else:
+            x_offset, y_offset, z_offset = 0, 0, 0
+
+        if hasattr(self, "rings"):
+            if hasattr(self.rings, "mlab_source"):
+                if hasattr(self, "rings") and self.body.has_rings:
+                    x = self.rings.mlab_source.x
+                    y = self.rings.mlab_source.y
+                    z = self.rings.mlab_source.z
+                    self.rings.mlab_source.set(x=x + x_offset, y=y + y_offset, z=z + z_offset)
+
+        # return x_offset[0], y_offset[0], z_offset[0]
+
+    def build_rings(self):
+        if not hasattr(self, "rings") or self.rings is None:
+
+            R = 2
+            r = 0.3
+            rings_scale = 0.5e5
+            resolution = 50
+            theta = np.linspace(0, 2 * np.pi, resolution)
+            phi = np.linspace(0, 2 * np.pi, resolution)
+            torus = np.zeros((3, resolution, resolution))
+
+            for i in range(0, resolution):
+                for j in range(0, resolution):  # size_scale=8.0e2
+                    torus[0][i][j] = (R + r * np.cos(phi[j])) * np.cos(theta[i]) * \
+                                     self.body.size_scale * rings_scale + self.body.position[0]
+                    torus[1][i][j] = (R + r * np.cos(phi[j])) * np.sin(theta[i]) * \
+                                     self.body.size_scale * rings_scale + self.body.position[1]
+                    torus[2][i][j] = 1 * np.sin(phi[j]) + self.body.position[2]
+            rings_color = (173 / 255, 121 / 255, 92 / 255)
+            if hasattr(self.body, "rings_color"):
+                rings_color = tuple(np.array(self.body.rings_color) / 255)
+            self.rings = mlab.mesh(torus[0], torus[1], torus[2], color=rings_color,
+                                   representation='surface')
+        return self.rings
 
     def appear(self):
         """
-        构建球体对象
+        天体显示的操作，比如：构建天体视图对象
         :return:
         """
         if not hasattr(self, "sphere") or self.sphere is None:
-            scale_factor = self.body.size_scale * self.body.raduis
+            scale_factor = self.body.size_scale * self.body.diameter
             sphere = mlab.points3d(self.body.position[0], self.body.position[1], self.body.position[2],
                                    scale_factor=scale_factor,
                                    color=self.color,
@@ -87,17 +98,30 @@ class MayaviView(BodyView):
 
         if self.texture is not None and self.texture != '':
             self.__set_texture(self.texture)
-            pass
-        return self.sphere
+
+        if self.body.has_rings:
+            self.build_rings()
+            return self.sphere, self.rings
+
+        return self.sphere,
 
     def __set_texture(self, image_file):
-        if not os.path.exists(image_file):
-            return
-        img = plt.imread(image_file)
+        """
+        设置纹理图片到天体
+        :param image_file:
+        :return:
+        """
         outfile = image_file.replace('.jpg', '_flipped.jpg')
-        img = img[::-1, ...]
-        plt.imsave(outfile, img)
-        image_file = outfile
+        if os.path.exists(outfile):
+            image_file = outfile
+        else:
+            if not os.path.exists(image_file):
+                return
+            img = plt.imread(image_file)
+            img = img[::-1, ...]
+            plt.imsave(outfile, img)
+            image_file = outfile
+
         img = tvtk.JPEGReader(file_name=image_file)
         texture = tvtk.Texture(input_connection=img.output_port, interpolate=0, repeat=0)
         self.sphere.actor.actor.texture = texture

simulators/viewers/mpl_view.py → simulators/views/mpl_view.py

@@ -10,7 +10,7 @@ import os
 import matplotlib.pyplot as plt
 from common.func import get_dominant_colors
 
-from simulators.viewers.body_view import BodyView
+from simulators.views.body_view import BodyView
 import numpy as np