232 lines
8.0 KiB
Python
232 lines
8.0 KiB
Python
#!/usr/bin/env python3
|
||
"""
|
||
从 CSV 导入 Alpha 数据到 SQLite（修复去重版）
|
||
"""
|
||
|
||
import sqlite3
|
||
import pandas as pd
|
||
import ast
|
||
import re
|
||
from pathlib import Path
|
||
|
||
# 路径配置
|
||
CSV_PATH = Path(__file__).parent / "alpha_list.csv"
|
||
SQLITE_PATH = Path(__file__).parent / "alpha_analysis.db"
|
||
|
||
# 84 个算子集合（同上，省略重复部分，你复制完整的）
|
||
OPERATORS_SET = {
|
||
'add', 'abs', 'log', 'subtract', 'signed_power', 'sign', 'reverse', 'power',
|
||
'multiply', 'min', 'max', 'inverse', 'sqrt', 's_log_1p', 'densify', 'divide',
|
||
'not', 'and', 'less', 'equal', 'or', 'not_equal', 'greater', 'greater_equal',
|
||
'less_equal', 'is_nan', 'if_else', 'ts_sum', 'ts_scale', 'ts_mean', 'ts_zscore',
|
||
'ts_std_dev', 'kth_element', 'inst_tvr', 'ts_corr', 'ts_count_nans',
|
||
'ts_target_tvr_decay', 'ts_median', 'ts_covariance', 'ts_decay_linear',
|
||
'ts_product', 'ts_regression', 'ts_delta_limit', 'ts_step', 'ts_decay_exp_window',
|
||
'ts_quantile', 'days_from_last_change', 'hump', 'last_diff_value', 'ts_arg_max',
|
||
'ts_arg_min', 'ts_av_diff', 'ts_backfill', 'ts_rank', 'ts_delay', 'ts_delta',
|
||
'winsorize', 'truncate', 'regression_neut', 'scale', 'rank', 'quantile',
|
||
'normalize', 'zscore', 'vec_min', 'vec_count', 'vec_stddev', 'vec_range',
|
||
'vec_avg', 'vec_sum', 'vec_max', 'left_tail', 'trade_when', 'right_tail',
|
||
'bucket', 'group_rank', 'group_cartesian_product', 'group_backfill', 'group_mean',
|
||
'group_neutralize', 'group_normalize', 'group_median', 'group_scale', 'group_zscore'
|
||
}
|
||
|
||
WINDOWS = {1, 2, 5, 10, 20, 30, 60, 90, 252}
|
||
|
||
|
||
def safe_parse_dict(s):
|
||
"""安全解析 Python 字典字符串"""
|
||
if pd.isna(s) or s == '':
|
||
return {}
|
||
try:
|
||
return ast.literal_eval(s)
|
||
except:
|
||
return {}
|
||
|
||
|
||
def extract_features(expression):
|
||
"""从表达式提取特征"""
|
||
if not expression:
|
||
return [], [], []
|
||
|
||
operators = set()
|
||
for match in re.finditer(r'\b([a-z_][a-z0-9_]*)\s*\(', expression):
|
||
op = match.group(1)
|
||
if op in OPERATORS_SET:
|
||
operators.add(op)
|
||
|
||
candidates = set(re.findall(r'\b([a-z][a-z0-9_]*)\b', expression))
|
||
fields = set()
|
||
for c in candidates:
|
||
if (c not in OPERATORS_SET and
|
||
not c.isdigit() and
|
||
len(c) > 2 and
|
||
c not in ['true', 'false', 'nan', 'null', 'constant', 'filter']):
|
||
fields.add(c)
|
||
|
||
windows = set()
|
||
for match in re.finditer(r'\b(\d+)\b', expression):
|
||
num = int(match.group(1))
|
||
if num in WINDOWS:
|
||
windows.add(num)
|
||
|
||
return list(operators), list(fields), list(windows)
|
||
|
||
|
||
def main():
|
||
print("=" * 60)
|
||
print("开始导入 CSV 到 SQLite")
|
||
print("=" * 60)
|
||
|
||
# 1. 读取 CSV
|
||
print("\n📖 读取 CSV...")
|
||
df = pd.read_csv(CSV_PATH)
|
||
print(f" 共 {len(df):,} 条记录")
|
||
|
||
# 2. 解析字段
|
||
print("\n🔍 解析字段...")
|
||
|
||
df['regular_dict'] = df['regular'].apply(safe_parse_dict)
|
||
df['expression'] = df['regular_dict'].apply(lambda x: x.get('code', ''))
|
||
|
||
df['settings_dict'] = df['settings'].apply(safe_parse_dict)
|
||
df['neutralization'] = df['settings_dict'].apply(lambda x: x.get('neutralization'))
|
||
df['universe'] = df['settings_dict'].apply(lambda x: x.get('universe'))
|
||
|
||
df['is_dict'] = df['is'].apply(safe_parse_dict)
|
||
df['sharpe'] = df['is_dict'].apply(lambda x: x.get('sharpe'))
|
||
df['fitness'] = df['is_dict'].apply(lambda x: x.get('fitness'))
|
||
df['drawdown'] = df['is_dict'].apply(lambda x: x.get('drawdown'))
|
||
df['turnover'] = df['is_dict'].apply(lambda x: x.get('turnover'))
|
||
|
||
# 过滤有效数据
|
||
df_valid = df[df['expression'].notna() & (df['expression'] != '') & df['fitness'].notna()].copy()
|
||
print(f" 有效记录: {len(df_valid):,} 条")
|
||
|
||
# 3. 去重：同一个 alpha_id 保留 fitness 最高的
|
||
print("\n🔄 去重（按 alpha_id 保留最高 fitness）...")
|
||
df_unique = df_valid.sort_values('fitness', ascending=False).drop_duplicates(subset=['id'], keep='first')
|
||
print(f" 去重后: {len(df_unique):,} 条")
|
||
|
||
# 4. 连接 SQLite
|
||
conn = sqlite3.connect(str(SQLITE_PATH))
|
||
cursor = conn.cursor()
|
||
|
||
# 清空旧数据
|
||
print("\n🗑️ 清空旧数据...")
|
||
cursor.execute("DELETE FROM alpha_feature_long")
|
||
cursor.execute("DELETE FROM alpha_success")
|
||
conn.commit()
|
||
|
||
# 5. 写入 alpha_success
|
||
print("\n💾 写入 alpha_success 表...")
|
||
|
||
for idx, row in df_unique.iterrows():
|
||
try:
|
||
cursor.execute("""
|
||
INSERT INTO alpha_success
|
||
(alpha_id, expression, sharpe, fitness, drawdown, turnover, neutralization, universe)
|
||
VALUES (?, ?, ?, ?, ?, ?, ?, ?)
|
||
""", (
|
||
row['id'],
|
||
row['expression'],
|
||
row['sharpe'],
|
||
row['fitness'],
|
||
row['drawdown'],
|
||
row['turnover'],
|
||
row['neutralization'],
|
||
row['universe']
|
||
))
|
||
except sqlite3.IntegrityError:
|
||
# 如果还有重复，跳过（理论上不应该）
|
||
print(f" 跳过重复: {row['id']}")
|
||
continue
|
||
|
||
conn.commit()
|
||
|
||
# 获取插入后的 id 映射
|
||
cursor.execute("SELECT alpha_id, rowid FROM alpha_success")
|
||
id_map = {row[0]: row[1] for row in cursor.fetchall()}
|
||
print(f" 已写入 {len(id_map):,} 条")
|
||
|
||
# 6. 提取特征并写入 alpha_feature_long
|
||
print("\n🔧 提取特征并写入 alpha_feature_long...")
|
||
|
||
feature_count = 0
|
||
processed = 0
|
||
|
||
for idx, row in df_unique.iterrows():
|
||
alpha_id = row['id']
|
||
if alpha_id not in id_map:
|
||
continue
|
||
|
||
sqlite_id = id_map[alpha_id]
|
||
expression = row['expression']
|
||
|
||
operators, fields, windows = extract_features(expression)
|
||
|
||
# 插入算子
|
||
for op in operators:
|
||
cursor.execute(
|
||
"INSERT INTO alpha_feature_long (alpha_id, feature_type, feature_name) VALUES (?, ?, ?)",
|
||
(sqlite_id, 'operator', op)
|
||
)
|
||
feature_count += 1
|
||
|
||
# 插入字段
|
||
for field in fields:
|
||
cursor.execute(
|
||
"INSERT INTO alpha_feature_long (alpha_id, feature_type, feature_name) VALUES (?, ?, ?)",
|
||
(sqlite_id, 'field', field)
|
||
)
|
||
feature_count += 1
|
||
|
||
# 插入窗口
|
||
for w in windows:
|
||
cursor.execute(
|
||
"INSERT INTO alpha_feature_long (alpha_id, feature_type, feature_name, feature_value) VALUES (?, ?, ?, ?)",
|
||
(sqlite_id, 'window', str(w), w)
|
||
)
|
||
feature_count += 1
|
||
|
||
# 插入中性化方式
|
||
if pd.notna(row['neutralization']):
|
||
cursor.execute(
|
||
"INSERT INTO alpha_feature_long (alpha_id, feature_type, feature_name) VALUES (?, ?, ?)",
|
||
(sqlite_id, 'neutralization', row['neutralization'])
|
||
)
|
||
feature_count += 1
|
||
|
||
processed += 1
|
||
if processed % 1000 == 0:
|
||
print(f" 已处理 {processed:,} 条，特征数 {feature_count:,}")
|
||
conn.commit() # 每 1000 条提交一次
|
||
|
||
conn.commit()
|
||
|
||
# 7. 最终统计
|
||
cursor.execute("SELECT COUNT(*) FROM alpha_success")
|
||
success_count = cursor.fetchone()[0]
|
||
cursor.execute("SELECT COUNT(*) FROM alpha_feature_long")
|
||
final_feature_count = cursor.fetchone()[0]
|
||
|
||
print(f"\n✅ 导入完成!")
|
||
print(f" alpha_success: {success_count:,} 条")
|
||
print(f" alpha_feature_long: {final_feature_count:,} 条特征")
|
||
|
||
# 8. 验证一条数据
|
||
print("\n🔍 验证前 10 条特征:")
|
||
cursor.execute("""
|
||
SELECT a.alpha_id, f.feature_type, f.feature_name
|
||
FROM alpha_success a
|
||
JOIN alpha_feature_long f ON a.rowid = f.alpha_id
|
||
LIMIT 10
|
||
""")
|
||
for row in cursor.fetchall():
|
||
print(f" {row[0]} | {row[1]}: {row[2]}")
|
||
|
||
conn.close()
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main() |