def setUpClass(self):
np.random.seed(8675309)
np.random.seed(0xDEADBEEF)
fake_data = np.random.randn(100, 10)
fake_coefficients = np.array([5, -4, 3] + [0] * 7)
random_walk = np.cumsum(np.random.randn(100))
noise = np.random.randn(100) * 1.5
y = random_walk + fake_data @ fake_coefficients + noise
data = pd.DataFrame(
fake_data,
columns=["V" + str(x + 1) for x in range(10)],
)
data["y"] = y
model = Bsts()
model.add_state(LocalLevelStateModel(y))
model.train(formula=f"y ~ {dot(data, omit='y')}", niter=100, data=data)
self._regression_model = model
data = np.log(AirPassengers)
model = Bsts()
model.add_state(LocalLinearTrendStateModel(data))
model.add_state(SeasonalStateModel(data, nseasons=12))
model.train(data, niter=100)
self._model = model
model = Bsts()
model.add_state(LocalLinearTrendStateModel(data))
model.train(data, niter=100)
self._model2 = model
class TestLocalLinearTrend(unittest.TestCase):
def setUp(self):
np.random.seed(8675309)
self.data = np.random.randn(100)
self.model = LocalLinearTrendStateModel(self.data)
self.model.set_state_index(1)
def tearDown(self):
if os.path.exists("llt.pkl") and os.path.isfile("llt.pkl"):
os.remove("llt.pkl")
def test_state_dimension(self):
self.assertEqual(2, self.model.state_dimension)
def test_storage(self):
time_dimension = 20
niter = 10
self.model.allocate_space(niter, time_dimension)
self.assertEqual(self.model.sigma_level.shape, (niter, ))
self.assertEqual(self.model.sigma_slope.shape, (niter, ))
self.assertEqual(self.model.state_contribution.shape,
(niter, time_dimension))
iteration = 4
self.model.sigma_level[iteration] = 2.6
self.model.sigma_slope[iteration] = 1.9
self.model.restore_state(iteration)
self.assertAlmostEqual(self.model._state_model.sigma_level, 2.6)
self.assertAlmostEqual(self.model._state_model.sigma_slope, 1.9)
state_matrix = np.random.randn(3, time_dimension)
iteration = 7
self.model.record_state(iteration, state_matrix)
self.assertAlmostEqual(self.model.sigma_level[iteration], 2.6)
self.assertAlmostEqual(self.model.sigma_slope[iteration], 1.9)
np.testing.assert_array_equal(
self.model._state_contribution[iteration, :], state_matrix[1, :])
with open("llt.pkl", "wb") as pkl:
pickle.dump(self.model, pkl)
with open("llt.pkl", "rb") as pkl:
m2 = pickle.load(pkl)
np.testing.assert_array_equal(self.model.sigma_level, m2.sigma_level)
np.testing.assert_array_equal(self.model.sigma_slope, m2.sigma_slope)
def test_local_level(self):
model = Bsts()
model.add_state(LocalLinearTrendStateModel(self._y))
model.add_state(TrigStateModel(self._y, period=12, frequencies=[1, 2]))
model.train(data=self._y, niter=1000)
fname = "trig.pkl"
with open(fname, "wb") as pkl:
pickle.dump(model, pkl)
with open(fname, "rb") as pkl:
m2 = pickle.load(pkl)
self.assertEqual(model.time_dimension, m2.time_dimension)
self.assertIsInstance(m2, Bsts)
def test_auto_ar(self):
model = Bsts()
model.add_state(LocalLinearTrendStateModel(self._y))
model.add_state(SeasonalStateModel(self._y, nseasons=12))
model.add_state(AutoArStateModel(self._y, lags=1))
model.train(data=self._y, niter=1000)
fname = "ar.pkl"
with open(fname, "wb") as pkl:
pickle.dump(model, pkl)
with open(fname, "rb") as pkl:
m2 = pickle.load(pkl)
self.assertEqual(model.time_dimension, m2.time_dimension)
self.assertIsInstance(m2, Bsts)
def test_mcmc(self):
model = Bsts()
model.add_state(LocalLinearTrendStateModel(AirPassengers))
model.add_state(GeneralSeasonalLLT(AirPassengers, nseasons=12))
model.train(data=AirPassengers, niter=1000)
fname = "seasonal_llt.pkl"
with open(fname, "wb") as pkl:
pickle.dump(model, pkl)
with open(fname, "rb") as pkl:
m2 = pickle.load(pkl)
self.assertEqual(model.time_dimension, m2.time_dimension)
self.assertIsInstance(m2, Bsts)
fig = model.plot("comp")
fig.show()
pred = model.predict(12)
fig2, ax = pred.plot()
fig2.show()
def test_basic_structural_model(self):
model = Bsts(family="student")
y = np.log(AirPassengers)
model.add_state(LocalLinearTrendStateModel(y))
model.add_state(SeasonalStateModel(y, nseasons=12))
model.train(data=y, niter=1000)
fig = model.plot("comp")
if _show_figs:
fig.show()
horizon = 24
predictions = model.predict(horizon)
target_quantiles = [0.025] + list(np.linspace(.05, 0.95, 19)) + [0.975]
prediction_quantiles = np.quantile(predictions.distribution,
target_quantiles,
axis=0)
self.assertIsInstance(prediction_quantiles, np.ndarray)
with open("bsm.pkl", "wb") as pkl:
pickle.dump(model, pkl)
with open("bsm.pkl", "rb") as pkl:
m2 = pickle.load(pkl)
pred2 = m2.predict(horizon)
fig, ax = plt.subplots(1, 2, figsize=(10, 5))
predictions.plot(ax=ax[0], original_series=60)
pred_plot_fig, pred_plot_ax = pred2.plot(ax=ax[1], original_series=60)
if _show_figs:
fig.show()
self.assertIsInstance(pred_plot_ax, plt.Axes)
np.testing.assert_array_almost_equal(predictions.distribution,
pred2.distribution)
def setUp(self):
np.random.seed(8675309)
self.data = np.random.randn(100)
self.model = LocalLinearTrendStateModel(self.data)
self.model.set_state_index(1)