seaborn: heatmap plot in the order of input-CodePudding

Consider the following example dataset (and the MWE):

import pandas as pd
import matplotlib.pyplot as plt
import matplotlib.colors as mcolors
import seaborn as sns

data = {'trajectory': [101,102,102,102,102,102,102,102,104,104,104,104,104,104,104,107,107,107,107,
          107,107,107,107,107,108,108,108,108,108,108,108,109,109,109,109,109,109,112,
         112,112,112,112,113,113,113,113,114,114,114,114],
 'segment': [1,1,1,1,2,2,3,3,1,1,2,2,2,3,3,1,1,2,2,2,2,3,3,3,1,1,1,
          2,2,2,2,1,1,1,2,2,2,1,1,2,2,2,1,2,2,3,1,2,2,2],
  'prediction': [3,0,0,1,3,3,2,2,0,0,4,4,2,0,0,0,0,2,2,2,3,0,0,2,0,0,1,1,
          1,1,0,1,2,1,3,3,3,1,1,4,4,2,1,4,4,3,0,3,3,2]}

df = pd.DataFrame(data)
df.head(2)
  trajectory segment prediction
0   101        1        3
1   102        1        0

Then I plot a heatmap of filtered data like so:

plot_data = (df.value_counts()
   .sort_values(ascending=False).reset_index()
   .drop_duplicates(['trajectory', 'segment'])
   .pivot_table(index='trajectory', columns='segment', values='prediction'))

cmap = mcolors.ListedColormap(['c', 'b', 'g', 'y','m',])

fig, ax = plt.subplots(figsize=(10,6))
ax.set_facecolor('black')
sns.heatmap(plot_data, vmin=-0.5, vmax=4.5,cmap=cmap, annot=True)

colorbar = ax.collections[0].colorbar
colorbar.set_ticks([0, 1., 2., 3, 4])
colorbar.set_ticklabels(['Foot', 'Bike', 'Bus', 'Car', 'Metro'])

To produce:

But then I want to plot the heatmap with input 0 first, followed by inputs 1, 2, etc..

Required output:

CodePudding user response：

You can

(in your desired result, the Bike values for trajectories 109, 112 and 113 should be in segment 1 instead of 2)