peak.py

# %%
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:]))

# %%