Skip to content

Commit 3dadb98

Browse files
Added treemap example (#531)
* Added treemap example * Update readme.md * Update readme
1 parent d5fb473 commit 3dadb98

File tree

5 files changed

+471
-0
lines changed

5 files changed

+471
-0
lines changed
112 KB
Loading
Lines changed: 130 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,130 @@
1+
import json
2+
import concurrent.futures
3+
from polygon import RESTClient
4+
5+
# Initialize Polygon API client
6+
client = RESTClient(
7+
trace=True
8+
) # Assuming you have POLYGON_API_KEY environment variable set up
9+
10+
# Initialize the data structure to hold SIC code groups
11+
sic_code_groups = {}
12+
13+
14+
# https://en.wikipedia.org/wiki/Standard_Industrial_Classification
15+
# https://www.investopedia.com/terms/s/sic_code.asp
16+
def sic_code_to_group(sic_code):
17+
"""
18+
Maps a given SIC code to the industry group.
19+
"""
20+
sic_code = int(sic_code)
21+
if 100 <= sic_code <= 999:
22+
return "Agriculture, Forestry and Fishing"
23+
elif 1000 <= sic_code <= 1499:
24+
return "Mining"
25+
elif 1500 <= sic_code <= 1799:
26+
return "Construction"
27+
# Note: 1800-1999 not used
28+
elif 2000 <= sic_code <= 3999:
29+
return "Manufacturing"
30+
elif 4000 <= sic_code <= 4999:
31+
return "Transportation and Public Utilities"
32+
elif 5000 <= sic_code <= 5199:
33+
return "Wholesale Trade"
34+
elif 5200 <= sic_code <= 5999:
35+
return "Retail Trade"
36+
elif 6000 <= sic_code <= 6799:
37+
return "Finance, Insurance and Real Estate"
38+
elif 7000 <= sic_code <= 8999:
39+
return "Services"
40+
elif 9100 <= sic_code <= 9729:
41+
return "Public Administration"
42+
elif 9900 <= sic_code <= 9999:
43+
return "Nonclassifiable"
44+
else:
45+
return None
46+
47+
48+
def process_ticker(ticker_snapshot):
49+
ticker = ticker_snapshot.ticker
50+
51+
try:
52+
details = client.get_ticker_details(ticker)
53+
54+
# Check if the type is 'CS' (common stock), if not, return early without processing this ticker
55+
# if getattr(details, 'type', None) != 'CS' or getattr(details, 'market_cap', None) != None:
56+
if (
57+
getattr(details, "type", None) != "CS"
58+
or getattr(details, "market_cap", None) is None
59+
):
60+
return
61+
62+
sic_code = details.sic_code
63+
sic_description = getattr(
64+
details, "sic_description", None
65+
) # Use getattr to avoid AttributeError if sic_description is not present
66+
market_cap = getattr(details, "market_cap", None)
67+
68+
# if sic_code:
69+
# sic_code = str(sic_code)[:1] # Extract first 1 digits
70+
71+
if sic_code:
72+
sic_group = sic_code_to_group(sic_code)
73+
if sic_group is None:
74+
return
75+
76+
# Check if the sic_code is already in the groups, if not create a new entry with sic_description and empty companies list
77+
# if sic_code not in sic_code_groups:
78+
# sic_code_groups[sic_code] = {"sic_description": sic_description, "companies": []}
79+
80+
if sic_group not in sic_code_groups:
81+
sic_code_groups[sic_group] = {
82+
"sic_description": sic_group,
83+
"companies": [],
84+
}
85+
86+
# Append the company details to the corresponding SIC code entry
87+
# sic_code_groups[sic_code]["companies"].append({
88+
# "ticker": ticker,
89+
# "market_cap": market_cap
90+
# })
91+
92+
sic_code_groups[sic_group]["companies"].append(
93+
{"ticker": ticker, "market_cap": market_cap}
94+
)
95+
96+
except Exception as e:
97+
print(f"Error processing ticker {ticker}: {e}")
98+
99+
100+
# Get snapshot data
101+
snapshot = client.get_snapshot_all("stocks")
102+
103+
# Execute the data processing in parallel, limited to 100 workers
104+
with concurrent.futures.ThreadPoolExecutor(max_workers=100) as executor:
105+
executor.map(process_ticker, snapshot)
106+
107+
# Modify the SIC Code Groups Dictionary to include the weights
108+
for sic_code, group_data in sic_code_groups.items():
109+
companies = group_data["companies"]
110+
total_market_cap = sum(
111+
company["market_cap"] for company in companies if company["market_cap"]
112+
)
113+
114+
# If total_market_cap is 0, we will skip weight calculation to avoid division by zero
115+
if total_market_cap == 0:
116+
continue
117+
118+
for company in companies:
119+
if company[
120+
"market_cap"
121+
]: # Avoid dividing by zero if a company's market cap is None or 0
122+
company["weight"] = company["market_cap"] / total_market_cap
123+
else:
124+
company["weight"] = 0 # You can also set to a default value if preferred
125+
126+
# Save the enhanced data structure to a JSON file
127+
with open("sic_code_groups.json", "w") as f:
128+
json.dump(sic_code_groups, f)
129+
130+
print("Data collection complete and saved to 'sic_code_groups.json'")

examples/tools/treemap/readme.md

Lines changed: 60 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,60 @@
1+
# Mapping Market Movements with Polygon.io and D3.js Treemap
2+
3+
This repository offers a tutorial on how to create a Treemap visualization of the current stock market conditions. Using D3.js Treemap, Polygon.io's [Snapshot API](https://polygon.io/docs/stocks/get_v2_snapshot_locale_us_markets_stocks_tickers), and the [python-client library](https://github.com/polygon-io/client-python), we'll guide you through building an interactive visualization. The Snapshot API allows us to fetch the most recent market data for all US-traded stocks, transforming them into color-coded nested rectangles within the Treemap. This presents an insightful and interactive snapshot of the market's current status.
4+
5+
![Treemap Visualization](./market-wide-treemap.png)
6+
7+
Please see the [tutorial](https://polygon.io/blog/market-movements-with-treemap) for more details.
8+
9+
## Structure
10+
11+
The repo consists of:
12+
13+
- `polygon_sic_code_data_gatherer.py`: Builds ticker to SIC code mapping for treemap groups.
14+
- `sic_code_groups.json`: Pre-built JSON file containing grouped ticker to SIC code data.
15+
- `treemap_server.py`: Simple server to host the treemap visualization (requires sic_code_groups.json).
16+
17+
For those interested in the underlying mechanics, the `polygon_sic_code_data_gatherer.py` script retrieves a snapshot of all ticker symbols, processes each one to obtain its SIC code via the Ticker Details API, and then saves these classifications into the file named `sic_code_groups.json`.
18+
19+
The logic of this SIC code-to-group enables us to transform a large dataset into a neatly structured visualization. This structured approach facilitates easy identification of market conditions, providing a snapshot of the market's overall health. You don't need to do anything since it is pre-built but we added the script if you wanted to modify anything.
20+
21+
## Getting Started
22+
23+
Setting up and visualizing the stock market's current conditions is straightforward. All you'll need to do is clone the repository, secure an API key from Polygon.io, install the required Python library, launch the visualization server example, and then dive into the visualization through your web browser.
24+
25+
### Prerequisites
26+
27+
- Python 3.x
28+
- Polygon.io account and API key
29+
30+
### Setup
31+
32+
1. Clone the repository:
33+
```
34+
git clone https://github.com/polygon-io/client-python.git
35+
```
36+
37+
2. Install the necessary Python packages.
38+
```
39+
pip install -U polygon-api-client
40+
```
41+
42+
3. Store your Polygon.io API key securely, or set it as an environment variable:
43+
```
44+
export POLYGON_API_KEY=YOUR_API_KEY_HERE
45+
```
46+
47+
### Running the Treemap Server
48+
49+
Change into the treemap example directory and execute the `treemap_server.py` script:
50+
```
51+
cd examples/tools/treemap
52+
python3 treemap_server.py
53+
```
54+
55+
Upon successful execution, the server will start, and you can view the treemap visualization by navigating to:
56+
```
57+
http://localhost:8889
58+
```
59+
60+
That’s it. You'll now see a Treemap that organizes over 4,000+ US-traded stocks into 10 distinct categories. Use the dropdown box in the top left corner to select different categories and delve deeper into the data.

examples/tools/treemap/sic_code_groups.json

Lines changed: 1 addition & 0 deletions
Large diffs are not rendered by default.

0 commit comments

Comments
 (0)