# Jupyter notebook style help + code suggestions for pandas

I wanted to open source some code to scrape and analyze publicly-filed stock buys and sells from U.S. senators. I'm not familiar with code style for Jupyter notebooks or pandas in general. Would it be possible for you to review my short notebook? The original can be found here.

Ideally I would get the scraping code reviewed as well, but for the sake of brevity I wanted to keep it to just the pandas and Jupyter notebook-related changes. Within scope is things like how the Jupyter notebook should be structured, general Python code style, and pandas conventions + optimizations.

I know that this is a larger ask, so I am also open to just high-level suggestions.

I have included the contents of the Jupyter notebook below. (I thought about removing the # In[ ]: comments, but realized they indicate where each Jupyter cell begins.) Thank you in advance!

# # Senator Filings Analysis

# ***

# ## Imports

# In[ ]:

from collections import defaultdict
import datetime as dt
from functools import lru_cache
import json
from os import path
import pickle

import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import yfinance as yf

# ## Introduction
#
# In this notebook, we explore stock orders that were publicly filed by U.S. senators. The filings are scraped from https://efdsearch.senate.gov/search/. We calculate the returns of each senator by mimicking their buys and sells.

# ***

#
# The senators.pickle file is scraped using the script in the root of the repository.

# In[ ]:

with open('senators.pickle', 'rb') as f:

# ## Data cleaning

# ### Filling in missing tickers

# In this section, we fill in as many of the missing ticker symbols as we can.

# In[ ]:

def tokenize(asset_name):
""" Convert an asset name into useful tokens. """
token_string = asset_name
.replace('(', '')
.replace(')', '')
.replace('-', ' ')
.replace('.', '')

def token_is_ticker(token, token_blacklist):
return len(token) <= 4 and token.upper() not in token_blacklist

# These generic words do not help us determine the ticker
with open('blacklist.json', 'r') as f:

missing_tickers = set(raw_senators_tx[
(raw_senators_tx['ticker'] == '--')
| (raw_senators_tx['ticker'] == '')
]['asset_name'])

ticker_map = {}
unmapped_tickers = set()
for m in missing_tickers:
tokens = tokenize(m)
if token_is_ticker(tokens[0], blacklist):
ticker_map[m] = tokens[0].upper()
elif token_is_ticker(tokens[-1], blacklist):
ticker_map[m] = tokens[-1].upper()
else:

# As a second pass, we assign tickers to asset names that have any of the specified keywords.

# In[ ]:

phrase_to_ticker = {
'FOX': 'FOX',
'AMAZON': 'AMZN',
'AARON': 'AAN',
'ALTRIA': 'MO',
'APPLE': 'AAPL',
'CHEVRON': 'CVX',
'DUPONT': 'DD',
'ALPHABET': 'GOOGL',
'GOOG': 'GOOGL',
'GENERAL ELECTRIC': 'GE',
'JOHNSON': 'JNJ',
'NEWELL': 'NWL',
'OWENS': 'OMI',
'PFIZER': 'PFE',
'TYSON': 'TSN',
'UNDER ARMOUR': 'UAA',
'VERIZON': 'VZ',
'WALT': 'DIS'
}

for m in unmapped_tickers:
for t in phrase_to_ticker:
if t in m.upper():
ticker_map[m] = phrase_to_ticker[t]

tx_with_tickers = raw_senators_tx.copy()
for a, t in ticker_map.items():
tx_with_tickers.loc[tx_with_tickers['asset_name'] == a, 'ticker'] = t

# ### Filtering rows and columns

# We filter out useless rows and missing symbols, and then add some useful columns for the final dataset.

# In[ ]:

filtered_tx = tx_with_tickers[tx_with_tickers['ticker'] != '--']
filtered_tx = filtered_tx.assign(
ticker=filtered_tx['ticker'].map(
lambda s: s.replace('--', '').replace('\n', '')))

filtered_tx = filtered_tx[filtered_tx['order_type'] != 'Exchange']

# In[ ]:

def parse_tx_amount(amt):
""" Get the lower bound for the transaction amount. """
return int(amt.replace('Over $50,000,000', '50000000') .split(' - ')[0] .replace(',', '') .replace('$', ''))

senators_tx = filtered_tx.assign(
tx_estimate=filtered_tx['tx_amount'].map(parse_tx_amount))
senators_tx = senators_tx.assign(
full_name=senators_tx['first_name']
.str
.cat(senators_tx['last_name'], sep=' ')
)
useful_cols = [
'file_date',
'tx_date',
'full_name',
'order_type',
'ticker',
'tx_estimate'
]
senators_tx = senators_tx[useful_cols]
senators_tx = senators_tx.assign(
tx_date=senators_tx['tx_date'].map(
lambda v: dt.datetime.strptime(v, '%m/%d/%Y')))
senators_tx = senators_tx.assign(
file_date=senators_tx['file_date'].map(
lambda v: dt.datetime.strptime(v, '%m/%d/%Y')))
senators_tx

# ## Returns calculation

# These cells help us download the market data for the specified tickers. We store the market data in files so we don't need to repeatedly download the same information.

# In[ ]:

""" Download a file of stock prices for this ticker to disk. """
if check_cache and path.exists('stocks/{0}.pickle'.format(ticker)):
return
d = yf.Ticker(ticker)
with open('stocks/{0}.pickle'.format(ticker), 'wb') as f:
pickle.dump({
'price': d.history(period='max').reset_index()
}, f)

""" Load the file of stock prices for this ticker. """
with open('stocks/{0}.pickle'.format(ticker), 'rb') as f:
raw = dump['price']
return raw[['Date', 'Close']]
.rename(columns={'Date': 'date', 'Close': 'price'})

def _price_for_date(df, date):
""" Helper function for ticker_at_date. """
df = df[df['date'] >= date].sort_values(by='date')
return df['price'].iloc[0]

@lru_cache(maxsize=128)
def ticker_at_date(ticker, date):
"""
Price of a ticker at a given date. Raise an IndexError if there is no
such price.
"""
try:
# Sell at the next opportunity possible
return _price_for_date(data, date)
except Exception:
# If any exception occurs, refresh the cache
return _price_for_date(data, date)

# In[ ]:

all_tickers = set(senators_tx['ticker'])
for i, t in enumerate(all_tickers):
if i % 100 == 0:
print('Working on ticker {0}'.format(i))
try:
except Exception as e:
print('Ticker {0} failed with exception: {1}'.format(t, e))

# ### Mimicking buy + sell orders
#
# We calculate a given senator's return by calculating the return between each buy or sell order, and then solving for the cumulative return. We convert that to a CAGR given the time period the senator was investing.
#
# We keep track of how many units of each stock a senator is holding. If we ever see a filing that indicates the senator sold more than we estimated they are holding, we just sell all of the units we have on record. (We do not allow the senator to go short.)

# In[ ]:

buckets = [
(1000, 15000),
(15000, 50000),
(50000, 100000),
(100000, 250000),
(250000, 500000),
(500000, 1000000),
(1000000, 5000000),
(5000000, 25000000),
(25000000, 50000000),
(50000000, float('inf'))
]

def same_bucket(dollar_value_a, dollar_value_b):
"""
If the dollar value of the stock units is roughly the same, sell all
units.
"""
for v1, v2 in buckets:
if dollar_value_a >= v1 and dollar_value_a < v2:
return dollar_value_b >= v1 and dollar_value_b < v2
return False

def portfolio_value(stocks, date):
"""
Value of a portfolio if each ticker has the specified number of units.
"""
v = 0
for s, units in stocks.items():
if units == 0:
continue
try:
v += ticker_at_date(s, date) * units
except IndexError as e:
# Swallow missing ticker data exception
pass
return v

def calculate_return(before_values,
after_values,
begin_date,
end_date,
tx_dates):
"""
Calculate cumulative return and CAGR given the senators portfolio
value over time.
"""
before_values.pop(0)
after_values.pop(-1)
# We calculate the total return by calculating the return
# between each transaction, and solving for the cumulative
# return.
growth = np.array(before_values) / np.array(after_values)
portfolio_return = np.prod(growth[~np.isnan(growth)])
years = (end_date - begin_date).days / 365
if years == 0:
cagr = 0
else:
cagr = portfolio_return**(1 / years)
# DataFrame of cumulative return
tx_dates.pop(0)
tx_dates = np.array(tx_dates)
tx_dates = tx_dates[~np.isnan(growth)]
cumulative_growth = np.cumprod(growth[~np.isnan(growth)])
growth_df = pd.DataFrame({
'date': tx_dates,
'cumulative_growth': cumulative_growth
})
return {
'portfolio_return': portfolio_return,
'annual_cagr': cagr,
'growth': growth_df
}

def return_for_senator(rows, date_col='tx_date'):
"""
Simulate a senator's buy and sell orders, and calculate the
return.
"""
stocks = defaultdict(int)
# Value of portfolio at various timepoints to calculate return
portfolio_value_before_tx = []
portfolio_value_after_tx = []
tx_dates = []
rows = rows.sort_values(by=date_col)
for _, row in rows.iterrows():
date = row[date_col]
if date_col == 'file_date':
# We can't execute the trade the same day
date += dt.timedelta(days=1)
try:
stock_price = ticker_at_date(row['ticker'], date)
except IndexError as e:
# Skip the row if we're missing ticker data
continue
value_before_tx = portfolio_value(stocks, date)
if 'Purchase' in row['order_type']:
tx_amt = row['tx_estimate']
n_units = tx_amt / ticker_at_date(row['ticker'], date)
stocks[row['ticker']] += n_units
elif 'Sale' in row['order_type']:
current_value = stock_price * stocks[row['ticker']]
if 'Full' in row['order_type'] or \
same_bucket(row['tx_estimate'], current_value):
stocks[row['ticker']] = 0
else:
new_n_units = stocks[row['ticker']] -\
row['tx_estimate'] / stock_price
stocks[row['ticker']] = max(0, new_n_units)
portfolio_value_before_tx.append(value_before_tx)
portfolio_value_after_tx.append(portfolio_value(stocks, date))
tx_dates.append(date)
return calculate_return(
portfolio_value_before_tx,
portfolio_value_after_tx,
begin_date=min(rows[date_col]),
end_date=max(rows[date_col]),
tx_dates=tx_dates
)

# In[ ]:

senator_returns = []
senator_tx_growth = {}
senator_file_growth = {}
senator_names = set(senators_tx['full_name'])

# The following cell took my laptop about three hours to run.

# In[ ]:

failed_senators = {}
print('{} senators total'.format(len(senator_names)))
for n in senator_names:
print('Starting {}'.format(n))
if n in senator_tx_growth:
# Don't re-calculate for a given senator
continue
try:
tx_return = return_for_senator(
senators_tx[senators_tx['full_name'] == n],
date_col='tx_date')
file_return = return_for_senator(
senators_tx[senators_tx['full_name'] == n],
date_col='file_date')
senator_returns.append({
'full_name': n,
'tx_total_return': tx_return['portfolio_return'],
'tx_cagr': tx_return['annual_cagr'],
'file_total_return': file_return['portfolio_return'],
'file_cagr': file_return['annual_cagr']
})
senator_tx_growth[n] = tx_return['growth']
senator_file_growth[n] = file_return['growth']
except Exception as e:
print('Failed senator {0} with exception {1}'.format(n, e))
failed_senators[n] = e

# We look at the results to see the senators that outperformed the market.

# In[ ]:

def plot_senator_growth(growth):
""" Plot the senator's portfolio growth against the S&P 500. """
plt.plot_date(growth['date'], growth['cumulative_growth'], '-')
spy = spy[(spy['date'] >= min(growth['date']))
& (spy['date'] <= max(growth['date']))]
spy_prices = spy['price']
spy_growth = np.cumprod(np.diff(spy_prices) / spy_prices[1:] + 1)
dates = spy['date'].iloc[1:]
plt.plot_date(dates, spy_growth, '-')
plt.show()
print('Earliest date: {}'.format(min(growth['date'])))
print('Latest date: {}'.format(max(growth['date'])))
print('Market return: {}'.format(
spy_prices.iloc[-1] / spy_prices.iloc[0]))
senator_growth = growth['cumulative_growth']
print('Senator return: {}'.format(
senator_growth.iloc[-1] / senator_growth.iloc[0]))

# In[ ]:

returns = pd.DataFrame(senator_returns)
returns = returns[(returns['tx_total_return'] > returns['tx_cagr'])
& (returns['tx_cagr'] > 0)]
returns.sort_values('tx_cagr')

# In[ ]:

plot_senator_growth(senator_tx_growth['Angus S King, Jr.'])

#
# Author: Neel Somani, Software Engineer
#
# Email: neel@berkeley.edu
#
# Website: https://www.ocf.berkeley.edu/~neel/
#
# Updated On: 2020-05-10

• I see that downloading the file served you well. Nice question :) I have a minor suggestion. I personally would remove the final section "About this notebook" to not expose your email to the public. If you don't mind that, then please ignore me :) – Peilonrayz May 9 at 12:57
• Thanks! I don't mind - think it's pretty easy for people to find my email these days. – Neel May 10 at 0:28
• A recent websiite I look at - senatestockwatcher.com - is that what you're trying to do? – C. Harley May 10 at 17:12
• I think that site is displaying the exact information from the STOCK act filings, but it doesn't do things like calculate returns or simulate portfolios on top of it. Useful website, though - thanks for sharing! – Neel May 11 at 1:49

First thing I notice is your imports.

from collections import defaultdict
import datetime as dt
from functools import lru_cache
import json
from os import path
import pickle

import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import yfinance as yf


The good thing about your imports is they're listed alphabetically and all imports are on their own line. Can we improve this further? Yes. PEP8 wants us to split it into 3 groups:

• Standard library imports
• Related third party imports
• Local and library-specific imports

But, honestly, I'd reorder them like this:

import json
import pickle

import datetime as dt
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import yfinance as yf

from collections import defaultdict
from functools import lru_cache
from os import path


They're still alphabetically sorted, but now also sorted by how they're imported. Looks cleaner to me.

And I understand this is Jupyter, where you create functions whenever you need them and not a moment sooner, but you first open your file and define a lot of functions right after it. Then you define a function mentioning a blacklist that's only explained after that function is defined and then comes the global that actually does something with the content of the file we read on the start.

That looks awkward at best.

useful_cols is not a useful name when in the global scope. Had it be part of a function or method, it would make more sense. Now, what columns are we referring to? It is no table, so it must be a list of column headers. From an input file? Output file? An intermediary result? Can't tell by the name. Going by the styling of the rest of your project, even tx_headers would've been better.

calculate_return is a bit of a mess, but I'm not sure how to improve it. Having to call data like

senators_tx[senators_tx['full_name'] == n]


and

returns = returns[(returns['tx_total_return'] > returns['tx_cagr'])
& (returns['tx_cagr'] > 0)]


looks odd as well, perhaps your data structure isn't optimized for what you're doing with it. If re-arranging your data takes 10 minutes extra and cuts the execution time of all other processes in half, you already have a massive gain. I'd definitely take a look in that direction. How much data do you get, how much do you actually use and is it in an actually useful format?

Some of the returns could be more succinct, but most of them are fine and your error handling isn't too bad either. I think it would be beneficial to you if you move more to your code into functions, but you've done most of that already. For a Jupyter project it really doesn't look that bad.