[FIXED] Can I plot a colorbar from a list of colors and for loop?

Issue

Here's part of the Python code I'm working on:

cm = LinearSegmentedColormap.from_list('defcol', ["#000000", "#FF0000"])
trace_color = cm(np.linspace(0,1,cycles))
for k, color in zip(range(cycles),trace_color):
        lis = KL_rest[k::cycles]
        plt.scatter(scanpoints, lis, color = color, marker = '^', alpha = 0.9)

Here I am generating the scatter plot using a for loop, and the colors come from the list trace_color. My question is if I can generate a colorbar, where the colors are from the trace_color and labels (scale) on the colorbar come from range(cycles). I tried to add plt.colorbar() after the for loop but that didn't work. Thanks!!

Solution

Matplotlib's colorbar needs a ScalarMappable object. By default, it is taken from what's plotted, e.g. a scatter plot that is created in one call. If you need to combine the results of multiple calls, you can create an own ScalarMappable. It needs a colormap and a norm.

from matplotlib import pyplot as plt
from matplotlib.colors import LinearSegmentedColormap
from matplotlib.cm import ScalarMappable
import numpy as np

cycles = 7
N = 100
KL_rest = np.sin(np.linspace(0, 10 * np.pi, cycles * N))
scanpoints = np.arange(N)
cm = LinearSegmentedColormap.from_list('defcol', ["#000000", "#FF0000"])
trace_color = cm(np.linspace(0, 1, cycles))
for k, color in zip(range(cycles), trace_color):
    lis = KL_rest[k::cycles]
    plt.scatter(scanpoints, lis, color=color, marker='^', alpha=0.9)
plt.colorbar(ScalarMappable(cmap=cm, norm=plt.Normalize(0, cycles - 1)), ticks=np.arange(cycles), label='cycles')
plt.show()

Note that in this case, you can create the scatter plot in one go, enabling a default color bar. For this to work, the scanpoints can be repeated for each cycle, and the colors can be indicated by tiling the cycle for each scan point.

If you only want to show the really used colors, you can add N=cycles in the creation of the color map. To put the tick marks for each number in the center of the cells, you can move the default limits by 0.5.

from matplotlib import pyplot as plt
from matplotlib.colors import LinearSegmentedColormap
import numpy as np

cycles = 7
N = 100
KL_rest = np.sin(np.linspace(0, 10 * np.pi, cycles * N))
scanpoints = np.arange(N)
cm = LinearSegmentedColormap.from_list('defcol', ["#000000", "#FF0000"], N=cycles)

plt.scatter(np.repeat(scanpoints, cycles), KL_rest,
            c=np.tile(range(cycles), len(scanpoints)),
            cmap=cm, norm=plt.Normalize(-0.5, cycles - 0.5),
            marker='^', alpha=0.9)
plt.colorbar(ticks=np.arange(cycles), label='cycles')
plt.show()

Answered By - JohanC