I am trying to plot a 2D line plot of two arrays with color as a function of a third array. My question is similar to this one: matplotlib: varying color of line to capture natural time parameterization in data I have successfully tried this method, below is a version of my current code:
import matplotlib.pyplot as plt
import matplotlib.colors as clr
import matplotlib.collections as col
import numpy as np
def XYPlot(A):
"""
Plots XY Plot of given trajectory, with color as a function of p
"""
x = A[0][A[7] != 0]
y = A[1][A[7] != 0]
p = A[7][A[7] != 0]
points = np.array([x,y]).T.reshape(-1,1,2)
segments = np.concatenate([points[:-1], points[1:]], axis=1)
lc = col.LineCollection(segments, cmap=plt.get_cmap('RdYlBu'),
norm=plt.Normalize(200, 1000))
lc.set_array(p)
lc.set_linewidth(1.5)
plt.gca().add_collection(lc)
plt.xlim(-25,25)
plt.ylim(-25,25)
plt.show()
pass
I was wondering, however, if there is a more convenient way of doing it, rather than working with the LineCollection.
In any case, I need to add a colorbar of that third variable to my plot. How do I create one?
Thanks!