5
\$\begingroup\$

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.

# ***

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

# In[ ]:

with open('senators.pickle', 'rb') as f:
    raw_senators_tx = pickle.load(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('.', '')
    return token_string.split(' ')

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:
    blacklist = set(json.load(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:
        unmapped_tickers.add(m)


# 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[ ]:

def download_for_ticker(ticker, check_cache=True):
    """ 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)

def load_for_ticker(ticker):
    """ Load the file of stock prices for this ticker. """
    with open('stocks/{0}.pickle'.format(ticker), 'rb') as f:
        dump = pickle.load(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:
        data = load_for_ticker(ticker)
        # Sell at the next opportunity possible
        return _price_for_date(data, date)
    except Exception:
        # If any exception occurs, refresh the cache
        download_for_ticker(ticker, check_cache=False)
        data = load_for_ticker(ticker)
        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:
        download_for_ticker(t)
    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'], '-')
    download_for_ticker('SPY')
    spy = load_for_ticker('SPY')
    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.'])


# ## About this notebook
#  
# Author: Neel Somani, Software Engineer
#  
# Email: neel@berkeley.edu
#  
# Website: https://www.ocf.berkeley.edu/~neel/
#  
# Updated On: 2020-05-10
\$\endgroup\$
  • \$\begingroup\$ 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 :) \$\endgroup\$ – Peilonrayz May 9 at 12:57
  • \$\begingroup\$ Thanks! I don't mind - think it's pretty easy for people to find my email these days. \$\endgroup\$ – Neel May 10 at 0:28
  • \$\begingroup\$ A recent websiite I look at - senatestockwatcher.com - is that what you're trying to do? \$\endgroup\$ – C. Harley May 10 at 17:12
  • \$\begingroup\$ 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! \$\endgroup\$ – Neel May 11 at 1:49
1
\$\begingroup\$

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.

| improve this answer | |
\$\endgroup\$

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Not the answer you're looking for? Browse other questions tagged or ask your own question.