请求：在太阳磁场流线图上叠加磁压Pm=B²/(8π)箭头可视化

需求说明 #

已实现太阳磁场流线图绘制，现需在同一张图上叠加磁压 ( P_m = \frac{B^2}{8\pi} ) 的箭头可视化效果（用箭头的颜色、长度映射磁压数值）。

修改后的完整代码 #

# 导入必要模块
import numpy as np
import matplotlib.pyplot as plt

# 定义磁场函数
def magnetic_field(x, y):
    Bx = y
    By = x
    return np.array([Bx, By])

# 设置坐标网格
x = np.linspace(-1, 1, 20)
y = np.linspace(-1, 1, 20)
X, Y = np.meshgrid(x, y)

# 计算每个网格点的磁场
Bx, By = magnetic_field(X, Y)

# 计算磁压 Pm = B²/(8π)
B_magnitude = np.sqrt(Bx**2 + By**2)
Pm = B_magnitude**2 / (8 * np.pi)

# 设置绘图尺寸
plt.figure(figsize=(8, 8))

# 绘制磁场流线图（浅灰色避免遮挡箭头）
plt.streamplot(X, Y, Bx, By, broken_streamlines=False, color='lightgray')

# 叠加磁压箭头可视化：用径向箭头+颜色/长度表示磁压大小
# 计算径向方向向量（避免原点除以0）
u = X / np.sqrt(X**2 + Y**2 + 1e-8)
v = Y / np.sqrt(X**2 + Y**2 + 1e-8)
# 绘制箭头，颜色映射磁压，scale控制箭头长度
q = plt.quiver(X, Y, u, v, Pm, scale=20, cmap='viridis', alpha=0.7)

# 设置坐标轴范围
plt.xlim(-0.5, 0.5)
plt.ylim(-0.5, 0.5)

# 添加标签、标题和磁压颜色条
plt.xlabel('x')
plt.ylabel('y')
plt.title('太阳磁场流线图与磁压分布')
plt.colorbar(q, label='磁压 $P_m = \\frac{B^2}{8\\pi}$')

# 显示图像
plt.show()

关键修改说明 #

• 磁压计算：先求磁场模长 ( |B| = \sqrt{B_x^2 + B_y^2} )，再代入公式得到磁压 ( P_m )
• 箭头可视化：用plt.quiver()实现，选择径向统一箭头方向，通过颜色（映射磁压数值）和长度（scale参数调整）双重展示磁压大小
• 图层优化：先绘制浅灰色流线图，再叠加箭头层，保证视觉层次清晰
• 细节修复：计算径向向量时加入1e-8，避免原点处除以0的报错

可选：用箭头方向表示磁压梯度 #

如果需要箭头方向代表磁压增长最快的方向（梯度方向），可替换箭头绘制部分代码：

# 计算磁压的x、y方向梯度
dPm_dx, dPm_dy = np.gradient(Pm, x, y)
# 箭头方向为梯度方向，颜色/长度表示磁压大小
q = plt.quiver(X, Y, dPm_dx, dPm_dy, Pm, scale=50, cmap='viridis', alpha=0.7)

内容的提问来源于stack exchange，提问作者Stef Athan