Alpha Mine

# Default parameters (user can override)
N_CANDIDATES = 50        # Number of candidates to generate
N_TOP = 10               # Number of top candidates to fully evaluate
HOLDING_PERIODS = [5, 10, 20]
IC_THRESHOLD = 0.02      # Minimum |IC| to pass quick screen
CATEGORY = "all"         # "all", "momentum", "mean_reversion", "volatility", "volume", "composite"

FIELDS = {
    "close": "close",           # 收盘价
    "open": "open_price",       # 开盘价 (if available)
    "high": "high",             # 最高价
    "low": "low",               # 最低价
    "volume": "volume",         # 成交量
    "returns": "close.pct_change(1)",  # 日收益率
}

TS_OPS = {
    "ts_mean":    lambda x, d: f"{x}.rolling({d}).mean()",
    "ts_std":     lambda x, d: f"{x}.rolling({d}).std()",
    "ts_max":     lambda x, d: f"{x}.rolling({d}).max()",
    "ts_min":     lambda x, d: f"{x}.rolling({d}).min()",
    "ts_rank":    lambda x, d: f"{x}.rolling({d}).apply(lambda s: s.rank(pct=True).iloc[-1])",
    "ts_delta":   lambda x, d: f"{x}.diff({d})",
    "ts_return":  lambda x, d: f"{x}.pct_change({d})",
    "ts_corr":    lambda x, y, d: f"{x}.rolling({d}).corr({y})",
    "ts_sum":     lambda x, d: f"{x}.rolling({d}).sum()",
    "ts_decay":   lambda x, d: f"({x} * pd.Series(range(1,{d}+1))).rolling({d}).sum() / pd.Series(range(1,{d}+1)).sum()",
}

CS_OPS = {
    "cs_rank":    lambda x: f"{x}.rank(axis=1, pct=True)",
    "cs_zscore":  lambda x: f"({x}.sub({x}.mean(axis=1), axis=0)).div({x}.std(axis=1), axis=0)",
}

ARITH = {
    "neg":   lambda x: f"-({x})",
    "abs":   lambda x: f"({x}).abs()",
    "log":   lambda x: f"np.log(({x}).clip(lower=1e-10))",
    "square": lambda x: f"({x})**2",
    "sign":  lambda x: f"np.sign({x})",
}

WINDOWS = [5, 10, 20, 40, 60]

TEMPLATES = {
    "momentum": [
        # 不同周期的动量
        ("ts_return(close, {w})", "Momentum {w}d"),
        # 跳跃动量（避免短期反转噪音）
        ("ts_return(close.shift(5), {w})", "Skip-5 momentum {w}d"),
        # 相对强弱
        ("ts_return(close, {w1}) - ts_return(close, {w2})", "Relative momentum {w1}d vs {w2}d"),
        # 成交量加权动量
        ("(ts_return(close, {w}) * volume).rolling({w}).sum() / volume.rolling({w}).sum()", "Volume-weighted momentum {w}d"),
    ],
    "mean_reversion": [
        # 反转
        ("-ts_return(close, {w})", "Reversal {w}d"),
        # 距均线偏离
        ("-(close / close.rolling({w}).mean() - 1)", "Mean reversion to MA{w}"),
        # RSI变体
        ("-(close.diff().clip(lower=0).rolling({w}).mean() / (-close.diff().clip(upper=0)).rolling({w}).mean())", "RSI-like {w}d"),
        # 布林带位置
        ("-((close - close.rolling({w}).mean()) / close.rolling({w}).std())", "Bollinger z-score {w}d"),
    ],
    "volatility": [
        # 已实现波动率
        ("-(close.pct_change().rolling({w}).std() * np.sqrt(252))", "Low volatility {w}d"),
        # 高低价比率
        ("-((high / low - 1).rolling({w}).mean())", "Low HL ratio {w}d"),
        # 波动率变化
        ("-(close.pct_change().rolling({w1}).std() / close.pct_change().rolling({w2}).std())", "Vol change {w1}d/{w2}d"),
        # 下行波动率
        ("-(close.pct_change().clip(upper=0).rolling({w}).std())", "Low downside vol {w}d"),
    ],
    "volume": [
        # 量价背离
        ("-(close.pct_change().rolling({w}).corr(volume.pct_change()))", "PV divergence {w}d"),
        # 换手率
        ("-(volume.rolling({w}).mean())", "Low turnover {w}d (proxy)"),
        # 异常成交量
        ("-(volume.rolling({w1}).mean() / volume.rolling({w2}).mean() - 1)", "Abnormal volume {w1}d/{w2}d"),
        # 成交量趋势
        ("volume.rolling({w1}).mean() / volume.rolling({w2}).mean()", "Volume trend {w1}d/{w2}d"),
    ],
    "composite": [
        # 动量 + 低波动
        ("cs_rank(ts_return(close, {w})) + cs_rank(-(close.pct_change().rolling({w}).std()))", "Momentum+LowVol {w}d"),
        # 反转 + 量价背离
        ("cs_rank(-ts_return(close, {w1})) + cs_rank(-(close.pct_change().rolling({w2}).corr(volume.pct_change())))", "Reversal+PVDiv {w1}d/{w2}d"),
        # 多维动量
        ("cs_rank(ts_return(close, {w1})) * 0.5 + cs_rank(ts_return(close, {w2})) * 0.3 + cs_rank(-(close.pct_change().rolling(20).std())) * 0.2", "MultiMom {w1}d+{w2}d+LowVol"),
    ],
}

import itertools
import random

def generate_candidates(category="all", n_candidates=50):
    """Generate candidate factor expressions"""
    candidates = []
    
    cats = TEMPLATES.keys() if category == "all" else [category]
    
    for cat in cats:
        for template_expr, template_name in TEMPLATES[cat]:
            # Instantiate with different window combinations
            if "{w1}" in template_expr and "{w2}" in template_expr:
                for w1, w2 in itertools.combinations(WINDOWS, 2):
                    expr = template_expr.format(w1=w1, w2=w2)
                    name = template_name.format(w1=w1, w2=w2)
                    candidates.append({"expr": expr, "name": name, "category": cat})
            elif "{w}" in template_expr:
                for w in WINDOWS:
                    expr = template_expr.format(w=w)
                    name = template_name.format(w=w)
                    candidates.append({"expr": expr, "name": name, "category": cat})
    
    # Shuffle and limit
    random.shuffle(candidates)
    return candidates[:n_candidates]

from scipy import stats

def quick_screen(expr, close, volume, high, low, forward_returns, ic_threshold=0.02):
    """
    Quick IC screen for a candidate expression.
    Returns (ic_mean, ic_std, passed) or None if computation fails.
    """
    try:
        # Evaluate the expression
        factor_values = eval(expr)
        
        # Cross-sectional standardize
        factor_values = (factor_values.sub(factor_values.mean(axis=1), axis=0)
                        .div(factor_values.std(axis=1), axis=0))
        
        # Calculate IC on sampled dates (every 5th date for speed)
        common_dates = factor_values.index.intersection(forward_returns.index)[::5]
        common_stocks = factor_values.columns.intersection(forward_returns.columns)
        
        ic_values = []
        for date in common_dates:
            f = factor_values.loc[date, common_stocks].dropna()
            r = forward_returns.loc[date, common_stocks].dropna()
            common = f.index.intersection(r.index)
            if len(common) < 30:
                continue
            fv, rv = f[common].values, r[common].values
            valid = np.isfinite(fv) & np.isfinite(rv)
            if valid.sum() < 30:
                continue
            corr, _ = stats.spearmanr(fv[valid], rv[valid])
            if np.isfinite(corr):
                ic_values.append(corr)
        
        if len(ic_values) < 10:
            return None
        
        ic_mean = np.mean(ic_values)
        ic_std = np.std(ic_values)
        icir = ic_mean / ic_std if ic_std > 0 else 0
        passed = abs(ic_mean) > ic_threshold
        
        return {"ic_mean": ic_mean, "ic_std": ic_std, "icir": icir, "passed": passed}
    except Exception:
        return None