Assign specific colors to array columns without explicit iteration when plotting in matplotlib

Question:

It is helpful sometimes to do plt.plot(x, y) when y is a 2D array due to every column of y will be plotted against x automatically in the same subplot. In such a case, line colors are set by default. But is it possible to customize colors with something similar to plt.plot(x, y, color=colors) where now colors is an iterable?

For example, let’s say I have three datasets that scatter around straight lines and want to plotted with fitting curves in such a way that each dataset and its fit share the same color.

np.random.seed(0)

# fake dataset
slope = [1, 2, 3]
X = np.arange(10) 
Y = slope * X[:,None] + np.random.randn(10,3)

# fitting lines
params = np.polyfit(X, Y, deg=1)
x = np.linspace(0, 10, 50)
y = np.polyval(params, x[:,None])

I would like to get the ouput of the following code without having to iterate manually.

colors = ['b', 'r', 'g']
for i in range(3):
    plt.plot(X, Y[:,i], '.', color=colors[i])
    plt.plot(x, y[:,i], color=colors[i])

enter image description here

Asked By: alpelito7

||

Answers:

enter image description here

"reset" the properties cycler

...
plt.plot(x, Y, '.')
plt.gca().set_prop_cycle(None)
plt.plot(x, y)
...

For an explanation, see my answer to another question.

Answered By: gboffi
Categories: questions Tags: , ,
Answers are sorted by their score. The answer accepted by the question owner as the best is marked with
at the top-right corner.