How to add a color bar to label the cells from which sample with `annadata.obs['samples']` in the `scanpy.pl.heatmap` plot?

Can anyone help me and tell me: How to add a color bar to label the cells from which sample with annadata.obs['samples'] in the scanpy.pl.heatmap plot (https://scanpy.readthedocs.io/en/stable/generated/scanpy.pl.heatmap.html)?

Thanks a lot.

I worked out this question and my solution is:

from anndata import AnnData
import numpy as np
import pandas as pd
    
#-- Create a reproducible example
gene_exp = np.array([[0,2,0,1], [1,2,1,3], [6,2,4,1], [5,2,2,1]])
adata = AnnData(gene_exp)
adata.obs['sample'] = ['sample1', 'sample2', 'sample3', 'sample4']
adata.var['gene'] = ['g'+str(i) for i in range(4)]

obs = adata.obs
sample1_idx = obs['sample'] == 'sample1'
sample2_idx = obs['sample'] == 'sample2'
sample3_idx = obs['sample'] == 'sample3'
sample4_idx = obs['sample'] == 'sample4'

data = (np.arange(obs.shape[0])+1)/sum(obs['sample'].value_counts())

import matplotlib.pyplot as plt

plt.rcParams["figure.figsize"] = (20,1)
plt.axis('off')
for a in data[sample1_idx]:
    plt.axvline(x=a, c = 'g',linewidth=2)
for b in data[sample2_idx]:
    plt.axvline(x=b, c = 'lime',linewidth=2)
for c in data[sample3_idx]:
    plt.axvline(x=c, c = 'yellow',linewidth=2)
for d in data[sample4_idx]:
    plt.axvline(x=d, c = 'y',linewidth=2)