From f206c8c598c16c0d0fb4309f67958b8cbcce3628 Mon Sep 17 00:00:00 2001 From: march3 Date: Fri, 14 Jul 2023 21:41:55 +0800 Subject: [PATCH] =?UTF-8?q?Python=E8=B6=85=E4=BA=BA-=E5=AE=87=E5=AE=99?= =?UTF-8?q?=E6=A8=A1=E6=8B=9F=E5=99=A8?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- common/system.py | 41 +++++++++++++++++-- .../science/jupiter_moon_protects_earth.py | 4 +- 2 files changed, 39 insertions(+), 6 deletions(-) diff --git a/common/system.py b/common/system.py index 48c0242..15111f5 100644 --- a/common/system.py +++ b/common/system.py @@ -163,7 +163,7 @@ class System(object): body1.acceleration = acceleration return True - def calculate_gravitational_accelerations(self, masses, positions): + def calculate_gravitational_accelerations_np(self, masses, positions): '''Params: - positions: numpy array of size (n,3) - masses: numpy array of size (n,) @@ -177,9 +177,42 @@ class System(object): forces = G * disps * mass_matrix / np.expand_dims(dists, 2) ** 3 return forces.sum(axis=1) / masses.reshape(-1, 1) + def calculate_gravitational_accelerations_cp(self, masses, positions): + '''Params: + - positions: numpy array of size (n,3) + - masses: numpy array of size (n,) + pip install -i https://pypi.douban.com/simple/ cupy + + nvcc -V + nvcc: NVIDIA (R) Cuda compiler driver + Copyright (c) 2005-2021 NVIDIA Corporation + Built on Fri_Dec_17_18:28:54_Pacific_Standard_Time_2021 + Cuda compilation tools, release 11.6, V11.6.55 + Build cuda_11.6.r11.6/compiler.30794723_0 + + pip install -i https://pypi.douban.com/simple/ cupy-cuda116 + cupy-11.6.0.tar.gz + https://pypi.doubanio.com/packages/70/e1/acc77e327548cce7cb28eb345b7f31ab85b6a3d99214479f9bcbe78e6e9b/cupy_cuda116-10.6.0-cp39-cp39-win_amd64.whl + https://pypi.doubanio.com/packages/e3/62/c808623b8000185efebd8b4542efdf76cc93d20dfd3f0a3eaeb5e5697430/cupy-11.6.0.tar.gz#sha256=53dbb840072bb32d4bfbaa6bfa072365a30c98b1fcd1f43e48969071ad98f1a7 + + ''' + import cupy as cp + masses = cp.array(masses) + positions = cp.array(positions) + mass_matrix = masses.reshape((1, -1, 1)) * masses.reshape((-1, 1, 1)) + disps = positions.reshape((1, -1, 3)) - positions.reshape((-1, 1, 3)) # displacements + dists = cp.linalg.norm(disps, axis=2) + dists[dists == 0] = 1 # Avoid divide by zero warnings + forces = G * disps * mass_matrix / cp.expand_dims(dists, 2) ** 3 + accelerations = forces.sum(axis=1) / masses.reshape(-1, 1) + return accelerations.get() + + def calculate_gravitational_accelerations(self, masses, positions): + return self.calculate_gravitational_accelerations_np(masses, positions) + def calc_bodies_acceleration_bak(self): """ - 计算加速度(使用矩阵的方式，性能提高很多) + 计算加速度(使用矩阵的方式，性能提高很多，不支持指定重力对某天体有效) @return: """ valid_bodies = list(filter(lambda b: not b.ignore_mass, self.bodies)) @@ -192,11 +225,11 @@ class System(object): accelerations = self.calculate_gravitational_accelerations(masses, positions) for idx, body in enumerate(valid_bodies): - body.acceleration = accelerations[idx]/1000 + body.acceleration = accelerations[idx] / 1000 def calc_bodies_acceleration(self): """ - 计算加速度（性能非常低，代码保留） + 计算加速度（性能非常低） @return: """ # 如果快速计算成功，则无需再计算 diff --git a/sim_scenes/science/jupiter_moon_protects_earth.py b/sim_scenes/science/jupiter_moon_protects_earth.py index eaeaf15..8b95851 100644 --- a/sim_scenes/science/jupiter_moon_protects_earth.py +++ b/sim_scenes/science/jupiter_moon_protects_earth.py @@ -19,7 +19,7 @@ import math import numpy as np -class JupiterProtectsEarthSim: +class JupiterMoonProtectsEarthSim: def __init__(self, comet_num=20): """ @@ -204,7 +204,7 @@ if __name__ == '__main__': # 设置计时器的最小时间单位为年 BodyTimer().min_unit = BodyTimer.MIN_UNIT_YEARS - sim = JupiterProtectsEarthSim(comet_num=20) + sim = JupiterMoonProtectsEarthSim(comet_num=30) # 运行前会触发 on_ready UrsinaEvent.on_ready_subscription(sim.on_ready) -- GitLab