Created
December 9, 2019 10:08
-
-
Save shellsong/28faa96dda8ee429a4c70734d5465be8 to your computer and use it in GitHub Desktop.
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
# %% | |
import scipy as sp | |
from scipy.signal import argrelextrema | |
import numpy as np | |
import pandas as pd | |
import tushare as ts | |
from bokeh.io import output_notebook, show | |
from bokeh.plotting import figure | |
output_notebook() | |
# %% | |
df = ts.get_hist_data('000001').pipe(lambda df: df.set_index(pd.to_datetime(df.index)).sort_index()) | |
df.index.name = 'index' | |
# %% | |
def restep(df): | |
pass | |
def plot_view(df): | |
df['interval'] = 24 * 60 * 60 * 1000 | |
df['color'] = 'white' | |
df.loc[df['open'] > df['close'], 'color'] = 'red' | |
df.loc[df['close'] > df['open'], 'color'] = 'green' | |
return df | |
def find_peak(df): | |
df['peak'] = 0 | |
df.loc[df.index[argrelextrema(df.low.values, np.less_equal, order=3)[0]], 'peak'] = -1 | |
df.loc[df.index[argrelextrema(df.high.values, np.greater_equal, order=3)[0]], 'peak'] = 1 | |
return df | |
# %% | |
def plot(df): | |
fig = figure(**{ | |
'x_axis_type': "datetime", | |
'plot_width': 12, | |
'plot_height': 6, | |
'sizing_mode': 'scale_width' | |
}) | |
fig.vbar(**{ | |
'x': 'index', | |
'width': 'interval', | |
'top': 'open', | |
'bottom': 'close', | |
'fill_color': 'color', | |
'line_color': 'white' | |
}, source=df) | |
fig.segment(**{ | |
'x0': 'index', | |
'x1': 'index', | |
'y0': 'high', | |
'y1': 'low', | |
'line_color': 'color' | |
}, source=df) | |
peak_min_indexer = df['peak'] == -1 | |
peak_max_indexer = df['peak'] == 1 | |
fig.circle(x=df[peak_min_indexer].index, y=df[peak_min_indexer].low, size=10, color='red') | |
fig.circle(x=df[peak_max_indexer].index, y=df[peak_max_indexer].high, size=10, color='green') | |
return fig | |
show(plot(df.pipe(find_peak).pipe(plot_view).iloc[-60:])) | |
# %% |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment