How can I synch up the colors for each iteration in a loop that plots successive curves in Matplotlib?

Question

I have the following code to plot the solutions to multiple values of a damped oscillator:

import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
sns.set()
x = np.linspace(0, 50, 1000)
plt.figure(figsize=(9,7.5))
for mu in range(40,100,14):
  plt.plot(x, np.exp(-(mu/500) * x), linestyle='--',alpha=0.4)
  plt.plot(x, -np.exp(-(mu/500) * x), linestyle='--',alpha=0.4)
  plt.plot(x, np.sin(np.sqrt(1 - ((mu/100)**2)) * x)*np.exp(-(mu/500) * x))

(See below for output figure)

The problem is that Pyplot seems to be assigning the colors for the 3 curves in each iteration (ie each value of mu ) randomly, but I want the colors to be "in synch" so to speak, that is for the first value of mu , the curves for np.exp(-(mu/500) , -np.exp(-(mu/500) * x) , and np.sin(np.sqrt(1 - ((mu/100)**2)) * x)*np.exp(-(mu/500) * x) to be the same (or close).

I kind of hacked it by setting the color through each iteration using the RGB color tuple, and then incrementing the values in the tuple each time, but that's tedious, and you have to tie the increments to the number of iterations your loop will go through?

Is there a more elegant and concise way of doing it that can work with any number of iterations?

Answer 1

From the matplotlib.pyplot.plot documentation :

By default, each line is assigned a different style specified by a 'style cycle'. The fmt and line property parameters are only necessary if you want explicit deviations from these defaults. Alternatively, you can also change the style cycle using the 'axes.prop_cycle' rcParam.

You can create a cycler object with repeating parameters for colors and pass it to to the axes object via set_prop_cycle for example:

import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
from cycler import cycler
sns.set()
x = np.linspace(0, 50, 1000)
fig, ax = plt.subplots(1, figsize=(9,7.5))
ax.set_prop_cycle(cycler('color', 'rrrbbb'))
for mu in range(40,40+28,14):
    ax.plot(x, np.exp(-(mu/500) * x), linestyle='--',alpha=0.4)
    ax.plot(x, -np.exp(-(mu/500) * x), linestyle='--',alpha=0.4)
    ax.plot(x, np.sin(np.sqrt(1 - ((mu/100)**2)) * x)*np.exp(-(mu/500) * x))

Answer 2

As an addition to the answers in @pink spikyhairman's comment, you can define a colormap and get colors by inputing mu/mumax into it:

# define mus outside of loop
mus = np.arange(40,100,14)
cmap = plt.cm.viridis

for mu in mus:

  # define color as colormap entry between 0 and 1
  color = cmap(mu/np.max(mus))  

  plt.plot(x, np.exp(-(mu/500) * x), linestyle='--',alpha=0.4, c=color)
  plt.plot(x, -np.exp(-(mu/500) * x), linestyle='--',alpha=0.4, c=color)
  plt.plot(x, np.sin(np.sqrt(1 - ((mu/100)**2)) * x)*np.exp(-(mu/500) * x), c=color)