Как реализовать Linecollection в dataframe, чтобы закрасить одну строку другим цветом?

Question

Мне нужно построить участки с разноцветными линиями в соответствии с условием. Для этого я стараюсь использовать LineCollection. Но с моими данными я получаю ошибки и не могу найти способ использовать их с моими данными. Я собираюсь с этим стеком, может кто-нибудь подсказать мне, как заставить его работать.

Мои данные выглядят следующим образом:

age_gps_data ref_station_id river_km    
1.0     2421.0          667.5144925407869   
0.5     2421.0          667.5144592533758   
1.0     2421.0          667.5144249505418   
0.5     2421.0          667.5143958295257   
1.0     2421.0          667.5143629441299   
0.5     2421.0          667.5143266152246   
1.0     2.0     667.5142970594003   
0.5     2.0     667.5142580343961   
1.0     2.0     667.5142211073334   
0.5     2.0     667.5141878187346   

Я пробую это решение, но у меня есть ошибка:

 filename = "G:\\ais_evaluation\\track_AIS_2route.csv"
 df2 = pd.read_csv(filename, delimiter=';')

 segments = []
 color = np.zeros(shape=(10,4))
 x = df2['river_km'].tolist()
 y = df2['age_gps_data'].tolist()

 i = 0
 z = df2.ref_station_id

 for x1, y1, x2, y2 in zip(x, x[1:], y, y[1:]):

    if z == 2.0:
       color[i] = colors.to_rgba('Crimson')
    else:
        color[i] = colors.to_rgba('slategray')
    segments.append([(x1, y1), (x2, y2)])
    i += 1 

 lc = mc.LineCollection(segments, colors=color, linewidths=2)
 fig, ax = pl.subplots()
 ax.add_collection(lc)
 ax.autoscale()
 ax.margins(0.1)
 pl.show()     

У меня есть эта ошибка:

 ValueError: The truth value of a Series is ambiguous. Use a.empty, a.bool(), a.item(), a.any() or a.all(

Как мне улучшить этот код, чтобы получить строку river_km против age_gps_data с разными цветами в соответствии с параметром ref_station_id? Я новичок в Python, и я был бы признателен за любую помощь.

Answer 1

Ваша проблема исходит от линии if z == 2.0:. Здесь z - это серия панд df2.ref_station_id), и поэтому она не может быть равна 2,0.

Я не совсем следую логике вашего кода, поэтому не могу точно сказать вам, что делать, но выглядитнапример, вам нужно перебирать значения z одновременно с x и y и проверять, равны ли последующие значения z, а не z 2,0

дляНапример:

for x1, y1, x2, y2, z0 in zip(x, x[1:], y, y[1:], z):
    if z0 == 2.0:
    (...)

РЕДАКТИРОВАТЬ Вот целый код функции:

import matplotlib.pyplot as plt
from matplotlib import colors
from matplotlib import collections as mc
from io import StringIO

d = """age_gps_data ref_station_id river_km    
1.0     2421.0          667.5144925407869   
0.5     2421.0          667.5144592533758   
1.0     2421.0          667.5144249505418   
0.5     2421.0          667.5143958295257   
1.0     2421.0          667.5143629441299   
0.5     2421.0          667.5143266152246   
1.0     2.0     667.5142970594003   
0.5     2.0     667.5142580343961   
1.0     2.0     667.5142211073334   
0.5     2.0     667.5141878187346   
"""
df2 = pd.read_csv(StringIO(d), delimiter='\\s+')

segments = []
color = np.zeros(shape=(10,4))
x = df2['river_km'].tolist()
y = df2['age_gps_data'].tolist()

i = 0
z = df2.ref_station_id

for x1, y1, x2, y2, z0 in zip(x, y, x[1:], y[1:], z):
    if z0 == 2.0:
        color[i] = colors.to_rgba('Crimson')
    else:
        color[i] = colors.to_rgba('slategray')
    segments.append([(x1, y1), (x2, y2)])
    i += 1 

lc = mc.LineCollection(segments, colors=color, linewidths=2)
fig, ax = plt.subplots()
ax.add_collection(lc)
ax.autoscale()
ax.margins(0.1)
plt.show()