Matplotlib subplots with different number of lines issues with colormap and legend

Question

I'm trying to do a subplot where each subplot could have a different number of dimensions (lines) to plot. I'd like to use the colormaps for colors and get the labels properly.

In the example below I'm oversimplifying the real data I'm dealing with but it's good enough to explain the structure.

Here's what I tried so far:

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

This gives:

Almost there, but in the chart with multiple lines I'm not able to get different colors (from colormap) and the legend gets the list and not the label. Any ideas on how to improve this?

Answer 1

I think I found a solution by defining all lists in the columns to plot dictionary and using set_prop_cycle inside the loop

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