How to draw a straight line between two points on a heatmap

Viewed 49

I am creating a 2D matplotlib plot (i and j coordinates) which contains 10 subplots. Each subplot contains 150 by 150 grid cell data. How can I insert a thick straight line arrow (color = red) that passes from one point (62, 62) to another point (58, 94) on each generated heatmap sub-plot across those 10 sub-plots? Any feedback is greatly appreciated.

    from matplotlib.patches import Rectangle
    import numpy as np
    import matplotlib.pyplot as plt
    from sklearn.metrics import r2_score, median_absolute_error
    import os
    import matplotlib.cm as cm
    from mpl_toolkits import axes_grid1
    import matplotlib.pyplot as plt
    #import seaborn as sns
    import matplotlib.pyplot as plt
    import matplotlib.colors as mcolors
    import matplotlib.colors
    import matplotlib.colors as colors
    
    fig, axes = plt.subplots(2, 5, figsize=(16, 8))
    row_count = 0
    col_count = 0

    for i,ax in enumerate(axes.flat):

    sub_plot_data = data[(i)*(150*150):(i+1)*150*150]

    x = 150
    y = 150
    #--------------------------- Define the map boundary ---------------------- 
    xmin = 1258096.6
    xmax = 1291155.0
    ymin = 11251941.6
    ymax = 11285000.0

    pmin = min(sub_plot_data)
    pmax = max(sub_plot_data) 
    # ---------------------------  define color bar for Discrete  color 
    bounds = np.linspace(-1, 1, 10)
    Discrete_colors = plt.get_cmap('jet')(np.linspace(0,1,len(bounds)+1))
    # create colormap without the outmost colors
    cmap = mcolors.ListedColormap(Discrete_colors[1:-1]) # 

    actual_2d = np.reshape(sub_plot_data,(y,x)) 

    #im = ax.imshow(actual_2d, interpolation=None, cmap=cmap, extent=(xmin, xmax, ymin, ymax), vmin=pmin, vmax=pmax)      
    im = ax.imshow(actual_2d, interpolation=None, cmap=cmap)      
    ax.text(actual_2d[62, 62], actual_2d[62, 62]-10, '%s' % 'Sale_1', 
        horizontalalignment='center', verticalalignment='center', color= 'black', fontsize=18)
    ax.set_title("Sale_Stores-%s - L: %s"%(i+1, 1), fontsize=14, pad=30, x=0.5, y=0.999)
    ax.set_aspect('auto')
    ax.add_patch(Rectangle((60, 60), 6, 6, edgecolor='red', facecolor='red', fill=True, lw=2))
    ax.text(62, 62, '%s' % 'Sale_1', ha='center', va='center', color='black', fontsize=14)

       
    fig.tight_layout(h_pad=10) 
    plt.subplots_adjust(left=0.02,
                    bottom=0.1, 
                    right=0.91, 
                    top=0.8, 
                    wspace=0.1, 
                    hspace=0.5)

    cbaxes = fig.add_axes([0.94, 0.05, 0.02, 0.8])
    cbar = fig.colorbar(im, ax=axes.flat, ticks=v, extend='both', cax=cbaxes)
    cbar.ax.tick_params(labelsize=10) 
    #cbar.set_ticks(v)
    cbar.ax.set_yticklabels([str(i) for i in v], fontsize=12)

   #plt.tick_params(left=False, labelleft=False, top=False, labeltop=False, right=False, 
   labelright=False, bottom=False, labelbottom=False)

   plt.show()

0 Answers
Related