2016 研究 缠论工具（笔， 线段）.ipynb

# 风险及免责提示：该策略由聚宽用户在聚宽社区分享，仅供学习交流使用。
# 原文一般包含策略说明，如有疑问请到原文和作者交流讨论。
# 原文网址：https://www.joinquant.com/view/community/detail/26842
# 标题：缠论工具（笔， 线段）

import numpy as np
import copy
import talib

from enum import Enum 

from numpy.lib.recfunctions import append_fields
from scipy.ndimage.interpolation import shift

class InclusionType(Enum):
    # output: 0 = no inclusion, 1 = first contains second, 2 second contains first
    noInclusion = 0
    firstCsecond = 2
    secondCfirst = 3
    
class TopBotType(Enum):
    noTopBot = 0
    bot2top = 0.5
    top = 1
    top2bot = -0.5
    bot = -1

    @classmethod
    def reverse(cls, tp):
        if tp == cls.top:
            return cls.bot
        elif tp == cls.bot:
            return cls.top
        elif tp == cls.top2bot:
            return cls.bot2top
        elif tp == cls.bot2top:
            return cls.top2bot
        else:
            return cls.noTopBot
    
    @classmethod
    def value2type(cls, val):
        if val == 0:
            return cls.noTopBot
        elif val == 0.5:
            return cls.bot2top
        elif val == 1:
            return cls.top
        elif val == -0.5:
            return cls.top2bot
        elif val == -1:
            return cls.bot
        else:
            return cls.noTopBot


######################## common method ###############################

def float_less(a, b):
    return a < b and not np.isclose(a, b)

def float_more(a, b):
    return a > b and not np.isclose(a, b)

def float_less_equal(a, b):
    return a < b or np.isclose(a, b)

def float_more_equal(a, b):
    return a > b or np.isclose(a, b)

######################## kBarprocessor #############################
GOLDEN_RATIO = 0.618
MIN_PRICE_UNIT=0.01
        
FEN_BI_COLUMNS = ['date', 'close', 'high', 'low', 'tb', 'real_loc']
FEN_DUAN_COLUMNS =  ['date', 'close', 'high', 'low', 'chan_price', 'tb', 'xd_tb', 'real_loc']

def gap_range_func(a):
    if float_more(a['low'] - a['high_s1'], MIN_PRICE_UNIT):
        return [a['high_s1'], a['low']]
    elif float_less(a['high'] - a['low_s1'], -MIN_PRICE_UNIT):
        return [a['high'], a['low_s1']]
    else:
        return [0, 0]

def get_previous_loc(loc, working_df):
    i = loc - 1
    while i >= 0:
        if working_df[i]['tb'] == TopBotType.top.value or working_df[i]['tb'] == TopBotType.bot.value:
            return i
        else:
            i = i - 1
    return None

def get_next_loc(loc, working_df):
    i = loc + 1
    while i < len(working_df):
        if working_df[i]['tb'] == TopBotType.top.value or working_df[i]['tb'] == TopBotType.bot.value:
            return i
        else:
            i = i + 1
    return None