python - 如何绘制条件为 (r < 1.5*rs) 的数组？

Question

我正在对绕黑洞运行的光子进行编码。当光子直接进入黑洞时，由于被零除，其轨迹非常奇怪。我想忽略 r < 1.5*rs 数组中的光子，但我不知道如何实现

我尝试过使用 while True 和 if，但没有成功

h0=[t0, r2, theta, phi2, pt2, pr2, ptheta, pphi2] 

T = np.linspace(0, 1000, 9000)

zz=odeint(func, h0, T, args=(rs,))

r22 = zz[:, 1]
theta22 = zz[:, 2]
phi22 = zz[:, 3]
pt22 = zz[:, 4]
pr22 = zz[:, 5]
pphi22 = zz[:, 7]

def sph2cart(r, phi, theta):
    X = r * np.cos(phi) * np.sin(theta)
    Y = r * np.sin(phi) * np.sin(theta)
    Z = r * np.cos(theta)
    return(X, Y, Z)

X2, Y2, Z2 = sph2cart(r22, phi22, theta22)

plt.plot(X2, Y2, Z2, 'g')

我认为你真的不需要代码来帮助我，但是有人知道如何绘制半径 r < 1.5*rs 的 X2、Y2、Z2 (rs 在代码中定义)吗？

Answer 1

Numpy 内置了运行良好的掩码

r=np.arange(20)
r_mask = np.ma.masked_where(r < 10 , r)

所以，因为你只想在 r<1.5*rs 时生成 X、Y 和 Z 值(我认为你在某一点上以相反的方式说过，但在任何情况下你只需要翻转符号)，您的代码可能如下所示

import numpy as np
h0=[t0, r2, theta, phi2, pt2, pr2, ptheta, pphi2] 

T = np.linspace(0, 1000, 9000)

zz=odeint(func, h0, T, args=(rs,))

r22 = zz[:, 1]
theta22 = zz[:, 2]
phi22 = zz[:, 3]
pt22 = zz[:, 4]
pr22 = zz[:, 5]
pphi22 = zz[:, 7]

def sph2cart(r, phi, theta):
    X = r * np.cos(phi) * np.sin(theta)
    Y = r * np.sin(phi) * np.sin(theta)
    Z = r * np.cos(theta)
    return(X, Y, Z)

r22_masked = np.ma.masked_where(r22 > 1.5*rs , r22)

X2, Y2, Z2 = sph2cart(r22_masked, phi22, theta22)

这只会为 r22 的未屏蔽值生成 X2、Y2 和 Z2 值(即 r22 小于 1.5*rs)