mirror of
https://ghfast.top/https://github.com/aeroxw/easy-tdx.git
synced 2026-09-12 16:54:17 +08:00
feat(portfolio): add optimizer, risk model, and rebalance engine
- WeightOptimizer base class with registry (equal, factor_weighted, risk_parity, mean_variance) - RiskModel with shrinkage covariance estimation and portfolio risk metrics - RebalanceEngine for multi-period backtesting with commission/slippage - 20 unit tests covering all components Co-Authored-By: Claude <noreply@anthropic.com>
This commit is contained in:
GitHub
@@ -0,0 +1 @@
|
||||
"""组合管理模块。"""
|
||||
@@ -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())}
|
||||
@@ -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},
|
||||
)
|
||||
@@ -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)}
|
||||
@@ -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]
|
||||
@@ -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")
|
||||
@@ -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
|
||||
@@ -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
|
||||
Reference in New Issue
Block a user