diff --git a/src/easy_tdx/portfolio/__init__.py b/src/easy_tdx/portfolio/__init__.py new file mode 100644 index 0000000..b17be90 --- /dev/null +++ b/src/easy_tdx/portfolio/__init__.py @@ -0,0 +1 @@ +"""组合管理模块。""" diff --git a/src/easy_tdx/portfolio/optimizer.py b/src/easy_tdx/portfolio/optimizer.py new file mode 100644 index 0000000..8f9cb4b --- /dev/null +++ b/src/easy_tdx/portfolio/optimizer.py @@ -0,0 +1,153 @@ +"""权重优化器。""" +from __future__ import annotations + +from abc import ABC, abstractmethod + +import numpy as np +import pandas as pd + + +class WeightOptimizer(ABC): + """权重优化器基类。""" + + @abstractmethod + def optimize( + self, + factor_scores: pd.DataFrame, + n_stocks: int = 50, + **kwargs: object, + ) -> dict[str, float]: + """返回 {code: weight},权重和为 1.0。""" + ... + + +_OPTIMIZER_REGISTRY: dict[str, type[WeightOptimizer]] = {} + + +def register_optimizer(name: str) -> type[WeightOptimizer]: + """注册优化器。""" + def wrapper(cls: type[WeightOptimizer]) -> type[WeightOptimizer]: + _OPTIMIZER_REGISTRY[name] = cls + return cls + return wrapper + + +def get_optimizer(name: str) -> WeightOptimizer: + """按名称获取优化器实例。""" + if name not in _OPTIMIZER_REGISTRY: + raise ValueError(f"未知优化器: {name!r}。可用: {sorted(_OPTIMIZER_REGISTRY.keys())}") + return _OPTIMIZER_REGISTRY[name]() + + +@register_optimizer("equal") +class EqualWeightOptimizer(WeightOptimizer): + """等权 — 取 top-N 等权分配。""" + + def optimize( + self, + factor_scores: pd.DataFrame, + n_stocks: int = 50, + **kwargs: object, + ) -> dict[str, float]: + if factor_scores.empty or "score" not in factor_scores.columns: + return {} + top = factor_scores.nlargest(n_stocks, "score") + if len(top) == 0: + return {} + w = 1.0 / len(top) + return {row["code"]: w for _, row in top.iterrows()} + + +@register_optimizer("factor_weighted") +class FactorWeightedOptimizer(WeightOptimizer): + """因子加权 — 按因子得分加权。""" + + def optimize( + self, + factor_scores: pd.DataFrame, + n_stocks: int = 50, + **kwargs: object, + ) -> dict[str, float]: + top = factor_scores.nlargest(n_stocks, "score") + if len(top) == 0: + return {} + scores = top["score"].to_numpy(dtype=np.float64) + if len(scores) > 10: + q95 = np.percentile(scores, 95) + q05 = np.percentile(scores, 5) + scores = np.clip(scores, q05, q95) + scores = scores - scores.min() + 1e-8 + total = scores.sum() + if total == 0: + w = 1.0 / len(top) + return {row["code"]: w for _, row in top.iterrows()} + weights = scores / total + return {row["code"]: float(weights[i]) for i, (_, row) in enumerate(top.iterrows())} + + +@register_optimizer("risk_parity") +class RiskParityOptimizer(WeightOptimizer): + """风险平价 — 每只股票贡献相等风险。""" + + def optimize( + self, + factor_scores: pd.DataFrame, + n_stocks: int = 50, + **kwargs: object, + ) -> dict[str, float]: + top = factor_scores.nlargest(n_stocks, "score") + if len(top) == 0: + return {} + if "volatility" in top.columns: + vol = top["volatility"].to_numpy(dtype=np.float64) + else: + scores = top["score"].abs().to_numpy(dtype=np.float64) + vol = 1.0 / (scores + 1e-8) + vol = np.maximum(vol, 1e-8) + inv_vol = 1.0 / vol + total = inv_vol.sum() + weights = inv_vol / total + return {row["code"]: float(weights[i]) for i, (_, row) in enumerate(top.iterrows())} + + +@register_optimizer("mean_variance") +class MeanVarianceOptimizer(WeightOptimizer): + """均值方差优化 — Markowitz 模型(可选 scipy)。""" + + def optimize( + self, + factor_scores: pd.DataFrame, + n_stocks: int = 50, + **kwargs: object, + ) -> dict[str, float]: + try: + from scipy.optimize import minimize # noqa: F401 + return self._optimize_with_scipy(factor_scores, n_stocks) + except ImportError: + fallback = EqualWeightOptimizer() + return fallback.optimize(factor_scores, n_stocks) + + def _optimize_with_scipy( + self, + factor_scores: pd.DataFrame, + n_stocks: int, + ) -> dict[str, float]: + from scipy.optimize import minimize as _minimize # type: ignore[import] + + top = factor_scores.nlargest(n_stocks, "score") + if len(top) == 0: + return {} + n = len(top) + scores = top["score"].to_numpy(dtype=np.float64) + variances = 1.0 / (np.abs(scores) + 1e-8) ** 2 + cov = np.diag(variances) + + def objective(w: np.ndarray) -> float: + return float(w @ cov @ w) + + constraints = {"type": "eq", "fun": lambda w: float(np.sum(w) - 1.0)} + bounds = [(0.0, 0.1)] * n + x0 = np.ones(n) / n + result = _minimize(objective, x0, method="SLSQP", bounds=bounds, constraints=constraints) + weights = result.x if result.success else np.ones(n) / n + return {row["code"]: float(weights[i]) for i, (_, row) in enumerate(top.iterrows())} diff --git a/src/easy_tdx/portfolio/rebalance.py b/src/easy_tdx/portfolio/rebalance.py new file mode 100644 index 0000000..5b58f06 --- /dev/null +++ b/src/easy_tdx/portfolio/rebalance.py @@ -0,0 +1,190 @@ +"""多期调仓回测引擎。""" +from __future__ import annotations + +import numpy as np +import pandas as pd + +from easy_tdx.factor.engine import FactorEngine +from easy_tdx.portfolio.optimizer import WeightOptimizer +from easy_tdx.portfolio.types import PortfolioState, RebalanceResult + + +class RebalanceEngine: + """多期调仓回测引擎。""" + + def __init__( + self, + optimizer: WeightOptimizer, + factor_name: str = "momentum_20d", + n_stocks: int = 50, + rebalance_freq: str = "M", + commission: float = 0.0003, + slippage: float = 0.001, + cash: float = 1_000_000, + ) -> None: + self._optimizer = optimizer + self._factor_name = factor_name + self._n_stocks = n_stocks + self._rebalance_freq = rebalance_freq + self._commission = commission + self._slippage = slippage + self._cash = cash + + def _get_rebalance_dates(self, dates: pd.DatetimeIndex) -> list[int]: + freq_map = {"W": "W-MON", "M": "M", "Q": "Q"} + freq = freq_map.get(self._rebalance_freq, "M") + series = pd.Series(dates, index=dates) + grouped = series.groupby(series.dt.to_period(freq)) + rebalance_dates = [group.iloc[-1] for _, group in grouped if len(group) > 0] + return [int(d.strftime("%Y%m%d")) for d in rebalance_dates] + + def run( + self, + data: dict[str, pd.DataFrame], + start_date: int | None = None, + end_date: int | None = None, + ) -> RebalanceResult: + """执行多期回测。""" + if not data: + return self._empty_result() + factor_engine = FactorEngine() + + all_dates: list[pd.Timestamp] = [] + for df in data.values(): + if "datetime" in df.columns: + all_dates.extend(df["datetime"].tolist()) + if not all_dates: + return self._empty_result() + all_dates = sorted(set(all_dates)) + + if start_date: + all_dates = [d for d in all_dates if int(d.strftime("%Y%m%d")) >= start_date] + if end_date: + all_dates = [d for d in all_dates if int(d.strftime("%Y%m%d")) <= end_date] + if not all_dates: + return self._empty_result() + + rebalance_dates = self._get_rebalance_dates(pd.DatetimeIndex(all_dates)) + rebalance_set = set(rebalance_dates) + + cash = self._cash + holdings: dict[str, float] = {} + states: list[PortfolioState] = [] + trades_list: list[dict[str, object]] = [] + equity_records: list[dict[str, object]] = [] + + for dt in all_dates: + date_int = int(dt.strftime("%Y%m%d")) + is_rebalance = date_int in rebalance_set + + prices: dict[str, float] = {} + for code, df in data.items(): + if "datetime" in df.columns: + row = df[df["datetime"] == dt] + if not row.empty: + prices[code] = float(row["close"].iloc[0]) + + position_value = sum(holdings.get(c, 0) * prices.get(c, 0) for c in holdings) + total_value = cash + position_value + + if is_rebalance and total_value > 0: + scores_df = factor_engine.compute_cross_section( + data, [self._factor_name], date=date_int + ) + if not scores_df.empty and scores_df[self._factor_name].notna().any(): + scores_df = scores_df.rename(columns={self._factor_name: "score"}) + scores_df = scores_df[["code", "score"]].dropna(subset=["score"]) + target_weights = self._optimizer.optimize(scores_df, n_stocks=self._n_stocks) + else: + target_weights = {} + + if target_weights: + trades_list, cash, holdings = self._rebalance( + target_weights, prices, total_value, date_int, trades_list, cash, holdings + ) + + position_value = sum(holdings.get(c, 0) * prices.get(c, 0) for c in holdings) + total_value = cash + position_value + weights: dict[str, float] = {} + if total_value > 0: + for c in holdings: + if holdings[c] > 0 and c in prices: + weights[c] = holdings[c] * prices[c] / total_value + + states.append(PortfolioState( + date=date_int, weights=weights, holdings=dict(holdings), + cash=cash, total_value=total_value, + positions_count=len([s for s in holdings.values() if s > 0]), + )) + equity_records.append({"datetime": date_int, "total": total_value, "cash": cash, "position_value": position_value}) + + equity_curve = pd.DataFrame(equity_records) + trades_df = pd.DataFrame(trades_list) if trades_list else pd.DataFrame( + columns=["datetime", "direction", "code", "shares", "price", "cost"] + ) + performance = self._compute_performance(equity_curve) + return RebalanceResult( + rebalance_dates=rebalance_dates, states=states, trades=trades_df, + equity_curve=equity_curve, performance=performance, + ) + + def _rebalance( + self, + target_weights: dict[str, float], + prices: dict[str, float], + total_value: float, + date_int: int, + trades_list: list[dict[str, object]], + cash: float, + holdings: dict[str, float], + ) -> tuple[list[dict[str, object]], float, dict[str, float]]: + for code in list(holdings.keys()): + if code not in target_weights and holdings[code] > 0: + price = prices.get(code, 0) + if price > 0: + sell_value = holdings[code] * price + cost = sell_value * (self._commission + self._slippage) + cash += sell_value - cost + trades_list.append({"datetime": date_int, "direction": "SELL", "code": code, "shares": holdings[code], "price": price, "cost": cost}) + del holdings[code] + + new_holdings: dict[str, float] = {} + for code, weight in target_weights.items(): + price = prices.get(code, 0) + if price <= 0: + continue + target_value = total_value * weight + shares = int(target_value / price / 100) * 100 + if shares > 0: + new_holdings[code] = shares + trade_value = shares * price + cost = trade_value * (self._commission + self._slippage) + trades_list.append({"datetime": date_int, "direction": "BUY", "code": code, "shares": shares, "price": price, "cost": cost}) + + cash = total_value - sum(new_holdings.get(c, 0) * prices.get(c, 0) for c in new_holdings) + holdings.clear() + holdings.update(new_holdings) + return trades_list, cash, holdings + + def _compute_performance(self, equity_curve: pd.DataFrame) -> dict[str, float]: + if len(equity_curve) < 2: + return {"total_return": 0.0, "annual_return": 0.0, "max_drawdown": 0.0, "sharpe": 0.0} + total = equity_curve["total"].to_numpy() + total_return = (total[-1] / total[0]) - 1 + n_days = len(total) + annual_return = (1 + total_return) ** (252 / max(n_days, 1)) - 1 + peak = np.maximum.accumulate(total) + drawdown = (total - peak) / peak + max_drawdown = float(np.min(drawdown)) + daily_ret = np.diff(total) / total[:-1] + daily_ret = daily_ret[~np.isnan(daily_ret)] + sharpe = float(np.mean(daily_ret) / np.std(daily_ret) * np.sqrt(252)) if len(daily_ret) > 1 and np.std(daily_ret) > 0 else 0.0 + return {"total_return": total_return, "annual_return": annual_return, "max_drawdown": max_drawdown, "sharpe": sharpe, "total_trades": len(equity_curve)} + + def _empty_result(self) -> RebalanceResult: + return RebalanceResult( + rebalance_dates=[], states=[], + trades=pd.DataFrame(columns=["datetime", "direction", "code", "shares", "price", "cost"]), + equity_curve=pd.DataFrame(columns=["datetime", "total", "cash", "position_value"]), + performance={"total_return": 0.0, "annual_return": 0.0, "max_drawdown": 0.0, "sharpe": 0.0}, + ) diff --git a/src/easy_tdx/portfolio/risk.py b/src/easy_tdx/portfolio/risk.py new file mode 100644 index 0000000..45ea35d --- /dev/null +++ b/src/easy_tdx/portfolio/risk.py @@ -0,0 +1,51 @@ +"""简化风险模型。""" +from __future__ import annotations + +import numpy as np +import pandas as pd + + +class RiskModel: + """简化风险模型 — A 股够用。""" + + def estimate_covariance( + self, + returns: pd.DataFrame, + method: str = "shrinkage", + shrinkage_intensity: float = 0.5, + window: int = 60, + ) -> pd.DataFrame: + """协方差矩阵估计。""" + if len(returns) < 2: + codes = returns.columns.tolist() if len(returns.columns) > 0 else [] + return pd.DataFrame(np.eye(len(codes)), index=codes, columns=codes) + if len(returns) > window: + returns = returns.iloc[-window:] + sample_cov = returns.cov() + if method == "shrinkage": + target = pd.DataFrame( + np.diag(np.diag(sample_cov.to_numpy())), + index=sample_cov.index, + columns=sample_cov.columns, + ) + return (1 - shrinkage_intensity) * sample_cov + shrinkage_intensity * target + return sample_cov + + def portfolio_risk( + self, + weights: dict[str, float], + cov_matrix: pd.DataFrame, + ) -> dict[str, float]: + """组合风险指标。""" + codes = [c for c in weights if c in cov_matrix.columns] + if not codes: + return {"total_volatility": 0.0, "max_risk_contribution": 0.0, "n_positions": 0} + w = np.array([weights[c] for c in codes]) + cov_sub = cov_matrix.loc[codes, codes].to_numpy() + var = float(w @ cov_sub @ w) + total_vol = np.sqrt(max(0, var)) * np.sqrt(252) + marginal = cov_sub @ w + risk_contrib = np.abs(w * marginal) + total_rc = risk_contrib.sum() + max_rc = float(risk_contrib.max() / total_rc) if total_rc > 0 else 0.0 + return {"total_volatility": total_vol, "max_risk_contribution": max_rc, "n_positions": len(codes)} diff --git a/src/easy_tdx/portfolio/types.py b/src/easy_tdx/portfolio/types.py new file mode 100644 index 0000000..2fe94f1 --- /dev/null +++ b/src/easy_tdx/portfolio/types.py @@ -0,0 +1,29 @@ +"""组合管理数据结构。""" +from __future__ import annotations + +from dataclasses import dataclass + +import pandas as pd + + +@dataclass +class PortfolioState: + """组合状态快照。""" + + date: int + weights: dict[str, float] + holdings: dict[str, float] + cash: float + total_value: float + positions_count: int + + +@dataclass +class RebalanceResult: + """再平衡结果。""" + + rebalance_dates: list[int] + states: list[PortfolioState] + trades: pd.DataFrame + equity_curve: pd.DataFrame + performance: dict[str, float] diff --git a/tests/unit/test_portfolio_optimizer.py b/tests/unit/test_portfolio_optimizer.py new file mode 100644 index 0000000..fc86ff9 --- /dev/null +++ b/tests/unit/test_portfolio_optimizer.py @@ -0,0 +1,71 @@ +"""Test portfolio optimizers.""" +from __future__ import annotations + +import numpy as np +import pandas as pd +import pytest + +from easy_tdx.portfolio.optimizer import ( + EqualWeightOptimizer, + FactorWeightedOptimizer, + RiskParityOptimizer, + get_optimizer, +) + + +def _make_scores(n: int = 20, seed: int = 42) -> pd.DataFrame: + rng = np.random.default_rng(seed) + return pd.DataFrame({ + "code": [f"{i:06d}" for i in range(n)], + "score": rng.normal(0.02, 0.05, n), + }) + + +class TestEqualWeight: + def test_weights_sum_to_one(self): + w = EqualWeightOptimizer().optimize(_make_scores(), n_stocks=10) + assert abs(sum(w.values()) - 1.0) < 1e-8 + + def test_n_stocks_selected(self): + w = EqualWeightOptimizer().optimize(_make_scores(), n_stocks=5) + assert len(w) == 5 + + def test_all_equal(self): + w = EqualWeightOptimizer().optimize(_make_scores(), n_stocks=10) + vals = list(w.values()) + assert all(abs(v - vals[0]) < 1e-8 for v in vals) + + def test_empty_input(self): + w = EqualWeightOptimizer().optimize(pd.DataFrame(columns=["code", "score"]), n_stocks=5) + assert len(w) == 0 + + +class TestFactorWeighted: + def test_weights_sum_to_one(self): + w = FactorWeightedOptimizer().optimize(_make_scores(), n_stocks=10) + assert abs(sum(w.values()) - 1.0) < 1e-6 + + def test_higher_score_higher_weight(self): + scores = pd.DataFrame({"code": ["A", "B", "C"], "score": [3.0, 2.0, 1.0]}) + w = FactorWeightedOptimizer().optimize(scores, n_stocks=3) + assert w["A"] > w["C"] + + +class TestRiskParity: + def test_weights_sum_to_one(self): + w = RiskParityOptimizer().optimize(_make_scores(), n_stocks=10) + assert abs(sum(w.values()) - 1.0) < 1e-6 + + def test_with_volatility_column(self): + scores = pd.DataFrame({"code": ["A", "B", "C"], "score": [1.0, 1.0, 1.0], "volatility": [0.1, 0.2, 0.4]}) + w = RiskParityOptimizer().optimize(scores, n_stocks=3) + assert w["A"] > w["C"] + + +class TestRegistry: + def test_get_optimizer(self): + assert isinstance(get_optimizer("equal"), EqualWeightOptimizer) + + def test_unknown_raises(self): + with pytest.raises(ValueError, match="未知优化器"): + get_optimizer("nonexistent") diff --git a/tests/unit/test_portfolio_rebalance.py b/tests/unit/test_portfolio_rebalance.py new file mode 100644 index 0000000..1fb03e5 --- /dev/null +++ b/tests/unit/test_portfolio_rebalance.py @@ -0,0 +1,55 @@ +"""Test RebalanceEngine.""" +from __future__ import annotations + +import numpy as np +import pandas as pd +import pytest + +from easy_tdx.portfolio.optimizer import EqualWeightOptimizer, FactorWeightedOptimizer +from easy_tdx.portfolio.rebalance import RebalanceEngine + + +def _make_market(n_stocks: int = 10, n_days: int = 120, seed: int = 42) -> dict[str, pd.DataFrame]: + rng = np.random.default_rng(seed) + data = {} + for i in range(n_stocks): + close = 10.0 + np.cumsum(rng.normal(0.01, 0.5, n_days)) + close = np.maximum(close, 1.0) + data[f"{i:06d}"] = pd.DataFrame({ + "datetime": pd.date_range("2024-01-01", periods=n_days, freq="D"), + "open": close, "high": close + 0.3, "low": close - 0.3, + "close": close, + "vol": rng.integers(1e5, 1e7, n_days).astype(float), + "amount": close * 1e6, + }) + return data + + +class TestRebalanceEngine: + def test_basic_run(self): + engine = RebalanceEngine(optimizer=EqualWeightOptimizer(), factor_name="momentum_20d", n_stocks=5, rebalance_freq="M", cash=1_000_000) + result = engine.run(_make_market(), start_date=20240101, end_date=20240430) + assert len(result.states) > 0 + assert len(result.rebalance_dates) > 0 + assert len(result.equity_curve) > 0 + assert "total_return" in result.performance + + def test_with_factor_weighted(self): + engine = RebalanceEngine(optimizer=FactorWeightedOptimizer(), factor_name="momentum_20d", n_stocks=5) + result = engine.run(_make_market(), start_date=20240101, end_date=20240430) + assert len(result.states) > 0 + + def test_empty_data(self): + result = RebalanceEngine(optimizer=EqualWeightOptimizer()).run({}) + assert result.performance["total_return"] == 0.0 + + def test_equity_curve_dates_sorted(self): + result = RebalanceEngine(optimizer=EqualWeightOptimizer(), rebalance_freq="M").run(_make_market(), start_date=20240101, end_date=20240430) + dates = result.equity_curve["datetime"].tolist() + assert dates == sorted(dates) + + def test_trades_recorded(self): + engine = RebalanceEngine(optimizer=EqualWeightOptimizer(), n_stocks=3, rebalance_freq="M") + result = engine.run(_make_market(), start_date=20240101, end_date=20240430) + assert len(result.trades) > 0 + assert "BUY" in result.trades["direction"].values diff --git a/tests/unit/test_portfolio_risk.py b/tests/unit/test_portfolio_risk.py new file mode 100644 index 0000000..778f5ba --- /dev/null +++ b/tests/unit/test_portfolio_risk.py @@ -0,0 +1,44 @@ +"""Test RiskModel.""" +from __future__ import annotations + +import numpy as np +import pandas as pd + +from easy_tdx.portfolio.risk import RiskModel + + +def _make_returns(n_dates: int = 100, n_stocks: int = 5, seed: int = 42) -> pd.DataFrame: + rng = np.random.default_rng(seed) + codes = [f"{i:06d}" for i in range(n_stocks)] + return pd.DataFrame(rng.normal(0.001, 0.02, (n_dates, n_stocks)), columns=codes) + + +class TestCovarianceEstimation: + def test_shape(self): + cov = RiskModel().estimate_covariance(_make_returns()) + assert cov.shape == (5, 5) + + def test_symmetric(self): + cov = RiskModel().estimate_covariance(_make_returns()) + assert np.allclose(cov.to_numpy(), cov.to_numpy().T) + + def test_shrinkage_reduces_offdiag(self): + rm = RiskModel() + ret = _make_returns() + shrunk = rm.estimate_covariance(ret, method="shrinkage") + sample = rm.estimate_covariance(ret, method="sample") + off_shrunk = shrunk.values[~np.eye(5, dtype=bool)] + off_sample = sample.values[~np.eye(5, dtype=bool)] + assert np.abs(off_shrunk).mean() <= np.abs(off_sample).mean() + + +class TestPortfolioRisk: + def test_total_volatility(self): + cov = RiskModel().estimate_covariance(_make_returns()) + risk = RiskModel().portfolio_risk({"000000": 0.5, "000001": 0.5}, cov) + assert risk["total_volatility"] > 0 + assert risk["n_positions"] == 2 + + def test_empty_weights(self): + risk = RiskModel().portfolio_risk({}, pd.DataFrame()) + assert risk["total_volatility"] == 0.0