2

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!

Community
  • 1
  • 1
meteosteph
  • 143
  • 1
  • 2
  • 9

0 Answers0