Plot Multicolored line based on conditional in python but end up with fragmented line

user_13530226 0
Viewed 133
import pandas as pd
import matplotlib.pyplot as plt
trend=[0,0,0,1,1,1,-1,-1,-1,1,1,1,0,0,0,1,1,1,-1,-1]
price= [1,2,3,4,5,6,5,4,3,4,5,6,7,8,9,8,7,6,5,4]
df1= pd.DataFrame({'Trend':trend,'Price':price})

def plot_func(group):
    global ax
    if (group.Trend ==-1).all() :
        color = 'r'
    elif (group.Trend ==1).all() :
        color ='g'
    elif (group.Trend == 0).all() :
        color='b'
    lw = 2.0
    ax.plot(group.index, group.Price, c=color, linewidth=lw) 

fig, ax = plt.subplots()
df1.groupby((df1.Trend.shift() != df1.Trend).cumsum()).apply(plot_func)

Hi guys,

Need some advice on my code. I tried to plot the price data using matplotlib with multicoloured lined based on the trend condition but somehow the lines are disconnected and I have no idea how to make them continuous. Pls help.

Thks a lot, John

My plot

matplotlib python
1 Answers

This is because the groups are separated and your code only plots within each group. Try (similar to the link in the comment):

for trend, prev, cur in zip(df1.Trend.iloc[1:], df1.index[:-1], df1.index[1:]):
    if trend==-1: 
        color='r'
    elif trend==1:
        color='g'
    else:
        color='b'
    
    plt.plot([prev,cur],df1.loc[[prev,cur],'Price'], c=color)

Output:

enter image description here

user_42384080
Related