def _train_kf(self,
data: torch.Tensor,
num_epochs: int = 8,
cls: Type['KalmanFilter'] = KalmanFilter):
kf = cls(measures=['y'],
processes=[
LocalLevel(id='local_level').add_measure('y'),
Season(id='day_in_week',
seasonal_period=7,
**self.config['season_spec']).add_measure('y')
])
kf.opt = LBFGS(kf.parameters())
start_datetimes = (
np.zeros(self.config['num_groups'], dtype='timedelta64') +
self.config['season_spec']['season_start'])
def closure():
kf.opt.zero_grad()
pred = kf(data, start_datetimes=start_datetimes)
loss = -pred.log_prob(data).mean()
loss.backward()
return loss
print(f"Will train for {num_epochs} epochs...")
loss = float('nan')
for i in range(num_epochs):
new_loss = kf.opt.step(closure)
print(
f"EPOCH {i}, LOSS {new_loss.item()}, DELTA {loss - new_loss.item()}"
)
loss = new_loss.item()
return kf(data, start_datetimes=start_datetimes).predictions
def name_to_proc(id: str, **kwargs) -> Process:
season_start = '2010-01-04'
if 'hour_in_day' in id:
out = FourierSeasonFixed(id=id,
seasonal_period=24, season_start=season_start, dt_unit='h',
**kwargs)
elif 'day_in_year' in id:
out = FourierSeasonFixed(id=id,
seasonal_period=24 * 364.25, season_start=season_start, dt_unit='h',
**kwargs)
elif 'local_level' in id:
out = LocalLevel(id=id, **kwargs)
elif 'local_trend' in id:
out = LocalTrend(id=id, **kwargs)
elif 'day_in_week' in id:
out = Season(id=id,
seasonal_period=7, season_duration=24,
season_start=season_start, dt_unit='h',
**kwargs)
elif 'nn_predictors' in id:
out = NN(id=id,
add_module_params_to_process=False, # so we can use a separate parameter group
model_mat_kwarg_name='predictors',
**kwargs)
elif 'predictors' in id:
out = LinearModel(id=id,
covariates=self.predictors,
model_mat_kwarg_name='predictors',
**kwargs)
else:
raise NotImplementedError(f"Unsure what process to use for `{id}`.")
return out
def simulate(num_groups: int,
num_timesteps: int,
season_spec: dict,
noise: float = 1.0) -> torch.Tensor:
# make kf:
processes = [
LocalLevel(id='local_level').add_measure('y'),
Season(id='day_in_week', seasonal_period=7, fixed=True,
**season_spec).add_measure('y'),
FourierSeasonFixed(id='day_in_month',
seasonal_period=30,
K=2,
**season_spec).add_measure('y')
]
kf = KalmanFilter(measures=['y'], processes=processes)
# make local-level less aggressive:
pcov = kf.design.process_covariance.create().data
pcov[0, 0] *= .1
kf.design.process_covariance.set(pcov)
# simulate:
start_datetimes = np.zeros(
num_groups, dtype='timedelta64') + season_spec['season_start']
with torch.no_grad():
dfb = kf.design.for_batch(num_groups=num_groups,
num_timesteps=num_timesteps,
start_datetimes=start_datetimes)
initial_state = kf.predict_initial_state(dfb)
simulated_trajectories = initial_state.simulate_trajectories(dfb)
sim_data = simulated_trajectories.sample_measurements(eps=noise)
return sim_data
def name_to_proc(id: str, **kwargs) -> Process:
season_start = '2010-01-04'
if 'hour_in_day' in id:
out = FourierSeasonFixed(id=id,
seasonal_period=24,
season_start=season_start,
dt_unit='h',
**kwargs)
elif 'day_in_year' in id:
out = FourierSeasonFixed(id=id,
seasonal_period=24 * 364.25,
season_start=season_start,
dt_unit='h',
**kwargs)
elif 'local_level' in id:
out = LocalLevel(id=id, **kwargs)
elif 'local_trend' in id:
out = LocalTrend(id=id, **kwargs)
elif 'day_in_week' in id:
out = Season(id=id,
seasonal_period=7,
season_duration=24,
season_start=season_start,
dt_unit='h',
**kwargs)
else:
raise NotImplementedError(f"Unsure what process to use for `{id}`.")
return out
def test_discrete_seasons(self):
# test seasons without durations
season = Season(id='day_of_week',
seasonal_period=7,
season_duration=1,
season_start='2018-01-01',
dt_unit='D')
season.add_measure('measure')
# need to include start_datetimes since included above
with self.assertRaises(ValueError) as cm:
season.for_batch(1, 1)
self.assertEqual(
cm.exception.args[0],
'Must pass `start_datetimes` to process `day_of_week`.')
design = Design(processes=[season], measures=['measure'])
process_kwargs = {
'day_of_week': {
'start_datetimes': array([datetime64('2018-01-01')])
}
}
batch_season = design.for_batch(1, 1, process_kwargs=process_kwargs)
# test transitions manually:
state_mean = torch.arange(0.0, 7.0)[:, None]
state_mean[0] = -state_mean[1:].sum()
for i in range(10):
state_mean_last = state_mean
state_mean = torch.mm(batch_season.F(0)[0], state_mean)
self.assertTrue((state_mean[1:] == state_mean_last[:-1]).all())
self.assertListEqual(
batch_season.H(0)[0].tolist(),
[[1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]])
def _simulate(num_groups: int, num_timesteps: int, dt_unit: str, noise: float = 1.0) -> torch.Tensor:
# make kf:
processes = [
LocalLevel(id='local_level').add_measure('y'),
Season(id='day_in_week', seasonal_period=7, fixed=True, dt_unit=dt_unit).add_measure('y'),
FourierSeason(id='day_in_year', seasonal_period=365.25, K=2, fixed=True, dt_unit=dt_unit).add_measure('y')
]
kf = KalmanFilter(measures=['y'], processes=processes)
# simulate:
start_datetimes = np.zeros(num_groups, dtype='timedelta64') + DEFAULT_START_DT
with torch.no_grad():
dfb = kf.design.for_batch(num_groups=num_groups, num_timesteps=num_timesteps, start_datetimes=start_datetimes)
initial_state = kf._predict_initial_state(dfb)
simulated_trajectories = initial_state.simulate_trajectories(dfb)
sim_data = simulated_trajectories.sample_measurements(eps=noise)
return sim_data
def test_discrete_seasons(self):
# test seasons without durations
season = Season(
id='day_of_week', seasonal_period=7, season_duration=1, dt_unit='D'
)
season.add_measure('measure')
# need to include start_datetimes since included above
with self.assertRaises(TypeError) as cm:
season.for_batch(1, 1)
self.assertIn('Missing argument `start_datetimes`', cm.exception.args[0])
design = Design(processes=[season], measures=['measure'])
batch_season = design.for_batch(1, 1, start_datetimes=np.array([np.datetime64('2018-01-01')]))
# test transitions manually:
state_mean = torch.arange(0.0, 7.0)[:, None]
state_mean[0] = -state_mean[1:].sum()
for i in range(10):
state_mean_last = state_mean
state_mean = torch.mm(batch_season.F(0)[0], state_mean)
self.assertTrue((state_mean[1:] == state_mean_last[:-1]).all())
self.assertListEqual(batch_season.H(0)[0].tolist(), [[1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]])