matplotlib绘制RSI折线按Y值变色触发ValueError系列真值模糊问题

报错原因 #

你代码中的rsi2是Pandas Series类型，本质是包含所有RSI计算结果的一维数组，直接将整个数组放入if/while的条件判断中时，Python无法判断是要所有元素满足条件、还是任意一个元素满足条件，因此抛出The truth value of a Series is ambiguous报错。

PineScript中的着色逻辑是逐元素判断：每个K线对应的RSI值单独判断所属区间分配颜色，Python中需要对齐这个逻辑，同时matplotlib原生plot方法不支持单条折线多色渲染，需要用LineCollection实现分段着色。

修复后完整代码 #

import matplotlib.pyplot as plt
import pandas as pd
import numpy as np
from matplotlib.collections import LineCollection

# 原有读取数据、计算RSI、topl/lowl/rsi2的逻辑保持不变
df = pd.read_csv('BTC1DRSI.csv')
x = df['time']
# 如果x是时间字符串，先转成数值格式避免LineCollection报错，取消下一行注释即可
# x = pd.to_datetime(x).apply(pd.Timestamp.toordinal)
rsi = df['RSI']
xx = np.array(range(0,4086))

# TOP LINE; coordinates
p21 = [2, 86]
p22 = [4081, 77.5]
slope2 = (p22[1] - p21[1]) / (p22[0] - p21[0])
topl = slope2 * xx + p21[1]

# BOTTOM LINE; coordinates
p11 = [2, 44.4]
p12 = [4081, 36.6]
slope1 = (p12[1] - p11[1]) / (p12[0] - p11[0])
lowl = slope1 * xx + p11[1]    

# FIX to 0 - 100
rang = topl - lowl
yyy = rsi - lowl
rsi2 = (yyy / rang) *100
# 去掉空值避免报错
rsi2 = rsi2.fillna(0)

# 区间划分，和PineScript逻辑完全对齐
totalr= 100 - 0
lr= totalr/10    
bins = [0, lr*1, lr*2, lr*3, lr*4, lr*5, lr*6, lr*7, lr*8, lr*9, 100]
# 颜色从低区间到高区间，完全匹配PineScript的色值
colors = [
    '#0000ff', '#00acff', '#00faff', '#00ff69', '#00ff00',
    '#ffff00', '#ffcb00', '#ff8a00', '#ff4100', '#ff0000'
]

# 给每个rsi2值分配对应颜色
color_labels = pd.cut(rsi2, bins=bins, labels=colors, include_lowest=True)

# 构造分段线段数组
points = np.array([x, rsi2]).T.reshape(-1, 1, 2)
segments = np.concatenate([points[:-1], points[1:]], axis=1)

# 创建LineCollection对象
lc = LineCollection(segments, colors=color_labels, linewidth=3)

# 绘图部分
fig, ax = plt.subplots()
# 添加彩色折线
ax.add_collection(lc)
# 绘制区间横线
for s in np.arange(0, 101, 10):
    ax.axhline(y=s, color='silver', linestyle='-')
# 自动适配坐标轴范围
ax.autoscale()
plt.tight_layout()
plt.show()

关键逻辑说明 #

• 用pd.cut实现逐元素的区间匹配，完全对齐PineScript从低到高的层级判断逻辑
• LineCollection将整条折线拆分为相邻两点组成的独立线段，每个线段单独设置颜色，实现和PineScript完全一致的着色效果

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