tca.py

"""Execution-quality metrics for simulated VWAP parent orders."""

from typing import Iterable, Optional, Tuple

import matplotlib.pyplot as plt
import numpy as np
import pandas as pd


def parent_execution_metrics(
    parents: pd.DataFrame,
    fills: pd.DataFrame,
    quote_snapshots: pd.DataFrame,
) -> pd.DataFrame:
    """Calculate implementation shortfall, VWAP slippage, and fill rate per parent."""

    _require_simulated(parents, "parents")
    _require_simulated(fills, "fills", allow_empty=True)
    _require_columns(
        parents,
        {
            "parent_id",
            "side",
            "requested_size",
            "start_seconds",
            "end_seconds",
            "arrival_mid",
            "end_mid",
            "executed_size",
            "average_fill_price",
        },
        "parents",
    )
    _require_columns(fills, {"parent_id", "size", "fill_price"}, "fills")
    _require_columns(
        quote_snapshots,
        {"data_label", "timestamp", "event_type", "price", "size", "mid", "spread"},
        "quote_snapshots",
    )
    market_labels = sorted(quote_snapshots["data_label"].dropna().unique())
    if len(market_labels) != 1 or market_labels[0] not in ("REAL", "SIMULATED"):
        raise ValueError("Quote snapshots must have one explicit REAL or SIMULATED label")

    rows = []
    for parent in parents.itertuples(index=False):
        side = int(parent.side)
        if side not in (-1, 1):
            raise ValueError("Parent side must be -1 or +1")
        requested_size = int(parent.requested_size)
        parent_fills = fills.loc[fills["parent_id"].eq(parent.parent_id)]
        executed_size = int(parent_fills["size"].sum())
        average_fill_price = (
            float(np.average(parent_fills["fill_price"], weights=parent_fills["size"]))
            if executed_size > 0
            else np.nan
        )
        _validate_engine_aggregates(parent, executed_size, average_fill_price)

        window = quote_snapshots.loc[
            quote_snapshots["timestamp"].ge(parent.start_seconds)
            & quote_snapshots["timestamp"].lt(parent.end_seconds)
        ]
        market_trades = window.loc[window["event_type"].isin((4, 5))]
        interval_volume = int(market_trades["size"].sum())
        interval_vwap = (
            float(np.average(market_trades["price"], weights=market_trades["size"]))
            if interval_volume > 0
            else np.nan
        )
        mean_spread = float(window["spread"].mean()) if not window.empty else np.nan
        mean_mid = float(window["mid"].mean()) if not window.empty else np.nan
        mean_spread_bps = 10_000.0 * mean_spread / mean_mid if mean_mid > 0 else np.nan
        volatility_bps = _realized_volatility_bps(window["mid"])

        arrival_mid = float(parent.arrival_mid)
        end_mid = float(parent.end_mid)
        execution_cost_bps = (
            10_000.0
            * side
            * (average_fill_price - arrival_mid)
            * executed_size
            / (arrival_mid * requested_size)
            if executed_size > 0
            else 0.0
        )
        unfilled_size = requested_size - executed_size
        opportunity_cost_bps = (
            10_000.0
            * side
            * (end_mid - arrival_mid)
            * unfilled_size
            / (arrival_mid * requested_size)
        )
        implementation_shortfall_bps = execution_cost_bps + opportunity_cost_bps
        vwap_slippage_bps = (
            10_000.0 * side * (average_fill_price - interval_vwap) / interval_vwap
            if executed_size > 0 and interval_vwap > 0
            else np.nan
        )
        rows.append(
            {
                "data_label": "SIMULATED",
                "market_data_label": market_labels[0],
                "parent_id": parent.parent_id,
                "side": side,
                "requested_size": requested_size,
                "start_seconds": float(parent.start_seconds),
                "end_seconds": float(parent.end_seconds),
                "executed_size": executed_size,
                "average_fill_price": average_fill_price,
                "arrival_mid": arrival_mid,
                "end_mid": end_mid,
                "interval_volume": interval_volume,
                "interval_vwap": interval_vwap,
                "size_participation": requested_size / interval_volume if interval_volume else np.nan,
                "mean_spread": mean_spread,
                "mean_spread_bps": mean_spread_bps,
                "volatility_bps": volatility_bps,
                "execution_cost_bps": execution_cost_bps,
                "opportunity_cost_bps": opportunity_cost_bps,
                "implementation_shortfall_bps": implementation_shortfall_bps,
                "vwap_slippage_bps": vwap_slippage_bps,
                "fill_rate": executed_size / requested_size,
            }
        )
    return pd.DataFrame(rows)


def add_regime_buckets(
    metrics: pd.DataFrame,
    size_edges: Iterable[float],
    spread_quantiles: Iterable[float],
    volatility_quantiles: Iterable[float],
) -> pd.DataFrame:
    """Assign parent orders to size, spread, and volatility regimes."""

    _require_simulated(metrics, "metrics")
    result = metrics.copy()
    edges = tuple(float(edge) for edge in size_edges)
    if len(edges) < 2 or any(right <= left for left, right in zip(edges, edges[1:])):
        raise ValueError("Size edges must be strictly increasing")
    size_labels = [f"{100 * left:g}-{100 * right:g}%" for left, right in zip(edges, edges[1:])]
    result["size_regime"] = pd.cut(
        result["size_participation"],
        bins=edges,
        labels=size_labels,
        include_lowest=True,
    )
    result["spread_regime"] = _quantile_buckets(
        result["mean_spread_bps"], spread_quantiles, "spread"
    )
    result["volatility_regime"] = _quantile_buckets(
        result["volatility_bps"], volatility_quantiles, "volatility"
    )
    return result


def summarize_cost_by_regime(bucketed_metrics: pd.DataFrame) -> pd.DataFrame:
    """Aggregate simulated costs separately for each configured regime dimension."""

    _require_simulated(bucketed_metrics, "bucketed_metrics")
    summaries = []
    for dimension in ("size", "spread", "volatility"):
        column = f"{dimension}_regime"
        _require_columns(bucketed_metrics, {column}, "bucketed_metrics")
        grouped = (
            bucketed_metrics.dropna(subset=[column])
            .groupby(column, observed=True)
            .agg(
                parent_count=("parent_id", "size"),
                mean_is_bps=("implementation_shortfall_bps", "mean"),
                mean_vwap_slippage_bps=("vwap_slippage_bps", "mean"),
                mean_fill_rate=("fill_rate", "mean"),
            )
            .reset_index()
            .rename(columns={column: "regime"})
        )
        grouped.insert(0, "dimension", dimension)
        summaries.append(grouped)
    result = pd.concat(summaries, ignore_index=True)
    result.insert(0, "data_label", "SIMULATED")
    return result


def plot_cost_vs_size(
    metrics: pd.DataFrame,
    ax: Optional[plt.Axes] = None,
) -> Tuple[plt.Figure, plt.Axes]:
    """Render mean implementation shortfall and VWAP slippage by parent size."""

    _require_simulated(metrics, "metrics")
    curve = (
        metrics.groupby("requested_size", as_index=False)
        .agg(
            implementation_shortfall_bps=("implementation_shortfall_bps", "mean"),
            vwap_slippage_bps=("vwap_slippage_bps", "mean"),
        )
        .sort_values("requested_size")
    )
    if ax is None:
        figure, ax = plt.subplots(figsize=(6, 4))
    else:
        figure = ax.figure
    ax.plot(
        curve["requested_size"],
        curve["implementation_shortfall_bps"],
        marker="o",
        label="Implementation shortfall",
    )
    ax.plot(
        curve["requested_size"],
        curve["vwap_slippage_bps"],
        marker="o",
        label="VWAP slippage",
    )
    ax.axhline(0.0, color="black", linewidth=0.8)
    ax.set(
        title="SIMULATED VWAP — cost by parent size",
        xlabel="Parent size (shares)",
        ylabel="Cost (bps)",
    )
    ax.legend()
    return figure, ax


def _validate_engine_aggregates(parent: object, executed_size: int, average_fill_price: float) -> None:
    if int(parent.executed_size) != executed_size:
        raise ValueError(f"Fill size does not match parent {parent.parent_id}")
    recorded_average = float(parent.average_fill_price)
    if executed_size > 0 and not np.isclose(recorded_average, average_fill_price):
        raise ValueError(f"Average fill price does not match parent {parent.parent_id}")


def _realized_volatility_bps(mids: pd.Series) -> float:
    positive = mids.loc[mids.gt(0)].to_numpy(dtype="float64")
    if len(positive) < 2:
        return np.nan
    returns = np.diff(np.log(positive))
    return float(10_000.0 * np.sqrt(np.square(returns).sum()))


def _quantile_buckets(
    values: pd.Series,
    quantiles: Iterable[float],
    prefix: str,
) -> pd.Series:
    quantiles = tuple(float(quantile) for quantile in quantiles)
    if len(quantiles) < 2 or quantiles[0] != 0.0 or quantiles[-1] != 1.0:
        raise ValueError("Quantiles must start at zero and end at one")
    if any(right <= left for left, right in zip(quantiles, quantiles[1:])):
        raise ValueError("Quantiles must be strictly increasing")
    thresholds = np.unique(values.dropna().quantile(quantiles).to_numpy())
    if len(thresholds) < 2:
        return pd.Series(f"{prefix}_all", index=values.index, dtype="object")
    labels = [f"{prefix}_q{index + 1}" for index in range(len(thresholds) - 1)]
    return pd.cut(values, bins=thresholds, labels=labels, include_lowest=True).astype("object")


def _require_simulated(frame: pd.DataFrame, name: str, allow_empty: bool = False) -> None:
    _require_columns(frame, {"data_label"}, name)
    labels = set(frame["data_label"].dropna().unique())
    if not labels and allow_empty and frame.empty:
        return
    if labels != {"SIMULATED"}:
        raise ValueError(f"{name} must be explicitly labelled SIMULATED")


def _require_columns(frame: pd.DataFrame, required: set, name: str) -> None:
    missing = required.difference(frame.columns)
    if missing:
        raise ValueError(f"{name} missing columns: {sorted(missing)}")