115 lines
3.8 KiB
Python
115 lines
3.8 KiB
Python
import unittest
|
|
|
|
import pandas as pd
|
|
|
|
|
|
class USAlphaPipelineTests(unittest.TestCase):
|
|
def test_build_equal_weight_portfolio_caps_holdings_under_ties(self):
|
|
from research.us_alpha_pipeline import _build_equal_weight_portfolio
|
|
|
|
dates = pd.date_range("2024-01-01", periods=2, freq="D")
|
|
score = pd.DataFrame(
|
|
{
|
|
"AAA": [0.9, 0.9],
|
|
"BBB": [0.9, 0.9],
|
|
"CCC": [0.9, 0.9],
|
|
},
|
|
index=dates,
|
|
)
|
|
tradable_mask = pd.DataFrame(True, index=dates, columns=score.columns)
|
|
regime = pd.Series([True, True], index=dates)
|
|
|
|
weights = _build_equal_weight_portfolio(score, tradable_mask, regime, top_n=2)
|
|
|
|
self.assertEqual(int((weights.iloc[-1] > 0).sum()), 2)
|
|
self.assertAlmostEqual(float(weights.iloc[-1].sum()), 1.0)
|
|
|
|
def test_equity_curve_uses_prior_day_weights_for_returns(self):
|
|
from research.us_alpha_pipeline import _equity_curve
|
|
|
|
dates = pd.date_range("2024-01-01", periods=3, freq="D")
|
|
close = pd.DataFrame({"AAA": [1.0, 2.0, 4.0]}, index=dates)
|
|
weights = pd.DataFrame({"AAA": [0.0, 1.0, 0.0]}, index=dates)
|
|
|
|
equity = _equity_curve(close, weights)
|
|
|
|
self.assertEqual(float(equity.iloc[1]), 1.0)
|
|
self.assertEqual(float(equity.iloc[2]), 2.0)
|
|
|
|
def test_run_alpha_pipeline_returns_expected_strategy_summary(self):
|
|
from research.us_alpha_pipeline import run_alpha_pipeline
|
|
|
|
dates = pd.date_range("2023-01-01", periods=400, freq="D")
|
|
|
|
aaa_close = [50.0 + 0.20 * i for i in range(400)]
|
|
bbb_close = [55.0 + 0.12 * i for i in range(400)]
|
|
ccc_close = [60.0 + 0.05 * i for i in range(400)]
|
|
close = pd.DataFrame(
|
|
{
|
|
"AAA": aaa_close,
|
|
"BBB": bbb_close,
|
|
"CCC": ccc_close,
|
|
},
|
|
index=dates,
|
|
)
|
|
high = pd.DataFrame(
|
|
{
|
|
"AAA": [value + 0.5 for value in aaa_close],
|
|
"BBB": [value + 1.0 for value in bbb_close],
|
|
"CCC": [value + 1.5 for value in ccc_close],
|
|
},
|
|
index=dates,
|
|
)
|
|
low = pd.DataFrame(
|
|
{
|
|
"AAA": [value - 0.5 for value in aaa_close],
|
|
"BBB": [value - 1.0 for value in bbb_close],
|
|
"CCC": [value - 1.5 for value in ccc_close],
|
|
},
|
|
index=dates,
|
|
)
|
|
volume = pd.DataFrame(
|
|
{
|
|
"AAA": [1_500_000.0] * 400,
|
|
"BBB": [1_400_000.0] * 400,
|
|
"CCC": [1_300_000.0] * 400,
|
|
},
|
|
index=dates,
|
|
)
|
|
volume.loc[dates[-2], "AAA"] = 4_000_000.0
|
|
|
|
etf_close = pd.DataFrame(
|
|
{
|
|
"SPY": [300.0 + 0.8 * i for i in range(400)],
|
|
"QQQ": [280.0 + 1.1 * i for i in range(400)],
|
|
"XLF": [200.0 + 0.4 * i for i in range(400)],
|
|
},
|
|
index=dates,
|
|
)
|
|
|
|
market_data = {
|
|
"close": close,
|
|
"high": high,
|
|
"low": low,
|
|
"volume": volume,
|
|
}
|
|
|
|
summary = run_alpha_pipeline(
|
|
market_data=market_data,
|
|
etf_close=etf_close,
|
|
pit_membership=None,
|
|
windows=(1,),
|
|
top_n=2,
|
|
)
|
|
|
|
required_columns = {"strategy", "window_years", "CAGR", "Sharpe", "MaxDD", "TotalRet"}
|
|
self.assertTrue(required_columns.issubset(summary.columns))
|
|
self.assertEqual(set(summary["strategy"]), {"breakout_regime", "rank_blend_regime"})
|
|
self.assertEqual(set(summary["window_years"]), {1})
|
|
self.assertEqual(len(summary), 2)
|
|
self.assertTrue(summary[["CAGR", "Sharpe", "MaxDD", "TotalRet"]].notna().all().all())
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|