from pandas_composition import UserFrame, UserSeries
import pandas as pd
import pandas.io.data as pdd
import IPython
IPython.core.pylabtools.figsize(20, 15)
from matplotlib.finance import candlestick
from itertools import izip
def ohlc(df, width=0.3, secondary_y=False,
*args, **kwargs):
"""
Takes a df and plots a candlestick.
"""
dates = df.index.to_pydatetime()
times = date2num(dates)
quotes = izip(times, df['Open'], df['Close'], df['High'], df['Low'])
ax = gca()
candlestick(ax, quotes, width=width, colorup='g')
return ax
We want a subclass that knows what datasource it was derived from and can plot accordingly.
class StockData(UserFrame):
def __init__(self, *args, **kwargs):
symbol = kwargs.get('symbol', None)
self.symbol = symbol
def ohlc(self):
" plot candlestick "
ax = ohlc(self, label=self.symbol)
ax.set_title(self.symbol)
return ax
def get_gaps(df, offset=0):
" when open is above or below the previous days range "
gap_up = df.Open > (df.High.shift(1) + offset)
gap_down = df.Open < (df.Low.shift(1) - offset)
gaps = gap_up | gap_down
return Indicator(gaps, source=df, name="Gaps")
def ma(self, period):
" moving average "
ma = pd.rolling_mean(df.Close, period)
ma = Indicator(ma, source=self, name='MA {period}'.format(period=period))
return ma
class Indicator(UserSeries):
def __init__(self, *args, **kwargs):
source = kwargs.pop('source', None)
self.source = source
def plot(self, source_col='Close', *args, **kwargs):
" plot the indicator. source_col only used for markers "
df = self.source
label = kwargs.get('label', None) or self.name # use Series.name
kwargs['label'] = label
if self.dtype == bool: # plot markers for bool
sig = self * df[source_col]
sig[sig == 0] = np.nan
plot_date(sig.index.to_pydatetime(), sig.values, *args, **kwargs)
else:
plot(self.index.to_pydatetime(), self, *args, **kwargs)
legend()
df = pdd.get_data_yahoo('AAPL')
df = StockData(df, symbol='AAPL')
df.tail(30).ohlc()
df.get_gaps().tail(30).plot('Open', color='pink', markersize=10)
df.ma(30).tail(30).plot()