chinese-equity-quant/pipeline/alpha/compute.py

"""Alpha computation and evaluation.

Alphas are position WEIGHTS — positive=long, negative=short. They are NOT
predictors of future returns. Concrete alphas are classes that live in
``pipeline/alpha/library/`` (or any external module) and are resolved by name
through :mod:`pipeline.alpha.registry`.
"""

import logging

import numpy as np
import pandas as pd

from pipeline.alpha.registry import get_alpha
from pipeline.common.schema import ALPHA_COLUMNS

logger = logging.getLogger(__name__)


def _pivot_close(df: pd.DataFrame) -> pd.DataFrame:
    """Pivot data to wide format: date index, columns = symbol_id, values = close."""
    pivot = df.pivot_table(
        index="date", columns="symbol_id", values="close", aggfunc="first"
    )
    return pivot.sort_index()


def _daily_returns(close: pd.DataFrame) -> pd.DataFrame:
    """Compute daily returns from wide close DataFrame."""
    return close.pct_change()


def compute_alpha(
    data: pd.DataFrame,
    alpha_name: str,
    alpha_type: str,
    **params,
) -> pd.DataFrame:
    """Compute alpha weights from raw data.

    Args:
        data: DataFrame with DATA_COLUMNS.
        alpha_name: Label stored in the ``alpha_name`` output column.
        alpha_type: Registry key of the alpha class (e.g. ``reversal``).
        **params: Constructor parameters for the alpha (e.g. ``lookback``,
            ``vol_window``). Only the params the alpha's ``__init__`` accepts are
            used; extras are ignored.

    Returns:
        DataFrame with ALPHA_COLUMNS.

    Raises:
        KeyError: If ``alpha_type`` is not registered.
    """
    alpha = get_alpha(alpha_type, **params)
    close = _pivot_close(data)
    weights = alpha.weights(close)

    # Melt to long format
    weights_melted = weights.reset_index().melt(
        id_vars="date", var_name="symbol_id", value_name="weight"
    )
    weights_melted["alpha_name"] = alpha_name
    weights_melted = weights_melted[ALPHA_COLUMNS]
    weights_melted = weights_melted.dropna(subset=["weight"])
    weights_melted = weights_melted.sort_values(["symbol_id", "date"]).reset_index(drop=True)

    logger.info(
        "Alpha '%s' (%r): %d symbols × %d dates, weight range [%.4f, %.4f]",
        alpha_name,
        alpha,
        weights_melted["symbol_id"].nunique(),
        weights_melted["date"].nunique(),
        weights_melted["weight"].min(),
        weights_melted["weight"].max(),
    )
    return weights_melted


def evaluate_alpha(alpha_df: pd.DataFrame, data_df: pd.DataFrame) -> dict:
    """Evaluate an alpha's performance as position weights.

    Computes return, annualized Sharpe, annualized turnover, max drawdown.

    Alpha is interpreted as POSITION WEIGHTS, not predictions.
    Return on date t = sum(weight[s,t] * realized_return[s,t]) / sum(abs(weight[s,t]))

    Args:
        alpha_df: DataFrame with ALPHA_COLUMNS.
        data_df: DataFrame with DATA_COLUMNS (for price data).

    Returns:
        Dict with metrics: cumulative_return, sharpe_annual, turnover_annual,
        max_drawdown, hit_rate, n_dates.
    """
    close = _pivot_close(data_df)
    returns = _daily_returns(close)

    # Pivot alpha weights to wide format
    weights = alpha_df.pivot_table(
        index="date", columns="symbol_id", values="weight", aggfunc="first"
    ).sort_index()

    # Align dates
    common_dates = weights.index.intersection(returns.index)
    weights = weights.loc[common_dates]
    returns = returns.loc[common_dates]

    if len(common_dates) < 2:
        return {
            "cumulative_return": 0.0,
            "sharpe_annual": 0.0,
            "turnover_annual": 0.0,
            "max_drawdown": 0.0,
            "hit_rate": 0.0,
            "n_dates": len(common_dates),
        }

    # Daily portfolio return = sum(w * r) / sum(|w|)  — normalized by gross exposure
    daily_returns = (weights * returns).sum(axis=1) / weights.abs().sum(axis=1)

    # Cumulative return
    cumulative_return = float((1.0 + daily_returns).prod() - 1.0)

    # Annualized Sharpe (sqrt(252) * mean / std)
    mu = daily_returns.mean()
    sigma = daily_returns.std()
    sharpe_annual = float(np.sqrt(252) * mu / sigma) if sigma > 0 else 0.0

    # Annualized turnover: avg daily turnover * 252
    # Daily turnover = sum(|w_t - w_{t-1}|) / sum(|w_{t-1}|)
    weight_change = weights.diff().abs().sum(axis=1)
    gross_exposure = weights.abs().sum(axis=1).shift(1)
    daily_turnover = weight_change / gross_exposure
    turnover_annual = float(daily_turnover.mean() * 252)

    # Max drawdown
    equity = (1.0 + daily_returns).cumprod()
    peak = equity.cummax()
    drawdown = (equity - peak) / peak
    max_drawdown = float(drawdown.min())

    # Hit rate
    hit_rate = float((daily_returns > 0).mean())

    return {
        "cumulative_return": cumulative_return,
        "sharpe_annual": sharpe_annual,
        "turnover_annual": turnover_annual,
        "max_drawdown": max_drawdown,
        "hit_rate": hit_rate,
        "n_dates": len(common_dates),
    }