calculate_beta.py

'''
Simple function to calculate BETA ,ALPHA and R^2 for stocks
http://blog.aboveindex.com/
'''
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt

# statsmodels is used to do OLS regression
import statsmodels.api as sm

# we're using yahoo finance data, pandas datareader will import the data we need
from pandas_datareader.data import DataReader

def get_monthly_returns(symbols,start_date):
    # df is the main dataframe that'll hold the Adjusted closing prices
    df = pd.DataFrame() 

    for symbol in symbols:
        dftemp = DataReader(symbol,"yahoo",start_date)
        # we only need the adjusted close price.
        df[symbol] = dftemp["Adj Close"]
    
    dfm = df.resample('M').last()
    dfm_returns = dfm.pct_change()
    return dfm_returns
  
# wrapping up the code in a simple function
def calc_stats(df):
    '''
        requires a DataFrame of monthly returns where 
        first column is the stock
        second column is the market index 
    '''
    X = sm.add_constant(df.iloc[:,1])
    model = sm.OLS(df.iloc[:,0],X).fit()
    
    beta = model.params["^GSPC"]
    alpha = model.params["const"]
    r2 = model.rsquared
    return beta, alpha, r2
def rolling_stats(df, window=5):
    #     dataframe to hold the results
    res = pd.DataFrame(index=df.index)
    
    for i in xrange(0,len(df.index)):
    
        if len(df) - i >= window:
            # break the df into smaller chunks
            chunk = df.iloc[i:window+i,:]
            # calc_stats is a function created from the code above, 
            # refer to the Gist at the end of the article.
            beta,alpha,r2 = calc_stats(chunk)
            res.set_value(chunk.tail(1).index[0],"beta",beta)
            res.set_value(chunk.tail(1).index[0],"alpha",alpha)
            res.set_value(chunk.tail(1).index[0],"r2",r2)
            # print "%s beta: %.4f \t alpha: %.4f" % (chunk.tail(1).index[0],b,a)
    res = res.dropna()
    return res
  
 #testing the rolling stats
dftest = get_monthly_returns(["AAPL","^GSPC"],"2008-01-01")
df_rolling = rolling_stats(dftest, window=10)

print df_rolling.head(20)

# plotting the rolling beta
ax = df_rolling[["beta"]].plot(title="AAPL Rolling Beta")
ax.grid(True)
ax.legend(loc='upper left', ncol=2, fontsize='small')
plt.show()