Matplotlib plotting custom colormap with the plot

Question

I have been following a tutorial on plotting F1 data over a circuit, color coded with the fastf1 library. I wanted to add some extra's to the script to utilize the official team colors. It works but the end result shows the colormap with the circuit covering the n bins 100 . In the picture above I used the same colormap as in the tutorial 'winter' so there is most certainly something wrong in my code.

However, the original tutorial gets a cleaner end result with only the circuit showing like this:

the tutorial in question uses a default colormap from matplotlib 'winter' . To get the team colors working I had to create a custom colormap from the 2 colors that are fetched from api.

Let's get into the code, I have tried so much and searched everywhere without success...

The custom colormap is build with this sequence of code I got from the matplotlib docs.

# Create custom colormap
teamcolor1 = to_rgb('{}'.format(team1_color))
teamcolor2 = to_rgb('{}'.format(team2_color))
colors = [teamcolor1, teamcolor2]
n_bins = [3, 6, 10, 100]
cmap_name = 'colors'
fig, axs = plt.subplots(2, 2, figsize=(6, 9))
fig.subplots_adjust(left=0.02, bottom=0.06, right=0.95, top=0.94, wspace=0.05)

x = np.arange(0, np.pi, 0.1)
y = np.arange(0, 2 * np.pi, 0.1)
X, Y = np.meshgrid(x, y)
Z = np.cos(X) * np.sin(Y) * 10
for n_bin, ax in zip(n_bins, axs.ravel()):
    colormap = LinearSegmentedColormap.from_list(cmap_name, colors, N=n_bin)
    im = ax.imshow(Z, interpolation='nearest', origin='lower', cmap=colormap)
    ax.set_title("N bins: %s" % n_bin)
    fig.colorbar(im, ax=ax)
    cm.register_cmap(cmap_name, colormap)

I register the colormap to easily call it later in the script with get_cmap .

The eventual plotting of the circuit is done in this piece of code:

x = np.array(telemetry['X'].values)
y = np.array(telemetry['Y'].values)

points = np.array([x, y]).T.reshape(-1, 1, 2)
segments = np.concatenate([points[:-1], points[1:]], axis=1)
fastest_driver_array = telemetry['Fastest_driver_int'].to_numpy().astype(float)

cmap = cm.get_cmap('winter', 2)
lc_comp = LineCollection(segments, norm=plt.Normalize(1, cmap.N+1), cmap=cmap)
lc_comp.set_array(fastest_driver_array)
lc_comp.set_linewidth(5)

plt.rcParams['figure.figsize'] = [18, 10]

plt.gca().add_collection(lc_comp)
plt.axis('equal')
plt.tick_params(labelleft=False, left=False, labelbottom=False, bottom=False)

cbar = plt.colorbar(mappable=lc_comp, boundaries=np.arange(1, 4))
cbar.set_ticks(np.arange(1.5, 9.5))
cbar.set_ticklabels(['{}'.format(driver1), '{}'.format(driver2)])

plt.savefig(
    '{}_'.format(year) + '{}_'.format(driver1) + '{}_'.format(driver2) + '{}_'.format(circuit) + '{}.png'.format(
        session), dpi=300)

plt.show()

This is where I think things go wrong, but I'm unsure of what is going wrong. I guess it has to do with how I use the colormap. But everything I changed broke the whole script. As I don't have a lot of experience with matplotlib, it's getting very complicated.

As I don't want this question to be overly long the whole code can be read here: https://gist.github.com/platinaCoder/7b5be22405f2003bd577189692a2b36b

Answer 1

Instead of creating a whole custome cmap, I got rid of this piece of code:

# Create custom colormap
teamcolor1 = to_rgb('{}'.format(team1_color))
teamcolor2 = to_rgb('{}'.format(team2_color))
colors = [teamcolor1, teamcolor2]
n_bins = [3, 6, 10, 100]
cmap_name = 'colors'
fig, axs = plt.subplots(2, 2, figsize=(6, 9))
fig.subplots_adjust(left=0.02, bottom=0.06, right=0.95, top=0.94, wspace=0.05)

x = np.arange(0, np.pi, 0.1)
y = np.arange(0, 2 * np.pi, 0.1)
X, Y = np.meshgrid(x, y)
Z = np.cos(X) * np.sin(Y) * 10
for n_bin, ax in zip(n_bins, axs.ravel()):
    colormap = LinearSegmentedColormap.from_list(cmap_name, colors, N=n_bin)
    im = ax.imshow(Z, interpolation='nearest', origin='lower', cmap=colormap)
    ax.set_title("N bins: %s" % n_bin)
    fig.colorbar(im, ax=ax)
    cm.register_cmap(cmap_name, colormap)

and replaced cmap = cm.get_cmap('colors', 2) with cmap = cm.colors.ListedColormap(['{}'.format(team1_color), '{}'.format(team2_color)])