Подзаголовки Matplotlib с различным количеством строк, проблемы с цветовой картой и легендой - PullRequest
Новая
       110

Подзаголовки Matplotlib с различным количеством строк, проблемы с цветовой картой и легендой

1 голос
/ 04 августа 2020

Я пытаюсь создать подзаговор, где каждый подзаголовок может иметь разное количество измерений (линий) для построения. Я хотел бы использовать палитру для цветов и правильно получить метки.

В приведенном ниже примере я слишком упрощаю реальные данные, с которыми я имею дело, но этого достаточно, чтобы объяснить структуру.

Вот что я пробовал до сих пор 

import collections

import matplotlib.pyplot as plt
import numpy as np

import pandas as pd

# convenience function to flatten inhomogeneous list
def flatten(l):
    for el in l:
        if isinstance(el, collections.Iterable) and not isinstance(el, (str, bytes)):
            yield from flatten(el)
        else:
            yield el

# some data
x = np.linspace(0, 2, 100)

linear = x
quadratic = x**2
cubic = x**3

noise = np.random.normal(0, 0.5, 100)
cubic_noise = x**3 + noise

# dataframe
df = pd.DataFrame({'X': x,
                   'linear': linear,
                   'quadratic': quadratic,
                   'cubic': cubic,
                   'cubic_noise': cubic_noise})

# dictionary with dataframe columns to plot
dict_ = {'Linear' : 'linear',
         'Quadratic' : 'quadratic',
         'Cubic_and_CubicNoise' : ['cubic', 'cubic_noise']}

keys = list(dict_.keys())
flat_keys = list(flatten(list(dict_.values())))

colourmap = plt.cm.viridis
colours = [colourmap(i) for i in np.linspace(0.05, 0.95, len(flat_keys))]

# do the subplots in a loop
fig, ax = plt.subplots(len(dict_), 1, sharex=True)

xcol = 'X'

for i, a, key in zip(range(len(keys)), ax, keys):

    ycols = dict_[key]

    a.grid(b=True, which='major', color='#666666', linestyle='-')

    a.set_title(key, fontsize=12)
    a.plot(df[xcol],
            df[ycols],
            marker='.',
            label=ycols,
            c=colourmap((i+1) / float(len(flat_keys))))

    a.legend()

plt.xlabel(xcol, fontsize=14)

title_ = 'Subplots with different dimensions to plot'
fig.suptitle(title_, fontsize=14)

fig.set_size_inches(15, 8)

plt.show()

Это дает:

участок

Почти готово, но в диаграмме с несколькими линиями я не могу получить разные цвета (из цветовой карты), и легенда получает список, а не метку. Есть идеи, как это улучшить? Большое спасибо.

1 Ответ

1 голос
/ 07 августа 2020

Я думаю, что нашел решение, определив все списки в столбцах для построения словаря и используя set_prop_cycle внутри l oop

import collections

import numpy as np
import pandas as pd

import matplotlib.pyplot as plt

# convenience function to flatten inhomogeneous list
def flatten(l):
    for el in l:
        if isinstance(el, collections.Iterable) and not isinstance(el, (str, bytes)):
            yield from flatten(el)
        else:
            yield el

# some data
x = np.linspace(0, 2, 100)

linear = x
quadratic = x**2
cubic = x**3

noise = np.random.normal(0, 0.5, 100)
cubic_noise = x**3 + noise

# dataframe
df = pd.DataFrame({'X': x,
                   'linear': linear,
                   'quadratic': quadratic,
                   'cubic': cubic,
                   'cubic_noise': cubic_noise})

# dictionary with dataframe columns to plot
dict_ = {'Linear' : ['linear'],
         'Quadratic' : ['quadratic'],
          'Cubic_and_CubicNoise' : ['cubic', 'cubic_noise']}

keys = list(dict_.keys())
flat_keys = list(flatten(list(dict_.values())))

colourmap = plt.cm.viridis

# do the subplots in a loop
fig, ax = plt.subplots(len(dict_), 1, sharex=True)

xcol = 'X'

for a, key in zip(ax, keys):
    print(key)
    ycols = dict_[key]

    idx = [flat_keys.index(i) for i in ycols if i in  flat_keys]
    colours = [colourmap(1.*(i+1)/len(flat_keys)) for i in idx]

    a.grid(b=True, which='major', color='#666666', linestyle='-')

    a.set_title(key, fontsize=12)
    a.set_prop_cycle('color', colours)

    a.plot(df[xcol],
            df[ycols],
            marker='.')

    a.legend(ycols)

plt.xlabel(xcol, fontsize=14)

title_ = 'Subplots with different dimensions to plot'
fig.suptitle(title_, fontsize=14)

fig.set_size_inches(15, 8)

plt.show()

участок

...