如何在Matplotlib极坐标图中根据theta值改变环线颜色

问题描述 #

我想要绘制麦克风的指向性图案，希望根据theta值为线条设置不同颜色：当theta大于90度且小于270度时，线条为红色；其余情况为蓝色。尝试了如下代码但无法实现需求，因为每次调用line.set_color()都会改变整条线的颜色：

theta = np.arange(0,2*np.pi,0.01)
r = abs(0.0 + 1*np.cos(theta))


fig, ax = plt.subplots(subplot_kw={'projection': 'polar'})
line, = ax.plot(theta,r)

for t in theta:
    if t > np.pi/2 and t < 3*np.pi/2:
        line.set_color('red')  
    else:
        line.set_color('blue')

解决方法 #

方法1：拆分数据段分别绘制（推荐，保持线条连续性） #

利用numpy的布尔索引拆分theta和r数据，将符合两种颜色条件的数据分别调用plot绘制，两段线条会保持连续且各自对应目标颜色：

import numpy as np
import matplotlib.pyplot as plt

theta = np.arange(0, 2*np.pi, 0.01)
r = abs(0.0 + 1*np.cos(theta))

fig, ax = plt.subplots(subplot_kw={'projection': 'polar'})

# 生成颜色筛选的布尔掩码
mask_red = (theta > np.pi/2) & (theta < 3*np.pi/2)
mask_blue = ~mask_red

# 分别绘制蓝色和红色线条
ax.plot(theta[mask_blue], r[mask_blue], color='blue', linewidth=2)
ax.plot(theta[mask_red], r[mask_red], color='red', linewidth=2)

plt.show()

方法2：使用scatter逐点着色（适合精细控色场景） #

如果需要对每个点单独设置颜色，可以用scatter函数，通过c参数传入颜色数组实现。调整点的大小和边框宽度可以近似线条效果：

import numpy as np
import matplotlib.pyplot as plt

theta = np.arange(0, 2*np.pi, 0.01)
r = abs(0.0 + 1*np.cos(theta))

fig, ax = plt.subplots(subplot_kw={'projection': 'polar'})

# 根据条件生成对应颜色的数组
colors = np.where((theta > np.pi/2) & (theta < 3*np.pi/2), 'red', 'blue')

# 绘制带颜色的散点，调整参数模拟线条效果
ax.scatter(theta, r, c=colors, s=6, linewidths=0.3, edgecolors=colors)

plt.show()

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