# 我是一些注释！！！
import numpy as np
import pyqtgraph as pg
import pyqtgraph.opengl as gl
from pyqtgraph.Qt import QtCore
from PyQt5.QtCore import *
from PyQt5.QtWidgets import *
from PyQt5.QtGui import QIcon
import csv
DATA=np.zeros((19))
[文档]def remove_i(x, i):
'''计算加速度要去除自身所在位置影响
参数:
- x:位置矩阵x
- i:第i个质点'''
shape = (x.shape[0]-1,) + x.shape[1:]
y = np.empty(shape, dtype=float)
y[:i] = x[:i]
y[i:] = x[i+1:]
return y
[文档]def a(i, x, G, m):
'''计算第i个质量点在位置矩阵X,引力常数G和质量为m的条件下的加速度
参数:
- i:第几个质点
- x:位置矩阵X
- G:引力常数G
- m:质量矩阵
返回值:
- result:第i个质量点的加速度
'''
x_i = x[i]
x_j = remove_i(x, i)
m_j = remove_i(m, i)
diff = x_j - x_i
# if abs(np.min(diff))<1e-6:
# raise Exception('精度失败！'+str(diff))
mag3 = np.sum(diff**2, axis=1)**1.5
result = G * np.sum(diff * (m_j / mag3)[:,np.newaxis], axis=0)
return result
[文档]def timestep(x0, v0, G, m, dt):
'''给定瞬时条件和时间间隔，为所有质量点计算下一瞬间位置和速度(标准离散时间运动)
参数:
- x0:初始位置
- v0:初始速度
- G:引力常数
- m:质量
- dt:时间微分（微分越小结果越精确，但会降低程序的运行速度）
返回值:
- x1:在一个dt后各点的末位置
- v1:在一个dt后各点的末速度
'''
N = len(x0)
x1 = np.empty(x0.shape, dtype=float)
v1 = np.empty(v0.shape, dtype=float)
for i in range(N):
a_i0 = a(i, x0, G, m)
# print('a%d='%(i+1)+str(a_i0))
v1[i] = a_i0 * dt + v0[i]
x1[i] = a_i0 * dt**2 + v0[i] * dt + x0[i]
return x1, v1
[文档]def timestep2(x0, v0, G, m, dt):
'''给定瞬时条件和时间间隔，为所有质量点计算下一瞬间位置和速度(直线匀速运动)
参数:
- x0:初始位置
- v0:初始速度
- G:引力常数
- m:质量
- dt:时间微分（微分越小结果越精确，但会降低程序的运行速度）
返回值:
- x1:在一个dt后各点的末位置
- v1:在一个dt后各点的末速度'''
N = len(x0)
x1 = np.empty(x0.shape, dtype=float)
v1 = np.empty(v0.shape, dtype=float)
for i in range(N):
v1[i] = [0,0,0]
x1[i] = x0[i]+0.1
return x1, v1
[文档]def initial_cond(N=3, D=3):
'''在空间中为N个质量点产生随机条件
参数:
- N:质量点个数，三体运动设为3
- D:空间维度，三维空间设为3
'''
global DATA
x0 = np.random.rand(N, D)*10
v0 = np.zeros((N, D), dtype=float)
m = np.ones(N, dtype=float)
DATA=np.asarray(
np.concatenate(
([0],
x0.flatten(),
v0.flatten(),)
)
)
# print(DATA)
return x0, v0, m
[文档]def simulate(N, D, S, G, dt):
'''数值模拟函数'''
x0, v0, m = initial_cond(N, D)
for s in range(S):
x1, v1 = timestep(x0, v0, G, m, dt)
x0, v0 = x1, v1
[文档]class myTable(QTableWidget):
'''自定义表格'''
pro='constant table'
[文档]class MainPrograme():
def __init__(self):
'''初始化函数，主要作用为声明各个组件'''
app = pg.mkQApp("三体问题")
self.win=QMainWindow()
self.cw=QWidget()
self.layout=QGridLayout()
self.w = gl.GLViewWidget()
self.axs=gl.GLAxisItem()
# print(self.axs.childItems())
self.axs.setSize(10,10,10)
self.axs.translate(-30,-30,0)
self.line1=gl.GLLinePlotItem(width=1,antialias=True,color=pg.mkColor((51,1.4)))
self.line2=gl.GLLinePlotItem(width=1,antialias=True,color=pg.mkColor((51,1.4)))
self.line3=gl.GLLinePlotItem(width=1,antialias=True,color=pg.mkColor((51,1.4)))
self.linedraw(False)
self.w.addItem(self.line1)
self.w.addItem(self.line2)
self.w.addItem(self.line3)
self.t =QtCore.QTimer()
self.timing=0
self.btn_stop=QPushButton('暂停/运动')
self.btn_stop.clicked.connect(self.animate_start_stop)
self.btn_table=QPushButton('设置参数')
self.btn_table.clicked.connect(self.parametersettings)
self.c=QRadioButton('是否绘制轨迹')
self.c.toggled.connect(self.linedraw)
self.slider=QSlider(orientation=Qt.Orientation.Horizontal)
self.label_for_slider=QLabel(text='speed='+self.slider.value().__str__())
self.Nbody_lable=QLabel(text='设置天体形状大小')
self.Nbodies=QSpinBox()
self.Nbodies.valueChanged.connect(self.change_size)
# 坐标网格
self.g = gl.GLGridItem()
self.ini_camera=QPushButton('摄像机归位')
self.ini_camera.clicked.connect(self.ini)
# self.poses = np.random.random(size=(2,3))
# self.poses[0] = (0,0,0)
self.points_color = np.tile(np.asarray([0,255,255,255]),(3,1))
self.x0, self.v0, self.m = initial_cond(3, 3)
# self.varibles
self.dt=0.01
self.G=1
self.x1, self.v1 = timestep(self.x0, self.v0, self.G, self.m, self.dt)
# 控制散点大小
# size = np.random.random(size=pos.shape[0])*10
self.point_size=10
self.sp2 = gl.GLScatterPlotItem(pos=self.x0, color=self.points_color, size=self.point_size)
[文档] def change_size(self):
'''改变散点外观大小'''
self.sp2.size=self.Nbodies.value()
self.update()
[文档] def ini(self):
'''初始化摄像机位置'''
self.w.reset()
self.w.setCameraPosition(distance=100)
[文档] def load(self):
'''加载各个组件，分配空间布局并为组件绑定事件监听函数'''
self.win.resize(880,600)
self.win.show()
self.win.setWindowTitle('三体问题数值计算及动画模拟')
self.win.setWindowIcon(QIcon('icon.jpeg'))
self.win.setCentralWidget(self.cw)
self.layout.setSpacing(20)
self.cw.setLayout(self.layout)
leftLayout=QVBoxLayout()
self.layout.addWidget(self.w,0,1,6,3)
self.layout.addLayout(leftLayout,1,0)
leftLayout.setSpacing(40)
leftLayout.addWidget(self.btn_stop)
leftLayout.addWidget(self.c)
leftLayout.addWidget(self.Nbodies)
leftLayout.addWidget(self.btn_table)
leftLayout.addWidget(self.Nbody_lable)
leftLayout.addWidget(self.ini_camera)
speed_layout=QVBoxLayout()
speed_layout.setSpacing(10)
speed_layout.addWidget(self.label_for_slider,alignment=Qt.AlignmentFlag.AlignCenter)
speed_layout.addWidget(self.slider)
leftLayout.addLayout(speed_layout)
self.g.scale(4,4,1)
self.Nbody_lable.setMaximumSize(130,15)
self.w.addItem(self.g)
self.w.addItem(self.sp2)
self.w.setCameraPosition(distance=50)
self.w.addItem(self.axs)
self.Nbodies.setMinimum(1)
self.Nbodies.setMaximum(10)
self.slider.setMaximumSize(150,40)
self.slider.setTickPosition(QSlider.TickPosition.TicksBelow)
self.slider.setTickInterval(100)
self.slider.setMinimum(10)
self.slider.setMaximum(1000)
self.slider.valueChanged.connect(self.changeSpeed)
self.t.timeout.connect(self.update)
self.statusbar=self.win.statusBar()
self.statusbar.showMessage('程序正在运行！')
qact_exit=QAction(text='exit')
qact_exit.setShortcut('ctrl+q')
qact_exit.setStatusTip('退出')
qact_exit.triggered.connect(qApp.quit)
qact_open=QAction(text='open')
qact_open.setShortcut('ctrl+n')
qact_open.setStatusTip('打开csv文件')
qact_open.triggered.connect(self.filedialog_open)
qact_save=QAction(text='save')
qact_save.setShortcut('ctrl+s')
qact_save.setStatusTip('保存csv文件')
qact_save.triggered.connect(self.filedialog_save)
qact_c1=QAction(text='调整一号天体颜色')
qact_c1.setStatusTip('打开调色板')
qact_c1.triggered.connect(self.colordiallog)
qact_c2=QAction(text='调整二号天体颜色')
qact_c2.setStatusTip('打开调色板')
qact_c2.triggered.connect(self.colordiallog)
qact_c3=QAction(text='调整三号天体颜色')
qact_c3.setStatusTip('打开调色板')
qact_c3.triggered.connect(self.colordiallog)
menu=self.win.menuBar()
filemenu=menu.addMenu('文件')
filemenu.addAction(qact_exit)
filemenu.addAction(qact_open)
filemenu.addAction(qact_save)
look=menu.addMenu('视图')
look.addAction(qact_c1)
look.addAction(qact_c2)
look.addAction(qact_c3)
helpact=QAction(text='帮助')
helpact.setStatusTip('如何使用...')
helpact.triggered.connect(self.helper)
authinfo=QAction(text='关于我')
authinfo.setStatusTip('作者信息')
authinfo.triggered.connect(self.info)
more=menu.addMenu('更多')
more.addAction(helpact)
more.addAction(authinfo)
self.x1, self.v1 = timestep(self.x0, self.v0, self.G, self.m, self.dt)
pg.exec()
[文档] def info(self):
'''弹出作者信息对话框'''
reply = QMessageBox.information(self.win, '作者信息','计科二班闫泽轩,个人网页home.ustc.edu.cn/~yzx_ustc',QMessageBox.Yes | QMessageBox.No,QMessageBox.Yes)
[文档] def helper(self):
'''弹出提示信息框'''
reply = QMessageBox.information(self.win, '帮助',
"ctrl+s可以保存当前状态信息成csv格式\n\
ctrl+n可以打开以csv格式保存的数据信息，详情参考我的个人网页中的样例\n\
设置参数后将丢失调参前信息，即使没有进行修改，请及时保存！\n\
ctrl+q退出程序，不要手贱哦~",
QMessageBox.Yes | QMessageBox.No,QMessageBox.Yes)
[文档] def linedraw(self,checked):
'''响应单选框是否绘制轨迹的事件请求
参数:
checked:单选框是否被选中'''
# print('hahaha'+str(checked))
if checked:
self.line1.show()
self.line2.show()
self.line3.show()
else:
self.line1.hide()
self.line2.hide()
self.line3.hide()
[文档] def myfmt(self,item):
'''保证参数设置表格数据符合格式要求，并更新至DATA之中
参数：
item:发生改变的表格元素,若数据不符合格式将被置0'''
global DATA
#DATA已保证初始化，仅有一维数组
sender=self.win.sender()
# print(item.column())
if sender.pro=='const':
if isinstance(item,QTableWidgetItem):
try:
var=np.float64(item.text())
item.setText(format(var,'.6f'))
except ValueError:
var=0
item.setText('0')
finally:
if item.column()==0:
# 引力常数
self.G=var
pass
elif item.column()==1:
# 时间微分
self.dt=var
pass
else:
# 质量设置
self.m[item.column()-2]=var
pass
else:
j=item.column()
i=item.row()
if isinstance(item,QTableWidgetItem):
try:
var=np.float64(item.text())
item.setText(format(var,'.6f'))
except ValueError:
var=0
item.setText('0')
finally:
DATA[3*j+i+1]=var
self.x0=DATA[1:10].reshape(3,3)
self.v0=DATA[10:19].reshape(3,3)
pass
# print(DATA)
pass
[文档] def parametersettings(self):
'''弹出参数设置表格对话框，执行此函数将清空原有DATA中的数据信息！'''
dialog=QDialog()
global DATA
self.t.stop()
dialog.setGeometry(self.win.geometry().left()+100,self.win.geometry().top()+100,400,300)
dialog.setFixedSize(700,300)
dialog.setWindowFlags( Qt.WindowCloseButtonHint)
dialog.setWindowTitle('参数设置')
dialog.setWindowModality(Qt.ApplicationModal)
vt=myTable(1,5)
vt.pro='const'
l=[self.G,self.dt,self.m[0],self.m[1],self.m[2]]
for i in range(5):
itemContent=l[i].__str__()
item=QTableWidgetItem(itemContent)
#为每个表格内添加数据
vt.setItem(0,i,item)
vt.itemChanged.connect(self.myfmt)
vt.setHorizontalHeaderLabels(['引力常数','时间微分','质量1','质量2','质量3'])
mt = myTable(3,6)
mt.pro='varib'
mt.itemChanged.connect(self.myfmt)
mt.setVerticalHeaderLabels(['x','y','z'])
mt.setHorizontalHeaderLabels(['x1','x2','x3','v1','v2','v3'])
if(DATA.shape[0]!=1):
DATA=np.asarray(
np.concatenate(
([0],
self.x0.flatten(),
self.v0.flatten(),
)
)
)
# print('重新初始化')
# DATA=DATA[-1]
# print(DATA)
# print(DATA[1:].reshape(6,3))
dts=DATA[1:].reshape(6,3)
for i in range(3):
for j in range(6):
itemContent=dts[j][i].__str__()
#为每个表格内添加数据
mt.setItem(i,j,QTableWidgetItem(itemContent))
lay=QGridLayout(dialog)
lay.addWidget(vt,0,0,5,1)
lay.addWidget(mt)
dialog.exec()
[文档] def changeSpeed(self,value):
'''改变图像刷新速度（注意不是时间微分dt!为获得更精确的结果请在参数设置表格中手动更改dt）
参数:
value:刷新速率'''
self.statusbar.showMessage(str(self.t.interval()))
self.label_for_slider.setText('speed='+str(value))
self.t.setInterval((int)(1000-value))
[文档] def update(self):
'''更新散点位置及轨迹'''
global DATA
# color = np.ones((self.poses.shape[0],4), dtype=np.ubyte)
# self.poses=np.random.random(size=(self.Nbodies.value(),3))*10
self.sp2.setData(pos=self.x1)
self.x1, self.v1 = timestep(self.x0, self.v0, self.G, self.m, self.dt)
self.timing+=self.dt
DATA=np.vstack(
(DATA,
np.concatenate(
(
[self.timing],
self.x1.flatten(),
self.v1.flatten(),)
)
)
)
self.line1.setData(pos=DATA[...,1:4])
self.line2.setData(pos=DATA[...,4:7])
self.line3.setData(pos=DATA[...,7:10])
# print(DATA[...,1:4])
self.x0, self.v0 = self.x1, self.v1
# print('x0 ='+str(self.x0))
# print('v0 ='+str(self.v0))
[文档] def animate_start_stop(self):
'''控制绘图的暂停和继续'''
if self.t.isActive():
self.t.stop()
else:
self.t.start(1000-self.slider.value())
[文档] def colordiallog(self):
'''弹出颜色选择框，为散点及其轨迹设置颜色'''
sender=self.win.sender()
# print(str(sender.text()))
color=QColorDialog.getColor()
id=1
if sender.text()=='调整一号天体颜色':
id=1
self.line1.color=color
elif sender.text()=='调整二号天体颜色':
id=2
self.line2.color=color
elif sender.text()=='调整三号天体颜色':
id=3
self.line3.color=color
if color.isValid():
self.points_color[id-1]=np.asarray(color.getRgbF())
self.sp2.color=self.points_color
self.statusbar.showMessage(str(id)+"选择了"+color.getRgb().__str__())
[文档] def filedialog_open(self):
'''弹出文件对话框，打开已有的csv格式文件，从中读取已有轨道信息！'''
filename = QFileDialog.getOpenFileName(self.win, "打开", "./",filter="csv(*.csv);;")[0]
if filename:
global DATA
with open(filename,encoding='utf-8') as f:
dts=np.loadtxt(f,delimiter=',',skiprows=1,max_rows=1,usecols=[0,1,2,3,4])
# print(dts)
DATA = np.loadtxt(f,delimiter = ",", skiprows = 1)#usecols=(i1,i2,i3...)
# print(DATA)
self.G=np.float64(dts[0])
self.dt=np.float64(dts[1])
self.m=np.float64(dts[2:])
self.x0=DATA[-1,1:10].reshape(3,3)
self.v0=DATA[-1,10:19].reshape(3,3)
self.update()
# print('文件数据读取完成！')
[文档] def filedialog_save(self):
'''弹出保存文件对话框，将现有DATA中的数据信息，以及各质点的质量，时间微分和引力常数信息保存至指定文件'''
fileName2, ok2 = QFileDialog.getSaveFileName(self.win,"文件保存","./","csv (*.csv)")
with open(fileName2,'w',encoding='utf-8',newline='') as f:
wr=csv.writer(f)
wr.writerow(['gravity constant G','differential time dt','Mass m1','Mass m2','Mass m3'])
wr.writerow([self.G,self.dt,self.m[0],self.m[1],self.m[2]])
wr.writerow(['time','X1_x','X1_y','X1_z','X2_x','X2_y','X2_z','X3_x','X3_y','X3_z','V1_x','V1_y','V1_z','V2_x','V2_y','V2_z','V3_x','V3_y','V3_z'])
for i in range(0,DATA.shape[0]):
wr.writerow(DATA[i])
if __name__ == '__main__':
'''执行主函数，声明并加载程序'''
sim=MainPrograme()
sim.inteval=50
sim.load()