Как я могу синхронизировать цвета для каждой итерации в al oop, которая строит последовательные кривые в Matplotlib?

Question

У меня есть следующий код для построения решения для нескольких значений демпфированного осциллятора:

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))

(см. Рисунок ниже)

Проблема в том, что Pyplot кажется назначая цвета для 3 кривых в каждой итерации (т.е. каждое значение mu) случайным образом, но я хочу, чтобы цвета были "синхронизированы", так сказать, для первого значения mu, кривые для np.exp(-(mu/500), -np.exp(-(mu/500) * x) и np.sin(np.sqrt(1 - ((mu/100)**2)) * x)*np.exp(-(mu/500) * x) должны быть одинаковыми (или близкими).

Я вроде как взломал его, установив цвет на каждой итерации, используя кортеж цвета RGB, и затем каждый раз увеличивая значения в кортеже, но это утомительно, и вам нужно t ie приращения до количество итераций, через которое l oop пройдет go?

Существует ли более элегантный и лаконичный способ сделать это, который может работать с любым количеством итераций?

Answer 1

В дополнение к ответам в комментарии @pink spikyhairman, вы можете определить цветовую карту и получить цвета, введя в нее mu/mumax:

# 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)

Answer 2

Из документации matplotlib.pyplot.plot :

По умолчанию каждой строке назначается отдельный стиль, заданный «циклом стиля». Параметры свойств fmt и line необходимы только в том случае, если вы хотите явных отклонений от этих значений по умолчанию. Кроме того, вы также можете изменить цикл стиля, используя 'axes.prop_cycle' rcParam.

Вы можете создать cycler объект с повторяющимися параметрами для цветов и передать его объекту осей через set_prop_cycle например:

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))