Matplotlib - Getting colorbar to line up with colors in plot and to use original values as labels

Question

I have a dataframe with two values x1 and x2, each of these points has a class label.

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from matplotlib import cm
from mpl_toolkits.axes_grid1 import make_axes_locatable
from sklearn.preprocessing import normalize

# Random data
np.random.seed(0)
test_df = pd.DataFrame(data=np.random.randint(1, 11, (100, 2)), columns=['x1', 'x2'])
test_df['class'] = np.random.randint(1, 25, 100)  # What I want to plot as colors
test_df['norm_class'] = normalize(test_df['class'].values.reshape(-1, 1), norm='max', axis=0) 
test_df.head()

Dataframe output:

    x1  x2  class   norm_class
0   6   1   24      1.000000
1   4   4   22      0.916667
2   8   10  20      0.833333
3   4   6   15      0.625000
4   3   5   21      0.875000

I want to plot x1 and x2 points in 2 dimensions and use the class label as a continuous color of the points, so I normalize class, create a color map from the normalized values, and plot it along with a colorbar:

# Create colormap 
response_norm_colors = [cm.viridis(x) for x in test_df['norm_class'].values]

# Plot x1 and x2 as 2 dimensions, with color as "third" dimension
fig = plt.figure(figsize=(10, 6))
ax = fig.add_subplot(1, 1, 1)

divider = make_axes_locatable(ax)
cax = divider.append_axes('right', size='5%', pad=0.05)

sctr = ax.scatter(test_df['x1'], test_df['x2'], c=response_norm_colors)

fig.colorbar(sctr, cax=cax, orientation='vertical')

plt.show()
plt.close()

Obviously, my colors don't match between the scatterplot and colorbar. I'd also like for the colorbar scale to be for the original class values before I normalized them.

1. How do I get the colormap of the colorbar to match up to the colormap I've used in my scatterplot (viridis)?
2. How can I put the original, non-normalized values (test_df['class']) to be listed as the scale of the colorbar, rather than the [0, 1] interval?

Answer 1

Your code is working perfectly fine at my end.

enter image description here