How to plot a column value with its index as axis

user_13021743 0
Viewed 109

I have a data frame df:

ColumnA  ColumnB  correlation
dog1     cat1     3.6
dog2     cat2     2.9
dog3     cat3     -2.5

In reality, I have 50 rows in the data frame. To make it simple I am representing it here with only 3 rows.

I am interested in illustrating the correlation between ColumnA and ColumnB that is given in df['correlation']. What would be the best possible way to do so?

One of the choices may be to plot a plot of this kind:

enter image description here

However, I am not sure how to plot the same using matplotlib or any other python module. Also, how to keep only one axis on the plot. For instance, keep only ColumnA axis on the left as it is and remove ColumnB axis from the right.

Other suggestions to represent the same in a better way are welcome.

seaborn dataframe matplotlib plot python
1 Answers

Using sns.heatmap we can plot a single column heatmap and then tweak the y-axes to put the additional labels on:

fig, (cbar_ax, box_ax_0) = plt.subplots(
    nrows=1, 
    ncols=2, 
    gridspec_kw={'width_ratios': [1, 19]},
)
sns.heatmap(
    data=df["correlation"].to_numpy().reshape(-1,1), 
    annot=True,
    fmt='.2f',
    vmin=-5, 
    vmax=5, 
    square=True, 
    ax=box_ax_0,
    cbar_ax=cbar_ax,
)
# rotate the first set of ticks
box_ax_0.set_yticklabels(df["ColumnA"], rotation="horizontal")

# remove x-axis ticks
box_ax_0.xaxis.set_ticks([])

# Duplicate the y-axis
box_ax_1 = box_ax_0.secondary_yaxis("right")
# copy ticks from the first y-axis
box_ax_1.set_yticks(box_ax_0.get_yticks())
# add `ColumnB` labels 
box_ax_1.set_yticklabels(df["ColumnB"])

# remove spines from secondary axis
box_ax_1.spines['right'].set_visible(False)

# tweak layout to bring the colour bar closer
plt.tight_layout()

Single column heatmap with secondary y-axis

If you want to reverse the colour map you can set cmap="rocket_r" in the sns.heatmap function call.

user_32797160
Related