123 lines
4.1 KiB
Python
123 lines
4.1 KiB
Python
import unittest
|
|
|
|
import numpy as np
|
|
import pandas as pd
|
|
|
|
from lib import utils
|
|
from lib.utils import StandardScaler
|
|
|
|
|
|
class MyTestCase(unittest.TestCase):
|
|
def test_separate_seasonal_trend_and_residual(self):
|
|
data = np.array([
|
|
[2, 1, 2, 3, 0, 1, 2, 1, 2, 3, 4, 3, 0, 3, 4, 1]
|
|
], dtype=np.float32).T
|
|
trends = np.array([
|
|
[1, 2, 3, 2, 1, 2, 3, 2, 1, 2, 3, 2, 1, 2, 3, 2]
|
|
], dtype=np.float32).T
|
|
residual = np.array([
|
|
[1, -1, -1, 1, -1, -1, -1, -1, 1, 1, 1, 1, -1, 1, 1, -1]
|
|
], dtype=np.float32).T
|
|
df = pd.DataFrame(data)
|
|
df_trend, df_residual = utils.separate_seasonal_trend_and_residual(df, period=4, test_ratio=0, null_val=-1)
|
|
self.assertTrue(np.array_equal(df_trend.values, trends))
|
|
self.assertTrue(np.array_equal(df_residual.values, residual))
|
|
|
|
def test_get_rush_hours_bool_index(self):
|
|
index = pd.date_range('2017-02-27', '2017-03-06', freq='1min')
|
|
data = np.zeros((len(index), 3))
|
|
df = pd.DataFrame(data, index=index)
|
|
ind = utils.get_rush_hours_bool_index(df)
|
|
df = df[ind]
|
|
self.assertEqual(6 * 5 * 60, df.shape[0])
|
|
|
|
|
|
class IODataPreparationTest(unittest.TestCase):
|
|
from lib import utils
|
|
def test_generate_io_data_with_time(self):
|
|
data = np.array([
|
|
[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15],
|
|
[1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16]
|
|
], dtype=np.float32).T
|
|
df = pd.DataFrame(data, index=pd.date_range('2017-10-18', '2017-10-19 23:59', freq='3h'))
|
|
xs, ys = utils.generate_io_data_with_time(df, batch_size=2, seq_len=3, horizon=3, output_type='range', )
|
|
self.assertTupleEqual(xs.shape, (3, 2, 9))
|
|
self.assertTupleEqual(ys.shape, (3, 2, 6))
|
|
|
|
def test_generate_graph_seq2seq_io_data_with_time(self):
|
|
data = np.array([
|
|
[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15],
|
|
[1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16]
|
|
], dtype=np.float32).T
|
|
df = pd.DataFrame(data, index=pd.date_range('2017-10-18', '2017-10-19 23:59', freq='3h'))
|
|
xs, ys = utils.generate_graph_seq2seq_io_data_with_time2(df, batch_size=2, seq_len=3, horizon=3, num_nodes=2)
|
|
self.assertTupleEqual(xs.shape, (5, 2, 3, 2, 2))
|
|
self.assertTupleEqual(ys.shape, (5, 2, 3, 2, 2))
|
|
|
|
|
|
class StandardScalerTest(unittest.TestCase):
|
|
def test_transform(self):
|
|
data = np.array([
|
|
[35., 0.],
|
|
[0., 17.5],
|
|
[70., 35.]]
|
|
)
|
|
expected_result = np.array([
|
|
[0., -1.],
|
|
[-1, -0.5],
|
|
[1., 0.]]
|
|
)
|
|
scaler = StandardScaler(mean=35., std=35.)
|
|
result = scaler.transform(data)
|
|
self.assertTrue(np.array_equal(expected_result, result))
|
|
|
|
def test_transform_df(self):
|
|
df = pd.DataFrame([
|
|
[35., 0.],
|
|
[0., 17.5],
|
|
[70., 35.]]
|
|
)
|
|
expected_result = np.array([
|
|
[0., -1.],
|
|
[-1, -0.5],
|
|
[1., 0.]]
|
|
)
|
|
scaler = StandardScaler(mean=35., std=35.)
|
|
result = scaler.transform(df)
|
|
|
|
self.assertTrue(np.array_equal(expected_result, result.values))
|
|
|
|
def test_reverse_transform(self):
|
|
data = np.array([
|
|
[0., -1.],
|
|
[-1, -0.5],
|
|
[1., 0.]]
|
|
)
|
|
expected_result = np.array([
|
|
[35., 0.],
|
|
[0., 17.5],
|
|
[70., 35.]]
|
|
)
|
|
scaler = StandardScaler(mean=35., std=35.)
|
|
result = scaler.inverse_transform(data)
|
|
self.assertTrue(np.array_equal(expected_result, result))
|
|
|
|
def test_reverse_transform_df(self):
|
|
df = pd.DataFrame([
|
|
[0., -1.],
|
|
[-1, -0.5],
|
|
[1., 0.]]
|
|
)
|
|
expected_result = np.array([
|
|
[35., 0.],
|
|
[0., 17.5],
|
|
[70., 35.]]
|
|
)
|
|
scaler = StandardScaler(mean=35., std=35.)
|
|
result = scaler.inverse_transform(df)
|
|
self.assertTrue(np.array_equal(expected_result, result.values))
|
|
|
|
|
|
if __name__ == '__main__':
|
|
unittest.main()
|