Create A Stock Market Heat Map Using Python

Sharing is caring!

Last Updated on July 14, 2022 by Jay

In this tutorial we’ll go through step by step how to create a stock market heat map using the Python plotly library.

The heatmap shows quite a lot of information:

• Companies are placed in different sectors
• The size of the boxes represents the market cap of each company
• The color of each stock presents the stock price change from the previous day
• Interactive and allows zooming in and out

Libraries

We’ll use the following three libraries: pandas, yfinance, and plotly.

pip install pandas yfinance plotly

Getting Stock Data

In this example, I’m going to work with 100 randomly selected companies from the S&P 500 companies.

Getting S&P 500 Company List

This Wikipedia page contains the S&P 500 companies. The list is in a table format, so we can use pandas.read_html to scrape that information directly.

Note there are two tables, and we just need the first one.

import pandas as pd
sp500 = pd.read_html(r'https://en.wikipedia.org/wiki/List_of_S%26P_500_companies')[0]['Symbol']

Download Individual Stock Information

We’ll be using the yfinance library to download stock information. Check out this tutorial for a quick introduction of the library.

We are interested in the following information:

• Stock ticker (already obtained from Wikipedia)
• Company sector
• Company market cap
• Change in stock price (from a trading day ago)

We’ll initialize a few lists to store the above information:

tickers = []
deltas = []
sectors =[]
market_caps = []

Then we’ll loop through the sp500 list (which contains 100 randomly selected companies) and download relevant information.

for ticker in sp500:
    
    try:
        ## create Ticker object
        stock = yf.Ticker(ticker)
        tickers.append(ticker)

        ## download info
        info = stock.info

        ## download sector
        sectors.append(info['sector'])

        ## download daily stock prices for 2 days
        hist = stock.history('2d')

        ## calculate change in stock price (from a trading day ago)
        deltas.append((hist['Close'][1]-hist['Close'][0])/hist['Close'][0])

        ## calculate market cap
        market_caps.append(info['sharesOutstanding'] * info['previousClose'])

        ## add print statement to ensure code is running
        print(f'downloaded {ticker}')
    except Exception as e:
        print(e)

Note that the yfinance library scrapes directly from Yahoo Finance website and can be quite slow. We included a print statement in the code to display messages so we can see that code is running instead of freezing. As a result, it depends on how many companies you want to download. To get all the S&P 500 companies, I suggest you run the code in the background while doing something else.

If you want just a subset of the S&P 500 companies, the pandas.sample() method can quickly select some random samples from the list of S&P 500 companies. For example, sample(100) will return 100 random samples.

sp500 = sp500.sample(100)

Plotting A Stock Heat Map Using Python Plotly

Now we have all the information, it’s time to do some data processing before plotting.

Data Preparation

First, we’ll put all the data into a pandas dataframe.

df = pd.DataFrame({'ticker':tickers,
                  'sector': sectors,
                  'delta': deltas,
                  'market_cap': market_caps,
                  })

Then, we’ll bin the data into different groups based on the change in stock price (i.e. delta). This is for assigning different colors to each company.

I picked some arbitrary thresholds: -100%, -2%, -1%, 0%, 1%, 2%, and 100%. You can pick any bands that fit your needs.

Then we bin the data using the pd.cut method. We also assign color names based on the bins:

color_bin = [-1,-0.02,-0.01,0, 0.01, 0.02,1]
df['colors'] = pd.cut(df['delta'], bins=color_bin, labels=['red','indianred','lightpink','lightgreen','lime','green'])
df

The resulting dataframe looks like this:

Plotting with plotly.express

The plotly library offers two different classes (and APIs) for plotting. The easier one is called plotly.express, and the more advanced one is called plotly.graph_objects. For this demo, we’ll use plotly.express. We’ll talk about how to use the plotly.graph_objects in another tutorial since that one requires more preparation work.

Plotly has a heatmap plot, but that’s not what we are going to use!

Instead, we’ll be using the treemap plot, which is good for visualizing hierarchical data such as companies and their sectors.

Note we need to specify:

• A path argument for the hierarchical data structure.
• The values argument determines the size of each box.
• The color argument tells the function which dataframe column (i.e. ‘colors’ column) to look for the color information.
• The color_discrete_map argument is a dictionary that maps each color value (in the ‘colors’ column) to a color name that can be rendered by plotly.
• Last but not least, the hover_data argument helps format the text on the tooltip. i.e. converting decimal numbers into percentages.

import plotly.express as px
fig = px.treemap(df, path=[px.Constant("all"), 'sector','ticker'], values = 'market_cap', color='colors',
                 color_discrete_map ={'(?)':'#262931', 'red':'red', 'indianred':'indianred','lightpink':'lightpink', 'lightgreen':'lightgreen','lime':'lime','green':'green'},

                hover_data = {'delta':':.2p'}
                ))
fig.show()

Additional Resources

How To Use Python To Get Stock Data

Binning Data With Pandas Cut Method