158 lines
5.0 KiB
Python
158 lines
5.0 KiB
Python
import requests
|
|
import pandas as pd
|
|
import numpy as np
|
|
|
|
CN_CODE_TYPE = {
|
|
"600":'sh',
|
|
"601":'sh',
|
|
"603":'sh',
|
|
"605":'sh',
|
|
"000":'sz',
|
|
"300":'sz',
|
|
"688":'sh',
|
|
"002":'sz',
|
|
"50":"sh",
|
|
"51":"sh",
|
|
"52":"sh",
|
|
"16":"sz",
|
|
"15":"sz",
|
|
"18":"sz"
|
|
}
|
|
|
|
## arbr指标
|
|
def calculate_brar(df: pd.DataFrame, N: int = 26) -> pd.DataFrame:
|
|
"""
|
|
计算AR和BR情绪指标
|
|
|
|
参数:
|
|
df (pd.DataFrame): 包含 'high', 'low', 'open', 'close' 列的DataFrame
|
|
N (int): 计算周期，默认通常为26
|
|
返回:
|
|
pd.DataFrame: 包含原始数据及计算出的 'ar' 和 'br' 列的DataFrame
|
|
"""
|
|
# 创建副本，避免修改原始数据
|
|
data = df.copy()
|
|
|
|
# --- 计算 AR (人气指标) ---
|
|
# 公式: AR = (N日内 (最高价 - 开盘价) 之和) / (N日内 (开盘价 - 最低价) 之和) * 100
|
|
high_open = data['high'] - data['open']
|
|
open_low = data['open'] - data['low']
|
|
|
|
data['ar'] = (high_open.rolling(window=N, min_periods=1).sum() /
|
|
open_low.rolling(window=N, min_periods=1).sum()) * 100
|
|
|
|
# --- 计算 BR (意愿指标) ---
|
|
# 公式: BR = (N日内 MAX(0, 最高价 - 前一日收盘价) 之和) /
|
|
# (N日内 MAX(0, 前一日收盘价 - 最低价) 之和) * 100
|
|
prev_close = data['close'].shift(1)
|
|
|
|
# 计算多头意愿：当日最高价 - 前一日收盘价，若小于0则记为0
|
|
high_prev_close = data['high'] - prev_close
|
|
high_prev_close = high_prev_close.where(high_prev_close > 0, 0)
|
|
|
|
# 计算空头意愿：前一日收盘价 - 当日最低价，若小于0则记为0
|
|
prev_close_low = prev_close - data['low']
|
|
prev_close_low = prev_close_low.where(prev_close_low > 0, 0)
|
|
|
|
data['br'] = (high_prev_close.rolling(window=N, min_periods=1).sum() /
|
|
prev_close_low.rolling(window=N, min_periods=1).sum()) * 100
|
|
|
|
return data
|
|
|
|
## kdj计算
|
|
def calc_kdj_tdx(df, n=9, m1=3, m2=3):
|
|
# 1. 计算 RSV
|
|
low_min = df['low'].rolling(n).min()
|
|
high_max = df['high'].rolling(n).max()
|
|
rsv = (df['close'] - low_min) / (high_max - low_min) * 100
|
|
rsv = rsv.fillna(0) # 或者处理 NaN
|
|
|
|
# 2. 自定义 SMA 函数（通达信风格）
|
|
def sma_tdx(series, period, weight):
|
|
# period=N, weight=M。公式：Y = (X*M + PREV_Y*(N-M))/N
|
|
result = np.zeros(len(series))
|
|
# 初始化第一个有效值。通常，如果数据足够，初始值为 series[0]
|
|
result[0] = series[0] if not np.isnan(series[0]) else 0
|
|
for i in range(1, len(series)):
|
|
# 如果 series[i] 是 NaN（例如，前 n 天），则保持 0 或向前填充，但通常 RSV 已填充。
|
|
val = series[i] if not np.isnan(series[i]) else 0
|
|
result[i] = (val * weight + result[i - 1] * (period - weight)) / period
|
|
return result
|
|
|
|
# 3. 计算 K 和 D
|
|
k_series = sma_tdx(rsv.values, m1, 1) # K = SMA(RSV, 3, 1)
|
|
d_series = sma_tdx(k_series, m2, 1) # D = SMA(K, 3, 1)
|
|
|
|
df['K'] = k_series
|
|
df['D'] = d_series
|
|
df['J'] = 3 * df['K'] - 2 * df['D']
|
|
|
|
return df
|
|
|
|
## 获取历史k线
|
|
def get_history_k(code, start_date='19700101', end_date=None):
|
|
|
|
if '.' in code:
|
|
code = code.split('.')[0]
|
|
|
|
stock_code = ''
|
|
|
|
if code[0:2] in CN_CODE_TYPE:
|
|
stock_code = code + '.' + CN_CODE_TYPE[code[0:2]]
|
|
elif code[0:3] in CN_CODE_TYPE:
|
|
stock_code = code + '.' + CN_CODE_TYPE[code[0:3]]
|
|
if stock_code != '':
|
|
|
|
url = "https://api.zhituapi.com/hs/history/%s/d/f?token=6E0E86BC-15AD-4275-8A95-B02D168D63C1&st=%s"%(stock_code, start_date)
|
|
|
|
if end_date is not None:
|
|
url += "&et=%s"%(end_date)
|
|
response = requests.get(url)
|
|
|
|
data = response.json()
|
|
df = pd.DataFrame(data)
|
|
|
|
df['t'] = df['t'].apply(pd.to_datetime)
|
|
df.rename({'t': 'date', 'o': 'open', 'h': 'high', 'l': 'low', 'c': 'close', 'v': 'volume', 'a': 'amount'},
|
|
axis='columns', inplace=True)
|
|
df = df[['date', 'open', 'high', 'low', 'close', 'volume', 'amount']].set_index('date')
|
|
|
|
else:
|
|
df = pd.DataFrame()
|
|
|
|
|
|
return df
|
|
|
|
|
|
## 获取当日5分钟的k线
|
|
def get_5min_k(code, start_date='19700101', end_date=None):
|
|
stock_code = ''
|
|
|
|
if code[0:2] in CN_CODE_TYPE:
|
|
stock_code = code + '.' + CN_CODE_TYPE[code[0:2]]
|
|
elif code[0:3] in CN_CODE_TYPE:
|
|
stock_code = code + '.' + CN_CODE_TYPE[code[0:3]]
|
|
|
|
if stock_code != '':
|
|
|
|
url = "https://api.zhituapi.com/hs/history/%s/5/n?token=6E0E86BC-15AD-4275-8A95-B02D168D63C1&st=%s"%(stock_code, start_date)
|
|
|
|
if end_date is not None:
|
|
url += "&end_date=%s"%(end_date)
|
|
|
|
response = requests.get(url)
|
|
|
|
data = response.json()
|
|
|
|
df = pd.DataFrame(data)
|
|
df['t'] = df['t'].apply(pd.to_datetime)
|
|
df.rename({'t': 'date', 'o': 'open', 'h': 'high', 'l': 'low', 'c': 'close', 'v': 'volume', 'a': 'amount'},
|
|
axis='columns', inplace=True)
|
|
df = df[['date','close', 'volume']].set_index('date')
|
|
|
|
else:
|
|
df = pd.DataFrame()
|
|
|
|
|
|
return df
|