def model(self, zero_data, covariates):
duration,data_dim = zero_data.shape
feature_dim = covariates.size(-1)
drift_stability = pyro.sample("drift_stability", dist.Uniform(1, 2))
drift_scale = pyro.sample("drift_scale", dist.LogNormal(-20, 5))
with pyro.plate("item", data_dim, dim=-2):
bias = pyro.sample("bias", dist.Normal(5.0, 10.0))
with pyro.plate('features',feature_dim,dim=-1):
weight = pyro.sample("weight", dist.Normal(0.0, 5))
# We'll sample a time-global scale parameter outside the time plate,
# then time-local iid noise inside the time plate.
with self.time_plate:
# We combine two different reparameterizers: the inner SymmetricStableReparam
# is needed for the Stable site, and the outer LocScaleReparam is optional but
# appears to improve inference.
with poutine.reparam(config={"drift": LocScaleReparam()}):
with poutine.reparam(config={"drift": SymmetricStableReparam()}):
drift = pyro.sample("drift",
dist.Stable(drift_stability, 0, drift_scale))
motion = drift.cumsum(-1) # A Brownian motion.
# The prediction now includes three terms.
regression = torch.matmul(covariates, weight[...,None])
prediction = motion + bias + regression.sum(axis=-1)
prediction = prediction.unsqueeze(-1).transpose(-1, -3)
# Finally we can construct a noise distribution.
# We will share parameters across all time series.
obs_scale = pyro.sample("obs_scale", dist.LogNormal(-5, 5))
noise_dist = dist.Normal(loc=0.0, scale=obs_scale.unsqueeze(-1))
self.predict(noise_dist, prediction)
def test_symmetric_stable(shape):
stability = torch.empty(shape).uniform_(1.6, 1.9).requires_grad_()
scale = torch.empty(shape).uniform_(0.5, 1.0).requires_grad_()
loc = torch.empty(shape).uniform_(-1.0, 1.0).requires_grad_()
params = [stability, scale, loc]
def model():
with pyro.plate_stack("plates", shape):
with pyro.plate("particles", 200000):
return pyro.sample("x", dist.Stable(stability, 0, scale, loc))
value = model()
expected_moments = get_moments(value)
reparam_model = poutine.reparam(model, {"x": SymmetricStableReparam()})
trace = poutine.trace(reparam_model).get_trace()
assert isinstance(trace.nodes["x"]["fn"], dist.Normal)
trace.compute_log_prob() # smoke test only
value = trace.nodes["x"]["value"]
actual_moments = get_moments(value)
assert_close(actual_moments, expected_moments, atol=0.05)
for actual_m, expected_m in zip(actual_moments, expected_moments):
expected_grads = grad(expected_m.sum(), params, retain_graph=True)
actual_grads = grad(actual_m.sum(), params, retain_graph=True)
assert_close(actual_grads[0], expected_grads[0], atol=0.2)
assert_close(actual_grads[1], expected_grads[1], atol=0.1)
assert_close(actual_grads[2], expected_grads[2], atol=0.1)
def test_stable_hmm_shape(skew, batch_shape, duration, hidden_dim, obs_dim):
stability = dist.Uniform(0.5, 2).sample(batch_shape)
init_dist = random_stable(batch_shape + (hidden_dim, ),
stability.unsqueeze(-1),
skew=skew).to_event(1)
trans_mat = torch.randn(batch_shape + (duration, hidden_dim, hidden_dim))
trans_dist = random_stable(batch_shape + (duration, hidden_dim),
stability.unsqueeze(-1).unsqueeze(-1),
skew=skew).to_event(1)
obs_mat = torch.randn(batch_shape + (duration, hidden_dim, obs_dim))
obs_dist = random_stable(batch_shape + (duration, obs_dim),
stability.unsqueeze(-1).unsqueeze(-1),
skew=skew).to_event(1)
hmm = dist.LinearHMM(init_dist, trans_mat, trans_dist, obs_mat, obs_dist)
def model(data=None):
with pyro.plate_stack("plates", batch_shape):
return pyro.sample("x", hmm, obs=data)
data = model()
rep = SymmetricStableReparam() if skew == 0 else StableReparam()
with poutine.trace() as tr:
with poutine.reparam(config={"x": LinearHMMReparam(rep, rep, rep)}):
model(data)
assert isinstance(tr.trace.nodes["x"]["fn"], dist.GaussianHMM)
tr.trace.compute_log_prob() # smoke test only
def test_stable_hmm_distribution(stability, skew, duration, hidden_dim,
obs_dim):
init_dist = random_stable((hidden_dim, ), stability, skew=skew).to_event(1)
trans_mat = torch.randn(duration, hidden_dim, hidden_dim)
trans_dist = random_stable((duration, hidden_dim), stability,
skew=skew).to_event(1)
obs_mat = torch.randn(duration, hidden_dim, obs_dim)
obs_dist = random_stable((duration, obs_dim), stability,
skew=skew).to_event(1)
hmm = dist.LinearHMM(init_dist, trans_mat, trans_dist, obs_mat, obs_dist)
num_samples = 200000
expected_samples = hmm.sample([num_samples
]).reshape(num_samples, duration * obs_dim)
expected_loc, expected_scale, expected_corr = get_hmm_moments(
expected_samples)
rep = SymmetricStableReparam() if skew == 0 else StableReparam()
with pyro.plate("samples", num_samples):
with poutine.reparam(config={"x": LinearHMMReparam(rep, rep, rep)}):
actual_samples = pyro.sample("x",
hmm).reshape(num_samples,
duration * obs_dim)
actual_loc, actual_scale, actual_corr = get_hmm_moments(actual_samples)
assert_close(actual_loc, expected_loc, atol=0.05, rtol=0.05)
assert_close(actual_scale, expected_scale, atol=0.05, rtol=0.05)
assert_close(actual_corr, expected_corr, atol=0.01)
def test_init_shape(skew, batch_shape, duration, hidden_dim, obs_dim):
stability = dist.Uniform(0.5, 2).sample(batch_shape)
init_dist = random_stable(batch_shape + (hidden_dim, ),
stability.unsqueeze(-1),
skew=skew).to_event(1)
trans_mat = torch.randn(batch_shape + (duration, hidden_dim, hidden_dim))
trans_dist = random_stable(
batch_shape + (duration, hidden_dim),
stability.unsqueeze(-1).unsqueeze(-1),
skew=skew,
).to_event(1)
obs_mat = torch.randn(batch_shape + (duration, hidden_dim, obs_dim))
obs_dist = random_stable(
batch_shape + (duration, obs_dim),
stability.unsqueeze(-1).unsqueeze(-1),
skew=skew,
).to_event(1)
hmm = dist.LinearHMM(init_dist,
trans_mat,
trans_dist,
obs_mat,
obs_dist,
duration=duration)
assert hmm.batch_shape == batch_shape
assert hmm.event_shape == (duration, obs_dim)
def model():
with pyro.plate_stack("plates", batch_shape):
return pyro.sample("x", hmm)
rep = SymmetricStableReparam() if skew == 0 else StableReparam()
check_init_reparam(model, LinearHMMReparam(rep, rep, rep))
def model(self, zero_data, covariates):
duration, data_dim = zero_data.shape
# Let's model each time series as a Levy stable process, and share process parameters
# across time series. To do that in Pyro, we'll declare the shared random variables
# outside of the "origin" plate:
drift_stability = pyro.sample("drift_stability", dist.Uniform(1, 2))
drift_scale = pyro.sample("drift_scale", dist.LogNormal(-20, 5))
with pyro.plate("origin", data_dim, dim=-2):
# Now inside of the origin plate we sample drift and seasonal components.
# All the time series inside the "origin" plate are independent,
# given the drift parameters above.
with self.time_plate:
# We combine two different reparameterizers: the inner SymmetricStableReparam
# is needed for the Stable site, and the outer LocScaleReparam is optional but
# appears to improve inference.
with poutine.reparam(config={"drift": LocScaleReparam()}):
with poutine.reparam(config={"drift": SymmetricStableReparam()}):
drift = pyro.sample("drift",
dist.Stable(drift_stability, 0, drift_scale))
with pyro.plate("hour_of_week", 24 * 7, dim=-1):
seasonal = pyro.sample("seasonal", dist.Normal(0, 5))
# Now outside of the time plate we can perform time-dependent operations like
# integrating over time. This allows us to create a motion with slow drift.
seasonal = periodic_repeat(seasonal, duration, dim=-1)
motion = drift.cumsum(dim=-1) # A Levy stable motion to model shocks.
prediction = motion + seasonal
# Next we do some reshaping. Pyro's forecasting framework assumes all data is
# multivariate of shape (duration, data_dim), but the above code uses an "origins"
# plate that is left of the time_plate. Our prediction starts off with shape
assert prediction.shape[-2:] == (data_dim, duration)
# We need to swap those dimensions but keep the -2 dimension intact, in case Pyro
# adds sample dimensions to the left of that.
prediction = prediction.unsqueeze(-1).transpose(-1, -3)
assert prediction.shape[-3:] == (1, duration, data_dim), prediction.shape
# Finally we can construct a noise distribution.
# We will share parameters across all time series.
obs_scale = pyro.sample("obs_scale", dist.LogNormal(-5, 5))
noise_dist = dist.Normal(0, obs_scale.unsqueeze(-1))
self.predict(noise_dist, prediction)