使用Numba加速引力模拟计算时Jupyter内核频繁崩溃重启问题求助
2026-9-26
引力计算Numba加速后内核崩溃问题排查
问题背景
使用Anaconda开展简单引力计算工作，为提升运算速度对原有代码进行了改造：将基于Python列表实现的逻辑全部改写为NumPy数组运算，再通过@jit/@njit装饰器调用Numba加速。短时长计算可以正常运行，但执行长时长计算或与其他代码联用时，会出现内核崩溃后自动重启的问题。
依赖模块
%matplotlib notebook import numpy as np from matplotlib import pyplot as plt import pandas as pd from timeit import default_timer as timer from matplotlib import colors from matplotlib.widgets import Slider, Button import re import os from numba import jit, njit
原版Python列表实现引力计算代码
def verlet_int(u,p,v,t_set,N,k,d,p_num): t = t_set[0] T = t_set[1] dt = t_set[2] p_prime = [] v_prime = [] o = [] a_ik = lambda u,p,k,i,j : (-1)*u[k]*((((p[i][0] - p[k][0])**2) + ((p[i][1] - p[k][1])**2) + ((p[i][2] - p[k][2])**2))**(-1.5))*(p[i][j] - p[k][j]) while t <= T: a_1 = [] a_2 = [] for i in np.arange(N): b_1 = [] b_2 = [] for j in [0,1,2]: if p_num == -1: b_1.append(p[i][j]) b_2.append(v[i][j]) elif i == p_num: a_1.append(p[i][j]) a_2.append(v[i][j]) v_half = v[i][j] + 0.5*sum([a_ik(u,p,k,i,j) for k in np.delete(np.arange(N),i)])*dt p[i][j] = p[i][j] + v_half*dt v[i][j] = v_half + 0.5*sum([a_ik(u,p,k,i,j) for k in np.delete(np.arange(N),i)])*dt if p_num == -1: a_1.append(b_1) a_2.append(b_2) elif i == p_num: pass if round(t/dt)%k == 0: p_prime.append(a_1) v_prime.append(a_2) else: pass t += dt t = np.round(t, decimals = d) o.append(p_prime) o.append(v_prime) return o
代码输出格式说明：o[位置(0)或速度(1)][选中的记录时间步][粒子编号0...N-1][x(0)、y(1)、z(2)]，当p_num = -1时记录所有粒子数据，p_num取值为0...N-1时仅记录对应编号粒子的位置和速度。
改造后带Numba JIT的代码
@jit(debug = True) def verlet_int_numpy(u,p_v_1,t_i,T,dt,k,dec,p_num): times = np.arange(t_i,T,dt) p_v = np.copy(p_v_1) if p_num == -1: storage = np.zeros((round(times.size/k),2,len(u),3)) else: storage = np.zeros((round(times.size/k),2,1,3)) for tau in range(times.size): if tau%k == 0: if p_num == -1: storage[tau] = p_v else: storage[tau][0][0] = p_v[0][p_num] storage[tau][0][1] = p_v[1][p_num] else: pass for j in range(u.size): a_g = np.array([0.0,0.0,0.0]) for i in np.delete(np.arange(u.size),j): a_g = a_g + u[i]*(((p_v[0][i] - p_v[0][j])@(p_v[0][i] - p_v[0][j]))**(-1.5))*(p_v[0][i] - p_v[0][j]) v_half = p_v[1][j] + 0.5*a_g*dt p_v[0][j] = p_v[0][j] + v_half*dt a_g = np.array([0.0,0.0,0.0]) for i in np.delete(np.arange(u.size),j): a_g = a_g + u[i]*(((p_v[0][i] - p_v[0][j])@(p_v[0][i] - p_v[0][j]))**(-1.5))*(p_v[0][i] - p_v[0][j]) p_v[1][j] = v_half + 0.5*a_g*dt return storage
CSV写入工具函数
def setup_csv(i_s,i_v,strings): r_cm = (1/np.sum(i_s[0]))*np.dot(i_s[0],i_v[0]) v_cm = (1/np.sum(i_s[0]))*np.dot(i_s[0],i_v[1]) i_v[0] = i_v[0] - r_cm i_v[1] = i_v[1] - v_cm O = verlet_int_numpy(np.array(i_s[0]),i_v,i_s[1],i_s[2],i_s[3],i_s[6],i_s[4],i_s[7]) if i_s[7] == -1: for j in np.arange(i_s[0].size): D_1 = pd.DataFrame(data = O[:,0,j]) D_2 = pd.DataFrame(data = O[:,1,j]) D_1.to_csv(r"/Volumes/Storage/Physics Simulations and Research/Gravitational Research/D_test/" + "particle_{}_{}_{}_".format(j,i_s[6],i_s[3]) + strings[0], mode = 'w', header = False, index = None) D_2.to_csv(r"/Volumes/Storage/Physics Simulations and Research/Gravitational Research/D_test/" + "particle_{}_{}_{}_".format(j,i_s[6],i_s[3]) + strings[1], mode = 'w', header = False, index = None) else: D_1 = pd.DataFrame(data = O[:,0,i_s[7]]) D_2 = pd.DataFrame(data = O[:,1,i_s[7]]) D_1.to_csv(r"/Volumes/Storage/Physics Simulations and Research/Gravitational Research/D_test/" + "particle_{}_{}_{}_".format(i_s[7],i_s[6],i_s[3]) + strings[0], mode = 'w', header = False, index = None) D_2.to_csv(r"/Volumes/Storage/Physics Simulations and Research/Gravitational Research/D_test/" + "particle_{}_{}_{}_".format(i_s[7],i_s[6],i_s[3]) + strings[1], mode = 'w', header = False, index = None)
原版参数说明
#p_pick gives the particle to be recorded from orbit_int_verlet(). It takes values 0,1,...,N-1. It can also take on #. . . -1 for recording all the possible particle positions/velocites. #s_scale gives the magnitude of the random unit vector. #t_i is the initial time. #Usually assume t_o = 0. #dt is the time step. #T = total length of time. #k decides the number of stored values by whether it evenly divides the number of time steps gone by. #u are the 'N' number of masses for the bodies. #d_places is the number of decimal places to round the time step every loop so as to not accumlate much numerical error. #i_s = inital_scalars = [u,t_i,T,dt,d_places,s_scale,k,p_pick] #strings = [] . . . set of string titles. The first two will be the comparison orbits; position then velocity. #. . . the last two will be where the output files for position then velocity will go #i_v = initial_vectors = np.array([ [[],[],...,[]], [[],[],...,[]] ]) this is 2xNx3 array.
Numba版本的改造点为将参数u也改为NumPy数组而非集合类型。
初始测试条件
u = np.array([0.97,0.02,0.01]) orbits = 0.05 dec = 3 T = round( (2*np.pi*(80**(1.5)))/(np.sqrt(0.98))*orbits) dt = 0.005 p_num = -1 s_scale = 0.1 k = 950 t_i = 0 i_s = [u,t_i,T,dt,dec,s_scale,k,p_num] i_v = np.array([[[20,0,0],[-40,0,0],[80,0,0]],[[0,0,0],[0,np.sqrt(0.85/40),0],[0,np.sqrt(0.85/80),0]]]) strings = ["comparison_p.csv","comparison_v.csv"]
报错信息
代码运行后短时间内就会弹出内核已死亡、将自动重启的提示，JIT编译过程无异常报错，无法定位问题原因。此前运行曾抛出如下错误：
--------------------------------------------------------------------------- AttributeError Traceback (most recent call last) /opt/anaconda3/lib/python3.7/site-packages/IPython/core/async_helpers.py in _pseudo_sync_runner(coro) 66 """ 67 try: ---> 68 coro.send(None) 69 except StopIteration as exc: 70 return exc.value /opt/anaconda3/lib/python3.7/site-packages/IPython/core/interactiveshell.py in run_cell_async(self, raw_cell, store_history, silent, shell_futures, transformed_cell, preprocessing_exc_tuple) 3166 # Write output to the database. Does nothing unless 3167 # history output logging is enabled. -> 3168 self.history_manager.store_output(self.execution_count) 3169 # Each cell is a *single* input, regardless of how many lines it has 3170 self.execution_count += 1 AttributeError: 'ZMQInteractiveShell' object has no attribute 'history_manager'
解决方向
- 修复数组越界逻辑：当前
storage数组长度为round(times.size/k)，但赋值时直接用storage[tau] = p_v，当tau大于等于storage长度时会触发内存越界，Numba不会抛出Python层异常，直接导致内核崩溃。修正方案：存储时使用save_idx = tau // k作为storage的第一维索引，替换当前的tau索引。 - 增加数值保护逻辑：当两个粒子距离过近时，
(距离)**(-1.5)会生成无穷大值，触发Numba运行时异常。建议在计算加速度前增加软ening因子，公式改为(距离**2 + eps**2)**(-1.5)，eps取1e-6即可，避免数值爆炸。 - 修复Numba使用逻辑：删除
debug=True参数，改为@njit(fastmath=True)，关闭调试模式避免内存泄漏，同时删除循环内的np.delete调用，直接用if i != j判断跳过当前粒子，减少临时数组生成带来的内存碎片化。 - 修复Jupyter环境兼容性问题：抛出的
AttributeError属于旧版本IPython和jupyter_client的兼容性bug，执行conda update ipython jupyter_client升级对应依赖即可，也可以直接在纯Python脚本环境下运行计算，规避Jupyter内核不稳定问题。 - 优化IO逻辑：长时运行时直接写入外接硬盘容易因IO阻塞、挂载异常触发崩溃，建议先将结果写入本地磁盘，计算完成后再转移到外接存储，同时增加IO异常捕获逻辑。
内容的提问来源于stack exchange，提问作者The victorious truther
相关产品推荐
相关产品推荐
