Spaces:
Sleeping
Sleeping
File size: 8,299 Bytes
9831ced a9e6751 9831ced | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 | import importlib.util
import unittest
import numpy as np
import pandas as pd
from features.feature_pipeline import compute_feature_frame
def _have(module_name: str) -> bool:
return importlib.util.find_spec(module_name) is not None
_HAVE_STATSMODELS = _have('statsmodels')
_HAVE_ARCH = _have('arch')
_HAVE_MOIRAI = _have('uni2ts') and _have('gluonts')
_HAVE_TIMESFM = _have('timesfm') and _have('torch')
def _synthetic_ohlcv(n: int, seed: int=42) -> pd.DataFrame:
rng = np.random.default_rng(seed)
steps = rng.normal(loc=0.0, scale=0.4, size=n)
close = 100 + np.cumsum(steps)
close = np.abs(close) + 5.0
open_ = close + rng.normal(0, 0.1, n)
high = np.maximum(open_, close) + np.abs(rng.normal(0, 0.2, n))
low = np.minimum(open_, close) - np.abs(rng.normal(0, 0.2, n))
volume = rng.integers(100, 10000, n).astype(float)
idx = pd.date_range('2024-01-01', periods=n, freq='min')
return pd.DataFrame({'Open': open_, 'High': high, 'Low': low, 'Close': close, 'Volume': volume}, index=idx)
def _warm_close_and_features(df: pd.DataFrame, feats: pd.DataFrame):
complete = ~feats.isna().any(axis=1).to_numpy()
start = int(np.argmax(complete)) if complete.any() else len(df)
return (df['Close'].reset_index(drop=True).iloc[start:], feats.iloc[start:].reset_index(drop=True))
class TestArimaFamilyAlignment(unittest.TestCase):
def test_returns_alignment_has_no_off_by_one(self):
n = 10
prices = pd.Series(100 + np.arange(n, dtype=float))
features = pd.Series([float(i) for i in range(n)])
log_returns = np.log(prices / prices.shift(1)).dropna().reset_index(drop=True)
exog = features.iloc[1:].reset_index(drop=True)
self.assertEqual(len(exog), len(log_returns))
for k in range(len(log_returns)):
self.assertEqual(exog[k], float(k + 1))
@unittest.skipUnless(_HAVE_STATSMODELS, 'statsmodels not installed in this environment')
class TestArimaWithFeatures(unittest.TestCase):
def setUp(self):
self.df = _synthetic_ohlcv(120)
self.feats = compute_feature_frame(self.df)
self.close, self.warm_feats = _warm_close_and_features(self.df, self.feats)
def test_baseline_unchanged_shape(self):
from models.arima_model import ArimaModel
m = ArimaModel(max_p=2, max_d=1, max_q=2)
out = m.predict(self.close, horizon=3)
self.assertEqual(len(out), 3)
self.assertTrue(all(np.isfinite(out)))
def test_with_features_runs_and_is_finite(self):
from models.arima_model import ArimaModel
m = ArimaModel(max_p=2, max_d=1, max_q=2)
out = m.predict(self.close, horizon=3, features=self.warm_feats)
self.assertEqual(len(out), 3)
self.assertTrue(all(np.isfinite(out)))
def test_mismatched_features_length_raises(self):
from models.arima_model import ArimaModel
m = ArimaModel(max_p=2, max_d=1, max_q=2)
short_window = self.warm_feats.iloc[:-5].reset_index(drop=True)
with self.assertRaises(ValueError):
m.predict(self.close, horizon=3, features=short_window)
@unittest.skipUnless(_HAVE_STATSMODELS, 'statsmodels not installed in this environment')
class TestAutoArimaWithFeatures(unittest.TestCase):
def setUp(self):
self.df = _synthetic_ohlcv(120)
self.feats = compute_feature_frame(self.df)
self.close, self.warm_feats = _warm_close_and_features(self.df, self.feats)
def test_with_features_runs_and_is_finite(self):
from models.auto_arima_model import AutoArimaModel
m = AutoArimaModel(max_p=3, max_q=3, max_d=1)
out = m.predict(self.close, horizon=2, features=self.warm_feats)
self.assertEqual(len(out), 2)
self.assertTrue(all(np.isfinite(out)))
@unittest.skipUnless(_HAVE_STATSMODELS and _HAVE_ARCH, 'statsmodels+arch not installed in this environment')
class TestArimaGarchWithFeatures(unittest.TestCase):
def setUp(self):
self.df = _synthetic_ohlcv(150)
self.feats = compute_feature_frame(self.df)
self.close, self.warm_feats = _warm_close_and_features(self.df, self.feats)
def test_with_features_runs_and_is_finite(self):
from models.arima_garch_model import ArimaGarchModel
m = ArimaGarchModel(max_p=2, max_q=2)
out = m.predict(self.close, horizon=2, features=self.warm_feats)
self.assertEqual(len(out), 2)
self.assertTrue(all(np.isfinite(out)))
self.assertIsNotNone(m.last_volatility_forecast)
@unittest.skipUnless(_HAVE_MOIRAI, 'uni2ts/gluonts not installed in this environment')
class TestMoiraiWithFeatures(unittest.TestCase):
def test_with_features_runs_and_is_finite(self):
from models.moirai_model import MoiraiModel
df = _synthetic_ohlcv(60)
feats = compute_feature_frame(df)
close, warm_feats = _warm_close_and_features(df, feats)
m = MoiraiModel(context_length=60, num_samples=10)
out = m.predict(close, horizon=2, features=warm_feats)
self.assertEqual(len(out), 2)
self.assertTrue(all(np.isfinite(out)))
@unittest.skipUnless(_HAVE_TIMESFM, 'timesfm/torch not installed in this environment')
class TestTimesFMWithFeatures(unittest.TestCase):
def test_dashboard_non_patch_multiple_window_uses_real_model(self):
from pipeline import RunConfig, new_configured_model
config = RunConfig(models=('TimesFM',), window=200)
model = new_configured_model('TimesFM', config)
assert model.max_context == 224
values = model.predict(_synthetic_ohlcv(200)['Close'], horizon=2)
assert len(values) == 2 and np.isfinite(values).all()
assert model.last_diagnostics['weights_loaded']
def test_with_features_runs_and_is_finite(self):
from models.timesfm_model import TimesFMModel
df = _synthetic_ohlcv(64)
feats = compute_feature_frame(df)
close, warm_feats = _warm_close_and_features(df, feats)
m = TimesFMModel(max_context=64)
out = m.predict(close, horizon=2, features=warm_feats)
self.assertEqual(len(out), 2)
self.assertTrue(all(np.isfinite(out)))
def test_short_input_below_compiled_context_is_finite(self):
from models.timesfm_model import TimesFMModel
df = _synthetic_ohlcv(60, seed=99)
close = 40000 + df['Close'].reset_index(drop=True)
m = TimesFMModel(max_context=512)
out = m.predict(close, horizon=2)
self.assertEqual(len(out), 2)
self.assertTrue(all(np.isfinite(out)), f'TimesFM NaN on short window: {out}')
def test_repeated_window_lengths_stay_finite(self):
from models.timesfm_model import TimesFMModel
df = _synthetic_ohlcv(120, seed=7)
close = 40000 + df['Close'].reset_index(drop=True)
m = TimesFMModel(max_context=512)
for w in (40, 100, 40):
out = m.predict(close.iloc[-w:], horizon=1)
self.assertEqual(len(out), 1)
self.assertTrue(np.isfinite(out[0]), f'TimesFM NaN at w={w}: {out[0]}')
class TestInputContractGuards(unittest.TestCase):
def test_nan_history_raises_for_all_models(self):
from models.registry import fresh_model
rng = np.random.default_rng(3)
good = pd.Series(100 + np.cumsum(rng.normal(0, 0.5, 200)))
for name in ('ARIMA', 'Auto-ARIMA', 'ARIMA-GARCH', 'Moirai'):
m = fresh_model(name)
h = good.copy()
h.iloc[50] = np.nan
with self.assertRaises(ValueError, msg=name):
m.predict(h, horizon=1)
m = fresh_model('TimesFM')
h = good.copy()
h.iloc[50] = np.nan
with self.assertRaises(ValueError):
m.predict(h, horizon=1)
def test_horizon_below_one_raises_for_all_models(self):
from models.registry import fresh_model
rng = np.random.default_rng(3)
good = pd.Series(100 + np.cumsum(rng.normal(0, 0.5, 200)))
for name in ('ARIMA', 'Auto-ARIMA', 'ARIMA-GARCH', 'Moirai', 'TimesFM'):
m = fresh_model(name)
for bad in (0, -1):
with self.assertRaises(ValueError, msg=f'{name} horizon={bad}'):
m.predict(good, horizon=bad)
if __name__ == '__main__':
unittest.main()
|