python finance 8 correlation matrix heat map

heat map shows if 2 stocks are related, negative related, or unrelated.

correlation matrix, close to 1=> related, close to -1 => negative related, 

close to 0 => unrelated

full heat map, all 500 companies

def visualize_data():
    df = pd.read_csv('sp500_joined_closes.csv')

    #build in function, generate correlation matrix
    df_corr = df.corr()
    print(df_corr.tail())

    data = df_corr.values
    fig = plt.figure()
    #3D plot
    ax = fig.add_subplot(1, 1, 1)

    #heatmap color: red, yellow, green
    heatmap = ax.pcolor(data, cmap=plt.cm.RdYlGn)
    fig.colorbar(heatmap)

    #labels are in the middle of the bar
    ax.set_xticks(np.arange(data.shape[0]) + 0.5, minor=False)
    ax.set_yticks(np.arange(data.shape[1]) + 0.5, minor=False)
    #yaxis flipped
    ax.invert_yaxis()
    ax.xaxis.tick_top()

    column_lables = df_corr.columns
    row_labels = df_corr.index

    ax.set_xticklabels(column_lables)
    ax.set_yticklabels(row_labels )
    #rotate xaxis label 90 degree
    plt.xticks(rotation=90)
    #correlation matrix data values are between -1 and 1, set heat map accordingly
    heatmap.set_clim(-1, 1)
    plt.tight_layout()
    plt.show()

visualize_data()

----------
reference:
https://www.youtube.com/watch?v=PxUzcDJBEZ4&list=PLQVvvaa0QuDcOdF96TBtRtuQksErCEBYZ&index=8