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
from pathlib import Path
from typing import Iterable

import numpy as np
import pandas as pd

from pipeline.alpha.registry import get_alpha
from pipeline.common.schema import ALPHA_COLUMNS
from pipeline.derived.compute import (
    DERIVED_KEY_COLUMNS,
    read_derived_frames,
    validate_derived_frame,
)

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 _pivot_open(df: pd.DataFrame) -> pd.DataFrame:
    """Pivot data to wide format: date index, columns = symbol_id, values = open."""
    pivot = df.pivot_table(
        index="date", columns="symbol_id", values="open", aggfunc="first"
    )
    return pivot.sort_index()


def join_feature_frames(
    data: pd.DataFrame,
    feature_frames: Iterable[pd.DataFrame],
) -> pd.DataFrame:
    """Left-join validated daily derived/feature frames onto long daily data."""
    out = data.copy()
    out["date"] = pd.to_datetime(out["date"])
    existing = set(out.columns)
    joined_cols: list[str] = []

    for frame in feature_frames:
        features = validate_derived_frame(frame)
        feature_cols = [col for col in features.columns if col not in DERIVED_KEY_COLUMNS]
        overlap = sorted(existing.intersection(feature_cols))
        if overlap:
            raise ValueError(
                f"Feature columns conflict with existing daily data columns: {overlap}"
            )
        out = out.merge(
            features,
            on=DERIVED_KEY_COLUMNS,
            how="left",
            validate="many_to_one",
        )
        existing.update(feature_cols)
        joined_cols.extend(feature_cols)

    if joined_cols:
        logger.info("Joined feature columns into daily data: %s", joined_cols)
    return out


def _forward_open_to_open_returns(open_: pd.DataFrame) -> pd.DataFrame:
    """Return earned by a close-formed signal after next-open execution.

    A weight formed after close on date t can first be traded at open[t+1].
    With daily retargeting it is then held until open[t+2], so the signal-date
    forward return is open[t+2] / open[t+1] - 1.
    """
    return open_.shift(-2).divide(open_.shift(-1)) - 1.0


def investable_universe_mask(
    data: pd.DataFrame,
    template: pd.DataFrame,
    *,
    top_n: int = 1000,
    min_history: int = 60,
    require_tradable: bool = True,
    exclude_st: bool = True,
) -> pd.DataFrame:
    """Build a per-date investable-universe mask aligned to ``template``.

    A ``(date, symbol_id)`` cell is ``True`` when the name is, on that date,
    seasoned (at least ``min_history`` prior closes), currently tradable
    (``tradestatus == 1``), not flagged ST (``isST == 0``), and inside the
    ``top_n`` most liquid names by trailing 20-day mean ``amount``. The mask is
    applied to the *signal* (computed on full contiguous prices), so it
    restricts only what is *held*, never the price history used to form the
    signal — that keeps ``pct_change`` correct and look-ahead free.

    Args:
        data: Long DataFrame with at least ``symbol_id``, ``date``, ``close``,
            ``amount``, ``isST``, ``tradestatus``.
        template: Wide signal (date index × ``symbol_id`` columns) to align to.
        top_n: Keep this many most-liquid names per date.
        min_history: Minimum number of observed closes before a name is eligible.
        require_tradable: Require ``tradestatus == 1`` on the date.
        exclude_st: Drop names flagged ``isST == 1``.

    Returns:
        Boolean wide DataFrame aligned to ``template``.
    """
    def _wide(col: str) -> pd.DataFrame:
        return (
            data.pivot_table(index="date", columns="symbol_id", values=col, aggfunc="first")
            .sort_index()
            .reindex(index=template.index, columns=template.columns)
        )

    close = _wide("close")
    mask = close.notna()

    seasoned = close.notna().cumsum() >= min_history
    mask &= seasoned

    if exclude_st and "isST" in data.columns:
        mask &= _wide("isST").fillna(1) == 0
    if require_tradable and "tradestatus" in data.columns:
        mask &= _wide("tradestatus").fillna(0) == 1

    amount = _wide("amount")
    amt_ma = amount.rolling(20, min_periods=10).mean()
    liquid_rank = amt_ma.rank(axis=1, ascending=False)
    mask &= liquid_rank <= top_n

    return mask.fillna(False)


def compute_alpha(
    data: pd.DataFrame,
    alpha_name: str,
    alpha_type: str,
    universe: dict | None = None,
    feature_paths: Iterable[str | Path] | None = None,
    feature_frames: Iterable[pd.DataFrame] | None = None,
    **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``).
        universe: Optional investable-universe filter. When given, the alpha's
            raw signal is masked to the investable set (see
            :func:`investable_universe_mask`) *before* it is turned into
            weights, so unheld names get weight 0. Keys are forwarded as keyword
            arguments to :func:`investable_universe_mask`.
        feature_paths: Optional parquet files/datasets keyed by ``symbol_id``
            and ``date``. Their numeric feature columns are left-joined onto
            ``data`` before alpha logic runs.
        feature_frames: Optional in-memory feature frames with the same schema.
        **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.
    """
    feature_inputs: list[pd.DataFrame] = []
    if feature_paths:
        feature_inputs.extend(read_derived_frames(feature_paths))
    if feature_frames:
        feature_inputs.extend(feature_frames)
    if feature_inputs:
        data = join_feature_frames(data, feature_inputs)

    alpha = get_alpha(alpha_type, **params)
    close = _pivot_close(data)
    signal = alpha.signal_from_data(data, close)
    if universe is None:
        weights = alpha.to_weights(signal)
    else:
        mask = investable_universe_mask(data, signal, **universe)
        weights = alpha.to_weights(signal.where(mask))

    # 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. A close-formed
    weight on date t is assumed tradable at open[t+1] and held until open[t+2].
    Return on signal date t = sum(weight[s,t] * open_to_open_return[s,t]) /
    sum(abs(weight[s,t])). This matches the execution convention without
    crediting the new signal for the overnight gap before it can be traded.

    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.
    """
    open_ = _pivot_open(data_df)
    fwd_returns_all = _forward_open_to_open_returns(open_)

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

    # Align weights to signal dates that exist on the market calendar. Compute
    # forward open-to-open returns on the full market calendar first, so sparse
    # signal grids still earn the next available open-to-open interval instead
    # of the next signal date.
    common_dates = weights.index.intersection(open_.index)
    weights = weights.loc[common_dates]
    fwd_returns = fwd_returns_all.reindex(common_dates)

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

    # Daily portfolio return = sum(w_t * r_open[t+1→t+2]) / sum(|w_t|).
    # The final two signal dates have no complete next-open holding interval
    # and are dropped below.
    gross = weights.abs().sum(axis=1)
    daily_returns = (
        (weights * fwd_returns).sum(axis=1, min_count=1)
        / gross.replace(0.0, np.nan)
    )
    daily_returns = daily_returns.dropna()
    if len(daily_returns) < 2:
        return {
            "cumulative_return": 0.0,
            "sharpe_annual": 0.0,
            "turnover_annual": 0.0,
            "max_drawdown": 0.0,
            "hit_rate": 0.0,
            "n_dates": int(len(daily_returns)),
        }

    # 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 = gross.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": int(len(daily_returns)),
    }