system.py 6.4 KB
Newer Older
M
march3 已提交
1
# -*- coding:utf-8 -*-
M
march3 已提交
2 3
# title           :天体系统
# description     :天体系统,多个天体就是一个系统
M
march3 已提交
4
# author          :Python超人
M
march3 已提交
5 6
# date            :2023-02-11
# link            :https://gitcode.net/pythoncr/
M
march3 已提交
7 8 9
# python_version  :3.8
# ==============================================================================
import numpy as np
M
march3 已提交
10
from common.consts import AU, G
M
march3 已提交
11
from bodies import Body, Sun, Mercury, Venus, Earth, Mars, Jupiter, Saturn, Uranus, Neptune, Pluto
M
march3 已提交
12
from common.func import calculate_distance
M
march3 已提交
13 14 15


class System(object):
M
march3 已提交
16 17 18 19
    """
    天体系统
    """

M
march3 已提交
20 21 22 23 24 25
    def __init__(self, bodies, max_distance=200 * AU):
        """

        :param bodies:
        :param max_distance:系统的最大范围,超出范围的天体就不显示了
        """
M
march3 已提交
26
        self.bodies = bodies
M
march3 已提交
27
        self.max_distance = max_distance
M
march3 已提交
28 29 30 31 32 33

    def add(self, body):
        self.bodies.append(body)

    def total_mass(self):
        """
M
march3 已提交
34
        总质量
M
march3 已提交
35 36
        :return:
        """
M
march3 已提交
37
        total_mass = 0.0
M
march3 已提交
38
        for body in self.bodies:
M
march3 已提交
39 40
            total_mass += body.mass
        return total_mass
M
march3 已提交
41

M
march3 已提交
42 43
    def __repr__(self):
        return 'System({})'.format(self.bodies)
M
march3 已提交
44

M
march3 已提交
45
    def center_of_mass(self):
M
march3 已提交
46
        """
M
march3 已提交
47
        质心
M
march3 已提交
48 49
        :return:
        """
M
march3 已提交
50
        r = np.zeros(2)
M
march3 已提交
51
        for body in self.bodies:
M
march3 已提交
52 53
            r = body.mass * body.position
        return r / self.total_mass()
M
march3 已提交
54

M
march3 已提交
55
    def evolve(self, dt):
M
march3 已提交
56 57
        """

M
march3 已提交
58
        :param dt:
M
march3 已提交
59 60
        :return:
        """
M
march3 已提交
61
        self.calc_bodies_acceleration()
M
march3 已提交
62 63

        for body in self.bodies:
M
march3 已提交
64 65 66 67
            # acceleration 加速度
            body.velocity += body.acceleration * dt
            # body.position += 0.5 * body.acceleration * (dt ** 2)
            body.position += body.velocity * dt
M
march3 已提交
68

三月三net's avatar
三月三net 已提交
69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98
    def save_to_json(self, json_file_name, params=None):
        """

        :param json_file_name:
        :param params:
        :return:
        """
        import json
        import os
        # json_file = os.path.join("../data", json_file_name)
        filed_names = ["name", "mass", "init_position", "init_velocity",
                       "density", "color", "texture",
                       "size_scale", "distance_scale",  # "parent"
                       "rotation_speed", "ignore_mass", "is_fixed_star"]
        bodies = []
        for b in self.bodies:
            body = {}
            for filed_name in filed_names:
                filed_value = getattr(b, filed_name)
                if type(filed_value) is np.ndarray:
                    filed_value = filed_value.tolist()
                body[filed_name] = filed_value
            bodies.append(body)
        data = {"bodies": bodies}
        if params is not None:
            data["params"] = params
        json_str = json.dumps(data, indent=2, ensure_ascii=False, separators=(',', ': '))
        with open(json_file_name, "w", encoding='utf-8') as f:
            f.write(json_str)

M
march3 已提交
99
    def calc_bodies_acceleration(self):
M
march3 已提交
100
        """
M
march3 已提交
101 102
        计算加速度
        :return:
M
march3 已提交
103
        """
M
march3 已提交
104 105 106

        def valid_body(body):
            """
M
march3 已提交
107
            判断是否为有效的天体
M
march3 已提交
108 109 110
            :param body:
            :return:
            """
M
march3 已提交
111
            if not body.appeared:  # 不显示
M
march3 已提交
112
                return False
三月三net's avatar
三月三net 已提交
113 114 115 116 117
            # if self.max_distance > 0:
            #     # 超过了 max_distance 距离,则不显示,并消失
            #     if calculate_distance(body.position) > self.max_distance:
            #         body.appeared = False
            #         return False
M
march3 已提交
118 119 120

            return True

M
march3 已提交
121
        # self.bodies = list(filter(valid_body, self.bodies))
M
march3 已提交
122

M
march3 已提交
123
        for body1 in self.bodies:
三月三net's avatar
三月三net 已提交
124 125
            if body1.ignore_mass:
                continue
M
march3 已提交
126 127
            if not valid_body(body1):
                continue
M
march3 已提交
128
            acceleration = np.zeros(3)
M
march3 已提交
129
            for body2 in self.bodies:
三月三net's avatar
三月三net 已提交
130 131
                if body2.ignore_mass:
                    continue
M
march3 已提交
132
                if self.max_distance > 0:
M
march3 已提交
133
                    if calculate_distance(body1.position) > self.max_distance:  # 超过了max_distance距离,则消失
M
march3 已提交
134
                        body1.appeared = False
M
march3 已提交
135
                    if calculate_distance(body2.position) > self.max_distance:  # 超过了max_distance距离,则消失
M
march3 已提交
136 137 138 139 140
                        body2.appeared = False

                if False == body1.appeared or body2.appeared == False:
                    continue

M
march3 已提交
141 142
                if body1 is body2:
                    continue
M
march3 已提交
143 144
                elif body1.ignore_gravity(body2) or body2.ignore_gravity(body1):
                    continue
M
march3 已提交
145 146

                r = body2.position - body1.position
M
march3 已提交
147 148 149 150
                # G = 6.67e-11 # 万有引力常数
                # m/s² = kg * m / m**3
                # km/s² = kg * m / m**3 / 1e9
                # acceleration = G * body2.mass * dx / (d ** 3)
M
march3 已提交
151
                acceleration += (G * body2.mass * r / np.linalg.norm(r) ** 3) / 1e9
M
march3 已提交
152

M
march3 已提交
153
            body1.acceleration = acceleration
M
march3 已提交
154 155


M
march3 已提交
156
if __name__ == '__main__':
三月三net's avatar
三月三net 已提交
157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195
    # body_sys = System([
    #     Sun(),  # 太阳
    #     Mercury(),  # 水星
    #     Venus(),  # 金星
    #     Earth(),  # 地球
    #     Mars(),  # 火星
    #     Jupiter(),  # 木星
    #     Saturn(),  # 土星
    #     Uranus(),  # 天王星
    #     Neptune(),  # 海王星
    #     Pluto()  # 冥王星(从太阳系的行星中排除)
    # ])
    import math

    mass = 2e30
    r = 2 * AU
    # p = 14.9
    p = 14.89
    bodies = [
        Sun(name="太阳A红色", mass=mass,
            init_position=[0, r * math.sqrt(3), 0],  # 位置
            init_velocity=[-p, 0, 0],  # 速度(km/s)
            size_scale=5e1, texture="sun2.jpg", color=(255, 0, 0)),  # 太阳放大 100 倍
        Sun(name="太阳B绿色", mass=mass,
            init_position=[-r, 0, 0],
            init_velocity=[1 / 2 * p, -math.sqrt(3) / 2 * p, 0],
            size_scale=5e1, texture="sun2.jpg", color=(0, 255, 0)),  # 太阳放大 100 倍
        Sun(name="太阳C蓝色", mass=mass,
            init_position=[r, 0, 0],
            init_velocity=[1 / 2 * p, math.sqrt(3) / 2 * p, 0],
            size_scale=5e1, texture="sun2.jpg", color=(0, 0, 255)),  # 太阳放大 100 倍
        Earth(name="地球",
              # init_position=[0, -AU * -2, 5 * AU],
              init_position=[0, math.sqrt(3) * r / 6, 5 * AU],
              init_velocity=[0, 0, -10],
              size_scale=4e3, distance_scale=1),  # 地球放大 4000 倍,距离保持不变
    ]
    body_sys = System(bodies)
    print(body_sys.save_to_json("../data/tri_bodies_sim_perfect_01.json"))