chinese-equity-quant/pipeline/portfolio/constraints.py

"""Trade constraints for the execution simulator.

A constraint answers: *given today's market state, how much may each name be
traded?* Each returns per-name signed delta bounds ``(low, high)`` in shares;
the simulator intersects them (``low = max(lows)``, ``high = min(highs)``) and
clips the desired trade. ``-inf/inf`` mean uncapped, ``0`` blocks that direction.

Constraints self-register via :func:`register_constraint` (mirroring the alpha
registry) so the CLI can select them by name. They consume the abstracted
:class:`~pipeline.portfolio.market_rules.LimitStatus` rather than raw prices,
leaving room for richer fill models (一字板, queues, partial fills) later.
"""

from __future__ import annotations

from abc import ABC, abstractmethod
from typing import TYPE_CHECKING, Type

import numpy as np

from pipeline.portfolio.market_rules import LimitStatus

if TYPE_CHECKING:
    from pipeline.portfolio.simulator import TradeContext


class TradeConstraint(ABC):
    """Per-name tradeability rule producing signed share-delta bounds."""

    name: str = ""

    @abstractmethod
    def delta_bounds(self, ctx: "TradeContext") -> tuple[np.ndarray, np.ndarray]:
        """Return ``(low, high)`` arrays: the min/max signed share delta allowed."""

    def adjust_targets(self, ctx: "TradeContext") -> np.ndarray | None:
        """Optional portfolio-level retargeting hook (e.g. neutrality).

        Returns a new target-shares array, or None to leave targets unchanged.
        Default: no adjustment. Future industry/beta neutrality can implement a
        cheap numpy least-squares projection here — no MIP needed.
        """
        return None


_CONSTRAINTS: dict[str, Type[TradeConstraint]] = {}


def register_constraint(cls: Type[TradeConstraint]) -> Type[TradeConstraint]:
    """Class decorator registering a constraint under its ``name``."""
    if not (isinstance(cls, type) and issubclass(cls, TradeConstraint)):
        raise TypeError(f"{cls!r} is not a TradeConstraint subclass")
    key = getattr(cls, "name", "")
    if not key:
        raise ValueError(f"{cls.__name__} must set a non-empty class attribute `name`")
    existing = _CONSTRAINTS.get(key)
    if existing is not None and existing is not cls:
        raise ValueError(f"Constraint name '{key}' already registered by {existing.__name__}")
    _CONSTRAINTS[key] = cls
    return cls


def available_constraints() -> list[str]:
    """Sorted names of all registered constraints."""
    return sorted(_CONSTRAINTS)


def get_constraint(name: str, **params) -> TradeConstraint:
    """Instantiate a registered constraint by name.

    Raises:
        KeyError: If ``name`` is not registered.
    """
    if name not in _CONSTRAINTS:
        raise KeyError(f"Unknown constraint '{name}'. Available: {sorted(_CONSTRAINTS)}")
    return _CONSTRAINTS[name](**params)


def _unbounded(n: int) -> tuple[np.ndarray, np.ndarray]:
    return np.full(n, -np.inf), np.full(n, np.inf)


@register_constraint
class SuspensionConstraint(TradeConstraint):
    """Suspended names (``tradestatus == 0``) cannot trade at all."""

    name = "suspension"

    def delta_bounds(self, ctx):
        n = len(ctx.target_shares)
        low, high = _unbounded(n)
        suspended = ctx.slice.tradestatus == 0
        low = np.where(suspended, 0.0, low)
        high = np.where(suspended, 0.0, high)
        return low, high


@register_constraint
class PriceLimitConstraint(TradeConstraint):
    """Limit-up blocks buys; limit-down blocks sells.

    Consumes ``ctx.slice.limit_status`` (NORMAL/UP_LIMIT/DOWN_LIMIT), not raw
    prices, so future states (e.g. limit-locked 一字板 with zero fill) plug in
    by extending the status set.
    """

    name = "price_limit"

    def delta_bounds(self, ctx):
        n = len(ctx.target_shares)
        low, high = _unbounded(n)
        status = ctx.slice.limit_status
        at_up = status == LimitStatus.UP_LIMIT.value
        at_down = status == LimitStatus.DOWN_LIMIT.value
        high = np.where(at_up, 0.0, high)   # cannot buy at the up limit
        low = np.where(at_down, 0.0, low)   # cannot sell at the down limit
        return low, high


@register_constraint
class VolumeCapConstraint(TradeConstraint):
    """Cap traded **value** at a fraction of the day's turnover value.

    ``|trade_value| ≤ max_frac · amount`` (amount = daily turnover in yuan),
    converted to a share cap via the execution price. Value-based, not
    share-count based.
    """

    name = "volume_cap"

    def __init__(self, max_frac: float = 0.10):
        self.max_frac = max_frac

    def delta_bounds(self, ctx):
        amount = np.asarray(ctx.slice.amount, dtype=np.float64)
        price = np.asarray(ctx.slice.price, dtype=np.float64)
        cap_value = self.max_frac * np.where(np.isfinite(amount), amount, 0.0)
        with np.errstate(divide="ignore", invalid="ignore"):
            cap_shares = np.where(price > 0, cap_value / price, 0.0)
        cap_shares = np.floor(cap_shares)
        return -cap_shares, cap_shares