3. 特征工程
任务目标
本案例以贵州茅台(600519)近一年的日线行情为原始数据,通过Python实现三类经典技术指标的计算:移动平均线(MA)、相对强弱指标(RSI)和指数平滑异同移动平均线(MACD)。完成计算后,将结果与原始行情合并输出为一份可直接用于后续机器学习建模的特征数据表。
环境与依赖
需要安装 pandas、numpy 和 akshare(用于获取真实行情数据)。若无法访问网络获取数据,可使用代码中内置的随机数据生成器作为替代输入,两种方式均已在代码中实现。
核心实现
代码包含三个核心函数:calc_ma 计算N日简单移动平均线;calc_rsi 采用Wilder平滑法计算RSI(比简单平均法更接近主流交易软件数值);calc_macd 计算DIF、DEA及柱状值。主函数 build_features 负责数据获取、指标计算与结果合并。
import pandas as pd
import numpy as np
def calc_ma(close, window):
"""计算简单移动平均线"""
return close.rolling(window=window, min_periods=window).mean()
def calc_rsi(close, period=14):
"""Wilder平滑法计算RSI"""
delta = close.diff()
gain = delta.clip(lower=0)
loss = -delta.clip(upper=0)
# Wilder平滑:使用ewm的alpha=1/period
avg_gain = gain.ewm(alpha=1/period, min_periods=period, adjust=False).mean()
avg_loss = loss.ewm(alpha=1/period, min_periods=period, adjust=False).mean()
rs = avg_gain / avg_loss
rsi = 100 - (100 / (1 + rs))
return rsi
def calc_macd(close, fast=12, slow=26, signal=9):
"""计算MACD的DIF、DEA和柱"""
ema_fast = close.ewm(span=fast, adjust=False).mean()
ema_slow = close.ewm(span=slow, adjust=False).mean()
dif = ema_fast - ema_slow
dea = dif.ewm(span=signal, adjust=False).mean()
macd_bar = 2 * (dif - dea)
return dif, dea, macd_bar
def build_features(data_source='akshare'):
"""主函数:获取数据并计算全部指标"""
if data_source == 'akshare':
import akshare as ak
df = ak.stock_zh_a_hist(symbol="600519", period="daily",
start_date="20240101", end_date="20241231",
adjust="qfq")
df = df.rename(columns={'日期': 'date', '收盘': 'close',
'开盘': 'open', '最高': 'high',
'最低': 'low', '成交量': 'volume'})
df = df[['date', 'open', 'high', 'low', 'close', 'volume']].copy()
else:
# 内置随机数据生成器(离线可用)
np.random.seed(42)
dates = pd.date_range('2024-01-01', periods=300, freq='D')
close = 1700 + np.cumsum(np.random.randn(300) * 15)
df = pd.DataFrame({'date': dates, 'close': close})
df['open'] = df['close'].shift(1).fillna(close[0])
df['high'] = df[['open', 'close']].max(axis=1) * 1.01
df['low'] = df[['open', 'close']].min(axis=1) * 0.99
df['volume'] = np.random.randint(30000, 80000, size=300)
# 计算指标
df['MA5'] = calc_ma(df['close'], 5)
df['MA10'] = calc_ma(df['close'], 10)
df['MA20'] = calc_ma(df['close'], 20)
df['RSI14'] = calc_rsi(df['close'], 14)
df['DIF'], df['DEA'], df['MACD_BAR'] = calc_macd(df['close'])
# 删除前20日因窗口不足产生的NaN行(MA20所需)
df_clean = df.iloc[20:].reset_index(drop=True)
return df_clean
if __name__ == "__main__":
# 默认使用akshare在线数据;若网络不可用,改为 build_features('random')
result = build_features('akshare')
print(result.tail(5))
print(f"\n特征矩阵形状: {result.shape}")
print(f"列名: {list(result.columns)}")
# 验证无NaN
assert result[['MA5', 'MA10', 'MA20', 'RSI14', 'DIF', 'DEA', 'MACD_BAR']].isnull().sum().sum() == 0
print("验证通过:所有指标列均无缺失值")
运行与验证
将代码保存为 feature_engineering.py,在终端执行 python feature_engineering.py。程序会输出最后5行数据预览、特征矩阵形状(应为280行×12列),并自动断言所有指标列无NaN值。若使用akshare数据源,输出中的 date 列应为2024年真实交易日;若使用随机数据,日期为连续自然日。
判定通过的三个条件:
- 程序无异常退出,打印的特征矩阵形状为
(280, 12); - 断言语句通过,即MA、RSI、MACD相关列不含缺失值;
- 目视检查
MA5与close列:MA5数值应大致跟随收盘价且滞后于价格变动,RSI14数值应始终在0~100区间内。
常见问题
若akshare数据获取失败(网络受限或接口变更),改用 build_features('random') 即可验证代码逻辑,但此时数据无真实市场含义。RSI计算采用Wilder平滑法,与部分财经网站使用简单平均法所得数值可能有1~3的差异,属正常现象。